Add opt-in low-VRAM mode for Wan generation

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

@@ -32,16 +32,27 @@ enable_partial_loading: false
 
 ## Details and fine-tuning
 
-Low-VRAM mode involves 4 features, each of which can be configured or fine-tuned:
+Low-VRAM mode involves 6 features, each of which can be configured or fine-tuned:
 
 - Partial model loading (`enable_partial_loading`)
 - PyTorch CUDA allocator config (`pytorch_cuda_alloc_conf`)
 - 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`)
 
 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 limits resident Wan transformer weights to about 2 GiB and streams remaining layers from RAM. 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.
+
 ### Partial model loading
 
 Invoke's partial model loading works by streaming model "layers" between RAM and VRAM as they are needed.

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 limit resident transformer weights to about 2 GiB, lower denoise activation memory, and stream untiled VAE decode directly to MP4. This can make generation substantially slower and requires enough system RAM for offloaded weights.
+
 * **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": "auto",

invokeai/app/invocations/wan_denoise.py

@@ -52,13 +52,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 Wan weights use aggressive layer streaming."""
+    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 +189,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,6 +200,8 @@ 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
@@ -242,7 +259,10 @@ 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()
+        if self._working_mem_bytes is None:
+            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 +276,17 @@ 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:
+            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)
+            partial_unload = getattr(cached_model, "partial_unload_from_vram", None)
+            if callable(cur_vram_bytes) and callable(partial_unload):
+                vram_bytes_to_free = max(0, cur_vram_bytes() - self._max_resident_model_bytes)
+                if vram_bytes_to_free > 0:
+                    partial_unload(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 +632,10 @@ 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: limiting resident transformer weights to about 2 GiB")
         with ExitStack() as exit_stack:
             swapper = _ExpertSwapper(
                 context=context,
@@ -611,6 +646,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 +680,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]
 

invokeai/app/invocations/wan_latents_to_video.py

@@ -36,6 +36,7 @@
 from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_wan
+from invokeai.backend.wan.vae_decode import iter_wan_vae_decode_chunks
 
 
 class _FrameWriter(Protocol):
@@ -117,13 +118,15 @@ def invoke(self, context: InvocationContext) -> VideoOutput:
         temporal_scale = getattr(vae_info.model.config, "scale_factor_temporal", None) or 4
         t_pixel = (t_lat - 1) * temporal_scale + 1
         h_pixel, w_pixel = h_lat * spatial_scale, w_lat * spatial_scale
+        optimize_memory = context.config.get().wan_memory_optimization
 
         estimated_working_memory = estimate_vae_working_memory_wan(
             operation="decode",
             vae=vae_info.model,
             pixel_height=h_pixel,
             pixel_width=w_pixel,
             pixel_frames=t_pixel,
+            streaming=optimize_memory,
         )
         # Long/high-res clips can need a working set no card fits. When the full-frame
         # estimate exceeds the execution device's total VRAM, fall back to spatial tiling
@@ -144,72 +147,82 @@ def invoke(self, context: InvocationContext) -> VideoOutput:
                         pixel_width=w_pixel,
                         pixel_frames=t_pixel,
                         tile_size=tile_size,
+                        streaming=False,
                     )
 
-        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
-            assert isinstance(vae, AutoencoderKLWan)
-            context.logger.info(
-                f"Running Wan VAE decode: {t_lat} latent frames -> {t_pixel} pixel frames at {w_pixel}x{h_pixel}"
-                + (" (tiled)" if use_tiling else "")
-            )
-            context.util.signal_progress("Running Wan VAE decode (video)")
-
-            vae_dtype = next(iter(vae.parameters())).dtype
-            latents = latents.to(device=get_effective_device(vae), dtype=vae_dtype)
-
-            TorchDevice.empty_cache()
-
-            if use_tiling:
-                vae.enable_tiling()
-            try:
-                with torch.inference_mode():
-                    # Denormalise from denoiser space back to VAE space.
-                    latents_mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latents)
-                    latents_std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1).to(latents)
-                    latents = latents * latents_std + latents_mean
-
-                    # [B, C=3, T_pixel, H, W] in [-1, 1] (roughly).
-                    decoded = vae.decode(latents, return_dict=False)[0]
-                    del latents, latents_mean, latents_std
-            finally:
-                if use_tiling:
-                    # The VAE instance is cached and shared; don't leak tiling into other nodes.
-                    vae.disable_tiling()
-
-            # Take batch 0 (we generate one video at a time) and move the clip off the
-            # accelerator now — MP4 encoding can take a while, and holding the full
-            # decoded clip in VRAM for its duration starves the next node's model load.
-            decoded = decoded[0].cpu()  # [C, T, H, W]
-
-        TorchDevice.empty_cache()
-
-        if context.util.is_canceled():
-            raise CanceledException
-
-        num_frames = decoded.shape[1]
-        if num_frames == 0:
-            raise ValueError("Wan VAE decode produced zero frames.")
-
-        height, width = decoded.shape[2:]
-        duration = num_frames / float(self.fps)
-
-        # Encode to a temporary MP4 (libx264 + yuv420p, exact frame dimensions —
-        # see make_mp4_writer for why macro_block_size matters).
         tmp = tempfile.NamedTemporaryFile(prefix="invokeai_wan_video_", suffix=".mp4", delete=False)
         tmp.close()
         tmp_path = Path(tmp.name)
         try:
-            context.logger.info(
-                f"Encoding MP4: {num_frames} frames @ {self.fps} fps ({duration:.2f}s) at {width}x{height} via libx264"
-            )
-            context.util.signal_progress(f"Encoding MP4 ({num_frames} frames @ {self.fps} fps)")
-            writer = make_mp4_writer(tmp_path, self.fps)
-            try:
-                _write_video_frames(writer, _iter_decoded_frames(decoded), context.util.is_canceled)
-            finally:
-                writer.close()
-            del decoded
+            stream_decode = optimize_memory and not use_tiling
+            decoded: torch.Tensor | None = None
+            num_frames = 0
+
+            with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
+                assert isinstance(vae, AutoencoderKLWan)
+                context.logger.info(
+                    f"Running Wan VAE decode: {t_lat} latent frames -> {t_pixel} pixel frames at {w_pixel}x{h_pixel}"
+                    + (" (tiled)" if use_tiling else " (streaming to MP4)" if stream_decode else "")
+                )
+                context.util.signal_progress("Running Wan VAE decode (video)")
+
+                vae_dtype = next(iter(vae.parameters())).dtype
+                latents = latents.to(device=get_effective_device(vae), dtype=vae_dtype)
+                TorchDevice.empty_cache()
+
+                if use_tiling:
+                    vae.enable_tiling()
+                try:
+                    with torch.inference_mode():
+                        # Denormalise from denoiser space back to VAE space.
+                        latents_mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latents)
+                        latents_std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1).to(latents)
+                        latents = latents * latents_std + latents_mean
+
+                        if stream_decode:
+                            writer = make_mp4_writer(tmp_path, self.fps)
+                            try:
+                                for chunk in iter_wan_vae_decode_chunks(vae, latents):
+                                    chunk = chunk[0].cpu()
+                                    num_frames += chunk.shape[1]
+                                    _write_video_frames(writer, _iter_decoded_frames(chunk), context.util.is_canceled)
+                            finally:
+                                writer.close()
+                        else:
+                            # [C=3, T_pixel, H, W] in [-1, 1] (roughly), on CPU.
+                            decoded = vae.decode(latents, return_dict=False)[0][0].cpu()
+                            num_frames = decoded.shape[1]
+                        del latents, latents_mean, latents_std
+                finally:
+                    if use_tiling:
+                        # The VAE instance is cached and shared; don't leak tiling into other nodes.
+                        vae.disable_tiling()
+
             TorchDevice.empty_cache()
+
+            if context.util.is_canceled():
+                raise CanceledException
+            if num_frames == 0:
+                raise ValueError("Wan VAE decode produced zero frames.")
+            if num_frames != t_pixel:
+                raise ValueError(f"Wan VAE decode produced {num_frames} frames; expected {t_pixel}.")
+
+            height, width = h_pixel, w_pixel
+            duration = num_frames / float(self.fps)
+            if decoded is not None:
+                context.logger.info(
+                    f"Encoding MP4: {num_frames} frames @ {self.fps} fps "
+                    f"({duration:.2f}s) at {width}x{height} via libx264"
+                )
+                context.util.signal_progress(f"Encoding MP4 ({num_frames} frames @ {self.fps} fps)")
+                writer = make_mp4_writer(tmp_path, self.fps)
+                try:
+                    _write_video_frames(writer, _iter_decoded_frames(decoded), context.util.is_canceled)
+                finally:
+                    writer.close()
+                del decoded
+                TorchDevice.empty_cache()
+
             encoded_bytes = tmp_path.stat().st_size
             context.logger.info(f"MP4 encode complete: {encoded_bytes / 1024:.1f} KB")
             video_dto = context.videos.save(