mstar.model.wan22.submodules

mstar.model.wan22.submodules#

Wan2.2-TI2V-5B node submodules (all four declare no engine resources).

text_encoder -> Wan22TextEncoderSubmodule (UMT5-XXL, thin HF wrapper) vae_encoder -> Wan22VaeEncoderSubmodule (Wan2.2-VAE encode, I2V only) dit -> Wan22DitSubmodule (native 5B DiT + inline UniPC step) vae_decoder -> Wan22VaeDecoderSubmodule (Wan2.2-VAE decode)

All four set disable_torch_compile, so the engine never wraps their forwards: the three wrapped nodes run once per request with no graph to amortize, and the dit’s whole forward trips Inductor on the CPU-resident UniPC sigma. Instead the dit compiles its inner transformer region alone (config.compile_dit), leaving the solver eager.

Numerics are governed by the checkpoint dtypes, not by the engine. Wan22Model.get_autocast_dtype returns None so the engine neither autocasts nor blanket-casts the modules, and each forward enters _inference_ctx itself — reference parity must not depend on what context wraps the call, and the v1 engine supplies neither the autocast nor the no-grad wrapper.

Functions

normalize_decode_tiling_mode(raw)

Canonicalize a tiling-policy string to one of _DECODE_TILING_MODES, else "auto".

Classes

Wan22DitSubmodule(transformer, config)

Dense 5B video DiT with the UniPC step run inline.

Wan22TextEncoderSubmodule(text_encoder, config)

UMT5-XXL prompt encoder.

Wan22VaeDecoderSubmodule(vae, config)

Wan2.2-VAE latent -> pixel decoder (tiled).

Wan22VaeEncoderSubmodule(vae, config)

Wan2.2-VAE first-frame encoder (I2V requests only).
class mstar.model.wan22.submodules.Wan22DitSubmodule(transformer, config)[source]#

Bases: _SingleRequestMixin, _Fp32IslandMixin, NodeSubmodule

Dense 5B video DiT with the UniPC step run inline.

One loop iteration is one batch-2 transformer forward (the positive and negative prompt embeddings stacked on the batch axis) plus one UniPC predictor/corrector update from components/unipc.py.

Loop-carried edges, all float32 over the latent grid:

latents current sample x_t time_index [1] int64, the 0-based step index k unipc_model_outputs ring buffer of the last two converted outputs;

see UniPCState

unipc_last_sample the sample the previous predictor was given

Nothing else is carried. The sigmas, the order ramp and the timestep pair all follow from time_index and the per-request tables, so they are recomputed rather than shipped across an edge.

At iteration 0 the conductor sends the four edges empty and this submodule seeds them (noise from the request’s seed, on a CPU generator as diffusers does). I2V also consumes the persisted image_latent: frame 0 of the input is replaced by it and its per-token timestep zeroed, and the conditioning is written back into the output latents only on the final iteration, matching the reference’s post-loop injection.

Parameters:
check_stop(request_id, request_info, outputs)[source]#

Stop the denoise loop after exactly num_inference_steps iterations.

While iteration k (0-based) is being postprocessed the iteration count still reads k, and the stop registered here ends the loop at the end of that iteration. So for N steps it must fire at k == N - 1, i.e. when k + 1 >= N. The dit node runs with async scheduling, so the conductor may dispatch a speculative iteration N before this stop lands; prepare_inputs vetoes that overshoot by returning None, and >= keeps the stop firing when the deferred count reads N.

Parameters:
Return type:

set[str]

disable_torch_compile: bool = True#
forward(graph_walk, engine_inputs, latents, time_index, unipc_model_outputs, unipc_last_sample, text_embeds_pos, text_embeds_neg, image_latent=None, **kwargs)[source]#

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:
Return type:

dict[str, list[Tensor]]

prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)[source]#
Parameters:
Return type:

NodeInputs | None

class mstar.model.wan22.submodules.Wan22TextEncoderSubmodule(text_encoder, config)[source]#

Bases: _SingleRequestMixin, _Fp32IslandMixin, NodeSubmodule

UMT5-XXL prompt encoder.

Consumes text_inputs = [positive_ids, negative_ids]; emits text_embeds_pos / text_embeds_neg, each [1, text_max_seq_len, text_dim], zero-padded past the true sequence length. The two prompts run as two batch-1 forwards, as the reference’s separate encode_prompt calls do. With CFG off (guidance_scale <= 1.0) the negative prompt is not encoded at all: the dit never reads it, so text_embeds_neg is a 1-element placeholder that only keeps the persisted edge populated.

Parameters:
disable_torch_compile: bool = True#
forward(graph_walk, engine_inputs, positive_ids, negative_ids=None, **kwargs)[source]#

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:
Return type:

dict[str, list[Tensor]]

prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)[source]#
Parameters:
Return type:

NodeInputs

class mstar.model.wan22.submodules.Wan22VaeDecoderSubmodule(vae, config)[source]#

Bases: _SingleRequestMixin, _Fp32IslandMixin, NodeSubmodule

Wan2.2-VAE latent -> pixel decoder (tiled).

Consumes the loop’s final latents and emits video_output [1, 3, num_frames, H, W] as uint8. Quantizing here, at the worker boundary, keeps the edge to the data worker at one byte per pixel instead of a 4x larger float tensor; the output is 8-bit mp4 either way, and Wan22Model.postprocess does the muxing.

The node owns its own VAE instance because the decode dtype depends on the cuDNN version (see _decode_dtype) while the I2V encode stays bf16 — one shared instance would re-cast the whole VAE whenever the two disagree.

Decode calls tiled_decode / decode directly rather than enable_tiling, whose flag is instance state: the encoder’s numerics must not depend on which nodes happen to be colocated. Tiling bounds the workspace by tile count; untiled wants a conv3d workspace that scales with output spatial area (~20 GiB at dense fp32, enough to OOM a 32 GiB card when the DiT is co-resident) but decodes faster with no tile-seam error. So the path is VRAM-gated (Feature 1): decode untiled when free VRAM at decode time comfortably exceeds the estimated untiled peak for the requested size, tile otherwise. The estimate is the calibrated formula above; the decision reads live free memory via torch.cuda.mem_get_info; config.vae_decode_tiling / WAN22_VAE_DECODE_TILING force either path; the chosen path is logged per request. The tiled/untiled numeric gap is priced by the tiled-vs-untiled test.

Parameters:
disable_torch_compile: bool = True#
forward(graph_walk, engine_inputs, latents, **kwargs)[source]#

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:
Return type:

dict[str, list[Tensor]]

prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)[source]#
Parameters:
Return type:

NodeInputs

class mstar.model.wan22.submodules.Wan22VaeEncoderSubmodule(vae, config)[source]#

Bases: _SingleRequestMixin, _Fp32IslandMixin, NodeSubmodule

Wan2.2-VAE first-frame encoder (I2V requests only).

Consumes image_inputs (one [C, H, W] image, float in [0, 1] or uint8, already at the request’s height/width — the server does not resize, and a mismatch is an error). Emits image_latent [1, vae_z_dim, 1, H/16, W/16]: the mode sample of the encoded frame, normalized with the checkpoint’s per-channel statistics in float32. The dit node injects it at frame 0 and rebuilds the first-frame mask itself, since the mask follows from the latent dims and does not need persisting.

Parameters:
disable_torch_compile: bool = True#
forward(graph_walk, engine_inputs, image, **kwargs)[source]#

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:
Return type:

dict[str, list[Tensor]]

prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)[source]#
Parameters:
Return type:

NodeInputs

mstar.model.wan22.submodules.normalize_decode_tiling_mode(raw)[source]#

Canonicalize a tiling-policy string to one of _DECODE_TILING_MODES, else “auto”. Shared by the decoder gate and the request config echo so both report the same resolved policy.

Parameters:

raw (str)

Return type:

str