mstar.model.wan22.components

Wan2.2-TI2V-5B component modules.

Native components: the video DiT (dit.py — patchify, 3D-RoPE attention blocks, adaLN head, loaded through wan22.weight_loader) and the inline UniPC solver (unipc.py). The UMT5 text encoder and the Wan2.2-VAE stay thin diffusers wrappers, constructed in Wan22Model.get_submodule.

class mstar.model.wan22.components.UniPCState(model_outputs, last_sample)

Bases: object

Loop-carried UniPC solver state for one request.

model_outputs is the order-2 ring buffer of converted outputs (slot 1 is the previous step, slot 0 the one before). last_sample is the sample the previous predictor was given, consumed by this step’s corrector. Both are zero before they are written, and the order ramp guarantees they are not read then.

Parameters:

• model_outputs (Tensor)

• last_sample (Tensor)

last_sample: Tensor

model_outputs: Tensor

class mstar.model.wan22.components.Wan22DiT(config)

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

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)

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:

• hidden_states (Tensor)

• timestep (Tensor)

• encoder_hidden_states (Tensor)

Return type:

Tensor

class mstar.model.wan22.components.Wan22DiTAttention(dim, num_heads, eps)

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:

• dim (int)

• num_heads (int)

• eps (float)

forward(hidden_states, encoder_hidden_states=None, rotary_emb=None)

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:

• hidden_states (Tensor)

• encoder_hidden_states (Tensor | None)

• rotary_emb (tuple[Tensor, Tensor] | None)

Return type:

Tensor

class mstar.model.wan22.components.Wan22DiTBlock(dim, ffn_dim, num_heads, eps)

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:

• dim (int)

• ffn_dim (int)

• num_heads (int)

• eps (float)

forward(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb)

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:

• hidden_states (Tensor)

• encoder_hidden_states (Tensor)

• timestep_proj (Tensor)

• rotary_emb (tuple[Tensor, Tensor])

Return type:

Tensor

class mstar.model.wan22.components.Wan22RoPE3D(attention_head_dim, max_seq_len)

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)

Parameters:

device (device)

Return type:

tuple[Tensor, Tensor]

class mstar.model.wan22.components.Wan22TimeTextEmbedding(dim, freq_dim, text_dim)

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:

• dim (int)

• freq_dim (int)

• text_dim (int)

forward(timestep, encoder_hidden_states)

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:

• timestep (Tensor)

• encoder_hidden_states (Tensor)

Return type:

tuple[Tensor, Tensor, Tensor]

mstar.model.wan22.components.make_unipc_tables(num_inference_steps, flow_shift)

Per-request sigma and timestep tables.

Returns sigmas float32 [N+1] with the terminal zero appended, and timesteps int64 [N]. Sigmas are computed float64 and cast at the end, and the timesteps are truncated by the int64 cast — both as the reference does, and both load-bearing for bit-exactness.

Parameters:

• num_inference_steps (int)

• flow_shift (float)

Return type:

tuple[Tensor, Tensor]

mstar.model.wan22.components.unipc_convert_model_output(model_output, sample, sigmas, step_index)

Flow-prediction x0 conversion.

sigma_t stays a CPU 0-dim tensor on purpose, as the reference leaves it. A CPU scalar takes the CUDA kernel’s full-precision Scalar path; a device 0-dim tensor would instead be type-promoted to the bf16 operand’s dtype before the multiply. That is a bitwise difference, and it breaks lockstep with the reference. (It is also why this node cannot be compiled — see Wan22DitSubmodule.)

Parameters:

• model_output (Tensor)

• sample (Tensor)

• sigmas (Tensor)

• step_index (int)

Return type:

Tensor

mstar.model.wan22.components.unipc_corrector_step(state, this_model_output, this_sample, sigmas, step_index, order)

UniC correction of x_k using step k’s converted output (source multistep_uni_c_bh_update). Runs before the ring shift, so state.model_outputs[1] is step k-1’s output and [0] step k-2’s; state.last_sample is the sample step k-1’s predictor started from.

Parameters:

• state (UniPCState)

• this_model_output (Tensor)

• this_sample (Tensor)

• sigmas (Tensor)

• step_index (int)

• order (int)

Return type:

Tensor

mstar.model.wan22.components.unipc_effective_order(step_index, num_inference_steps)

Predictor order at step k.

Capped both by the steps remaining (so the last step is order 1) and by the warmup ramp. Step k’s corrector reuses step k-1’s value.

Parameters:

• step_index (int)

• num_inference_steps (int)

Return type:

int

mstar.model.wan22.components.unipc_predictor_step(state, sample, sigmas, step_index, order)

UniP update x_k -> x_{k+1}.

state.model_outputs[1] must already hold step k’s converted output: the ring buffer is shifted before the predictor runs.

Parameters:

• state (UniPCState)

• sample (Tensor)

• sigmas (Tensor)

• step_index (int)

• order (int)

Return type:

Tensor

Modules

dit	Native Wan2.2-TI2V-5B video DiT (dense 5B transformer).
unipc	Inline UniPC (bh2, order 2, flow prediction) for the Wan2.2 denoise loop.