Aluode/PerceptionLabPortable
0
1# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨2# This file was automatically generated from src/transformers/models/emu3/modular_emu3.py.3# Do NOT edit this file manually as any edits will be overwritten by the generation of4# the file from the modular. If any change should be done, please apply the change to the5# modular_emu3.py file directly. One of our CI enforces this.6# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨7# coding=utf-88# Copyright 2024 HuggingFace Inc. team. All rights reserved.9#10#11# Licensed under the Apache License, Version 2.0 (the "License");12# you may not use this file except in compliance with the License.13# You may obtain a copy of the License at14#15# http://www.apache.org/licenses/LICENSE-2.016#17# Unless required by applicable law or agreed to in writing, software18# distributed under the License is distributed on an "AS IS" BASIS,19# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.20# See the License for the specific language governing permissions and21# limitations under the License.22 23import math24from functools import cached_property25from typing import Callable, Optional, Union26 27import torch28import torch.nn as nn29import torch.nn.functional as F30 31from ...activations import ACT2FN32from ...cache_utils import Cache, DynamicCache33from ...generation import GenerationMixin34from ...integrations import use_kernel_forward_from_hub35from ...masking_utils import create_causal_mask36from ...modeling_layers import GradientCheckpointingLayer37from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast38from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update39from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel40from ...processing_utils import Unpack41from ...utils import TransformersKwargs, auto_docstring, can_return_tuple42from ...utils.deprecation import deprecate_kwarg43from ...utils.generic import check_model_inputs44from .configuration_emu3 import Emu3Config, Emu3TextConfig, Emu3VQVAEConfig45 46 47def rotate_half(x):48 """Rotates half the hidden dims of the input."""49 x1 = x[..., : x.shape[-1] // 2]50 x2 = x[..., x.shape[-1] // 2 :]51 return torch.cat((-x2, x1), dim=-1)52 53 54def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):55 """Applies Rotary Position Embedding to the query and key tensors.56 57 Args:58 q (`torch.Tensor`): The query tensor.59 k (`torch.Tensor`): The key tensor.60 cos (`torch.Tensor`): The cosine part of the rotary embedding.61 sin (`torch.Tensor`): The sine part of the rotary embedding.62 position_ids (`torch.Tensor`, *optional*):63 Deprecated and unused.64 unsqueeze_dim (`int`, *optional*, defaults to 1):65 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and66 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note67 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and68 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes69 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have70 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.71 Returns:72 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.73 """74 cos = cos.unsqueeze(unsqueeze_dim)75 sin = sin.unsqueeze(unsqueeze_dim)76 q_embed = (q * cos) + (rotate_half(q) * sin)77 k_embed = (k * cos) + (rotate_half(k) * sin)78 return q_embed, k_embed79 80 81def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:82 """83 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,84 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)85 """86 batch, num_key_value_heads, slen, head_dim = hidden_states.shape87 if n_rep == 1:88 return hidden_states89 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)90 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)91 92 93def eager_attention_forward(94 module: nn.Module,95 query: torch.Tensor,96 key: torch.Tensor,97 value: torch.Tensor,98 attention_mask: Optional[torch.Tensor],99 scaling: float,100 dropout: float = 0.0,101 **kwargs: Unpack[TransformersKwargs],102):103 key_states = repeat_kv(key, module.num_key_value_groups)104 value_states = repeat_kv(value, module.num_key_value_groups)105 106 attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling107 if attention_mask is not None:108 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]109 attn_weights = attn_weights + causal_mask110 111 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)112 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)113 attn_output = torch.matmul(attn_weights, value_states)114 attn_output = attn_output.transpose(1, 2).contiguous()115 116 return attn_output, attn_weights117 118 119class Emu3Attention(nn.Module):120 """Multi-headed attention from 'Attention Is All You Need' paper"""121 122 def __init__(self, config: Emu3Config, layer_idx: int):123 super().__init__()124 self.config = config125 self.layer_idx = layer_idx126 self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)127 self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads128 self.scaling = self.head_dim**-0.5129 self.attention_dropout = config.attention_dropout130 self.is_causal = True131 132 self.q_proj = nn.Linear(133 config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias134 )135 self.k_proj = nn.Linear(136 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias137 )138 self.v_proj = nn.Linear(139 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias140 )141 self.o_proj = nn.Linear(142 config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias143 )144 145 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")146 def forward(147 self,148 hidden_states: torch.Tensor,149 position_embeddings: tuple[torch.Tensor, torch.Tensor],150 attention_mask: Optional[torch.Tensor],151 past_key_values: Optional[Cache] = None,152 cache_position: Optional[torch.LongTensor] = None,153 **kwargs: Unpack[TransformersKwargs],154 ) -> tuple[torch.Tensor, torch.Tensor]:155 input_shape = hidden_states.shape[:-1]156 hidden_shape = (*input_shape, -1, self.head_dim)157 158 query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)159 key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)160 value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)161 162 cos, sin = position_embeddings163 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)164 165 if past_key_values is not None:166 # sin and cos are specific to RoPE models; cache_position needed for the static cache167 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}168 key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)169 170 attention_interface: Callable = eager_attention_forward171 if self.config._attn_implementation != "eager":172 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]173 174 attn_output, attn_weights = attention_interface(175 self,176 query_states,177 key_states,178 value_states,179 attention_mask,180 dropout=0.0 if not self.training else self.attention_dropout,181 scaling=self.scaling,182 **kwargs,183 )184 185 attn_output = attn_output.reshape(*input_shape, -1).contiguous()186 attn_output = self.o_proj(attn_output)187 return attn_output, attn_weights188 189 190@use_kernel_forward_from_hub("RMSNorm")191class Emu3RMSNorm(nn.Module):192 def __init__(self, hidden_size, eps=1e-6):193 """194 Emu3RMSNorm is equivalent to T5LayerNorm195 """196 super().__init__()197 self.weight = nn.Parameter(torch.ones(hidden_size))198 self.variance_epsilon = eps199 200 def forward(self, hidden_states):201 input_dtype = hidden_states.dtype202 hidden_states = hidden_states.to(torch.float32)203 variance = hidden_states.pow(2).mean(-1, keepdim=True)204 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)205 return self.weight * hidden_states.to(input_dtype)206 207 def extra_repr(self):208 return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"209 210 211class Emu3MLP(nn.Module):212 def __init__(self, config):213 super().__init__()214 self.config = config215 self.hidden_size = config.hidden_size216 self.intermediate_size = config.intermediate_size217 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)218 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)219 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)220 self.act_fn = ACT2FN[config.hidden_act]221 222 def forward(self, x):223 down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))224 return down_proj225 226 227class Emu3DecoderLayer(GradientCheckpointingLayer):228 def __init__(self, config: Emu3Config, layer_idx: int):229 super().__init__()230 self.hidden_size = config.hidden_size231 232 self.self_attn = Emu3Attention(config=config, layer_idx=layer_idx)233 234 self.mlp = Emu3MLP(config)235 self.input_layernorm = Emu3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)236 self.post_attention_layernorm = Emu3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)237 self.dropout = nn.Dropout(config.attention_dropout)238 239 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")240 def forward(241 self,242 hidden_states: torch.Tensor,243 attention_mask: Optional[torch.Tensor] = None,244 position_ids: Optional[torch.LongTensor] = None,245 past_key_values: Optional[Cache] = None,246 use_cache: Optional[bool] = False,247 cache_position: Optional[torch.LongTensor] = None,248 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,249 **kwargs: Unpack[TransformersKwargs],250 ) -> torch.Tensor:251 residual = hidden_states252 hidden_states = self.input_layernorm(hidden_states)253 254 hidden_states, _ = self.self_attn(255 hidden_states=hidden_states,256 attention_mask=attention_mask,257 position_ids=position_ids,258 past_key_values=past_key_values,259 use_cache=use_cache,260 cache_position=cache_position,261 position_embeddings=position_embeddings,262 **kwargs,263 )264 hidden_states = residual + self.dropout(hidden_states)265 266 residual = hidden_states267 hidden_states = self.post_attention_layernorm(hidden_states)268 hidden_states = self.mlp(hidden_states)269 hidden_states = residual + self.dropout(hidden_states)270 return hidden_states271 272 273class Emu3VQVAEVectorQuantizer(nn.Module):274 """275 A module for vector quantization using learned embedding vectors.276 277 This module implements the quantization process similar to te one described in278 the VQ-VAE (Vector Quantized Variational AutoEncoder) paper. It quantizes continuous279 input vectors into discrete codebook vectors, which are learned during training.280 Current implementation improves over previous ones by avoiding costly matrix multiplications281 and allowing for post-hoc remapping of indices.282 """283 284 def __init__(self, config: Emu3VQVAEConfig):285 super().__init__()286 self.embedding = nn.Embedding(config.codebook_size, config.embed_dim)287 self.embedding.weight.data.uniform_(-1.0 / config.codebook_size, 1.0 / config.codebook_size)288 289 def forward(self, hidden_state: torch.Tensor):290 batch_size, temporal, channels, height, width = hidden_state.shape291 hidden_state = hidden_state.permute(0, 1, 3, 4, 2).contiguous()292 hidden_state_flattened = hidden_state.view(-1, channels)293 294 # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z295 hidden_state_sum = torch.sum(hidden_state_flattened**2, dim=1, keepdim=True)296 embedding_sum = torch.sum(self.embedding.weight**2, dim=1)297 298 # "bd,dn->bn",299 distances = 2 * torch.matmul(hidden_state_flattened, self.embedding.weight.transpose(0, 1))300 distances = hidden_state_sum + embedding_sum - distances301 302 min_encoding_indices = torch.argmin(distances, dim=1)303 min_encoding_indices = min_encoding_indices.view(batch_size, temporal, height, width)304 return min_encoding_indices305 306 307class Emu3VQVAEEncoderConvDownsample(nn.Module):308 def __init__(self, in_channels):309 super().__init__()310 self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)311 312 def forward(self, hidden_states):313 # no asymmetric padding in torch conv, must do it ourselves314 hidden_states = F.pad(hidden_states, pad=(0, 1, 0, 1), mode="constant", value=0)315 hidden_states = self.conv(hidden_states)316 return hidden_states317 318 319class Emu3VQVAEEncoderConvUpsample(nn.Module):320 def __init__(self, in_channels):321 super().__init__()322 self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)323 324 def forward(self, hidden_states):325 hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")326 hidden_states = self.conv(hidden_states)327 return hidden_states328 329 330class Emu3VQVAEConv3d(nn.Module):331 def __init__(332 self,333 in_channel: int,334 out_channel: int,335 kernel_size: tuple[int],336 stride: tuple[int],337 ):338 super().__init__()339 340 padding_sizes = [one_kernel - one_stride for one_kernel, one_stride in zip(kernel_size[1:], stride[1:])]341 self.padding = ()342 for pad_size in padding_sizes[::-1]:343 self.padding += (pad_size // 2 + pad_size % 2, pad_size // 2)344 self.padding += (2, 0)345 346 self.conv = nn.Conv3d(347 in_channel,348 out_channel,349 kernel_size,350 stride=stride,351 )352 353 def forward(self, hidden_states: torch.Tensor):354 hidden_states = F.pad(hidden_states, self.padding)355 hidden_states = self.conv(hidden_states)356 return hidden_states357 358 359class Emu3VQVAESpatialNorm(nn.Module):360 def __init__(361 self,362 in_channels: int,363 out_channels: int,364 ):365 super().__init__()366 self.norm_layer = nn.GroupNorm(367 num_channels=out_channels,368 num_groups=32,369 eps=1e-6,370 affine=True,371 )372 373 self.conv_y = nn.Conv2d(374 in_channels,375 out_channels,376 kernel_size=1,377 stride=1,378 padding=0,379 )380 self.conv_b = nn.Conv2d(381 in_channels,382 out_channels,383 kernel_size=1,384 stride=1,385 padding=0,386 )387 388 def forward(self, hidden_states: torch.Tensor, quant_states: torch.Tensor):389 quant_states = F.interpolate(quant_states, size=hidden_states.shape[-2:], mode="nearest")390 hidden_states = self.norm_layer(hidden_states)391 hidden_states = hidden_states * self.conv_y(quant_states) + self.conv_b(quant_states)392 return hidden_states393 394 395class Emu3VQVAETemporalUpsample(nn.Module):396 def __init__(397 self,398 in_channel: int,399 out_channel: int,400 ):401 super().__init__()402 self.conv = Emu3VQVAEConv3d(403 in_channel,404 out_channel,405 kernel_size=(3, 3, 3),406 stride=(1, 1, 1),407 )408 409 def forward(self, hidden_states: torch.Tensor):410 batch_size, channels, temporal, height, width = hidden_states.shape411 hidden_states = hidden_states.permute(0, 1, 3, 4, 2).contiguous().view(batch_size, -1, temporal)412 hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")413 hidden_states = hidden_states.view(batch_size, channels, height, width, -1).permute(0, 1, 4, 2, 3).contiguous()414 hidden_states = self.conv(hidden_states)415 return hidden_states416 417 418class Emu3VQVAETemporalDownsample(nn.Module):419 def __init__(420 self,421 in_channel: int,422 out_channel: int,423 ):424 super().__init__()425 self.conv = Emu3VQVAEConv3d(426 in_channel,427 out_channel,428 kernel_size=(4, 3, 3),429 stride=(2, 1, 1),430 )431 432 def forward(self, hidden_states: torch.Tensor):433 hidden_states = self.conv(hidden_states)434 return hidden_states435 436 437class Emu3VQVAETemporalResnetBlock(nn.Module):438 def __init__(439 self,440 in_channels,441 out_channels=None,442 ):443 super().__init__()444 self.in_channels = in_channels445 self.out_channels = in_channels if out_channels is None else out_channels446 447 self.norm1 = nn.BatchNorm3d(in_channels)448 self.conv1 = Emu3VQVAEConv3d(449 in_channels,450 out_channels,451 kernel_size=(3, 3, 3),452 stride=(1, 1, 1),453 )454 self.norm2 = nn.BatchNorm3d(out_channels)455 self.conv2 = Emu3VQVAEConv3d(456 out_channels,457 out_channels,458 kernel_size=(3, 3, 3),459 stride=(1, 1, 1),460 )461 if self.in_channels != self.out_channels:462 self.nin_shortcut = nn.Conv3d(463 in_channels,464 out_channels,465 kernel_size=1,466 stride=1,467 padding=0,468 )469 470 def forward(self, hidden_states):471 residual = hidden_states472 hidden_states = self.norm1(hidden_states)473 hidden_states *= torch.sigmoid(hidden_states)474 hidden_states = self.conv1(hidden_states)475 476 hidden_states = self.norm2(hidden_states)477 hidden_states *= torch.sigmoid(hidden_states)478 hidden_states = self.conv2(hidden_states)479 480 if self.in_channels != self.out_channels:481 residual = self.nin_shortcut(residual)482 483 return residual + hidden_states484 485 486class Emu3VQVAEResnetBlock(nn.Module):487 def __init__(488 self,489 in_channels: int,490 out_channels: Optional[int] = None,491 quant_channels: Optional[int] = None,492 ):493 super().__init__()494 self.in_channels = in_channels495 out_channels = in_channels if out_channels is None else out_channels496 self.out_channels = out_channels497 self.quant_channels = quant_channels498 499 if quant_channels is None:500 self.norm1 = nn.GroupNorm(num_channels=in_channels, num_groups=32, eps=1e-6, affine=True)501 self.norm2 = nn.GroupNorm(num_channels=out_channels, num_groups=32, eps=1e-6, affine=True)502 else:503 self.norm1 = Emu3VQVAESpatialNorm(quant_channels, in_channels)504 self.norm2 = Emu3VQVAESpatialNorm(quant_channels, out_channels)505 506 self.conv1 = nn.Conv2d(507 in_channels,508 out_channels,509 kernel_size=3,510 stride=1,511 padding=1,512 )513 514 self.conv2 = nn.Conv2d(515 out_channels,516 out_channels,517 kernel_size=3,518 stride=1,519 padding=1,520 )521 522 if self.in_channels != self.out_channels:523 self.nin_shortcut = nn.Conv2d(524 in_channels,525 out_channels,526 kernel_size=1,527 stride=1,528 padding=0,529 )530 531 def forward(self, hidden_states: torch.Tensor, quant_channels: Optional[torch.Tensor] = None):532 norm_args = () if self.quant_channels is None else (quant_channels,)533 534 residual = hidden_states535 hidden_states = self.norm1(hidden_states, *norm_args)536 hidden_states *= torch.sigmoid(hidden_states)537 hidden_states = self.conv1(hidden_states)538 539 hidden_states = self.norm2(hidden_states, *norm_args)540 hidden_states *= torch.sigmoid(hidden_states)541 hidden_states = self.conv2(hidden_states)542 543 if self.in_channels != self.out_channels:544 residual = self.nin_shortcut(residual)545 546 return residual + hidden_states547 548 549class Emu3VQVAEAttentionBlock(nn.Module):550 """Multi-headed attention from 'Attention Is All You Need' paper"""551 552 def __init__(self, config: Emu3VQVAEConfig):553 super().__init__()554 self.config = config555 self.embed_dim = config.hidden_size556 self.num_heads = config.num_attention_heads557 self.head_dim = self.embed_dim // self.num_heads558 if self.head_dim * self.num_heads != self.embed_dim:559 raise ValueError(560 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"561 f" {self.num_heads})."562 )563 self.scale = self.head_dim**-0.5564 self.dropout = config.attention_dropout565 self.is_causal = False566 567 self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)568 self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)569 self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)570 self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)571 572 # for compatibility with the attention interface573 self.num_key_value_groups = 1574 575 def forward(576 self,577 hidden_states: torch.Tensor,578 attention_mask: Optional[torch.Tensor] = None,579 **kwargs,580 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:581 """Input shape: Batch x Time x Channel"""582 583 batch_size, seq_length, embed_dim = hidden_states.shape584 585 queries = self.q_proj(hidden_states)586 keys = self.k_proj(hidden_states)587 values = self.v_proj(hidden_states)588 589 queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)590 keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)591 values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)592 593 attention_interface: Callable = eager_attention_forward594 if self.config._attn_implementation != "eager":595 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]596 597 attn_output, attn_weights = attention_interface(598 self,599 queries,600 keys,601 values,602 attention_mask,603 is_causal=self.is_causal,604 scaling=self.scale,605 dropout=0.0 if not self.training else self.dropout,606 )607 608 attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous()609 attn_output = self.out_proj(attn_output)610 611 return attn_output, attn_weights612 613 614class Emu3VQVAEGroupNorm(nn.GroupNorm):615 """616 Same as the torch GroupNorm with the only difference that this ones accepts617 an optional kwarg `quant_states` which is not used. This class makes it easier to618 use SpatialNorm or GroupNorm without conditionals619 """620 621 def __init__(self, **kwargs):622 super().__init__(**kwargs)623 624 def forward(self, input, quant_states=None):625 return F.group_norm(input, self.num_groups, self.weight, self.bias, self.eps)626 627 628class Emu3VQVAEMiddleBlock(nn.Module):629 def __init__(self, config, in_channels, quant_channels=None):630 super().__init__()631 632 self.block_1 = Emu3VQVAEResnetBlock(633 in_channels=in_channels,634 out_channels=in_channels,635 quant_channels=quant_channels,636 )637 self.attn_1 = Emu3VQVAEAttentionBlock(config)638 if quant_channels is None:639 self.attn_norm = Emu3VQVAEGroupNorm(num_channels=in_channels, num_groups=32, eps=1e-6, affine=True)640 else:641 self.attn_norm = Emu3VQVAESpatialNorm(quant_channels, in_channels)642 643 self.block_2 = Emu3VQVAEResnetBlock(644 in_channels=in_channels,645 out_channels=in_channels,646 quant_channels=quant_channels,647 )648 649 def forward(self, hidden_states: torch.FloatTensor, quant_states: Optional[torch.FloatTensor] = None):650 hidden_states = self.block_1(hidden_states, quant_states)651 residual = hidden_states652 hidden_states = self.attn_norm(hidden_states, quant_states)653 batch_size, channels, height, width = hidden_states.shape654 hidden_states = hidden_states.view(batch_size, channels, height * width).transpose(1, 2)655 hidden_states = self.attn_1(hidden_states)[0]656 hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)657 hidden_states = residual + hidden_states658 hidden_states = self.block_2(hidden_states, quant_states)659 return hidden_states660 661 662class Emu3VQVAEDownBlock(nn.Module):663 def __init__(self, config):664 super().__init__()665 666 self.num_resolutions = len(config.channel_multiplier)667 self.num_res_blocks = config.num_res_blocks668 base_channels = config.base_channels669 channel_multiplier = config.channel_multiplier670 671 in_channel_multiplier = (1,) + tuple(channel_multiplier)672 self.in_channel_multiplier = in_channel_multiplier673 self.down = nn.ModuleList()674 for i_level in range(self.num_resolutions):675 block = nn.ModuleList()676 attn = nn.ModuleList()677 attn_norms = nn.ModuleList()678 block_in = base_channels * in_channel_multiplier[i_level]679 block_out = base_channels * channel_multiplier[i_level]680 for i_block in range(self.num_res_blocks):681 block.append(682 Emu3VQVAEResnetBlock(683 in_channels=block_in,684 out_channels=block_out,685 )686 )687 block_in = block_out688 if config.attn_resolutions is not None and i_level in config.attn_resolutions:689 attn.append(Emu3VQVAEAttentionBlock(config))690 attn_norms.append(nn.GroupNorm(num_channels=block_in, num_groups=32, eps=1e-6, affine=True))691 692 down = nn.Module()693 down.block = block694 down.attn = attn695 down.attn_norms = attn_norms696 if i_level != self.num_resolutions - 1:697 down.downsample = Emu3VQVAEEncoderConvDownsample(block_in)698 self.down.append(down)699 700 def forward(self, hidden_states: torch.FloatTensor):701 for i_level, blocks in enumerate(self.down):702 for i_block in range(self.num_res_blocks):703 hidden_states = blocks.block[i_block](hidden_states)704 if len(blocks.attn) > 0:705 residual = hidden_states706 hidden_states = blocks.attn_norms[i_block](hidden_states)707 708 batch_size, channels, height, width = hidden_states.shape709 hidden_states = hidden_states.view(batch_size, channels, height * width).transpose(1, 2)710 hidden_states = blocks.attn[i_block](hidden_states)[0]711 712 hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)713 hidden_states = residual + hidden_states714 715 if i_level != self.num_resolutions - 1:716 hidden_states = blocks.downsample(hidden_states)717 718 return hidden_states719 720 721class Emu3VQVAEUpBlock(nn.Module):722 def __init__(self, config):723 super().__init__()724 725 self.num_resolutions = len(config.channel_multiplier)726 self.num_res_blocks = config.num_res_blocks727 728 quant_channels = config.embed_dim729 block_in = config.base_channels * config.channel_multiplier[-1]730 731 self.up = nn.ModuleList()732 for i_level in reversed(range(self.num_resolutions)):733 block = nn.ModuleList()734 attn = nn.ModuleList()735 attn_norms = nn.ModuleList()736 block_out = config.base_channels * config.channel_multiplier[i_level]737 for i_block in range(self.num_res_blocks + 1):738 block.append(739 Emu3VQVAEResnetBlock(740 in_channels=block_in,741 out_channels=block_out,742 quant_channels=quant_channels,743 )744 )745 block_in = block_out746 if i_level in config.attn_resolutions:747 attn.append(Emu3VQVAEAttentionBlock(config))748 attn_norms.append(Emu3VQVAESpatialNorm(quant_channels, block_in))749 750 up = nn.Module()751 up.block = block752 up.attn = attn753 up.attn_norms = attn_norms754 if i_level != 0:755 up.upsample = Emu3VQVAEEncoderConvUpsample(block_in)756 757 self.up.insert(0, up)758 759 def forward(self, hidden_states: torch.FloatTensor, quant_states: torch.FloatTensor):760 for i_level, blocks in enumerate(self.up[::-1]):761 for i_block in range(self.num_res_blocks + 1):762 hidden_states = blocks.block[i_block](hidden_states, quant_states)763 if len(blocks.attn) > 0:764 residual = hidden_states765 hidden_states = blocks.attn_norms[i_block](hidden_states, quant_states)766 767 batch_size, channels, height, width = hidden_states.shape768 hidden_states = hidden_states.view(batch_size, channels, height * width).transpose(1, 2)769 hidden_states = blocks.attn[i_block](hidden_states)[0]770 771 hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)772 hidden_states = residual + hidden_states773 if i_level != len(self.up) - 1:774 hidden_states = blocks.upsample(hidden_states)775 776 return hidden_states777 778 779class Emu3VQVAEEncoder(nn.Module):780 def __init__(self, config):781 super().__init__()782 783 base_channels = config.base_channels784 in_channels = config.in_channels785 double_latent = config.double_latent786 latent_channels = config.latent_channels787 channel_multiplier = config.channel_multiplier788 out_channels = 2 * latent_channels if double_latent else latent_channels789 block_in = base_channels * channel_multiplier[-1]790 791 self.conv_in = torch.nn.Conv2d(in_channels, base_channels, kernel_size=3, stride=1, padding=1)792 self.down_block = Emu3VQVAEDownBlock(config)793 self.middle_block = Emu3VQVAEMiddleBlock(config, block_in)794 795 self.norm_out = torch.nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)796 self.conv_out = torch.nn.Conv2d(797 block_in,798 out_channels,799 kernel_size=3,800 stride=1,801 padding=1,802 )803 804 temporal_down_blocks = int(math.log2(config.temporal_downsample_factor))805 self.time_conv = nn.ModuleList()806 self.time_res_stack = nn.ModuleList()807 808 for i in range(temporal_down_blocks):809 conv = Emu3VQVAETemporalDownsample(out_channels, out_channels)810 self.time_conv.append(conv)811 812 for _ in range(config.num_res_blocks):813 time_res_conv = Emu3VQVAETemporalResnetBlock(814 in_channels=out_channels,815 out_channels=out_channels,816 )817 self.time_res_stack.append(time_res_conv)818 819 def forward(self, pixel_values: torch.LongTensor):820 temporal_dim = pixel_values.shape[1]821 pixel_values = pixel_values.reshape(-1, *pixel_values.shape[2:])822 823 # downsampling & middle824 hidden_states = self.conv_in(pixel_values)825 hidden_states = self.down_block(hidden_states)826 hidden_states = self.middle_block(hidden_states)827 828 # end829 hidden_states = self.norm_out(hidden_states)830 hidden_states *= torch.sigmoid(hidden_states)831 hidden_states = self.conv_out(hidden_states)832 833 hidden_states = hidden_states.reshape(-1, temporal_dim, *hidden_states.shape[1:])834 hidden_states = hidden_states.permute(0, 2, 1, 3, 4)835 836 # temporal convs837 for conv in self.time_conv:838 hidden_states = conv(hidden_states)839 hidden_states *= torch.sigmoid(hidden_states)840 841 for layer in self.time_res_stack:842 hidden_states = layer(hidden_states)843 844 hidden_states = hidden_states.permute(0, 2, 1, 3, 4)845 846 return hidden_states847 848 849class Emu3VQVAEDecoder(nn.Module):850 def __init__(self, config: Emu3VQVAEConfig):851 super().__init__()852 853 quant_channels = config.embed_dim854 block_in = config.base_channels * config.channel_multiplier[-1]855 self.time_res_stack = nn.ModuleList()856 for _ in range(config.num_res_blocks):857 time_res_conv = Emu3VQVAETemporalResnetBlock(858 in_channels=config.latent_channels, out_channels=config.latent_channels859 )860 self.time_res_stack.append(time_res_conv)861 862 temp_upsample_block_num = int(math.log2(config.temporal_downsample_factor))863 self.time_conv = nn.ModuleList()864 for i in range(temp_upsample_block_num):865 conv = Emu3VQVAETemporalUpsample(config.latent_channels, config.latent_channels)866 self.time_conv.append(conv)867 868 self.conv_in = nn.Conv2d(869 config.latent_channels,870 block_in,871 kernel_size=3,872 stride=1,873 padding=1,874 )875 876 self.middle_block = Emu3VQVAEMiddleBlock(config, block_in, quant_channels=quant_channels)877 self.up_block = Emu3VQVAEUpBlock(config)878 879 block_in = config.base_channels * config.channel_multiplier[0]880 self.norm_out = Emu3VQVAESpatialNorm(quant_channels, block_in)881 self.conv_out = nn.Conv2d(882 block_in,883 config.out_channels,884 kernel_size=3,885 stride=1,886 padding=1,887 )888 889 def forward(self, hidden_states: torch.Tensor, quant_states: torch.Tensor):890 hidden_quant_states = torch.cat((hidden_states, quant_states), dim=0)891 hidden_quant_states = hidden_quant_states.permute(0, 2, 1, 3, 4)892 893 # temporal convs894 for layer in self.time_res_stack:895 hidden_quant_states = layer(hidden_quant_states)896 897 for layer in self.time_conv:898 hidden_quant_states = layer(hidden_quant_states)899 hidden_quant_states *= torch.sigmoid(hidden_quant_states)900 901 hidden_quant_states = hidden_quant_states.permute(0, 2, 1, 3, 4)902 hidden_states, quant_states = torch.chunk(hidden_quant_states, 2, dim=0)903 hidden_states = hidden_states.reshape(-1, *hidden_states.shape[2:])904 quant_states = quant_states.reshape(-1, *quant_states.shape[2:])905 906 hidden_states = self.conv_in(hidden_states)907 908 # middle & upsampling909 hidden_states = self.middle_block(hidden_states, quant_states)910 hidden_states = self.up_block(hidden_states, quant_states)911 912 hidden_states = self.norm_out(hidden_states, quant_states)913 hidden_states *= torch.sigmoid(hidden_states)914 hidden_states = self.conv_out(hidden_states)915 916 return hidden_states917 918 919@auto_docstring(920 custom_intro="""921 The VQ-VAE model used in Emu3 for encoding/decoding images into discrete tokens.922 This model follows the "Make-a-scene: Scene-based text-to-image generation with human priors" paper from923 [ Oran Gafni, Adam Polyak, Oron Ashual, Shelly Sheynin, Devi Parikh, and Yaniv924 Taigman](https://huggingface.co/papers/2203.13131).925 """926)927class Emu3VQVAE(PreTrainedModel):928 config: Emu3VQVAEConfig929 base_model_prefix = "emuvideovq"930 main_input_name = "pixel_values"931 _supports_sdpa = True932 _supports_flash_attn = True933 _supports_flex_attn = True934 _supports_attention_backend = True935 _no_split_modules = [936 "Emu3VQVAETemporalResnetBlock",937 "Emu3VQVAEAttentionBlock",938 "Emu3VQVAEResnetBlock",939 "Emu3VQVAEVectorQuantizer",940 ]941 942 def _init_weights(self, module):943 if isinstance(module, (nn.Conv2d, nn.Conv3d)):944 nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")945 if module.bias is not None:946 fan_in, _ = nn.init._calculate_fan_in_and_fan_out(module.weight)947 bound = 1 / math.sqrt(fan_in)948 nn.init.uniform_(module.bias, -bound, bound)949 elif isinstance(module, nn.Linear):950 nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5))951 if module.bias is not None:952 fan_in, _ = nn.init._calculate_fan_in_and_fan_out(module.weight)953 bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0954 nn.init.uniform_(module.bias, -bound, bound)955 elif isinstance(module, (nn.BatchNorm2d, nn.BatchNorm3d, nn.GroupNorm)):956 nn.init.constant_(module.weight, 1.0)957 nn.init.constant_(module.bias, 0.0)958 elif isinstance(module, nn.Embedding):959 module.weight.data.normal_()960 if module.padding_idx is not None:961 module.weight.data[module.padding_idx].zero_()962 963 def __init__(self, config: Emu3VQVAEConfig):964 super().__init__(config)965 966 self.config = config967 968 self.encoder = Emu3VQVAEEncoder(config)969 self.decoder = Emu3VQVAEDecoder(config)970 self.quantize = Emu3VQVAEVectorQuantizer(config)971 self.vision_spatial_factor = 2 ** (len(config.channel_multiplier) - 1)972 973 self.quant_conv = Emu3VQVAEConv3d(974 config.latent_channels, config.embed_dim, kernel_size=(3, 1, 1), stride=(1, 1, 1)975 )976 self.post_quant_conv = Emu3VQVAEConv3d(977 config.embed_dim, config.latent_channels, kernel_size=(3, 1, 1), stride=(1, 1, 1)978 )979 self.spatial_scale_factor = 2 ** (len(config.channel_multiplier) - 1)980 self.eval() # Emu3's VQ model is frozen981 982 self.post_init()983 984 def encode(self, pixel_values: torch.Tensor, image_sizes: torch.Tensor):985 is_image = pixel_values.ndim == 4986 if is_image:987 temporal = self.config.temporal_downsample_factor988 batch_size, channels, height, width = pixel_values.shape989 pixel_values = pixel_values.unsqueeze(1).repeat(1, temporal, 1, 1, 1)990 else:991 batch_size, temporal, channels, height, width = pixel_values.shape992 993 hidden_states = self.encoder(pixel_values)994 995 # b t c h w -> b c t h w996 hidden_states = hidden_states.permute(0, 2, 1, 3, 4)997 hidden_states = self.quant_conv(hidden_states)998 999 # b c t h w -> b t c h w1000 hidden_states = hidden_states.permute(0, 2, 1, 3, 4)1001 codes = self.quantize(hidden_states)1002 1003 image_tokens = codes.squeeze(1) if is_image else codes1004 1005 image_tokens = [1006 single_image[: int(size[0] / self.vision_spatial_factor), : int(size[1] / self.vision_spatial_factor)]1007 for single_image, size in zip(image_tokens, image_sizes)1008 ]1009 1010 return image_tokens1011 1012 def decode(self, hidden_states: torch.Tensor):1013 is_image = hidden_states.ndim == 31014 if is_image:1015 hidden_states = hidden_states.unsqueeze(1)1016 1017 batch_size, temporal, height, width = hidden_states.shape1018 quant = self.quantize.embedding(hidden_states.flatten())1019 1020 channels = quant.shape[-1]1021 quant = quant.view(batch_size, temporal, height, width, channels).permute(0, 4, 1, 2, 3).contiguous()1022 post_quant = self.post_quant_conv(quant)1023 1024 quant = quant.permute(0, 2, 1, 3, 4)1025 post_quant = post_quant.permute(0, 2, 1, 3, 4)1026 1027 video = self.decoder(post_quant, quant)1028 video = video.reshape(1029 batch_size,1030 temporal * self.config.temporal_downsample_factor,1031 self.config.out_channels,1032 height * self.spatial_scale_factor,1033 width * self.spatial_scale_factor,1034 )1035 return video[:, 0] if is_image else video1036 1037 1038class Emu3ImageVocabularyMapping:1039 """1040 A class for mapping discrete image tokens from VQGAN to BPE tokens.1041 """1042 1043 def __init__(self, vocab_map):1044 self.vocab_map = vocab_map1045 self.eol_token_id = vocab_map.get("<|extra_200|>")1046 self.image_token_id = vocab_map.get("<image>")1047 1048 @cached_property1049 def image_tokens(self):1050 return sorted([val for name, val in self.vocab_map.items() if name.startswith("<|visual token")])1051 1052 @cached_property1053 def image_tokens_str(self):1054 return sorted([name for name, val in self.vocab_map.items() if name.startswith("<|visual token")])1055 1056 @cached_property1057 def img2bpe(self):1058 return {int(token[-8:-2]): self.vocab_map[token] for token in self.image_tokens_str}1059 1060 @cached_property1061 def bpe2img(self):1062 return {v: k for k, v in self.img2bpe.items()}1063 1064 @cached_property1065 def bpe2img_mapping_tensor(self):1066 mapping = torch.zeros(max(self.bpe2img.keys()) + 1, dtype=torch.int)1067 for k, v in self.bpe2img.items():1068 mapping[k] = v1069 return mapping1070 1071 @cached_property1072 def img2bpe_mapping_tensor(self):1073 mapping = torch.zeros(max(self.img2bpe.keys()) + 1, dtype=torch.int)1074 for k, v in self.img2bpe.items():1075 mapping[k] = v1076 return mapping1077 1078 def convert_img2bpe(self, img_batch: list[torch.Tensor]) -> torch.Tensor:1079 device = img_batch.device1080 eol_row = torch.ones((img_batch.shape[0], 1), dtype=torch.int) * self.eol_token_id1081 img_tokens = self.img2bpe_mapping_tensor[img_batch.to("cpu")]1082 img_tokens = torch.cat([img_tokens, eol_row], dim=-1)1083 return img_tokens.to(device)1084 1085 def convert_bpe2img(self, img_batch: torch.Tensor) -> torch.Tensor:1086 device = img_batch.device1087 img_batch = img_batch[..., :-1] # remove last row of EOL tokens1088 img_tokens = self.bpe2img_mapping_tensor[img_batch.to("cpu")]1089 return img_tokens.to(device)1090 1091 1092@auto_docstring1093class Emu3PreTrainedModel(PreTrainedModel):1094 config: Emu3Config1095 base_model_prefix = "model"1096 supports_gradient_checkpointing = True1097 _no_split_modules = [1098 "Emu3DecoderLayer",1099 ]1100 _skip_keys_device_placement = ["past_key_values", "causal_mask"]1101 _supports_flash_attn = True1102 _supports_sdpa = True1103 1104 _can_compile_fullgraph = True1105 _supports_param_buffer_assignment = False1106 _supports_flex_attn = True1107 _supports_attention_backend = True1108 1109 1110class Emu3RotaryEmbedding(nn.Module):1111 inv_freq: torch.Tensor # fix linting for `register_buffer`1112 1113 def __init__(self, config: Emu3Config, device=None):1114 super().__init__()1115 # BC: "rope_type" was originally "type"1116 if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):1117 self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))1118 else:1119 self.rope_type = "default"1120 self.max_seq_len_cached = config.max_position_embeddings1121 self.original_max_seq_len = config.max_position_embeddings1122 1123 self.config = config1124 self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]1125 1126 inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)1127 self.register_buffer("inv_freq", inv_freq, persistent=False)1128 self.original_inv_freq = self.inv_freq1129 1130 @torch.no_grad()1131 @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)1132 def forward(self, x, position_ids):1133 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)1134 position_ids_expanded = position_ids[:, None, :].float()1135 1136 device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"1137 with torch.autocast(device_type=device_type, enabled=False): # Force float321138 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)1139 emb = torch.cat((freqs, freqs), dim=-1)1140 cos = emb.cos() * self.attention_scaling1141 sin = emb.sin() * self.attention_scaling1142 1143 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)1144 1145 1146@auto_docstring1147class Emu3TextModel(Emu3PreTrainedModel):1148 _can_record_outputs = {1149 "hidden_states": Emu3DecoderLayer,1150 "attentions": Emu3Attention,1151 }1152 1153 def __init__(self, config: Emu3Config):1154 super().__init__(config)1155 self.padding_idx = config.pad_token_id1156 self.vocab_size = config.vocab_size1157 1158 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)1159 self.layers = nn.ModuleList(1160 [Emu3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]1161 )1162 self.norm = Emu3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)1163 self.rotary_emb = Emu3RotaryEmbedding(config=config)1164 self.gradient_checkpointing = False1165 1166 # Initialize weights and apply final processing1167 self.post_init()1168 1169 @check_model_inputs()1170 @auto_docstring1171 def forward(1172 self,1173 input_ids: Optional[torch.LongTensor] = None,1174 attention_mask: Optional[torch.Tensor] = None,1175 position_ids: Optional[torch.LongTensor] = None,1176 past_key_values: Optional[Cache] = None,1177 inputs_embeds: Optional[torch.FloatTensor] = None,1178 cache_position: Optional[torch.LongTensor] = None,1179 use_cache: Optional[bool] = None,1180 **kwargs: Unpack[TransformersKwargs],1181 ) -> BaseModelOutputWithPast:1182 if (input_ids is None) ^ (inputs_embeds is not None):1183 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")1184 1185 if inputs_embeds is None:1186 inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)1187 1188 if use_cache and past_key_values is None:1189 past_key_values = DynamicCache(config=self.config)1190 1191 if cache_position is None:1192 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 01193 cache_position: torch.Tensor = torch.arange(1194 past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device1195 )1196 1197 if position_ids is None:1198 position_ids = cache_position.unsqueeze(0)1199 1200 causal_mask = create_causal_mask(