`fastvideo.models.dits.hunyuanvideo15.DistributedAttention`

DistributedAttention(num_heads: int, head_size: int, num_kv_heads: int | None = None, softmax_scale: float | None = None, causal: bool = False, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, prefix: str = '', **extra_impl_args)

Distributed attention layer.

`fastvideo.models.dits.hunyuanvideo15.DistributedAttention.forward`

forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, original_seq_len: int | None = None, replicated_q: torch.Tensor | None = None, replicated_k: torch.Tensor | None = None, replicated_v: torch.Tensor | None = None, freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, attn_metadata: AttentionMetadata | None = None) -> tuple[torch.Tensor, torch.Tensor | None]

Forward pass for distributed attention.

Parameters

q: torch.Tensor: Query tensor [batch_size, seq_len, num_heads, head_dim]
k: torch.Tensor: Key tensor [batch_size, seq_len, num_heads, head_dim]
v: torch.Tensor: Value tensor [batch_size, seq_len, num_heads, head_dim]
original_seq_len: int: Original (unpadded) full sequence length
replicated_q: Optional[torch.Tensor]: Replicated query tensor, typically for text tokens
replicated_k: Optional[torch.Tensor]: Replicated key tensor
replicated_v: Optional[torch.Tensor]: Replicated value tensor

Returns

tuple[torch.Tensor, torch.Tensor | None]: Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing: - o (torch.Tensor): Output tensor after attention for the main sequence - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided

`fastvideo.models.dits.hunyuanvideo15.LocalAttention`

LocalAttention(num_heads: int, head_size: int, num_kv_heads: int | None = None, softmax_scale: float | None = None, causal: bool = False, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, default_backend: AttentionBackendEnum | None = None, **extra_impl_args)

Attention layer.

`fastvideo.models.dits.hunyuanvideo15.LocalAttention.forward`

forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor

Apply local attention between query, key and value tensors.

Parameters

q: torch.Tensor: Query tensor of shape [batch_size, seq_len, num_heads, head_dim]
k: torch.Tensor: Key tensor of shape [batch_size, seq_len, num_heads, head_dim]
v: torch.Tensor: Value tensor of shape [batch_size, seq_len, num_heads, head_dim]

Returns

torch.Tensor: torch.Tensor: Output tensor after local attention

`fastvideo.models.dits.hunyuanvideo15.set_forward_context`

set_forward_context(current_timestep, attn_metadata, forward_batch: Optional[ForwardBatch] = None)

A context manager that stores the current forward context, can be attention metadata, etc. Here we can inject common logic for every model forward pass.

`fastvideo.models.dits.hunyuanvideo15.sequence_model_parallel_all_gather_with_unpad`

sequence_model_parallel_all_gather_with_unpad(input_: torch.Tensor, original_seq_len: int, dim: int = -1) -> torch.Tensor

All-gather the input tensor and remove padding.

Parameters

input_: torch.Tensor: Sharded (and possibly padded) tensor to gather
original_seq_len: int: Original sequence length before padding
dim: int: Dimension to gather along (default: -1)

Returns

Tensor: torch.Tensor: Gathered and unpadded tensor

`fastvideo.models.dits.hunyuanvideo15.sequence_model_parallel_shard`

sequence_model_parallel_shard(input_: torch.Tensor, dim: int = 1) -> tuple[torch.Tensor, int]

Shard the input tensor across model parallel group with optional padding.

Parameters

input_: torch.Tensor: Input tensor to shard
dim: int: Dimension to shard along (default: 1)

Returns

tuple: tuple[torch.Tensor, int]: (sharded_tensor, original_seq_len) - sharded_tensor: The sharded (and possibly padded) tensor - original_seq_len: Original sequence length before padding

`fastvideo.models.dits.hunyuanvideo15.HunyuanVideo15Config`

HunyuanVideo15Config(arch_config: DiTArchConfig = HunyuanVideo15ArchConfig(), prefix: str = 'Hunyuan15', quant_config: QuantizationConfig | None = None, *, _resolved_attention_backend: AttentionBackendEnum | None = None)

`fastvideo.models.dits.hunyuanvideo15.LayerNormScaleShift`

LayerNormScaleShift(hidden_size: int, norm_type: str = 'rms', eps: float = 1e-06, elementwise_affine: bool = False, dtype: torch.dtype = torch.float32, compute_dtype: torch.dtype | None = None, prefix: str = '')

Fused operation that combines LayerNorm with scale and shift operations. This reduces memory bandwidth by combining memory-bound operations.

`fastvideo.models.dits.hunyuanvideo15.LayerNormScaleShift.forward`

forward(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor, convert_modulation_dtype: bool = False) -> torch.Tensor

Apply ln followed by scale and shift in a single fused operation.

`fastvideo.models.dits.hunyuanvideo15.ScaleResidual`

ScaleResidual(prefix: str = '')

Applies gated residual connection.

`fastvideo.models.dits.hunyuanvideo15.ScaleResidual.forward`

forward(residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor

Apply gated residual connection.

`fastvideo.models.dits.hunyuanvideo15.ScaleResidualLayerNormScaleShift`

ScaleResidualLayerNormScaleShift(hidden_size: int, norm_type: str = 'rms', eps: float = 1e-06, elementwise_affine: bool = False, dtype: torch.dtype = torch.float32, compute_dtype: torch.dtype | None = None, prefix: str = '', fuse_inference: bool = False)

Fused operation that combines: 1. Gated residual connection 2. LayerNorm 3. Scale and shift operations

This reduces memory bandwidth by combining memory-bound operations.

`fastvideo.models.dits.hunyuanvideo15.ScaleResidualLayerNormScaleShift.forward`

forward(residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor | int, shift: torch.Tensor | None, scale: torch.Tensor | None, convert_modulation_dtype: bool = False) -> tuple[torch.Tensor, torch.Tensor]

Apply gated residual connection, followed by layernorm and scale/shift in a single fused operation.

Pass shift=None, scale=None (both, not just one) to skip the modulation entirely and return the normalized tensor as-is. Callers that need norm-only behavior (e.g. Wan's post-self-attn site) would otherwise pay two full-tensor elementwise passes computing normalized * (1.0 + 0) + 0.

Returns

torch.Tensor

Tuple containing:

torch.Tensor

normalized and modulated output

tuple[torch.Tensor, torch.Tensor]

residual value (value after residual connection but before normalization)

`fastvideo.models.dits.hunyuanvideo15.ReplicatedLinear`

ReplicatedLinear(input_size: int, output_size: int, bias: bool = True, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, prefix: str = '')

Replicated linear layer.

Parameters

input_size: int: input dimension of the linear layer.
output_size: int: output dimension of the linear layer.
bias: bool: If true, add bias.
skip_bias_add: bool: If true, skip adding bias but instead return it.
params_dtype: torch.dtype | None: Data type for the parameters.
quant_config: QuantizationConfig | None: Quantization configure.
prefix: str: The name of the layer in the state dict, including all parents (e.g. model.layers.0.qkv_proj)

`fastvideo.models.dits.hunyuanvideo15.ReplicatedLinear.weight_loader`

weight_loader(param: Parameter, loaded_weight: torch.Tensor) -> None

`fastvideo.models.dits.hunyuanvideo15.ReplicatedLinear.forward`

forward(x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]

`fastvideo.models.dits.hunyuanvideo15.ReplicatedLinear.extra_repr`

extra_repr() -> str

`fastvideo.models.dits.hunyuanvideo15.ReplicatedLinear.get_shape_mapping`

get_shape_mapping() -> dict

Get the mapping from (input_shape, output_shape) to layer types.

`fastvideo.models.dits.hunyuanvideo15.ReplicatedLinear.reset_shape_tracking`

reset_shape_tracking() -> None

Clear tracked shapes and layer type mappings.

`fastvideo.models.dits.hunyuanvideo15.ReplicatedLinear.print_shape_summary`

print_shape_summary() -> None

Log a summary of all unique shapes and their layer types.

`fastvideo.models.dits.hunyuanvideo15.MLP`

MLP(input_dim: int, mlp_hidden_dim: int, output_dim: int | None = None, bias: bool = True, act_type: str = 'gelu_pytorch_tanh', dtype: torch.dtype | None = None, prefix: str = '', quant_config: QuantizationConfig | None = None)

MLP for DiT blocks, NO gated linear units

`fastvideo.models.dits.hunyuanvideo15.MLP.forward`

forward(x: torch.Tensor) -> torch.Tensor

`fastvideo.models.dits.hunyuanvideo15.get_rotary_pos_embed`

get_rotary_pos_embed(rope_sizes, hidden_size, heads_num, rope_dim_list, rope_theta, theta_rescale_factor = 1.0, interpolation_factor = 1.0, shard_dim: int = 0, do_sp_sharding: bool = False, dtype: torch.dtype = torch.float32, start_frame: int = 0, use_real: bool = True) -> tuple[torch.Tensor, torch.Tensor]

Generate rotary positional embeddings for the given sizes.

Parameters

rope_sizes: Tuple of dimensions (t, h, w)
hidden_size: Hidden dimension size
heads_num: Number of attention heads
rope_dim_list: List of dimensions for each axis, or None
rope_theta: Base for frequency calculations
theta_rescale_factor: Rescale factor for theta. Defaults to 1.0
interpolation_factor: Factor to scale positions. Defaults to 1.0
shard_dim: int: Which dimension to shard for sequence parallelism. Defaults to 0.
do_sp_sharding: bool: Whether to shard the positional embeddings for sequence parallelism. Defaults to False.
use_real: bool: If True, output full head_dim for rotate_half style RoPE. Defaults to True.

Returns

tuple[torch.Tensor, torch.Tensor]: Tuple of (cos, sin) tensors for rotary embeddings. Shape [S, D] if use_real, [S, D/2] otherwise.

`fastvideo.models.dits.hunyuanvideo15.ModulateProjection`

ModulateProjection(hidden_size: int, factor: int = 2, act_layer: str = 'silu', dtype: torch.dtype | None = None, prefix: str = '')

Modulation layer for DiT blocks.

`fastvideo.models.dits.hunyuanvideo15.ModulateProjection.forward`

forward(x: torch.Tensor) -> torch.Tensor

`fastvideo.models.dits.hunyuanvideo15.PatchEmbed`

PatchEmbed(patch_size = 16, in_chans = 3, embed_dim = 768, norm_layer = None, flatten = True, bias = True, dtype = None, prefix: str = '')

2D Image to Patch Embedding

Image to Patch Embedding using Conv2d

A convolution based approach to patchifying a 2D image w/ embedding projection.

Based on the impl in https://github.com/google-research/vision_transformer

Remove the _assert function in forward function to be compatible with multi-resolution images.

`fastvideo.models.dits.hunyuanvideo15.PatchEmbed.forward`

forward(x)

`fastvideo.models.dits.hunyuanvideo15.TimestepEmbedder`

TimestepEmbedder(hidden_size, act_layer = 'silu', frequency_embedding_size = 256, max_period = 10000, dtype = None, freq_dtype = torch.float32, prefix: str = '')

Embeds scalar timesteps into vector representations.

`fastvideo.models.dits.hunyuanvideo15.TimestepEmbedder.forward`

forward(t: torch.Tensor, timestep_seq_len: int | None = None) -> torch.Tensor

`fastvideo.models.dits.hunyuanvideo15.unpatchify`

unpatchify(x, t, h, w, patch_size, channels) -> torch.Tensor

Convert patched representation back to image space.

Parameters

x: Tensor of shape [B, THW, CP_tP_h*P_w]
t, h, w: Temporal and spatial dimensions

Returns

torch.Tensor: Unpatchified tensor of shape [B, C, TP_t, HP_h, W*P_w]

`fastvideo.models.dits.hunyuanvideo15.BaseDiT`

BaseDiT(config: DiTConfig, hf_config: dict[str, Any], **kwargs)

`fastvideo.models.dits.hunyuanvideo15.BaseDiT.forward`

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, guidance = None, **kwargs) -> torch.Tensor

`fastvideo.models.dits.hunyuanvideo15.BaseDiT.preprocess_loaded_weight`

preprocess_loaded_weight(name: str, tensor: torch.Tensor) -> torch.Tensor

Adapt a checkpoint tensor before sharding; default models are unchanged.

`fastvideo.models.dits.hunyuanvideo15.BaseDiT.inference_config_overrides`

inference_config_overrides() -> dict[str, Any]

Extra config.json keys a deployable checkpoint of this model needs.

`fastvideo.models.dits.hunyuanvideo15.BaseDiT.inference_export_expected_shape`

inference_export_expected_shape(output_key: str, base_shape: tuple[int, ...]) -> tuple[int, ...]

The exported shape of output_key given the base checkpoint's shape.

`fastvideo.models.dits.hunyuanvideo15.AttentionBackendEnum`

AttentionBackendEnum()

`fastvideo.models.dits.hunyuanvideo15.init_logger`

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

`fastvideo.models.dits.hunyuanvideo15.get_sp_world_size`

get_sp_world_size() -> int

Return world size for the sequence model parallel group.

`fastvideo.models.dits.hunyuanvideo15.HunyuanRMSNorm`

HunyuanRMSNorm(dim: int, elementwise_affine = True, eps: float = 1e-06, device = None, dtype = None)

Initialize the RMSNorm normalization layer.

Parameters

dim: int: The dimension of the input tensor.
eps: float: A small value added to the denominator for numerical stability. Default is 1e-6.

Attributes

eps: float: A small value added to the denominator for numerical stability.
weight: nn.Parameter: Learnable scaling parameter.

Source

class HunyuanRMSNorm(nn.Module):

    def __init__(
        self,
        dim: int,
        elementwise_affine=True,
        eps: float = 1e-6,
        device=None,
        dtype=None,
    ):
        """
        Initialize the RMSNorm normalization layer.

        Args:
            dim (int): The dimension of the input tensor.
            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.

        Attributes:
            eps (float): A small value added to the denominator for numerical stability.
            weight (nn.Parameter): Learnable scaling parameter.

        """
        factory_kwargs = {"device": device, "dtype": dtype}
        super().__init__()
        self.eps = eps
        if elementwise_affine:
            self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))

    def _norm(self, x) -> torch.Tensor:
        """
        Apply the RMSNorm normalization to the input tensor.

        Args:
            x (torch.Tensor): The input tensor.

        Returns:
            torch.Tensor: The normalized tensor.

        """
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        """
        Forward pass through the RMSNorm layer.

        Args:
            x (torch.Tensor): The input tensor.

        Returns:
            torch.Tensor: The output tensor after applying RMSNorm.

        """
        output = self._norm(x.float()).type_as(x)
        if hasattr(self, "weight"):
            output = output * self.weight
        return output

`fastvideo.models.dits.hunyuanvideo15.HunyuanRMSNorm.forward`

forward(x)

Forward pass through the RMSNorm layer.

Parameters

x: torch.Tensor: The input tensor.

Returns

: torch.Tensor: The output tensor after applying RMSNorm.

Source

def forward(self, x):
    """
    Forward pass through the RMSNorm layer.

    Args:
        x (torch.Tensor): The input tensor.

    Returns:
        torch.Tensor: The output tensor after applying RMSNorm.

    """
    output = self._norm(x.float()).type_as(x)
    if hasattr(self, "weight"):
        output = output * self.weight
    return output

`fastvideo.models.dits.hunyuanvideo15.HunyuanVideo15TimeEmbedding`

HunyuanVideo15TimeEmbedding(embedding_dim: int, use_meanflow: bool = False)

Time embedding for HunyuanVideo 1.5.

Supports standard timestep embedding and optional reference timestep embedding for MeanFlow-based super-resolution models.

Parameters

embedding_dim: `int`: The dimension of the output embedding.

Source

class HunyuanVideo15TimeEmbedding(nn.Module):
    r"""
    Time embedding for HunyuanVideo 1.5.

    Supports standard timestep embedding and optional reference timestep embedding for MeanFlow-based super-resolution
    models.

    Args:
        embedding_dim (`int`):
            The dimension of the output embedding.
    """

    def __init__(self, embedding_dim: int, use_meanflow: bool = False):
        super().__init__()

        self.timestep_embedder = TimestepEmbedder(hidden_size=embedding_dim)

        self.use_meanflow = use_meanflow
        self.time_proj_r = None
        self.timestep_embedder_r = None
        if use_meanflow:
            self.timestep_embedder_r = TimestepEmbedder(hidden_size=embedding_dim)

    def forward(
        self,
        timestep: torch.Tensor,
        timestep_r: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        timesteps_emb = self.timestep_embedder(timestep)

        if timestep_r is not None:
            timesteps_emb_r = self.timestep_embedder_r(timestep_r)
            timesteps_emb = timesteps_emb + timesteps_emb_r

        return timesteps_emb

`fastvideo.models.dits.hunyuanvideo15.HunyuanVideo15TimeEmbedding.forward`

forward(timestep: torch.Tensor, timestep_r: Optional[torch.Tensor] = None) -> torch.Tensor

Source

def forward(
    self,
    timestep: torch.Tensor,
    timestep_r: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    timesteps_emb = self.timestep_embedder(timestep)

    if timestep_r is not None:
        timesteps_emb_r = self.timestep_embedder_r(timestep_r)
        timesteps_emb = timesteps_emb + timesteps_emb_r

    return timesteps_emb

`fastvideo.models.dits.hunyuanvideo15.HunyuanVideo15ByT5TextProjection`

HunyuanVideo15ByT5TextProjection(in_features: int, hidden_size: int, out_features: int)

Source

class HunyuanVideo15ByT5TextProjection(nn.Module):

    def __init__(self, in_features: int, hidden_size: int, out_features: int):
        super().__init__()
        self.norm = nn.LayerNorm(in_features)
        self.linear_1 = nn.Linear(in_features, hidden_size)
        self.linear_2 = nn.Linear(hidden_size, hidden_size)
        self.linear_3 = nn.Linear(hidden_size, out_features)
        self.act_fn = nn.GELU()

    def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.norm(encoder_hidden_states)
        hidden_states = self.linear_1(hidden_states)
        hidden_states = self.act_fn(hidden_states)
        hidden_states = self.linear_2(hidden_states)
        hidden_states = self.act_fn(hidden_states)
        hidden_states = self.linear_3(hidden_states)
        return hidden_states

`fastvideo.models.dits.hunyuanvideo15.HunyuanVideo15ByT5TextProjection.forward`

forward(encoder_hidden_states: torch.Tensor) -> torch.Tensor

Source

def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
    hidden_states = self.norm(encoder_hidden_states)
    hidden_states = self.linear_1(hidden_states)
    hidden_states = self.act_fn(hidden_states)
    hidden_states = self.linear_2(hidden_states)
    hidden_states = self.act_fn(hidden_states)
    hidden_states = self.linear_3(hidden_states)
    return hidden_states

`fastvideo.models.dits.hunyuanvideo15.HunyuanVideo15ImageProjection`

HunyuanVideo15ImageProjection(in_channels: int, hidden_size: int)

Source

class HunyuanVideo15ImageProjection(nn.Module):

    def __init__(self, in_channels: int, hidden_size: int):
        super().__init__()
        self.norm_in = nn.LayerNorm(in_channels)
        self.linear_1 = nn.Linear(in_channels, in_channels)
        self.act_fn = nn.GELU()
        self.linear_2 = nn.Linear(in_channels, hidden_size)
        self.norm_out = nn.LayerNorm(hidden_size)

    def forward(self, image_embeds: torch.Tensor) -> torch.Tensor:
        hidden_states = self.norm_in(image_embeds)
        hidden_states = self.linear_1(hidden_states)
        hidden_states = self.act_fn(hidden_states)
        hidden_states = self.linear_2(hidden_states)
        hidden_states = self.norm_out(hidden_states)
        return hidden_states

`fastvideo.models.dits.hunyuanvideo15.HunyuanVideo15ImageProjection.forward`

forward(image_embeds: torch.Tensor) -> torch.Tensor

Source

def forward(self, image_embeds: torch.Tensor) -> torch.Tensor:
    hidden_states = self.norm_in(image_embeds)
    hidden_states = self.linear_1(hidden_states)
    hidden_states = self.act_fn(hidden_states)
    hidden_states = self.linear_2(hidden_states)
    hidden_states = self.norm_out(hidden_states)
    return hidden_states

`fastvideo.models.dits.hunyuanvideo15.MMDoubleStreamBlock`

MMDoubleStreamBlock(hidden_size: int, num_attention_heads: int, mlp_ratio: float, dtype: torch.dtype | None = None, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, prefix: str = '')

A multimodal DiT block with separate modulation for text and image/video, using distributed attention and linear layers.

Source

class MMDoubleStreamBlock(nn.Module):
    """
    A multimodal DiT block with separate modulation for text and image/video,
    using distributed attention and linear layers.
    """

    def __init__(
        self,
        hidden_size: int,
        num_attention_heads: int,
        mlp_ratio: float,
        dtype: torch.dtype | None = None,
        supported_attention_backends: tuple[AttentionBackendEnum, ...]
        | None = None,
        prefix: str = "",
    ):
        super().__init__()

        self.deterministic = False
        self.num_attention_heads = num_attention_heads
        head_dim = hidden_size // num_attention_heads
        mlp_hidden_dim = int(hidden_size * mlp_ratio)

        # Image modulation components
        self.img_mod = ModulateProjection(
            hidden_size,
            factor=6,
            act_layer="silu",
            dtype=dtype,
            prefix=f"{prefix}.img_mod",
        )

        # Fused operations for image stream
        self.img_attn_norm = LayerNormScaleShift(hidden_size, norm_type="layer", elementwise_affine=False, dtype=dtype)
        self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(hidden_size,
                                                                           norm_type="layer",
                                                                           elementwise_affine=False,
                                                                           dtype=dtype)
        self.img_mlp_residual = ScaleResidual()

        # Image attention components
        self.img_attn_qkv = ReplicatedLinear(hidden_size,
                                             hidden_size * 3,
                                             bias=True,
                                             params_dtype=dtype,
                                             prefix=f"{prefix}.img_attn_qkv")

        self.img_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
        self.img_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)

        self.img_attn_proj = ReplicatedLinear(hidden_size,
                                              hidden_size,
                                              bias=True,
                                              params_dtype=dtype,
                                              prefix=f"{prefix}.img_attn_proj")

        self.img_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype, prefix=f"{prefix}.img_mlp")

        # Text modulation components
        self.txt_mod = ModulateProjection(
            hidden_size,
            factor=6,
            act_layer="silu",
            dtype=dtype,
            prefix=f"{prefix}.txt_mod",
        )

        # Fused operations for text stream
        self.txt_attn_norm = LayerNormScaleShift(hidden_size, norm_type="layer", elementwise_affine=False, dtype=dtype)
        self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(hidden_size,
                                                                           norm_type="layer",
                                                                           elementwise_affine=False,
                                                                           dtype=dtype)
        self.txt_mlp_residual = ScaleResidual()

        # Text attention components
        self.txt_attn_qkv = ReplicatedLinear(hidden_size, hidden_size * 3, bias=True, params_dtype=dtype)

        # QK norm layers for text
        self.txt_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
        self.txt_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)

        self.txt_attn_proj = ReplicatedLinear(hidden_size, hidden_size, bias=True, params_dtype=dtype)

        self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype)

        # Distributed attention
        self.attn = DistributedAttention(num_heads=num_attention_heads,
                                         head_size=head_dim,
                                         causal=False,
                                         supported_attention_backends=supported_attention_backends,
                                         prefix=f"{prefix}.attn")

    def forward(
        self,
        img: torch.Tensor,
        txt: torch.Tensor,
        vec: torch.Tensor,
        freqs_cis: tuple,
        original_seq_len: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # Broadcast per-example conditioning over tokens, not over the batch.
        img_mod_outputs = self.img_mod(vec).unsqueeze(1)
        (
            img_attn_shift,
            img_attn_scale,
            img_attn_gate,
            img_mlp_shift,
            img_mlp_scale,
            img_mlp_gate,
        ) = torch.chunk(img_mod_outputs, 6, dim=-1)

        txt_mod_outputs = self.txt_mod(vec).unsqueeze(1)
        (
            txt_attn_shift,
            txt_attn_scale,
            txt_attn_gate,
            txt_mlp_shift,
            txt_mlp_scale,
            txt_mlp_gate,
        ) = torch.chunk(txt_mod_outputs, 6, dim=-1)

        # Prepare image for attention using fused operation
        img_attn_input = self.img_attn_norm(img, img_attn_shift, img_attn_scale)
        # Get QKV for image
        img_qkv, _ = self.img_attn_qkv(img_attn_input)
        batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1]

        # Split QKV
        img_qkv = img_qkv.view(batch_size, image_seq_len, 3, self.num_attention_heads, -1)
        img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, 2]

        # Apply QK-Norm if needed
        img_q = self.img_attn_q_norm(img_q).to(img_v)
        img_k = self.img_attn_k_norm(img_k).to(img_v)

        # Prepare text for attention using fused operation
        txt_attn_input = self.txt_attn_norm(txt, txt_attn_shift, txt_attn_scale)

        # Get QKV for text
        txt_qkv, _ = self.txt_attn_qkv(txt_attn_input)
        batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1]

        # Split QKV
        txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3, self.num_attention_heads, -1)
        txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, 2]
        # Apply QK-Norm if needed
        txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype)
        txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype)

        img_attn, txt_attn = self.attn(img_q, img_k, img_v, original_seq_len, txt_q, txt_k, txt_v, freqs_cis=freqs_cis)

        img_attn_out, _ = self.img_attn_proj(img_attn.reshape(batch_size, image_seq_len, -1))
        # Use fused operation for residual connection, normalization, and modulation
        img_mlp_input, img_residual = self.img_attn_residual_mlp_norm(img, img_attn_out, img_attn_gate, img_mlp_shift,
                                                                      img_mlp_scale)

        # Process image MLP
        img_mlp_out = self.img_mlp(img_mlp_input)
        img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate)

        # Process text attention output
        txt_attn_out, _ = self.txt_attn_proj(txt_attn.reshape(batch_size, text_seq_len, -1))

        # Use fused operation for residual connection, normalization, and modulation
        txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm(txt, txt_attn_out, txt_attn_gate, txt_mlp_shift,
                                                                      txt_mlp_scale)

        # Process text MLP
        txt_mlp_out = self.txt_mlp(txt_mlp_input)
        txt = self.txt_mlp_residual(txt_residual, txt_mlp_out, txt_mlp_gate)

        return img, txt

`fastvideo.models.dits.hunyuanvideo15.MMDoubleStreamBlock.forward`

forward(img: torch.Tensor, txt: torch.Tensor, vec: torch.Tensor, freqs_cis: tuple, original_seq_len: int) -> tuple[torch.Tensor, torch.Tensor]

Source

def forward(
    self,
    img: torch.Tensor,
    txt: torch.Tensor,
    vec: torch.Tensor,
    freqs_cis: tuple,
    original_seq_len: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    # Broadcast per-example conditioning over tokens, not over the batch.
    img_mod_outputs = self.img_mod(vec).unsqueeze(1)
    (
        img_attn_shift,
        img_attn_scale,
        img_attn_gate,
        img_mlp_shift,
        img_mlp_scale,
        img_mlp_gate,
    ) = torch.chunk(img_mod_outputs, 6, dim=-1)

    txt_mod_outputs = self.txt_mod(vec).unsqueeze(1)
    (
        txt_attn_shift,
        txt_attn_scale,
        txt_attn_gate,
        txt_mlp_shift,
        txt_mlp_scale,
        txt_mlp_gate,
    ) = torch.chunk(txt_mod_outputs, 6, dim=-1)

    # Prepare image for attention using fused operation
    img_attn_input = self.img_attn_norm(img, img_attn_shift, img_attn_scale)
    # Get QKV for image
    img_qkv, _ = self.img_attn_qkv(img_attn_input)
    batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1]

    # Split QKV
    img_qkv = img_qkv.view(batch_size, image_seq_len, 3, self.num_attention_heads, -1)
    img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, 2]

    # Apply QK-Norm if needed
    img_q = self.img_attn_q_norm(img_q).to(img_v)
    img_k = self.img_attn_k_norm(img_k).to(img_v)

    # Prepare text for attention using fused operation
    txt_attn_input = self.txt_attn_norm(txt, txt_attn_shift, txt_attn_scale)

    # Get QKV for text
    txt_qkv, _ = self.txt_attn_qkv(txt_attn_input)
    batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1]

    # Split QKV
    txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3, self.num_attention_heads, -1)
    txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, 2]
    # Apply QK-Norm if needed
    txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype)
    txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype)

    img_attn, txt_attn = self.attn(img_q, img_k, img_v, original_seq_len, txt_q, txt_k, txt_v, freqs_cis=freqs_cis)

    img_attn_out, _ = self.img_attn_proj(img_attn.reshape(batch_size, image_seq_len, -1))
    # Use fused operation for residual connection, normalization, and modulation
    img_mlp_input, img_residual = self.img_attn_residual_mlp_norm(img, img_attn_out, img_attn_gate, img_mlp_shift,
                                                                  img_mlp_scale)

    # Process image MLP
    img_mlp_out = self.img_mlp(img_mlp_input)
    img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate)

    # Process text attention output
    txt_attn_out, _ = self.txt_attn_proj(txt_attn.reshape(batch_size, text_seq_len, -1))

    # Use fused operation for residual connection, normalization, and modulation
    txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm(txt, txt_attn_out, txt_attn_gate, txt_mlp_shift,
                                                                  txt_mlp_scale)

    # Process text MLP
    txt_mlp_out = self.txt_mlp(txt_mlp_input)
    txt = self.txt_mlp_residual(txt_residual, txt_mlp_out, txt_mlp_gate)

    return img, txt

`fastvideo.models.dits.hunyuanvideo15.HunyuanVideo15Transformer3DModel`

HunyuanVideo15Transformer3DModel(config: HunyuanVideo15Config, hf_config: dict[str, Any])

A Transformer model for video-like data used in HunyuanVideo1.5.

Source

class HunyuanVideo15Transformer3DModel(BaseDiT):
    r"""
    A Transformer model for video-like data used in [HunyuanVideo1.5](https://huggingface.co/tencent/HunyuanVideo1.5).
    """

    # shard single stream, double stream blocks, and refiner_blocks
    _fsdp_shard_conditions = HunyuanVideo15Config()._fsdp_shard_conditions
    _compile_conditions = HunyuanVideo15Config()._compile_conditions
    _supported_attention_backends = HunyuanVideo15Config()._supported_attention_backends
    param_names_mapping = HunyuanVideo15Config().param_names_mapping
    reverse_param_names_mapping = HunyuanVideo15Config().reverse_param_names_mapping
    lora_param_names_mapping = HunyuanVideo15Config().lora_param_names_mapping

    def __init__(
        self,
        config: HunyuanVideo15Config,
        hf_config: dict[str, Any],
    ) -> None:
        super().__init__(config=config, hf_config=hf_config)

        self.hidden_size = config.hidden_size
        self.num_attention_heads = config.num_attention_heads
        self.num_channels_latents = config.num_channels_latents
        self.out_channels = config.out_channels or config.in_channels
        self.patch_size = (config.patch_size_t, config.patch_size, config.patch_size)

        # 1. Latent and condition embedders
        self.img_in = PatchEmbed(self.patch_size,
                                 config.in_channels,
                                 self.hidden_size,
                                 prefix=f"{config.prefix}.img_in")
        self.image_embedder = HunyuanVideo15ImageProjection(config.image_embed_dim, self.hidden_size)

        self.txt_in = SingleTokenRefiner(config.text_embed_dim,
                                         self.hidden_size,
                                         config.num_attention_heads,
                                         depth=config.num_refiner_layers,
                                         dtype=None,
                                         prefix=f"{config.prefix}.txt_in")

        self.txt_in_2 = HunyuanVideo15ByT5TextProjection(config.text_embed_2_dim, 2048, self.hidden_size)

        self.time_in = HunyuanVideo15TimeEmbedding(self.hidden_size, use_meanflow=config.use_meanflow)

        self.cond_type_embed = nn.Embedding(3, self.hidden_size)

        # 3. Dual stream transformer blocks

        self.double_blocks = nn.ModuleList([
            MMDoubleStreamBlock(hidden_size=self.hidden_size,
                                num_attention_heads=config.num_attention_heads,
                                mlp_ratio=config.mlp_ratio,
                                dtype=None,
                                supported_attention_backends=self._supported_attention_backends,
                                prefix=f"{config.prefix}.double_blocks.{i}") for i in range(config.num_layers)
        ])

        # 5. Output projection
        self.final_layer = FinalLayer(self.hidden_size,
                                      self.patch_size,
                                      self.out_channels,
                                      prefix=f"{config.prefix}.final_layer")

        self.gradient_checkpointing = False

        self.__post_init__()

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: List[torch.Tensor],
        timestep: torch.LongTensor,
        encoder_hidden_states_image: List[torch.Tensor],
        guidance: Optional[torch.Tensor] = None,
        timestep_r: Optional[torch.LongTensor] = None,
        attention_kwargs: Optional[Dict[str, Any]] = None,
    ):
        encoder_hidden_states_image = encoder_hidden_states_image[0]
        encoder_hidden_states, encoder_hidden_states_2 = encoder_hidden_states
        batch_size, num_channels, num_frames, height, width = hidden_states.shape
        p_t, p_h, p_w = self.config.patch_size_t, self.config.patch_size, self.config.patch_size
        post_patch_num_frames = num_frames // p_t
        post_patch_height = height // p_h
        post_patch_width = width // p_w

        # 1. RoPE
        # Get rotary embeddings
        freqs_cos, freqs_sin = get_rotary_pos_embed((post_patch_num_frames, post_patch_height, post_patch_width),
                                                    self.hidden_size, self.num_attention_heads,
                                                    self.config.rope_axes_dim, self.config.rope_theta)
        freqs_cos = freqs_cos.to(hidden_states.device)
        freqs_sin = freqs_sin.to(hidden_states.device)
        freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None

        # 2. Conditional embeddings
        temb = self.time_in(timestep, timestep_r=timestep_r)

        hidden_states = self.img_in(hidden_states)
        hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)

        current_seq_len = hidden_states.shape[1]
        sp_world_size = get_sp_world_size()
        padded_seq_len = current_seq_len * sp_world_size

        encoder_hidden_states = self.txt_in(encoder_hidden_states, timestep)

        encoder_hidden_states_cond_emb = self.cond_type_embed(
            torch.zeros_like(encoder_hidden_states[:, :, 0], dtype=torch.long))
        encoder_hidden_states = encoder_hidden_states + encoder_hidden_states_cond_emb

        # byt5 text embedding
        encoder_hidden_states_2 = self.txt_in_2(encoder_hidden_states_2)

        encoder_hidden_states_2_cond_emb = self.cond_type_embed(
            torch.ones_like(encoder_hidden_states_2[:, :, 0], dtype=torch.long))
        encoder_hidden_states_2 = encoder_hidden_states_2 + encoder_hidden_states_2_cond_emb

        # image embed
        encoder_hidden_states_3 = self.image_embedder(encoder_hidden_states_image)
        is_t2v = torch.all(encoder_hidden_states_image == 0)
        encoder_hidden_states_3_cond_emb = self.cond_type_embed(2 * torch.ones_like(
            encoder_hidden_states_3[:, :, 0],
            dtype=torch.long,
        ))
        encoder_hidden_states_3 = encoder_hidden_states_3 + encoder_hidden_states_3_cond_emb

        encoder_hidden_states = torch.cat(
            [encoder_hidden_states_2, encoder_hidden_states], dim=1) if is_t2v else torch.cat(
                [encoder_hidden_states_3, encoder_hidden_states_2, encoder_hidden_states], dim=1)

        # 4. Transformer blocks
        if torch.is_grad_enabled() and self.gradient_checkpointing:
            for block in self.double_blocks:
                hidden_states, encoder_hidden_states = self._gradient_checkpointing_func(
                    block, hidden_states, encoder_hidden_states, temb, freqs_cis, original_seq_len)

        else:
            for block in self.double_blocks:
                hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, freqs_cis,
                                                             original_seq_len)

        # Final layer processing
        if get_sp_world_size() > 1:
            hidden_states = hidden_states.contiguous()
        hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1)
        hidden_states = self.final_layer(hidden_states, temb)
        # Unpatchify to get original shape
        hidden_states = unpatchify(hidden_states, post_patch_num_frames, post_patch_height, post_patch_width,
                                   self.patch_size, self.out_channels)

        return hidden_states

`fastvideo.models.dits.hunyuanvideo15.HunyuanVideo15Transformer3DModel.forward`

forward(hidden_states: torch.Tensor, encoder_hidden_states: List[torch.Tensor], timestep: torch.LongTensor, encoder_hidden_states_image: List[torch.Tensor], guidance: Optional[torch.Tensor] = None, timestep_r: Optional[torch.LongTensor] = None, attention_kwargs: Optional[Dict[str, Any]] = None)

Source

def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: List[torch.Tensor],
    timestep: torch.LongTensor,
    encoder_hidden_states_image: List[torch.Tensor],
    guidance: Optional[torch.Tensor] = None,
    timestep_r: Optional[torch.LongTensor] = None,
    attention_kwargs: Optional[Dict[str, Any]] = None,
):
    encoder_hidden_states_image = encoder_hidden_states_image[0]
    encoder_hidden_states, encoder_hidden_states_2 = encoder_hidden_states
    batch_size, num_channels, num_frames, height, width = hidden_states.shape
    p_t, p_h, p_w = self.config.patch_size_t, self.config.patch_size, self.config.patch_size
    post_patch_num_frames = num_frames // p_t
    post_patch_height = height // p_h
    post_patch_width = width // p_w

    # 1. RoPE
    # Get rotary embeddings
    freqs_cos, freqs_sin = get_rotary_pos_embed((post_patch_num_frames, post_patch_height, post_patch_width),
                                                self.hidden_size, self.num_attention_heads,
                                                self.config.rope_axes_dim, self.config.rope_theta)
    freqs_cos = freqs_cos.to(hidden_states.device)
    freqs_sin = freqs_sin.to(hidden_states.device)
    freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None

    # 2. Conditional embeddings
    temb = self.time_in(timestep, timestep_r=timestep_r)

    hidden_states = self.img_in(hidden_states)
    hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)

    current_seq_len = hidden_states.shape[1]
    sp_world_size = get_sp_world_size()
    padded_seq_len = current_seq_len * sp_world_size

    encoder_hidden_states = self.txt_in(encoder_hidden_states, timestep)

    encoder_hidden_states_cond_emb = self.cond_type_embed(
        torch.zeros_like(encoder_hidden_states[:, :, 0], dtype=torch.long))
    encoder_hidden_states = encoder_hidden_states + encoder_hidden_states_cond_emb

    # byt5 text embedding
    encoder_hidden_states_2 = self.txt_in_2(encoder_hidden_states_2)

    encoder_hidden_states_2_cond_emb = self.cond_type_embed(
        torch.ones_like(encoder_hidden_states_2[:, :, 0], dtype=torch.long))
    encoder_hidden_states_2 = encoder_hidden_states_2 + encoder_hidden_states_2_cond_emb

    # image embed
    encoder_hidden_states_3 = self.image_embedder(encoder_hidden_states_image)
    is_t2v = torch.all(encoder_hidden_states_image == 0)
    encoder_hidden_states_3_cond_emb = self.cond_type_embed(2 * torch.ones_like(
        encoder_hidden_states_3[:, :, 0],
        dtype=torch.long,
    ))
    encoder_hidden_states_3 = encoder_hidden_states_3 + encoder_hidden_states_3_cond_emb

    encoder_hidden_states = torch.cat(
        [encoder_hidden_states_2, encoder_hidden_states], dim=1) if is_t2v else torch.cat(
            [encoder_hidden_states_3, encoder_hidden_states_2, encoder_hidden_states], dim=1)

    # 4. Transformer blocks
    if torch.is_grad_enabled() and self.gradient_checkpointing:
        for block in self.double_blocks:
            hidden_states, encoder_hidden_states = self._gradient_checkpointing_func(
                block, hidden_states, encoder_hidden_states, temb, freqs_cis, original_seq_len)

    else:
        for block in self.double_blocks:
            hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, freqs_cis,
                                                         original_seq_len)

    # Final layer processing
    if get_sp_world_size() > 1:
        hidden_states = hidden_states.contiguous()
    hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1)
    hidden_states = self.final_layer(hidden_states, temb)
    # Unpatchify to get original shape
    hidden_states = unpatchify(hidden_states, post_patch_num_frames, post_patch_height, post_patch_width,
                               self.patch_size, self.out_channels)

    return hidden_states

`fastvideo.models.dits.hunyuanvideo15.SingleTokenRefiner`

SingleTokenRefiner(in_channels, hidden_size, num_attention_heads, depth = 2, qkv_bias = True, dtype = None, prefix: str = '')

A token refiner that processes text embeddings with attention to improve their representation for cross-attention with image features.

Source

class SingleTokenRefiner(nn.Module):
    """
    A token refiner that processes text embeddings with attention to improve
    their representation for cross-attention with image features.
    """

    def __init__(
        self,
        in_channels,
        hidden_size,
        num_attention_heads,
        depth=2,
        qkv_bias=True,
        dtype=None,
        prefix: str = "",
    ) -> None:
        super().__init__()

        # Input projection
        # self.input_embedder = ReplicatedLinear(
        #     in_channels,
        #     hidden_size,
        #     bias=True,
        #     params_dtype=dtype,
        #     prefix=f"{prefix}.input_embedder")
        self.input_embedder = nn.Linear(in_channels, hidden_size, bias=True)

        # Timestep embedding
        self.t_embedder = TimestepEmbedder(hidden_size, act_layer="silu", dtype=dtype, prefix=f"{prefix}.t_embedder")

        # Context embedding
        self.c_embedder = MLP(in_channels,
                              hidden_size,
                              hidden_size,
                              act_type="silu",
                              dtype=dtype,
                              prefix=f"{prefix}.c_embedder")

        # Refiner blocks
        self.refiner_blocks = nn.ModuleList([
            IndividualTokenRefinerBlock(
                hidden_size,
                num_attention_heads,
                qkv_bias=qkv_bias,
                dtype=dtype,
                prefix=f"{prefix}.refiner_blocks.{i}",
            ) for i in range(depth)
        ])

    def forward(self, x, t, mask=None):
        # Get timestep embeddings
        timestep_aware_representations = self.t_embedder(t)

        original_dtype = x.dtype
        if mask is None:
            context_aware_representations = x.mean(dim=1)
        else:
            mask_float = mask.float().unsqueeze(-1)  # [B, L, 1]
            context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1)

        context_aware_representations = self.c_embedder(context_aware_representations)
        c = timestep_aware_representations + context_aware_representations
        # Project input
        x = self.input_embedder(x)
        # Process through refiner blocks
        for block in self.refiner_blocks:
            x = block(x, c, mask=mask)
        return x

`fastvideo.models.dits.hunyuanvideo15.SingleTokenRefiner.forward`

forward(x, t, mask = None)

Source

def forward(self, x, t, mask=None):
    # Get timestep embeddings
    timestep_aware_representations = self.t_embedder(t)

    original_dtype = x.dtype
    if mask is None:
        context_aware_representations = x.mean(dim=1)
    else:
        mask_float = mask.float().unsqueeze(-1)  # [B, L, 1]
        context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1)

    context_aware_representations = self.c_embedder(context_aware_representations)
    c = timestep_aware_representations + context_aware_representations
    # Project input
    x = self.input_embedder(x)
    # Process through refiner blocks
    for block in self.refiner_blocks:
        x = block(x, c, mask=mask)
    return x

`fastvideo.models.dits.hunyuanvideo15.IndividualTokenRefinerBlock`

IndividualTokenRefinerBlock(hidden_size, num_attention_heads, mlp_ratio = 4.0, qkv_bias = True, dtype = None, prefix: str = '')

A transformer block for refining individual tokens with self-attention.

Source

class IndividualTokenRefinerBlock(nn.Module):
    """
    A transformer block for refining individual tokens with self-attention.
    """

    def __init__(
        self,
        hidden_size,
        num_attention_heads,
        mlp_ratio=4.0,
        qkv_bias=True,
        dtype=None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.num_attention_heads = num_attention_heads
        mlp_hidden_dim = int(hidden_size * mlp_ratio)

        # Normalization and attention
        self.norm1 = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=True, dtype=dtype)

        self.self_attn_qkv = ReplicatedLinear(hidden_size,
                                              hidden_size * 3,
                                              bias=qkv_bias,
                                              params_dtype=dtype,
                                              prefix=f"{prefix}.self_attn_qkv")

        self.self_attn_proj = ReplicatedLinear(hidden_size,
                                               hidden_size,
                                               bias=qkv_bias,
                                               params_dtype=dtype,
                                               prefix=f"{prefix}.self_attn_proj")

        # MLP
        self.norm2 = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=True, dtype=dtype)
        self.mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, act_type="silu", dtype=dtype, prefix=f"{prefix}.mlp")

        # Modulation
        self.adaLN_modulation = ModulateProjection(hidden_size,
                                                   factor=2,
                                                   act_layer="silu",
                                                   dtype=dtype,
                                                   prefix=f"{prefix}.adaLN_modulation")

        # Scaled dot product attention
        self.attn = LocalAttention(
            num_heads=num_attention_heads,
            head_size=hidden_size // num_attention_heads,
            # TODO: remove hardcode
            supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA),
        )

    def forward(self, x, c, mask=None):

        # Get modulation parameters
        gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=-1)
        # Self-attention
        norm_x = self.norm1(x)
        qkv, _ = self.self_attn_qkv(norm_x)

        batch_size, seq_len = qkv.shape[0], qkv.shape[1]
        qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1)
        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
        if mask is None:
            attn_output = self.attn(q, k, v)
        else:
            from fastvideo.attention.backends.sdpa import SDPAMetadataBuilder
            attn_metadata = SDPAMetadataBuilder().build(
                current_timestep=0,
                attn_mask=mask,
            )
            # Run distributed attention
            with set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
                attn_output = self.attn(q, k, v)  # [B, L, H, D]
        attn_output = attn_output.reshape(batch_size, seq_len, -1)

        # Project and apply residual connection with gating
        attn_out, _ = self.self_attn_proj(attn_output)
        x = x + attn_out * gate_msa.unsqueeze(1)

        # MLP
        mlp_out = self.mlp(self.norm2(x))
        x = x + mlp_out * gate_mlp.unsqueeze(1)

        return x

`fastvideo.models.dits.hunyuanvideo15.IndividualTokenRefinerBlock.forward`

forward(x, c, mask = None)

Source

def forward(self, x, c, mask=None):

    # Get modulation parameters
    gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=-1)
    # Self-attention
    norm_x = self.norm1(x)
    qkv, _ = self.self_attn_qkv(norm_x)

    batch_size, seq_len = qkv.shape[0], qkv.shape[1]
    qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1)
    q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
    if mask is None:
        attn_output = self.attn(q, k, v)
    else:
        from fastvideo.attention.backends.sdpa import SDPAMetadataBuilder
        attn_metadata = SDPAMetadataBuilder().build(
            current_timestep=0,
            attn_mask=mask,
        )
        # Run distributed attention
        with set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
            attn_output = self.attn(q, k, v)  # [B, L, H, D]
    attn_output = attn_output.reshape(batch_size, seq_len, -1)

    # Project and apply residual connection with gating
    attn_out, _ = self.self_attn_proj(attn_output)
    x = x + attn_out * gate_msa.unsqueeze(1)

    # MLP
    mlp_out = self.mlp(self.norm2(x))
    x = x + mlp_out * gate_mlp.unsqueeze(1)

    return x

`fastvideo.models.dits.hunyuanvideo15.FinalLayer`

FinalLayer(hidden_size, patch_size, out_channels, dtype = None, prefix: str = '')

The final layer of DiT that projects features to pixel space.

Source

class FinalLayer(nn.Module):
    """
    The final layer of DiT that projects features to pixel space.
    """

    def __init__(self, hidden_size, patch_size, out_channels, dtype=None, prefix: str = "") -> None:
        super().__init__()

        # Normalization
        self.norm_final = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False, dtype=dtype)

        output_dim = patch_size[0] * patch_size[1] * patch_size[2] * out_channels

        self.linear = ReplicatedLinear(hidden_size,
                                       output_dim,
                                       bias=True,
                                       params_dtype=dtype,
                                       prefix=f"{prefix}.linear")

        # Modulation
        self.adaLN_modulation = ModulateProjection(hidden_size,
                                                   factor=2,
                                                   act_layer="silu",
                                                   dtype=dtype,
                                                   prefix=f"{prefix}.adaLN_modulation")

    def forward(self, x, c):
        # What the heck HF? Why you change the scale and shift order here???
        scale, shift = self.adaLN_modulation(c).chunk(2, dim=-1)
        x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
        x, _ = self.linear(x)
        return x

`fastvideo.models.dits.hunyuanvideo15.FinalLayer.forward`

forward(x, c)

Source

def forward(self, x, c):
    # What the heck HF? Why you change the scale and shift order here???
    scale, shift = self.adaLN_modulation(c).chunk(2, dim=-1)
    x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
    x, _ = self.linear(x)
    return x