mstar.model.wan22.components.dit

mstar.model.wan22.components.dit#

Native Wan2.2-TI2V-5B video DiT (dense 5B transformer).

Exact port of diffusers 0.39.0 WanTransformer3DModel, restricted to the TI2V-5B checkpoint’s configuration: no image-conditioning branch, and the per-token timestep path (the dit submodule always feeds a [batch, seq] timestep grid).

Numerics contract (bf16 weights with fp32 islands):

  • fp32 islands, matching the reference’s _keep_in_fp32_modules: the time embedder’s two Linears, every scale_shift_table, and each block’s norm2 affine. cast_serving_dtypes() pins them; the rest is bf16.

  • The fp32 -> bf16 boundaries are the reference’s type_as calls, and they must stay where they are: the timestep embedding and modulation run fp32 and cast back per block, RoPE is applied fp32, the self-attention and FFN residual adds upcast to fp32, and the cross-attention residual add stays bf16.

  • The RoPE cos/sin tables are DERIVED state, not checkpoint state. Non- persistent buffers hold garbage after to_empty, so they are built on CPU at init (device-independent values, copied lazily per device) rather than registered as buffers.

Two shared components are deliberately not reused. components.norm.RMSNorm dispatches to a FlashInfer kernel that rejects sm_120; the reference op is a plain torch.nn.RMSNorm across heads. components.attention.Attention is built around a KV-cache handle and a token-flat layout, and a stateless bidirectional video DiT with 3D RoPE would have to override all of it. The two-linear MLPs ARE the shared components.mlp.MLP.

Classes

Wan22DiT(config)

The dense 5B TI2V video DiT: Conv3d patchify -> 30 blocks -> adaLN output head + unpatchify.

Wan22DiTAttention(dim, num_heads, eps)

Multi-head attention with across-heads RMSNorm on q/k (the checkpoint's rms_norm_across_heads: one 3072-wide norm before the head split, not per-head) and SDPA.

Wan22DiTBlock(dim, ffn_dim, num_heads, eps)

One transformer block: adaLN-modulated self-attention, cross-attention to the text stream, adaLN-modulated FFN (reference WanTransformerBlock, per-token temb.ndim == 4 branch).

Wan22RoPE3D(attention_head_dim, max_seq_len)

3D-factorized rotary tables for the post-patchify (t, h, w) grid.

Wan22TimeTextEmbedding(dim, freq_dim, text_dim)

Condition embedder: per-token timestep embedding + text projection (reference WanTimeTextImageEmbedding, image branch absent on TI2V-5B).

WanFP32LayerNorm(normalized_shape[, eps, ...])

LayerNorm computed in fp32 regardless of input dtype (diffusers FP32LayerNorm): upcast input, normalize with fp32 weights, cast back to the input dtype.
class mstar.model.wan22.components.dit.Wan22DiT(config)[source]#

Bases: Module

The dense 5B TI2V video DiT: Conv3d patchify -> 30 blocks -> adaLN output head + unpatchify. Built on the meta device and materialized by weight_loader.build_wan22_dit (meta -> cast_serving_dtypes -> to_empty(device) -> checkpoint load).

Parameters:

config (Wan22Config)

cast_serving_dtypes()[source]#

bf16 everywhere except the fp32 islands (the checkpoint’s _keep_in_fp32_modules: time_embedder, scale_shift_tables, block norm2 affines). Called on the meta module BEFORE to_empty so storage is allocated directly in the serving dtypes.

Return type:

Wan22DiT

property dtype: dtype#

Bulk compute dtype (the non-island weights); callers cast inputs to this, mirroring diffusers ModelMixin.dtype.

forward(hidden_states, timestep, encoder_hidden_states)[source]#

hidden_states [B, C, F, H, W] (bf16 latents), timestep [B, post-patch seq] per-token grid (expand_timesteps — the only timestep form mstar’s dit submodule produces), text embeds [B, 512, text_dim] bf16. Returns [B, C_out, F, H, W] bf16.

Parameters:
Return type:

Tensor

class mstar.model.wan22.components.dit.Wan22DiTAttention(dim, num_heads, eps)[source]#

Bases: Module

Multi-head attention with across-heads RMSNorm on q/k (the checkpoint’s rms_norm_across_heads: one 3072-wide norm before the head split, not per-head) and SDPA. Self-attention applies 3D RoPE; cross-attention reads k/v from the 512-token text stream and skips RoPE. All four projections carry bias (reference WanAttention).

Parameters:
forward(hidden_states, encoder_hidden_states=None, rotary_emb=None)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:
Return type:

Tensor

class mstar.model.wan22.components.dit.Wan22DiTBlock(dim, ffn_dim, num_heads, eps)[source]#

Bases: Module

One transformer block: adaLN-modulated self-attention, cross-attention to the text stream, adaLN-modulated FFN (reference WanTransformerBlock, per-token temb.ndim == 4 branch).

The modulation runs fp32 (the scale-shift table is an fp32 island), the self-attention and FFN residual adds upcast to fp32, and the cross-attention residual stays bf16. All three follow the reference exactly.

Parameters:
forward(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:
Return type:

Tensor

class mstar.model.wan22.components.dit.Wan22RoPE3D(attention_head_dim, max_seq_len)[source]#

Bases: object

3D-factorized rotary tables for the post-patchify (t, h, w) grid.

The head dim splits into three axis bands — h = w = 2 * (head_dim // 6), t takes the remainder (44/42/42 for head_dim 128) — each with its own 1D table over max_seq_len positions, computed in float64 and stored fp32 concatenated to [max_seq_len, head_dim] (reference WanRotaryPosEmbed.__init__). The CPU tables are built at init (bit-identical regardless of the eventual device) and copied to each device on first use there.

Deliberately not an nn.Module: the tables are derived state that must never ride through to_empty/state_dict, and keeping them out of the module tree means no buffer for the loader or .to() to corrupt.

Parameters:

  • attention_head_dim (int)

  • max_seq_len (int)

tables(device)[source]#
Parameters:

device (device)

Return type:

tuple[Tensor, Tensor]

class mstar.model.wan22.components.dit.Wan22TimeTextEmbedding(dim, freq_dim, text_dim)[source]#

Bases: Module

Condition embedder: per-token timestep embedding + text projection (reference WanTimeTextImageEmbedding, image branch absent on TI2V-5B).

time_embedder is an fp32 island: the sinusoidal embedding and both its Linears run fp32; temb crosses to bf16 at type_as(text), so the downstream time_proj modulation projection is a bf16 matmul.

Parameters:
forward(timestep, encoder_hidden_states)[source]#

timestep [B, seq] (per-token, expand_timesteps); returns temb [B, seq, dim], timestep_proj [B, seq, 6, dim], projected text [B, 512, dim] — all in the text embeds’ dtype.

Parameters:
Return type:

tuple[Tensor, Tensor, Tensor]

class mstar.model.wan22.components.dit.WanFP32LayerNorm(normalized_shape, eps=1e-5, elementwise_affine=True, bias=True, device=None, dtype=None)[source]#

Bases: LayerNorm

LayerNorm computed in fp32 regardless of input dtype (diffusers FP32LayerNorm): upcast input, normalize with fp32 weights, cast back to the input dtype.

Parameters:
forward(inputs)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

inputs (Tensor)

Return type:

Tensor