Add Wan2.2: text/image-to-video (issue #126) - #165
Open
Gaurav-Shah05 wants to merge 8 commits into
Open
Conversation
…ar-project#126) Adds the bottom layer of Wan2.2-TI2V-5B support: the config dataclass, the native video DiT (30 layers, 3072 hidden, batched-CFG forward), the inline UniPC solver port (bh2, order 2, flow prediction), and the mstar weight loader that maps the checkpoint onto them. The DiT is native rather than a diffusers wrapper so it can be built on meta device and materialized straight onto the GPU through mstar's loader (no CPU staging copy of 5B weights), and so the serving dtypes are ours to pin: bf16 weights with fp32 islands (time embedder, scale-shift tables, norm affines) that the reference keeps in fp32 and that the numerics depend on. UniPC is ported inline instead of held in a scheduler object because the solver state has to travel with the request (see the model commit). Architecture values are facts of the Wan-AI/Wan2.2-TI2V-5B-Diffusers checkpoint and are hardcoded so that constructing the model never touches the network. The loader checks them against the checkpoint at load time and refuses to serve on any drift, and it requires a bijection between checkpoint keys and native parameters: an unexpected or unloaded key is an error, not a warning.
… (issue mstar-project#126) The four NodeSubmodules the engine executes. The dit submodule owns one UniPC step: it seeds the latents from the request's seed at iteration 0, runs the batched-CFG forward, advances the solver, and carries (latents, step index, UniPC history, last_sample) on the loop's own edges. UMT5-XXL and the Wan2.2-VAE are thin HF/diffusers wrappers — there is nothing to gain from reimplementing them and a great deal of numerics to lose. Two dtype seams worth knowing about: * the VAE decode dtype is chosen from the live cuDNN version, because cuDNN ships fast bf16 conv3d for this decode only from 9.16 and older builds are much slower in bf16 than in fp32; * the decoder holds its own VAE instance in the checkpoint dtype while the I2V encoder holds a bf16 one, so neither has to re-cast the full VAE when the two dtypes disagree. All four nodes run eager. The transformer alone compiles cleanly, but the engine compiles the whole node, and Inductor then fuses the UniPC step's sigma arithmetic into a Triton kernel and passes the deliberately CPU-resident sigma as a device pointer. That sigma is what keeps the solver bit-exact against the reference, so numerics win and the node stays eager. The VAE decode is tiled unconditionally: the untiled conv3d workspace OOMs a 32 GiB card at the dense tier. Tiling is not free (it decodes ~2.5x the pixels it needs); choosing it on free VRAM instead of always is a possible follow-up.
…sue mstar-project#126) The Model ABC implementation and the registry/config entries that make "wan22" servable. Four stateless nodes (text_encoder, vae_encoder, dit, vae_decoder) and four graph walks; T2V runs encode_text -> video_gen, I2V inserts encode_image and swaps in a walk whose dit additionally consumes the persisted first-frame latent. Output is muxed to mp4 by PyAV, at the request's fps (a playback rate, not a generation knob). The denoise loop is dynamic: Loop.max_iters is the config ceiling and the per-request num_inference_steps stops it via check_stop, so steps is a request knob rather than a server one. Why UniPC state rides the loop edges instead of living in a scheduler object, which is the one deliberate structural difference from cosmos3: the solver's history buffer and corrector state are per-request, so putting them on the request's own edges keeps requests independent, keeps the node stateless, and lets the loop resume on whichever rank picks it up. It also means the dual-DiT A14B variants can reuse the loop as-is later. Request validation lives in process_prompt, on the DATA WORKER, because that is the seam where a ValueError becomes a 400 for the one offending request. The same raise at the conductor, or on the result path, is swallowed and the client hangs until timeout, so those checks are backstops only. Validated here: output modality, prompt presence, a declared-but-missing image, and all four geometry knobs — * height/width must be multiples of 32 (VAE downsample 16 x DiT patch 2). 720x1280 is impossible for this model (720/32 = 22.5); the 720p-class tier is 704x1280. Unaligned sizes previously reached the DiT as a mid-forward shape mismatch, which faults the WORKER rather than failing the request — observed as a non-terminating error loop that hung the client, blocked the queue and grew the log to gigabytes. * num_frames must be 4k+1 (the VAE compresses time by 4 around an anchor frame). Non-positive values made the latent time extent <= 0 and crashed the worker on a negative tensor dim; unaligned ones did not fail at all, which is worse — 32 frames requested was silently floored to 29 returned. * fps must be positive. postprocess also rejects it, but postprocess runs after the whole video is generated and its raise is swallowed, so the client burned a full generation and then hung. Each rejection names the rule and the nearest valid values.
Registers wan22 on /v1/videos/generations. The adapter splits the OpenAI-shaped
"size" ("WxH") into the model's width/height kwargs, passes seed/num_frames/fps
as first-class fields and the rest (guidance_scale, num_inference_steps,
negative_prompt) through extra_body, and turns an "image" into an
image-to-video request.
"video" conditioning is rejected outright rather than silently ignored:
Wan2.2 has no video-conditioned mode, and a request that asks for one is
asking for something it will not get.
The native /generate route needs no adapter and serves the same model.
…order (issue mstar-project#126) Three suites and the recorder that makes the third one reproducible. * test_wan22_model.py — dummy mode (no weights, no GPU, no network): graph walks, loop-carried state, walk transitions, check_stop boundaries, mp4 postprocess, and the full request-rejection matrix (bad size, bad num_frames, bad fps, wrong output modality, promptless, image declared but absent). * test_wan22_native_dit_equivalence.py — the native DiT against the diffusers module it replaces: 825/825 parameters bitwise identical under a remap asserted to be a true bijection (injective AND onto — equal counts alone would not catch two checkpoint keys colliding onto one native key), and bitwise-identical per-layer outputs. This is what licenses reimplementing the DiT at all. * test_wan22_reference_equivalence.py — the whole serving path against a recorded diffusers oracle: text embeddings, the 50-step trajectory, the inline UniPC port in lockstep with the reference scheduler, CFG batched vs sequential, final latents, decode PSNR, and tiled vs untiled decode. record_oracle.py records that oracle, and it ships because without it the suite cannot be run on a new machine. Its docstring carries the rule the thresholds depend on: the oracle must be recorded under the serving process's numerics flag (set_float32_matmul_precision("high")) and on the same torch build as the server under test. Record it without the flag and the reference itself moves — by 1.3e-3 at step 0, growing to 4.2 by step 49. An oracle is therefore only valid for the machine it was recorded on, and it is not distributable, so there is no baked-in path: point WAN22_ORACLE_DIR at what you recorded. Unset, the suite skips with instructions rather than silently comparing against someone else's numerics. Thresholds are derived from a matched run and are NOT to be loosened to make a run pass; a regression there is a finding.
…ect#126) A phase-instrumented t2v/i2v benchmark for wan22, with vllm-omni and SGLang-Diffusion as baselines and stock diffusers as the compute floor. The per-phase breakdown reuses mstar's own request profiler (--log-stats-file) rather than adding instrumentation: wan22's graph has one node per phase and the dit node runs once per UniPC step, so dit.total/exec_count IS the mean per-step wall. Per-step p50/p95 needs a distribution the profiler discards, so it comes from an opt-in step log, enabled by an env var, off by default and numerically inert when off. The measurement rules encoded here were each paid for with a wrong number: * Cross-system LATENCY requires one system at a time, alone on one GPU. Serving each system on its own GPU concurrently is valid for throughput and invalid for latency: cards settle at different clocks under power cap (a 6.7% inter-GPU spread was measured), which is larger than the effect being measured. reproduce.sh gains a preflight that tears down the process TREE and then VERIFIES the card is empty — orphaned workers survive a port-based kill, keep their CUDA contexts, and are then counted as the next run's peak. * Peak VRAM is labelled with WHICH metric it is. nvidia-smi memory.used is whole-GPU and is not a working set; torch.cuda.max_memory_allocated is process-scoped but unreadable from a client, so only the in-process diffusers baseline reports it. The two differ by GiB on one identical run. * torch/cuDNN are probed from the SERVER's interpreter (--server-python), not the client's. Baselines run in their own venvs, so a client-stamped torch on a baseline row records a stack that never ran the work. * The baselines' poll interval is 20ms, not 250ms: it biases their e2e upward by half its value, and 125ms is the same order as the differences measured. Every row carries the engine's timing_boundary so two e2e numbers are never differenced without it. * Per-phase columns at concurrency>1 are batch-shared, not per-request (three distinct requests report bit-identical node timings), and are withdrawn. * An i2v cell on a baseline is refused rather than run: those clients send no conditioning frame, so the cell would be a t2v run labelled i2v. reproduce.sh's serve commands encode two baseline gotchas the obvious recipe gets wrong: * vllm-omni: `vllm serve --omni` serves no video routes and /v1/videos/sync 405s; the real path is a different entrypoint and an async job API. It must run vllm 0.18.0 — 0.20.0+ are cu13-only wheels — which pins torch 2.10 on the baseline side (disclosed, since it differs from mstar's). * SGLang: most of its documented diffusion flags are absent from the shipped build and kill the server on unrecognized arguments; its client is a dispatch stub because SGLang-Diffusion is blocked upstream. The measured prices for the deferred work (compile, CUDA graphs, VRAM-conditional tiling) live in the PR description, cited to the recorded artifacts. The layout matches benchmark/cosmos3/: the measurement engine, its profile parser, and reproduce.sh, whose header carries the run order and the rules above.
Adds wan22 to the model tables in docs/models.rst and README.md, and a docs/models.rst section covering the supported variant and checkpoint, both routes with working curl examples, the generation knobs, the size and frame-count rules, and the inline-UniPC design note. Also registers wan22 in the CLI's DEFAULT_CONFIGS, without which "mstar serve wan22" fails despite configs/wan22.yaml existing.
Gaurav-Shah05
marked this pull request as ready for review
July 15, 2026 02:47
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Wan2.2-TI2V-5B: text/image-to-video (issue #126)
Adds
wan22- Wan2.2-TI2V-5B (Wan-AI/Wan2.2-TI2V-5B-Diffusers) - as a served model:text-to-video and image-to-video on
/v1/videos/generationsand/generate.What's in it:
built on the meta device (shapes only, no memory), then checkpoint weights load straight
into GPU memory with no CPU staging - bf16, with fp32 kept exactly where the reference
numerics need it.
and risks numerics.
Wan2.2 also ships larger 14B variants ("A14B") that run two DiTs per video, one for the
early noisy steps and one for the late refinement steps. This PR serves only the 5B; the
config rejects the 14B variants rather than half-supporting them.
The graph
Four stateless nodes (no KV cache):
text_encoder,vae_encoder,dit,vae_decoder.Three walks:
encode_textruns UMT5 once and persists the embeddings;encode_image(I2Vonly) encodes the conditioning frame;
video_genloops the dit node for the requestednumber of steps, then decodes. The denoise loop is a named loop, so each request sets its
own step count through
check_stop, and everything the solver needs between steps travelson the loop's own edges.
That last part is the one deliberate departure from cosmos3, which holds its scheduler as a
per-request Python object. Solver state is part of the request, so wan22 carries it on the
request's graph edges next to the latents. The dit node stays stateless, requests cannot
leak into each other, and the state follows the request to whichever rank runs the next
step - which is what will let the 14B variants swap their two DiTs mid-loop without
touching the loop. The port matches the reference scheduler exactly, verified
step-for-step in the tests.
What this PR optimizes, and what it doesn't
The inner DiT region (patchify -> blocks -> head) is torch.compiled while the
UniPC solver stays eager, since its CPU-resident sigma is what the solver's bit-exactness
depends on; the two CFG branches run as one batched forward instead of two calls; the
meta-device weight loading described above; the VAE decodes in fp32 or bf16 depending on
which the installed cuDNN runs faster; VAE decode is VRAM-gated - untiled when free memory
at decode time comfortably clears the estimated untiled peak, tiled otherwise, and the
tiled path is what still makes a 32 GB card work; frames cross the worker boundary as
uint8; and one resolved-settings line is logged per request, so an operator reads back the
knobs the server resolved instead of asserting them (the vllm-style arg dump).
Not done here, priced below: CUDA-graph capture, batching across requests, CFG parallelism
(the two branches share one GPU today; running them on separate ranks is expressible in the
graph but no such config ships), and multi-rank placement in general - the shipped config
is single-GPU.
Correctness
mstar's full path - tokenize, UMT5, per-token timesteps, native DiT, inline UniPC - matches
stock diffusers bit for bit at every one of the 50 denoise steps, under matched numerics.
The native DiT is verified layer by layer against the diffusers module (identical outputs
at every one of the 30 blocks), and the weight loader maps each of the 825 checkpoint
weights to exactly one native weight, with the tests failing if any is missing, duplicated,
or left over. Decode PSNR against the reference video is 35.0 dB on H100.
One place our computation differs from the reference: diffusers computes the guided and
unguided predictions in two model calls; we compute both in one batched call. On H100 the
results are identical. On the RTX 5090 dev GPU they differ by that card's rounding noise,
which the test bounds allow for.
The tests compare against recordings made from stock diffusers. Recordings depend on the
GPU and torch build, so they cannot be shipped;
test/wan22/record_oracle.pyregeneratesthem in a few minutes, and without one the suite skips and prints instructions.
Performance
1x H100, every system alone on the same GPU, run one after another, 10 runs per cell.
Dense tier, 480x832 x 33 frames, concurrency 1:
mstar is the fastest of the three at every tier measured: 1.4x vllm-omni's best previous
numbers and 1.76-1.83x stock diffusers at the dense tier. The two optimizations account for
the gain additively (both off: 8.89 s). Throughput plateaus at ~0.22 req/s from concurrency
2 - one request saturates the GPU.
Pending
Issue #126 checklist
Known limitations
construction (the conditioning image is inserted where and when the reference inserts it)
and by inspecting outputs, but has no recorded reference trajectory yet.
cleanup dependency; ASCII prompts are identical).
Request validation
Height and width must be multiples of 32 (the VAE shrinks each side 16x, then the DiT
patches 2x2), num_frames must be 4k+1 (the VAE compresses time 4x around an anchor frame),
and fps must be positive. All three fail fast as a 400 naming the rule and the nearest
valid values.
Reproduce