diff --git a/docs/src/content/docs/configuration/low-vram-mode.mdx b/docs/src/content/docs/configuration/low-vram-mode.mdx index 8bf6b3320f9..f2b970eca6b 100644 --- a/docs/src/content/docs/configuration/low-vram-mode.mdx +++ b/docs/src/content/docs/configuration/low-vram-mode.mdx @@ -39,10 +39,27 @@ Low-VRAM mode and related workload-specific optimizations include: - Dynamic RAM and VRAM cache sizes (`max_cache_ram_gb`, `max_cache_vram_gb`) - Working memory (`device_working_mem_gb`) - Keeping a RAM weight copy (`keep_ram_copy_of_weights`) +- Wan video memory optimization (`wan_memory_optimization`) - PiD decode activation chunking (`pid_memory_optimization`) Read on to learn about these features and understand how to fine-tune them for your system and use-cases. +### Wan video memory optimization + +Wan video generation has an additional opt-in memory optimization: + +```yaml +wan_memory_optimization: true +``` + +This reserves VRAM so partial-load Wan transformer weights target about 2 GiB resident and stream remaining layers from RAM when partial model loading is enabled (the default). The explicit residency trim may be a no-op when cache admission already reaches that target. If `enable_partial_loading: false`, the activation, timestep, and VAE optimizations still apply, but transformer weights remain fully resident and the 2 GiB residency target is unavailable. It also chunks pointwise transformer activations, compacts TI2V per-token timestep conditioning, and streams untiled VAE decode chunks directly to MP4. It reduces peak VRAM during both denoise and decode, but generation can be substantially slower and requires enough system RAM for offloaded weights. Spatially tiled VAE decode continues to use its existing full-tile path. + +The optimized BF16 transformer path can produce small numerical differences because chunked matrix operations accumulate in a different order. Same-seed output is not guaranteed to be bit-identical; set `wan_memory_optimization: false` for the baseline path. + +Developers with a CUDA or ROCm device can validate Wan VAE memory estimates with `python scripts/calibrate_wan_vae_working_memory.py --vae `. The script reports allocated and reserved memory deltas; its implied scaling constant uses allocated memory to match the shipped estimator, while reserved memory shows allocator headroom. Use `--tiling` to measure the spatially tiled full-decode fallback; it overrides streaming mode. Use `--tile-size ` to override the VAE's default tile size. + +Direct MP4 streaming keeps the VAE cache lock while chunks are decoded and written. Wan's causal decoder state and weights must remain live for the sequence; releasing the lock would require a separate bounded decode and encode queue. + ### Partial model loading Invoke's partial model loading works by streaming model "layers" between RAM and VRAM as they are needed. @@ -102,7 +119,7 @@ max_cache_vram_gb: 16 ``` :::caution[Max safe value for `max_cache_vram_gb`] - Most users should not manually configure the `max_cache_vram_gb`. This configuration value takes precedence over the `device_working_mem_gb` and any operations that explicitly reserve additional working memory (e.g. VAE decode). As such, manually configuring it increases the likelihood of encountering out-of-memory errors. + Most users should not manually configure the `max_cache_vram_gb`. This configuration value caps model-cache residency; `device_working_mem_gb` and operation-specific reservations (e.g. VAE decode) are still subtracted from that cap for every model-cache operation, not only when Wan memory optimization is enabled. A cap below the active working-memory reservation can force aggressive model offloading. For users who wish to configure `max_cache_vram_gb`, the max safe value can be determined by subtracting `device_working_mem_gb` from your GPU's VRAM. As described below, the default for `device_working_mem_gb` is 3GB. diff --git a/docs/src/content/docs/features/video-generation.mdx b/docs/src/content/docs/features/video-generation.mdx index c8f205e2d44..ac2e2e38a56 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 target about 2 GiB of resident transformer weights when `enable_partial_loading` is enabled, lower denoise activation memory, and stream untiled VAE decode directly to MP4. The explicit residency trim may be a no-op when cache admission already reaches that target. If partial loading is disabled, the activation, timestep, and VAE optimizations remain active but transformer weights are fully resident. This can make generation substantially slower and requires enough system RAM for offloaded weights. Optimized BF16 execution may produce small numerical differences from the baseline path. + * **Drop resolution before frame count.** Going from 1280×720 to 832×480 is a ~2.4× memory drop and visually subtle in most content. Going from 81 frames to 65 only saves ~20%. * **TI2V-5B before A14B.** TI2V-5B Q4_K_M peaks around ~6–8 GB at 832×480, versus ~12–14 GB for A14B Q4_K_M. If you're at the OOM edge, switch model family. * **OOM at the *reference image encoder* step** is usually allocator fragmentation from a previous run rather than absolute memory pressure. Restart the dev server and try again; if it recurs reproducibly, file an issue. diff --git a/docs/src/generated/settings.json b/docs/src/generated/settings.json index 031badde195..9b221f16293 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": false, diff --git a/invokeai/app/invocations/wan_denoise.py b/invokeai/app/invocations/wan_denoise.py index 4a2c9c76c97..a7336ede9ed 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 partial-load Wan weights target about 2 GiB resident.""" + if not enabled or device.type != "cuda": + return None + + total_vram = torch.cuda.get_device_properties(device).total_memory + if total_vram <= WAN_MAX_RESIDENT_TRANSFORMER_BYTES: + return None + return total_vram - WAN_MAX_RESIDENT_TRANSFORMER_BYTES def _resolve_variant(context: InvocationContext, transformer_field: WanTransformerField) -> WanVariantType: @@ -176,6 +190,8 @@ def __init__( low_lora_factory: LoRAIteratorFactory | None = None, high_is_quantized: bool = False, low_is_quantized: bool = False, + working_mem_bytes: int | None = None, + max_resident_model_bytes: int | None = None, ) -> None: self._context = context self._high_model = high_model @@ -185,11 +201,14 @@ def __init__( self._low_lora_factory = low_lora_factory self._high_is_quantized = high_is_quantized self._low_is_quantized = low_is_quantized + self._working_mem_bytes = working_mem_bytes + self._max_resident_model_bytes = max_resident_model_bytes self._active_label: str | None = None self._active_info: Any | None = None self._active_device_ctx: Any | None = None self._active_lora_ctx: Any | None = None self._active_model: Any | None = None + self._warned_partial_loading_unavailable = False def get(self, label: str) -> Any: if label not in (self.HIGH, self.LOW): @@ -203,11 +222,10 @@ def get(self, label: str) -> Any: # Capture the outgoing expert's cache record before _release() drops our handle. # We need it to force-unload below. outgoing_cached_model = None + outgoing_info = self._active_info if self._active_info is not None: - # ``LoadedModel`` exposes its cache_record only via a private attribute. There - # is no public ``unload_from_vram`` on the LoadedModel today, and we don't want - # to take on a broader backend refactor in this fix; tolerate AttributeError - # so a future refactor doesn't break the swap. + # ``LoadedModel`` keeps the cache record private, but exposes + # ``unload_from_vram`` so cache error handling stays in one place. outgoing_cached_model = getattr(self._active_info, "_cache_record", None) if outgoing_cached_model is not None: outgoing_cached_model = getattr(outgoing_cached_model, "cached_model", None) @@ -229,7 +247,14 @@ def get(self, label: str) -> Any: # and now — the cached_model object still owns the tensors. if outgoing_cached_model is not None: try: - outgoing_cached_model.full_unload_from_vram() + unload_from_vram = getattr(outgoing_info, "unload_from_vram", None) + if callable(unload_from_vram): + unload_from_vram(outgoing_cached_model.total_bytes()) + else: + # Keep compatibility with old LoadedModel handles while preserving + # the process-global register_parameter guard. + with MODEL_LOAD_LOCK.read_lock(): + outgoing_cached_model.full_unload_from_vram() except Exception: pass @@ -242,7 +267,11 @@ def get(self, label: str) -> Any: # always fresh — see class docstring for the cache-eviction reasoning. model_id = self._high_model if label == self.HIGH else self._low_model info = self._context.models.load(model_id) - device_ctx = info.model_on_device() + supports_partial_loading = getattr(info, "supports_partial_loading", None) + if self._working_mem_bytes is None or supports_partial_loading is False: + device_ctx = info.model_on_device() + else: + device_ctx = info.model_on_device(working_mem_bytes=self._working_mem_bytes) cached_weights, model = device_ctx.__enter__() # Stash the device-context state immediately. If anything below fails (most @@ -256,6 +285,25 @@ def get(self, label: str) -> Any: self._active_device_ctx = device_ctx self._active_model = model + if self._max_resident_model_bytes is not None: + if supports_partial_loading is False: + if not self._warned_partial_loading_unavailable: + self._context.logger.warning( + "Wan memory optimization cannot limit resident transformer weights because " + "partial model loading is disabled." + ) + self._warned_partial_loading_unavailable = True + else: + cache_record = getattr(info, "_cache_record", None) + cached_model = getattr(cache_record, "cached_model", None) + cur_vram_bytes = getattr(cached_model, "cur_vram_bytes", None) + unload_from_vram = getattr(info, "unload_from_vram", None) + if callable(cur_vram_bytes) and callable(unload_from_vram): + vram_bytes_to_free = max(0, cur_vram_bytes() - self._max_resident_model_bytes) + if vram_bytes_to_free > 0: + unload_from_vram(vram_bytes_to_free, keep_required_weights_in_vram=True) + TorchDevice.empty_cache() + # Apply LoRA patches for this expert. GGUF transformers need sidecar # patching since direct patching of GGMLTensors isn't supported. lora_factory = self._high_lora_factory if label == self.HIGH else self._low_lora_factory @@ -601,6 +649,13 @@ def high_lora_factory() -> Iterable[PatchSpec]: def low_lora_factory() -> Iterable[PatchSpec]: return self._lora_iterator(context, low_loras) + optimize_memory = context.config.get().wan_memory_optimization + working_mem_bytes = _get_wan_transformer_working_mem_bytes(device, enabled=optimize_memory) + if working_mem_bytes is not None: + context.logger.info( + "Wan memory optimization: targeting about 2 GiB of resident transformer weights when partial " + "loading is available" + ) with ExitStack() as exit_stack: swapper = _ExpertSwapper( context=context, @@ -611,6 +666,10 @@ def low_lora_factory() -> Iterable[PatchSpec]: low_lora_factory=low_lora_factory if low_loras else None, high_is_quantized=high_is_quantized, low_is_quantized=low_is_quantized, + working_mem_bytes=working_mem_bytes, + max_resident_model_bytes=( + WAN_MAX_RESIDENT_TRANSFORMER_BYTES if working_mem_bytes is not None else None + ), ) exit_stack.callback(swapper.close) @@ -641,25 +700,26 @@ def low_lora_factory() -> Iterable[PatchSpec]: if ref_condition is not None: latent_model_input = torch.cat([latent_model_input, ref_condition], dim=1) - noise_pred_cond = transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=pos_cond.prompt_embeds.unsqueeze(0), - attention_kwargs=None, - return_dict=False, - )[0] - - if neg_cond is not None and active_cfg != 1.0: - noise_pred_uncond = transformer( + with wan_memory_optimization(transformer, enabled=optimize_memory): + noise_pred_cond = transformer( hidden_states=latent_model_input, timestep=timestep, - encoder_hidden_states=neg_cond.prompt_embeds.unsqueeze(0), + encoder_hidden_states=pos_cond.prompt_embeds.unsqueeze(0), attention_kwargs=None, return_dict=False, )[0] - noise_pred = noise_pred_uncond + active_cfg * (noise_pred_cond - noise_pred_uncond) - else: - noise_pred = noise_pred_cond + + if neg_cond is not None and active_cfg != 1.0: + noise_pred_uncond = transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=neg_cond.prompt_embeds.unsqueeze(0), + attention_kwargs=None, + return_dict=False, + )[0] + noise_pred = noise_pred_uncond + active_cfg * (noise_pred_cond - noise_pred_uncond) + else: + noise_pred = noise_pred_cond latents = scheduler.step(noise_pred, t, latents, return_dict=False)[0] diff --git a/invokeai/app/invocations/wan_latents_to_video.py b/invokeai/app/invocations/wan_latents_to_video.py index 3a245b0164b..25d2e63e92e 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 @@ -147,72 +150,95 @@ 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) + 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 + + # Keep the VAE cache lock while the MP4 writer consumes chunks. The causal decoder's + # feature cache and weights must remain live for the whole sequence; releasing the + # lock would require a separate bounded decode/encode queue and would add buffering. + 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)") - TorchDevice.empty_cache() + 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 ad81a153ac2..4facf97f404 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`, `xpu`, `cuda:N`, `xpu: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. pid_memory_optimization: Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path. 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` @@ -223,6 +224,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.") pid_memory_optimization: bool = Field(default=False, description="Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.") 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".') diff --git a/invokeai/backend/model_manager/load/load_base.py b/invokeai/backend/model_manager/load/load_base.py index 7225fd1402f..9df213f90a8 100644 --- a/invokeai/backend/model_manager/load/load_base.py +++ b/invokeai/backend/model_manager/load/load_base.py @@ -131,6 +131,11 @@ def compute_device(self) -> torch.device: """ return self._cache_record.cached_model.compute_device + @property + def supports_partial_loading(self) -> bool: + """Whether this model can stream individual weights between RAM and the compute device.""" + return isinstance(self._cache_record.cached_model, CachedModelWithPartialLoad) + def repair_required_tensors_on_device(self) -> int: """Repair required tensors that should be resident on the cached model's execution device.""" cached_model = self._cache_record.cached_model @@ -143,6 +148,20 @@ 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. + Full-load-only entries ignore the requested byte count and unload all weights. + """ + 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 213e4bbd2a2..4904ff5d1e3 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -1165,11 +1165,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() @@ -1180,18 +1190,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 a62170dbac9..bc224e50765 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -9,6 +9,10 @@ from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR from invokeai.backend.flux.modules.autoencoder import AutoEncoder +_WAN_VAE_SINGLE_FRAME_DECODE_SCALING_CONSTANT = 2900 +_WAN_VAE_VIDEO_DECODE_SCALING_CONSTANT_A14B = 6500 +_WAN_VAE_VIDEO_DECODE_SCALING_CONSTANT_TI2V = 7000 + def estimate_vae_working_memory_sd15_sdxl( operation: Literal["encode", "decode"], @@ -132,34 +136,54 @@ 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. - Generalizes the single-frame Wan 2.1 calibration (see - estimate_vae_working_memory_anima) to multi-frame clips and to the TI2V-5B VAE's - 16x spatial compression — callers pass *pixel-space* dimensions, so the VAE's - spatial scale factor is already applied. The Wan VAE processes the clip causally, - one latent frame at a time with cached features, so the conv working set scales - with a single frame's pixels; what grows with clip length is the full RGB clip, - which diffusers keeps resident on the execution device for the whole operation. + Callers pass pixel-space dimensions, so the VAE's spatial scale factor is already + applied. Single-frame decode and encode use the original Wan 2.1 calibration; + multi-frame decode uses conservative, VAE-variant-specific calibrations because + causal-convolution state makes the single-frame value unsafe at video resolutions. + The Wan VAE processes the clip causally, one latent frame at a time with cached + features. In streaming mode, only one temporal-upscale chunk of the RGB output is + kept on the execution device; otherwise the full output clip and its transient copy + are budgeted. """ element_size = next(vae.parameters()).element_size() - # Per-frame conv working set: ~2900 bytes per output pixel per element byte for a - # full-frame decode, encode ~50% (calibrated empirically on a Wan 2.1 fp16 decode). - scaling_constant = 2900 if operation == "decode" else 1450 + # The original 2900-byte calibration covers a single Wan 2.1 frame. Multi-frame video + # decodes retain causal-convolution state that makes that constant unsafe at video + # resolutions. These conservative constants are based on measured allocated-memory + # peaks with allocator headroom: 6500 for the z_dim=16 A14B VAE and 7000 for the + # larger z_dim=48 TI2V VAE. Keep the single-frame value for image decode and the + # existing encode calibration. + if operation == "decode" and pixel_frames > 1: + try: + z_dim = int(getattr(vae.config, "z_dim", 16)) + except (TypeError, ValueError): + z_dim = 48 + scaling_constant = ( + _WAN_VAE_VIDEO_DECODE_SCALING_CONSTANT_TI2V if z_dim >= 32 else _WAN_VAE_VIDEO_DECODE_SCALING_CONSTANT_A14B + ) + else: + scaling_constant = _WAN_VAE_SINGLE_FRAME_DECODE_SCALING_CONSTANT if operation == "decode" else 1450 if tile_size is not None: # Add 25% for tile overlap. per_frame = tile_size * tile_size * element_size * scaling_constant * 1.25 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 163ce53ed01..92d3e561c22 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -48908,6 +48908,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 + }, "pid_memory_optimization": { "type": "boolean", "title": "Pid Memory Optimization", @@ -49196,7 +49202,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`, `xpu`, `cuda:N`, `xpu: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 pid_memory_optimization: Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.\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 allow_private_download_urls: Allow the download queue to fetch from loopback, link-local and private-network addresses. Disabled by default so that a download URL cannot be used to reach services that are only reachable from the server. Enable this only if you install models from a mirror on your own network.\n download_proxy: Optional HTTP proxy for model downloads. The proxy must enforce the public-address policy because proxy-side DNS cannot be checked by InvokeAI.\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`, `xpu`, `cuda:N`, `xpu: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 pid_memory_optimization: Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.\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 allow_private_download_urls: Allow the download queue to fetch from loopback, link-local and private-network addresses. Disabled by default so that a download URL cannot be used to reach services that are only reachable from the server. Enable this only if you install models from a mirror on your own network.\n download_proxy: Optional HTTP proxy for model downloads. The proxy must enforce the public-address policy because proxy-side DNS cannot be checked by InvokeAI.\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 5fe1513f593..e52e01e4234 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -18814,6 +18814,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`, `xpu`, `cuda:N`, `xpu: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. * pid_memory_optimization: Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path. * 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` @@ -19158,6 +19159,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; /** * Pid Memory Optimization * @description Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path. diff --git a/scripts/calibrate_wan_vae_working_memory.py b/scripts/calibrate_wan_vae_working_memory.py new file mode 100644 index 00000000000..59cb83c992d --- /dev/null +++ b/scripts/calibrate_wan_vae_working_memory.py @@ -0,0 +1,225 @@ +"""Measure Wan VAE decode memory and compare it with the shipped estimate. + +Run on a CUDA or ROCm device with either a local diffusers VAE directory or a +single Wan ``.safetensors`` checkpoint: + + python scripts/calibrate_wan_vae_working_memory.py --vae /path/to/vae-or-checkpoint + +The default shape matches the 12 GiB-card calibration point. Use ``--no-streaming`` +to measure the full-frame decode path, or ``--tiling`` to measure the spatially +tiled path used as a low-VRAM fallback. Tiling overrides streaming. The script +reports allocated and reserved deltas; the implied scaling constant uses allocated +memory to match the shipped estimator, while reserved memory shows allocator headroom. +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import torch +from diffusers.models.autoencoders import AutoencoderKLWan + +# Direct script execution puts ``scripts/`` on sys.path, not the repository root. +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from invokeai.backend.model_manager.load.model_loaders.vae import _wan_vae_init_kwargs_for # noqa: E402 +from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_wan # noqa: E402 +from invokeai.backend.wan.vae_decode import iter_wan_vae_decode_chunks # noqa: E402 + +DTYPES = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32} + + +def _load_vae(path: Path, dtype: torch.dtype) -> AutoencoderKLWan: + if path.is_dir(): + vae = AutoencoderKLWan.from_pretrained(path, local_files_only=True, torch_dtype=dtype) + vae.eval() + return vae + if not path.is_file(): + raise ValueError("--vae must point to a diffusers directory or a Wan .safetensors file") + + import accelerate + from safetensors.torch import load_file + + state_dict = load_file(str(path), device="cpu") + try: + latent_channels = int(state_dict["decoder.conv_in.weight"].shape[1]) + except (KeyError, IndexError, TypeError, ValueError) as exc: + raise ValueError("Wan checkpoint is missing decoder.conv_in.weight or has an invalid shape") from exc + + with accelerate.init_empty_weights(): + vae = AutoencoderKLWan(**_wan_vae_init_kwargs_for(latent_channels)) + for key, tensor in state_dict.items(): + if tensor.is_floating_point(): + state_dict[key] = tensor.to(dtype=dtype) + vae.load_state_dict(state_dict, strict=True, assign=True) + vae.eval() + return vae + + +@torch.inference_mode() +def _measure( + vae: AutoencoderKLWan, + pixel_height: int, + pixel_width: int, + pixel_frames: int, + streaming: bool, + tiling: bool = False, + tile_size: int | None = None, +) -> dict[str, int | float | bool | str | None]: + temporal_scale = int(getattr(vae.config, "scale_factor_temporal", None) or 4) + spatial_scale = int(getattr(vae.config, "scale_factor_spatial", None) or 8) + if pixel_frames < 1 or (pixel_frames - 1) % temporal_scale != 0: + raise ValueError(f"pixel_frames must satisfy (frames - 1) % {temporal_scale} == 0") + if pixel_height < 1 or pixel_width < 1 or pixel_height % spatial_scale or pixel_width % spatial_scale: + raise ValueError(f"height and width must be positive multiples of {spatial_scale}") + + device = torch.device("cuda") + vae.to(device=device) + if tiling: + streaming = False + if tile_size is None: + tile_size = int(getattr(vae, "tile_sample_min_height", 256)) + if tile_size < spatial_scale or tile_size % spatial_scale: + raise ValueError(f"tile_size must be a positive multiple of {spatial_scale}") + vae.enable_tiling(tile_sample_min_height=tile_size, tile_sample_min_width=tile_size) + else: + tile_size = None + vae.disable_tiling() + element_size = next(vae.parameters()).element_size() + latent_frames = (pixel_frames - 1) // temporal_scale + 1 + latent_height = pixel_height // spatial_scale + latent_width = pixel_width // spatial_scale + latents = torch.randn( + 1, + int(getattr(vae.config, "z_dim", 16)), + latent_frames, + latent_height, + latent_width, + device=device, + dtype=next(vae.parameters()).dtype, + ) + + estimate = estimate_vae_working_memory_wan( + operation="decode", + vae=vae, + pixel_height=pixel_height, + pixel_width=pixel_width, + pixel_frames=pixel_frames, + tile_size=tile_size, + streaming=streaming, + ) + if streaming: + resident_frames = min(pixel_frames, temporal_scale) + clip_copies = 1 + else: + resident_frames = pixel_frames + clip_copies = 2 + clip_bytes = clip_copies * 3 * resident_frames * pixel_height * pixel_width * element_size + + torch.cuda.synchronize() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats(device) + baseline_allocated = torch.cuda.memory_allocated(device) + baseline_reserved = torch.cuda.memory_reserved(device) + if tiling: + assert tile_size is not None + scaling_basis_bytes = tile_size**2 * element_size * 1.25 + else: + scaling_basis_bytes = pixel_height * pixel_width * element_size + try: + if streaming: + for chunk in iter_wan_vae_decode_chunks(vae, latents): + chunk = chunk[0].cpu() + else: + vae.decode(latents, return_dict=False)[0].cpu() + finally: + if tiling: + vae.disable_tiling() + torch.cuda.synchronize() + peak_allocated = torch.cuda.max_memory_allocated(device) + peak_reserved = torch.cuda.max_memory_reserved(device) + measured_allocated_delta = peak_allocated - baseline_allocated + measured_reserved_delta = peak_reserved - baseline_reserved + implied_constant = (measured_allocated_delta - clip_bytes) / scaling_basis_bytes + return { + "device": torch.cuda.get_device_name(device), + "backend": "ROCm" if torch.version.hip is not None else "CUDA", + "dtype": str(next(vae.parameters()).dtype), + "streaming": streaming, + "tiling": tiling, + "tile_size": tile_size, + "pixel_height": pixel_height, + "pixel_width": pixel_width, + "pixel_frames": pixel_frames, + "estimate_bytes": estimate, + "measured_allocated_delta_bytes": measured_allocated_delta, + "measured_reserved_delta_bytes": measured_reserved_delta, + "implied_scaling_constant": implied_constant, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument( + "--vae", type=Path, required=True, help="Diffusers AutoencoderKLWan directory or .safetensors checkpoint." + ) + parser.add_argument("--height", type=int, default=704, help="Pixel height. Default: 704.") + parser.add_argument("--width", type=int, default=1280, help="Pixel width. Default: 1280.") + parser.add_argument("--frames", type=int, default=81, help="Pixel frame count. Default: 81.") + parser.add_argument("--dtype", choices=list(DTYPES), default="float16", help="VAE dtype. Default: float16.") + parser.add_argument( + "--streaming", + action=argparse.BooleanOptionalAction, + default=True, + help="Measure chunked streaming decode. Use --no-streaming for full decode.", + ) + parser.add_argument( + "--tiling", + action="store_true", + help="Measure spatially tiled full decode. Overrides --streaming; use --tile-size to override the tile size.", + ) + parser.add_argument( + "--tile-size", + type=int, + default=None, + help="Spatial tile size in pixels. Requires --tiling; defaults to the VAE tile size.", + ) + args = parser.parse_args() + + if args.tile_size is not None and not args.tiling: + parser.error("--tile-size requires --tiling") + if args.tile_size is not None and args.tile_size <= 0: + parser.error("--tile-size must be positive") + + if not torch.cuda.is_available(): + raise SystemExit("CUDA or ROCm device required") + vae = _load_vae(args.vae, DTYPES[args.dtype]) + try: + result = _measure( + vae, + args.height, + args.width, + args.frames, + args.streaming, + tiling=args.tiling, + tile_size=args.tile_size, + ) + except torch.cuda.OutOfMemoryError as exc: + raise SystemExit("VAE decode ran out of device memory; reduce --height, --width, or --frames") from exc + + gib = 2**30 + print(f"device: {result['device']} ({result['backend']})") + print(f"dtype: {result['dtype']}; streaming: {result['streaming']}; tiling: {result['tiling']}") + if result["tiling"]: + print(f"tile size: {result['tile_size']}px") + print(f"shape: {result['pixel_height']}x{result['pixel_width']}x{result['pixel_frames']}") + print(f"estimate: {result['estimate_bytes'] / gib:.3f} GiB") + print(f"measured allocated delta: {result['measured_allocated_delta_bytes'] / gib:.3f} GiB") + print(f"measured reserved delta: {result['measured_reserved_delta_bytes'] / gib:.3f} GiB") + print(f"implied scaling constant (allocated): {result['implied_scaling_constant']:.1f}") + + +if __name__ == "__main__": + main() diff --git a/tests/app/invocations/test_wan_denoise.py b/tests/app/invocations/test_wan_denoise.py index 4dbde079db1..acf3d904f9c 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,108 @@ 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() + loaded.supports_partial_loading = True + 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_expert_swapper_does_not_trim_when_residency_is_already_targeted(self) -> None: + transformer = _ZeroTransformer() + loaded = MagicMock() + loaded.supports_partial_loading = True + loaded._cache_record.cached_model.cur_vram_bytes.return_value = 2 * 2**30 + 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_not_called() + + def test_expert_swapper_skips_residency_trim_without_partial_loading(self) -> None: + transformer = _ZeroTransformer() + loaded = MagicMock() + loaded.supports_partial_loading = False + loaded._cache_record.cached_model.cur_vram_bytes.return_value = 5 * 2**30 + 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 + swapper = _ExpertSwapper( + context=context, + high_model=MagicMock(), + low_model=None, + inference_dtype=torch.bfloat16, + working_mem_bytes=22 * 2**30, + 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() + loaded.unload_from_vram.assert_not_called() + context.logger.warning.assert_called_once() + 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 +393,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..dd4deacc84c 100644 --- a/tests/app/invocations/test_wan_working_memory.py +++ b/tests/app/invocations/test_wan_working_memory.py @@ -71,25 +71,60 @@ def test_additional_frames_add_only_clip_bytes(self): grow the resident clip (2 copies at decode peak), not the conv working set.""" vae = _mock_wan_vae() one = estimate_vae_working_memory_wan( - operation="decode", vae=vae, pixel_height=128, pixel_width=128, pixel_frames=1 + operation="decode", vae=vae, pixel_height=128, pixel_width=128, pixel_frames=2 ) many = estimate_vae_working_memory_wan( operation="decode", vae=vae, pixel_height=128, pixel_width=128, pixel_frames=81 ) - assert many - one == 2 * 3 * 80 * 128 * 128 * 2 + assert many - one == 2 * 3 * 79 * 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=2, + 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 == 2 * 3 * 128 * 128 * 2 def test_tile_size_bounds_the_per_frame_term(self): vae = _mock_wan_vae() tiled = estimate_vae_working_memory_wan( operation="decode", vae=vae, pixel_height=1920, pixel_width=1080, pixel_frames=17, tile_size=256 ) - expected = int(256 * 256 * 2 * 2900 * 1.25 + 2 * 3 * 17 * 1920 * 1080 * 2) + expected = int(256 * 256 * 2 * 6500 * 1.25 + 2 * 3 * 17 * 1920 * 1080 * 2) assert tiled == expected full = estimate_vae_working_memory_wan( operation="decode", vae=vae, pixel_height=1920, pixel_width=1080, pixel_frames=17 ) assert tiled < full + def test_multi_frame_decode_estimate_triggers_tiling_on_12gb_cards(self): + """Conservative video estimates must engage the tiling fallback for both Wan VAEs.""" + for z_dim, spatial_scale in ((16, 8), (48, 16)): + vae = _mock_wan_vae(z_dim=z_dim, spatial_scale=spatial_scale) + estimate = estimate_vae_working_memory_wan( + operation="decode", + vae=vae, + pixel_height=704, + pixel_width=1280, + pixel_frames=81, + streaming=True, + ) + assert estimate > 0.9 * 12 * 2**30 + class TestWanInvocationsRequestWorkingMemory: """Every Wan VAE path must reserve its estimated working memory with the model cache.""" @@ -187,6 +222,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 +246,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 3efe2d2fe12..c734e96def1 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 @@ -415,6 +415,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/model_manager/load/test_loaded_model.py b/tests/backend/model_manager/load/test_loaded_model.py index b792475b842..a8dcfa86679 100644 --- a/tests/backend/model_manager/load/test_loaded_model.py +++ b/tests/backend/model_manager/load/test_loaded_model.py @@ -62,6 +62,21 @@ def test_model_on_device_leaves_full_load_models_unchanged(): assert all(param.device.type == "cpu" for param in model.parameters()) +def test_supports_partial_loading_reflects_cached_model_type(): + partial_model = CachedModelWithPartialLoad( + model=torch.nn.Linear(4, 4), compute_device=torch.device("meta"), keep_ram_copy=False + ) + full_model = CachedModelOnlyFullLoad( + model=torch.nn.Linear(4, 4), compute_device=torch.device("meta"), total_bytes=1, keep_ram_copy=False + ) + + loaded_partial_model = LoadedModelWithoutConfig(CacheRecord(key="partial", cached_model=partial_model), cache=None) + assert loaded_partial_model.supports_partial_loading + assert not LoadedModelWithoutConfig( + CacheRecord(key="full", cached_model=full_model), cache=None + ).supports_partial_loading + + def test_enter_unlocks_if_repair_raises(): class BrokenCachedModel(CachedModelWithPartialLoad): def repair_required_tensors_on_compute_device(self) -> int: 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_calibrate_wan_vae_working_memory.py b/tests/test_calibrate_wan_vae_working_memory.py new file mode 100644 index 00000000000..4c0a59c6d78 --- /dev/null +++ b/tests/test_calibrate_wan_vae_working_memory.py @@ -0,0 +1,138 @@ +from contextlib import nullcontext +from importlib.util import module_from_spec, spec_from_file_location +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + + +def _load_calibration_script(): + path = Path(__file__).parents[1] / "scripts" / "calibrate_wan_vae_working_memory.py" + spec = spec_from_file_location("calibrate_wan_vae_working_memory", path) + assert spec is not None and spec.loader is not None + module = module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_load_vae_accepts_single_wan_safetensors_checkpoint(tmp_path, monkeypatch): + script = _load_calibration_script() + checkpoint = tmp_path / "wan.safetensors" + checkpoint.touch() + state_dict = {"decoder.conv_in.weight": torch.zeros(2, 48, 1, 1, 1)} + fake_vae = MagicMock() + fake_autoencoder = MagicMock(return_value=fake_vae) + + monkeypatch.setattr(script, "AutoencoderKLWan", fake_autoencoder) + monkeypatch.setattr(script, "_wan_vae_init_kwargs_for", lambda latent_channels: {"z_dim": latent_channels}) + monkeypatch.setattr("safetensors.torch.load_file", lambda path, device: state_dict) + monkeypatch.setattr("accelerate.init_empty_weights", lambda: nullcontext()) + + result = script._load_vae(checkpoint, torch.float16) + + assert result is fake_vae + fake_autoencoder.assert_called_once_with(z_dim=48) + fake_vae.load_state_dict.assert_called_once_with(state_dict, strict=True, assign=True) + fake_vae.eval.assert_called_once_with() + + +def test_load_vae_accepts_diffusers_directory(tmp_path, monkeypatch): + script = _load_calibration_script() + directory = tmp_path / "vae" + directory.mkdir() + fake_vae = MagicMock() + monkeypatch.setattr(script.AutoencoderKLWan, "from_pretrained", MagicMock(return_value=fake_vae)) + + result = script._load_vae(directory, torch.bfloat16) + + assert result is fake_vae + script.AutoencoderKLWan.from_pretrained.assert_called_once_with( + directory, local_files_only=True, torch_dtype=torch.bfloat16 + ) + fake_vae.eval.assert_called_once_with() + + +def test_measure_tiling_uses_full_decode_and_tile_estimate(monkeypatch): + script = _load_calibration_script() + parameter = torch.nn.Parameter(torch.zeros(1, dtype=torch.bfloat16)) + fake_vae = MagicMock() + fake_vae.config = SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8, z_dim=16) + fake_vae.parameters.side_effect = lambda: iter([parameter]) + fake_vae.tile_sample_min_height = 256 + fake_vae.tile_sample_min_width = 256 + fake_vae.decode.return_value = (torch.zeros(1, 3, 4, 64, 64),) + + monkeypatch.setattr(script.torch, "randn", lambda *args, **kwargs: torch.zeros(*args, dtype=kwargs["dtype"])) + monkeypatch.setattr(script.torch.cuda, "synchronize", lambda *args, **kwargs: None) + monkeypatch.setattr(script.torch.cuda, "empty_cache", lambda: None) + monkeypatch.setattr(script.torch.cuda, "reset_peak_memory_stats", lambda *args, **kwargs: None) + monkeypatch.setattr(script.torch.cuda, "memory_reserved", lambda device: 100) + monkeypatch.setattr(script.torch.cuda, "max_memory_reserved", lambda device: 200) + monkeypatch.setattr(script.torch.cuda, "memory_allocated", lambda device: 50) + monkeypatch.setattr(script.torch.cuda, "max_memory_allocated", lambda device: 150) + monkeypatch.setattr(script.torch.cuda, "get_device_name", lambda device: "test-device") + estimate = MagicMock(return_value=123) + monkeypatch.setattr(script, "estimate_vae_working_memory_wan", estimate) + monkeypatch.setattr(script, "iter_wan_vae_decode_chunks", MagicMock(side_effect=AssertionError)) + + result = script._measure(fake_vae, 512, 512, 81, streaming=True, tiling=True, tile_size=128) + + assert result["streaming"] is False + assert result["tiling"] is True + assert result["tile_size"] == 128 + fake_vae.enable_tiling.assert_called_once_with(tile_sample_min_height=128, tile_sample_min_width=128) + fake_vae.disable_tiling.assert_called_once_with() + fake_vae.decode.assert_called_once() + estimate.assert_called_once_with( + operation="decode", + vae=fake_vae, + pixel_height=512, + pixel_width=512, + pixel_frames=81, + tile_size=128, + streaming=False, + ) + + +def test_measure_tiling_implied_constant_uses_tiled_area(monkeypatch): + script = _load_calibration_script() + parameter = torch.nn.Parameter(torch.zeros(1, dtype=torch.bfloat16)) + fake_vae = MagicMock() + fake_vae.config = SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8, z_dim=16) + fake_vae.parameters.side_effect = lambda: iter([parameter]) + fake_vae.decode.return_value = (torch.zeros(1, 3, 81, 64, 64),) + + tile_size = 128 + constant = 4321.0 + element_size = parameter.element_size() + pixel_height = pixel_width = 512 + pixel_frames = 81 + clip_bytes = 2 * 3 * pixel_frames * pixel_height * pixel_width * element_size + measured_delta = int(tile_size**2 * element_size * constant * 1.25 + clip_bytes) + + monkeypatch.setattr(script.torch, "randn", lambda *args, **kwargs: torch.zeros(*args, dtype=kwargs["dtype"])) + monkeypatch.setattr(script.torch.cuda, "synchronize", lambda *args, **kwargs: None) + monkeypatch.setattr(script.torch.cuda, "empty_cache", lambda: None) + monkeypatch.setattr(script.torch.cuda, "reset_peak_memory_stats", lambda *args, **kwargs: None) + monkeypatch.setattr(script.torch.cuda, "memory_reserved", lambda device: 100) + monkeypatch.setattr(script.torch.cuda, "max_memory_reserved", lambda device: measured_delta + 100 + 12345) + monkeypatch.setattr(script.torch.cuda, "memory_allocated", lambda device: 0) + monkeypatch.setattr(script.torch.cuda, "max_memory_allocated", lambda device: measured_delta) + monkeypatch.setattr(script.torch.cuda, "get_device_name", lambda device: "test-device") + monkeypatch.setattr(script, "estimate_vae_working_memory_wan", lambda **kwargs: measured_delta) + + result = script._measure( + fake_vae, + pixel_height, + pixel_width, + pixel_frames, + streaming=True, + tiling=True, + tile_size=tile_size, + ) + + assert result["measured_allocated_delta_bytes"] == measured_delta + assert result["measured_reserved_delta_bytes"] == measured_delta + 12345 + assert result["implied_scaling_constant"] == pytest.approx(constant, abs=0.01) diff --git a/tests/test_config.py b/tests/test_config.py index 78d5dbe3466..b3853f16c4e 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"