--- title: gen3c ---

GEN3C is a video-conditioned diffusion model that uses a 3D cache for camera control.

Key Features: - Extends Cosmos 2.5 architecture with video conditioning inputs - condition_video_input_mask: Binary mask indicating conditioning frames - condition_video_pose: VAE-encoded 3D cache buffers (rendered warped images/masks) - Augment sigma embedding for conditioning noise augmentation - 3D RoPE with learnable per-axis positional embeddings

Reference: https://arxiv.org/abs/2503.03751

fastvideo.models.dits.gen3c.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.gen3c.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.gen3c.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.gen3c.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.gen3c.Gen3CVideoConfig

Gen3CVideoConfig(arch_config: DiTArchConfig = Gen3CArchConfig(), prefix: str = 'Gen3C', quant_config: QuantizationConfig | None = None, *, _resolved_attention_backend: AttentionBackendEnum | None = None)

Configuration for GEN3C video generation model.

fastvideo.models.dits.gen3c.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.gen3c.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.gen3c.get_sp_world_size

get_sp_world_size() -> int

Return world size for the sequence model parallel group.

fastvideo.models.dits.gen3c.get_forward_context

get_forward_context() -> ForwardContext

Get the current forward context.

fastvideo.models.dits.gen3c.RMSNorm

RMSNorm(hidden_size: int, eps: float = 1e-06, dtype: torch.dtype = torch.float32, var_hidden_size: int | None = None, has_weight: bool = True)

Root mean square normalization.

Computes x -> w * x / sqrt(E[x^2] + eps) where w is the learned weight. Refer to https://arxiv.org/abs/1910.07467

fastvideo.models.dits.gen3c.RMSNorm.forward_native

forward_native(x: torch.Tensor, residual: torch.Tensor | None = None) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]

PyTorch-native implementation equivalent to forward().

fastvideo.models.dits.gen3c.RMSNorm.extra_repr

extra_repr() -> str

fastvideo.models.dits.gen3c.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.gen3c.ReplicatedLinear.weight_loader

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

fastvideo.models.dits.gen3c.ReplicatedLinear.forward

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

fastvideo.models.dits.gen3c.ReplicatedLinear.extra_repr

extra_repr() -> str

fastvideo.models.dits.gen3c.ReplicatedLinear.get_shape_mapping

get_shape_mapping() -> dict

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

fastvideo.models.dits.gen3c.ReplicatedLinear.reset_shape_tracking

reset_shape_tracking() -> None

Clear tracked shapes and layer type mappings.

fastvideo.models.dits.gen3c.ReplicatedLinear.print_shape_summary

print_shape_summary() -> None

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

fastvideo.models.dits.gen3c.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.gen3c.MLP.forward

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

fastvideo.models.dits.gen3c.apply_rotary_emb

apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor | tuple[torch.Tensor, torch.Tensor], use_real: bool = True, use_real_unbind_dim: int = -1, sequence_dim: int = 2) -> torch.Tensor

Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are returned as real tensors. Args: x (torch.Tensor): Query or key tensor to apply rotary embeddings. [B, H, S, D] if sequence_dim=2 else [B, S, H, D]. freqs_cis (Tuple[torch.Tensor]): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],) sequence_dim: 1 = sequence at dim 1 (cos [1,S,1,D], x [B,S,H,D]); 2 = sequence at dim 2 (cos [1,1,S,D], x [B,H,S,D]). Returns: Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.

fastvideo.models.dits.gen3c.Timesteps

Timesteps(num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1)

fastvideo.models.dits.gen3c.Timesteps.forward

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

fastvideo.models.dits.gen3c.BaseDiT

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

fastvideo.models.dits.gen3c.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.gen3c.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.gen3c.BaseDiT.inference_config_overrides

inference_config_overrides() -> dict[str, Any]

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

fastvideo.models.dits.gen3c.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.gen3c.AttentionBackendEnum

AttentionBackendEnum()

fastvideo.models.dits.gen3c.Gen3CPatchEmbed

Gen3CPatchEmbed(in_channels: int, out_channels: int, patch_size: tuple[int, int, int])

GEN3C patch embedding - converts video (B, C, T, H, W) to patches (B, T', H', W', D). Uses linear projection after rearranging patches.

Input channels include: - VAE latent (16 channels) - condition_video_input_mask (1 channel) - condition_video_pose (frame_buffer_max * 32 channels) - padding_mask (1 channel, if concat_padding_mask=True)

Source
class Gen3CPatchEmbed(nn.Module):
    """
    GEN3C patch embedding - converts video (B, C, T, H, W) to patches (B, T', H', W', D).
    Uses linear projection after rearranging patches.

    Input channels include:
    - VAE latent (16 channels)
    - condition_video_input_mask (1 channel)
    - condition_video_pose (frame_buffer_max * 32 channels)
    - padding_mask (1 channel, if concat_padding_mask=True)
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        patch_size: tuple[int, int, int],
    ) -> None:
        super().__init__()
        self.patch_size = patch_size
        self.dim = in_channels * patch_size[0] * patch_size[1] * patch_size[2]

        self.proj = ReplicatedLinear(self.dim, out_channels, bias=False)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        Args:
            hidden_states: (B, C, T, H, W)
        Returns:
            (B, T', H', W', D) where T'=T//pt, H'=H//ph, W'=W//pw
        """
        batch_size, num_channels, num_frames, height, width = hidden_states.shape
        p_t, p_h, p_w = self.patch_size

        # Rearrange: b c (t pt) (h ph) (w pw) -> b t h w (c pt ph pw)
        hidden_states = hidden_states.reshape(batch_size, num_channels, num_frames // p_t, p_t, height // p_h, p_h,
                                              width // p_w, p_w)
        hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5, 7)
        hidden_states = hidden_states.flatten(4, 7)  # Flatten patch dimensions

        # Project to model dimension
        hidden_states, _ = self.proj(hidden_states)
        return hidden_states

fastvideo.models.dits.gen3c.Gen3CPatchEmbed.forward

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

Parameters

hidden_states: torch.Tensor

(B, C, T, H, W)

Returns: (B, T', H', W', D) where T'=T//pt, H'=H//ph, W'=W//pw

Source
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """
    Args:
        hidden_states: (B, C, T, H, W)
    Returns:
        (B, T', H', W', D) where T'=T//pt, H'=H//ph, W'=W//pw
    """
    batch_size, num_channels, num_frames, height, width = hidden_states.shape
    p_t, p_h, p_w = self.patch_size

    # Rearrange: b c (t pt) (h ph) (w pw) -> b t h w (c pt ph pw)
    hidden_states = hidden_states.reshape(batch_size, num_channels, num_frames // p_t, p_t, height // p_h, p_h,
                                          width // p_w, p_w)
    hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5, 7)
    hidden_states = hidden_states.flatten(4, 7)  # Flatten patch dimensions

    # Project to model dimension
    hidden_states, _ = self.proj(hidden_states)
    return hidden_states

fastvideo.models.dits.gen3c.Gen3CTimestepEmbedding

Gen3CTimestepEmbedding(in_features: int, out_features: int, use_adaln_lora: bool = True, adaln_lora_dim: int = 256)

GEN3C timestep embedding with AdaLN-LoRA support. Generates both standard embedding and AdaLN-LoRA parameters.

Source
class Gen3CTimestepEmbedding(nn.Module):
    """
    GEN3C timestep embedding with AdaLN-LoRA support.
    Generates both standard embedding and AdaLN-LoRA parameters.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        use_adaln_lora: bool = True,
        adaln_lora_dim: int = 256,
    ) -> None:
        super().__init__()
        self.use_adaln_lora = use_adaln_lora

        self.linear_1 = ReplicatedLinear(in_features, out_features, bias=False)
        self.activation = nn.SiLU()

        if use_adaln_lora:
            self.linear_2 = ReplicatedLinear(out_features, 3 * out_features, bias=False)
        else:
            self.linear_2 = ReplicatedLinear(out_features, out_features, bias=False)

    def forward(self, sample: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]:
        """
        Returns:
            emb: Standard embedding (B, D) - the original sinusoidal input
            adaln_lora: AdaLN-LoRA parameters (B, 3D) or None

        Note: When use_adaln_lora=True, the standard embedding is the INPUT
        (sinusoidal timestep embedding), not the processed output. The processed
        output (linear_2) is used exclusively for AdaLN-LoRA parameters.
        This matches the official GEN3C implementation.
        """
        emb, _ = self.linear_1(sample)
        emb = self.activation(emb)
        emb, _ = self.linear_2(emb)

        if self.use_adaln_lora:
            adaln_lora = emb  # (B, 3D) - full processed embedding for LoRA
            emb_standard = sample  # (B, D) - original input as the standard embedding
        else:
            emb_standard = emb
            adaln_lora = None

        return emb_standard, adaln_lora

fastvideo.models.dits.gen3c.Gen3CTimestepEmbedding.forward

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

Returns

emb: torch.Tensor

Standard embedding (B, D) - the original sinusoidal input

adaln_lora: torch.Tensor | None

AdaLN-LoRA parameters (B, 3D) or None

Note: When use_adaln_lora=True, the standard embedding is the INPUT (sinusoidal timestep embedding), not the processed output. The processed output (linear_2) is used exclusively for AdaLN-LoRA parameters. This matches the official GEN3C implementation.

Source
def forward(self, sample: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]:
    """
    Returns:
        emb: Standard embedding (B, D) - the original sinusoidal input
        adaln_lora: AdaLN-LoRA parameters (B, 3D) or None

    Note: When use_adaln_lora=True, the standard embedding is the INPUT
    (sinusoidal timestep embedding), not the processed output. The processed
    output (linear_2) is used exclusively for AdaLN-LoRA parameters.
    This matches the official GEN3C implementation.
    """
    emb, _ = self.linear_1(sample)
    emb = self.activation(emb)
    emb, _ = self.linear_2(emb)

    if self.use_adaln_lora:
        adaln_lora = emb  # (B, 3D) - full processed embedding for LoRA
        emb_standard = sample  # (B, D) - original input as the standard embedding
    else:
        emb_standard = emb
        adaln_lora = None

    return emb_standard, adaln_lora

fastvideo.models.dits.gen3c.Gen3CEmbedding

Gen3CEmbedding(embedding_dim: int, condition_dim: int, use_adaln_lora: bool = True, adaln_lora_dim: int = 256)

GEN3C timestep conditioning embedding. Generates sinusoidal embeddings and processes them through MLP.

Source
class Gen3CEmbedding(nn.Module):
    """
    GEN3C timestep conditioning embedding.
    Generates sinusoidal embeddings and processes them through MLP.
    """

    def __init__(
        self,
        embedding_dim: int,
        condition_dim: int,
        use_adaln_lora: bool = True,
        adaln_lora_dim: int = 256,
    ) -> None:
        super().__init__()

        self.time_proj = Timesteps(embedding_dim, flip_sin_to_cos=True, downscale_freq_shift=0.0)
        self.t_embedder = Gen3CTimestepEmbedding(
            embedding_dim,
            condition_dim,
            use_adaln_lora=use_adaln_lora,
            adaln_lora_dim=adaln_lora_dim,
        )

    def forward(
        self,
        timestep: torch.Tensor,
        dtype: torch.dtype,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """
        Args:
            timestep: (B,) tensor of timesteps
            dtype: Target dtype for embeddings

        Returns:
            embedded_timestep: Timestep embedding (B, D)
            adaln_lora: AdaLN-LoRA parameters (B, 3D) or None
        """
        # timestep should be 1D (B,)
        timestep = timestep.flatten()
        timesteps_proj = self.time_proj(timestep).to(dtype)  # (B, D)
        embedded_timestep, adaln_lora = self.t_embedder(timesteps_proj)

        return embedded_timestep, adaln_lora

fastvideo.models.dits.gen3c.Gen3CEmbedding.forward

forward(timestep: torch.Tensor, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor | None]

Parameters

timestep: torch.Tensor

(B,) tensor of timesteps

dtype: torch.dtype

Target dtype for embeddings

Returns

embedded_timestep: torch.Tensor

Timestep embedding (B, D)

adaln_lora: torch.Tensor | None

AdaLN-LoRA parameters (B, 3D) or None

Source
def forward(
    self,
    timestep: torch.Tensor,
    dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """
    Args:
        timestep: (B,) tensor of timesteps
        dtype: Target dtype for embeddings

    Returns:
        embedded_timestep: Timestep embedding (B, D)
        adaln_lora: AdaLN-LoRA parameters (B, 3D) or None
    """
    # timestep should be 1D (B,)
    timestep = timestep.flatten()
    timesteps_proj = self.time_proj(timestep).to(dtype)  # (B, D)
    embedded_timestep, adaln_lora = self.t_embedder(timesteps_proj)

    return embedded_timestep, adaln_lora

fastvideo.models.dits.gen3c.Gen3CAdaLayerNormZero

Gen3CAdaLayerNormZero(in_features: int)

GEN3C Adaptive Layer Normalization with zero initialization and gate. This is a simplified version that expects pre-computed shift/scale/gate parameters.

Source
class Gen3CAdaLayerNormZero(nn.Module):
    """
    GEN3C Adaptive Layer Normalization with zero initialization and gate.
    This is a simplified version that expects pre-computed shift/scale/gate parameters.
    """

    def __init__(
        self,
        in_features: int,
    ) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(in_features, elementwise_affine=False, eps=1e-6)

    def forward(
        self,
        hidden_states: torch.Tensor,
        shift: torch.Tensor,
        scale: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            hidden_states: Input tensor
            shift: Shift parameter for modulation
            scale: Scale parameter for modulation

        Returns:
            normalized_hidden_states: Modulated normalized hidden states
        """
        # Apply layer norm and modulation
        hidden_states = self.norm(hidden_states)
        hidden_states = hidden_states * (1 + scale) + shift
        return hidden_states

fastvideo.models.dits.gen3c.Gen3CAdaLayerNormZero.forward

forward(hidden_states: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor

Parameters

hidden_states: torch.Tensor

Input tensor

shift: torch.Tensor

Shift parameter for modulation

scale: torch.Tensor

Scale parameter for modulation

Returns

normalized_hidden_states: torch.Tensor

Modulated normalized hidden states

Source
def forward(
    self,
    hidden_states: torch.Tensor,
    shift: torch.Tensor,
    scale: torch.Tensor,
) -> torch.Tensor:
    """
    Args:
        hidden_states: Input tensor
        shift: Shift parameter for modulation
        scale: Scale parameter for modulation

    Returns:
        normalized_hidden_states: Modulated normalized hidden states
    """
    # Apply layer norm and modulation
    hidden_states = self.norm(hidden_states)
    hidden_states = hidden_states * (1 + scale) + shift
    return hidden_states

fastvideo.models.dits.gen3c.Gen3CSelfAttention

Gen3CSelfAttention(dim: int, num_heads: int, qk_norm: bool = True, eps: float = 1e-06, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None)

GEN3C self-attention with QK normalization and RoPE.

Source
class Gen3CSelfAttention(nn.Module):
    """
    GEN3C self-attention with QK normalization and RoPE.
    """

    def __init__(
        self,
        dim: int,
        num_heads: int,
        qk_norm: bool = True,
        eps: float = 1e-6,
        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
    ) -> None:
        assert dim % num_heads == 0
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads

        self.to_q = ReplicatedLinear(dim, dim, bias=False)
        self.to_k = ReplicatedLinear(dim, dim, bias=False)
        self.to_v = ReplicatedLinear(dim, dim, bias=False)
        self.to_out = ReplicatedLinear(dim, dim, bias=False)

        self.norm_q = RMSNorm(self.head_dim, eps=eps) if qk_norm else nn.Identity()
        self.norm_k = RMSNorm(self.head_dim, eps=eps) if qk_norm else nn.Identity()

        if supported_attention_backends is None:
            supported_attention_backends = (AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA)

        self.attn = DistributedAttention(num_heads=num_heads,
                                         head_size=self.head_dim,
                                         causal=False,
                                         supported_attention_backends=supported_attention_backends,
                                         prefix="self_attn")

    def forward(
        self,
        hidden_states: torch.Tensor,
        rope_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
        original_seq_len: int | None = None,
    ) -> torch.Tensor:
        """
        Args:
            hidden_states: (B, S, D) where S = T*H*W
            rope_emb: Tuple of (cos, sin) for RoPE
            original_seq_len: Original (unpadded) full sequence length
        """
        # Get QKV
        query, _ = self.to_q(hidden_states)
        key, _ = self.to_k(hidden_states)
        value, _ = self.to_v(hidden_states)

        # Reshape for multi-head attention: (B, S, D) -> (B, S, H, D_h) -> (B, H, S, D_h)
        query = query.unflatten(-1, (self.num_heads, self.head_dim)).transpose(1, 2)
        key = key.unflatten(-1, (self.num_heads, self.head_dim)).transpose(1, 2)
        value = value.unflatten(-1, (self.num_heads, self.head_dim)).transpose(1, 2)

        # Apply QK normalization
        query = self.norm_q(query)
        key = self.norm_k(key)

        # Apply RoPE if provided (query/key are now in (B, H, S, D_h) format)
        if rope_emb is not None:
            cos, sin = rope_emb
            query = apply_rotary_emb(query, (cos, sin), use_real=True, use_real_unbind_dim=-2)
            key = apply_rotary_emb(key, (cos, sin), use_real=True, use_real_unbind_dim=-2)

        # Attention computation
        query = query.transpose(1, 2)  # (B, H, S, D_h) -> (B, S, H, D_h)
        key = key.transpose(1, 2)
        value = value.transpose(1, 2)

        attn_output, _ = self.attn(query, key, value, original_seq_len=original_seq_len)
        attn_output = attn_output.flatten(-2, -1)  # (B, S, H*D_h)

        # Output projection
        attn_output, _ = self.to_out(attn_output)
        return attn_output

fastvideo.models.dits.gen3c.Gen3CSelfAttention.forward

forward(hidden_states: torch.Tensor, rope_emb: tuple[torch.Tensor, torch.Tensor] | None = None, original_seq_len: int | None = None) -> torch.Tensor

Parameters

hidden_states: torch.Tensor

(B, S, D) where S = THW

rope_emb: tuple[torch.Tensor, torch.Tensor] | None

Tuple of (cos, sin) for RoPE

original_seq_len: int | None

Original (unpadded) full sequence length

Source
def forward(
    self,
    hidden_states: torch.Tensor,
    rope_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
    original_seq_len: int | None = None,
) -> torch.Tensor:
    """
    Args:
        hidden_states: (B, S, D) where S = T*H*W
        rope_emb: Tuple of (cos, sin) for RoPE
        original_seq_len: Original (unpadded) full sequence length
    """
    # Get QKV
    query, _ = self.to_q(hidden_states)
    key, _ = self.to_k(hidden_states)
    value, _ = self.to_v(hidden_states)

    # Reshape for multi-head attention: (B, S, D) -> (B, S, H, D_h) -> (B, H, S, D_h)
    query = query.unflatten(-1, (self.num_heads, self.head_dim)).transpose(1, 2)
    key = key.unflatten(-1, (self.num_heads, self.head_dim)).transpose(1, 2)
    value = value.unflatten(-1, (self.num_heads, self.head_dim)).transpose(1, 2)

    # Apply QK normalization
    query = self.norm_q(query)
    key = self.norm_k(key)

    # Apply RoPE if provided (query/key are now in (B, H, S, D_h) format)
    if rope_emb is not None:
        cos, sin = rope_emb
        query = apply_rotary_emb(query, (cos, sin), use_real=True, use_real_unbind_dim=-2)
        key = apply_rotary_emb(key, (cos, sin), use_real=True, use_real_unbind_dim=-2)

    # Attention computation
    query = query.transpose(1, 2)  # (B, H, S, D_h) -> (B, S, H, D_h)
    key = key.transpose(1, 2)
    value = value.transpose(1, 2)

    attn_output, _ = self.attn(query, key, value, original_seq_len=original_seq_len)
    attn_output = attn_output.flatten(-2, -1)  # (B, S, H*D_h)

    # Output projection
    attn_output, _ = self.to_out(attn_output)
    return attn_output

fastvideo.models.dits.gen3c.Gen3CCrossAttention

Gen3CCrossAttention(dim: int, cross_attention_dim: int, num_heads: int, qk_norm: bool = True, eps: float = 1e-06, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None)

GEN3C cross-attention for text conditioning.

Source
class Gen3CCrossAttention(nn.Module):
    """
    GEN3C cross-attention for text conditioning.
    """

    def __init__(
        self,
        dim: int,
        cross_attention_dim: int,
        num_heads: int,
        qk_norm: bool = True,
        eps: float = 1e-6,
        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
    ) -> None:
        assert dim % num_heads == 0
        super().__init__()
        self.dim = dim
        self.cross_attention_dim = cross_attention_dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads

        self.to_q = ReplicatedLinear(dim, dim, bias=False)
        self.to_k = ReplicatedLinear(cross_attention_dim, dim, bias=False)
        self.to_v = ReplicatedLinear(cross_attention_dim, dim, bias=False)
        self.to_out = ReplicatedLinear(dim, dim, bias=False)

        self.norm_q = RMSNorm(self.head_dim, eps=eps) if qk_norm else nn.Identity()
        self.norm_k = RMSNorm(self.head_dim, eps=eps) if qk_norm else nn.Identity()

        if supported_attention_backends is None:
            supported_attention_backends = (AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA)

        self.attn = LocalAttention(
            num_heads=num_heads,
            head_size=self.head_dim,
            causal=False,
            supported_attention_backends=supported_attention_backends,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """
        Args:
            hidden_states: (B, S, D)
            encoder_hidden_states: (B, N, D_text)
        """
        # Get QKV
        query, _ = self.to_q(hidden_states)
        key, _ = self.to_k(encoder_hidden_states)
        value, _ = self.to_v(encoder_hidden_states)

        # Reshape for multi-head attention
        query = query.unflatten(-1, (self.num_heads, self.head_dim))
        key = key.unflatten(-1, (self.num_heads, self.head_dim))
        value = value.unflatten(-1, (self.num_heads, self.head_dim))

        # Apply QK normalization
        query = self.norm_q(query)
        key = self.norm_k(key)

        attn_output = self.attn(query, key, value)
        attn_output = attn_output.flatten(-2, -1)

        # Output projection
        attn_output, _ = self.to_out(attn_output)
        return attn_output

fastvideo.models.dits.gen3c.Gen3CCrossAttention.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor

Parameters

hidden_states: torch.Tensor

(B, S, D)

encoder_hidden_states: torch.Tensor

(B, N, D_text)

Source
def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Args:
        hidden_states: (B, S, D)
        encoder_hidden_states: (B, N, D_text)
    """
    # Get QKV
    query, _ = self.to_q(hidden_states)
    key, _ = self.to_k(encoder_hidden_states)
    value, _ = self.to_v(encoder_hidden_states)

    # Reshape for multi-head attention
    query = query.unflatten(-1, (self.num_heads, self.head_dim))
    key = key.unflatten(-1, (self.num_heads, self.head_dim))
    value = value.unflatten(-1, (self.num_heads, self.head_dim))

    # Apply QK normalization
    query = self.norm_q(query)
    key = self.norm_k(key)

    attn_output = self.attn(query, key, value)
    attn_output = attn_output.flatten(-2, -1)

    # Output projection
    attn_output, _ = self.to_out(attn_output)
    return attn_output

fastvideo.models.dits.gen3c.Gen3CTransformerBlock

Gen3CTransformerBlock(num_attention_heads: int, attention_head_dim: int, cross_attention_dim: int, mlp_ratio: float = 4.0, adaln_lora_dim: int = 256, use_adaln_lora: bool = True, qk_norm: bool = True, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None)

GEN3C transformer block with self-attention, cross-attention, and MLP. Uses AdaLN-LoRA for conditioning.

Source
class Gen3CTransformerBlock(nn.Module):
    """
    GEN3C transformer block with self-attention, cross-attention, and MLP.
    Uses AdaLN-LoRA for conditioning.
    """

    def __init__(
        self,
        num_attention_heads: int,
        attention_head_dim: int,
        cross_attention_dim: int,
        mlp_ratio: float = 4.0,
        adaln_lora_dim: int = 256,
        use_adaln_lora: bool = True,
        qk_norm: bool = True,
        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
    ) -> None:
        super().__init__()

        hidden_size = num_attention_heads * attention_head_dim
        self.use_adaln_lora = use_adaln_lora

        # Layer norms
        self.norm1 = Gen3CAdaLayerNormZero(hidden_size)
        self.norm2 = Gen3CAdaLayerNormZero(hidden_size)
        self.norm3 = Gen3CAdaLayerNormZero(hidden_size)

        # Attention and MLP layers
        self.attn1 = Gen3CSelfAttention(
            dim=hidden_size,
            num_heads=num_attention_heads,
            qk_norm=qk_norm,
            supported_attention_backends=supported_attention_backends,
        )
        self.attn2 = Gen3CCrossAttention(
            dim=hidden_size,
            cross_attention_dim=cross_attention_dim,
            num_heads=num_attention_heads,
            qk_norm=qk_norm,
            supported_attention_backends=supported_attention_backends,
        )
        self.mlp = MLP(hidden_size, int(hidden_size * mlp_ratio), act_type="gelu", bias=False)

        # AdaLN modulation layers
        if use_adaln_lora:
            self.adaln_modulation_self_attn = nn.ModuleList([
                nn.SiLU(),
                ReplicatedLinear(hidden_size, adaln_lora_dim, bias=False),
                ReplicatedLinear(adaln_lora_dim, 3 * hidden_size, bias=False),
            ])
            self.adaln_modulation_cross_attn = nn.ModuleList([
                nn.SiLU(),
                ReplicatedLinear(hidden_size, adaln_lora_dim, bias=False),
                ReplicatedLinear(adaln_lora_dim, 3 * hidden_size, bias=False),
            ])
            self.adaln_modulation_mlp = nn.ModuleList([
                nn.SiLU(),
                ReplicatedLinear(hidden_size, adaln_lora_dim, bias=False),
                ReplicatedLinear(adaln_lora_dim, 3 * hidden_size, bias=False),
            ])
        else:
            self.adaln_modulation_self_attn = nn.ModuleList(
                [nn.SiLU(), ReplicatedLinear(hidden_size, 3 * hidden_size, bias=False)])
            self.adaln_modulation_cross_attn = nn.ModuleList(
                [nn.SiLU(), ReplicatedLinear(hidden_size, 3 * hidden_size, bias=False)])
            self.adaln_modulation_mlp = nn.ModuleList(
                [nn.SiLU(), ReplicatedLinear(hidden_size, 3 * hidden_size, bias=False)])

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        affine_emb: torch.Tensor,
        adaln_lora: torch.Tensor | None = None,
        rope_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
        extra_pos_emb: torch.Tensor | None = None,
        original_seq_len: int | None = None,
    ) -> torch.Tensor:
        """
        Args:
            hidden_states: (B, S, D) where S = T*H*W (sequence may be sharded with SP)
            encoder_hidden_states: (B, N, D_text)
            affine_emb: (B, D) affine embedding
            adaln_lora: (B, 3D) AdaLN-LoRA parameters
            rope_emb: Tuple of (cos, sin) for RoPE
            extra_pos_emb: Optional learnable positional embeddings (B, S, D)
            original_seq_len: Original (unpadded) full sequence length
        """
        # Add extra positional embeddings if provided
        if extra_pos_emb is not None:
            hidden_states = hidden_states + extra_pos_emb

        B, S, D = hidden_states.shape

        # Compute modulation parameters
        if self.use_adaln_lora and adaln_lora is not None:
            shift_self_attn, scale_self_attn, gate_self_attn = (
                _run_linear_stack(self.adaln_modulation_self_attn, affine_emb) + adaln_lora).chunk(3, dim=-1)
            shift_cross_attn, scale_cross_attn, gate_cross_attn = (
                _run_linear_stack(self.adaln_modulation_cross_attn, affine_emb) + adaln_lora).chunk(3, dim=-1)
            shift_mlp, scale_mlp, gate_mlp = (_run_linear_stack(self.adaln_modulation_mlp, affine_emb) +
                                              adaln_lora).chunk(3, dim=-1)
        else:
            shift_self_attn, scale_self_attn, gate_self_attn = _run_linear_stack(self.adaln_modulation_self_attn,
                                                                                 affine_emb).chunk(3, dim=-1)
            shift_cross_attn, scale_cross_attn, gate_cross_attn = _run_linear_stack(self.adaln_modulation_cross_attn,
                                                                                    affine_emb).chunk(3, dim=-1)
            shift_mlp, scale_mlp, gate_mlp = _run_linear_stack(self.adaln_modulation_mlp, affine_emb).chunk(3, dim=-1)

        # Reshape modulation parameters for broadcasting: (B, D) -> (B, 1, D)
        shift_self_attn = shift_self_attn.unsqueeze(1).type_as(hidden_states)
        scale_self_attn = scale_self_attn.unsqueeze(1).type_as(hidden_states)
        gate_self_attn = gate_self_attn.unsqueeze(1).type_as(hidden_states)

        shift_cross_attn = shift_cross_attn.unsqueeze(1).type_as(hidden_states)
        scale_cross_attn = scale_cross_attn.unsqueeze(1).type_as(hidden_states)
        gate_cross_attn = gate_cross_attn.unsqueeze(1).type_as(hidden_states)

        shift_mlp = shift_mlp.unsqueeze(1).type_as(hidden_states)
        scale_mlp = scale_mlp.unsqueeze(1).type_as(hidden_states)
        gate_mlp = gate_mlp.unsqueeze(1).type_as(hidden_states)

        # Self-attention block
        norm_hidden_states = self.norm1(hidden_states, shift_self_attn, scale_self_attn)

        attn_output = self.attn1(
            norm_hidden_states,
            rope_emb=rope_emb,
            original_seq_len=original_seq_len,
        )
        hidden_states = hidden_states + gate_self_attn * attn_output

        # Cross-attention block
        norm_hidden_states = self.norm2(hidden_states, shift_cross_attn, scale_cross_attn)

        attn_output = self.attn2(
            norm_hidden_states,
            encoder_hidden_states=encoder_hidden_states,
        )
        hidden_states = hidden_states + gate_cross_attn * attn_output

        # MLP block
        norm_hidden_states = self.norm3(hidden_states, shift_mlp, scale_mlp)

        mlp_output = self.mlp(norm_hidden_states)
        hidden_states = hidden_states + gate_mlp * mlp_output

        return hidden_states

fastvideo.models.dits.gen3c.Gen3CTransformerBlock.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, affine_emb: torch.Tensor, adaln_lora: torch.Tensor | None = None, rope_emb: tuple[torch.Tensor, torch.Tensor] | None = None, extra_pos_emb: torch.Tensor | None = None, original_seq_len: int | None = None) -> torch.Tensor

Parameters

hidden_states: torch.Tensor

(B, S, D) where S = THW (sequence may be sharded with SP)

encoder_hidden_states: torch.Tensor

(B, N, D_text)

affine_emb: torch.Tensor

(B, D) affine embedding

adaln_lora: torch.Tensor | None

(B, 3D) AdaLN-LoRA parameters

rope_emb: tuple[torch.Tensor, torch.Tensor] | None

Tuple of (cos, sin) for RoPE

extra_pos_emb: torch.Tensor | None

Optional learnable positional embeddings (B, S, D)

original_seq_len: int | None

Original (unpadded) full sequence length

Source
def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    affine_emb: torch.Tensor,
    adaln_lora: torch.Tensor | None = None,
    rope_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
    extra_pos_emb: torch.Tensor | None = None,
    original_seq_len: int | None = None,
) -> torch.Tensor:
    """
    Args:
        hidden_states: (B, S, D) where S = T*H*W (sequence may be sharded with SP)
        encoder_hidden_states: (B, N, D_text)
        affine_emb: (B, D) affine embedding
        adaln_lora: (B, 3D) AdaLN-LoRA parameters
        rope_emb: Tuple of (cos, sin) for RoPE
        extra_pos_emb: Optional learnable positional embeddings (B, S, D)
        original_seq_len: Original (unpadded) full sequence length
    """
    # Add extra positional embeddings if provided
    if extra_pos_emb is not None:
        hidden_states = hidden_states + extra_pos_emb

    B, S, D = hidden_states.shape

    # Compute modulation parameters
    if self.use_adaln_lora and adaln_lora is not None:
        shift_self_attn, scale_self_attn, gate_self_attn = (
            _run_linear_stack(self.adaln_modulation_self_attn, affine_emb) + adaln_lora).chunk(3, dim=-1)
        shift_cross_attn, scale_cross_attn, gate_cross_attn = (
            _run_linear_stack(self.adaln_modulation_cross_attn, affine_emb) + adaln_lora).chunk(3, dim=-1)
        shift_mlp, scale_mlp, gate_mlp = (_run_linear_stack(self.adaln_modulation_mlp, affine_emb) +
                                          adaln_lora).chunk(3, dim=-1)
    else:
        shift_self_attn, scale_self_attn, gate_self_attn = _run_linear_stack(self.adaln_modulation_self_attn,
                                                                             affine_emb).chunk(3, dim=-1)
        shift_cross_attn, scale_cross_attn, gate_cross_attn = _run_linear_stack(self.adaln_modulation_cross_attn,
                                                                                affine_emb).chunk(3, dim=-1)
        shift_mlp, scale_mlp, gate_mlp = _run_linear_stack(self.adaln_modulation_mlp, affine_emb).chunk(3, dim=-1)

    # Reshape modulation parameters for broadcasting: (B, D) -> (B, 1, D)
    shift_self_attn = shift_self_attn.unsqueeze(1).type_as(hidden_states)
    scale_self_attn = scale_self_attn.unsqueeze(1).type_as(hidden_states)
    gate_self_attn = gate_self_attn.unsqueeze(1).type_as(hidden_states)

    shift_cross_attn = shift_cross_attn.unsqueeze(1).type_as(hidden_states)
    scale_cross_attn = scale_cross_attn.unsqueeze(1).type_as(hidden_states)
    gate_cross_attn = gate_cross_attn.unsqueeze(1).type_as(hidden_states)

    shift_mlp = shift_mlp.unsqueeze(1).type_as(hidden_states)
    scale_mlp = scale_mlp.unsqueeze(1).type_as(hidden_states)
    gate_mlp = gate_mlp.unsqueeze(1).type_as(hidden_states)

    # Self-attention block
    norm_hidden_states = self.norm1(hidden_states, shift_self_attn, scale_self_attn)

    attn_output = self.attn1(
        norm_hidden_states,
        rope_emb=rope_emb,
        original_seq_len=original_seq_len,
    )
    hidden_states = hidden_states + gate_self_attn * attn_output

    # Cross-attention block
    norm_hidden_states = self.norm2(hidden_states, shift_cross_attn, scale_cross_attn)

    attn_output = self.attn2(
        norm_hidden_states,
        encoder_hidden_states=encoder_hidden_states,
    )
    hidden_states = hidden_states + gate_cross_attn * attn_output

    # MLP block
    norm_hidden_states = self.norm3(hidden_states, shift_mlp, scale_mlp)

    mlp_output = self.mlp(norm_hidden_states)
    hidden_states = hidden_states + gate_mlp * mlp_output

    return hidden_states

fastvideo.models.dits.gen3c.Gen3CRotaryPosEmbed

Gen3CRotaryPosEmbed(hidden_size: int, max_size: tuple[int, int, int] = (128, 240, 240), patch_size: tuple[int, int, int] = (1, 2, 2), base_fps: int = 24, rope_scale: tuple[float, float, float] = (1.0, 1.0, 1.0), enable_fps_modulation: bool = True)

GEN3C 3D Rotary Position Embedding with NTK-aware extrapolation.

Source
class Gen3CRotaryPosEmbed(nn.Module):
    """
    GEN3C 3D Rotary Position Embedding with NTK-aware extrapolation.
    """

    def __init__(
            self,
            hidden_size: int,
            max_size: tuple[int, int, int] = (128, 240, 240),
            patch_size: tuple[int, int, int] = (1, 2, 2),
            base_fps: int = 24,
            rope_scale: tuple[float, float, float] = (1.0, 1.0, 1.0),
            enable_fps_modulation: bool = True,
    ) -> None:
        super().__init__()

        self.max_size = [size // patch for size, patch in zip(max_size, patch_size, strict=True)]
        self.patch_size = patch_size
        self.base_fps = base_fps
        self.enable_fps_modulation = enable_fps_modulation

        # Split dimensions: 1/3 for T, 1/3 for H, 1/3 for W
        self.dim_h = hidden_size // 6 * 2
        self.dim_w = hidden_size // 6 * 2
        self.dim_t = hidden_size - self.dim_h - self.dim_w

        # NTK-aware extrapolation factors
        self.h_ntk_factor = rope_scale[1]**(self.dim_h / (self.dim_h - 2))
        self.w_ntk_factor = rope_scale[2]**(self.dim_w / (self.dim_w - 2))
        self.t_ntk_factor = rope_scale[0]**(self.dim_t / (self.dim_t - 2))

    def forward(self, hidden_states: torch.Tensor, fps: int | None = None) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Generate 3D RoPE embeddings.
        """
        batch_size, T, H, W, input_dim = hidden_states.shape
        device = hidden_states.device

        # Generate frequency scales with NTK
        h_theta = 10000.0 * self.h_ntk_factor
        w_theta = 10000.0 * self.w_ntk_factor
        t_theta = 10000.0 * self.t_ntk_factor

        seq = torch.arange(max(self.max_size), device=device, dtype=torch.float32)

        dim_h_range = torch.arange(0, self.dim_h, 2, device=device,
                                   dtype=torch.float32)[:(self.dim_h // 2)] / self.dim_h
        dim_w_range = torch.arange(0, self.dim_w, 2, device=device,
                                   dtype=torch.float32)[:(self.dim_w // 2)] / self.dim_w
        dim_t_range = torch.arange(0, self.dim_t, 2, device=device,
                                   dtype=torch.float32)[:(self.dim_t // 2)] / self.dim_t

        h_spatial_freqs = 1.0 / (h_theta**dim_h_range)
        w_spatial_freqs = 1.0 / (w_theta**dim_w_range)
        temporal_freqs = 1.0 / (t_theta**dim_t_range)

        # Generate positional embeddings
        half_emb_h = torch.outer(seq[:H], h_spatial_freqs)
        half_emb_w = torch.outer(seq[:W], w_spatial_freqs)

        if self.enable_fps_modulation and fps is not None:
            half_emb_t = torch.outer(seq[:T] / fps * self.base_fps, temporal_freqs)
        else:
            half_emb_t = torch.outer(seq[:T], temporal_freqs)

        # Broadcast and concatenate embeddings
        emb_t = half_emb_t[:, None, None, :].repeat(1, H, W, 1)
        emb_h = half_emb_h[None, :, None, :].repeat(T, 1, W, 1)
        emb_w = half_emb_w[None, None, :, :].repeat(T, H, 1, 1)

        # Concatenate [t, h, w, t, h, w] for sin/cos pairs
        freqs = torch.cat([emb_t, emb_h, emb_w] * 2, dim=-1)
        freqs = freqs.flatten(0, 2).float()  # (THW, D)

        cos = torch.cos(freqs)
        sin = torch.sin(freqs)

        return cos, sin

fastvideo.models.dits.gen3c.Gen3CRotaryPosEmbed.forward

forward(hidden_states: torch.Tensor, fps: int | None = None) -> tuple[torch.Tensor, torch.Tensor]

Generate 3D RoPE embeddings.

Source
def forward(self, hidden_states: torch.Tensor, fps: int | None = None) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Generate 3D RoPE embeddings.
    """
    batch_size, T, H, W, input_dim = hidden_states.shape
    device = hidden_states.device

    # Generate frequency scales with NTK
    h_theta = 10000.0 * self.h_ntk_factor
    w_theta = 10000.0 * self.w_ntk_factor
    t_theta = 10000.0 * self.t_ntk_factor

    seq = torch.arange(max(self.max_size), device=device, dtype=torch.float32)

    dim_h_range = torch.arange(0, self.dim_h, 2, device=device,
                               dtype=torch.float32)[:(self.dim_h // 2)] / self.dim_h
    dim_w_range = torch.arange(0, self.dim_w, 2, device=device,
                               dtype=torch.float32)[:(self.dim_w // 2)] / self.dim_w
    dim_t_range = torch.arange(0, self.dim_t, 2, device=device,
                               dtype=torch.float32)[:(self.dim_t // 2)] / self.dim_t

    h_spatial_freqs = 1.0 / (h_theta**dim_h_range)
    w_spatial_freqs = 1.0 / (w_theta**dim_w_range)
    temporal_freqs = 1.0 / (t_theta**dim_t_range)

    # Generate positional embeddings
    half_emb_h = torch.outer(seq[:H], h_spatial_freqs)
    half_emb_w = torch.outer(seq[:W], w_spatial_freqs)

    if self.enable_fps_modulation and fps is not None:
        half_emb_t = torch.outer(seq[:T] / fps * self.base_fps, temporal_freqs)
    else:
        half_emb_t = torch.outer(seq[:T], temporal_freqs)

    # Broadcast and concatenate embeddings
    emb_t = half_emb_t[:, None, None, :].repeat(1, H, W, 1)
    emb_h = half_emb_h[None, :, None, :].repeat(T, 1, W, 1)
    emb_w = half_emb_w[None, None, :, :].repeat(T, H, 1, 1)

    # Concatenate [t, h, w, t, h, w] for sin/cos pairs
    freqs = torch.cat([emb_t, emb_h, emb_w] * 2, dim=-1)
    freqs = freqs.flatten(0, 2).float()  # (THW, D)

    cos = torch.cos(freqs)
    sin = torch.sin(freqs)

    return cos, sin

fastvideo.models.dits.gen3c.Gen3CLearnablePositionalEmbed

Gen3CLearnablePositionalEmbed(hidden_size: int, max_size: tuple[int, int, int], patch_size: tuple[int, int, int], eps: float = 1e-06)

GEN3C learnable absolute positional embeddings.

Source
class Gen3CLearnablePositionalEmbed(nn.Module):
    """
    GEN3C learnable absolute positional embeddings.
    """

    def __init__(
        self,
        hidden_size: int,
        max_size: tuple[int, int, int],
        patch_size: tuple[int, int, int],
        eps: float = 1e-6,
    ) -> None:
        super().__init__()

        self.max_size = [size // patch for size, patch in zip(max_size, patch_size, strict=True)]
        self.patch_size = patch_size
        self.eps = eps

        self.pos_emb_t = nn.Parameter(torch.zeros(self.max_size[0], hidden_size))
        self.pos_emb_h = nn.Parameter(torch.zeros(self.max_size[1], hidden_size))
        self.pos_emb_w = nn.Parameter(torch.zeros(self.max_size[2], hidden_size))

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        Args:
            hidden_states: (B, T, H, W, D)
        Returns:
            pos_emb: (B, T, H, W, D)
        """
        B, T, H, W, D = hidden_states.shape

        emb_t = self.pos_emb_t[:T][None, :, None, None, :].repeat(B, 1, H, W, 1)
        emb_h = self.pos_emb_h[:H][None, None, :, None, :].repeat(B, T, 1, W, 1)
        emb_w = self.pos_emb_w[:W][None, None, None, :, :].repeat(B, T, H, 1, 1)

        emb = emb_t + emb_h + emb_w

        # Normalize
        norm = torch.linalg.vector_norm(emb, dim=-1, keepdim=True, dtype=torch.float32)
        norm = torch.add(self.eps, norm, alpha=np.sqrt(norm.numel() / emb.numel()))
        return (emb / norm).type_as(hidden_states)

fastvideo.models.dits.gen3c.Gen3CLearnablePositionalEmbed.forward

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

Parameters

hidden_states: torch.Tensor

(B, T, H, W, D)

Returns: pos_emb: (B, T, H, W, D)

Source
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """
    Args:
        hidden_states: (B, T, H, W, D)
    Returns:
        pos_emb: (B, T, H, W, D)
    """
    B, T, H, W, D = hidden_states.shape

    emb_t = self.pos_emb_t[:T][None, :, None, None, :].repeat(B, 1, H, W, 1)
    emb_h = self.pos_emb_h[:H][None, None, :, None, :].repeat(B, T, 1, W, 1)
    emb_w = self.pos_emb_w[:W][None, None, None, :, :].repeat(B, T, H, 1, 1)

    emb = emb_t + emb_h + emb_w

    # Normalize
    norm = torch.linalg.vector_norm(emb, dim=-1, keepdim=True, dtype=torch.float32)
    norm = torch.add(self.eps, norm, alpha=np.sqrt(norm.numel() / emb.numel()))
    return (emb / norm).type_as(hidden_states)

fastvideo.models.dits.gen3c.Gen3CFinalLayer

Gen3CFinalLayer(hidden_size: int, out_channels: int, patch_size: tuple[int, int, int], adaln_lora_dim: int = 256, use_adaln_lora: bool = True)

GEN3C final layer with AdaLN modulation and unpatchification.

Source
class Gen3CFinalLayer(nn.Module):
    """
    GEN3C final layer with AdaLN modulation and unpatchification.
    """

    def __init__(
        self,
        hidden_size: int,
        out_channels: int,
        patch_size: tuple[int, int, int],
        adaln_lora_dim: int = 256,
        use_adaln_lora: bool = True,
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.use_adaln_lora = use_adaln_lora

        self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)

        # AdaLN modulation
        if use_adaln_lora:
            self.adaln_modulation = nn.ModuleList([
                nn.SiLU(),
                ReplicatedLinear(hidden_size, adaln_lora_dim, bias=False),
                ReplicatedLinear(adaln_lora_dim, 2 * hidden_size, bias=False),
            ])
        else:
            self.adaln_modulation = nn.ModuleList([
                nn.SiLU(),
                ReplicatedLinear(hidden_size, 2 * hidden_size, bias=False),
            ])

        # Output projection
        output_dim = out_channels * patch_size[0] * patch_size[1] * patch_size[2]
        self.proj_out = ReplicatedLinear(hidden_size, output_dim, bias=False)

    def forward(
        self,
        hidden_states: torch.Tensor,
        affine_emb: torch.Tensor,
        adaln_lora: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """
        Args:
            hidden_states: (B, T, H, W, D)
            affine_emb: (B, D)
            adaln_lora: (B, 3D) or None
        """
        # Generate modulation parameters
        modulation = _run_linear_stack(self.adaln_modulation, affine_emb)

        if self.use_adaln_lora and adaln_lora is not None:
            modulation = modulation + adaln_lora[..., :2 * self.hidden_size]

        shift, scale = modulation.chunk(2, dim=-1)

        # Apply normalization and modulation
        hidden_states = self.norm(hidden_states)

        # Broadcast over all non-channel dimensions. Supports both:
        # - (B, T, H, W, D)
        # - (B, S, D) when sequence-parallel path is active.
        mod_shape = [hidden_states.shape[0]] + [1] * (hidden_states.dim() - 2) + [self.hidden_size]
        shift = shift.view(*mod_shape)
        scale = scale.view(*mod_shape)

        hidden_states = hidden_states * (1 + scale) + shift

        # Project to output
        hidden_states, _ = self.proj_out(hidden_states)

        return hidden_states

fastvideo.models.dits.gen3c.Gen3CFinalLayer.forward

forward(hidden_states: torch.Tensor, affine_emb: torch.Tensor, adaln_lora: torch.Tensor | None = None) -> torch.Tensor

Parameters

hidden_states: torch.Tensor

(B, T, H, W, D)

affine_emb: torch.Tensor

(B, D)

adaln_lora: torch.Tensor | None

(B, 3D) or None

Source
def forward(
    self,
    hidden_states: torch.Tensor,
    affine_emb: torch.Tensor,
    adaln_lora: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Args:
        hidden_states: (B, T, H, W, D)
        affine_emb: (B, D)
        adaln_lora: (B, 3D) or None
    """
    # Generate modulation parameters
    modulation = _run_linear_stack(self.adaln_modulation, affine_emb)

    if self.use_adaln_lora and adaln_lora is not None:
        modulation = modulation + adaln_lora[..., :2 * self.hidden_size]

    shift, scale = modulation.chunk(2, dim=-1)

    # Apply normalization and modulation
    hidden_states = self.norm(hidden_states)

    # Broadcast over all non-channel dimensions. Supports both:
    # - (B, T, H, W, D)
    # - (B, S, D) when sequence-parallel path is active.
    mod_shape = [hidden_states.shape[0]] + [1] * (hidden_states.dim() - 2) + [self.hidden_size]
    shift = shift.view(*mod_shape)
    scale = scale.view(*mod_shape)

    hidden_states = hidden_states * (1 + scale) + shift

    # Project to output
    hidden_states, _ = self.proj_out(hidden_states)

    return hidden_states

fastvideo.models.dits.gen3c.Gen3CTransformer3DModel

Gen3CTransformer3DModel(config: Gen3CVideoConfig, hf_config: dict[str, Any])

GEN3C DiT - Video-conditioned diffusion transformer with 3D cache support.

Key features: - AdaLN-LoRA conditioning - 3D RoPE with NTK-aware extrapolation - Learnable positional embeddings - QK normalization - Video conditioning with 3D cache buffers - Augment sigma embedding for conditioning noise augmentation

Source
class Gen3CTransformer3DModel(BaseDiT):
    """
    GEN3C DiT - Video-conditioned diffusion transformer with 3D cache support.

    Key features:
    - AdaLN-LoRA conditioning
    - 3D RoPE with NTK-aware extrapolation
    - Learnable positional embeddings
    - QK normalization
    - Video conditioning with 3D cache buffers
    - Augment sigma embedding for conditioning noise augmentation
    """

    _fsdp_shard_conditions = Gen3CVideoConfig()._fsdp_shard_conditions
    _compile_conditions = Gen3CVideoConfig()._compile_conditions
    param_names_mapping = Gen3CVideoConfig().param_names_mapping
    lora_param_names_mapping = Gen3CVideoConfig().lora_param_names_mapping

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

        inner_dim = config.num_attention_heads * config.attention_head_dim
        self.hidden_size = inner_dim
        self.num_attention_heads = config.num_attention_heads
        self.in_channels = config.in_channels
        self.out_channels = config.out_channels
        self.num_channels_latents = config.num_channels_latents
        self.patch_size = config.patch_size
        self.max_size = config.max_size
        self.rope_scale = config.rope_scale
        self.concat_padding_mask = config.concat_padding_mask
        self.use_adaln_lora = getattr(config, "use_adaln_lora", True)
        self.adaln_lora_dim = getattr(config, "adaln_lora_dim", 256)
        self.extra_pos_embed_type = getattr(config, "extra_pos_embed_type", "learnable")
        self.affine_emb_norm = getattr(config, "affine_emb_norm", True)

        # GEN3C-specific
        self.frame_buffer_max = getattr(config, "frame_buffer_max", 2)
        self.buffer_channels = self.frame_buffer_max * 32  # Each buffer: 16 (frame) + 16 (mask)
        self.add_augment_sigma_embedding = getattr(config, "add_augment_sigma_embedding", True)

        # 1. Patch Embedding
        # Input channels: VAE (16) + condition_mask (1) + pose_buffers (64 for 2 buffers) + padding (1)
        patch_embed_in_channels = config.in_channels + 1 + self.buffer_channels
        if config.concat_padding_mask:
            patch_embed_in_channels += 1

        self.patch_embed = Gen3CPatchEmbed(patch_embed_in_channels, inner_dim, config.patch_size)

        # 2. Positional Embeddings
        self.rope = Gen3CRotaryPosEmbed(
            hidden_size=config.attention_head_dim,
            max_size=config.max_size,
            patch_size=config.patch_size,
            rope_scale=config.rope_scale,
            enable_fps_modulation=getattr(config, "rope_enable_fps_modulation", True),
        )

        self.learnable_pos_embed = None
        if self.extra_pos_embed_type == "learnable":
            self.learnable_pos_embed = Gen3CLearnablePositionalEmbed(
                hidden_size=inner_dim,
                max_size=config.max_size,
                patch_size=config.patch_size,
            )

        # 3. Time Embedding
        self.time_embed = Gen3CEmbedding(
            inner_dim,
            inner_dim,
            use_adaln_lora=self.use_adaln_lora,
            adaln_lora_dim=self.adaln_lora_dim,
        )

        # 4. Augment Sigma Embedding (GEN3C-specific)
        if self.add_augment_sigma_embedding:
            self.augment_sigma_embed = Gen3CEmbedding(
                inner_dim,
                inner_dim,
                use_adaln_lora=self.use_adaln_lora,
                adaln_lora_dim=self.adaln_lora_dim,
            )

        # 5. Affine Embedding Normalization
        if self.affine_emb_norm:
            self.affine_norm = RMSNorm(inner_dim, eps=config.eps)
        else:
            self.affine_norm = nn.Identity()

        # 6. Transformer Blocks
        self.transformer_blocks = nn.ModuleList([
            Gen3CTransformerBlock(
                num_attention_heads=config.num_attention_heads,
                attention_head_dim=config.attention_head_dim,
                cross_attention_dim=config.text_embed_dim,
                mlp_ratio=config.mlp_ratio,
                adaln_lora_dim=self.adaln_lora_dim,
                use_adaln_lora=self.use_adaln_lora,
                qk_norm=(config.qk_norm == "rms_norm"),
                supported_attention_backends=config._supported_attention_backends,
            ) for i in range(config.num_layers)
        ])

        # 7. Final Layer
        self.final_layer = Gen3CFinalLayer(
            hidden_size=inner_dim,
            out_channels=config.out_channels,
            patch_size=config.patch_size,
            adaln_lora_dim=self.adaln_lora_dim,
            use_adaln_lora=self.use_adaln_lora,
        )

        self.gradient_checkpointing = False
        self.__post_init__()

    def forward(
        self,
        hidden_states: torch.Tensor,
        timestep: torch.Tensor,
        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
        attention_mask: torch.Tensor | None = None,
        fps: int | None = None,
        condition_video_input_mask: torch.Tensor | None = None,
        condition_video_pose: torch.Tensor | None = None,
        condition_video_augment_sigma: torch.Tensor | None = None,
        padding_mask: torch.Tensor | None = None,
        **kwargs,
    ) -> torch.Tensor:
        """
        Args:
            hidden_states: (B, C, T, H, W) latent video (16 channels)
            timestep: (B,) diffusion timesteps
            encoder_hidden_states: (B, N, D_text) text embeddings
            attention_mask: Optional attention mask
            fps: Frames per second
            condition_video_input_mask: (B, 1, T, H, W) conditioning mask for video frames
            condition_video_pose: (B, buffer_channels, T, H, W) VAE-encoded 3D cache buffers
            condition_video_augment_sigma: (B,) sigma for conditioning noise augmentation
            padding_mask: (B, 1, H, W) padding mask
        """
        if not isinstance(encoder_hidden_states, torch.Tensor):
            encoder_hidden_states = encoder_hidden_states[0]

        batch_size, num_channels, num_frames, height, width = hidden_states.shape

        # 1. Concatenate condition_video_input_mask
        if condition_video_input_mask is not None:
            hidden_states = torch.cat([hidden_states, condition_video_input_mask], dim=1)
        else:
            # Default: zeros mask (no conditioning)
            condition_video_input_mask = torch.zeros(batch_size,
                                                     1,
                                                     num_frames,
                                                     height,
                                                     width,
                                                     device=hidden_states.device,
                                                     dtype=hidden_states.dtype)
            hidden_states = torch.cat([hidden_states, condition_video_input_mask], dim=1)

        # 2. Concatenate condition_video_pose (3D cache buffers)
        if condition_video_pose is not None:
            hidden_states = torch.cat([hidden_states, condition_video_pose], dim=1)
        else:
            # Default: zeros for pose buffers
            pose_zeros = torch.zeros(batch_size,
                                     self.buffer_channels,
                                     num_frames,
                                     height,
                                     width,
                                     device=hidden_states.device,
                                     dtype=hidden_states.dtype)
            hidden_states = torch.cat([hidden_states, pose_zeros], dim=1)

        # 3. Concatenate padding mask if needed
        if self.concat_padding_mask:
            if padding_mask is None:
                padding_mask = torch.ones(batch_size,
                                          1,
                                          height,
                                          width,
                                          device=hidden_states.device,
                                          dtype=hidden_states.dtype)
            padding_mask = transforms.functional.resize(
                padding_mask,
                list(hidden_states.shape[-2:]),
                interpolation=transforms.InterpolationMode.NEAREST,
            )
            hidden_states = torch.cat(
                [hidden_states, padding_mask.unsqueeze(2).repeat(1, 1, num_frames, 1, 1)],
                dim=1,
            )

        # 4. Patchify input
        p_t, p_h, p_w = self.patch_size
        hidden_states = self.patch_embed(hidden_states)  # (B, T', H', W', D)
        post_patch_num_frames, post_patch_height, post_patch_width = hidden_states.shape[1:4]

        # 5. Generate RoPE embeddings
        rope_emb = self.rope(hidden_states, fps=fps)

        # 6. Generate learnable positional embeddings
        extra_pos_emb = None
        if self.learnable_pos_embed is not None:
            extra_pos_emb = self.learnable_pos_embed(hidden_states)

        # Flatten to sequence representation for transformer blocks.
        hidden_states = hidden_states.flatten(1, 3)  # (B, S, D)
        if extra_pos_emb is not None:
            extra_pos_emb = extra_pos_emb.flatten(1, 3)

        # Sequence parallel sharding with optional padding.
        # This is a no-op when SP world size is 1.
        sp_world_size = get_sp_world_size()
        original_seq_len = hidden_states.shape[1]
        if sp_world_size > 1:
            hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)
            if extra_pos_emb is not None:
                extra_pos_emb, _ = sequence_model_parallel_shard(extra_pos_emb, dim=1)
            rope_cos, _ = sequence_model_parallel_shard(rope_emb[0], dim=0)
            rope_sin, _ = sequence_model_parallel_shard(rope_emb[1], dim=0)
            rope_emb = (rope_cos, rope_sin)

        # 7. Timestep embeddings
        affine_emb, adaln_lora = self.time_embed(timestep, hidden_states.dtype)

        # 8. Augment sigma embedding (GEN3C-specific)
        if self.add_augment_sigma_embedding:
            if condition_video_augment_sigma is None:
                condition_video_augment_sigma = torch.zeros_like(timestep)
            augment_emb, _ = self.augment_sigma_embed(condition_video_augment_sigma, hidden_states.dtype)
            affine_emb = affine_emb + augment_emb

        # 9. Apply affine normalization
        affine_emb = self.affine_norm(affine_emb)

        # 10. Transformer blocks
        if torch.is_grad_enabled() and self.gradient_checkpointing:
            for block in self.transformer_blocks:
                hidden_states = self._gradient_checkpointing_func(
                    block,
                    hidden_states,
                    encoder_hidden_states,
                    affine_emb,
                    adaln_lora,
                    rope_emb,
                    extra_pos_emb,
                    original_seq_len,
                )
        else:
            for block in self.transformer_blocks:
                hidden_states = block(
                    hidden_states=hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    affine_emb=affine_emb,
                    adaln_lora=adaln_lora,
                    rope_emb=rope_emb,
                    extra_pos_emb=extra_pos_emb,
                    original_seq_len=original_seq_len,
                )

        # 11. Final layer
        hidden_states = self.final_layer(hidden_states, affine_emb, adaln_lora)

        if sp_world_size > 1:
            hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1)

        # 12. Unpatchify: (B, S, P) -> (B, T', H', W', P) -> (B, C, T, H, W)
        hidden_states = hidden_states.unflatten(1, (post_patch_num_frames, post_patch_height, post_patch_width))
        hidden_states = hidden_states.unflatten(-1, (p_t, p_h, p_w, self.out_channels))
        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
        hidden_states = hidden_states.flatten(2, 3).flatten(3, 4).flatten(4, 5)

        return hidden_states

fastvideo.models.dits.gen3c.Gen3CTransformer3DModel.forward

forward(hidden_states: torch.Tensor, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], attention_mask: torch.Tensor | None = None, fps: int | None = None, condition_video_input_mask: torch.Tensor | None = None, condition_video_pose: torch.Tensor | None = None, condition_video_augment_sigma: torch.Tensor | None = None, padding_mask: torch.Tensor | None = None, **kwargs) -> torch.Tensor

Parameters

hidden_states: torch.Tensor

(B, C, T, H, W) latent video (16 channels)

timestep: torch.Tensor

(B,) diffusion timesteps

encoder_hidden_states: torch.Tensor | list[torch.Tensor]

(B, N, D_text) text embeddings

attention_mask: torch.Tensor | None

Optional attention mask

fps: int | None

Frames per second

condition_video_input_mask: torch.Tensor | None

(B, 1, T, H, W) conditioning mask for video frames

condition_video_pose: torch.Tensor | None

(B, buffer_channels, T, H, W) VAE-encoded 3D cache buffers

condition_video_augment_sigma: torch.Tensor | None

(B,) sigma for conditioning noise augmentation

padding_mask: torch.Tensor | None

(B, 1, H, W) padding mask

Source
def forward(
    self,
    hidden_states: torch.Tensor,
    timestep: torch.Tensor,
    encoder_hidden_states: torch.Tensor | list[torch.Tensor],
    attention_mask: torch.Tensor | None = None,
    fps: int | None = None,
    condition_video_input_mask: torch.Tensor | None = None,
    condition_video_pose: torch.Tensor | None = None,
    condition_video_augment_sigma: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    **kwargs,
) -> torch.Tensor:
    """
    Args:
        hidden_states: (B, C, T, H, W) latent video (16 channels)
        timestep: (B,) diffusion timesteps
        encoder_hidden_states: (B, N, D_text) text embeddings
        attention_mask: Optional attention mask
        fps: Frames per second
        condition_video_input_mask: (B, 1, T, H, W) conditioning mask for video frames
        condition_video_pose: (B, buffer_channels, T, H, W) VAE-encoded 3D cache buffers
        condition_video_augment_sigma: (B,) sigma for conditioning noise augmentation
        padding_mask: (B, 1, H, W) padding mask
    """
    if not isinstance(encoder_hidden_states, torch.Tensor):
        encoder_hidden_states = encoder_hidden_states[0]

    batch_size, num_channels, num_frames, height, width = hidden_states.shape

    # 1. Concatenate condition_video_input_mask
    if condition_video_input_mask is not None:
        hidden_states = torch.cat([hidden_states, condition_video_input_mask], dim=1)
    else:
        # Default: zeros mask (no conditioning)
        condition_video_input_mask = torch.zeros(batch_size,
                                                 1,
                                                 num_frames,
                                                 height,
                                                 width,
                                                 device=hidden_states.device,
                                                 dtype=hidden_states.dtype)
        hidden_states = torch.cat([hidden_states, condition_video_input_mask], dim=1)

    # 2. Concatenate condition_video_pose (3D cache buffers)
    if condition_video_pose is not None:
        hidden_states = torch.cat([hidden_states, condition_video_pose], dim=1)
    else:
        # Default: zeros for pose buffers
        pose_zeros = torch.zeros(batch_size,
                                 self.buffer_channels,
                                 num_frames,
                                 height,
                                 width,
                                 device=hidden_states.device,
                                 dtype=hidden_states.dtype)
        hidden_states = torch.cat([hidden_states, pose_zeros], dim=1)

    # 3. Concatenate padding mask if needed
    if self.concat_padding_mask:
        if padding_mask is None:
            padding_mask = torch.ones(batch_size,
                                      1,
                                      height,
                                      width,
                                      device=hidden_states.device,
                                      dtype=hidden_states.dtype)
        padding_mask = transforms.functional.resize(
            padding_mask,
            list(hidden_states.shape[-2:]),
            interpolation=transforms.InterpolationMode.NEAREST,
        )
        hidden_states = torch.cat(
            [hidden_states, padding_mask.unsqueeze(2).repeat(1, 1, num_frames, 1, 1)],
            dim=1,
        )

    # 4. Patchify input
    p_t, p_h, p_w = self.patch_size
    hidden_states = self.patch_embed(hidden_states)  # (B, T', H', W', D)
    post_patch_num_frames, post_patch_height, post_patch_width = hidden_states.shape[1:4]

    # 5. Generate RoPE embeddings
    rope_emb = self.rope(hidden_states, fps=fps)

    # 6. Generate learnable positional embeddings
    extra_pos_emb = None
    if self.learnable_pos_embed is not None:
        extra_pos_emb = self.learnable_pos_embed(hidden_states)

    # Flatten to sequence representation for transformer blocks.
    hidden_states = hidden_states.flatten(1, 3)  # (B, S, D)
    if extra_pos_emb is not None:
        extra_pos_emb = extra_pos_emb.flatten(1, 3)

    # Sequence parallel sharding with optional padding.
    # This is a no-op when SP world size is 1.
    sp_world_size = get_sp_world_size()
    original_seq_len = hidden_states.shape[1]
    if sp_world_size > 1:
        hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)
        if extra_pos_emb is not None:
            extra_pos_emb, _ = sequence_model_parallel_shard(extra_pos_emb, dim=1)
        rope_cos, _ = sequence_model_parallel_shard(rope_emb[0], dim=0)
        rope_sin, _ = sequence_model_parallel_shard(rope_emb[1], dim=0)
        rope_emb = (rope_cos, rope_sin)

    # 7. Timestep embeddings
    affine_emb, adaln_lora = self.time_embed(timestep, hidden_states.dtype)

    # 8. Augment sigma embedding (GEN3C-specific)
    if self.add_augment_sigma_embedding:
        if condition_video_augment_sigma is None:
            condition_video_augment_sigma = torch.zeros_like(timestep)
        augment_emb, _ = self.augment_sigma_embed(condition_video_augment_sigma, hidden_states.dtype)
        affine_emb = affine_emb + augment_emb

    # 9. Apply affine normalization
    affine_emb = self.affine_norm(affine_emb)

    # 10. Transformer blocks
    if torch.is_grad_enabled() and self.gradient_checkpointing:
        for block in self.transformer_blocks:
            hidden_states = self._gradient_checkpointing_func(
                block,
                hidden_states,
                encoder_hidden_states,
                affine_emb,
                adaln_lora,
                rope_emb,
                extra_pos_emb,
                original_seq_len,
            )
    else:
        for block in self.transformer_blocks:
            hidden_states = block(
                hidden_states=hidden_states,
                encoder_hidden_states=encoder_hidden_states,
                affine_emb=affine_emb,
                adaln_lora=adaln_lora,
                rope_emb=rope_emb,
                extra_pos_emb=extra_pos_emb,
                original_seq_len=original_seq_len,
            )

    # 11. Final layer
    hidden_states = self.final_layer(hidden_states, affine_emb, adaln_lora)

    if sp_world_size > 1:
        hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1)

    # 12. Unpatchify: (B, S, P) -> (B, T', H', W', P) -> (B, C, T, H, W)
    hidden_states = hidden_states.unflatten(1, (post_patch_num_frames, post_patch_height, post_patch_width))
    hidden_states = hidden_states.unflatten(-1, (p_t, p_h, p_w, self.out_channels))
    hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
    hidden_states = hidden_states.flatten(2, 3).flatten(3, 4).flatten(4, 5)

    return hidden_states