daVinci-MagiHuman DiT (base variant).

Ported from https://github.com/GAIR-NLP/daVinci-MagiHuman (inference/model/dit/dit_module.py, ~950 lines in the reference).

Architecture summary (verified against GAIR/daVinci-MagiHuman/base/ weights):

40 transformer layers, hidden 5120, head_dim 128.
GQA with 40 query heads and 8 KV heads.
Multi-modality "sandwich": layers 0..3 and 36..39 use 3-way modality experts (video/audio/text) packed inside each linear as weight[..., out * 3, in]. Middle layers share a single expert.
Per-head attention gating: the QKV projection emits an extra num_heads_q channels that are sigmoid-gated onto the attention output.
Activation is GELU7 on layers 0..3 (non-gated, intermediate=4hidden) and SwiGLU7 elsewhere (gated, intermediate=int(hidden42/3)//44).
Position encoding is an element-wise Fourier embedding over 9-column coords (t,h,w + original TxHxW + reference TxHxW), not a standard 1D/3D RoPE.
Forward takes a flat concatenated token stream (video first, then audio, then text) plus a modality map; the internal ModalityDispatcher permutes by modality before each linear so per-expert chunks line up.

Deviations from the "use fastvideo.layers primitives everywhere" guideline in the add-model skill:

The packed-expert linears store weight as [out * num_experts, in]. FastVideo's ReplicatedLinear does not model this layout; we use raw nn.Parameter with a small wrapper below. This is deliberate and scoped to this DiT: ReplicatedLinear still handles the adapter.* embedders and final_linear_{video,audio} (single-expert) here.
Self-attention is full-sequence and crosses modalities inside the flat concat stream; DistributedAttention assumes a clean spatial-sequence layout, so for the first port we use torch SDPA. Multi-GPU sequence parallelism is a follow-up.
torch.compile via magi_compiler is replaced with a plain nn.Module.

For the full history and shape-by-shape verification notes, see .claude/skills/add-model/SKILL.md and the scaffold PR description.

`fastvideo.models.dits.magi_human.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.magi_human.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.magi_human.MagiHumanArchConfig`

MagiHumanArchConfig(stacked_params_mapping: list[tuple[str, str, str]] = list(), _fsdp_shard_conditions: list = (lambda: [_is_block_layer])(), _compile_conditions: list = list(), param_names_mapping: dict = dict(), reverse_param_names_mapping: dict = dict(), lora_param_names_mapping: dict = dict(), cast_prompt_embeds_to_dit_dtype: bool = False, _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (AttentionBackendEnum.SAGE_ATTN, AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.SAGE_ATTN_THREE, AttentionBackendEnum.ATTN_QAT_INFER, AttentionBackendEnum.ATTN_QAT_TRAIN, AttentionBackendEnum.SLA_ATTN, AttentionBackendEnum.SAGE_SLA_ATTN), hidden_size: int = 5120, num_attention_heads: int = 0, num_channels_latents: int = 0, in_channels: int = 0, out_channels: int = 0, exclude_lora_layers: list[str] = list(), boundary_ratio: float | None = None, num_layers: int = 40, head_dim: int = 128, num_query_groups: int = 8, video_in_channels: int = 192, audio_in_channels: int = 64, text_in_channels: int = 3584, mm_layers: tuple[int, ...] = (0, 1, 2, 3, 36, 37, 38, 39), local_attn_layers: tuple[int, ...] = (), gelu7_layers: tuple[int, ...] = (0, 1, 2, 3), post_norm_layers: tuple[int, ...] = (), enable_attn_gating: bool = True, activation_type: str = 'swiglu7', patch_size: tuple[int, int, int] = (1, 2, 2), spatial_rope_interpolation: str = 'extra', tread_selection_rate: float = 0.5, tread_start_layer_idx: int = 2, tread_end_layer_idx: int = 25, num_heads_kv: int = 0)

MagiHuman base DiT architecture constants.

Scope contract: fields here must match the transformer/config.json emitted by scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py 1:1, and both are sourced from the upstream Python reference inference/common/config.py::ModelConfig (the HF root config.json is empty so the Python source is canonical). Pipeline-level knobs (VAE stride, fps, num_inference_steps, CFG scales, flow_shift, t5_gemma_target_length) and data-proxy knobs (coords_style, frame_receptive_field, ref_audio_offset, text_offset) live on MagiHumanBaseConfig, NOT here.

param_names_mapping is intentionally empty: the FastVideo implementation keeps the same module tree as the reference (adapter.*, block.layers.<i>.*, final_linear_{video,audio}.*, final_norm_{video,audio}.*), so converted weights load directly.

`fastvideo.models.dits.magi_human.MagiHumanVideoConfig`

MagiHumanVideoConfig(arch_config: DiTArchConfig = MagiHumanArchConfig(), prefix: str = 'magi_human', quant_config: QuantizationConfig | None = None, *, _resolved_attention_backend: AttentionBackendEnum | None = None)

`fastvideo.models.dits.magi_human.BaseDiT`

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

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

inference_config_overrides() -> dict[str, Any]

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

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

AttentionBackendEnum()

`fastvideo.models.dits.magi_human.Modality`

Modality()

Source

class Modality(IntEnum):
    VIDEO = 0
    AUDIO = 1
    TEXT = 2

`fastvideo.models.dits.magi_human.swiglu7`

swiglu7(x: torch.Tensor, alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor

Gated swish-GLU with OpenAI-OSS-style limits and +1 linear bias.

Source

def swiglu7(x: torch.Tensor, alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor:
    """Gated swish-GLU with OpenAI-OSS-style limits and +1 linear bias."""
    in_dtype = x.dtype
    x = x.to(torch.float32)
    x_glu, x_linear = x[..., ::2], x[..., 1::2]
    x_glu = x_glu.clamp(max=limit)
    x_linear = x_linear.clamp(min=-limit, max=limit)
    out_glu = x_glu * torch.sigmoid(alpha * x_glu)
    return (out_glu * (x_linear + 1)).to(in_dtype)

`fastvideo.models.dits.magi_human.gelu7`

gelu7(x: torch.Tensor, alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor

Source

def gelu7(x: torch.Tensor, alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor:
    in_dtype = x.dtype
    x = x.to(torch.float32).clamp(max=limit)
    return (x * torch.sigmoid(alpha * x)).to(in_dtype)

`fastvideo.models.dits.magi_human.ModalityDispatcher`

ModalityDispatcher(modality_mapping: torch.Tensor, num_modalities: int)

Permute a flat token stream so same-modality tokens are contiguous.

The DiT's multi-expert linears apply a different weight chunk per modality. Instead of carrying a branch inside each Linear, we pre-permute tokens so each chunk sees a contiguous slice, then un-permute before computing RoPE/attention across the full sequence.

Source

class ModalityDispatcher:
    """Permute a flat token stream so same-modality tokens are contiguous.

    The DiT's multi-expert linears apply a different weight chunk per modality.
    Instead of carrying a branch inside each Linear, we pre-permute tokens so
    each chunk sees a contiguous slice, then un-permute before computing
    RoPE/attention across the full sequence.
    """

    def __init__(self, modality_mapping: torch.Tensor, num_modalities: int):
        self.modality_mapping = modality_mapping
        self.num_modalities = num_modalities
        self.permute_mapping = torch.argsort(modality_mapping)
        self.inv_permute_mapping = torch.argsort(self.permute_mapping)
        permuted = modality_mapping[self.permute_mapping]
        self.group_size = torch.bincount(permuted, minlength=num_modalities).to(torch.int32)
        self.group_size_cpu: list[int] = [int(x) for x in self.group_size.cpu().tolist()]

    def dispatch(self, x: torch.Tensor) -> list[torch.Tensor]:
        return list(torch.split(x, self.group_size_cpu, dim=0))

    def undispatch(self, *chunks: torch.Tensor) -> torch.Tensor:
        return torch.cat(chunks, dim=0)

    @staticmethod
    def permute(x: torch.Tensor, permute_mapping: torch.Tensor) -> torch.Tensor:
        return x[permute_mapping]

    @staticmethod
    def inv_permute(x: torch.Tensor, inv_permute_mapping: torch.Tensor) -> torch.Tensor:
        return x[inv_permute_mapping]

`fastvideo.models.dits.magi_human.ModalityDispatcher.dispatch`

dispatch(x: torch.Tensor) -> list[torch.Tensor]

Source

def dispatch(self, x: torch.Tensor) -> list[torch.Tensor]:
    return list(torch.split(x, self.group_size_cpu, dim=0))

`fastvideo.models.dits.magi_human.ModalityDispatcher.undispatch`

undispatch(*chunks: torch.Tensor) -> torch.Tensor

Source

def undispatch(self, *chunks: torch.Tensor) -> torch.Tensor:
    return torch.cat(chunks, dim=0)

`fastvideo.models.dits.magi_human.ModalityDispatcher.permute`

permute(x: torch.Tensor, permute_mapping: torch.Tensor) -> torch.Tensor

Source

@staticmethod
def permute(x: torch.Tensor, permute_mapping: torch.Tensor) -> torch.Tensor:
    return x[permute_mapping]

`fastvideo.models.dits.magi_human.ModalityDispatcher.inv_permute`

inv_permute(x: torch.Tensor, inv_permute_mapping: torch.Tensor) -> torch.Tensor

Source

@staticmethod
def inv_permute(x: torch.Tensor, inv_permute_mapping: torch.Tensor) -> torch.Tensor:
    return x[inv_permute_mapping]

`fastvideo.models.dits.magi_human.MultiModalityRMSNorm`

MultiModalityRMSNorm(dim: int, eps: float = 1e-06, num_modality: int = 1)

RMSNorm with optional per-modality scale.

When num_modality == 1, behaves identically to a standard RMSNorm with weight initialized to zero (effective weight is 1 + weight, hence the learnable +1 offset baked into the forward path). When num_modality > 1, the weight tensor packs per-modality scales along its flat axis and the dispatcher selects the right chunk per modality.

Source

class MultiModalityRMSNorm(nn.Module):
    """RMSNorm with optional per-modality scale.

    When num_modality == 1, behaves identically to a standard RMSNorm with
    weight initialized to zero (effective weight is 1 + weight, hence the
    learnable +1 offset baked into the forward path). When num_modality > 1,
    the weight tensor packs per-modality scales along its flat axis and the
    dispatcher selects the right chunk per modality.
    """

    def __init__(self, dim: int, eps: float = 1e-6, num_modality: int = 1):
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.num_modality = num_modality
        # Always stored in fp32; matches the reference initialization.
        self.weight = nn.Parameter(torch.zeros(dim * num_modality, dtype=torch.float32))

    def _rms(self, x: torch.Tensor) -> torch.Tensor:
        t = x.float()
        return t * torch.rsqrt(t.pow(2).mean(dim=-1, keepdim=True) + self.eps)

    def forward(
        self,
        x: torch.Tensor,
        modality_dispatcher: Optional[ModalityDispatcher] = None,
    ) -> torch.Tensor:
        original_dtype = x.dtype
        t = self._rms(x)
        if self.num_modality == 1:
            return (t * (self.weight + 1)).to(original_dtype)
        assert modality_dispatcher is not None, ("MultiModalityRMSNorm with num_modality>1 requires a dispatcher")
        weight_chunks = self.weight.chunk(self.num_modality, dim=0)
        parts = modality_dispatcher.dispatch(t)
        for i in range(self.num_modality):
            parts[i] = parts[i] * (weight_chunks[i] + 1)
        return modality_dispatcher.undispatch(*parts).to(original_dtype)

`fastvideo.models.dits.magi_human.MultiModalityRMSNorm.forward`

forward(x: torch.Tensor, modality_dispatcher: Optional[ModalityDispatcher] = None) -> torch.Tensor

Source

def forward(
    self,
    x: torch.Tensor,
    modality_dispatcher: Optional[ModalityDispatcher] = None,
) -> torch.Tensor:
    original_dtype = x.dtype
    t = self._rms(x)
    if self.num_modality == 1:
        return (t * (self.weight + 1)).to(original_dtype)
    assert modality_dispatcher is not None, ("MultiModalityRMSNorm with num_modality>1 requires a dispatcher")
    weight_chunks = self.weight.chunk(self.num_modality, dim=0)
    parts = modality_dispatcher.dispatch(t)
    for i in range(self.num_modality):
        parts[i] = parts[i] * (weight_chunks[i] + 1)
    return modality_dispatcher.undispatch(*parts).to(original_dtype)

`fastvideo.models.dits.magi_human.ElementWiseFourierEmbed`

ElementWiseFourierEmbed(dim: int, temperature: float = 10000.0, dtype: torch.dtype = torch.float32)

Element-wise Fourier embedding over 9-column coords (t, h, w, T, H, W, ref_T, ref_H, ref_W). Produces a per-token positional embedding that acts as the RoPE angle input for attention.

Weight: bands of shape [dim // 8] (fixed at init via freq_bands).

Source

class ElementWiseFourierEmbed(nn.Module):
    """Element-wise Fourier embedding over 9-column coords (t, h, w, T, H, W,
    ref_T, ref_H, ref_W). Produces a per-token positional embedding that
    acts as the RoPE angle input for attention.

    Weight: `bands` of shape `[dim // 8]` (fixed at init via freq_bands).
    """

    def __init__(
        self,
        dim: int,
        temperature: float = 10000.0,
        dtype: torch.dtype = torch.float32,
    ):
        super().__init__()
        self.dim = dim
        bands = _freq_bands(dim // 8, temperature=temperature).to(dtype)
        # `register_buffer` so state_dict keeps it, matching upstream naming.
        self.register_buffer("bands", bands)

    def forward(self, coords: torch.Tensor) -> torch.Tensor:
        # coords: [L, 9] = (t, h, w, T, H, W, ref_T, ref_H, ref_W)
        coords_xyz = coords[:, :3]
        sizes = coords[:, 3:6]
        refs = coords[:, 6:9]

        scales = (refs - 1) / (sizes - 1)
        scales[(refs == 1) & (sizes == 1)] = 1
        # Center H and W (leave time uncentered).
        centers = (sizes - 1) / 2
        centers[:, 0] = 0
        coords_xyz = coords_xyz - centers

        proj = coords_xyz.unsqueeze(-1) * scales.unsqueeze(-1) * self.bands  # [L, 3, B]
        sin_proj = proj.sin()
        cos_proj = proj.cos()
        return torch.cat((sin_proj, cos_proj), dim=1).flatten(1)

`fastvideo.models.dits.magi_human.ElementWiseFourierEmbed.forward`

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

Source

def forward(self, coords: torch.Tensor) -> torch.Tensor:
    # coords: [L, 9] = (t, h, w, T, H, W, ref_T, ref_H, ref_W)
    coords_xyz = coords[:, :3]
    sizes = coords[:, 3:6]
    refs = coords[:, 6:9]

    scales = (refs - 1) / (sizes - 1)
    scales[(refs == 1) & (sizes == 1)] = 1
    # Center H and W (leave time uncentered).
    centers = (sizes - 1) / 2
    centers[:, 0] = 0
    coords_xyz = coords_xyz - centers

    proj = coords_xyz.unsqueeze(-1) * scales.unsqueeze(-1) * self.bands  # [L, 3, B]
    sin_proj = proj.sin()
    cos_proj = proj.cos()
    return torch.cat((sin_proj, cos_proj), dim=1).flatten(1)

`fastvideo.models.dits.magi_human.PackedExpertLinear`

PackedExpertLinear(in_features: int, out_features: int, num_experts: int = 1, bias: bool = False, dtype: torch.dtype = torch.bfloat16)

Linear where the weight is packed per-modality along the output axis.

Shapes

weight: [out_features * num_experts, in_features] bias: [out_features * num_experts] (optional)

When num_experts == 1, behaves exactly like nn.Linear. When num_experts > 1, forward dispatches the input via the supplied ModalityDispatcher, applies the per-modality weight/bias chunk, and gathers the outputs in original order.

Why not use ReplicatedLinear? Because the packed-expert layout is not what ReplicatedLinear (or any other fastvideo.layers.linear) is wired for. Using raw nn.Parameter keeps weight loading trivial (names map 1:1 to the upstream checkpoint) and avoids quantization-path assumptions that don't match this layout.

Source

class PackedExpertLinear(nn.Module):
    """Linear where the weight is packed per-modality along the output axis.

    Shapes:
        weight: [out_features * num_experts, in_features]
        bias:   [out_features * num_experts]  (optional)

    When `num_experts == 1`, behaves exactly like `nn.Linear`. When
    `num_experts > 1`, `forward` dispatches the input via the supplied
    `ModalityDispatcher`, applies the per-modality weight/bias chunk, and
    gathers the outputs in original order.

    Why not use `ReplicatedLinear`? Because the packed-expert layout is not
    what ReplicatedLinear (or any other fastvideo.layers.linear) is wired
    for. Using raw `nn.Parameter` keeps weight loading trivial (names map
    1:1 to the upstream checkpoint) and avoids quantization-path assumptions
    that don't match this layout.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        num_experts: int = 1,
        bias: bool = False,
        dtype: torch.dtype = torch.bfloat16,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.num_experts = num_experts
        self.use_bias = bias
        self.weight = nn.Parameter(torch.empty(out_features * num_experts, in_features, dtype=dtype))
        if bias:
            self.bias = nn.Parameter(torch.empty(out_features * num_experts, dtype=dtype))
        else:
            self.register_parameter("bias", None)

    def forward(
        self,
        x: torch.Tensor,
        modality_dispatcher: Optional[ModalityDispatcher] = None,
    ) -> torch.Tensor:
        if self.num_experts == 1:
            return F.linear(x, self.weight, self.bias)
        assert modality_dispatcher is not None, ("PackedExpertLinear with num_experts>1 requires a dispatcher")
        parts = modality_dispatcher.dispatch(x)
        w_chunks = self.weight.chunk(self.num_experts, dim=0)
        b_chunks = (self.bias.chunk(self.num_experts, dim=0) if self.bias is not None else [None] * self.num_experts)
        for i in range(self.num_experts):
            parts[i] = F.linear(parts[i], w_chunks[i], b_chunks[i])
        return modality_dispatcher.undispatch(*parts)

`fastvideo.models.dits.magi_human.PackedExpertLinear.forward`

forward(x: torch.Tensor, modality_dispatcher: Optional[ModalityDispatcher] = None) -> torch.Tensor

Source

def forward(
    self,
    x: torch.Tensor,
    modality_dispatcher: Optional[ModalityDispatcher] = None,
) -> torch.Tensor:
    if self.num_experts == 1:
        return F.linear(x, self.weight, self.bias)
    assert modality_dispatcher is not None, ("PackedExpertLinear with num_experts>1 requires a dispatcher")
    parts = modality_dispatcher.dispatch(x)
    w_chunks = self.weight.chunk(self.num_experts, dim=0)
    b_chunks = (self.bias.chunk(self.num_experts, dim=0) if self.bias is not None else [None] * self.num_experts)
    for i in range(self.num_experts):
        parts[i] = F.linear(parts[i], w_chunks[i], b_chunks[i])
    return modality_dispatcher.undispatch(*parts)

`fastvideo.models.dits.magi_human.AttentionSubConfig`

AttentionSubConfig(hidden_size: int, num_heads_q: int, num_heads_kv: int, head_dim: int, num_modality: int, enable_attn_gating: bool, use_local_attn: bool = False, frame_receptive_field: int = 11)

Source

@dataclass
class AttentionSubConfig:
    hidden_size: int
    num_heads_q: int
    num_heads_kv: int
    head_dim: int
    num_modality: int
    enable_attn_gating: bool
    use_local_attn: bool = False
    frame_receptive_field: int = 11

`fastvideo.models.dits.magi_human.MagiAttention`

MagiAttention(cfg: AttentionSubConfig)

Self-attention with GQA + optional per-head sigmoid gating.

Source

class MagiAttention(nn.Module):
    """Self-attention with GQA + optional per-head sigmoid gating."""

    def __init__(self, cfg: AttentionSubConfig):
        super().__init__()
        self.cfg = cfg
        self.gating_size = cfg.num_heads_q if cfg.enable_attn_gating else 0
        qkv_out = (cfg.num_heads_q * cfg.head_dim + 2 * cfg.num_heads_kv * cfg.head_dim + self.gating_size)
        self.pre_norm = MultiModalityRMSNorm(cfg.hidden_size, num_modality=cfg.num_modality)
        self.linear_qkv = PackedExpertLinear(
            cfg.hidden_size,
            qkv_out,
            num_experts=cfg.num_modality,
            bias=False,
        )
        self.linear_proj = PackedExpertLinear(
            cfg.num_heads_q * cfg.head_dim,
            cfg.hidden_size,
            num_experts=cfg.num_modality,
            bias=False,
        )
        self.q_norm = MultiModalityRMSNorm(cfg.head_dim, num_modality=cfg.num_modality)
        self.k_norm = MultiModalityRMSNorm(cfg.head_dim, num_modality=cfg.num_modality)

        self.q_size = cfg.num_heads_q * cfg.head_dim
        self.kv_size = cfg.num_heads_kv * cfg.head_dim

        self.attn = LocalAttention(
            num_heads=cfg.num_heads_q,
            head_size=cfg.head_dim,
            num_kv_heads=cfg.num_heads_kv,
            causal=False,
            supported_attention_backends=(
                AttentionBackendEnum.FLASH_ATTN,
                AttentionBackendEnum.TORCH_SDPA,
            ),
        )

    def configure_local_attention(
        self,
        *,
        enabled: bool,
        frame_receptive_field: int = 11,
    ) -> None:
        self.cfg.use_local_attn = enabled
        self.cfg.frame_receptive_field = frame_receptive_field

    def _sdpa(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        """Run SDPA on [L, H, D] tensors and return [L, Hq, D]."""
        if q.numel() == 0:
            return q.new_empty(q.shape)
        out = F.scaled_dot_product_attention(
            q.transpose(0, 1).unsqueeze(0).contiguous(),
            k.transpose(0, 1).unsqueeze(0).contiguous(),
            v.transpose(0, 1).unsqueeze(0).contiguous(),
            enable_gqa=self.cfg.num_heads_q != self.cfg.num_heads_kv,
        )
        return out.squeeze(0).transpose(0, 1).contiguous()

    def _local_window_attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        *,
        num_video_tokens: int,
        num_frames: int,
    ) -> torch.Tensor:
        """Approximate upstream FFAHandler block accumulation with SDPA.

        SR-1080p's reference kernel sums three independently-normalized
        attention contributions:

        * video frame queries -> local-window video keys;
        * all video queries -> all audio+text keys;
        * all audio+text queries -> full sequence keys.

        This method mirrors that accumulator semantics with ordinary SDPA
        slices. It is intentionally scoped to single-process inference; layers
        without ``use_local_attn`` keep the existing full LocalAttention path.
        """
        if num_frames <= 0 or num_video_tokens <= 0:
            return self._sdpa(q, k, v)
        if num_video_tokens % num_frames != 0:
            raise ValueError(f"MagiHuman local attention expects video tokens divisible by "
                             f"frames, got {num_video_tokens=} and {num_frames=}.")

        token_per_frame = num_video_tokens // num_frames
        out = torch.zeros(
            q.shape[0],
            self.cfg.num_heads_q,
            self.cfg.head_dim,
            device=q.device,
            dtype=q.dtype,
        )
        rf = int(self.cfg.frame_receptive_field)

        q_video = q[:num_video_tokens]
        k_video = k[:num_video_tokens]
        v_video = v[:num_video_tokens]
        for frame_idx in range(num_frames):
            q_start = frame_idx * token_per_frame
            q_end = q_start + token_per_frame
            k_start = max(0, (frame_idx - rf) * token_per_frame)
            k_end = min(num_video_tokens, (frame_idx + rf + 1) * token_per_frame)
            out[q_start:q_end] = self._sdpa(
                q_video[q_start:q_end],
                k_video[k_start:k_end],
                v_video[k_start:k_end],
            )

        if num_video_tokens < q.shape[0]:
            k_at = k[num_video_tokens:]
            v_at = v[num_video_tokens:]
            out[:num_video_tokens] = out[:num_video_tokens] + self._sdpa(
                q[:num_video_tokens],
                k_at,
                v_at,
            )
            out[num_video_tokens:] = self._sdpa(
                q[num_video_tokens:],
                k,
                v,
            )
        return out

    def forward(
        self,
        hidden_states: torch.Tensor,
        rope: torch.Tensor,
        permute_mapping: torch.Tensor,
        inv_permute_mapping: torch.Tensor,
        modality_dispatcher: ModalityDispatcher,
        num_video_tokens: int | None = None,
        num_frames: int | None = None,
    ) -> torch.Tensor:
        orig_dtype = self.linear_qkv.weight.dtype
        h = self.pre_norm(hidden_states, modality_dispatcher=modality_dispatcher).to(orig_dtype)
        qkv = self.linear_qkv(h, modality_dispatcher=modality_dispatcher).float()
        q, k, v, g = torch.split(
            qkv,
            [self.q_size, self.kv_size, self.kv_size, self.gating_size],
            dim=-1,
        )
        q = q.view(-1, self.cfg.num_heads_q, self.cfg.head_dim)
        k = k.view(-1, self.cfg.num_heads_kv, self.cfg.head_dim)
        v = v.view(-1, self.cfg.num_heads_kv, self.cfg.head_dim)
        g = g.view(-1, self.cfg.num_heads_q, 1) if self.gating_size else None

        q = self.q_norm(q, modality_dispatcher=modality_dispatcher)
        k = self.k_norm(k, modality_dispatcher=modality_dispatcher)

        # Un-permute before RoPE + attention so positional order reflects
        # the original (video, audio, text) concat — matches reference.
        q = ModalityDispatcher.inv_permute(q, inv_permute_mapping)
        k = ModalityDispatcher.inv_permute(k, inv_permute_mapping)
        v = ModalityDispatcher.inv_permute(v, inv_permute_mapping)
        if g is not None:
            g = ModalityDispatcher.inv_permute(g, inv_permute_mapping)

        # Element-wise Fourier embed packs sin/cos of 3 axes into a single
        # `rope` tensor. Match reference's split:
        #   sin_emb, cos_emb = rope.tensor_split(2, -1)
        # Reference passes (cos_emb, sin_emb) but splits sin first — replicated
        # exactly so weight parity holds. Partial RoPE: rope dim is
        # 6 * (head_dim // 8) = 96 < head_dim (128), so the trailing 32
        # head_dim positions stay unrotated, matching the reference.
        sin_emb, cos_emb = rope.tensor_split(2, -1)
        rot_dim = cos_emb.shape[-1] * 2
        q_rot = _apply_rotary_emb(q[..., :rot_dim], cos_emb, sin_emb, is_neox_style=True)
        k_rot = _apply_rotary_emb(k[..., :rot_dim], cos_emb, sin_emb, is_neox_style=True)
        if rot_dim < q.shape[-1]:
            q = torch.cat([q_rot, q[..., rot_dim:]], dim=-1)
            k = torch.cat([k_rot, k[..., rot_dim:]], dim=-1)
        else:
            q, k = q_rot, k_rot

        # Run SDPA via FastVideo's LocalAttention so the backend selection
        # (SDPA / FlashAttn / SLA / SageAttn) flows through the standard
        # configurable path. GQA is handled inside the SDPA backend via
        # `enable_gqa=True` when num_heads_q != num_heads_kv, so we no
        # longer need the manual `repeat_interleave` here.
        # Attention math runs at orig_dtype (bf16 in production and in the
        # parity test, since PackedExpertLinear's default is bf16, matching
        # upstream BaseLinear at dit_module.py:330). The gating multiply
        # promotes back to fp32 implicitly via PyTorch's type-promotion
        # rules: bf16_attn_out * sigmoid(fp32_g) -> fp32, mirroring upstream
        # dit_module.py:649.
        q = q.to(orig_dtype)
        k = k.to(orig_dtype)
        v = v.to(orig_dtype)
        if self.cfg.use_local_attn:
            if num_video_tokens is None or num_frames is None:
                raise ValueError("MagiHuman local attention requires video token/frame metadata.")
            out = self._local_window_attention(
                q,
                k,
                v,
                num_video_tokens=num_video_tokens,
                num_frames=num_frames,
            )
        else:
            out = self.attn(q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0)).squeeze(0)

        out = ModalityDispatcher.permute(out, permute_mapping)
        if g is not None:
            g = ModalityDispatcher.permute(g, permute_mapping)
            out = out * torch.sigmoid(g)
        out = out.reshape(-1, self.cfg.num_heads_q * self.cfg.head_dim).to(orig_dtype)
        return self.linear_proj(out, modality_dispatcher=modality_dispatcher)

`fastvideo.models.dits.magi_human.MagiAttention.configure_local_attention`

configure_local_attention(*, enabled: bool, frame_receptive_field: int = 11) -> None

Source

def configure_local_attention(
    self,
    *,
    enabled: bool,
    frame_receptive_field: int = 11,
) -> None:
    self.cfg.use_local_attn = enabled
    self.cfg.frame_receptive_field = frame_receptive_field

`fastvideo.models.dits.magi_human.MagiAttention.forward`

forward(hidden_states: torch.Tensor, rope: torch.Tensor, permute_mapping: torch.Tensor, inv_permute_mapping: torch.Tensor, modality_dispatcher: ModalityDispatcher, num_video_tokens: int | None = None, num_frames: int | None = None) -> torch.Tensor

Source

def forward(
    self,
    hidden_states: torch.Tensor,
    rope: torch.Tensor,
    permute_mapping: torch.Tensor,
    inv_permute_mapping: torch.Tensor,
    modality_dispatcher: ModalityDispatcher,
    num_video_tokens: int | None = None,
    num_frames: int | None = None,
) -> torch.Tensor:
    orig_dtype = self.linear_qkv.weight.dtype
    h = self.pre_norm(hidden_states, modality_dispatcher=modality_dispatcher).to(orig_dtype)
    qkv = self.linear_qkv(h, modality_dispatcher=modality_dispatcher).float()
    q, k, v, g = torch.split(
        qkv,
        [self.q_size, self.kv_size, self.kv_size, self.gating_size],
        dim=-1,
    )
    q = q.view(-1, self.cfg.num_heads_q, self.cfg.head_dim)
    k = k.view(-1, self.cfg.num_heads_kv, self.cfg.head_dim)
    v = v.view(-1, self.cfg.num_heads_kv, self.cfg.head_dim)
    g = g.view(-1, self.cfg.num_heads_q, 1) if self.gating_size else None

    q = self.q_norm(q, modality_dispatcher=modality_dispatcher)
    k = self.k_norm(k, modality_dispatcher=modality_dispatcher)

    # Un-permute before RoPE + attention so positional order reflects
    # the original (video, audio, text) concat — matches reference.
    q = ModalityDispatcher.inv_permute(q, inv_permute_mapping)
    k = ModalityDispatcher.inv_permute(k, inv_permute_mapping)
    v = ModalityDispatcher.inv_permute(v, inv_permute_mapping)
    if g is not None:
        g = ModalityDispatcher.inv_permute(g, inv_permute_mapping)

    # Element-wise Fourier embed packs sin/cos of 3 axes into a single
    # `rope` tensor. Match reference's split:
    #   sin_emb, cos_emb = rope.tensor_split(2, -1)
    # Reference passes (cos_emb, sin_emb) but splits sin first — replicated
    # exactly so weight parity holds. Partial RoPE: rope dim is
    # 6 * (head_dim // 8) = 96 < head_dim (128), so the trailing 32
    # head_dim positions stay unrotated, matching the reference.
    sin_emb, cos_emb = rope.tensor_split(2, -1)
    rot_dim = cos_emb.shape[-1] * 2
    q_rot = _apply_rotary_emb(q[..., :rot_dim], cos_emb, sin_emb, is_neox_style=True)
    k_rot = _apply_rotary_emb(k[..., :rot_dim], cos_emb, sin_emb, is_neox_style=True)
    if rot_dim < q.shape[-1]:
        q = torch.cat([q_rot, q[..., rot_dim:]], dim=-1)
        k = torch.cat([k_rot, k[..., rot_dim:]], dim=-1)
    else:
        q, k = q_rot, k_rot

    # Run SDPA via FastVideo's LocalAttention so the backend selection
    # (SDPA / FlashAttn / SLA / SageAttn) flows through the standard
    # configurable path. GQA is handled inside the SDPA backend via
    # `enable_gqa=True` when num_heads_q != num_heads_kv, so we no
    # longer need the manual `repeat_interleave` here.
    # Attention math runs at orig_dtype (bf16 in production and in the
    # parity test, since PackedExpertLinear's default is bf16, matching
    # upstream BaseLinear at dit_module.py:330). The gating multiply
    # promotes back to fp32 implicitly via PyTorch's type-promotion
    # rules: bf16_attn_out * sigmoid(fp32_g) -> fp32, mirroring upstream
    # dit_module.py:649.
    q = q.to(orig_dtype)
    k = k.to(orig_dtype)
    v = v.to(orig_dtype)
    if self.cfg.use_local_attn:
        if num_video_tokens is None or num_frames is None:
            raise ValueError("MagiHuman local attention requires video token/frame metadata.")
        out = self._local_window_attention(
            q,
            k,
            v,
            num_video_tokens=num_video_tokens,
            num_frames=num_frames,
        )
    else:
        out = self.attn(q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0)).squeeze(0)

    out = ModalityDispatcher.permute(out, permute_mapping)
    if g is not None:
        g = ModalityDispatcher.permute(g, permute_mapping)
        out = out * torch.sigmoid(g)
    out = out.reshape(-1, self.cfg.num_heads_q * self.cfg.head_dim).to(orig_dtype)
    return self.linear_proj(out, modality_dispatcher=modality_dispatcher)

`fastvideo.models.dits.magi_human.MLPSubConfig`

MLPSubConfig(hidden_size: int, intermediate_size: int, activation: str, num_modality: int, gated: bool)

Source

@dataclass
class MLPSubConfig:
    hidden_size: int
    intermediate_size: int
    activation: str  # "swiglu7" or "gelu7"
    num_modality: int
    gated: bool

`fastvideo.models.dits.magi_human.MagiMLP`

MagiMLP(cfg: MLPSubConfig)

Source

class MagiMLP(nn.Module):

    def __init__(self, cfg: MLPSubConfig):
        super().__init__()
        self.cfg = cfg
        self.pre_norm = MultiModalityRMSNorm(cfg.hidden_size, num_modality=cfg.num_modality)
        up_out = cfg.intermediate_size * 2 if cfg.gated else cfg.intermediate_size
        self.up_gate_proj = PackedExpertLinear(
            cfg.hidden_size,
            up_out,
            num_experts=cfg.num_modality,
            bias=False,
        )
        self.down_proj = PackedExpertLinear(
            cfg.intermediate_size,
            cfg.hidden_size,
            num_experts=cfg.num_modality,
            bias=False,
        )
        self._act = swiglu7 if cfg.activation == "swiglu7" else gelu7

    def forward(
        self,
        x: torch.Tensor,
        modality_dispatcher: ModalityDispatcher,
    ) -> torch.Tensor:
        orig_dtype = self.up_gate_proj.weight.dtype
        x = self.pre_norm(x, modality_dispatcher=modality_dispatcher).to(orig_dtype)
        x = self.up_gate_proj(x, modality_dispatcher=modality_dispatcher).float()
        x = self._act(x).to(orig_dtype)
        x = self.down_proj(x, modality_dispatcher=modality_dispatcher).float()
        return x

`fastvideo.models.dits.magi_human.MagiMLP.forward`

forward(x: torch.Tensor, modality_dispatcher: ModalityDispatcher) -> torch.Tensor

Source

def forward(
    self,
    x: torch.Tensor,
    modality_dispatcher: ModalityDispatcher,
) -> torch.Tensor:
    orig_dtype = self.up_gate_proj.weight.dtype
    x = self.pre_norm(x, modality_dispatcher=modality_dispatcher).to(orig_dtype)
    x = self.up_gate_proj(x, modality_dispatcher=modality_dispatcher).float()
    x = self._act(x).to(orig_dtype)
    x = self.down_proj(x, modality_dispatcher=modality_dispatcher).float()
    return x

`fastvideo.models.dits.magi_human.MagiTransformerLayer`

MagiTransformerLayer(arch: MagiHumanArchConfig, layer_idx: int)

Source

class MagiTransformerLayer(nn.Module):

    def __init__(self, arch: MagiHumanArchConfig, layer_idx: int):
        super().__init__()
        num_modality = 3 if layer_idx in arch.mm_layers else 1
        self.post_norm = layer_idx in arch.post_norm_layers
        self.layer_idx = layer_idx

        self.attention = MagiAttention(
            AttentionSubConfig(
                hidden_size=arch.hidden_size,
                num_heads_q=arch.num_attention_heads,
                num_heads_kv=arch.num_heads_kv,
                head_dim=arch.head_dim,
                num_modality=num_modality,
                enable_attn_gating=arch.enable_attn_gating,
                use_local_attn=layer_idx in arch.local_attn_layers,
            ))

        is_gelu7 = layer_idx in arch.gelu7_layers
        if is_gelu7:
            intermediate = arch.hidden_size * 4
            gated = False
            activation = "gelu7"
        else:
            intermediate = (arch.hidden_size * 4 * 2 // 3) // 4 * 4
            gated = True
            activation = "swiglu7"

        self.mlp = MagiMLP(
            MLPSubConfig(
                hidden_size=arch.hidden_size,
                intermediate_size=intermediate,
                activation=activation,
                num_modality=num_modality,
                gated=gated,
            ))

        if self.post_norm:
            self.attn_post_norm = MultiModalityRMSNorm(arch.hidden_size, num_modality=num_modality)
            self.mlp_post_norm = MultiModalityRMSNorm(arch.hidden_size, num_modality=num_modality)

    def forward(
        self,
        hidden_states: torch.Tensor,
        rope: torch.Tensor,
        permute_mapping: torch.Tensor,
        inv_permute_mapping: torch.Tensor,
        modality_dispatcher: ModalityDispatcher,
        num_video_tokens: int | None = None,
        num_frames: int | None = None,
    ) -> torch.Tensor:
        attn_out = self.attention(
            hidden_states,
            rope,
            permute_mapping,
            inv_permute_mapping,
            modality_dispatcher,
            num_video_tokens=num_video_tokens,
            num_frames=num_frames,
        )
        if self.post_norm:
            attn_out = self.attn_post_norm(attn_out, modality_dispatcher=modality_dispatcher)
        hidden_states = hidden_states + attn_out

        mlp_out = self.mlp(hidden_states, modality_dispatcher=modality_dispatcher)
        if self.post_norm:
            mlp_out = self.mlp_post_norm(mlp_out, modality_dispatcher=modality_dispatcher)
        return hidden_states + mlp_out

`fastvideo.models.dits.magi_human.MagiTransformerLayer.forward`

forward(hidden_states: torch.Tensor, rope: torch.Tensor, permute_mapping: torch.Tensor, inv_permute_mapping: torch.Tensor, modality_dispatcher: ModalityDispatcher, num_video_tokens: int | None = None, num_frames: int | None = None) -> torch.Tensor

Source

def forward(
    self,
    hidden_states: torch.Tensor,
    rope: torch.Tensor,
    permute_mapping: torch.Tensor,
    inv_permute_mapping: torch.Tensor,
    modality_dispatcher: ModalityDispatcher,
    num_video_tokens: int | None = None,
    num_frames: int | None = None,
) -> torch.Tensor:
    attn_out = self.attention(
        hidden_states,
        rope,
        permute_mapping,
        inv_permute_mapping,
        modality_dispatcher,
        num_video_tokens=num_video_tokens,
        num_frames=num_frames,
    )
    if self.post_norm:
        attn_out = self.attn_post_norm(attn_out, modality_dispatcher=modality_dispatcher)
    hidden_states = hidden_states + attn_out

    mlp_out = self.mlp(hidden_states, modality_dispatcher=modality_dispatcher)
    if self.post_norm:
        mlp_out = self.mlp_post_norm(mlp_out, modality_dispatcher=modality_dispatcher)
    return hidden_states + mlp_out

`fastvideo.models.dits.magi_human.MagiAdapter`

MagiAdapter(arch: MagiHumanArchConfig)

Source

class MagiAdapter(nn.Module):

    def __init__(self, arch: MagiHumanArchConfig):
        super().__init__()
        # Embedders stay in fp32 to match the reference dtype exactly.
        self.video_embedder = nn.Linear(
            arch.video_in_channels,
            arch.hidden_size,
            bias=True,
            dtype=torch.float32,
        )
        self.text_embedder = nn.Linear(
            arch.text_in_channels,
            arch.hidden_size,
            bias=True,
            dtype=torch.float32,
        )
        self.audio_embedder = nn.Linear(
            arch.audio_in_channels,
            arch.hidden_size,
            bias=True,
            dtype=torch.float32,
        )
        self.rope = ElementWiseFourierEmbed(arch.head_dim)
        # RoPE cache: coords_mapping is the same tensor object across timesteps
        # in the denoising loop, so data_ptr()+shape+dtype+device is a fast,
        # collision-free key that avoids recomputing the Fourier embed each step.
        self._cached_rope: Optional[torch.Tensor] = None
        self._cached_rope_key: Optional[tuple] = None

    def _rope_cache_key(self, t: torch.Tensor) -> tuple:
        return (t.data_ptr(), t.shape, t.dtype, t.device)

    def forward(
        self,
        x: torch.Tensor,
        coords_mapping: torch.Tensor,
        video_mask: torch.Tensor,
        audio_mask: torch.Tensor,
        text_mask: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        key = self._rope_cache_key(coords_mapping)
        if key != self._cached_rope_key:
            self._cached_rope = self.rope(coords_mapping)
            self._cached_rope_key = key
        rope = self._cached_rope
        # Embedder dtypes may differ from x's dtype when FastVideo's FSDP
        # loader casts all weights to `pipeline_config.precision` (bf16).
        # Match the weight dtype per modality.
        v_w = self.video_embedder.weight
        a_w = self.audio_embedder.weight
        t_w = self.text_embedder.weight
        out = torch.zeros(
            x.shape[0],
            self.video_embedder.out_features,
            device=x.device,
            dtype=v_w.dtype,
        )
        out[text_mask] = self.text_embedder(x[text_mask, :self.text_embedder.in_features].to(t_w.dtype)).to(out.dtype)
        out[audio_mask] = self.audio_embedder(x[audio_mask, :self.audio_embedder.in_features].to(a_w.dtype)).to(
            out.dtype)
        out[video_mask] = self.video_embedder(x[video_mask, :self.video_embedder.in_features].to(v_w.dtype)).to(
            out.dtype)
        return out, rope

`fastvideo.models.dits.magi_human.MagiAdapter.forward`

forward(x: torch.Tensor, coords_mapping: torch.Tensor, video_mask: torch.Tensor, audio_mask: torch.Tensor, text_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]

Source

def forward(
    self,
    x: torch.Tensor,
    coords_mapping: torch.Tensor,
    video_mask: torch.Tensor,
    audio_mask: torch.Tensor,
    text_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    key = self._rope_cache_key(coords_mapping)
    if key != self._cached_rope_key:
        self._cached_rope = self.rope(coords_mapping)
        self._cached_rope_key = key
    rope = self._cached_rope
    # Embedder dtypes may differ from x's dtype when FastVideo's FSDP
    # loader casts all weights to `pipeline_config.precision` (bf16).
    # Match the weight dtype per modality.
    v_w = self.video_embedder.weight
    a_w = self.audio_embedder.weight
    t_w = self.text_embedder.weight
    out = torch.zeros(
        x.shape[0],
        self.video_embedder.out_features,
        device=x.device,
        dtype=v_w.dtype,
    )
    out[text_mask] = self.text_embedder(x[text_mask, :self.text_embedder.in_features].to(t_w.dtype)).to(out.dtype)
    out[audio_mask] = self.audio_embedder(x[audio_mask, :self.audio_embedder.in_features].to(a_w.dtype)).to(
        out.dtype)
    out[video_mask] = self.video_embedder(x[video_mask, :self.video_embedder.in_features].to(v_w.dtype)).to(
        out.dtype)
    return out, rope

`fastvideo.models.dits.magi_human.MagiHumanDiT`

MagiHumanDiT(config: MagiHumanVideoConfig, hf_config: dict | None = None, **kwargs)

Top-level DiT for daVinci-MagiHuman (base).

Forward signature mirrors the reference DiTModel.forward: it takes a flat token stream, its per-token coords and modality mapping, and returns per-modality outputs packed into a max-channel-width tensor.

This scaffold is single-GPU only; the ulysses_scheduler().dispatch(...) sequence-parallel wrapping in the reference has no equivalent here yet.

Source

class MagiHumanDiT(BaseDiT):
    """Top-level DiT for daVinci-MagiHuman (base).

    Forward signature mirrors the reference `DiTModel.forward`: it takes a
    flat token stream, its per-token coords and modality mapping, and
    returns per-modality outputs packed into a max-channel-width tensor.

    This scaffold is single-GPU only; the `ulysses_scheduler().dispatch(...)`
    sequence-parallel wrapping in the reference has no equivalent here yet.
    """

    # BaseDiT requires these class attrs. Source them from the config so
    # they stay in sync with MagiHumanVideoConfig edits.
    _fsdp_shard_conditions = _CFG._fsdp_shard_conditions
    _compile_conditions = _CFG._compile_conditions
    _supported_attention_backends = _CFG._supported_attention_backends
    param_names_mapping = _CFG.param_names_mapping
    reverse_param_names_mapping = _CFG.reverse_param_names_mapping
    lora_param_names_mapping = _CFG.lora_param_names_mapping

    def __init__(self, config: MagiHumanVideoConfig, hf_config: dict | None = None, **kwargs):
        super().__init__(config=config, hf_config=hf_config or {})
        arch: MagiHumanArchConfig = getattr(config, "arch_config", config)
        self.arch = arch

        # BaseDiT contract instance vars.
        self.hidden_size = arch.hidden_size
        self.num_attention_heads = arch.num_attention_heads
        self.num_channels_latents = arch.num_channels_latents

        self.adapter = MagiAdapter(arch)
        self.block = _TransformerBlock(arch)
        self.final_norm_video = MultiModalityRMSNorm(arch.hidden_size)
        self.final_norm_audio = MultiModalityRMSNorm(arch.hidden_size)
        self.final_linear_video = nn.Linear(
            arch.hidden_size,
            arch.video_in_channels,
            bias=False,
            dtype=torch.float32,
        )
        self.final_linear_audio = nn.Linear(
            arch.hidden_size,
            arch.audio_in_channels,
            bias=False,
            dtype=torch.float32,
        )
        # Dispatcher + mask cache: modality_mapping is the same tensor object
        # across all timesteps in the denoising loop; data_ptr()+shape+dtype+device
        # is a fast, collision-free key that avoids rebuilding ModalityDispatcher
        # (which calls argsort + bincount) on every forward call.
        self._cached_dispatcher: Optional[ModalityDispatcher] = None
        self._cached_video_mask: Optional[torch.Tensor] = None
        self._cached_audio_mask: Optional[torch.Tensor] = None
        self._cached_text_mask: Optional[torch.Tensor] = None
        self._cached_modality_key: Optional[tuple] = None

    def configure_local_attention(
        self,
        local_attn_layers: tuple[int, ...] | list[int],
        frame_receptive_field: int = 11,
    ) -> None:
        layers = tuple(int(layer) for layer in local_attn_layers)
        self.arch.local_attn_layers = layers
        self.block.configure_local_attention(layers, frame_receptive_field)

    def _modality_cache_key(self, t: torch.Tensor) -> tuple:
        return (t.data_ptr(), t.shape, t.dtype, t.device)

    def forward(
        self,
        x: torch.Tensor,
        coords_mapping: torch.Tensor,
        modality_mapping: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            x:                [L, max(V_ch, A_ch, T_ch)]
            coords_mapping:   [L, 9]
            modality_mapping: [L]  (int in {VIDEO, AUDIO, TEXT})
        Returns:
            out: [L, max(V_ch, A_ch)] with video channels in video slots and
                 audio channels in audio slots; text slots are zero.
        """
        key = self._modality_cache_key(modality_mapping)
        if key != self._cached_modality_key:
            self._cached_dispatcher = ModalityDispatcher(modality_mapping, num_modalities=3)
            self._cached_video_mask = modality_mapping == Modality.VIDEO
            self._cached_audio_mask = modality_mapping == Modality.AUDIO
            self._cached_text_mask = modality_mapping == Modality.TEXT
            self._cached_modality_key = key
        dispatcher = self._cached_dispatcher
        video_mask = self._cached_video_mask
        audio_mask = self._cached_audio_mask
        text_mask = self._cached_text_mask
        num_video_tokens = int(video_mask.sum().item())
        if num_video_tokens:
            num_frames = int(coords_mapping[:num_video_tokens, 0].max().item()) + 1
        else:
            num_frames = 0

        x, rope = self.adapter(x, coords_mapping, video_mask, audio_mask, text_mask)
        # Keep the residual stream in adapter dtype (fp32) entering the block.
        # Upstream daVinci-MagiHuman dit_module.py:923 casts to params_dtype,
        # which is fp32 by default; each layer's pre_norm.to(bf16) handles
        # the bf16 internal-compute boundary, and linear_proj outputs bf16
        # which gets promoted back to fp32 by the residual addition. Casting
        # the residual to bf16 here degrades the cross-layer accumulator and
        # compounds visibly over 40 layers in pipeline parity.
        x = ModalityDispatcher.permute(x, dispatcher.permute_mapping)

        x = self.block(
            x,
            rope,
            permute_mapping=dispatcher.permute_mapping,
            inv_permute_mapping=dispatcher.inv_permute_mapping,
            modality_dispatcher=dispatcher,
            num_video_tokens=num_video_tokens,
            num_frames=num_frames,
        )
        x = ModalityDispatcher.inv_permute(x, dispatcher.inv_permute_mapping)

        x_video = x[video_mask].to(self.final_norm_video.weight.dtype)
        x_video = self.final_norm_video(x_video)
        x_video = self.final_linear_video(x_video)

        x_audio = x[audio_mask].to(self.final_norm_audio.weight.dtype)
        x_audio = self.final_norm_audio(x_audio)
        x_audio = self.final_linear_audio(x_audio)

        max_ch = max(self.arch.video_in_channels, self.arch.audio_in_channels)
        out = torch.zeros(x.shape[0], max_ch, device=x.device, dtype=x.dtype)
        out[video_mask, :self.arch.video_in_channels] = x_video.to(out.dtype)
        out[audio_mask, :self.arch.audio_in_channels] = x_audio.to(out.dtype)
        return out

`fastvideo.models.dits.magi_human.MagiHumanDiT.configure_local_attention`

configure_local_attention(local_attn_layers: tuple[int, ...] | list[int], frame_receptive_field: int = 11) -> None

Source

def configure_local_attention(
    self,
    local_attn_layers: tuple[int, ...] | list[int],
    frame_receptive_field: int = 11,
) -> None:
    layers = tuple(int(layer) for layer in local_attn_layers)
    self.arch.local_attn_layers = layers
    self.block.configure_local_attention(layers, frame_receptive_field)

`fastvideo.models.dits.magi_human.MagiHumanDiT.forward`

forward(x: torch.Tensor, coords_mapping: torch.Tensor, modality_mapping: torch.Tensor) -> torch.Tensor

Parameters

x: torch.Tensor: [L, max(V_ch, A_ch, T_ch)]
coords_mapping: torch.Tensor: [L, 9]
modality_mapping: torch.Tensor: [L] (int in {VIDEO, AUDIO, TEXT})

Returns: out: [L, max(V_ch, A_ch)] with video channels in video slots and audio channels in audio slots; text slots are zero.

Source

def forward(
    self,
    x: torch.Tensor,
    coords_mapping: torch.Tensor,
    modality_mapping: torch.Tensor,
) -> torch.Tensor:
    """
    Args:
        x:                [L, max(V_ch, A_ch, T_ch)]
        coords_mapping:   [L, 9]
        modality_mapping: [L]  (int in {VIDEO, AUDIO, TEXT})
    Returns:
        out: [L, max(V_ch, A_ch)] with video channels in video slots and
             audio channels in audio slots; text slots are zero.
    """
    key = self._modality_cache_key(modality_mapping)
    if key != self._cached_modality_key:
        self._cached_dispatcher = ModalityDispatcher(modality_mapping, num_modalities=3)
        self._cached_video_mask = modality_mapping == Modality.VIDEO
        self._cached_audio_mask = modality_mapping == Modality.AUDIO
        self._cached_text_mask = modality_mapping == Modality.TEXT
        self._cached_modality_key = key
    dispatcher = self._cached_dispatcher
    video_mask = self._cached_video_mask
    audio_mask = self._cached_audio_mask
    text_mask = self._cached_text_mask
    num_video_tokens = int(video_mask.sum().item())
    if num_video_tokens:
        num_frames = int(coords_mapping[:num_video_tokens, 0].max().item()) + 1
    else:
        num_frames = 0

    x, rope = self.adapter(x, coords_mapping, video_mask, audio_mask, text_mask)
    # Keep the residual stream in adapter dtype (fp32) entering the block.
    # Upstream daVinci-MagiHuman dit_module.py:923 casts to params_dtype,
    # which is fp32 by default; each layer's pre_norm.to(bf16) handles
    # the bf16 internal-compute boundary, and linear_proj outputs bf16
    # which gets promoted back to fp32 by the residual addition. Casting
    # the residual to bf16 here degrades the cross-layer accumulator and
    # compounds visibly over 40 layers in pipeline parity.
    x = ModalityDispatcher.permute(x, dispatcher.permute_mapping)

    x = self.block(
        x,
        rope,
        permute_mapping=dispatcher.permute_mapping,
        inv_permute_mapping=dispatcher.inv_permute_mapping,
        modality_dispatcher=dispatcher,
        num_video_tokens=num_video_tokens,
        num_frames=num_frames,
    )
    x = ModalityDispatcher.inv_permute(x, dispatcher.inv_permute_mapping)

    x_video = x[video_mask].to(self.final_norm_video.weight.dtype)
    x_video = self.final_norm_video(x_video)
    x_video = self.final_linear_video(x_video)

    x_audio = x[audio_mask].to(self.final_norm_audio.weight.dtype)
    x_audio = self.final_norm_audio(x_audio)
    x_audio = self.final_linear_audio(x_audio)

    max_ch = max(self.arch.video_in_channels, self.arch.audio_in_channels)
    out = torch.zeros(x.shape[0], max_ch, device=x.device, dtype=x.dtype)
    out[video_mask, :self.arch.video_in_channels] = x_video.to(out.dtype)
    out[audio_mask, :self.arch.audio_in_channels] = x_audio.to(out.dtype)
    return out