diff --git a/docs/src/content/docs/configuration/low-vram-mode.mdx b/docs/src/content/docs/configuration/low-vram-mode.mdx index fa15cef8735..09f8f50a14a 100644 --- a/docs/src/content/docs/configuration/low-vram-mode.mdx +++ b/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. @@ -101,7 +112,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. 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. diff --git a/docs/src/content/docs/features/video-generation.mdx b/docs/src/content/docs/features/video-generation.mdx index c8f205e2d44..3ba4b25b9ae 100644 --- a/docs/src/content/docs/features/video-generation.mdx +++ b/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. diff --git a/docs/src/generated/settings.json b/docs/src/generated/settings.json index 2c183f55400..ca16a76f11b 100644 --- a/docs/src/generated/settings.json +++ b/docs/src/generated/settings.json @@ -561,6 +561,17 @@ "type": "", "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": "", + "validation": {} + }, { "category": "GENERATION", "default": "auto", diff --git a/invokeai/app/invocations/wan_denoise.py b/invokeai/app/invocations/wan_denoise.py index 4a2c9c76c97..53f17f62d5e 100644 --- a/invokeai/app/invocations/wan_denoise.py +++ b/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,6 +53,7 @@ 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. @@ -59,6 +61,18 @@ # 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 +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,6 +201,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 @@ -203,11 +221,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 +246,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 +266,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 +283,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) + 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 +639,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 +653,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 +687,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] diff --git a/invokeai/app/invocations/wan_latents_to_video.py b/invokeai/app/invocations/wan_latents_to_video.py index 305a13a11eb..3302c9fbe70 100644 --- a/invokeai/app/invocations/wan_latents_to_video.py +++ b/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,6 +118,7 @@ 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", @@ -124,6 +126,7 @@ def invoke(self, context: InvocationContext) -> VideoOutput: 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,92 @@ 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() + tmp = tempfile.NamedTemporaryFile(prefix="invokeai_wan_video_", suffix=".mp4", delete=False) + tmp.close() + tmp_path = Path(tmp.name) + try: + 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 - - # [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.enable_tiling() + else: + # AutoencoderKLWan is cached and shared with Anima. Clear any + # tiling state left by a prior image decode before streaming. vae.disable_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: + duration = t_pixel / float(self.fps) + context.logger.info( + f"Encoding MP4: {t_pixel} frames @ {self.fps} fps " + f"({duration:.2f}s) at {w_pixel}x{h_pixel} via libx264" + ) + context.util.signal_progress(f"Encoding MP4 ({t_pixel} frames @ {self.fps} fps)") + 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: + # The VAE instance is cached and shared; don't leak tiling into other nodes. + if use_tiling: + 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.") + TorchDevice.empty_cache() - height, width = decoded.shape[2:] - duration = num_frames / float(self.fps) + 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() - # 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 - 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( diff --git a/invokeai/app/invocations/wan_video_denoise.py b/invokeai/app/invocations/wan_video_denoise.py index c0761eefaa2..0d76d27c108 100644 --- a/invokeai/app/invocations/wan_video_denoise.py +++ b/invokeai/app/invocations/wan_video_denoise.py @@ -28,8 +28,10 @@ from invokeai.app.invocations.model import WanTransformerField from invokeai.app.invocations.primitives import LatentsOutput from invokeai.app.invocations.wan_denoise import ( + WAN_MAX_RESIDENT_TRANSFORMER_BYTES, WanDenoiseInvocation, _ExpertSwapper, + _get_wan_transformer_working_mem_bytes, _resolve_variant, _validate_ref_condition_shape, _validate_spatial_dimensions, @@ -40,6 +42,7 @@ 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_default_latent_channels, get_spatial_scale_factor, @@ -288,6 +291,10 @@ def high_lora_factory() -> Iterable[PatchSpec]: def low_lora_factory() -> Iterable[PatchSpec]: return proxy._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, @@ -298,6 +305,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) @@ -336,25 +347,26 @@ def low_lora_factory() -> Iterable[PatchSpec]: # T2V (any variant): scalar timestep per batch. timestep = t.expand(latents.shape[0]) - 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] diff --git a/invokeai/app/services/config/config_default.py b/invokeai/app/services/config/config_default.py index ae1e38e0ccc..fd6f1ceba32 100644 --- a/invokeai/app/services/config/config_default.py +++ b/invokeai/app/services/config/config_default.py @@ -113,6 +113,7 @@ class InvokeAIAppConfig(BaseSettings): device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `cuda:N` (where N is a device number) precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.
Valid values: `auto`, `float16`, `bfloat16`, `float32` sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements. + wan_memory_optimization: Enable experimental Wan memory optimizations at the cost of slower generation. attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp` attention_slice_size: Slice size, valid when attention_type=="sliced".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8` force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). @@ -220,6 +221,7 @@ class InvokeAIAppConfig(BaseSettings): # GENERATION sequential_guidance: bool = Field(default=False, description="Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.") + wan_memory_optimization: bool = Field(default=False, description="Enable experimental Wan memory optimizations at the cost of slower generation.") attention_type: ATTENTION_TYPE = Field(default="auto", description="Attention type.") attention_slice_size: ATTENTION_SLICE_SIZE = Field(default="auto", description='Slice size, valid when attention_type=="sliced".') force_tiled_decode: bool = Field(default=False, description="Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).") diff --git a/invokeai/backend/model_manager/load/load_base.py b/invokeai/backend/model_manager/load/load_base.py index 7225fd1402f..85f0e13b135 100644 --- a/invokeai/backend/model_manager/load/load_base.py +++ b/invokeai/backend/model_manager/load/load_base.py @@ -143,6 +143,19 @@ def repair_required_tensors_on_device(self) -> int: with MODEL_LOAD_LOCK.read_lock(): return cached_model.repair_required_tensors_on_compute_device() + def unload_from_vram(self, vram_bytes_to_free: int, keep_required_weights_in_vram: bool = False) -> int: + """Unload model weights through the cache's failure-safe path. + + The model may be partially resident. The caller must keep its model handle + alive while unloading; the cache entry can be evicted independently. + """ + with MODEL_LOAD_LOCK.read_lock(): + return self._cache.unload_model_from_vram( + self._cache_record, + vram_bytes_to_free, + keep_required_weights_in_vram=keep_required_weights_in_vram, + ) + class LoadedModel(LoadedModelWithoutConfig): """Context manager object that mediates transfer from RAM<->VRAM.""" diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index 8f8296b8674..9c056fd438c 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -1065,11 +1065,21 @@ def _move_model_to_vram(self, cache_entry: CacheRecord, vram_available: int) -> self._delete_cache_entry(cache_entry) raise - def _move_model_to_ram(self, cache_entry: CacheRecord, vram_bytes_to_free: int) -> int: + def _move_model_to_ram( + self, + cache_entry: CacheRecord, + vram_bytes_to_free: int, + keep_required_weights_in_vram: bool | None = None, + ) -> int: try: if isinstance(cache_entry.cached_model, CachedModelWithPartialLoad): return cache_entry.cached_model.partial_unload_from_vram( - vram_bytes_to_free, keep_required_weights_in_vram=cache_entry.is_locked + vram_bytes_to_free, + keep_required_weights_in_vram=( + cache_entry.is_locked + if keep_required_weights_in_vram is None + else keep_required_weights_in_vram + ), ) elif isinstance(cache_entry.cached_model, CachedModelOnlyFullLoad): # type: ignore return cache_entry.cached_model.full_unload_from_vram() @@ -1080,18 +1090,38 @@ def _move_model_to_ram(self, cache_entry: CacheRecord, vram_bytes_to_free: int) self._delete_cache_entry(cache_entry) raise + @synchronized + def unload_model_from_vram( + self, + cache_entry: CacheRecord, + vram_bytes_to_free: int, + keep_required_weights_in_vram: bool = False, + ) -> int: + """Unload model weights through cache error handling. + + Caller must hold the model's usage lock when unloading a model that is in use. + The cache entry may already have been evicted; the cached model remains safe to + operate on while its owning handle is still alive. + """ + return self._move_model_to_ram( + cache_entry, + vram_bytes_to_free, + keep_required_weights_in_vram=keep_required_weights_in_vram, + ) + def _get_vram_available(self, working_mem_bytes: Optional[int]) -> int: """Calculate the amount of additional VRAM available for the cache to use (takes into account the working memory). """ - # If self._max_vram_cache_size_gb is set, then it overrides the default logic. - if self._max_vram_cache_size_gb is not None: - vram_total_available_to_cache = int(self._max_vram_cache_size_gb * GB) - return vram_total_available_to_cache - self._get_vram_in_use() - working_mem_bytes_default = int(self._execution_device_working_mem_gb * GB) working_mem_bytes = max(working_mem_bytes or working_mem_bytes_default, working_mem_bytes_default) + # An explicit cache cap limits model residency, but operation-specific working + # memory still must remain free for activations and temporary tensors. + if self._max_vram_cache_size_gb is not None: + vram_total_available_to_cache = int(self._max_vram_cache_size_gb * GB) - working_mem_bytes + return vram_total_available_to_cache - self._get_vram_in_use() + if self._execution_device.type == "cuda": # TODO(ryand): It is debatable whether we should use memory_reserved() or memory_allocated() here. # memory_reserved() includes memory reserved by the torch CUDA memory allocator that may or may not be diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index fbd1c0e1280..302ef1376f4 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -132,6 +132,7 @@ def estimate_vae_working_memory_wan( pixel_width: int, pixel_frames: int, tile_size: int | None = None, + streaming: bool = False, ) -> int: """Estimate the working memory required to encode or decode with a Wan VAE. @@ -154,12 +155,17 @@ def estimate_vae_working_memory_wan( else: per_frame = pixel_height * pixel_width * element_size * scaling_constant - # The full RGB clip stays on the execution device regardless of tiling (decode - # output / encode input). Decode accumulates frames with torch.cat, whose final - # iterations transiently hold both the accumulated clip and its copy — ~2x the - # clip bytes at peak. Encode consumes the input clip without duplicating it. - clip_copies = 2 if operation == "decode" else 1 - clip_bytes = clip_copies * 3 * pixel_frames * pixel_height * pixel_width * element_size + # Streaming decode moves each causal decoder chunk to CPU immediately. Only one + # temporal-upscale chunk remains on the execution device, instead of the full RGB + # clip plus the transient copy created by torch.cat. + if operation == "decode" and streaming: + temporal_scale = int(getattr(vae.config, "scale_factor_temporal", None) or 4) + resident_frames = min(pixel_frames, temporal_scale) + clip_copies = 1 + else: + resident_frames = pixel_frames + clip_copies = 2 if operation == "decode" else 1 + clip_bytes = clip_copies * 3 * resident_frames * pixel_height * pixel_width * element_size return int(per_frame + clip_bytes) diff --git a/invokeai/backend/wan/memory_optimization.py b/invokeai/backend/wan/memory_optimization.py new file mode 100644 index 00000000000..1b3d68683be --- /dev/null +++ b/invokeai/backend/wan/memory_optimization.py @@ -0,0 +1,253 @@ +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from types import MethodType +from typing import Any + +import torch +from diffusers.models.modeling_outputs import Transformer2DModelOutput + +WAN_ACTIVATION_CHUNK_SIZE = 1024 + + +@dataclass +class _CompactTimestepConditioning: + timestep_embeddings: torch.Tensor + modulation: torch.Tensor + indices: torch.Tensor + + +def _get_modulation( + block: torch.nn.Module, + temb: torch.Tensor | _CompactTimestepConditioning, + start: int, + end: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if isinstance(temb, _CompactTimestepConditioning): + modulation = temb.modulation[temb.indices[:, start:end]].to(device=device, dtype=torch.float32) + modulation = block.scale_shift_table.unsqueeze(0).to(device) + modulation # type: ignore[attr-defined] + chunks = modulation.chunk(6, dim=2) + return tuple(chunk.squeeze(2) for chunk in chunks) # type: ignore[return-value] + if temb.ndim == 4: + modulation = block.scale_shift_table.unsqueeze(0).to(device) + temb[:, start:end].to( # type: ignore[attr-defined] + device=device, dtype=torch.float32 + ) + chunks = modulation.chunk(6, dim=2) + return tuple(chunk.squeeze(2) for chunk in chunks) # type: ignore[return-value] + + modulation = block.scale_shift_table.to(device) + temb.to(device=device, dtype=torch.float32) # type: ignore[attr-defined] + return modulation.chunk(6, dim=1) # type: ignore[return-value] + + +def _optimized_wan_transformer_forward( + transformer: torch.nn.Module, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_image: torch.Tensor | None = None, + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, +) -> Any: + original_forward = transformer._invokeai_original_forward # type: ignore[attr-defined] + # The custom transformer path is only needed to compact TI2V's per-token + # timesteps. Keep Diffusers' decorated forward for scalar timesteps and + # non-default attention kwargs. + if torch.is_grad_enabled() or timestep.ndim != 2 or attention_kwargs: + return original_forward( + hidden_states=hidden_states, + timestep=timestep, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_image=encoder_hidden_states_image, + return_dict=return_dict, + attention_kwargs=attention_kwargs, + ) + + batch_size, _, num_frames, height, width = hidden_states.shape + patch_frames, patch_height, patch_width = transformer.config.patch_size # type: ignore[attr-defined] + post_patch_num_frames = num_frames // patch_frames + post_patch_height = height // patch_height + post_patch_width = width // patch_width + + rotary_emb = transformer.rope(hidden_states) # type: ignore[attr-defined] + hidden_states = transformer.patch_embedding(hidden_states).flatten(2).transpose(1, 2) # type: ignore[attr-defined] + + unique_timesteps, inverse_indices = torch.unique(timestep, sorted=False, return_inverse=True) + temb, timestep_projection, encoder_hidden_states, encoder_hidden_states_image = transformer.condition_embedder( # type: ignore[attr-defined] + unique_timesteps, + encoder_hidden_states, + encoder_hidden_states_image, + timestep_seq_len=None, + ) + compact_timestep = _CompactTimestepConditioning( + timestep_embeddings=temb, + modulation=timestep_projection.unflatten(1, (6, -1)), + indices=inverse_indices.view_as(timestep), + ) + if encoder_hidden_states_image is not None: + encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) + + for block in transformer.blocks: # type: ignore[attr-defined] + hidden_states = block(hidden_states, encoder_hidden_states, compact_timestep, rotary_emb) + + sequence_length = hidden_states.shape[1] + chunk_size: int = transformer._invokeai_activation_chunk_size # type: ignore[attr-defined] + projected = hidden_states.new_empty( + (batch_size, sequence_length, transformer.proj_out.out_features) # type: ignore[attr-defined] + ) + for start in range(0, sequence_length, chunk_size): + end = min(start + chunk_size, sequence_length) + timestep_chunk = compact_timestep.timestep_embeddings[compact_timestep.indices[:, start:end]].to( + hidden_states.device + ) + shift, scale = ( + transformer.scale_shift_table.unsqueeze(0).to(hidden_states.device) + timestep_chunk.unsqueeze(2) # type: ignore[attr-defined] + ).chunk(2, dim=2) + normalized_chunk = ( + transformer.norm_out(hidden_states[:, start:end].float()) * (1 + scale.squeeze(2)) # type: ignore[attr-defined] + + shift.squeeze(2) + ).type_as(hidden_states) + projected[:, start:end].copy_(transformer.proj_out(normalized_chunk)) # type: ignore[attr-defined] + hidden_states = projected + + hidden_states = hidden_states.reshape( + batch_size, + post_patch_num_frames, + post_patch_height, + post_patch_width, + patch_frames, + patch_height, + patch_width, + -1, + ) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) + + +def _optimized_wan_block_forward( + block: torch.nn.Module, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor | _CompactTimestepConditioning, + rotary_emb: torch.Tensor, +) -> torch.Tensor: + original_forward = block._invokeai_original_forward # type: ignore[attr-defined] + chunk_size: int = block._invokeai_activation_chunk_size # type: ignore[attr-defined] + sequence_length = hidden_states.shape[1] + + if torch.is_grad_enabled() or ( + sequence_length <= chunk_size and not isinstance(temb, _CompactTimestepConditioning) + ): + return original_forward(hidden_states, encoder_hidden_states, temb, rotary_emb) + + normalized = torch.empty_like(hidden_states) + for start in range(0, sequence_length, chunk_size): + end = min(start + chunk_size, sequence_length) + shift_msa, scale_msa, _, _, _, _ = _get_modulation(block, temb, start, end, hidden_states.device) + normalized_chunk = ( + block.norm1(hidden_states[:, start:end].float()) * (1 + scale_msa) + shift_msa # type: ignore[attr-defined] + ).type_as(hidden_states) + normalized[:, start:end].copy_(normalized_chunk) + + attention_output = block.attn1(normalized, None, None, rotary_emb) # type: ignore[attr-defined] + del normalized + + updated = torch.empty_like(hidden_states) + for start in range(0, sequence_length, chunk_size): + end = min(start + chunk_size, sequence_length) + _, _, gate_msa, _, _, _ = _get_modulation(block, temb, start, end, hidden_states.device) + updated_chunk = (hidden_states[:, start:end].float() + attention_output[:, start:end] * gate_msa).type_as( + hidden_states + ) + updated[:, start:end].copy_(updated_chunk) + hidden_states = updated + del attention_output + + normalized = torch.empty_like(hidden_states) + for start in range(0, sequence_length, chunk_size): + end = min(start + chunk_size, sequence_length) + normalized[:, start:end].copy_( + block.norm2(hidden_states[:, start:end].float()).type_as(hidden_states) # type: ignore[attr-defined] + ) + attention_output = block.attn2(normalized, encoder_hidden_states, None, None) # type: ignore[attr-defined] + hidden_states = hidden_states + attention_output + del normalized, attention_output + + output = torch.empty_like(hidden_states) + for start in range(0, sequence_length, chunk_size): + end = min(start + chunk_size, sequence_length) + _, _, _, shift_mlp, scale_mlp, gate_mlp = _get_modulation(block, temb, start, end, hidden_states.device) + normalized_chunk = ( + block.norm3(hidden_states[:, start:end].float()) * (1 + scale_mlp) + shift_mlp # type: ignore[attr-defined] + ).type_as(hidden_states) + feed_forward_output = block.ffn(normalized_chunk) # type: ignore[attr-defined] + output_chunk = (hidden_states[:, start:end].float() + feed_forward_output.float() * gate_mlp).type_as( + hidden_states + ) + output[:, start:end].copy_(output_chunk) + + return output + + +@contextmanager +def wan_memory_optimization( + transformer: torch.nn.Module, + *, + enabled: bool, + activation_chunk_size: int = WAN_ACTIVATION_CHUNK_SIZE, +) -> Iterator[None]: + """Temporarily chunk Wan transformer pointwise activations during inference.""" + if not enabled: + yield + return + if activation_chunk_size <= 0: + raise ValueError("activation_chunk_size must be positive") + + blocks: Any = getattr(transformer, "blocks", None) + if blocks is None: + raise TypeError(f"Expected a Wan transformer with blocks, got {type(transformer).__name__}.") + blocks = list(blocks) + if hasattr(transformer, "_invokeai_original_forward") or any( + hasattr(block, "_invokeai_original_forward") for block in blocks + ): + raise RuntimeError("Wan memory optimization context cannot be nested.") + + patched_blocks: list[tuple[torch.nn.Module, Any, bool]] = [] + original_transformer_forward = transformer.forward + transformer_had_instance_forward = "forward" in transformer.__dict__ + patch_transformer_forward = all( + hasattr(transformer, name) + for name in ("condition_embedder", "patch_embedding", "proj_out", "rope", "scale_shift_table") + ) + try: + if patch_transformer_forward: + transformer._invokeai_original_forward = original_transformer_forward + transformer._invokeai_activation_chunk_size = activation_chunk_size + transformer.forward = MethodType(_optimized_wan_transformer_forward, transformer) + for block in blocks: + original_forward = block.forward + had_instance_forward = "forward" in block.__dict__ + block._invokeai_original_forward = original_forward + block._invokeai_activation_chunk_size = activation_chunk_size + block.forward = MethodType(_optimized_wan_block_forward, block) + patched_blocks.append((block, original_forward, had_instance_forward)) + yield + finally: + for block, original_forward, had_instance_forward in patched_blocks: + if had_instance_forward: + block.forward = original_forward + else: + del block.forward + del block._invokeai_original_forward + del block._invokeai_activation_chunk_size + if patch_transformer_forward: + if transformer_had_instance_forward: + transformer.forward = original_transformer_forward + else: + del transformer.forward + del transformer._invokeai_original_forward + del transformer._invokeai_activation_chunk_size diff --git a/invokeai/backend/wan/vae_decode.py b/invokeai/backend/wan/vae_decode.py new file mode 100644 index 00000000000..dbdf956199f --- /dev/null +++ b/invokeai/backend/wan/vae_decode.py @@ -0,0 +1,31 @@ +from collections.abc import Iterator + +import torch +from diffusers.models.autoencoders import AutoencoderKLWan +from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify + + +def iter_wan_vae_decode_chunks(vae: AutoencoderKLWan, latents: torch.Tensor) -> Iterator[torch.Tensor]: + """Decode one latent frame at a time while preserving Wan causal-convolution state.""" + _, _, num_frames, height, width = latents.shape + tile_latent_min_height = vae.tile_sample_min_height // vae.spatial_compression_ratio + tile_latent_min_width = vae.tile_sample_min_width // vae.spatial_compression_ratio + 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.") + + vae.clear_cache() + try: + hidden_states = vae.post_quant_conv(latents) + for frame_index in range(num_frames): + vae._conv_idx = [0] + decoded = vae.decoder( + hidden_states[:, :, frame_index : frame_index + 1], + feat_cache=vae._feat_map, + feat_idx=vae._conv_idx, + first_chunk=frame_index == 0, + ) + if vae.config.patch_size is not None: + decoded = unpatchify(decoded, patch_size=vae.config.patch_size) + yield decoded.clamp(-1.0, 1.0) + finally: + vae.clear_cache() diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 0499e6f426d..9b5f3ff772e 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -49027,6 +49027,12 @@ "description": "Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.", "default": false }, + "wan_memory_optimization": { + "type": "boolean", + "title": "Wan Memory Optimization", + "description": "Enable experimental Wan memory optimizations at the cost of slower generation.", + "default": false + }, "attention_type": { "type": "string", "enum": ["auto", "normal", "xformers", "sliced", "torch-sdp"], @@ -49291,7 +49297,7 @@ "additionalProperties": false, "type": "object", "title": "InvokeAIAppConfig", - "description": "Invoke's global app configuration.\n\nTypically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.\n\nAttributes:\n host: IP address to bind to. Use `0.0.0.0` to serve to your local network.\n port: Port to bind to.\n allow_origins: Allowed CORS origins.\n allow_credentials: Allow CORS credentials.\n allow_methods: Methods allowed for CORS.\n allow_headers: Headers allowed for CORS.\n ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n log_tokenization: Enable logging of parsed prompt tokens.\n patchmatch: Enable patchmatch inpaint code.\n models_dir: Path to the models directory.\n convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).\n download_cache_dir: Path to the directory that contains dynamically downloaded models.\n legacy_conf_dir: Path to directory of legacy checkpoint config files.\n db_dir: Path to InvokeAI databases directory.\n outputs_dir: Path to directory for outputs.\n image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash`\n custom_nodes_dir: Path to directory for custom nodes.\n style_presets_dir: Path to directory for style presets.\n workflow_thumbnails_dir: Path to directory for workflow thumbnails.\n log_handlers: Log handler. Valid options are \"console\", \"file=\", \"syslog=path|address:host:port\", \"http=\".\n log_format: Log format. Use \"plain\" for text-only, \"color\" for colorized output, \"legacy\" for 2.3-style logging and \"syslog\" for syslog-style.
Valid values: `plain`, `color`, `syslog`, `legacy`\n log_level: Emit logging messages at this level or higher.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.\n log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n use_memory_db: Use in-memory database. Useful for development.\n dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.\n profile_graphs: Enable graph profiling using `cProfile`.\n profile_prefix: An optional prefix for profile output files.\n profiles_dir: Path to profiles output directory.\n max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.\n max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.\n log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.\n model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.\n device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.\n enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.\n keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.\n ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.\n pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.\n device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `cuda:N` (where N is a device number)\n precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.
Valid values: `auto`, `float16`, `bfloat16`, `float32`\n sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.\n attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`\n attention_slice_size: Slice size, valid when attention_type==\"sliced\".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`\n force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).\n pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.\n max_queue_size: Maximum number of items in the session queue.\n session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin`\n clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.\n max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.\n allow_nodes: List of nodes to allow. Omit to allow all.\n deny_nodes: List of nodes to deny. Omit to deny none.\n node_cache_size: How many cached nodes to keep in memory.\n hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.
Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`\n remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.\n scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.\n unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.\n allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.\n multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.\n strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.\n external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.\n external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.\n external_gemini_api_key: API key for Gemini image generation.\n external_openai_api_key: API key for OpenAI image generation.\n external_gemini_base_url: Base URL override for Gemini image generation.\n external_openai_base_url: Base URL override for OpenAI image generation.\n external_seedream_api_key: API key for Seedream image generation.\n external_seedream_base_url: Base URL override for Seedream image generation.\n base_url: Public base path when running behind a reverse proxy under a sub-path, e.g. `/invoke`. Set only when the proxy PRESERVES the sub-path (the backend receives `/invoke/api/...`). Leave unset when the proxy strips the sub-path or when serving at the domain root.\n forwarded_allow_ips: Comma-separated list of IPs (or `*`) allowed to set X-Forwarded-* headers. Set to the reverse proxy's IP. Only used when `base_url` is set." + "description": "Invoke's global app configuration.\n\nTypically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.\n\nAttributes:\n host: IP address to bind to. Use `0.0.0.0` to serve to your local network.\n port: Port to bind to.\n allow_origins: Allowed CORS origins.\n allow_credentials: Allow CORS credentials.\n allow_methods: Methods allowed for CORS.\n allow_headers: Headers allowed for CORS.\n ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n log_tokenization: Enable logging of parsed prompt tokens.\n patchmatch: Enable patchmatch inpaint code.\n models_dir: Path to the models directory.\n convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).\n download_cache_dir: Path to the directory that contains dynamically downloaded models.\n legacy_conf_dir: Path to directory of legacy checkpoint config files.\n db_dir: Path to InvokeAI databases directory.\n outputs_dir: Path to directory for outputs.\n image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash`\n custom_nodes_dir: Path to directory for custom nodes.\n style_presets_dir: Path to directory for style presets.\n workflow_thumbnails_dir: Path to directory for workflow thumbnails.\n log_handlers: Log handler. Valid options are \"console\", \"file=\", \"syslog=path|address:host:port\", \"http=\".\n log_format: Log format. Use \"plain\" for text-only, \"color\" for colorized output, \"legacy\" for 2.3-style logging and \"syslog\" for syslog-style.
Valid values: `plain`, `color`, `syslog`, `legacy`\n log_level: Emit logging messages at this level or higher.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.\n log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n use_memory_db: Use in-memory database. Useful for development.\n dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.\n profile_graphs: Enable graph profiling using `cProfile`.\n profile_prefix: An optional prefix for profile output files.\n profiles_dir: Path to profiles output directory.\n max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.\n max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.\n log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.\n model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.\n device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.\n enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.\n keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.\n ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.\n pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.\n device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `cuda:N` (where N is a device number)\n precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.
Valid values: `auto`, `float16`, `bfloat16`, `float32`\n sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.\n wan_memory_optimization: Enable experimental Wan memory optimizations at the cost of slower generation.\n attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`\n attention_slice_size: Slice size, valid when attention_type==\"sliced\".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`\n force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).\n pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.\n max_queue_size: Maximum number of items in the session queue.\n session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin`\n clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.\n max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.\n allow_nodes: List of nodes to allow. Omit to allow all.\n deny_nodes: List of nodes to deny. Omit to deny none.\n node_cache_size: How many cached nodes to keep in memory.\n hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.
Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`\n remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.\n scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.\n unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.\n allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.\n multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.\n strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.\n external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.\n external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.\n external_gemini_api_key: API key for Gemini image generation.\n external_openai_api_key: API key for OpenAI image generation.\n external_gemini_base_url: Base URL override for Gemini image generation.\n external_openai_base_url: Base URL override for OpenAI image generation.\n external_seedream_api_key: API key for Seedream image generation.\n external_seedream_base_url: Base URL override for Seedream image generation.\n base_url: Public base path when running behind a reverse proxy under a sub-path, e.g. `/invoke`. Set only when the proxy PRESERVES the sub-path (the backend receives `/invoke/api/...`). Leave unset when the proxy strips the sub-path or when serving at the domain root.\n forwarded_allow_ips: Comma-separated list of IPs (or `*`) allowed to set X-Forwarded-* headers. Set to the reverse proxy's IP. Only used when `base_url` is set." }, "InvokeAIAppConfigWithSetFields": { "properties": { diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 2928b7170d6..45ccb943e24 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -18918,6 +18918,7 @@ export type components = { * device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `cuda:N` (where N is a device number) * precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.
Valid values: `auto`, `float16`, `bfloat16`, `float32` * sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements. + * wan_memory_optimization: Enable experimental Wan memory optimizations at the cost of slower generation. * attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp` * attention_slice_size: Slice size, valid when attention_type=="sliced".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8` * force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). @@ -19259,6 +19260,12 @@ export type components = { * @default false */ sequential_guidance?: boolean; + /** + * Wan Memory Optimization + * @description Enable experimental Wan memory optimizations at the cost of slower generation. + * @default false + */ + wan_memory_optimization?: boolean; /** * Attention Type * @description Attention type. diff --git a/tests/app/invocations/test_wan_denoise.py b/tests/app/invocations/test_wan_denoise.py index 4dbde079db1..54d8a656d8c 100644 --- a/tests/app/invocations/test_wan_denoise.py +++ b/tests/app/invocations/test_wan_denoise.py @@ -22,7 +22,11 @@ from invokeai.app.invocations.fields import ImageField, LatentsField, WanConditioningField, WanRefImageConditioningField from invokeai.app.invocations.model import ModelIdentifierField, VAEField, WanTransformerField -from invokeai.app.invocations.wan_denoise import WanDenoiseInvocation +from invokeai.app.invocations.wan_denoise import ( + WanDenoiseInvocation, + _ExpertSwapper, + _get_wan_transformer_working_mem_bytes, +) from invokeai.app.invocations.wan_ref_image_encoder import WanRefImageEncoderInvocation from invokeai.app.invocations.wan_video_denoise import WanVideoDenoiseInvocation from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType, WanVariantType @@ -134,6 +138,7 @@ def _load_conditioning(name: str) -> ConditioningFieldData: context.util.signal_progress = MagicMock() context.util.sd_step_callback = MagicMock() context.logger = MagicMock() + context.config.get.return_value.wan_memory_optimization = False return context @@ -259,6 +264,51 @@ def test_run_diffusion_returns_4d_finite( # Step callback invoked once per step. assert ctx.util.sd_step_callback.call_count == 4 + def test_memory_optimization_reserves_vram_for_streamed_transformer_weights(self, monkeypatch) -> None: + total_vram = 24 * 2**30 + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda _device: MagicMock(total_memory=total_vram), + ) + + working_mem_bytes = _get_wan_transformer_working_mem_bytes(torch.device("cuda"), enabled=True) + + assert working_mem_bytes == 22 * 2**30 + + def test_memory_optimization_does_not_change_cpu_or_disabled_loading(self) -> None: + assert _get_wan_transformer_working_mem_bytes(torch.device("cpu"), enabled=True) is None + assert _get_wan_transformer_working_mem_bytes(torch.device("cuda"), enabled=False) is None + + def test_expert_swapper_passes_aggressive_working_memory_to_model_cache(self) -> None: + transformer = _ZeroTransformer() + loaded = MagicMock() + cached_model = MagicMock() + cached_model.cur_vram_bytes.return_value = 5 * 2**30 + loaded._cache_record.cached_model = cached_model + device_context = MagicMock() + device_context.__enter__.return_value = (None, transformer) + loaded.model_on_device.return_value = device_context + context = MagicMock() + context.models.load.return_value = loaded + working_mem_bytes = 22 * 2**30 + swapper = _ExpertSwapper( + context=context, + high_model=MagicMock(), + low_model=None, + inference_dtype=torch.bfloat16, + working_mem_bytes=working_mem_bytes, + max_resident_model_bytes=2 * 2**30, + ) + + try: + assert swapper.get(_ExpertSwapper.HIGH) is transformer + finally: + swapper.close() + + loaded.model_on_device.assert_called_once_with(working_mem_bytes=working_mem_bytes) + loaded.unload_from_vram.assert_called_once_with(3 * 2**30, keep_required_weights_in_vram=True) + def test_cfg_doubles_transformer_calls(self, fake_model_root) -> None: """With cfg_scale != 1.0 and a negative prompt, each step runs the model twice.""" transformer = _ZeroTransformer() @@ -286,6 +336,38 @@ def test_cfg_doubles_transformer_calls(self, fake_model_root) -> None: # 3 steps × 2 (cond + uncond) = 6 forward calls. assert len(transformer.calls) == 6 + def test_memory_optimization_config_wraps_each_active_expert_step(self, fake_model_root, monkeypatch) -> None: + transformer = _ZeroTransformer() + ctx = _build_context( + transformer, + variant=WanVariantType.T2V_A14B, + model_root=fake_model_root, + pos_cond=_make_conditioning(), + neg_cond=None, + ) + ctx.config.get.return_value.wan_memory_optimization = True + enabled_calls: list[bool] = [] + + @contextmanager + def record_memory_optimization(_transformer, *, enabled: bool): + enabled_calls.append(enabled) + yield + + monkeypatch.setattr("invokeai.app.invocations.wan_denoise.wan_memory_optimization", record_memory_optimization) + inv = _make_invocation( + transformer_field=_wan_transformer_field(), + pos_field=WanConditioningField(conditioning_name="pos"), + neg_field=None, + width=64, + height=64, + steps=3, + guidance_scale=1.0, + ) + + inv._run_diffusion(ctx) + + assert enabled_calls == [True, True, True] + def test_ti2v_image_rejects_dimensions_not_divisible_by_32(self, fake_model_root: Path) -> None: context = _build_context( _ZeroTransformer(), diff --git a/tests/app/invocations/test_wan_expert_swapper.py b/tests/app/invocations/test_wan_expert_swapper.py index b5f910a9ac2..977dfc1a8a5 100644 --- a/tests/app/invocations/test_wan_expert_swapper.py +++ b/tests/app/invocations/test_wan_expert_swapper.py @@ -58,6 +58,9 @@ def full_unload_from_vram(self) -> int: self.unload_calls += 1 return 0 + def total_bytes(self) -> int: + return 0 + class _FakeCacheRecord: def __init__(self, cached_model: _FakeCachedModel) -> None: @@ -77,6 +80,10 @@ def __init__(self, label: str, model: nn.Module, log: list[str]) -> None: def model_on_device(self): return _FakeModelOnDevice(self._label, self._model, self._log) + def unload_from_vram(self, _vram_bytes_to_free, keep_required_weights_in_vram=False): + assert keep_required_weights_in_vram is False + return self._cache_record.cached_model.full_unload_from_vram() + class _FakeContext: """Mocks ``InvocationContext.models.load`` returning a fresh ``_FakeInfo`` @@ -442,13 +449,13 @@ def test_empty_cache_called_on_swap(): def test_outgoing_expert_force_unloaded_from_vram(): """Regression: on swap, the previous expert's weights must be explicitly forced - off VRAM via ``cached_model.full_unload_from_vram()``. + off VRAM via the loaded model's cache-routed unload API. A14B users observed the high-noise transformer continuing to occupy ~9 GB of VRAM during the low-noise step, because the cache's automatic offload heuristic underestimated how much room the new expert needed when workspace memory from the previous denoise step was still allocated. The swapper sidesteps that by - invoking full_unload_from_vram on the outgoing expert directly.""" + invoking the cache-routed unload on the outgoing expert directly.""" log: list[str] = [] high_info = _FakeInfo("HIGH", nn.Linear(1, 1), log) low_info = _FakeInfo("LOW", nn.Linear(1, 1), log) diff --git a/tests/app/invocations/test_wan_working_memory.py b/tests/app/invocations/test_wan_working_memory.py index 5a84df10d6f..30bb585cc2d 100644 --- a/tests/app/invocations/test_wan_working_memory.py +++ b/tests/app/invocations/test_wan_working_memory.py @@ -78,6 +78,27 @@ def test_additional_frames_add_only_clip_bytes(self): ) assert many - one == 2 * 3 * 80 * 128 * 128 * 2 + def test_streaming_decode_bounds_resident_clip_to_one_temporal_chunk(self): + vae = _mock_wan_vae(temporal_scale=4) + one = estimate_vae_working_memory_wan( + operation="decode", + vae=vae, + pixel_height=128, + pixel_width=128, + pixel_frames=1, + streaming=True, + ) + many = estimate_vae_working_memory_wan( + operation="decode", + vae=vae, + pixel_height=128, + pixel_width=128, + pixel_frames=81, + streaming=True, + ) + + assert many - one == 3 * 3 * 128 * 128 * 2 + def test_tile_size_bounds_the_per_frame_term(self): vae = _mock_wan_vae() tiled = estimate_vae_working_memory_wan( @@ -187,6 +208,7 @@ def _video_context(self, vae_info: MagicMock, t_lat: int = 5) -> MagicMock: mock_context.models.load.return_value = vae_info mock_context.tensors.load.return_value = torch.zeros(1, 16, t_lat, 32, 32) mock_context.util.is_canceled.return_value = False + mock_context.config.get.return_value.wan_memory_optimization = False return mock_context def test_latents_to_video_requests_decode_memory_for_all_frames(self): @@ -210,6 +232,46 @@ def test_latents_to_video_requests_decode_memory_for_all_frames(self): assert mock_estimate.call_args.kwargs["pixel_height"] == 256 vae_info.model_on_device.assert_called_once_with(working_mem_bytes=5678) + def test_latents_to_video_streams_decode_chunks_directly_to_mp4(self): + vae = _mock_wan_vae() + vae.use_tiling = True # Simulate a prior Anima tiled decode on the shared VAE. + vae_info = _mock_vae_info(vae) + mock_context = self._video_context(vae_info, t_lat=2) + mock_context.config.get.return_value.wan_memory_optimization = True + writer = MagicMock() + chunks = [ + torch.zeros(1, 3, 1, 256, 256), + torch.zeros(1, 3, 4, 256, 256), + ] + expected_output = MagicMock() + + with ( + patch( + "invokeai.app.invocations.wan_latents_to_video.estimate_vae_working_memory_wan", + return_value=5678, + ) as mock_estimate, + patch( + "invokeai.app.invocations.wan_latents_to_video.iter_wan_vae_decode_chunks", + return_value=iter(chunks), + ) as mock_decode_chunks, + patch("invokeai.app.invocations.wan_latents_to_video.make_mp4_writer", return_value=writer), + patch("invokeai.app.invocations.wan_latents_to_video.VideoOutput.build", return_value=expected_output), + patch.object(TorchDevice, "choose_torch_device", return_value=torch.device("cpu")), + patch.object(TorchDevice, "empty_cache"), + ): + invocation = WanLatentsToVideoInvocation.model_construct( + latents=MagicMock(latents_name="l"), vae=MagicMock(vae=MagicMock()), fps=16 + ) + actual_output = invocation.invoke(mock_context) + + assert actual_output is expected_output + assert mock_estimate.call_args.kwargs["streaming"] is True + mock_decode_chunks.assert_called_once() + assert writer.append_data.call_count == 5 + writer.close.assert_called_once() + vae.disable_tiling.assert_called_once() + vae.decode.assert_not_called() + def test_latents_to_video_falls_back_to_tiling_when_estimate_exceeds_vram(self): vae = _mock_wan_vae() vae_info = _mock_vae_info(vae) diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py b/tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py index d57e1970d0a..b8c553c9154 100644 --- a/tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py @@ -270,6 +270,45 @@ def test_get_vram_in_use_queries_this_caches_execution_device(mock_logger): cache.shutdown() +def test_max_vram_cache_reserves_per_operation_working_memory(mock_logger): + """An explicit cache cap must still leave room for decode/diffusion activations. + + The capped-branch arithmetic is device-independent, so use a CPU cache to keep this test + runnable on CI hosts without a CUDA driver. + """ + cache = ModelCache( + execution_device_working_mem_gb=3.0, + enable_partial_loading=True, + keep_ram_copy_of_weights=True, + max_vram_cache_size_gb=16.0, + execution_device="cpu", + storage_device="cpu", + logger=mock_logger, + ) + try: + with patch.object(cache, "_get_vram_in_use", return_value=2 * GB): + assert cache._get_vram_available(working_mem_bytes=5 * GB) == 9 * GB + finally: + cache.shutdown() + + +def test_loaded_model_unload_drops_cache_entry_when_move_fails(mock_logger): + """A failed expert trim must not leave a half-moved model cache hit.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + try: + cache.put("broken", DummyModule()) + record = cache.get("broken") + loaded_model = LoadedModelWithoutConfig(cache_record=record, cache=cache) + with patch.object(record.cached_model, "full_unload_from_vram", side_effect=RuntimeError("move failed")): + with pytest.raises(RuntimeError, match="move failed"): + loaded_model.unload_from_vram(1) + assert "broken" not in cache._cached_models + finally: + cache.shutdown() + + def test_cuda_cache_init_queries_total_vram_without_mem_get_info(mock_logger): """CUDA cache sizing must not call the VRAM-holding mem_get_info API during idle startup.""" import torch diff --git a/tests/backend/wan/test_memory_optimization.py b/tests/backend/wan/test_memory_optimization.py new file mode 100644 index 00000000000..88fb97676c0 --- /dev/null +++ b/tests/backend/wan/test_memory_optimization.py @@ -0,0 +1,203 @@ +import math +from contextlib import nullcontext + +import pytest +import torch +from diffusers.models.transformers.transformer_wan import WanTransformer3DModel, WanTransformerBlock + +from invokeai.backend.wan.memory_optimization import wan_memory_optimization + + +def _build_block() -> WanTransformerBlock: + return WanTransformerBlock( + dim=8, + ffn_dim=16, + num_heads=2, + cross_attn_norm=True, + ).eval() + + +class _Transformer(torch.nn.Module): + def __init__(self, block: WanTransformerBlock) -> None: + super().__init__() + self.blocks = torch.nn.ModuleList([block]) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + ) -> torch.Tensor: + return self.blocks[0](hidden_states, encoder_hidden_states, temb, rotary_emb=None) + + +@pytest.mark.parametrize("per_token_timestep", [False, True]) +@pytest.mark.parametrize("use_autocast", [False, True]) +def test_wan_memory_optimization_matches_original_and_bounds_ffn_sequence( + per_token_timestep: bool, use_autocast: bool +) -> None: + torch.manual_seed(0) + original = _Transformer(_build_block()) + optimized = _Transformer(_build_block()) + optimized.load_state_dict(original.state_dict()) + + batch_size = 2 + sequence_length = 7 + hidden_states = torch.randn(batch_size, sequence_length, 8) + encoder_hidden_states = torch.randn(batch_size, 5, 8) + if per_token_timestep: + temb = torch.randn(batch_size, sequence_length, 6, 8) + else: + temb = torch.randn(batch_size, 6, 8) + + chunk_size = 3 + ffn_sequence_lengths: list[int] = [] + + def record_ffn_sequence_length(_module: torch.nn.Module, inputs: tuple[torch.Tensor, ...]) -> None: + ffn_sequence_lengths.append(inputs[0].shape[1]) + + handle = optimized.blocks[0].ffn.register_forward_pre_hook(record_ffn_sequence_length) + try: + autocast_context = torch.autocast("cpu", dtype=torch.bfloat16) if use_autocast else nullcontext() + with torch.no_grad(), autocast_context: + expected = original(hidden_states, encoder_hidden_states, temb) + with wan_memory_optimization(optimized, enabled=True, activation_chunk_size=chunk_size): + actual = optimized(hidden_states, encoder_hidden_states, temb) + finally: + handle.remove() + + torch.testing.assert_close(actual, expected) + assert max(ffn_sequence_lengths) <= chunk_size + assert len(ffn_sequence_lengths) == math.ceil(sequence_length / chunk_size) + + +def test_wan_memory_optimization_is_not_sticky_between_calls() -> None: + transformer = _Transformer(_build_block()) + hidden_states = torch.randn(1, 5, 8) + encoder_hidden_states = torch.randn(1, 3, 8) + temb = torch.randn(1, 6, 8) + ffn_sequence_lengths: list[int] = [] + + def record_ffn_sequence_length(_module: torch.nn.Module, inputs: tuple[torch.Tensor, ...]) -> None: + ffn_sequence_lengths.append(inputs[0].shape[1]) + + handle = transformer.blocks[0].ffn.register_forward_pre_hook(record_ffn_sequence_length) + try: + with torch.no_grad(): + with wan_memory_optimization(transformer, enabled=True, activation_chunk_size=2): + transformer(hidden_states, encoder_hidden_states, temb) + optimized_call_count = len(ffn_sequence_lengths) + transformer(hidden_states, encoder_hidden_states, temb) + finally: + handle.remove() + + assert ffn_sequence_lengths[:optimized_call_count] == [2, 2, 1] + assert ffn_sequence_lengths[optimized_call_count:] == [hidden_states.shape[1]] + assert "forward" not in transformer.blocks[0].__dict__ + + +def test_wan_memory_optimization_restores_blocks_after_exception() -> None: + transformer = _Transformer(_build_block()) + original_forward = transformer.blocks[0].forward + + with pytest.raises(RuntimeError, match="boom"): + with wan_memory_optimization(transformer, enabled=True, activation_chunk_size=2): + raise RuntimeError("boom") + + assert transformer.blocks[0].forward == original_forward + + +def test_wan_memory_optimization_rejects_nesting_without_corrupting_outer_context() -> None: + transformer = _Transformer(_build_block()) + original_forward = transformer.blocks[0].forward + + with wan_memory_optimization(transformer, enabled=True, activation_chunk_size=2): + optimized_forward = transformer.blocks[0].forward + with pytest.raises(RuntimeError, match="cannot be nested"): + with wan_memory_optimization(transformer, enabled=True, activation_chunk_size=2): + pass + assert transformer.blocks[0].forward == optimized_forward + + assert transformer.blocks[0].forward == original_forward + + +def test_wan_memory_optimization_rejects_non_positive_chunk_size() -> None: + transformer = _Transformer(_build_block()) + + with pytest.raises(ValueError, match="activation_chunk_size must be positive"): + with wan_memory_optimization(transformer, enabled=True, activation_chunk_size=0): + pass + + +def test_wan_memory_optimization_uses_original_path_with_gradients() -> None: + transformer = _Transformer(_build_block()) + hidden_states = torch.randn(1, 5, 8, requires_grad=True) + encoder_hidden_states = torch.randn(1, 3, 8) + temb = torch.randn(1, 6, 8) + ffn_sequence_lengths: list[int] = [] + + def record_ffn_sequence_length(_module: torch.nn.Module, inputs: tuple[torch.Tensor, ...]) -> None: + ffn_sequence_lengths.append(inputs[0].shape[1]) + + handle = transformer.blocks[0].ffn.register_forward_pre_hook(record_ffn_sequence_length) + try: + with wan_memory_optimization(transformer, enabled=True, activation_chunk_size=2): + output = transformer(hidden_states, encoder_hidden_states, temb) + output.sum().backward() + finally: + handle.remove() + + assert ffn_sequence_lengths == [hidden_states.shape[1]] + assert hidden_states.grad is not None + + +def test_wan_memory_optimization_compacts_per_token_timesteps() -> None: + torch.manual_seed(0) + original = WanTransformer3DModel( + patch_size=(1, 2, 2), + num_attention_heads=2, + attention_head_dim=12, + in_channels=4, + out_channels=4, + text_dim=16, + freq_dim=8, + ffn_dim=32, + num_layers=1, + rope_max_seq_len=16, + ).eval() + optimized = WanTransformer3DModel.from_config(original.config).eval() + optimized.load_state_dict(original.state_dict()) + + hidden_states = torch.randn(1, 4, 3, 4, 4) + sequence_length = 3 * 2 * 2 + timestep = torch.full((1, sequence_length), 500.0) + timestep[:, :4] = 0 + encoder_hidden_states = torch.randn(1, 5, 16) + embedded_timestep_counts: list[int] = [] + + def record_timestep_count(_module: torch.nn.Module, inputs: tuple[torch.Tensor, ...]) -> None: + embedded_timestep_counts.append(inputs[0].shape[0]) + + handle = optimized.condition_embedder.time_embedder.register_forward_pre_hook(record_timestep_count) + try: + with torch.no_grad(): + expected = original( + hidden_states=hidden_states, + timestep=timestep, + encoder_hidden_states=encoder_hidden_states, + return_dict=False, + )[0] + # A chunk larger than the sequence verifies that compact conditioning + # remains valid when block activation chunking is not otherwise needed. + with wan_memory_optimization(optimized, enabled=True, activation_chunk_size=100): + actual = optimized( + hidden_states=hidden_states, + timestep=timestep, + encoder_hidden_states=encoder_hidden_states, + return_dict=False, + )[0] + finally: + handle.remove() + + torch.testing.assert_close(actual, expected) + assert embedded_timestep_counts == [2] diff --git a/tests/backend/wan/test_vae_decode.py b/tests/backend/wan/test_vae_decode.py new file mode 100644 index 00000000000..37441c3d7db --- /dev/null +++ b/tests/backend/wan/test_vae_decode.py @@ -0,0 +1,34 @@ +import torch +from diffusers.models.autoencoders import AutoencoderKLWan + +from invokeai.backend.wan.vae_decode import iter_wan_vae_decode_chunks + + +def _build_tiny_vae() -> AutoencoderKLWan: + return AutoencoderKLWan( + base_dim=2, + z_dim=2, + dim_mult=[1, 1], + num_res_blocks=1, + attn_scales=[], + temperal_downsample=[True], + latents_mean=[0.0, 0.0], + latents_std=[1.0, 1.0], + scale_factor_temporal=2, + scale_factor_spatial=2, + ).eval() + + +def test_iter_wan_vae_decode_chunks_matches_full_decode() -> None: + torch.manual_seed(0) + vae = _build_tiny_vae() + latents = torch.randn(1, 2, 3, 4, 4) + + with torch.inference_mode(): + expected = vae.decode(latents, return_dict=False)[0] + chunks = list(iter_wan_vae_decode_chunks(vae, latents)) + + actual = torch.cat(chunks, dim=2) + torch.testing.assert_close(actual, expected) + assert len(chunks) == latents.shape[2] + assert max(chunk.shape[2] for chunk in chunks) <= vae.config.scale_factor_temporal diff --git a/tests/test_config.py b/tests/test_config.py index 85eaaef5328..879d39bb2b5 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -80,6 +80,15 @@ def test_path_resolution_root_not_set(patch_rootdir: None): assert config.root_path == expected_root +def test_wan_memory_optimization_defaults_to_false_and_loads_from_yaml(tmp_path: Path, patch_rootdir: None) -> None: + assert InvokeAIAppConfig().wan_memory_optimization is False + + temp_config_file = tmp_path / "temp_invokeai.yaml" + temp_config_file.write_text('schema_version: "4.0.3"\nwan_memory_optimization: true\n') + + assert load_and_migrate_config(temp_config_file).wan_memory_optimization is True + + def test_read_config_from_file(tmp_path: Path, patch_rootdir: None): """Test reading configuration from a file.""" temp_config_file = tmp_path / "temp_invokeai.yaml"