--- title: cosmos ---

fastvideo.models.dits.cosmos.CosmosVideoConfig

CosmosVideoConfig(arch_config: DiTArchConfig = CosmosArchConfig(), prefix: str = 'Cosmos', quant_config: QuantizationConfig | None = None, *, _resolved_attention_backend: AttentionBackendEnum | None = None)

fastvideo.models.dits.cosmos.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.cosmos.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.cosmos.RMSNorm.extra_repr

extra_repr() -> str

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

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

fastvideo.models.dits.cosmos.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.cosmos.Timesteps

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

fastvideo.models.dits.cosmos.Timesteps.forward

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

fastvideo.models.dits.cosmos.BaseDiT

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

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

inference_config_overrides() -> dict[str, Any]

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

fastvideo.models.dits.cosmos.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.cosmos.CosmosPatchEmbed

CosmosPatchEmbed(in_channels: int, out_channels: int, patch_size: tuple[int, int, int], bias: bool = True)
Source
class CosmosPatchEmbed(nn.Module):

    def __init__(self,
                 in_channels: int,
                 out_channels: int,
                 patch_size: tuple[int, int, int],
                 bias: bool = True) -> None:
        super().__init__()
        self.patch_size = patch_size

        self.proj = nn.Linear(in_channels * patch_size[0] * patch_size[1] * patch_size[2], out_channels, bias=bias)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        batch_size, num_channels, num_frames, height, width = hidden_states.shape
        p_t, p_h, p_w = self.patch_size
        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).flatten(4, 7)
        hidden_states = self.proj(hidden_states)
        return hidden_states

fastvideo.models.dits.cosmos.CosmosPatchEmbed.forward

forward(hidden_states: torch.Tensor) -> torch.Tensor
Source
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    batch_size, num_channels, num_frames, height, width = hidden_states.shape
    p_t, p_h, p_w = self.patch_size
    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).flatten(4, 7)
    hidden_states = self.proj(hidden_states)
    return hidden_states

fastvideo.models.dits.cosmos.CosmosTimestepEmbedding

CosmosTimestepEmbedding(in_features: int, out_features: int)
Source
class CosmosTimestepEmbedding(nn.Module):

    def __init__(self, in_features: int, out_features: int) -> None:
        super().__init__()
        self.linear_1 = nn.Linear(in_features, out_features, bias=False)
        self.activation = nn.SiLU()
        self.linear_2 = nn.Linear(out_features, 3 * out_features, bias=False)

    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
        emb = self.linear_1(timesteps)
        emb = self.activation(emb)
        emb = self.linear_2(emb)
        return emb

fastvideo.models.dits.cosmos.CosmosTimestepEmbedding.forward

forward(timesteps: torch.Tensor) -> torch.Tensor
Source
def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
    emb = self.linear_1(timesteps)
    emb = self.activation(emb)
    emb = self.linear_2(emb)
    return emb

fastvideo.models.dits.cosmos.CosmosEmbedding

CosmosEmbedding(embedding_dim: int, condition_dim: int)
Source
class CosmosEmbedding(nn.Module):

    def __init__(self, embedding_dim: int, condition_dim: int) -> None:
        super().__init__()

        self.time_proj = Timesteps(embedding_dim, flip_sin_to_cos=True, downscale_freq_shift=0.0)
        self.t_embedder = CosmosTimestepEmbedding(embedding_dim, condition_dim)
        self.norm = RMSNorm(embedding_dim, eps=1e-6)

    def forward(self, hidden_states: torch.Tensor, timestep: torch.LongTensor) -> torch.Tensor:
        timesteps_proj = self.time_proj(timestep).type_as(hidden_states)
        temb = self.t_embedder(timesteps_proj)
        embedded_timestep = self.norm(timesteps_proj)
        return temb, embedded_timestep

fastvideo.models.dits.cosmos.CosmosEmbedding.forward

forward(hidden_states: torch.Tensor, timestep: torch.LongTensor) -> torch.Tensor
Source
def forward(self, hidden_states: torch.Tensor, timestep: torch.LongTensor) -> torch.Tensor:
    timesteps_proj = self.time_proj(timestep).type_as(hidden_states)
    temb = self.t_embedder(timesteps_proj)
    embedded_timestep = self.norm(timesteps_proj)
    return temb, embedded_timestep

fastvideo.models.dits.cosmos.CosmosAdaLayerNorm

CosmosAdaLayerNorm(in_features: int, hidden_features: int)
Source
class CosmosAdaLayerNorm(nn.Module):

    def __init__(self, in_features: int, hidden_features: int) -> None:
        super().__init__()
        self.embedding_dim = in_features

        self.activation = nn.SiLU()
        self.norm = nn.LayerNorm(in_features, elementwise_affine=False, eps=1e-6)
        self.linear_1 = nn.Linear(in_features, hidden_features, bias=False)
        self.linear_2 = nn.Linear(hidden_features, 2 * in_features, bias=False)

    def forward(self,
                hidden_states: torch.Tensor,
                embedded_timestep: torch.Tensor,
                temb: torch.Tensor | None = None) -> torch.Tensor:
        embedded_timestep = self.activation(embedded_timestep)
        embedded_timestep = self.linear_1(embedded_timestep)
        embedded_timestep = self.linear_2(embedded_timestep)

        if temb is not None:
            embedded_timestep = embedded_timestep + temb[..., :2 * self.embedding_dim]

        shift, scale = embedded_timestep.chunk(2, dim=-1)
        with torch.autocast(device_type=hidden_states.device.type, enabled=False):
            hidden_states = self.norm(hidden_states)

        if embedded_timestep.ndim == 2:
            shift, scale = (x.unsqueeze(1) for x in (shift, scale))

        hidden_states = hidden_states * (1 + scale) + shift
        return hidden_states

fastvideo.models.dits.cosmos.CosmosAdaLayerNorm.forward

forward(hidden_states: torch.Tensor, embedded_timestep: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor
Source
def forward(self,
            hidden_states: torch.Tensor,
            embedded_timestep: torch.Tensor,
            temb: torch.Tensor | None = None) -> torch.Tensor:
    embedded_timestep = self.activation(embedded_timestep)
    embedded_timestep = self.linear_1(embedded_timestep)
    embedded_timestep = self.linear_2(embedded_timestep)

    if temb is not None:
        embedded_timestep = embedded_timestep + temb[..., :2 * self.embedding_dim]

    shift, scale = embedded_timestep.chunk(2, dim=-1)
    with torch.autocast(device_type=hidden_states.device.type, enabled=False):
        hidden_states = self.norm(hidden_states)

    if embedded_timestep.ndim == 2:
        shift, scale = (x.unsqueeze(1) for x in (shift, scale))

    hidden_states = hidden_states * (1 + scale) + shift
    return hidden_states

fastvideo.models.dits.cosmos.CosmosAdaLayerNormZero

CosmosAdaLayerNormZero(in_features: int, hidden_features: int | None = None)
Source
class CosmosAdaLayerNormZero(nn.Module):

    def __init__(self, in_features: int, hidden_features: int | None = None) -> None:
        super().__init__()

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

        if hidden_features is None:
            self.linear_1 = nn.Identity()
        else:
            self.linear_1 = nn.Linear(in_features, hidden_features, bias=False)

        self.linear_2 = nn.Linear(hidden_features, 3 * in_features, bias=False)

    def forward(
        self,
        hidden_states: torch.Tensor,
        embedded_timestep: torch.Tensor,
        temb: torch.Tensor | None = None,
    ) -> torch.Tensor:
        embedded_timestep = self.activation(embedded_timestep)
        embedded_timestep = self.linear_1(embedded_timestep)
        embedded_timestep = self.linear_2(embedded_timestep)

        if temb is not None:
            embedded_timestep = embedded_timestep + temb

        shift, scale, gate = embedded_timestep.chunk(3, dim=-1)

        with torch.autocast(device_type=hidden_states.device.type, enabled=False):
            hidden_states = self.norm(hidden_states)

        if embedded_timestep.ndim == 2:
            shift, scale, gate = (x.unsqueeze(1) for x in (shift, scale, gate))

        hidden_states = hidden_states * (1 + scale) + shift
        return hidden_states, gate

fastvideo.models.dits.cosmos.CosmosAdaLayerNormZero.forward

forward(hidden_states: torch.Tensor, embedded_timestep: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor
Source
def forward(
    self,
    hidden_states: torch.Tensor,
    embedded_timestep: torch.Tensor,
    temb: torch.Tensor | None = None,
) -> torch.Tensor:
    embedded_timestep = self.activation(embedded_timestep)
    embedded_timestep = self.linear_1(embedded_timestep)
    embedded_timestep = self.linear_2(embedded_timestep)

    if temb is not None:
        embedded_timestep = embedded_timestep + temb

    shift, scale, gate = embedded_timestep.chunk(3, dim=-1)

    with torch.autocast(device_type=hidden_states.device.type, enabled=False):
        hidden_states = self.norm(hidden_states)

    if embedded_timestep.ndim == 2:
        shift, scale, gate = (x.unsqueeze(1) for x in (shift, scale, gate))

    hidden_states = hidden_states * (1 + scale) + shift
    return hidden_states, gate

fastvideo.models.dits.cosmos.CosmosSelfAttention

CosmosSelfAttention(dim: int, num_heads: int, qk_norm = True, eps = 1e-06, prefix: str = '')
Source
class CosmosSelfAttention(nn.Module):

    def __init__(self, dim: int, num_heads: int, qk_norm=True, eps=1e-6, prefix: str = "") -> None:
        assert dim % num_heads == 0
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.qk_norm = qk_norm
        self.eps = eps

        # layers - use standard PyTorch layers when using torch backend
        self.to_q = nn.Linear(dim, dim, bias=False)
        self.to_k = nn.Linear(dim, dim, bias=False)
        self.to_v = nn.Linear(dim, dim, bias=False)
        self.to_out = nn.Linear(dim, dim, bias=False)
        self.dropout = nn.Dropout(0.0)

        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()

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

        if encoder_hidden_states is None:
            encoder_hidden_states = hidden_states

        # 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(2, (self.num_heads, -1)).transpose(1, 2)
        key = key.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
        value = value.unflatten(2, (self.num_heads, -1)).transpose(1, 2)

        # Apply normalization
        if self.norm_q is not None:
            query = self.norm_q(query)
        if self.norm_k is not None:
            key = self.norm_k(key)

        # Apply RoPE if provided
        if image_rotary_emb is not None:
            query = apply_rotary_emb(query, image_rotary_emb, use_real=True, use_real_unbind_dim=-2)
            key = apply_rotary_emb(key, image_rotary_emb, use_real=True, use_real_unbind_dim=-2)

        # Prepare for GQA (Grouped Query Attention)
        if torch.onnx.is_in_onnx_export():
            query_idx = torch.tensor(query.size(3), device=query.device)
            key_idx = torch.tensor(key.size(3), device=key.device)
            value_idx = torch.tensor(value.size(3), device=value.device)
        else:
            query_idx = query.size(3)
            key_idx = key.size(3)
            value_idx = value.size(3)
        key = key.repeat_interleave(query_idx // key_idx, dim=3)
        value = value.repeat_interleave(query_idx // value_idx, dim=3)

        # Attention computation
        attn_output = torch.nn.functional.scaled_dot_product_attention(query,
                                                                       key,
                                                                       value,
                                                                       attn_mask=attention_mask,
                                                                       dropout_p=0.0,
                                                                       is_causal=False)
        attn_output = attn_output.transpose(1, 2).flatten(2, 3).type_as(query)

        # Output projection
        attn_output = self.to_out(attn_output)
        attn_output = self.dropout(attn_output)

        return attn_output

fastvideo.models.dits.cosmos.CosmosSelfAttention.forward

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

    if encoder_hidden_states is None:
        encoder_hidden_states = hidden_states

    # 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(2, (self.num_heads, -1)).transpose(1, 2)
    key = key.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
    value = value.unflatten(2, (self.num_heads, -1)).transpose(1, 2)

    # Apply normalization
    if self.norm_q is not None:
        query = self.norm_q(query)
    if self.norm_k is not None:
        key = self.norm_k(key)

    # Apply RoPE if provided
    if image_rotary_emb is not None:
        query = apply_rotary_emb(query, image_rotary_emb, use_real=True, use_real_unbind_dim=-2)
        key = apply_rotary_emb(key, image_rotary_emb, use_real=True, use_real_unbind_dim=-2)

    # Prepare for GQA (Grouped Query Attention)
    if torch.onnx.is_in_onnx_export():
        query_idx = torch.tensor(query.size(3), device=query.device)
        key_idx = torch.tensor(key.size(3), device=key.device)
        value_idx = torch.tensor(value.size(3), device=value.device)
    else:
        query_idx = query.size(3)
        key_idx = key.size(3)
        value_idx = value.size(3)
    key = key.repeat_interleave(query_idx // key_idx, dim=3)
    value = value.repeat_interleave(query_idx // value_idx, dim=3)

    # Attention computation
    attn_output = torch.nn.functional.scaled_dot_product_attention(query,
                                                                   key,
                                                                   value,
                                                                   attn_mask=attention_mask,
                                                                   dropout_p=0.0,
                                                                   is_causal=False)
    attn_output = attn_output.transpose(1, 2).flatten(2, 3).type_as(query)

    # Output projection
    attn_output = self.to_out(attn_output)
    attn_output = self.dropout(attn_output)

    return attn_output

fastvideo.models.dits.cosmos.CosmosCrossAttention

CosmosCrossAttention(dim: int, cross_attention_dim: int, num_heads: int, qk_norm = True, eps = 1e-06, prefix: str = '')
Source
class CosmosCrossAttention(nn.Module):

    def __init__(self,
                 dim: int,
                 cross_attention_dim: int,
                 num_heads: int,
                 qk_norm=True,
                 eps=1e-6,
                 prefix: str = "") -> 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.qk_norm = qk_norm
        self.eps = eps

        self.to_q = nn.Linear(dim, dim, bias=False)
        self.to_k = nn.Linear(cross_attention_dim, dim, bias=False)
        self.to_v = nn.Linear(cross_attention_dim, dim, bias=False)
        self.to_out = nn.Linear(dim, dim, bias=False)
        self.dropout = nn.Dropout(0.0)

        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()

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

        # 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(2, (self.num_heads, -1)).transpose(1, 2)
        key = key.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
        value = value.unflatten(2, (self.num_heads, -1)).transpose(1, 2)

        # Apply normalization
        if self.norm_q is not None:
            query = self.norm_q(query)
        if self.norm_k is not None:
            key = self.norm_k(key)

        # Prepare for GQA (Grouped Query Attention)
        if torch.onnx.is_in_onnx_export():
            query_idx = torch.tensor(query.size(3), device=query.device)
            key_idx = torch.tensor(key.size(3), device=key.device)
            value_idx = torch.tensor(value.size(3), device=value.device)
        else:
            query_idx = query.size(3)
            key_idx = key.size(3)
            value_idx = value.size(3)
        key = key.repeat_interleave(query_idx // key_idx, dim=3)
        value = value.repeat_interleave(query_idx // value_idx, dim=3)

        # Attention computation
        attn_output = torch.nn.functional.scaled_dot_product_attention(query,
                                                                       key,
                                                                       value,
                                                                       attn_mask=attention_mask,
                                                                       dropout_p=0.0,
                                                                       is_causal=False)
        attn_output = attn_output.transpose(1, 2).flatten(2, 3).type_as(query)

        # Output projection
        attn_output = self.to_out(attn_output)
        attn_output = self.dropout(attn_output)

        return attn_output

fastvideo.models.dits.cosmos.CosmosCrossAttention.forward

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

    # 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(2, (self.num_heads, -1)).transpose(1, 2)
    key = key.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
    value = value.unflatten(2, (self.num_heads, -1)).transpose(1, 2)

    # Apply normalization
    if self.norm_q is not None:
        query = self.norm_q(query)
    if self.norm_k is not None:
        key = self.norm_k(key)

    # Prepare for GQA (Grouped Query Attention)
    if torch.onnx.is_in_onnx_export():
        query_idx = torch.tensor(query.size(3), device=query.device)
        key_idx = torch.tensor(key.size(3), device=key.device)
        value_idx = torch.tensor(value.size(3), device=value.device)
    else:
        query_idx = query.size(3)
        key_idx = key.size(3)
        value_idx = value.size(3)
    key = key.repeat_interleave(query_idx // key_idx, dim=3)
    value = value.repeat_interleave(query_idx // value_idx, dim=3)

    # Attention computation
    attn_output = torch.nn.functional.scaled_dot_product_attention(query,
                                                                   key,
                                                                   value,
                                                                   attn_mask=attention_mask,
                                                                   dropout_p=0.0,
                                                                   is_causal=False)
    attn_output = attn_output.transpose(1, 2).flatten(2, 3).type_as(query)

    # Output projection
    attn_output = self.to_out(attn_output)
    attn_output = self.dropout(attn_output)

    return attn_output

fastvideo.models.dits.cosmos.CosmosTransformerBlock

CosmosTransformerBlock(num_attention_heads: int, attention_head_dim: int, cross_attention_dim: int, mlp_ratio: float = 4.0, adaln_lora_dim: int = 256, qk_norm: str = 'rms_norm', out_bias: bool = False, prefix: str = '')
Source
class CosmosTransformerBlock(nn.Module):

    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,
        qk_norm: str = "rms_norm",
        out_bias: bool = False,
        prefix: str = "",
    ) -> None:
        super().__init__()

        hidden_size = num_attention_heads * attention_head_dim

        self.norm1 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim)
        self.attn1 = CosmosSelfAttention(dim=hidden_size,
                                         num_heads=num_attention_heads,
                                         qk_norm=(qk_norm == "rms_norm"),
                                         eps=1e-5,
                                         prefix=f"{prefix}.attn1")

        self.norm2 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim)
        self.attn2 = CosmosCrossAttention(dim=hidden_size,
                                          cross_attention_dim=cross_attention_dim,
                                          num_heads=num_attention_heads,
                                          qk_norm=(qk_norm == "rms_norm"),
                                          eps=1e-5,
                                          prefix=f"{prefix}.attn2")

        self.norm3 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim)
        self.ff = MLP(hidden_size, int(hidden_size * mlp_ratio), act_type="gelu", bias=False)

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        embedded_timestep: torch.Tensor,
        temb: torch.Tensor | None = None,
        image_rotary_emb: torch.Tensor | None = None,
        extra_pos_emb: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if extra_pos_emb is not None:
            hidden_states = hidden_states + extra_pos_emb

        norm_hidden_states, gate = self.norm1(hidden_states, embedded_timestep, temb)

        attn_output = self.attn1(norm_hidden_states, image_rotary_emb=image_rotary_emb)
        hidden_states = hidden_states + gate * attn_output

        norm_hidden_states, gate = self.norm2(hidden_states, embedded_timestep, temb)
        attn_output = self.attn2(norm_hidden_states,
                                 encoder_hidden_states=encoder_hidden_states,
                                 attention_mask=attention_mask)

        hidden_states = hidden_states + gate * attn_output

        norm_hidden_states, gate = self.norm3(hidden_states, embedded_timestep, temb)
        ff_output = self.ff(norm_hidden_states)
        hidden_states = hidden_states + gate * ff_output

        return hidden_states

fastvideo.models.dits.cosmos.CosmosTransformerBlock.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, embedded_timestep: torch.Tensor, temb: torch.Tensor | None = None, image_rotary_emb: torch.Tensor | None = None, extra_pos_emb: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None) -> torch.Tensor
Source
def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    embedded_timestep: torch.Tensor,
    temb: torch.Tensor | None = None,
    image_rotary_emb: torch.Tensor | None = None,
    extra_pos_emb: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    if extra_pos_emb is not None:
        hidden_states = hidden_states + extra_pos_emb

    norm_hidden_states, gate = self.norm1(hidden_states, embedded_timestep, temb)

    attn_output = self.attn1(norm_hidden_states, image_rotary_emb=image_rotary_emb)
    hidden_states = hidden_states + gate * attn_output

    norm_hidden_states, gate = self.norm2(hidden_states, embedded_timestep, temb)
    attn_output = self.attn2(norm_hidden_states,
                             encoder_hidden_states=encoder_hidden_states,
                             attention_mask=attention_mask)

    hidden_states = hidden_states + gate * attn_output

    norm_hidden_states, gate = self.norm3(hidden_states, embedded_timestep, temb)
    ff_output = self.ff(norm_hidden_states)
    hidden_states = hidden_states + gate * ff_output

    return hidden_states

fastvideo.models.dits.cosmos.CosmosRotaryPosEmbed

CosmosRotaryPosEmbed(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] = (2.0, 1.0, 1.0))
Source
class CosmosRotaryPosEmbed(nn.Module):

    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] = (2.0, 1.0, 1.0),
    ) -> None:
        super().__init__()

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

        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

        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]:
        fps = 16
        batch_size, num_channels, num_frames, height, width = hidden_states.shape
        pe_size = [num_frames // self.patch_size[0], height // self.patch_size[1], width // self.patch_size[2]]
        device = hidden_states.device

        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)

        emb_h = torch.outer(seq[:pe_size[1]], h_spatial_freqs)[None, :, None, :].repeat(pe_size[0], 1, pe_size[2], 1)
        emb_w = torch.outer(seq[:pe_size[2]], w_spatial_freqs)[None, None, :, :].repeat(pe_size[0], pe_size[1], 1, 1)

        if fps is None:
            emb_t = torch.outer(seq[:pe_size[0]], temporal_freqs)
        else:
            temporal_scale = seq[:pe_size[0]] / fps * self.base_fps
            emb_t = torch.outer(temporal_scale, temporal_freqs)

        emb_t = emb_t[:, None, None, :].repeat(1, pe_size[1], pe_size[2], 1)
        freqs = torch.cat([emb_t, emb_h, emb_w] * 2, dim=-1).flatten(0, 2).float()
        cos = torch.cos(freqs)
        sin = torch.sin(freqs)
        return cos, sin

fastvideo.models.dits.cosmos.CosmosRotaryPosEmbed.forward

forward(hidden_states: torch.Tensor, fps: int | None = None) -> tuple[torch.Tensor, torch.Tensor]
Source
def forward(self, hidden_states: torch.Tensor, fps: int | None = None) -> tuple[torch.Tensor, torch.Tensor]:
    fps = 16
    batch_size, num_channels, num_frames, height, width = hidden_states.shape
    pe_size = [num_frames // self.patch_size[0], height // self.patch_size[1], width // self.patch_size[2]]
    device = hidden_states.device

    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)

    emb_h = torch.outer(seq[:pe_size[1]], h_spatial_freqs)[None, :, None, :].repeat(pe_size[0], 1, pe_size[2], 1)
    emb_w = torch.outer(seq[:pe_size[2]], w_spatial_freqs)[None, None, :, :].repeat(pe_size[0], pe_size[1], 1, 1)

    if fps is None:
        emb_t = torch.outer(seq[:pe_size[0]], temporal_freqs)
    else:
        temporal_scale = seq[:pe_size[0]] / fps * self.base_fps
        emb_t = torch.outer(temporal_scale, temporal_freqs)

    emb_t = emb_t[:, None, None, :].repeat(1, pe_size[1], pe_size[2], 1)
    freqs = torch.cat([emb_t, emb_h, emb_w] * 2, dim=-1).flatten(0, 2).float()
    cos = torch.cos(freqs)
    sin = torch.sin(freqs)
    return cos, sin

fastvideo.models.dits.cosmos.CosmosLearnablePositionalEmbed

CosmosLearnablePositionalEmbed(hidden_size: int, max_size: tuple[int, int, int], patch_size: tuple[int, int, int], eps: float = 1e-06)
Source
class CosmosLearnablePositionalEmbed(nn.Module):

    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=False)]
        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:
        batch_size, num_channels, num_frames, height, width = hidden_states.shape
        pe_size = [num_frames // self.patch_size[0], height // self.patch_size[1], width // self.patch_size[2]]

        emb_t = self.pos_emb_t[:pe_size[0]][None, :, None, None, :].repeat(batch_size, 1, pe_size[1], pe_size[2], 1)
        emb_h = self.pos_emb_h[:pe_size[1]][None, None, :, None, :].repeat(batch_size, pe_size[0], 1, pe_size[2], 1)
        emb_w = self.pos_emb_w[:pe_size[2]][None, None, None, :, :].repeat(batch_size, pe_size[0], pe_size[1], 1, 1)
        emb = emb_t + emb_h + emb_w
        emb = emb.flatten(1, 3)

        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.cosmos.CosmosLearnablePositionalEmbed.forward

forward(hidden_states: torch.Tensor) -> torch.Tensor
Source
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    batch_size, num_channels, num_frames, height, width = hidden_states.shape
    pe_size = [num_frames // self.patch_size[0], height // self.patch_size[1], width // self.patch_size[2]]

    emb_t = self.pos_emb_t[:pe_size[0]][None, :, None, None, :].repeat(batch_size, 1, pe_size[1], pe_size[2], 1)
    emb_h = self.pos_emb_h[:pe_size[1]][None, None, :, None, :].repeat(batch_size, pe_size[0], 1, pe_size[2], 1)
    emb_w = self.pos_emb_w[:pe_size[2]][None, None, None, :, :].repeat(batch_size, pe_size[0], pe_size[1], 1, 1)
    emb = emb_t + emb_h + emb_w
    emb = emb.flatten(1, 3)

    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.cosmos.CosmosTransformer3DModel

CosmosTransformer3DModel(config: CosmosVideoConfig, hf_config: dict[str, Any])
Source
class CosmosTransformer3DModel(BaseDiT):
    _fsdp_shard_conditions = CosmosVideoConfig()._fsdp_shard_conditions
    _compile_conditions = CosmosVideoConfig()._compile_conditions
    # _supported_attention_backends = CosmosVideoConfig()._supported_attention_backends
    param_names_mapping = CosmosVideoConfig().param_names_mapping
    lora_param_names_mapping = CosmosVideoConfig().lora_param_names_mapping

    def __init__(self, config: CosmosVideoConfig, 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 = config.hidden_size
        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.extra_pos_embed_type = config.extra_pos_embed_type

        # 1. Patch Embedding
        # config.in_channels already includes the condition_mask channel
        # (HF config: in_channels=17 = 16 latent + 1 condition_mask).
        # Only add +1 for the padding_mask when concat_padding_mask=True.
        patch_embed_in_channels = config.in_channels + (1 if config.concat_padding_mask else 0)
        self.patch_embed = CosmosPatchEmbed(patch_embed_in_channels, inner_dim, config.patch_size, bias=False)

        # 2. Positional Embedding
        self.rope = CosmosRotaryPosEmbed(hidden_size=config.attention_head_dim,
                                         max_size=config.max_size,
                                         patch_size=config.patch_size,
                                         rope_scale=config.rope_scale)

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

        # 3. Time Embedding
        self.time_embed = CosmosEmbedding(inner_dim, inner_dim)

        # 4. Transformer Blocks
        self.transformer_blocks = nn.ModuleList([
            CosmosTransformerBlock(
                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=config.adaln_lora_dim,
                qk_norm=config.qk_norm,
                out_bias=False,
                prefix=f"{config.prefix}.transformer_blocks.{i}",
            ) for i in range(config.num_layers)
        ])

        # 5. Output norm & projection
        self.norm_out = CosmosAdaLayerNorm(inner_dim, config.adaln_lora_dim)
        self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size), bias=False)

        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_mask: torch.Tensor | None = None,
                padding_mask: torch.Tensor | None = None,
                **kwargs) -> torch.Tensor:
        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

        # Defensive dtype alignment: the Cosmos checkpoint is bf16 but
        # FSDP-wrapped training copies may report fp32 via
        # `next(parameters()).dtype`, which disables autocast in the
        # shared denoising stage and feeds fp32 tensors into bf16
        # weights. Cast every external input to the patch_embed weight
        # dtype so the model forward is robust regardless of caller.
        _target_dtype = self.patch_embed.proj.weight.dtype
        if hidden_states.dtype != _target_dtype:
            hidden_states = hidden_states.to(_target_dtype)
        if condition_mask is not None and condition_mask.dtype != _target_dtype:
            condition_mask = condition_mask.to(_target_dtype)
        if padding_mask is not None and padding_mask.dtype != _target_dtype:
            padding_mask = padding_mask.to(_target_dtype)
        if isinstance(encoder_hidden_states, torch.Tensor):
            if encoder_hidden_states.dtype != _target_dtype:
                encoder_hidden_states = encoder_hidden_states.to(_target_dtype)
        else:
            encoder_hidden_states = [
                t.to(_target_dtype) if t.dtype != _target_dtype else t for t in encoder_hidden_states
            ]

        # 1. Concatenate padding mask if needed & prepare attention mask
        if condition_mask is not None:
            hidden_states = torch.cat([hidden_states, condition_mask], dim=1)

        if self.concat_padding_mask:
            from torchvision import transforms
            padding_mask = transforms.functional.resize(padding_mask,
                                                        list(hidden_states.shape[-2:]),
                                                        interpolation=transforms.InterpolationMode.NEAREST)
            # torchvision.resize may upcast bf16/fp16 → fp32; restore
            # hidden_states' dtype so the subsequent cat doesn't promote
            # everything and break patch_embed (bf16 weights).
            padding_mask = padding_mask.to(hidden_states.dtype)
            hidden_states = torch.cat(
                [hidden_states, padding_mask.unsqueeze(2).repeat(batch_size, 1, num_frames, 1, 1)], dim=1)

        if attention_mask is not None:
            attention_mask = attention_mask.unsqueeze(1).unsqueeze(1)  # [B, 1, 1, S]

        # 2. Generate positional embeddings
        image_rotary_emb = self.rope(hidden_states, fps=fps)
        extra_pos_emb = self.learnable_pos_embed(hidden_states) if self.extra_pos_embed_type == "learnable" else None

        # 3. Patchify input
        p_t, p_h, p_w = self.patch_size
        post_patch_num_frames = num_frames // p_t
        post_patch_height = height // p_h
        post_patch_width = width // p_w
        hidden_states = self.patch_embed(hidden_states)
        hidden_states = hidden_states.flatten(1, 3)  # [B, T, H, W, C] -> [B, THW, C] codespell:ignore

        # 4. Timestep embeddings
        if timestep.ndim == 1:
            temb, embedded_timestep = self.time_embed(hidden_states, timestep)
        elif timestep.ndim == 5:
            assert timestep.shape == (batch_size, 1, num_frames, 1,
                                      1), (f"Expected timestep to have shape [B, 1, T, 1, 1], but got {timestep.shape}")
            timestep = timestep.flatten()
            temb, embedded_timestep = self.time_embed(hidden_states, timestep)
            # We can do this because num_frames == post_patch_num_frames, as p_t is 1
            temb, embedded_timestep = (
                x.view(batch_size, post_patch_num_frames, 1, 1, -1).expand(-1, -1, post_patch_height, post_patch_width,
                                                                           -1).flatten(1, 3)
                for x in (temb, embedded_timestep)
            )  # [BT, C] -> [B, T, 1, 1, C] -> [B, T, H, W, C] -> [B, THW, C] codespell:ignore
        else:
            raise ValueError(f"Unsupported timestep shape: {timestep.shape}")

        # 6. Transformer blocks
        if torch.is_grad_enabled() and self.gradient_checkpointing:
            for i, block in enumerate(self.transformer_blocks):
                hidden_states = self._gradient_checkpointing_func(
                    block,
                    hidden_states,
                    encoder_hidden_states,
                    embedded_timestep,
                    temb,
                    image_rotary_emb,
                    extra_pos_emb,
                    attention_mask,
                )
        else:
            for i, block in enumerate(self.transformer_blocks):
                hidden_states = block(
                    hidden_states=hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    embedded_timestep=embedded_timestep,
                    temb=temb,
                    image_rotary_emb=image_rotary_emb,
                    extra_pos_emb=extra_pos_emb,
                    attention_mask=attention_mask,
                )

        # 7. Output norm & projection & unpatchify
        hidden_states = self.norm_out(hidden_states, embedded_timestep, temb)
        hidden_states = self.proj_out(hidden_states)
        hidden_states = hidden_states.unflatten(2, (p_h, p_w, p_t, -1))
        hidden_states = hidden_states.unflatten(1, (post_patch_num_frames, post_patch_height, post_patch_width))
        # NOTE: The permutation order here is not the inverse operation of what happens when patching as usually expected.
        # It might be a source of confusion to the reader, but this is correct
        hidden_states = hidden_states.permute(0, 7, 1, 6, 2, 4, 3, 5)
        hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)

        # Return as tuple for compatibility with callers that
        # do `transformer(..., return_dict=False)[0]` (diffusers
        # convention used by CosmosDenoisingStage).
        if not kwargs.get("return_dict", True):
            return (hidden_states, )
        return hidden_states

fastvideo.models.dits.cosmos.CosmosTransformer3DModel.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_mask: torch.Tensor | None = None, padding_mask: torch.Tensor | None = None, **kwargs) -> torch.Tensor
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_mask: torch.Tensor | None = None,
            padding_mask: torch.Tensor | None = None,
            **kwargs) -> torch.Tensor:
    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

    # Defensive dtype alignment: the Cosmos checkpoint is bf16 but
    # FSDP-wrapped training copies may report fp32 via
    # `next(parameters()).dtype`, which disables autocast in the
    # shared denoising stage and feeds fp32 tensors into bf16
    # weights. Cast every external input to the patch_embed weight
    # dtype so the model forward is robust regardless of caller.
    _target_dtype = self.patch_embed.proj.weight.dtype
    if hidden_states.dtype != _target_dtype:
        hidden_states = hidden_states.to(_target_dtype)
    if condition_mask is not None and condition_mask.dtype != _target_dtype:
        condition_mask = condition_mask.to(_target_dtype)
    if padding_mask is not None and padding_mask.dtype != _target_dtype:
        padding_mask = padding_mask.to(_target_dtype)
    if isinstance(encoder_hidden_states, torch.Tensor):
        if encoder_hidden_states.dtype != _target_dtype:
            encoder_hidden_states = encoder_hidden_states.to(_target_dtype)
    else:
        encoder_hidden_states = [
            t.to(_target_dtype) if t.dtype != _target_dtype else t for t in encoder_hidden_states
        ]

    # 1. Concatenate padding mask if needed & prepare attention mask
    if condition_mask is not None:
        hidden_states = torch.cat([hidden_states, condition_mask], dim=1)

    if self.concat_padding_mask:
        from torchvision import transforms
        padding_mask = transforms.functional.resize(padding_mask,
                                                    list(hidden_states.shape[-2:]),
                                                    interpolation=transforms.InterpolationMode.NEAREST)
        # torchvision.resize may upcast bf16/fp16 → fp32; restore
        # hidden_states' dtype so the subsequent cat doesn't promote
        # everything and break patch_embed (bf16 weights).
        padding_mask = padding_mask.to(hidden_states.dtype)
        hidden_states = torch.cat(
            [hidden_states, padding_mask.unsqueeze(2).repeat(batch_size, 1, num_frames, 1, 1)], dim=1)

    if attention_mask is not None:
        attention_mask = attention_mask.unsqueeze(1).unsqueeze(1)  # [B, 1, 1, S]

    # 2. Generate positional embeddings
    image_rotary_emb = self.rope(hidden_states, fps=fps)
    extra_pos_emb = self.learnable_pos_embed(hidden_states) if self.extra_pos_embed_type == "learnable" else None

    # 3. Patchify input
    p_t, p_h, p_w = self.patch_size
    post_patch_num_frames = num_frames // p_t
    post_patch_height = height // p_h
    post_patch_width = width // p_w
    hidden_states = self.patch_embed(hidden_states)
    hidden_states = hidden_states.flatten(1, 3)  # [B, T, H, W, C] -> [B, THW, C] codespell:ignore

    # 4. Timestep embeddings
    if timestep.ndim == 1:
        temb, embedded_timestep = self.time_embed(hidden_states, timestep)
    elif timestep.ndim == 5:
        assert timestep.shape == (batch_size, 1, num_frames, 1,
                                  1), (f"Expected timestep to have shape [B, 1, T, 1, 1], but got {timestep.shape}")
        timestep = timestep.flatten()
        temb, embedded_timestep = self.time_embed(hidden_states, timestep)
        # We can do this because num_frames == post_patch_num_frames, as p_t is 1
        temb, embedded_timestep = (
            x.view(batch_size, post_patch_num_frames, 1, 1, -1).expand(-1, -1, post_patch_height, post_patch_width,
                                                                       -1).flatten(1, 3)
            for x in (temb, embedded_timestep)
        )  # [BT, C] -> [B, T, 1, 1, C] -> [B, T, H, W, C] -> [B, THW, C] codespell:ignore
    else:
        raise ValueError(f"Unsupported timestep shape: {timestep.shape}")

    # 6. Transformer blocks
    if torch.is_grad_enabled() and self.gradient_checkpointing:
        for i, block in enumerate(self.transformer_blocks):
            hidden_states = self._gradient_checkpointing_func(
                block,
                hidden_states,
                encoder_hidden_states,
                embedded_timestep,
                temb,
                image_rotary_emb,
                extra_pos_emb,
                attention_mask,
            )
    else:
        for i, block in enumerate(self.transformer_blocks):
            hidden_states = block(
                hidden_states=hidden_states,
                encoder_hidden_states=encoder_hidden_states,
                embedded_timestep=embedded_timestep,
                temb=temb,
                image_rotary_emb=image_rotary_emb,
                extra_pos_emb=extra_pos_emb,
                attention_mask=attention_mask,
            )

    # 7. Output norm & projection & unpatchify
    hidden_states = self.norm_out(hidden_states, embedded_timestep, temb)
    hidden_states = self.proj_out(hidden_states)
    hidden_states = hidden_states.unflatten(2, (p_h, p_w, p_t, -1))
    hidden_states = hidden_states.unflatten(1, (post_patch_num_frames, post_patch_height, post_patch_width))
    # NOTE: The permutation order here is not the inverse operation of what happens when patching as usually expected.
    # It might be a source of confusion to the reader, but this is correct
    hidden_states = hidden_states.permute(0, 7, 1, 6, 2, 4, 3, 5)
    hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)

    # Return as tuple for compatibility with callers that
    # do `transformer(..., return_dict=False)[0]` (diffusers
    # convention used by CosmosDenoisingStage).
    if not kwargs.get("return_dict", True):
        return (hidden_states, )
    return hidden_states