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

[JPPhoto]

@lstein Thanks for the review.

• Blocking 1 resolved: invokeai/app/invocations/wan_latents_to_video.py: WanLatentsToVideoInvocation.invoke clears shared VAE tiling before non-tiled decode. Regression test simulates Anima residue; passes.

• Blocking 2 resolved: invokeai/app/invocations/wan_denoise.py: _ExpertSwapper.get routes unloads through LoadedModelWithoutConfig.unload_from_vram(), guarded by MODEL_LOAD_LOCK.

• Non-blocking 3 resolved: ModelCache.unload_model_from_vram() preserves cache failure cleanup; failed unload removes the cache entry. Regression test passes.

• Non-blocking 4 resolved: explicit max_cache_vram_gb now reserves requested working memory. Documentation updated at docs/src/content/docs/configuration/low-vram-mode.mdx.

I also addressed your minor notes:

• Instance-level forward leakage resolved. Cleanup now uses del block.forward / del transformer.forward; regression test passes.

• Streaming MP4 progress/log signal added in wan_latents_to_video.py.

• Gradient re-enable hazard was not changed. It remains hypothetical; optimization already falls back when gradients are enabled at entry. No in-tree trigger exists.

[Pfannkuchensack]

Speed is an explicitly accepted tradeoff in this PR, so I have left it out of the findings.
The measurements below were taken on a simulated 12 GB card (allocator capped at 11.0 GiB
on a 4090) because that is the stated target hardware.

Findings

High — invokeai/backend/util/vae_working_memory.py:128-170: the transformer now fits on a 12 GB card, but the VAE decode still OOMs and the tiling fallback never fires.

Simulated 12 GB card, TI2V-5B, 704x1280 / 81 frames:

denoise step   flag OFF -> CUDA OOM
               flag ON  -> peak 1.49 GiB   OK      <- the PR does its job here

VAE decode     estimate(streaming=False) = 5.68 GiB   tiling? (est > 0.9*12GiB = 10.80 GiB) False
               estimate(streaming=True ) = 4.89 GiB   tiling? False
               full decode      -> CUDA OOM
               streaming decode -> CUDA OOM

The flag gets a 12 GB user past the transformer only to fail at wan_l2v. The safety net at invokeai/app/invocations/wan_latents_to_video.py:140 cannot engage because estimate_vae_working_memory_wan under-predicts the real peak by roughly 2x. Backed out from measured peaks (streaming decode, bf16):

VAE                                    resolution        measured   estimate   ratio   implied const
TI2V-5B (Wan 2.2, z=48, 16x, patch=2)  704x1280 x33      11.27 GiB  4.89 GiB   2.31x        6703
A14B    (Wan 2.1, z=16,  8x, patch=1)  704x1280 x33       9.74 GiB  4.89 GiB   1.99x        5792

scaling_constant = 2900 at invokeai/backend/util/vae_working_memory.py:151 is documented as "calibrated empirically on a Wan 2.1 fp16 decode", but it under-predicts on the Wan 2.1 VAE too, not just on TI2V. The under-prediction is pre-existing; what this PR adds is a second consumer of the same number, and it moves it in the wrong direction for the target hardware: streaming=True drops the reservation a further 0.79 GiB (5.68 -> 4.89 GiB) for a real saving of only 0.38 GiB (measured 10562 -> 10186 MB at 81 frames on an unconstrained card, -3.6 percent). The peak is dominated by the per-frame causal-conv working set, which streaming does not touch; the RGB clip it removes is a small share on GPU.

The genuine win from streaming is system RAM: a 418 MB resident clip becomes a 21 MB rolling chunk. That is worth keeping.

The new decoder itself is correct: streaming output is bit-identical to vae.decode (max_abs_diff 0.0000e+00) on the real patch_size=2 VAE.

To expose this issue, add a test that asserts estimate_vae_working_memory_wan(operation="decode", ...) is not below a measured reference peak for a known Wan VAE and resolution, so the tiling fallback threshold is exercised against reality rather than against the estimator.

---

High — invokeai/app/invocations/wan_denoise.py:286-295: the residency trim moves the whole transformer to CPU when partial loading is off, crashing the first denoise step.

Measured with the real 5B TI2V transformer, a real ModelCache, and the unmodified 2 GiB cap:

enable_partial_loading=True   CachedModelWithPartialLoad  resident after trim 0.42 GiB  FORWARD OK
enable_partial_loading=False  CachedModelOnlyFullLoad     resident after trim 0.00 GiB  FORWARD FAILED
  RuntimeError: Input type (CUDABFloat16Type) and weight type (CPUBFloat16Type) should be the same

Chain: invokeai/backend/model_manager/load/model_cache/model_cache.py:602 picks CachedModelOnlyFullLoad when enable_partial_loading is false. The trim's guard only checks callable(cur_vram_bytes), which that class satisfies (invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py:104). invokeai/backend/model_manager/load/load_base.py:146 routes into _move_model_to_ram, whose CachedModelOnlyFullLoad branch (invokeai/backend/model_manager/load/model_cache/model_cache.py:1084-1085) calls full_unload_from_vram() and ignores both vram_bytes_to_free and keep_required_weights_in_vram. Only CachedModelWithPartialLoad enables device autocasting (invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py:151), so there is no fallback.

This combination is plausible: nothing in the code or in docs/src/content/docs/configuration/low-vram-mode.mdx says the flag requires enable_partial_loading: true, and the docs list it as one of six independent low-VRAM features.

test_expert_swapper_passes_aggressive_working_memory_to_model_cache cannot catch this: it asserts unload_from_vram by call signature on a MagicMock and never reaches _move_model_to_ram.

To expose this issue, add a test that runs the trim against a real ModelCache built with enable_partial_loading=False and asserts the model is not fully evicted, or that the swapper skips the trim for full-load-only cache entries.

---

Medium — invokeai/backend/model_manager/load/model_cache/model_cache.py:1112-1123: the max_cache_vram_gb semantics change is global and is not gated by the opt-in flag.

Previously the cap short-circuited before working memory was considered, documented as taking "precedence over the device_working_mem_gb and any operations that explicitly reserve additional working memory (e.g. VAE decode)". Now cap - working_mem_bytes applies to every model and every node with wan_memory_optimization at its default false. A user with max_cache_vram_gb: 8 running a decode that reserves 4 GB with 4.5 GB allocated gets a negative budget, so invokeai/backend/model_manager/load/model_cache/model_cache.py:1014-1017 partially unloads the model it is locking. The caution box in docs/src/content/docs/configuration/low-vram-mode.mdx is updated, so the change is deliberate, but it is a behavior change to a global tuning knob shipped inside an opt-in Wan PR, and no test pins the previous contract.

---

Low — invokeai/app/invocations/wan_denoise.py:286-295 is dead code on the supported path, and the 2 GiB figure in the log and docs is not what happens.

_get_vram_available derives its budget from free VRAM (invokeai/backend/model_manager/load/model_cache/model_cache.py:1130-1141), so passing working_mem_bytes = total_memory - 2 GiB leaves 2 GiB - (total - free) for weights. Measured on an otherwise idle card, residency came out at 0.28 GiB of 9.33 GiB (3 percent), already far below the 2 GiB cap, so vram_bytes_to_free is 0 and the clamp never fires when partial loading is on. Its only observable effect is the crash in the finding above. Being more aggressive than advertised is fine for the goal, but the log line at invokeai/app/invocations/wan_denoise.py:645 / invokeai/app/invocations/wan_video_denoise.py:295 and the matching sentence in docs/src/content/docs/configuration/low-vram-mode.mdx claim "about 2 GiB", and the PR's QA step 4 asks testers to confirm that line. Either drop the clamp and reword, or compute the reservation from free VRAM so the number means what it says.

---

Low — invokeai/app/invocations/wan_latents_to_video.py:186-200: the MP4 encode now runs inside the VAE model lock.

The replaced code deliberately closed the model context before encoding ("MP4 encoding can take a while, and holding the full decoded clip in VRAM for its duration starves the next node's model load"). Streaming solves the clip-in-VRAM half but holds the VAE cache record locked and non-evictable for the whole libx264 run. Also, writer.close() at line 200 runs before the num_frames == 0 check at line 216, so a zero-chunk decode surfaces as an ffmpeg error rather than the intended ValueError.

Open Questions

• invokeai/backend/wan/vae_decode.py:18-27 bypasses the @apply_forward_hook decorator on AutoencoderKLWan.decode. Correct for InvokeAI's own cache; would silently skip an accelerate-style offload hook if one were ever attached. None exists in the tree today.
• Enabling the flag changes generation output. The optimized transformer path is arithmetically exact (see Verification), but bf16 rounding through 30 layers produces a systematic ~2.4 percent mean deviation per forward. QA step 7's "confirm existing behavior is unchanged" will not hold bit-for-bit and same-seed videos will differ. Worth a sentence in the docs.

Verification

RTX 4090 24 GiB, torch 2.7.1+cu128, diffusers 0.39.0, PR head c1d62bfda3, real Wan2.2-TI2V-5B and T2V-A14B diffusers weights. 12 GB behaviour simulated with torch.cuda.set_per_process_memory_fraction capped at 11.0 GiB. Import precedence verified so the PR code, not an editable install, was under test.

What works, measured:

• Weight streaming achieves the PR's goal for denoise. 704x1280 / 81 frames on the simulated 12 GB card: OOM without the flag, 1.49 GiB peak with it.
• Activation chunking is effective. At 18480 tokens, peak activation memory 3532 MB -> 1217 MB (-65.5 percent) for a 1.02x slowdown on an unconstrained card.
• The optimized forward is arithmetically exact. Same real weights truncated to 2 blocks in fp32: mean deviation 0.0001 percent of RMS, max 9.4e-06. Isolating the two mechanisms in bf16, chunking alone and timestep compaction alone each produce ~2.4 percent mean deviation, and disabling both returns exactly 0.0 - the delta is accumulation-order rounding from changed GEMM shapes, not a logic error. Both paths are bit-exactly reproducible run to run.
• Timestep compaction confirmed against real weights: diffusers embeds 18480 rows, the PR embeds 2.
• Streaming VAE decode is bit-identical to vae.decode on the real patch_size=2 VAE.
• Backend tests for the touched areas: 145 passed.

Caveats on the numbers: absolute peaks vary by a few hundred MB with caching-allocator state, so the ~2x estimator ratio is the reliable signal rather than any single peak. The A14B VAE was measured standalone, not through a full A14B graph.

Not covered: everything was driven at invocation level rather than through the web server, so a full graph run (model manager, LoRA sidecar patching under a trimmed transformer, dual-expert A14B swap with the flag on, MP4 written to disk) is unverified. The dual-expert swap is where the Finding 2 trim fires repeatedly and is the highest-value manual QA. GGUF-quantized experts are untested here - all measurements used the diffusers bf16 builds.

[JPPhoto]

@Pfannkuchensack I made changes to address your issues:

• Updated invokeai/app/invocations/wan_denoise.py:
  • Describes the 2 GiB value as a target, not guaranteed trim work.
  • Logs partial-loading availability accurately.
  • Added regression coverage for no-op trims and full-load safety.
• Updated docs:
  • Explains that trim may already be satisfied during cache admission.
  • Clarifies partial-loading dependency and global cache-cap behavior.
  • Documents BF16 numerical differences.
  • Documents intentional VAE lock during streaming.
• Added scripts/calibrate_wan_vae_working_memory.py:
  • Measures real CUDA/ROCm reserved-memory peaks.
  • Compares measured usage with estimates.
  • Reports implied scaling constants.
  • Example: python scripts/calibrate_wan_vae_working_memory.py --vae <path>
• Responses to open questions:
  • apply_forward_hook: not applicable; no such API or path exists.
  • BF16 differences: documented; bit-identical output is not promised.
  • MP4 lock: retained intentionally because causal decoder state and weights must stay live. Releasing it requires a bounded decode/encode queue.
  • Global cache-cap semantics: intentionally global, documented, and covered by existing tests.

[JPPhoto]

@Pfannkuchensack Ready for you to look at again!

[Pfannkuchensack]

Re-review of the updated PR (head d19323dd50, +3 commits). Verified on an RTX 4090
with real Wan weights; 12 GB and 16 GB cards simulated by capping the allocator.

Previously raised, now resolved

• Crash with enable_partial_loading: false — fixed. The new supports_partial_loading property (invokeai/backend/model_manager/load/load_base.py:134-138) now gates both the working-memory reservation and the trim (invokeai/app/invocations/wan_denoise.py:270, :288-305). Verified by driving the real _ExpertSwapper against a real ModelCache holding the real 5B transformer:

  enable_partial_loading=True   supports_partial_loading=True   resident 0.42/9.33 GiB  FORWARD OK
  enable_partial_loading=False  supports_partial_loading=False  resident 9.33/9.33 GiB  FORWARD OK
                                                                warning logged, trim skipped
  

• VAE decode OOM on the target hardware — fixed. Variant-specific constants (invokeai/backend/util/vae_working_memory.py:12-14, :158-170) now push the tiling fallback over its threshold where it matters. Same 704x1280 / 81-frame clip that OOMed before:

  12 GB card  TI2V  est 11.77 GiB > 10.80 -> TILED      alloc_peak  3.25 GiB  OK   (was OOM)
  12 GB card  A14B  est 10.93 GiB > 10.80 -> TILED      alloc_peak  2.15 GiB  OK   (was OOM)
  16 GB card  TI2V  est 11.77 GiB < 14.40 -> STREAMING  alloc_peak 11.28 GiB  OK
  16 GB card  A14B  est 10.93 GiB < 14.40 -> STREAMING  alloc_peak  9.75 GiB  OK
  

  Calibration against allocated peak is now good and slightly conservative: TI2V 11.77 estimated vs 11.28 measured (1.04x), A14B 10.93 vs 9.75 (1.12x). The pixel_frames > 1 split keeping 2900 for single-frame decode is correct - one chunk accumulates no causal state, matching the image path.

• Log line and docs reworded to stop promising a hard 2 GiB; bf16 output divergence now documented; MP4-inside-the-lock documented with its rationale; the max_cache_vram_gb caution now states explicitly that the subtraction applies to every model-cache operation, not only Wan.

All 628 tests in the touched areas pass, including the three new swapper tests and the new supports_partial_loading test.

Findings

Medium — scripts/calibrate_wan_vae_working_memory.py (outside invokeai/), _measure: the --tiling implied constant is computed on the wrong basis and is off by ~11x.

implied_constant = (measured_delta - clip_bytes) / (pixel_height * pixel_width * element_size) is correct for the untiled path, where the estimate's per-frame term is H * W * element_size * constant. Under --tiling the estimate uses tile_size**2 * element_size * constant * 1.25 (invokeai/backend/util/vae_working_memory.py:172-174), so dividing by H * W yields a number that cannot be compared to the shipped constants:

--tiling, TI2V, 704x1280x81, tile 256:
  estimate 1.884 GiB   measured reserved delta 2.398 GiB   implied scaling constant: 942.9
  recomputed on the tiled basis ((measured - clip) / (256^2 * 2 * 1.25)): ~10120

942.9 against a shipped 7000 reads as "the constant is 7x too high". docs/src/content/docs/configuration/low-vram-mode.mdx explicitly directs developers to this flag ("Use --tiling to measure the spatially tiled full-decode fallback"), so the trap is on the documented path. Fix: when tiling is set, divide by tile_size**2 * element_size * 1.25.

To expose this issue, add a test that runs _measure with tiling=True against a stubbed peak equal to the estimate and asserts the returned implied_scaling_constant equals the constant the estimator used.

---

Medium — scripts/calibrate_wan_vae_working_memory.py (outside invokeai/): the script measures memory_reserved while the shipped constants track memory_allocated, so it disagrees with the PR's own values on the PR's own default shape.

--vae <TI2V dir> --dtype bfloat16   (the documented "12 GiB-card calibration point")
  estimate 11.769 GiB   measured reserved delta 16.727 GiB   implied constant 9953.4
--vae <TI2V checkpoint> 480x832x33
  estimate  5.216 GiB   measured reserved delta  7.453 GiB   implied constant 10007.4

Shipped is 7000. A maintainer recalibrating from this tool would raise the constants ~1.4x and push cards into tiling that do not need it. The reserved figure is allocator greed rather than a requirement - the same TI2V decode that reports a 16.73 GiB reserved delta on an unconstrained 24 GB card completes under a 15 GiB cap with a 13.25 GiB reserved delta and an 11.28 GiB allocated peak. Either report both metrics and say which one the constants are fit to, or switch the implied-constant arithmetic to max_memory_allocated and keep reserved as an informational headroom line.

Both load paths otherwise work: the diffusers-directory and single-.safetensors branches produced consistent results (9953 vs 10007), so _wan_vae_init_kwargs_for plus assign=True loading is sound.

Residual Risk

• The two new swapper tests assert the guard on MagicMocks, so they pin the branch but not CachedModelOnlyFullLoad's underlying "ignores vram_bytes_to_free" semantics. That behaviour at invokeai/backend/model_manager/load/model_cache/model_cache.py:1084-1085 is still a sharp edge for any future caller of unload_from_vram; only the Wan call site is now guarded. Worth a docstring note on LoadedModelWithoutConfig.unload_from_vram that partial byte counts are ignored for full-load-only entries.
• Card sizes between the tiling threshold and the real requirement were not swept exhaustively; 12 GB and 16 GB were checked at 704x1280 / 81 frames. 10 GB and 8 GB cards, and A14B at higher frame counts, are unverified.
• Still driven at invocation level rather than through the web server. A full graph run - dual-expert A14B swap with the flag on, LoRA sidecar patching over a trimmed transformer, MP4 written to disk - remains the highest-value manual QA. GGUF experts untested; all measurements used diffusers bf16 builds.
• Peak figures move by a few hundred MB with allocator state and with pytorch_cuda_alloc_conf; measurements above used the default allocator, not the backend:cudaMallocAsync some users configure.

[JPPhoto]

@Pfannkuchensack Thanks for the re-review. Addressed the remaining findings:

• Tiled calibration now derives its implied constant from tile_size^2 * element_size * 1.25, matching the shipped estimator. Added regression coverage.
• Calibration reports allocated and reserved deltas separately. The implied constant uses peak allocated memory; reserved memory remains available as allocator-headroom diagnostics. Added coverage that distinguishes the two.
• Documented that full-load-only cache entries ignore partial unload byte requests and unload all weights.
• Updated calibration documentation and Wan estimator comments.

[Pfannkuchensack]

All good now.