# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from Helios (https://github.com/BestWishYsh/Helios) import math from collections import OrderedDict from collections.abc import Iterable from functools import lru_cache from typing import TYPE_CHECKING, Any import torch import torch.nn as nn import torch.nn.functional as F from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import FP32LayerNorm from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) from vllm.logger import init_logger from vllm.model_executor.layers.conv import Conv3dLayer from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, SequenceParallelOutput, ) if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, ) logger = init_logger(__name__) def pad_for_3d_conv(x, kernel_size): b, c, t, h, w = x.shape pt, ph, pw = kernel_size pad_t = (pt - (t % pt)) % pt pad_h = (ph - (h % ph)) % ph pad_w = (pw - (w % pw)) % pw return F.pad(x, (0, pad_w, 0, pad_h, 0, pad_t), mode="replicate") def center_down_sample_3d(x, kernel_size): return F.avg_pool3d(x, kernel_size, stride=kernel_size) def apply_rotary_emb_helios( hidden_states: torch.Tensor, freqs_cis: torch.Tensor, ) -> torch.Tensor: """Apply Helios-style rotary embeddings. freqs_cis contains [cos_t, cos_y, cos_x, sin_t, sin_y, sin_x] concatenated along the last dimension, with shape [B, seq, D*2] where D = DT+DY+DX. hidden_states has shape [B, seq, H, head_dim]. """ x_1, x_2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1) # Use stack+flatten instead of strided slice assignment for contiguous # memory layout and better performance on GPU/NPU (#2436, cf. PR #2393). rotated = torch.stack( ( x_1 * cos[..., 0::2] - x_2 * sin[..., 1::2], x_1 * sin[..., 1::2] + x_2 * cos[..., 0::2], ), dim=-1, ) return rotated.flatten(-2, -1).type_as(hidden_states) class DistributedRMSNorm(nn.Module): """RMSNorm that computes global RMS across tensor parallel ranks.""" def __init__(self, hidden_size: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(hidden_size)) def forward(self, x: torch.Tensor) -> torch.Tensor: tp_size = get_tensor_model_parallel_world_size() input_dtype = x.dtype x_float = x.float() local_sum_sq = (x_float**2).sum(dim=-1, keepdim=True) local_count = x.shape[-1] if tp_size > 1: global_sum_sq = local_sum_sq.clone() tensor_model_parallel_all_reduce(global_sum_sq) global_count = local_count * tp_size else: global_sum_sq = local_sum_sq global_count = local_count rms = torch.sqrt(global_sum_sq / global_count + self.eps) output = (x_float / rms) * self.weight.float() return output.to(input_dtype) class ColumnParallelGELU(nn.Module): """Column parallel linear with GELU activation.""" def __init__( self, dim_in: int, dim_out: int, *, approximate: str = "tanh", bias: bool = True, quant_config: "QuantizationConfig | None" = None, ): super().__init__() self.proj = ColumnParallelLinear( dim_in, dim_out, bias=bias, gather_output=False, return_bias=False, quant_config=quant_config, ) self.approximate = approximate def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.proj(x) return F.gelu(x, approximate=self.approximate) class HeliosFeedForward(nn.Module): """TP-enabled FeedForward network for Helios.""" def __init__( self, dim: int, inner_dim: int, dim_out: int | None = None, bias: bool = True, quant_config: "QuantizationConfig | None" = None, ) -> None: super().__init__() dim_out = dim_out or dim self.net_0 = ColumnParallelGELU(dim, inner_dim, approximate="tanh", bias=bias, quant_config=quant_config) self.net_1 = nn.Identity() self.net_2 = RowParallelLinear( inner_dim, dim_out, bias=bias, input_is_parallel=True, return_bias=False, quant_config=quant_config ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.net_0(hidden_states) hidden_states = self.net_1(hidden_states) hidden_states = self.net_2(hidden_states) return hidden_states class HeliosRotaryPosEmbed(nn.Module): """Helios-style 3D rotary position embeddings using explicit frame indices.""" def __init__(self, rope_dim: tuple[int, int, int], theta: float): super().__init__() self.DT, self.DY, self.DX = rope_dim self.theta = theta self.register_buffer("freqs_base_t", self._get_freqs_base(self.DT), persistent=False) self.register_buffer("freqs_base_y", self._get_freqs_base(self.DY), persistent=False) self.register_buffer("freqs_base_x", self._get_freqs_base(self.DX), persistent=False) def _get_freqs_base(self, dim): return 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32)[: (dim // 2)] / dim)) @torch.no_grad() def get_frequency_batched(self, freqs_base, pos): freqs = torch.einsum("d,bthw->dbthw", freqs_base, pos) freqs = freqs.repeat_interleave(2, dim=0) return freqs.cos(), freqs.sin() @torch.no_grad() @lru_cache(maxsize=32) def _get_spatial_meshgrid(self, height, width, device_str): device = torch.device(device_str) grid_y_coords = torch.arange(height, device=device, dtype=torch.float32) grid_x_coords = torch.arange(width, device=device, dtype=torch.float32) grid_y, grid_x = torch.meshgrid(grid_y_coords, grid_x_coords, indexing="ij") return grid_y, grid_x @torch.no_grad() def forward(self, frame_indices, height, width, device): batch_size = frame_indices.shape[0] num_frames = frame_indices.shape[1] frame_indices = frame_indices.to(device=device, dtype=torch.float32) grid_y, grid_x = self._get_spatial_meshgrid(height, width, str(device)) grid_t = frame_indices[:, :, None, None].expand(batch_size, num_frames, height, width) grid_y_batch = grid_y[None, None, :, :].expand(batch_size, num_frames, -1, -1) grid_x_batch = grid_x[None, None, :, :].expand(batch_size, num_frames, -1, -1) freqs_cos_t, freqs_sin_t = self.get_frequency_batched(self.freqs_base_t, grid_t) freqs_cos_y, freqs_sin_y = self.get_frequency_batched(self.freqs_base_y, grid_y_batch) freqs_cos_x, freqs_sin_x = self.get_frequency_batched(self.freqs_base_x, grid_x_batch) result = torch.cat( [freqs_cos_t, freqs_cos_y, freqs_cos_x, freqs_sin_t, freqs_sin_y, freqs_sin_x], dim=0, ) return result.permute(1, 0, 2, 3, 4) class HeliosTimeTextEmbedding(nn.Module): """Combined time and text condition embeddings for Helios.""" def __init__( self, dim: int, time_freq_dim: int, time_proj_dim: int, text_embed_dim: int, ): super().__init__() self.timesteps_proj = Timesteps(num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0) self.time_embedder = TimestepEmbedding(in_channels=time_freq_dim, time_embed_dim=dim) self.act_fn = nn.SiLU() self.time_proj = nn.Linear(dim, time_proj_dim) self.text_embedder = PixArtAlphaTextProjection(text_embed_dim, dim, act_fn="gelu_tanh") def forward( self, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None, is_return_encoder_hidden_states: bool = True, ): timestep = self.timesteps_proj(timestep) time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8: timestep = timestep.to(time_embedder_dtype) temb = self.time_embedder(timestep).type_as(encoder_hidden_states) timestep_proj = self.time_proj(self.act_fn(temb)) if encoder_hidden_states is not None and is_return_encoder_hidden_states: encoder_hidden_states = self.text_embedder(encoder_hidden_states) return temb, timestep_proj, encoder_hidden_states class HeliosOutputNorm(nn.Module): """Output normalization that extracts only original_context_length tokens.""" def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.scale_shift_table = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) self.norm = FP32LayerNorm(dim, eps, elementwise_affine=False) def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor, original_context_length: int): temb = temb[:, -original_context_length:, :] shift, scale = (self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)).chunk(2, dim=2) shift, scale = shift.squeeze(2).to(hidden_states.device), scale.squeeze(2).to(hidden_states.device) hidden_states = hidden_states[:, -original_context_length:, :] hidden_states = (self.norm(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states) return hidden_states class HeliosSelfAttention(nn.Module): """Optimized self-attention for Helios with history amplification support.""" def __init__( self, dim: int, num_heads: int, head_dim: int, eps: float = 1e-5, dropout: float = 0.0, is_amplify_history: bool = False, history_scale_mode: str = "per_head", quant_config: "QuantizationConfig | None" = None, ): super().__init__() self.dim = dim self.num_heads = num_heads self.head_dim = head_dim self.inner_dim = num_heads * head_dim self.to_qkv = QKVParallelLinear( hidden_size=dim, head_size=head_dim, total_num_heads=num_heads, bias=True, quant_config=quant_config, ) self.num_heads = self.to_qkv.num_heads self.num_kv_heads = self.to_qkv.num_kv_heads self.tp_inner_dim = self.num_heads * head_dim self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) self.to_out = RowParallelLinear( self.inner_dim, dim, bias=True, input_is_parallel=True, return_bias=False, quant_config=quant_config, ) self.dropout = nn.Dropout(dropout) self.attn = Attention( num_heads=self.num_heads, head_size=head_dim, num_kv_heads=self.num_kv_heads, softmax_scale=1.0 / (head_dim**0.5), causal=False, ) self.is_amplify_history = is_amplify_history if is_amplify_history: tp_size = get_tensor_model_parallel_world_size() local_heads = num_heads // tp_size if history_scale_mode == "scalar": self.history_key_scale = nn.Parameter(torch.ones(1)) elif history_scale_mode == "per_head": self.history_key_scale = nn.Parameter(torch.ones(local_heads)) else: raise ValueError(f"Unknown history_scale_mode: {history_scale_mode}") self.history_scale_mode = history_scale_mode self.max_scale = 10.0 def forward( self, hidden_states: torch.Tensor, rotary_emb: torch.Tensor | None = None, original_context_length: int | None = None, ) -> torch.Tensor: qkv, _ = self.to_qkv(hidden_states) q_size = self.num_heads * self.head_dim kv_size = self.num_kv_heads * self.head_dim query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1) query = self.norm_q(query) key = self.norm_k(key) query = query.unflatten(2, (self.num_heads, self.head_dim)) key = key.unflatten(2, (self.num_kv_heads, self.head_dim)) value = value.unflatten(2, (self.num_kv_heads, self.head_dim)) if rotary_emb is not None: query = apply_rotary_emb_helios(query, rotary_emb) key = apply_rotary_emb_helios(key, rotary_emb) if self.is_amplify_history and original_context_length is not None: history_seq_len = hidden_states.shape[1] - original_context_length if history_seq_len > 0: scale_key = 1.0 + torch.sigmoid(self.history_key_scale) * (self.max_scale - 1.0) if self.history_scale_mode == "per_head": scale_key = scale_key.view(1, 1, -1, 1) key = torch.cat( [key[:, :history_seq_len] * scale_key, key[:, history_seq_len:]], dim=1, ) hidden_states = self.attn(query, key, value) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.type_as(query) hidden_states = self.to_out(hidden_states) hidden_states = self.dropout(hidden_states) return hidden_states class HeliosCrossAttention(nn.Module): """Optimized cross-attention for Helios.""" def __init__( self, dim: int, num_heads: int, head_dim: int, eps: float = 1e-5, dropout: float = 0.0, quant_config: "QuantizationConfig | None" = None, ): super().__init__() self.dim = dim self.num_heads = num_heads self.head_dim = head_dim self.inner_dim = num_heads * head_dim self.to_q = ColumnParallelLinear( dim, self.inner_dim, bias=True, gather_output=False, return_bias=False, quant_config=quant_config, ) self.to_k = ColumnParallelLinear( dim, self.inner_dim, bias=True, gather_output=False, return_bias=False, quant_config=quant_config, ) self.to_v = ColumnParallelLinear( dim, self.inner_dim, bias=True, gather_output=False, return_bias=False, quant_config=quant_config, ) tp_size = get_tensor_model_parallel_world_size() self.num_heads = num_heads // tp_size self.tp_inner_dim = self.num_heads * head_dim self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) self.to_out = RowParallelLinear( self.inner_dim, dim, bias=True, input_is_parallel=True, return_bias=False, quant_config=quant_config, ) self.dropout = nn.Dropout(dropout) self.attn = Attention( num_heads=self.num_heads, head_size=head_dim, num_kv_heads=self.num_heads, softmax_scale=1.0 / (head_dim**0.5), causal=False, ) def project_kv(self, encoder_hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: key = self.to_k(encoder_hidden_states) value = self.to_v(encoder_hidden_states) key = self.norm_k(key) key = key.unflatten(2, (self.num_heads, self.head_dim)) value = value.unflatten(2, (self.num_heads, self.head_dim)) return key, value def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None, encoder_key_value: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: query = self.to_q(hidden_states) query = self.norm_q(query) if encoder_key_value is None: if encoder_hidden_states is None: raise ValueError("encoder_hidden_states is required when encoder_key_value is not provided.") key, value = self.project_kv(encoder_hidden_states) else: key, value = encoder_key_value query = query.unflatten(2, (self.num_heads, self.head_dim)) hidden_states = self.attn(query, key, value) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.type_as(query) hidden_states = self.to_out(hidden_states) hidden_states = self.dropout(hidden_states) return hidden_states class HeliosTransformerBlock(nn.Module): """Transformer block with guidance cross-attention and history support.""" def __init__( self, dim: int, ffn_dim: int, num_heads: int, eps: float = 1e-6, cross_attn_norm: bool = False, guidance_cross_attn: bool = False, is_amplify_history: bool = False, history_scale_mode: str = "per_head", quant_config: "QuantizationConfig | None" = None, ): super().__init__() head_dim = dim // num_heads # 1. Self-attention self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) self.attn1 = HeliosSelfAttention( dim=dim, num_heads=num_heads, head_dim=head_dim, eps=eps, is_amplify_history=is_amplify_history, history_scale_mode=history_scale_mode, quant_config=quant_config, ) # 2. Cross-attention self.attn2 = HeliosCrossAttention( dim=dim, num_heads=num_heads, head_dim=head_dim, eps=eps, quant_config=quant_config ) self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() # 3. Feed-forward self.ffn = HeliosFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim, quant_config=quant_config) self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) # 4. Guidance cross-attention flag self.guidance_cross_attn = guidance_cross_attn def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, rotary_emb: torch.Tensor, original_context_length: int | None = None, cross_attn_key_value: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: if temb.ndim == 4: shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( self.scale_shift_table.unsqueeze(0) + temb.float() ).chunk(6, dim=2) shift_msa = shift_msa.squeeze(2) scale_msa = scale_msa.squeeze(2) gate_msa = gate_msa.squeeze(2) c_shift_msa = c_shift_msa.squeeze(2) c_scale_msa = c_scale_msa.squeeze(2) c_gate_msa = c_gate_msa.squeeze(2) else: shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( self.scale_shift_table + temb.float() ).chunk(6, dim=1) # 1. Self-attention norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states) attn_output = self.attn1(norm_hidden_states, rotary_emb, original_context_length) hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states) # 2. Cross-attention (with optional guidance: only current chunk attends to text) if self.guidance_cross_attn and original_context_length is not None: history_seq_len = hidden_states.shape[1] - original_context_length history_hidden_states, current_hidden_states = ( hidden_states[:, :history_seq_len], hidden_states[:, history_seq_len:], ) norm_hidden_states = self.norm2(current_hidden_states.float()).type_as(current_hidden_states) attn_output = self.attn2( norm_hidden_states, encoder_hidden_states, encoder_key_value=cross_attn_key_value, ) current_hidden_states = current_hidden_states + attn_output hidden_states = torch.cat([history_hidden_states, current_hidden_states], dim=1) else: norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states) attn_output = self.attn2( norm_hidden_states, encoder_hidden_states, encoder_key_value=cross_attn_key_value, ) hidden_states = hidden_states + attn_output # 3. Feed-forward norm_hidden_states = (self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as( hidden_states ) ff_output = self.ffn(norm_hidden_states) hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states) return hidden_states class HeliosTransformer3DModel(nn.Module): """Optimized Helios Transformer model for video generation using vLLM layers. Helios extends the Wan2.2 architecture with multi-term memory patches, guidance cross-attention, and chunked video generation support. """ _repeated_blocks = ["HeliosTransformerBlock"] _layerwise_offload_blocks_attrs = ["blocks"] packed_modules_mapping = { "to_qkv": ["to_q", "to_k", "to_v"], } @staticmethod def _is_transformer_block(name: str, module) -> bool: return "blocks" in name and name.split(".")[-1].isdigit() _hsdp_shard_conditions = [_is_transformer_block] _sp_plan = { "rope": { 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), }, "blocks.0": { "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True), }, "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), } def __init__( self, patch_size: tuple[int, ...] = (1, 2, 2), num_attention_heads: int = 40, attention_head_dim: int = 128, in_channels: int = 16, out_channels: int = 16, text_dim: int = 4096, freq_dim: int = 256, ffn_dim: int = 13824, num_layers: int = 40, cross_attn_norm: bool = True, qk_norm: str | None = "rms_norm_across_heads", eps: float = 1e-6, added_kv_proj_dim: int | None = None, rope_dim: tuple[int, ...] = (44, 42, 42), rope_theta: float = 10000.0, guidance_cross_attn: bool = True, zero_history_timestep: bool = True, has_multi_term_memory_patch: bool = True, is_amplify_history: bool = False, history_scale_mode: str = "per_head", quant_config: "QuantizationConfig | None" = None, ) -> None: super().__init__() inner_dim = num_attention_heads * attention_head_dim out_channels = out_channels or in_channels self.config = type( "Config", (), { "patch_size": patch_size, "num_attention_heads": num_attention_heads, "attention_head_dim": attention_head_dim, "in_channels": in_channels, "out_channels": out_channels, "text_dim": text_dim, "freq_dim": freq_dim, "ffn_dim": ffn_dim, "num_layers": num_layers, "cross_attn_norm": cross_attn_norm, "qk_norm": qk_norm, "eps": eps, "added_kv_proj_dim": added_kv_proj_dim, "rope_dim": rope_dim, "rope_theta": rope_theta, "guidance_cross_attn": guidance_cross_attn, "zero_history_timestep": zero_history_timestep, "has_multi_term_memory_patch": has_multi_term_memory_patch, "is_amplify_history": is_amplify_history, "history_scale_mode": history_scale_mode, }, )() self.inner_dim = inner_dim self.zero_history_timestep = zero_history_timestep self.has_multi_term_memory_patch = has_multi_term_memory_patch # 1. Patch & position embedding self.rope = HeliosRotaryPosEmbed(rope_dim=rope_dim, theta=rope_theta) self.patch_embedding = Conv3dLayer( in_channels=in_channels, out_channels=inner_dim, kernel_size=patch_size, stride=patch_size, ) # 2. Multi-term memory patches if has_multi_term_memory_patch: self.patch_short = Conv3dLayer( in_channels=in_channels, out_channels=inner_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2), ) self.patch_mid = Conv3dLayer( in_channels=in_channels, out_channels=inner_dim, kernel_size=(2, 4, 4), stride=(2, 4, 4), ) self.patch_long = Conv3dLayer( in_channels=in_channels, out_channels=inner_dim, kernel_size=(4, 8, 8), stride=(4, 8, 8), ) # 3. Condition embeddings self.condition_embedder = HeliosTimeTextEmbedding( dim=inner_dim, time_freq_dim=freq_dim, time_proj_dim=inner_dim * 6, text_embed_dim=text_dim, ) # 4. Transformer blocks self.blocks = nn.ModuleList( [ HeliosTransformerBlock( inner_dim, ffn_dim, num_attention_heads, eps, cross_attn_norm, guidance_cross_attn=guidance_cross_attn, is_amplify_history=is_amplify_history, history_scale_mode=history_scale_mode, quant_config=quant_config, ) for _ in range(num_layers) ] ) # 5. Output norm & projection self.norm_out = HeliosOutputNorm(inner_dim, eps) self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) self._projected_encoder_cache: OrderedDict[tuple, torch.Tensor] = OrderedDict() self._cross_attn_kv_cache: OrderedDict[tuple, list[tuple[torch.Tensor, torch.Tensor]]] = OrderedDict() self._cross_attn_cache_size = 2 @property def dtype(self) -> torch.dtype: return next(self.parameters()).dtype @staticmethod def _tensor_cache_key(tensor: torch.Tensor) -> tuple: try: version = tensor._version except RuntimeError: version = None return ( tensor.data_ptr(), tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype, tensor.device.type, tensor.device.index, version, ) @staticmethod def _get_from_lru(cache: OrderedDict, key: tuple): value = cache.get(key) if value is not None: cache.move_to_end(key) return value def _put_lru(self, cache: OrderedDict, key: tuple, value) -> None: cache[key] = value cache.move_to_end(key) while len(cache) > self._cross_attn_cache_size: cache.popitem(last=False) def clear_cross_attention_cache(self) -> None: self._projected_encoder_cache.clear() self._cross_attn_kv_cache.clear() def _cache_enabled(self) -> bool: return not self.training and not torch.is_grad_enabled() and not torch.compiler.is_compiling() def _project_encoder_hidden_states(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor: if not self._cache_enabled(): return self.condition_embedder.text_embedder(encoder_hidden_states) cache_key = self._tensor_cache_key(encoder_hidden_states) cached = self._get_from_lru(self._projected_encoder_cache, cache_key) if cached is not None: return cached projected = self.condition_embedder.text_embedder(encoder_hidden_states) self._put_lru(self._projected_encoder_cache, cache_key, projected) return projected def _get_cross_attn_key_values( self, encoder_hidden_states: torch.Tensor, ) -> list[tuple[torch.Tensor, torch.Tensor]] | None: if not self._cache_enabled(): return None cache_key = self._tensor_cache_key(encoder_hidden_states) cached = self._get_from_lru(self._cross_attn_kv_cache, cache_key) if cached is not None: return cached key_values = [block.attn2.project_kv(encoder_hidden_states) for block in self.blocks] self._put_lru(self._cross_attn_kv_cache, cache_key, key_values) return key_values def forward( self, hidden_states: torch.Tensor, timestep: torch.LongTensor, encoder_hidden_states: torch.Tensor, indices_hidden_states: torch.Tensor | None = None, indices_latents_history_short: torch.Tensor | None = None, indices_latents_history_mid: torch.Tensor | None = None, indices_latents_history_long: torch.Tensor | None = None, latents_history_short: torch.Tensor | None = None, latents_history_mid: torch.Tensor | None = None, latents_history_long: torch.Tensor | None = None, return_dict: bool = True, attention_kwargs: dict[str, Any] | None = None, ) -> torch.Tensor | Transformer2DModelOutput: batch_size = hidden_states.shape[0] p_t, p_h, p_w = self.config.patch_size # 1. Process noisy latents hidden_states = self.patch_embedding(hidden_states) _, _, post_patch_num_frames, post_patch_height, post_patch_width = hidden_states.shape if indices_hidden_states is None: indices_hidden_states = torch.arange(0, post_patch_num_frames).unsqueeze(0).expand(batch_size, -1) hidden_states = hidden_states.flatten(2).transpose(1, 2) rotary_emb = self.rope( frame_indices=indices_hidden_states, height=post_patch_height, width=post_patch_width, device=hidden_states.device, ) rotary_emb = rotary_emb.flatten(2).transpose(1, 2) original_context_length = hidden_states.shape[1] # 2. Process short history latents if latents_history_short is not None and indices_latents_history_short is not None: latents_history_short = latents_history_short.to(hidden_states) latents_history_short = self.patch_short(latents_history_short) _, _, _, H1, W1 = latents_history_short.shape latents_history_short = latents_history_short.flatten(2).transpose(1, 2) rotary_emb_history_short = self.rope( frame_indices=indices_latents_history_short, height=H1, width=W1, device=latents_history_short.device, ) rotary_emb_history_short = rotary_emb_history_short.flatten(2).transpose(1, 2) hidden_states = torch.cat([latents_history_short, hidden_states], dim=1) rotary_emb = torch.cat([rotary_emb_history_short, rotary_emb], dim=1) # 3. Process mid history latents if latents_history_mid is not None and indices_latents_history_mid is not None: latents_history_mid = latents_history_mid.to(hidden_states) latents_history_mid = pad_for_3d_conv(latents_history_mid, (2, 4, 4)) latents_history_mid = self.patch_mid(latents_history_mid) latents_history_mid = latents_history_mid.flatten(2).transpose(1, 2) rotary_emb_history_mid = self.rope( frame_indices=indices_latents_history_mid, height=H1, width=W1, device=latents_history_mid.device, ) rotary_emb_history_mid = pad_for_3d_conv(rotary_emb_history_mid, (2, 2, 2)) rotary_emb_history_mid = center_down_sample_3d(rotary_emb_history_mid, (2, 2, 2)) rotary_emb_history_mid = rotary_emb_history_mid.flatten(2).transpose(1, 2) hidden_states = torch.cat([latents_history_mid, hidden_states], dim=1) rotary_emb = torch.cat([rotary_emb_history_mid, rotary_emb], dim=1) # 4. Process long history latents if latents_history_long is not None and indices_latents_history_long is not None: latents_history_long = latents_history_long.to(hidden_states) latents_history_long = pad_for_3d_conv(latents_history_long, (4, 8, 8)) latents_history_long = self.patch_long(latents_history_long) latents_history_long = latents_history_long.flatten(2).transpose(1, 2) rotary_emb_history_long = self.rope( frame_indices=indices_latents_history_long, height=H1, width=W1, device=latents_history_long.device, ) rotary_emb_history_long = pad_for_3d_conv(rotary_emb_history_long, (4, 4, 4)) rotary_emb_history_long = center_down_sample_3d(rotary_emb_history_long, (4, 4, 4)) rotary_emb_history_long = rotary_emb_history_long.flatten(2).transpose(1, 2) hidden_states = torch.cat([latents_history_long, hidden_states], dim=1) rotary_emb = torch.cat([rotary_emb_history_long, rotary_emb], dim=1) history_context_length = hidden_states.shape[1] - original_context_length # 5. Compute timestep embeddings if indices_hidden_states is not None and self.zero_history_timestep: timestep_t0 = torch.zeros((1), dtype=timestep.dtype, device=timestep.device) temb_t0, timestep_proj_t0, _ = self.condition_embedder( timestep_t0, encoder_hidden_states, is_return_encoder_hidden_states=False ) temb_t0 = temb_t0.unsqueeze(1).expand(batch_size, history_context_length, -1) timestep_proj_t0 = ( timestep_proj_t0.unflatten(-1, (6, -1)) .view(1, 6, 1, -1) .expand(batch_size, -1, history_context_length, -1) ) temb, timestep_proj, _ = self.condition_embedder( timestep, encoder_hidden_states, is_return_encoder_hidden_states=False, ) encoder_hidden_states = self._project_encoder_hidden_states(encoder_hidden_states) timestep_proj = timestep_proj.unflatten(-1, (6, -1)) if indices_hidden_states is not None and not self.zero_history_timestep: main_repeat_size = hidden_states.shape[1] else: main_repeat_size = original_context_length temb = temb.view(batch_size, 1, -1).expand(batch_size, main_repeat_size, -1) timestep_proj = timestep_proj.view(batch_size, 6, 1, -1).expand(batch_size, 6, main_repeat_size, -1) if indices_hidden_states is not None and self.zero_history_timestep: temb = torch.cat([temb_t0, temb], dim=1) timestep_proj = torch.cat([timestep_proj_t0, timestep_proj], dim=2) if timestep_proj.ndim == 4: timestep_proj = timestep_proj.permute(0, 2, 1, 3) # 6. Transformer blocks hidden_states = hidden_states.contiguous() encoder_hidden_states = encoder_hidden_states.contiguous() rotary_emb = rotary_emb.contiguous() cross_attn_key_values = self._get_cross_attn_key_values(encoder_hidden_states) for block_idx, block in enumerate(self.blocks): cross_attn_key_value = None if cross_attn_key_values is None else cross_attn_key_values[block_idx] hidden_states = block( hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, original_context_length, cross_attn_key_value, ) # 7. Output normalization hidden_states = self.norm_out(hidden_states, temb, original_context_length) hidden_states = self.proj_out(hidden_states) # 8. Unpatchify hidden_states = hidden_states.reshape( batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1 ) hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) if not return_dict: return (output,) return Transformer2DModelOutput(sample=output) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights with QKV fusion, FFN remapping, and TP norm sharding.""" tp_rank = get_tensor_model_parallel_rank() tp_size = get_tensor_model_parallel_world_size() stacked_params_mapping = [ (".attn1.to_qkv", ".attn1.to_q", "q"), (".attn1.to_qkv", ".attn1.to_k", "k"), (".attn1.to_qkv", ".attn1.to_v", "v"), ] params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: original_name = name lookup_name = name # Handle QKV fusion for self-attention for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in original_name: continue lookup_name = original_name.replace(weight_name, param_name) param = params_dict[lookup_name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) break else: # FFN remapping: net.0 -> net_0, net.2 -> net_2 if ".ffn.net.0." in lookup_name: lookup_name = lookup_name.replace(".ffn.net.0.", ".ffn.net_0.") elif ".ffn.net.2." in lookup_name: lookup_name = lookup_name.replace(".ffn.net.2.", ".ffn.net_2.") # Output projection: to_out.0 -> to_out if ".to_out.0." in lookup_name: lookup_name = lookup_name.replace(".to_out.0.", ".to_out.") if lookup_name not in params_dict: logger.warning(f"Skipping weight {original_name} -> {lookup_name}") continue param = params_dict[lookup_name] # TP sharding for RMSNorm weights if tp_size > 1 and any( norm_name in lookup_name for norm_name in [ ".attn1.norm_q.", ".attn1.norm_k.", ".attn2.norm_q.", ".attn2.norm_k.", ] ): shard_size = loaded_weight.shape[0] // tp_size loaded_weight = loaded_weight[tp_rank * shard_size : (tp_rank + 1) * shard_size] # TP sharding for history_key_scale (per_head mode) if tp_size > 1 and ".history_key_scale" in lookup_name and loaded_weight.dim() == 1: shard_size = loaded_weight.shape[0] // tp_size if shard_size > 0: loaded_weight = loaded_weight[tp_rank * shard_size : (tp_rank + 1) * shard_size] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(original_name) loaded_params.add(lookup_name) return loaded_params