Add opt-in low-VRAM mode for Wan generation

docs/src/content/docs/configuration/low-vram-mode.mdx

@@ -39,10 +39,27 @@ Low-VRAM mode and related workload-specific optimizations include:
 - Dynamic RAM and VRAM cache sizes (`max_cache_ram_gb`, `max_cache_vram_gb`)
 - Working memory (`device_working_mem_gb`)
 - Keeping a RAM weight copy (`keep_ram_copy_of_weights`)
+- Wan video memory optimization (`wan_memory_optimization`)
 - PiD decode activation chunking (`pid_memory_optimization`)
 
 Read on to learn about these features and understand how to fine-tune them for your system and use-cases.
 
+### Wan video memory optimization
+
+Wan video generation has an additional opt-in memory optimization:
+
+```yaml
+wan_memory_optimization: true
+```
+
+This reserves VRAM so partial-load Wan transformer weights target about 2 GiB resident and stream remaining layers from RAM when partial model loading is enabled (the default). The explicit residency trim may be a no-op when cache admission already reaches that target. If `enable_partial_loading: false`, the activation, timestep, and VAE optimizations still apply, but transformer weights remain fully resident and the 2 GiB residency target is unavailable. It also chunks pointwise transformer activations, compacts TI2V per-token timestep conditioning, and streams untiled VAE decode chunks directly to MP4. It reduces peak VRAM during both denoise and decode, but generation can be substantially slower and requires enough system RAM for offloaded weights. Spatially tiled VAE decode continues to use its existing full-tile path.
+
+The optimized BF16 transformer path can produce small numerical differences because chunked matrix operations accumulate in a different order. Same-seed output is not guaranteed to be bit-identical; set `wan_memory_optimization: false` for the baseline path.
+
+Developers with a CUDA or ROCm device can validate Wan VAE memory estimates with `python scripts/calibrate_wan_vae_working_memory.py --vae <directory-or-safetensors-file>`. The script reports allocated and reserved memory deltas; its implied scaling constant uses allocated memory to match the shipped estimator, while reserved memory shows allocator headroom. Use `--tiling` to measure the spatially tiled full-decode fallback; it overrides streaming mode. Use `--tile-size <pixels>` to override the VAE's default tile size.
+
+Direct MP4 streaming keeps the VAE cache lock while chunks are decoded and written. Wan's causal decoder state and weights must remain live for the sequence; releasing the lock would require a separate bounded decode and encode queue.
+
 ### Partial model loading
 
 Invoke's partial model loading works by streaming model "layers" between RAM and VRAM as they are needed.
@@ -102,7 +119,7 @@ max_cache_vram_gb: 16
 ```
 
 :::caution[Max safe value for `max_cache_vram_gb`]
-  Most users should not manually configure the `max_cache_vram_gb`. This configuration value takes precedence over the `device_working_mem_gb` and any operations that explicitly reserve additional working memory (e.g. VAE decode). As such, manually configuring it increases the likelihood of encountering out-of-memory errors.
+  Most users should not manually configure the `max_cache_vram_gb`. This configuration value caps model-cache residency; `device_working_mem_gb` and operation-specific reservations (e.g. VAE decode) are still subtracted from that cap for every model-cache operation, not only when Wan memory optimization is enabled. A cap below the active working-memory reservation can force aggressive model offloading.
 
   For users who wish to configure `max_cache_vram_gb`, the max safe value can be determined by subtracting `device_working_mem_gb` from your GPU's VRAM. As described below, the default for `device_working_mem_gb` is 3GB.
 

docs/src/content/docs/features/video-generation.mdx

@@ -212,6 +212,8 @@ A real failure mode of long chains: each iteration's reference image is itself a
 
 Video denoise is memory-intensive — attention scales roughly as `(T_lat × H/16 × W/16)²`, so resolution and frame count both quadratically affect peak VRAM.
 
+Add `wan_memory_optimization: true` to `invokeai.yaml` and restart Invoke to target about 2 GiB of resident transformer weights when `enable_partial_loading` is enabled, lower denoise activation memory, and stream untiled VAE decode directly to MP4. The explicit residency trim may be a no-op when cache admission already reaches that target. If partial loading is disabled, the activation, timestep, and VAE optimizations remain active but transformer weights are fully resident. This can make generation substantially slower and requires enough system RAM for offloaded weights. Optimized BF16 execution may produce small numerical differences from the baseline path.
+
 * **Drop resolution before frame count.** Going from 1280×720 to 832×480 is a ~2.4× memory drop and visually subtle in most content. Going from 81 frames to 65 only saves ~20%.
 * **TI2V-5B before A14B.** TI2V-5B Q4_K_M peaks around ~6–8 GB at 832×480, versus ~12–14 GB for A14B Q4_K_M. If you're at the OOM edge, switch model family.
 * **OOM at the *reference image encoder* step** is usually allocator fragmentation from a previous run rather than absolute memory pressure. Restart the dev server and try again; if it recurs reproducibly, file an issue.

docs/src/generated/settings.json

@@ -561,6 +561,17 @@
       "type": "<class 'bool'>",
       "validation": {}
     },
+    {
+      "category": "GENERATION",
+      "default": false,
+      "description": "Enable experimental Wan memory optimizations at the cost of slower generation.",
+      "env_var": "INVOKEAI_WAN_MEMORY_OPTIMIZATION",
+      "literal_values": [],
+      "name": "wan_memory_optimization",
+      "required": false,
+      "type": "<class 'bool'>",
+      "validation": {}
+    },
     {
       "category": "GENERATION",
       "default": false,

invokeai/app/invocations/wan_denoise.py

@@ -44,6 +44,7 @@
 from invokeai.app.invocations.model import LoRAField, WanTransformerField
 from invokeai.app.invocations.primitives import LatentsOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.model_manager.load.model_cache.model_cache import MODEL_LOAD_LOCK
 from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, WanVariantType
 from invokeai.backend.patches.layer_patcher import LayerPatcher, PatchSpec
 from invokeai.backend.patches.lora_conversions.wan_lora_constants import WAN_LORA_TRANSFORMER_PREFIX
@@ -52,13 +53,26 @@
 from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import WanConditioningInfo
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.wan.memory_optimization import wan_memory_optimization
 from invokeai.backend.wan.sampling_utils import get_spatial_scale_factor, make_noise
 
 # Type alias: a factory that produces a fresh iterator of LoRA patch specs each time it is called.
 # We need fresh iterators because the patcher
 # consumes the iterator once per ``apply_smart_model_patches`` invocation, and
 # the expert may be swapped (and re-entered) multiple times in a render.
 LoRAIteratorFactory = Callable[[], Iterable[PatchSpec]]
+WAN_MAX_RESIDENT_TRANSFORMER_BYTES = 2 * 2**30
+
+
+def _get_wan_transformer_working_mem_bytes(device: torch.device, *, enabled: bool) -> int | None:
+    """Reserve all but 2 GiB of VRAM so partial-load Wan weights target about 2 GiB resident."""
+    if not enabled or device.type != "cuda":
+        return None
+
+    total_vram = torch.cuda.get_device_properties(device).total_memory
+    if total_vram <= WAN_MAX_RESIDENT_TRANSFORMER_BYTES:
+        return None
+    return total_vram - WAN_MAX_RESIDENT_TRANSFORMER_BYTES
 
 
 def _resolve_variant(context: InvocationContext, transformer_field: WanTransformerField) -> WanVariantType:
@@ -176,6 +190,8 @@ def __init__(
         low_lora_factory: LoRAIteratorFactory | None = None,
         high_is_quantized: bool = False,
         low_is_quantized: bool = False,
+        working_mem_bytes: int | None = None,
+        max_resident_model_bytes: int | None = None,
     ) -> None:
         self._context = context
         self._high_model = high_model
@@ -185,11 +201,14 @@ def __init__(
         self._low_lora_factory = low_lora_factory
         self._high_is_quantized = high_is_quantized
         self._low_is_quantized = low_is_quantized
+        self._working_mem_bytes = working_mem_bytes
+        self._max_resident_model_bytes = max_resident_model_bytes
         self._active_label: str | None = None
         self._active_info: Any | None = None
         self._active_device_ctx: Any | None = None
         self._active_lora_ctx: Any | None = None
         self._active_model: Any | None = None
+        self._warned_partial_loading_unavailable = False
 
     def get(self, label: str) -> Any:
         if label not in (self.HIGH, self.LOW):
@@ -203,11 +222,10 @@ def get(self, label: str) -> Any:
         # Capture the outgoing expert's cache record before _release() drops our handle.
         # We need it to force-unload below.
         outgoing_cached_model = None
+        outgoing_info = self._active_info
         if self._active_info is not None:
-            # ``LoadedModel`` exposes its cache_record only via a private attribute. There
-            # is no public ``unload_from_vram`` on the LoadedModel today, and we don't want
-            # to take on a broader backend refactor in this fix; tolerate AttributeError
-            # so a future refactor doesn't break the swap.
+            # ``LoadedModel`` keeps the cache record private, but exposes
+            # ``unload_from_vram`` so cache error handling stays in one place.
             outgoing_cached_model = getattr(self._active_info, "_cache_record", None)
             if outgoing_cached_model is not None:
                 outgoing_cached_model = getattr(outgoing_cached_model, "cached_model", None)
@@ -229,7 +247,14 @@ def get(self, label: str) -> Any:
         # and now — the cached_model object still owns the tensors.
         if outgoing_cached_model is not None:
             try:
-                outgoing_cached_model.full_unload_from_vram()
+                unload_from_vram = getattr(outgoing_info, "unload_from_vram", None)
+                if callable(unload_from_vram):
+                    unload_from_vram(outgoing_cached_model.total_bytes())
+                else:
+                    # Keep compatibility with old LoadedModel handles while preserving
+                    # the process-global register_parameter guard.
+                    with MODEL_LOAD_LOCK.read_lock():
+                        outgoing_cached_model.full_unload_from_vram()
             except Exception:
                 pass
 
@@ -242,7 +267,11 @@ def get(self, label: str) -> Any:
         # always fresh — see class docstring for the cache-eviction reasoning.
         model_id = self._high_model if label == self.HIGH else self._low_model
         info = self._context.models.load(model_id)
-        device_ctx = info.model_on_device()
+        supports_partial_loading = getattr(info, "supports_partial_loading", None)
+        if self._working_mem_bytes is None or supports_partial_loading is False:
+            device_ctx = info.model_on_device()
+        else:
+            device_ctx = info.model_on_device(working_mem_bytes=self._working_mem_bytes)
         cached_weights, model = device_ctx.__enter__()
 
         # Stash the device-context state immediately. If anything below fails (most
@@ -256,6 +285,25 @@ def get(self, label: str) -> Any:
         self._active_device_ctx = device_ctx
         self._active_model = model
 
+        if self._max_resident_model_bytes is not None:
+            if supports_partial_loading is False:
+                if not self._warned_partial_loading_unavailable:
+                    self._context.logger.warning(
+                        "Wan memory optimization cannot limit resident transformer weights because "
+                        "partial model loading is disabled."
+                    )
+                    self._warned_partial_loading_unavailable = True
+            else:
+                cache_record = getattr(info, "_cache_record", None)
+                cached_model = getattr(cache_record, "cached_model", None)
+                cur_vram_bytes = getattr(cached_model, "cur_vram_bytes", None)
+                unload_from_vram = getattr(info, "unload_from_vram", None)
+                if callable(cur_vram_bytes) and callable(unload_from_vram):
+                    vram_bytes_to_free = max(0, cur_vram_bytes() - self._max_resident_model_bytes)
+                    if vram_bytes_to_free > 0:
+                        unload_from_vram(vram_bytes_to_free, keep_required_weights_in_vram=True)
+                        TorchDevice.empty_cache()
+
         # Apply LoRA patches for this expert. GGUF transformers need sidecar
         # patching since direct patching of GGMLTensors isn't supported.
         lora_factory = self._high_lora_factory if label == self.HIGH else self._low_lora_factory
@@ -601,6 +649,13 @@ def high_lora_factory() -> Iterable[PatchSpec]:
         def low_lora_factory() -> Iterable[PatchSpec]:
             return self._lora_iterator(context, low_loras)
 
+        optimize_memory = context.config.get().wan_memory_optimization
+        working_mem_bytes = _get_wan_transformer_working_mem_bytes(device, enabled=optimize_memory)
+        if working_mem_bytes is not None:
+            context.logger.info(
+                "Wan memory optimization: targeting about 2 GiB of resident transformer weights when partial "
+                "loading is available"
+            )
         with ExitStack() as exit_stack:
             swapper = _ExpertSwapper(
                 context=context,
@@ -611,6 +666,10 @@ def low_lora_factory() -> Iterable[PatchSpec]:
                 low_lora_factory=low_lora_factory if low_loras else None,
                 high_is_quantized=high_is_quantized,
                 low_is_quantized=low_is_quantized,
+                working_mem_bytes=working_mem_bytes,
+                max_resident_model_bytes=(
+                    WAN_MAX_RESIDENT_TRANSFORMER_BYTES if working_mem_bytes is not None else None
+                ),
             )
             exit_stack.callback(swapper.close)
 
@@ -641,25 +700,26 @@ def low_lora_factory() -> Iterable[PatchSpec]:
                 if ref_condition is not None:
                     latent_model_input = torch.cat([latent_model_input, ref_condition], dim=1)
 
-                noise_pred_cond = transformer(
-                    hidden_states=latent_model_input,
-                    timestep=timestep,
-                    encoder_hidden_states=pos_cond.prompt_embeds.unsqueeze(0),
-                    attention_kwargs=None,
-                    return_dict=False,
-                )[0]
-
-                if neg_cond is not None and active_cfg != 1.0:
-                    noise_pred_uncond = transformer(
+                with wan_memory_optimization(transformer, enabled=optimize_memory):
+                    noise_pred_cond = transformer(
                         hidden_states=latent_model_input,
                         timestep=timestep,
-                        encoder_hidden_states=neg_cond.prompt_embeds.unsqueeze(0),
+                        encoder_hidden_states=pos_cond.prompt_embeds.unsqueeze(0),
                         attention_kwargs=None,
                         return_dict=False,
                     )[0]
-                    noise_pred = noise_pred_uncond + active_cfg * (noise_pred_cond - noise_pred_uncond)
-                else:
-                    noise_pred = noise_pred_cond
+
+                    if neg_cond is not None and active_cfg != 1.0:
+                        noise_pred_uncond = transformer(
+                            hidden_states=latent_model_input,
+                            timestep=timestep,
+                            encoder_hidden_states=neg_cond.prompt_embeds.unsqueeze(0),
+                            attention_kwargs=None,
+                            return_dict=False,
+                        )[0]
+                        noise_pred = noise_pred_uncond + active_cfg * (noise_pred_cond - noise_pred_uncond)
+                    else:
+                        noise_pred = noise_pred_cond
 
                 latents = scheduler.step(noise_pred, t, latents, return_dict=False)[0]