Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2024 HuggingFace Inc. team. All rights reserved.3#4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9# http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16 17import math18from functools import cached_property19from typing import Optional, Union20 21import torch22import torch.nn as nn23import torch.nn.functional as F24 25from ...cache_utils import Cache26from ...generation import GenerationMixin27from ...modeling_outputs import CausalLMOutputWithPast28from ...modeling_utils import PreTrainedModel29from ...processing_utils import Unpack30from ...utils import auto_docstring, can_return_tuple, logging31from ...utils.deprecation import deprecate_kwarg32from ..chameleon.modeling_chameleon import (33 ChameleonPreTrainedModel,34 ChameleonVQVAEEncoderConvDownsample,35)36from ..llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaForCausalLM, LlamaModel, TransformersKwargs37from ..siglip.modeling_siglip import SiglipAttention38from .configuration_emu3 import Emu3Config, Emu3TextConfig, Emu3VQVAEConfig39 40 41logger = logging.get_logger(__name__)42 43 44class Emu3Attention(LlamaAttention):45 pass46 47 48# Has extra dropout which no other model in the library has49class Emu3DecoderLayer(LlamaDecoderLayer):50 def __init__(self, config: Emu3Config, layer_idx: int):51 super().__init__(config, layer_idx)52 self.dropout = nn.Dropout(config.attention_dropout)53 54 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")55 def forward(56 self,57 hidden_states: torch.Tensor,58 attention_mask: Optional[torch.Tensor] = None,59 position_ids: Optional[torch.LongTensor] = None,60 past_key_values: Optional[Cache] = None,61 use_cache: Optional[bool] = False,62 cache_position: Optional[torch.LongTensor] = None,63 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,64 **kwargs: Unpack[TransformersKwargs],65 ) -> torch.Tensor:66 residual = hidden_states67 hidden_states = self.input_layernorm(hidden_states)68 69 hidden_states, _ = self.self_attn(70 hidden_states=hidden_states,71 attention_mask=attention_mask,72 position_ids=position_ids,73 past_key_values=past_key_values,74 use_cache=use_cache,75 cache_position=cache_position,76 position_embeddings=position_embeddings,77 **kwargs,78 )79 hidden_states = residual + self.dropout(hidden_states)80 81 residual = hidden_states82 hidden_states = self.post_attention_layernorm(hidden_states)83 hidden_states = self.mlp(hidden_states)84 hidden_states = residual + self.dropout(hidden_states)85 return hidden_states86 87 88class Emu3VQVAEVectorQuantizer(nn.Module):89 """90 A module for vector quantization using learned embedding vectors.91 92 This module implements the quantization process similar to te one described in93 the VQ-VAE (Vector Quantized Variational AutoEncoder) paper. It quantizes continuous94 input vectors into discrete codebook vectors, which are learned during training.95 Current implementation improves over previous ones by avoiding costly matrix multiplications96 and allowing for post-hoc remapping of indices.97 """98 99 def __init__(self, config: Emu3VQVAEConfig):100 super().__init__()101 self.embedding = nn.Embedding(config.codebook_size, config.embed_dim)102 self.embedding.weight.data.uniform_(-1.0 / config.codebook_size, 1.0 / config.codebook_size)103 104 def forward(self, hidden_state: torch.Tensor):105 batch_size, temporal, channels, height, width = hidden_state.shape106 hidden_state = hidden_state.permute(0, 1, 3, 4, 2).contiguous()107 hidden_state_flattened = hidden_state.view(-1, channels)108 109 # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z110 hidden_state_sum = torch.sum(hidden_state_flattened**2, dim=1, keepdim=True)111 embedding_sum = torch.sum(self.embedding.weight**2, dim=1)112 113 # "bd,dn->bn",114 distances = 2 * torch.matmul(hidden_state_flattened, self.embedding.weight.transpose(0, 1))115 distances = hidden_state_sum + embedding_sum - distances116 117 min_encoding_indices = torch.argmin(distances, dim=1)118 min_encoding_indices = min_encoding_indices.view(batch_size, temporal, height, width)119 return min_encoding_indices120 121 122class Emu3VQVAEEncoderConvDownsample(ChameleonVQVAEEncoderConvDownsample):123 pass124 125 126class Emu3VQVAEEncoderConvUpsample(nn.Module):127 def __init__(self, in_channels):128 super().__init__()129 self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)130 131 def forward(self, hidden_states):132 hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")133 hidden_states = self.conv(hidden_states)134 return hidden_states135 136 137class Emu3VQVAEConv3d(nn.Module):138 def __init__(139 self,140 in_channel: int,141 out_channel: int,142 kernel_size: tuple[int],143 stride: tuple[int],144 ):145 super().__init__()146 147 padding_sizes = [one_kernel - one_stride for one_kernel, one_stride in zip(kernel_size[1:], stride[1:])]148 self.padding = ()149 for pad_size in padding_sizes[::-1]:150 self.padding += (pad_size // 2 + pad_size % 2, pad_size // 2)151 self.padding += (2, 0)152 153 self.conv = nn.Conv3d(154 in_channel,155 out_channel,156 kernel_size,157 stride=stride,158 )159 160 def forward(self, hidden_states: torch.Tensor):161 hidden_states = F.pad(hidden_states, self.padding)162 hidden_states = self.conv(hidden_states)163 return hidden_states164 165 166class Emu3VQVAESpatialNorm(nn.Module):167 def __init__(168 self,169 in_channels: int,170 out_channels: int,171 ):172 super().__init__()173 self.norm_layer = nn.GroupNorm(174 num_channels=out_channels,175 num_groups=32,176 eps=1e-6,177 affine=True,178 )179 180 self.conv_y = nn.Conv2d(181 in_channels,182 out_channels,183 kernel_size=1,184 stride=1,185 padding=0,186 )187 self.conv_b = nn.Conv2d(188 in_channels,189 out_channels,190 kernel_size=1,191 stride=1,192 padding=0,193 )194 195 def forward(self, hidden_states: torch.Tensor, quant_states: torch.Tensor):196 quant_states = F.interpolate(quant_states, size=hidden_states.shape[-2:], mode="nearest")197 hidden_states = self.norm_layer(hidden_states)198 hidden_states = hidden_states * self.conv_y(quant_states) + self.conv_b(quant_states)199 return hidden_states200 201 202class Emu3VQVAETemporalUpsample(nn.Module):203 def __init__(204 self,205 in_channel: int,206 out_channel: int,207 ):208 super().__init__()209 self.conv = Emu3VQVAEConv3d(210 in_channel,211 out_channel,212 kernel_size=(3, 3, 3),213 stride=(1, 1, 1),214 )215 216 def forward(self, hidden_states: torch.Tensor):217 batch_size, channels, temporal, height, width = hidden_states.shape218 hidden_states = hidden_states.permute(0, 1, 3, 4, 2).contiguous().view(batch_size, -1, temporal)219 hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")220 hidden_states = hidden_states.view(batch_size, channels, height, width, -1).permute(0, 1, 4, 2, 3).contiguous()221 hidden_states = self.conv(hidden_states)222 return hidden_states223 224 225class Emu3VQVAETemporalDownsample(nn.Module):226 def __init__(227 self,228 in_channel: int,229 out_channel: int,230 ):231 super().__init__()232 self.conv = Emu3VQVAEConv3d(233 in_channel,234 out_channel,235 kernel_size=(4, 3, 3),236 stride=(2, 1, 1),237 )238 239 def forward(self, hidden_states: torch.Tensor):240 hidden_states = self.conv(hidden_states)241 return hidden_states242 243 244class Emu3VQVAETemporalResnetBlock(nn.Module):245 def __init__(246 self,247 in_channels,248 out_channels=None,249 ):250 super().__init__()251 self.in_channels = in_channels252 self.out_channels = in_channels if out_channels is None else out_channels253 254 self.norm1 = nn.BatchNorm3d(in_channels)255 self.conv1 = Emu3VQVAEConv3d(256 in_channels,257 out_channels,258 kernel_size=(3, 3, 3),259 stride=(1, 1, 1),260 )261 self.norm2 = nn.BatchNorm3d(out_channels)262 self.conv2 = Emu3VQVAEConv3d(263 out_channels,264 out_channels,265 kernel_size=(3, 3, 3),266 stride=(1, 1, 1),267 )268 if self.in_channels != self.out_channels:269 self.nin_shortcut = nn.Conv3d(270 in_channels,271 out_channels,272 kernel_size=1,273 stride=1,274 padding=0,275 )276 277 def forward(self, hidden_states):278 residual = hidden_states279 hidden_states = self.norm1(hidden_states)280 hidden_states *= torch.sigmoid(hidden_states)281 hidden_states = self.conv1(hidden_states)282 283 hidden_states = self.norm2(hidden_states)284 hidden_states *= torch.sigmoid(hidden_states)285 hidden_states = self.conv2(hidden_states)286 287 if self.in_channels != self.out_channels:288 residual = self.nin_shortcut(residual)289 290 return residual + hidden_states291 292 293class Emu3VQVAEResnetBlock(nn.Module):294 def __init__(295 self,296 in_channels: int,297 out_channels: Optional[int] = None,298 quant_channels: Optional[int] = None,299 ):300 super().__init__()301 self.in_channels = in_channels302 out_channels = in_channels if out_channels is None else out_channels303 self.out_channels = out_channels304 self.quant_channels = quant_channels305 306 if quant_channels is None:307 self.norm1 = nn.GroupNorm(num_channels=in_channels, num_groups=32, eps=1e-6, affine=True)308 self.norm2 = nn.GroupNorm(num_channels=out_channels, num_groups=32, eps=1e-6, affine=True)309 else:310 self.norm1 = Emu3VQVAESpatialNorm(quant_channels, in_channels)311 self.norm2 = Emu3VQVAESpatialNorm(quant_channels, out_channels)312 313 self.conv1 = nn.Conv2d(314 in_channels,315 out_channels,316 kernel_size=3,317 stride=1,318 padding=1,319 )320 321 self.conv2 = nn.Conv2d(322 out_channels,323 out_channels,324 kernel_size=3,325 stride=1,326 padding=1,327 )328 329 if self.in_channels != self.out_channels:330 self.nin_shortcut = nn.Conv2d(331 in_channels,332 out_channels,333 kernel_size=1,334 stride=1,335 padding=0,336 )337 338 def forward(self, hidden_states: torch.Tensor, quant_channels: Optional[torch.Tensor] = None):339 norm_args = () if self.quant_channels is None else (quant_channels,)340 341 residual = hidden_states342 hidden_states = self.norm1(hidden_states, *norm_args)343 hidden_states *= torch.sigmoid(hidden_states)344 hidden_states = self.conv1(hidden_states)345 346 hidden_states = self.norm2(hidden_states, *norm_args)347 hidden_states *= torch.sigmoid(hidden_states)348 hidden_states = self.conv2(hidden_states)349 350 if self.in_channels != self.out_channels:351 residual = self.nin_shortcut(residual)352 353 return residual + hidden_states354 355 356class Emu3VQVAEAttentionBlock(SiglipAttention):357 def __init__(self, config: Emu3VQVAEConfig):358 super().__init__(config)359 360 # for compatibility with the attention interface361 self.num_key_value_groups = 1362 363 364class Emu3VQVAEGroupNorm(nn.GroupNorm):365 """366 Same as the torch GroupNorm with the only difference that this ones accepts367 an optional kwarg `quant_states` which is not used. This class makes it easier to368 use SpatialNorm or GroupNorm without conditionals369 """370 371 def __init__(self, **kwargs):372 super().__init__(**kwargs)373 374 def forward(self, input, quant_states=None):375 return F.group_norm(input, self.num_groups, self.weight, self.bias, self.eps)376 377 378class Emu3VQVAEMiddleBlock(nn.Module):379 def __init__(self, config, in_channels, quant_channels=None):380 super().__init__()381 382 self.block_1 = Emu3VQVAEResnetBlock(383 in_channels=in_channels,384 out_channels=in_channels,385 quant_channels=quant_channels,386 )387 self.attn_1 = Emu3VQVAEAttentionBlock(config)388 if quant_channels is None:389 self.attn_norm = Emu3VQVAEGroupNorm(num_channels=in_channels, num_groups=32, eps=1e-6, affine=True)390 else:391 self.attn_norm = Emu3VQVAESpatialNorm(quant_channels, in_channels)392 393 self.block_2 = Emu3VQVAEResnetBlock(394 in_channels=in_channels,395 out_channels=in_channels,396 quant_channels=quant_channels,397 )398 399 def forward(self, hidden_states: torch.FloatTensor, quant_states: Optional[torch.FloatTensor] = None):400 hidden_states = self.block_1(hidden_states, quant_states)401 residual = hidden_states402 hidden_states = self.attn_norm(hidden_states, quant_states)403 batch_size, channels, height, width = hidden_states.shape404 hidden_states = hidden_states.view(batch_size, channels, height * width).transpose(1, 2)405 hidden_states = self.attn_1(hidden_states)[0]406 hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)407 hidden_states = residual + hidden_states408 hidden_states = self.block_2(hidden_states, quant_states)409 return hidden_states410 411 412class Emu3VQVAEDownBlock(nn.Module):413 def __init__(self, config):414 super().__init__()415 416 self.num_resolutions = len(config.channel_multiplier)417 self.num_res_blocks = config.num_res_blocks418 base_channels = config.base_channels419 channel_multiplier = config.channel_multiplier420 421 in_channel_multiplier = (1,) + tuple(channel_multiplier)422 self.in_channel_multiplier = in_channel_multiplier423 self.down = nn.ModuleList()424 for i_level in range(self.num_resolutions):425 block = nn.ModuleList()426 attn = nn.ModuleList()427 attn_norms = nn.ModuleList()428 block_in = base_channels * in_channel_multiplier[i_level]429 block_out = base_channels * channel_multiplier[i_level]430 for i_block in range(self.num_res_blocks):431 block.append(432 Emu3VQVAEResnetBlock(433 in_channels=block_in,434 out_channels=block_out,435 )436 )437 block_in = block_out438 if config.attn_resolutions is not None and i_level in config.attn_resolutions:439 attn.append(Emu3VQVAEAttentionBlock(config))440 attn_norms.append(nn.GroupNorm(num_channels=block_in, num_groups=32, eps=1e-6, affine=True))441 442 down = nn.Module()443 down.block = block444 down.attn = attn445 down.attn_norms = attn_norms446 if i_level != self.num_resolutions - 1:447 down.downsample = Emu3VQVAEEncoderConvDownsample(block_in)448 self.down.append(down)449 450 def forward(self, hidden_states: torch.FloatTensor):451 for i_level, blocks in enumerate(self.down):452 for i_block in range(self.num_res_blocks):453 hidden_states = blocks.block[i_block](hidden_states)454 if len(blocks.attn) > 0:455 residual = hidden_states456 hidden_states = blocks.attn_norms[i_block](hidden_states)457 458 batch_size, channels, height, width = hidden_states.shape459 hidden_states = hidden_states.view(batch_size, channels, height * width).transpose(1, 2)460 hidden_states = blocks.attn[i_block](hidden_states)[0]461 462 hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)463 hidden_states = residual + hidden_states464 465 if i_level != self.num_resolutions - 1:466 hidden_states = blocks.downsample(hidden_states)467 468 return hidden_states469 470 471class Emu3VQVAEUpBlock(nn.Module):472 def __init__(self, config):473 super().__init__()474 475 self.num_resolutions = len(config.channel_multiplier)476 self.num_res_blocks = config.num_res_blocks477 478 quant_channels = config.embed_dim479 block_in = config.base_channels * config.channel_multiplier[-1]480 481 self.up = nn.ModuleList()482 for i_level in reversed(range(self.num_resolutions)):483 block = nn.ModuleList()484 attn = nn.ModuleList()485 attn_norms = nn.ModuleList()486 block_out = config.base_channels * config.channel_multiplier[i_level]487 for i_block in range(self.num_res_blocks + 1):488 block.append(489 Emu3VQVAEResnetBlock(490 in_channels=block_in,491 out_channels=block_out,492 quant_channels=quant_channels,493 )494 )495 block_in = block_out496 if i_level in config.attn_resolutions:497 attn.append(Emu3VQVAEAttentionBlock(config))498 attn_norms.append(Emu3VQVAESpatialNorm(quant_channels, block_in))499 500 up = nn.Module()501 up.block = block502 up.attn = attn503 up.attn_norms = attn_norms504 if i_level != 0:505 up.upsample = Emu3VQVAEEncoderConvUpsample(block_in)506 507 self.up.insert(0, up)508 509 def forward(self, hidden_states: torch.FloatTensor, quant_states: torch.FloatTensor):510 for i_level, blocks in enumerate(self.up[::-1]):511 for i_block in range(self.num_res_blocks + 1):512 hidden_states = blocks.block[i_block](hidden_states, quant_states)513 if len(blocks.attn) > 0:514 residual = hidden_states515 hidden_states = blocks.attn_norms[i_block](hidden_states, quant_states)516 517 batch_size, channels, height, width = hidden_states.shape518 hidden_states = hidden_states.view(batch_size, channels, height * width).transpose(1, 2)519 hidden_states = blocks.attn[i_block](hidden_states)[0]520 521 hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)522 hidden_states = residual + hidden_states523 if i_level != len(self.up) - 1:524 hidden_states = blocks.upsample(hidden_states)525 526 return hidden_states527 528 529class Emu3VQVAEEncoder(nn.Module):530 def __init__(self, config):531 super().__init__()532 533 base_channels = config.base_channels534 in_channels = config.in_channels535 double_latent = config.double_latent536 latent_channels = config.latent_channels537 channel_multiplier = config.channel_multiplier538 out_channels = 2 * latent_channels if double_latent else latent_channels539 block_in = base_channels * channel_multiplier[-1]540 541 self.conv_in = torch.nn.Conv2d(in_channels, base_channels, kernel_size=3, stride=1, padding=1)542 self.down_block = Emu3VQVAEDownBlock(config)543 self.middle_block = Emu3VQVAEMiddleBlock(config, block_in)544 545 self.norm_out = torch.nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)546 self.conv_out = torch.nn.Conv2d(547 block_in,548 out_channels,549 kernel_size=3,550 stride=1,551 padding=1,552 )553 554 temporal_down_blocks = int(math.log2(config.temporal_downsample_factor))555 self.time_conv = nn.ModuleList()556 self.time_res_stack = nn.ModuleList()557 558 for i in range(temporal_down_blocks):559 conv = Emu3VQVAETemporalDownsample(out_channels, out_channels)560 self.time_conv.append(conv)561 562 for _ in range(config.num_res_blocks):563 time_res_conv = Emu3VQVAETemporalResnetBlock(564 in_channels=out_channels,565 out_channels=out_channels,566 )567 self.time_res_stack.append(time_res_conv)568 569 def forward(self, pixel_values: torch.LongTensor):570 temporal_dim = pixel_values.shape[1]571 pixel_values = pixel_values.reshape(-1, *pixel_values.shape[2:])572 573 # downsampling & middle574 hidden_states = self.conv_in(pixel_values)575 hidden_states = self.down_block(hidden_states)576 hidden_states = self.middle_block(hidden_states)577 578 # end579 hidden_states = self.norm_out(hidden_states)580 hidden_states *= torch.sigmoid(hidden_states)581 hidden_states = self.conv_out(hidden_states)582 583 hidden_states = hidden_states.reshape(-1, temporal_dim, *hidden_states.shape[1:])584 hidden_states = hidden_states.permute(0, 2, 1, 3, 4)585 586 # temporal convs587 for conv in self.time_conv:588 hidden_states = conv(hidden_states)589 hidden_states *= torch.sigmoid(hidden_states)590 591 for layer in self.time_res_stack:592 hidden_states = layer(hidden_states)593 594 hidden_states = hidden_states.permute(0, 2, 1, 3, 4)595 596 return hidden_states597 598 599class Emu3VQVAEDecoder(nn.Module):600 def __init__(self, config: Emu3VQVAEConfig):601 super().__init__()602 603 quant_channels = config.embed_dim604 block_in = config.base_channels * config.channel_multiplier[-1]605 self.time_res_stack = nn.ModuleList()606 for _ in range(config.num_res_blocks):607 time_res_conv = Emu3VQVAETemporalResnetBlock(608 in_channels=config.latent_channels, out_channels=config.latent_channels609 )610 self.time_res_stack.append(time_res_conv)611 612 temp_upsample_block_num = int(math.log2(config.temporal_downsample_factor))613 self.time_conv = nn.ModuleList()614 for i in range(temp_upsample_block_num):615 conv = Emu3VQVAETemporalUpsample(config.latent_channels, config.latent_channels)616 self.time_conv.append(conv)617 618 self.conv_in = nn.Conv2d(619 config.latent_channels,620 block_in,621 kernel_size=3,622 stride=1,623 padding=1,624 )625 626 self.middle_block = Emu3VQVAEMiddleBlock(config, block_in, quant_channels=quant_channels)627 self.up_block = Emu3VQVAEUpBlock(config)628 629 block_in = config.base_channels * config.channel_multiplier[0]630 self.norm_out = Emu3VQVAESpatialNorm(quant_channels, block_in)631 self.conv_out = nn.Conv2d(632 block_in,633 config.out_channels,634 kernel_size=3,635 stride=1,636 padding=1,637 )638 639 def forward(self, hidden_states: torch.Tensor, quant_states: torch.Tensor):640 hidden_quant_states = torch.cat((hidden_states, quant_states), dim=0)641 hidden_quant_states = hidden_quant_states.permute(0, 2, 1, 3, 4)642 643 # temporal convs644 for layer in self.time_res_stack:645 hidden_quant_states = layer(hidden_quant_states)646 647 for layer in self.time_conv:648 hidden_quant_states = layer(hidden_quant_states)649 hidden_quant_states *= torch.sigmoid(hidden_quant_states)650 651 hidden_quant_states = hidden_quant_states.permute(0, 2, 1, 3, 4)652 hidden_states, quant_states = torch.chunk(hidden_quant_states, 2, dim=0)653 hidden_states = hidden_states.reshape(-1, *hidden_states.shape[2:])654 quant_states = quant_states.reshape(-1, *quant_states.shape[2:])655 656 hidden_states = self.conv_in(hidden_states)657 658 # middle & upsampling659 hidden_states = self.middle_block(hidden_states, quant_states)660 hidden_states = self.up_block(hidden_states, quant_states)661 662 hidden_states = self.norm_out(hidden_states, quant_states)663 hidden_states *= torch.sigmoid(hidden_states)664 hidden_states = self.conv_out(hidden_states)665 666 return hidden_states667 668 669@auto_docstring(670 custom_intro="""671 The VQ-VAE model used in Emu3 for encoding/decoding images into discrete tokens.672 This model follows the "Make-a-scene: Scene-based text-to-image generation with human priors" paper from673 [ Oran Gafni, Adam Polyak, Oron Ashual, Shelly Sheynin, Devi Parikh, and Yaniv674 Taigman](https://huggingface.co/papers/2203.13131).675 """676)677class Emu3VQVAE(PreTrainedModel):678 config: Emu3VQVAEConfig679 base_model_prefix = "emuvideovq"680 main_input_name = "pixel_values"681 _supports_sdpa = True682 _supports_flash_attn = True683 _supports_flex_attn = True684 _supports_attention_backend = True685 _no_split_modules = [686 "Emu3VQVAETemporalResnetBlock",687 "Emu3VQVAEAttentionBlock",688 "Emu3VQVAEResnetBlock",689 "Emu3VQVAEVectorQuantizer",690 ]691 692 def _init_weights(self, module):693 if isinstance(module, (nn.Conv2d, nn.Conv3d)):694 nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")695 if module.bias is not None:696 fan_in, _ = nn.init._calculate_fan_in_and_fan_out(module.weight)697 bound = 1 / math.sqrt(fan_in)698 nn.init.uniform_(module.bias, -bound, bound)699 elif isinstance(module, nn.Linear):700 nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5))701 if module.bias is not None:702 fan_in, _ = nn.init._calculate_fan_in_and_fan_out(module.weight)703 bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0704 nn.init.uniform_(module.bias, -bound, bound)705 elif isinstance(module, (nn.BatchNorm2d, nn.BatchNorm3d, nn.GroupNorm)):706 nn.init.constant_(module.weight, 1.0)707 nn.init.constant_(module.bias, 0.0)708 elif isinstance(module, nn.Embedding):709 module.weight.data.normal_()710 if module.padding_idx is not None:711 module.weight.data[module.padding_idx].zero_()712 713 def __init__(self, config: Emu3VQVAEConfig):714 super().__init__(config)715 716 self.config = config717 718 self.encoder = Emu3VQVAEEncoder(config)719 self.decoder = Emu3VQVAEDecoder(config)720 self.quantize = Emu3VQVAEVectorQuantizer(config)721 self.vision_spatial_factor = 2 ** (len(config.channel_multiplier) - 1)722 723 self.quant_conv = Emu3VQVAEConv3d(724 config.latent_channels, config.embed_dim, kernel_size=(3, 1, 1), stride=(1, 1, 1)725 )726 self.post_quant_conv = Emu3VQVAEConv3d(727 config.embed_dim, config.latent_channels, kernel_size=(3, 1, 1), stride=(1, 1, 1)728 )729 self.spatial_scale_factor = 2 ** (len(config.channel_multiplier) - 1)730 self.eval() # Emu3's VQ model is frozen731 732 self.post_init()733 734 def encode(self, pixel_values: torch.Tensor, image_sizes: torch.Tensor):735 is_image = pixel_values.ndim == 4736 if is_image:737 temporal = self.config.temporal_downsample_factor738 batch_size, channels, height, width = pixel_values.shape739 pixel_values = pixel_values.unsqueeze(1).repeat(1, temporal, 1, 1, 1)740 else:741 batch_size, temporal, channels, height, width = pixel_values.shape742 743 hidden_states = self.encoder(pixel_values)744 745 # b t c h w -> b c t h w746 hidden_states = hidden_states.permute(0, 2, 1, 3, 4)747 hidden_states = self.quant_conv(hidden_states)748 749 # b c t h w -> b t c h w750 hidden_states = hidden_states.permute(0, 2, 1, 3, 4)751 codes = self.quantize(hidden_states)752 753 image_tokens = codes.squeeze(1) if is_image else codes754 755 image_tokens = [756 single_image[: int(size[0] / self.vision_spatial_factor), : int(size[1] / self.vision_spatial_factor)]757 for single_image, size in zip(image_tokens, image_sizes)758 ]759 760 return image_tokens761 762 def decode(self, hidden_states: torch.Tensor):763 is_image = hidden_states.ndim == 3764 if is_image:765 hidden_states = hidden_states.unsqueeze(1)766 767 batch_size, temporal, height, width = hidden_states.shape768 quant = self.quantize.embedding(hidden_states.flatten())769 770 channels = quant.shape[-1]771 quant = quant.view(batch_size, temporal, height, width, channels).permute(0, 4, 1, 2, 3).contiguous()772 post_quant = self.post_quant_conv(quant)773 774 quant = quant.permute(0, 2, 1, 3, 4)775 post_quant = post_quant.permute(0, 2, 1, 3, 4)776 777 video = self.decoder(post_quant, quant)778 video = video.reshape(779 batch_size,780 temporal * self.config.temporal_downsample_factor,781 self.config.out_channels,782 height * self.spatial_scale_factor,783 width * self.spatial_scale_factor,784 )785 return video[:, 0] if is_image else video786 787 788class Emu3ImageVocabularyMapping:789 """790 A class for mapping discrete image tokens from VQGAN to BPE tokens.791 """792 793 def __init__(self, vocab_map):794 self.vocab_map = vocab_map795 self.eol_token_id = vocab_map.get("<|extra_200|>")796 self.image_token_id = vocab_map.get("<image>")797 798 @cached_property799 def image_tokens(self):800 return sorted([val for name, val in self.vocab_map.items() if name.startswith("<|visual token")])801 802 @cached_property803 def image_tokens_str(self):804 return sorted([name for name, val in self.vocab_map.items() if name.startswith("<|visual token")])805 806 @cached_property807 def img2bpe(self):808 return {int(token[-8:-2]): self.vocab_map[token] for token in self.image_tokens_str}809 810 @cached_property811 def bpe2img(self):812 return {v: k for k, v in self.img2bpe.items()}813 814 @cached_property815 def bpe2img_mapping_tensor(self):816 mapping = torch.zeros(max(self.bpe2img.keys()) + 1, dtype=torch.int)817 for k, v in self.bpe2img.items():818 mapping[k] = v819 return mapping820 821 @cached_property822 def img2bpe_mapping_tensor(self):823 mapping = torch.zeros(max(self.img2bpe.keys()) + 1, dtype=torch.int)824 for k, v in self.img2bpe.items():825 mapping[k] = v826 return mapping827 828 def convert_img2bpe(self, img_batch: list[torch.Tensor]) -> torch.Tensor:829 device = img_batch.device830 eol_row = torch.ones((img_batch.shape[0], 1), dtype=torch.int) * self.eol_token_id831 img_tokens = self.img2bpe_mapping_tensor[img_batch.to("cpu")]832 img_tokens = torch.cat([img_tokens, eol_row], dim=-1)833 return img_tokens.to(device)834 835 def convert_bpe2img(self, img_batch: torch.Tensor) -> torch.Tensor:836 device = img_batch.device837 img_batch = img_batch[..., :-1] # remove last row of EOL tokens838 img_tokens = self.bpe2img_mapping_tensor[img_batch.to("cpu")]839 return img_tokens.to(device)840 841 842class Emu3PreTrainedModel(ChameleonPreTrainedModel, Emu3VQVAE):843 _no_split_modules = [844 "Emu3DecoderLayer",845 ]846 _supports_flex_attn = True847 _supports_attention_backend = True848 849 850class Emu3TextModel(LlamaModel, Emu3PreTrainedModel):851 _can_record_outputs = {852 "hidden_states": Emu3DecoderLayer,853 "attentions": Emu3Attention,854 }855 856 def __init__(self, config: Emu3Config):857 super().__init__(config)858 self.layers = nn.ModuleList(859 [Emu3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]860 )861 862 863class Emu3ForCausalLM(LlamaForCausalLM, Emu3PreTrainedModel, GenerationMixin):864 config: Emu3TextConfig865 866 def __init__(self, config):867 super().__init__(config)868 self.model = Emu3TextModel(config)869 870 def forward(**super_kwargs):871 r"""872 Example:873 874 ```python875 >>> from transformers import Emu3Processor, Emu3ForConditionalGeneration876 >>> import torch877 >>> import requests878 >>> from PIL import Image879 880 >>> model = Emu3ForCausalLM.from_pretrained("BAAI/Emu3-Chat-hf", dtype=torch.bfloat16)881 >>> processor = Emu3Processor.from_pretrained("BAAI/Emu3-Chat-hf")882 883 >>> inputs = processor(text=["Can you write me a poem about winter."], return_tensors="pt").to(model.device)884 885 >>> generated_ids = model.generate(**inputs, max_new_tokens=100, do_sample=False)886 >>> processor.batch_decode(generated_ids, skip_special_tokens=True)[0]887 ```"""888 super().forward()889 890 891class Emu3Model(Emu3PreTrainedModel):892 _checkpoint_conversion_mapping = {"text_model.model": "text_model"}893 894 def __init__(self, config):895 super().__init__(config)896 self.text_model = Emu3TextModel._from_config(config.text_config)897 self.vqmodel = Emu3VQVAE(config.vq_config)898 self.vocabulary_mapping = Emu3ImageVocabularyMapping(config.vocabulary_map)899 900 # Initialize weights and apply final processing901 self.post_init()902 903 def get_input_embeddings(self):904 return self.text_model.get_input_embeddings()905 906 def set_input_embeddings(self, value):907 self.text_model.set_input_embeddings(value)908 909 def set_decoder(self, decoder):910 self.text_model = decoder911 912 def get_decoder(self):913 return self.text_model914 915 def get_image_tokens(self, pixel_values: torch.FloatTensor, image_sizes: torch.LongTensor):916 """917 Tokenizes images into discrete tokens with VQGAN module. Converts918 obtained image tokens into BPE tokens and wraps with "boi" and "eoi"919 special tokens.920 921 Args:922 pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):923 The tensors corresponding to the input images.924 image_sizes (`torch.LongTensor` of shape `(batch_size, 2)`):925 The sizes of the images in the batch, being (height, width) for each image.926 """927 image_tokens_list = self.vqmodel.encode(pixel_values, image_sizes)928 bpe_tokens_list = [self.vocabulary_mapping.convert_img2bpe(tokens).flatten() for tokens in image_tokens_list]929 bpe_tokens = torch.cat(bpe_tokens_list)930 return bpe_tokens931 932 def get_image_features(self, pixel_values: torch.FloatTensor, image_sizes: torch.LongTensor):933 """934 Tokenizes images into discrete tokens with VQGAN module and embeds935 them with text embeddings layer936 937 Args:938 pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)):939 The tensors corresponding to the input images.940 """941 image_tokens = self.get_image_tokens(pixel_values, image_sizes)942 split_sizes = [943 (height // self.vqmodel.vision_spatial_factor) * (width // self.vqmodel.vision_spatial_factor + 1)944 for height, width in image_sizes945 ]946 image_features = self.get_input_embeddings()(image_tokens)947 image_features = torch.split(image_features, split_sizes)948 return image_features949 950 @torch.no_grad951 def decode_image_tokens(self, image_tokens: torch.LongTensor, height: int, width: int):952 """953 Decodes generated image tokens from language model to continuous pixel values954 with VQGAN module via upsampling.955 956 Args:957 image_tokens (`torch.LongTensor` of shape `(batch_size, num_of_tokens)`):958 The tensors corresponding to the input images.959 height (`int`):960 Height of the generated image before upsampling.961 width (`int`):962 Width of the generated image before upsampling.963 """964 sequences = image_tokens[:, :-3].view(-1, height, width + 1)965 image_tokens = self.vocabulary_mapping.convert_bpe2img(sequences)966 image = self.vqmodel.decode(image_tokens)967 return image968 969 def get_placeholder_mask(970 self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor971 ):972 """973 Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is974 equal to the length of multimodal features. If the lengths are different, an error is raised.975 """976 if input_ids is None:977 special_image_mask = inputs_embeds == self.get_input_embeddings()(978 torch.tensor(self.vocabulary_mapping.image_token_id, dtype=torch.long, device=inputs_embeds.device)979 )980 special_image_mask = special_image_mask.all(-1)981 else:982 special_image_mask = input_ids == self.vocabulary_mapping.image_token_id983 984 n_image_tokens = special_image_mask.sum()985 special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)986 n_image_features = image_features.shape[0] * image_features.shape[1]987 if inputs_embeds[special_image_mask].numel() != image_features.numel():988 raise ValueError(989 f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"990 )991 return special_image_mask992 993 @can_return_tuple994 @auto_docstring995 def forward(996 self,997 input_ids: Optional[torch.LongTensor] = None,998 pixel_values: Optional[torch.FloatTensor] = None,999 image_sizes: Optional[torch.Tensor] = None,1000 attention_mask: Optional[torch.Tensor] = None,1001 position_ids: Optional[torch.LongTensor] = None,1002 past_key_values: Optional[Cache] = None,1003 inputs_embeds: Optional[torch.FloatTensor] = None,1004 use_cache: Optional[bool] = None,1005 cache_position: Optional[torch.LongTensor] = None,1006 **kwargs: Unpack[TransformersKwargs],1007 ) -> Union[tuple, CausalLMOutputWithPast]:1008 r"""1009 image_sizes (`torch.LongTensor` of shape `(batch_size, 2)`):1010 The sizes of the images in the batch, being (height, width) for each image. Image sizes can be obtained using1011 [`AutoImageProcessor`]. See [`Emu3ImageProcessor.__call__`] for details ([]`Emu3Processor`] uses1012 [`Emu3ImageProcessor`] for processing images).1013 """1014 if (input_ids is None) ^ (inputs_embeds is not None):1015 raise ValueError(1016 "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"1017 )1018 1019 if inputs_embeds is None:1020 inputs_embeds = self.get_input_embeddings()(input_ids)1021 1022 if pixel_values is not None:1023 image_embeds = self.get_image_features(pixel_values, image_sizes)1024 image_embeds = torch.cat(image_embeds, dim=0)1025 special_image_mask = self.get_placeholder_mask(1026 input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds1027 )1028 inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_embeds)1029 1030 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)1031 outputs = self.text_model(1032 attention_mask=attention_mask,1033 position_ids=position_ids,1034 past_key_values=past_key_values,1035 inputs_embeds=inputs_embeds,1036 use_cache=use_cache,1037 cache_position=cache_position,1038 **kwargs,1039 )1040 1041 return outputs1042 1043 1044class Emu3ForConditionalGeneration(Emu3PreTrainedModel, GenerationMixin):1045 base_model_prefix = ""1046 _tied_weights_keys = ["lm_head.weight"]1047 _checkpoint_conversion_mapping = {1048 "^text_model.model": "model.text_model",1049 "^vqmodel": "model.vqmodel",1050 "^text_model.lm_head": "lm_head",1051 }1052 1053 def __init__(self, config):1054 super().__init__(config)1055 self.model = Emu3Model(config)1056 self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)1057 1058 self.post_init()1059 1060 def get_input_embeddings(self):1061 return self.model.get_input_embeddings()1062 1063 def set_input_embeddings(self, value):1064 self.model.set_input_embeddings(value)1065 1066 def get_output_embeddings(self) -> nn.Module:1067 return self.lm_head1068 1069 def set_decoder(self, decoder):1070 self.model.set_decoder(decoder)1071 1072 def get_decoder(self):1073 return self.model.get_decoder()1074 1075 # Make modules available through conditional class for BC1076 @property1077 def text_model(self):1078 return self.model.text_model1079 1080 @property1081 def vqmodel(self):1082 return self.model.vqmodel1083 1084 @property1085 def vocabulary_mapping(self):1086 return self.model.vocabulary_mapping1087 1088 def decode_image_tokens(self, **kwargs):1089 return self.model.decode_image_tokens(**kwargs)1090 1091 @can_return_tuple1092 @auto_docstring1093 def forward(1094 self,1095 input_ids: Optional[torch.LongTensor] = None,1096 pixel_values: Optional[torch.FloatTensor] = None,1097 image_sizes: Optional[torch.Tensor] = None,1098 attention_mask: Optional[torch.Tensor] = None,1099 position_ids: Optional[torch.LongTensor] = None,1100 past_key_values: Optional[Cache] = None,1101 inputs_embeds: Optional[torch.FloatTensor] = None,1102 use_cache: Optional[bool] = None,1103 cache_position: Optional[torch.LongTensor] = None,1104 labels: Optional[torch.LongTensor] = None,1105 logits_to_keep: Union[int, torch.Tensor] = 0,1106 **kwargs: Unpack[TransformersKwargs],1107 ) -> Union[tuple, CausalLMOutputWithPast]:1108 r"""1109 image_sizes (`torch.LongTensor` of shape `(batch_size, 2)`):1110 The sizes of the images in the batch, being (height, width) for each image. Image sizes can be obtained using1111 [`AutoImageProcessor`]. See [`Emu3ImageProcessor.__call__`] for details ([]`Emu3Processor`] uses1112 [`Emu3ImageProcessor`] for processing images).1113 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1114 Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,1115 config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored1116 (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.1117 1118 Example:1119 1120 ```python1121 >>> from transformers import Emu3Processor, Emu3ForConditionalGeneration1122 >>> import torch1123 >>> import requests1124 >>> from PIL import Image1125 1126 >>> model = Emu3ForConditionalGeneration.from_pretrained("BAAI/Emu3-Chat-hf", dtype=torch.bfloat16)1127 >>> processor = Emu3Processor.from_pretrained("BAAI/Emu3-Chat-hf")1128 1129 >>> conversation = [1130 ... {1131 ... "role": "system",1132 ... "content": [1133 ... {"type": "text", "text": "You are a helpful assistant."},1134 ... ],1135 ... },1136 ... {1137 ... "role": "user",1138 ... "content": [1139 ... {"type": "image"},1140 ... {"type": "text", "text": "Please describe the image."},1141 ... ],1142 ... },1143 ... ]1144 1145 >>> prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)1146 >>> image = Image.open(requests.get("https://www.ilankelman.org/stopsigns/australia.jpg", stream=True).raw)1147 1148 >>> inputs = processor(images=[image], text=[prompt], return_tensors="pt").to(model.device, torch.bfloat16)1149 1150 >>> generated_ids = model.generate(**inputs, max_new_tokens=100, do_sample=False)1151 >>> processor.batch_decode(generated_ids, skip_special_tokens=True)[0]1152 ```"""1153 outputs = self.model(1154 input_ids=input_ids,1155 attention_mask=attention_mask,1156 position_ids=position_ids,1157 past_key_values=past_key_values,1158 inputs_embeds=inputs_embeds,1159 use_cache=use_cache,1160 cache_position=cache_position,1161 **kwargs,1162 )1163 1164 hidden_states = outputs[0]1165 # Only compute necessary logits, and do not upcast them to float if we are not computing the loss1166 slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep1167 logits = self.lm_head(hidden_states[:, slice_indices, :])1168 1169 loss = None1170 if labels is not None:1171 loss = self.loss_function(1172 logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs1173 )1174 1175 return CausalLMOutputWithPast(1176 loss=loss,1177 logits=logits,1178 past_key_values=outputs.past_key_values,1179 hidden_states=outputs.hidden_states,1180 attentions=outputs.attentions,1181 )1182 1183 def prepare_inputs_for_generation(1184 self,1185 input_ids,1186 past_key_values=None,1187 attention_mask=None,1188 inputs_embeds=None,1189 cache_position=None,1190 position_ids=None,1191 use_cache=True,1192 pixel_values=None,1193 **kwargs,1194 ):1195 # Overwritten -- in specific circumstances we don't want to forward image inputs to the model1196 1197 model_inputs = super().prepare_inputs_for_generation(1198 input_ids,1199 past_key_values=past_key_values,1200 attention_mask=attention_mask,