feat(wan2.2): memory_mode=relay — run A14B bf16 on 48GB (one expert resident at a time)

[lBroth]

Up-front honesty, in the spirit of #20: this work was done with heavy AI assistance (Claude — Fable 5) and I'm not a professional ML engineer. I've validated everything I could on-device (numbers below are all real measured runs I can share), and I'd genuinely welcome expert scrutiny on anything I may have missed.

Overview

The Wan2.2 A14B dual model switches experts once, at a deterministic timestep boundary — but both experts are loaded eagerly and stay resident for the whole denoise (~2× weights). That forces quantization on 48GB machines, and quantization currently costs real quality (#40: the quantized CFG path also produces black clips at 81f).

This PR adds memory_mode="relay": build the high-noise expert, run the high phase, free it at the boundary, build the low-noise expert. Peak ≈ one expert + activations.

Headline: Wan2.2 A14B bf16 goes from impossible-on-48GB to 36.8GB peak (832×480, 81f, Lightning) / 44.3GB with CFG — with no wall-clock penalty (bf16 relay was the fastest config tested: 27.9 min at 20 steps vs 28.7–29.1 for the quantized variants; the boundary reload costs ~3s from a warm page cache).

Modes

• auto (default) — relay for dual models, unchanged behavior for single
• relay — one expert resident, swapped at the boundary
• parallel — both resident (previous behavior)
• typos fail fast with ValueError (a silently-unknown mode would degenerate into both-resident with no warning)

Correctness contract — full seed reproducibility

Relay changes only when weights are resident, never the math. Verified on-device, all three directions:

• current main vs this branch (parallel): byte-identical mp4 for the same seed (MD5-equal files) — zero behavior change for existing users;
• parallel vs relay: bit-identical pre-VAE latents (MD5; opt-in pytest contract test included, set WAN22_A14B_DIR);
• therefore relay reproduces exactly what previous releases produced.

How: model construction consumes the global PRNG stream (keyless layer inits AND the nn.quantize stub constructors), and the noise is drawn after both experts are built. Relay defers the real builds, so it pre-consumes the same draws with two discarded lazy replicas of the loader path (never evaluated, near-zero cost). Known exception, documented in code: quantized + LoRA in relay mode stays deterministic per-mode but is not seed-identical to parallel (the LoRA dequant-merge constructs extra layers whose RNG use depends on the LoRA configs).

Second commit (standalone fix)

fix(wan2.2): release transformer weights before VAE decode — the cleanup deletes the model variables but not the loop locals (_call — the mx.compile wrapper — retains the model through its traced graph), so the last-used transformer stayed fully resident through the whole VAE decode. Applies to stock parallel mode too.

Measured (M5 Pro 48GB, macOS 26.5.1, mlx 0.31.2, seed 42, 832×480×81f)

config	peak	wall	notes
bf16 relay, Lightning 4-step, guide=1	36.8GB	7.3 min	previously impossible on 48GB
bf16 relay, 20 steps guide=1	36.8GB	27.9 min	fastest 20-step config tested
bf16 relay, 20 steps CFG 3.5,3.5	44.3GB	53 min	CFG works (quantized CFG at 81f is black — #40)
Q4 parallel (baseline)	36.8GB	28.7 min

Quality context (same seed/prompt across a precision ladder): rendered TEXT with quantization ≤8-bit came out as broken glyphs in every mix we tried; bf16 renders real letters.

Before/after clips (5s, seed 42):

• before — Q4 + Lightning, the 48GB status quo: glyph-soup signage
• after — bf16 relay + CFG: legible neon sign · bf16 relay + Lightning, 7.3 min (faster than Q4) · bf16 relay + CFG, face

Usage

python -m mlx_video.models.wan_2.generate \
  --model-dir Wan2.2-I2V-A14B-MLX-bf16 --image key.png \
  --prompt "..." --num-frames 81 --steps 20 --guide-scale 3.5,3.5 \
  --memory-mode relay --tiling aggressive

Tests

• tests/test_wan_relay.py: mode validation, CLI surface, opt-in bit-identity contract test (WAN22_A14B_DIR)
• local run: 42 passed, 1 skipped (note: tests/test_generate_dev.py fails collection on main too — it imports mlx_video.generate_dev, which no longer exists; happy to fix in a follow-up if wanted)
• regression matrix on this branch (M5 Pro 48GB): 5B single-model CFG, 5B T2V (no image), Q4+Lightning parallel (LoRA merge path), Q4+CFG relay, bf16 relay no_compile — all green, no black clips. (trim_first_frames + I2V crashes identically on current main — pre-existing, unrelated.)