Add opt-in low-VRAM mode for Wan generation

[JPPhoto]

Summary

Adds an opt-in Wan memory optimization mode:

wan_memory_optimization: true

The option defaults to false.

When enabled, Invoke:

• Limits resident Wan transformer weights to about 2 GiB and streams remaining layers from RAM.
• Enforces the residency limit after model-cache loading, including already-resident experts and explicit VRAM cache configurations.
• Chunks pointwise transformer operations while preserving global self-attention and cross-attention.
• Compacts TI2V per-token timestep conditioning to its unique timestep states.
• Streams untiled causal VAE decode chunks directly into the MP4 writer instead of retaining the complete RGB video in GPU or system memory.
• Uses a streaming-aware VAE working-memory estimate.

This applies to Wan image and video denoise, including dual-expert A14B models.

The main tradeoff is speed: aggressive weight streaming can make generation substantially slower and requires enough system RAM for offloaded weights. Spatially tiled VAE decode retains its existing path.

Related Issues / Discussions

Related design reference: #9460.

QA Instructions

1. Add this setting to invokeai.yaml:

   wan_memory_optimization: true

2. Restart Invoke.

3. Run Wan image and video generation with representative configurations:

   • TI2V-5B
   • A14B T2V or I2V
   • A14B dual-expert generation
   • CFG enabled and disabled

4. Confirm the log contains:

   Wan memory optimization: limiting resident transformer weights to about 2 GiB
   

5. Confirm the model-cache log reports substantially reduced resident transformer weights and peak VRAM is lower.

6. Confirm generated videos have the expected dimensions, frame count, duration, and playback.

7. Repeat with wan_memory_optimization: false and confirm existing behavior is unchanged.

Merge Plan

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

[lstein]

I ran an adversarial review of this at a25ef93 (fresh-context review agents attacking the diff, every finding re-verified against the code by hand). The numerical core held up impressively: I bit-compared the chunked/compacted transformer forward against diffusers 0.39.0 on real WanTransformer3DModel instances (fp32 maxdiff ≤ 3.6e-7, bf16 within one ulp, including per-token TI2V timesteps with mixed unique values, I2V image embeds, GGUF-style bf16 scale_shift_table, and every chunk-size edge), and the streaming VAE decode matches vae.decode exactly with correct frame accounting. The test suite here is genuinely strong. I did find two things I'd like fixed before merge, plus two smaller ones.

Blocking

1. Anima's tiling leak deterministically breaks the streaming decode

anima_latents_to_image.py enables tiling on the shared cached Wan VAE and leaves it enabled — its stated convention (see the comment at line ~140, "always set the tiling state explicitly rather than leaving it as-is") is that each consumer sets the state itself, and its OOM-retry path (lines ~171-177) also exits with tiling on. The new iter_wan_vae_decode_chunks honors the leaked flag but raises instead of resetting:

if vae.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height):
    raise ValueError("Streaming Wan VAE decode does not support spatial tiling.")

Trigger (deterministic): with wan_memory_optimization: true, render one high-res Anima image (Anima and A14B share the same Wan 2.1 VAE record), then render an A14B video at 832×480. The leaked tile_sample_min_height=512 gives tile_latent_min_width = 64; latent width 104 > 64, so the whole video generation dies with this ValueError at the final node, after the full denoise. Reproduced live on a tiny AutoencoderKLWan after simulating the Anima leak.

On main the same leak merely caused a silent tiled vae.decode; the hard failure is new to this PR. One-line fix: call vae.disable_tiling() in the non-tiled branch of wan_latents_to_video.py before decoding (matching Anima's explicit-set convention).

2. partial_unload_from_vram runs outside MODEL_LOAD_LOCK — meta-tensor stranding race

The new trim block in _ExpertSwapper.activate (wan_denoise.py:279-288) runs after model_on_device.__enter__ returns, and load_base.py:97-98 holds the MODEL_LOAD_LOCK read lock only around the inner cache.lock() call — not the context body. partial_unload_from_vram ends in load_state_dict(assign=True) → register_parameter, which is exactly the operation the lock's docstring (model_cache.py:102-126) says must never overlap a concurrent model construction: construction installs accelerate's process-global register_parameter → meta monkey-patch.

Concrete interleaving (two session workers / multi-GPU — the topology the lock exists for):

1. Worker A (flag on) enters model_on_device(...) for a Wan expert; the read lock is released once lock() returns.
2. Worker B starts a cold load of any model (T5, VAE, the other expert) and takes the write lock — granted, since A holds nothing — then enters init_empty_weights(), patching register_parameter process-wide for the duration of construction.
3. Worker A's trim rebinds ~GBs of Wan parameters via load_state_dict(assign=True); each rebind is hijacked onto meta. The transformer loses its real weights.
4. A's next weight access (or next partial_load_to_vram) raises "Cannot copy out of meta tensor; no data!", and the cached model is corrupt for its remaining cache lifetime.

Every existing tensor-mover takes the read lock (model_on_device, repair_required_tensors_on_device, LayerPatcher.apply_smart_model_patches); this is the only path that doesn't. Fix is one line: wrap the trim in MODEL_LOAD_LOCK.read_lock() (importable from model_cache, as layer_patcher.py already does). Read locks are shared, so there's no new contention.

(Note: the pre-existing outgoing_cached_model.full_unload_from_vram() at wan_denoise.py:249 has the same defect, but that's on main and out of scope here — happy to see it fixed in the same pass, though.)

Non-blocking, but worth addressing

3. Raw partial_unload_from_vram bypasses the cache's delete-on-error contract

Every cache-internal unload goes through _move_model_to_ram, which deletes the cache entry on any exception because a half-moved model is in an undefined state (model_cache.py:1078-1081) — and partial_unload_from_vram only updates _cur_vram_bytes after the whole conversion succeeds. The invocation calls the method raw, so an exception partway through (e.g. host-RAM OOM with keep_ram_copy_of_weights: false, which allocates a CPU copy per module) fails the invocation but leaves the half-moved transformer cached with over-reported cur_vram_bytes — and the next session gets a cache hit on it. Either invalidate/delete the entry on exception, or route the trim through the cache so the existing contract applies.

4. max_cache_vram_gb (or legacy vram:) silently defeats the reservation

_get_vram_available returns early and ignores working_mem_bytes entirely when _max_vram_cache_size_gb is set (model_cache.py:1088-1090). Under that config, every expert swap fully loads the transformer up to the cache cap and then immediately trims it back to 2 GiB: GB-scale H2D churn per swap, and the load itself transiently occupies the very VRAM peak the flag promises to avoid — on exactly the low-VRAM-tuned installs likely to carry that setting. The "limiting resident transformer weights to about 2 GiB" INFO log still prints, which will make user reports confusing. Worth either honoring per-lock working_mem_bytes under the override or documenting the exclusion in low-vram-mode.mdx.

Minor notes (your call)

• The unpatch restores forward by assignment, permanently leaving a 'forward' entry in the instance __dict__ of the cached transformer and every block (the attribute wasn't there before patching). Verified harmless today — output bit-identical afterward — but any future class-level forward patch would be silently shadowed on models that ever ran with the flag. delattr on restore is cleaner.
• The streaming path never emits the "Encoding MP4 (...)" progress signal/log line — cosmetic regression vs the buffered path.
• Latent hazard note: if anything re-enables grad between transformer entry and a block call (e.g. a future torch.enable_grad() forward hook), the original block forward receives a _CompactTimestepConditioning and crashes on temb.ndim. Nothing in-tree can trigger it; just flagging for the file's future maintainers.

For completeness, attacks that failed: the torch.unique timestep compaction and per-token modulation math (empirically exact), CFG double-call and per-step patch churn, expert-swap-while-patched (impossible — the context never spans a swap), GGUF/quantized paths, LoRA over partially-offloaded weights (sidecar routing protects the canonical CPU weights), generator abandonment on cancel (the finally-driven clear_cache runs promptly at unwind — verified empirically), tmp-file cleanup on every error path, frame accounting vs t_pixel, and the estimate/tiling branch ordering (no path uses the small streaming estimate for a full non-streaming decode).