perf(qwen-image): add a tiling option to the Qwen-Image VAE nodes

invokeai/app/invocations/qwen_image_image_to_latents.py

@@ -14,7 +14,12 @@
 from invokeai.app.invocations.model import VAEField
 from invokeai.app.invocations.primitives import LatentsOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.krea2.vae_compat import as_qwen_image_vae
+from invokeai.backend.krea2.vae_compat import (
+    QWEN_IMAGE_VAE_MIN_TILE_SIZE,
+    as_qwen_image_vae,
+    patch_qwen_image_vae_tiling,
+    resolve_qwen_image_vae_tile_size,
+)
 from invokeai.backend.model_manager.load.load_base import LoadedModel
 from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
 from invokeai.backend.util.devices import TorchDevice
@@ -26,14 +31,23 @@
     title="Image to Latents - Qwen Image",
     tags=["image", "latents", "vae", "i2l", "qwen_image"],
     category="image",
-    version="1.0.0",
+    version="1.1.0",
     classification=Classification.Prototype,
 )
 class QwenImageImageToLatentsInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Generates latents from an image using the Qwen Image VAE."""
 
     image: ImageField = InputField(description="The image to encode.")
     vae: VAEField = InputField(description=FieldDescriptions.vae, input=Input.Connection)
+    tiled: bool = InputField(default=False, description=FieldDescriptions.tiled)
+    # NOTE: tile_size = 0 is a special value meaning "use the model's default", matching the
+    # SD/SDXL i2l node. `int | None` is avoided because the workflow UI does not handle it well.
+    tile_size: int = InputField(
+        default=0,
+        multiple_of=8,
+        description=f"{FieldDescriptions.vae_tile_size} Values between 1 and "
+        f"{QWEN_IMAGE_VAE_MIN_TILE_SIZE} are raised to {QWEN_IMAGE_VAE_MIN_TILE_SIZE}.",
+    )
     width: int | None = InputField(
         default=None,
         description="Resize the image to this width before encoding. If not set, encodes at the image's original size.",
@@ -44,24 +58,36 @@ class QwenImageImageToLatentsInvocation(BaseInvocation, WithMetadata, WithBoard)
     )
 
     @staticmethod
-    def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tensor:
+    def vae_encode(
+        vae_info: LoadedModel, image_tensor: torch.Tensor, tiled: bool = False, tile_size: int = 0
+    ) -> torch.Tensor:
         # NOTE: vae_info.model may be an AutoencoderKLWan (a native-layout qwen_image_vae single file is
         # classified with the Anima base); it is reinterpreted as AutoencoderKLQwenImage inside the
         # model_on_device context below. The working-memory estimate only reads tensor shape + element
         # size, so it is safe to run on either class here.
+        # Resolve tile_size=0 ("model default") before estimating, so the reserved working memory
+        # matches the tiles the VAE will actually use. Resolved against a constant rather than the
+        # module's current tile_sample_min_height, which a previous invocation may have overwritten.
+        effective_tile_size = resolve_qwen_image_vae_tile_size(tile_size) if tiled else None
+
         estimated_working_memory = estimate_vae_working_memory_qwen_image(
             operation="encode",
             image_tensor=image_tensor,
             vae=vae_info.model,
+            tile_size=effective_tile_size,
         )
         with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
             # Reinterpret an Anima-classified Wan VAE as AutoencoderKLQwenImage (identical weights).
             vae = as_qwen_image_vae(vae)
 
-            vae.disable_tiling()
-
             image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae.dtype)
-            with torch.inference_mode():
+
+            # Tiling bounds the encode's peak memory to a single tile, which is what makes large
+            # inputs (e.g. a 2560x1440 upscale round-trip) encodable while a multi-GB transformer
+            # is still resident. Off by default: full-frame is faster and avoids tile blending.
+            # The tiling state is scoped to this block: the VAE module belongs to the model cache
+            # and is shared with later invocations (and with the Anima decode node).
+            with torch.inference_mode(), patch_qwen_image_vae_tiling(vae, effective_tile_size):
                 # The Qwen Image VAE expects 5D input: (B, C, num_frames, H, W)
                 if image_tensor.dim() == 4:
                     image_tensor = image_tensor.unsqueeze(2)
@@ -91,7 +117,13 @@ def invoke(self, context: InvocationContext) -> LatentsOutput:
 
         # If target dimensions are specified, resize the image BEFORE encoding
         # (matching the diffusers pipeline which resizes in pixel space, not latent space).
-        if self.width is not None and self.height is not None:
+        #
+        # `width`/`height` are `int | None`, but the workflow UI cannot represent None in a number
+        # input and sends 0 for "unset" — which `is not None`, so a naive check reached
+        # `resize((0, 0))` and raised "height and width must be > 0". Treat any non-positive value
+        # as unset, which is also how `tile_size` uses 0. Note this means a half-filled pair (e.g.
+        # width=1024, height=0) encodes at the original size rather than raising.
+        if (self.width or 0) > 0 and (self.height or 0) > 0:
             image = image.convert("RGB").resize((self.width, self.height), resample=PILImage.LANCZOS)
 
         # multiple_of=16 ensures the post-VAE latents (vae_scale_factor=8) have even
@@ -102,7 +134,12 @@ def invoke(self, context: InvocationContext) -> LatentsOutput:
 
         vae_info = context.models.load(self.vae.vae)
 
-        latents = self.vae_encode(vae_info=vae_info, image_tensor=image_tensor)
+        latents = self.vae_encode(
+            vae_info=vae_info,
+            image_tensor=image_tensor,
+            tiled=self.tiled or context.config.get().force_tiled_decode,
+            tile_size=self.tile_size,
+        )
 
         latents = latents.to("cpu")
         name = context.tensors.save(tensor=latents)

invokeai/app/invocations/qwen_image_latents_to_image.py

@@ -1,5 +1,3 @@
-from contextlib import nullcontext
-
 import torch
 from einops import rearrange
 from PIL import Image
@@ -16,7 +14,12 @@
 from invokeai.app.invocations.model import VAEField
 from invokeai.app.invocations.primitives import ImageOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.krea2.vae_compat import as_qwen_image_vae
+from invokeai.backend.krea2.vae_compat import (
+    QWEN_IMAGE_VAE_MIN_TILE_SIZE,
+    as_qwen_image_vae,
+    patch_qwen_image_vae_tiling,
+    resolve_qwen_image_vae_tile_size,
+)
 from invokeai.backend.stable_diffusion.extensions.seamless import SeamlessExt
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_qwen_image
@@ -27,20 +30,38 @@
     title="Latents to Image - Qwen Image",
     tags=["latents", "image", "vae", "l2i", "qwen_image"],
     category="latents",
-    version="1.0.0",
+    version="1.1.0",
     classification=Classification.Prototype,
 )
 class QwenImageLatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Generates an image from latents using the Qwen Image VAE."""
 
     latents: LatentsField = InputField(description=FieldDescriptions.latents, input=Input.Connection)
     vae: VAEField = InputField(description=FieldDescriptions.vae, input=Input.Connection)
+    tiled: bool = InputField(default=False, description=FieldDescriptions.tiled)
+    # NOTE: tile_size = 0 is a special value meaning "use the model's default", matching the
+    # SD/SDXL l2i node. `int | None` is avoided because the workflow UI does not handle it well.
+    tile_size: int = InputField(
+        default=0,
+        multiple_of=8,
+        description=f"{FieldDescriptions.vae_tile_size} Values between 1 and "
+        f"{QWEN_IMAGE_VAE_MIN_TILE_SIZE} are raised to {QWEN_IMAGE_VAE_MIN_TILE_SIZE}.",
+    )
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ImageOutput:
         latents = context.tensors.load(self.latents.latents_name)
 
         vae_info = context.models.load(self.vae.vae)
+        tiled = self.tiled or context.config.get().force_tiled_decode
+        # Resolve tile_size=0 ("model default") before estimating, so the memory the cache reserves
+        # matches the tiles the VAE will actually use. Without this the estimate stays at the
+        # full-frame figure (~21 GB at 2560x1440 on CUDA) and tiling frees nothing: the VAE is
+        # bounded, but the cache still evicts other models to honour the reservation.
+        # Resolved against a constant rather than the module's current tile_sample_min_height, which a
+        # previous invocation (including the Anima decode node, which shares this VAE instance) may
+        # have overwritten.
+        effective_tile_size = resolve_qwen_image_vae_tile_size(self.tile_size) if tiled else None
         # NOTE: vae_info.model may be an AutoencoderKLWan (a native-layout qwen_image_vae single file is
         # classified with the Anima base); it is reinterpreted as AutoencoderKLQwenImage inside the
         # model_on_device context below. The working-memory estimate only reads tensor shape + element
@@ -49,6 +70,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
             operation="decode",
             image_tensor=latents,
             vae=vae_info.model,
+            tile_size=effective_tile_size,
         )
         with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
             context.util.signal_progress("Running VAE")
@@ -62,17 +84,13 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
                 # which would wrongly place the latents (and thus the whole decode) on the CPU (see #9373).
                 latents = latents.to(device=vae_info.compute_device, dtype=vae.dtype)
 
-                # Honor the global force_tiled_decode setting, like the SD/SDXL l2i node. Tiling bounds the
-                # VAE's per-tile memory, which is the scalable way to decode very large outputs that would
-                # exceed VRAM even after offloading the transformer/text encoder. For normal sizes, leave
-                # it off (faster, no tile blending) — the reserved working memory offloads other models so
-                # the full-frame decode fits.
-                if context.config.get().force_tiled_decode:
-                    vae.enable_tiling()
-                else:
-                    vae.disable_tiling()
-
-                tiling_context = nullcontext()
+                # Tiling bounds the VAE's per-tile memory, which is the scalable way to decode very
+                # large outputs that would exceed VRAM even after offloading the transformer/text
+                # encoder. For normal sizes, leave it off (faster, no tile blending) — the reserved
+                # working memory offloads other models so the full-frame decode fits.
+                # The tiling state is scoped to this block: the VAE module belongs to the model cache
+                # and is shared with later invocations (and with the Anima decode node).
+                tiling_context = patch_qwen_image_vae_tiling(vae, effective_tile_size)
 
                 TorchDevice.empty_cache()
 

invokeai/backend/krea2/vae_compat.py

@@ -6,8 +6,13 @@
 share the exact same diffusers state-dict (identical keys and shapes), so a Wan-loaded VAE can be
 used through the same encode/decode path without rebuilding it. Both default configs carry the same
 Qwen-Image ``latents_mean`` / ``latents_std`` / ``z_dim`` values read by the Qwen encode/decode nodes.
+
+Also holds the tiling helpers those nodes share, since the tile geometry has to be applied identically
+on both classes and restored afterwards — see ``patch_qwen_image_vae_tiling``.
 """
 
+from collections.abc import Iterator
+from contextlib import contextmanager
 from typing import Any
 
 from diffusers.models.autoencoders import AutoencoderKLWan
@@ -55,3 +60,91 @@ def as_qwen_image_vae(model: Any) -> QwenImageCompatibleVAE:
         )
 
     return model
+
+
+# The stock AutoencoderKLQwenImage tile geometry: 256px tiles advancing in 192px steps, i.e. a 3/4
+# stride ratio with a 64px blend band. Both nodes resolve tile_size=0 to QWEN_IMAGE_VAE_DEFAULT_TILE_SIZE
+# rather than reading the module's current value, which another invocation may have overwritten.
+QWEN_IMAGE_VAE_DEFAULT_TILE_SIZE = 256
+_QWEN_IMAGE_VAE_TILE_STRIDE_NUMERATOR = 3
+_QWEN_IMAGE_VAE_TILE_STRIDE_DENOMINATOR = 4
+
+# diffusers derives the latent-space stride as `tile_sample_stride // spatial_compression_ratio` and
+# uses it as the step of the tile loop, so a stride below 8 makes that step 0 (ValueError) or leaves the
+# pixel and latent steps inconsistent. 64px is the smallest tile whose 3/4 stride is still a clean
+# multiple of 8, and tiles below it are not useful in practice anyway.
+QWEN_IMAGE_VAE_MIN_TILE_SIZE = 64
+
+
+def resolve_qwen_image_vae_tile_size(tile_size: int) -> int:
+    """Resolve a node's ``tile_size`` field to the tile size the VAE will actually use.
+
+    ``tile_size <= 0`` is the nodes' "use the default" sentinel (the workflow UI cannot represent
+    ``None`` in a number input and sends 0). Values below ``QWEN_IMAGE_VAE_MIN_TILE_SIZE`` are clamped
+    rather than rejected, because the field also has to accept the 0 sentinel and so cannot carry a
+    pydantic lower bound.
+    """
+    if tile_size <= 0:
+        return QWEN_IMAGE_VAE_DEFAULT_TILE_SIZE
+    return max(tile_size, QWEN_IMAGE_VAE_MIN_TILE_SIZE)
+
+
+def _tile_stride_for(tile_size: int) -> int:
+    """Return the tile stride to pair with ``tile_size``, keeping the stock 3/4 ratio.
+
+    Rounded down to a multiple of the VAE's 8x spatial compression: ``tiled_encode``/``tiled_decode``
+    step the tile loop in one space (pixels for encode, latents for decode) while slicing the
+    accumulated tile in the other, so the pixel stride must be exactly 8x the latent stride or the
+    two disagree and the output is misaligned.
+    """
+    stride = tile_size * _QWEN_IMAGE_VAE_TILE_STRIDE_NUMERATOR // _QWEN_IMAGE_VAE_TILE_STRIDE_DENOMINATOR
+    return max(_QWEN_IMAGE_VAE_SPATIAL_SCALE, stride // _QWEN_IMAGE_VAE_SPATIAL_SCALE * _QWEN_IMAGE_VAE_SPATIAL_SCALE)
+
+
+@contextmanager
+def patch_qwen_image_vae_tiling(vae: QwenImageCompatibleVAE, tile_size: int | None) -> Iterator[None]:
+    """Set the VAE's tiling state for the duration of the block, then restore it.
+
+    Two things make this a context manager rather than a bare ``enable_tiling()`` call:
+
+    - ``enable_tiling`` writes the tile geometry straight onto the module, and ``disable_tiling`` only
+      clears ``use_tiling`` — it does not restore the sizes. The module here is the model cache's own
+      instance (``as_qwen_image_vae`` deliberately returns it unchanged to keep partial-loading hooks
+      intact), so without a restore a tile size set once would persist for the lifetime of the cache
+      entry and leak into later invocations — including ``anima_latents_to_image``, which shares the
+      same VAE instance when a native-layout ``qwen_image_vae`` single file is loaded.
+    - All four parameters are always passed explicitly. ``enable_tiling`` falls back to the module's
+      current value for any argument left out, and its ``min``/``stride`` pair must stay consistent:
+      the tile loops advance by *stride* but slice each accumulated tile to *min*. A ``min`` below the
+      inherited 192px stride silently drops whole bands of the image, and a ``min`` above it grows
+      every tile without removing any, making compute scale with ``tile_size**2``.
+
+    ``tile_size=None`` disables tiling for the block.
+    """
+    original = (
+        vae.use_tiling,
+        vae.tile_sample_min_height,
+        vae.tile_sample_min_width,
+        vae.tile_sample_stride_height,
+        vae.tile_sample_stride_width,
+    )
+    try:
+        if tile_size is None:
+            vae.disable_tiling()
+        else:
+            stride = _tile_stride_for(tile_size)
+            vae.enable_tiling(
+                tile_sample_min_height=tile_size,
+                tile_sample_min_width=tile_size,
+                tile_sample_stride_height=stride,
+                tile_sample_stride_width=stride,
+            )
+        yield
+    finally:
+        (
+            vae.use_tiling,
+            vae.tile_sample_min_height,
+            vae.tile_sample_min_width,
+            vae.tile_sample_stride_height,
+            vae.tile_sample_stride_width,
+        ) = original

invokeai/backend/util/vae_working_memory.py

@@ -189,14 +189,25 @@ def estimate_vae_working_memory_wan(
 
 
 def estimate_vae_working_memory_qwen_image(
-    operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoencoderKLQwenImage
+    operation: Literal["encode", "decode"],
+    image_tensor: torch.Tensor,
+    vae: AutoencoderKLQwenImage,
+    tile_size: int | None = None,
 ) -> int:
     """Estimate the working memory required by the invocation in bytes.
 
     The Qwen Image VAE is a video-style autoencoder that operates on 5D tensors of shape
-    (B, C, num_frames, H, W). Tiling is not used, so peak working memory scales with the full
-    spatial output. The two trailing dimensions are the spatial H/W in latent space (decode) or
-    pixel space (encode), matching the convention used by the other estimators here.
+    (B, C, num_frames, H, W). The two trailing dimensions are the spatial H/W in latent space
+    (decode) or pixel space (encode), matching the convention used by the other estimators here.
+
+    Without tiling, peak working memory scales with the full spatial extent. With tiling it is
+    bounded by a single tile instead, so the estimate must follow suit — otherwise the cache keeps
+    reserving the full-frame figure (~11.8 GB for a 2560x1440 encode on CUDA) and tiling buys
+    nothing. Mirrors ``estimate_vae_working_memory_wan``: one tile plus 25% for the tile overlap,
+    plus the pixel-space buffers, which stay resident on the execution device either way.
+
+    ``tile_size`` is the resolved tile size (the nodes' 0 sentinel already substituted), and assumes
+    the 4:3 tile-to-stride ratio applied by ``patch_qwen_image_vae_tiling``.
     """
     latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1
 
@@ -240,7 +251,22 @@ def estimate_vae_working_memory_qwen_image(
     else:  # encode
         scaling_constant = 6300 if is_rocm else 1600
 
-    working_memory = h * w * element_size * scaling_constant
+    if tile_size is not None and tile_size > 0:
+        # Bounded by one tile (plus overlap) rather than the full frame.
+        working_memory = tile_size * tile_size * element_size * scaling_constant * 1.25
+        # The full RGB image is the encode input / decode output and stays resident regardless. Unlike
+        # the per-tile term this scales with the output area, so it is the term that decides whether the
+        # estimate still holds at the resolutions tiling exists for.
+        #
+        # `tiled_decode` holds several pixel-space copies at once: every decoded tile in `rows`
+        # ((tile_min / tile_stride)^2 ~ 1.8 frames at the 4:3 ratio the nodes set), the blended and
+        # cropped `result_rows` (~1 frame) and the final `torch.cat` output (~1 frame). Measured at
+        # ~5 frames on a 2560x1440 fp16 decode. Encode consumes its input image without duplicating it,
+        # and accumulates only latents (16 channels at 1/64 the area — negligible).
+        image_copies = 5 if operation == "decode" else 1
+        working_memory += image_copies * 3 * h * w * element_size
+    else:
+        working_memory = h * w * element_size * scaling_constant
 
     return int(working_memory)