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

invokeai/app/invocations/qwen_image_image_to_latents.py

@@ -26,14 +26,18 @@
     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=FieldDescriptions.vae_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,21 +48,39 @@ 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.
+        effective_tile_size = None
+        if tiled:
+            effective_tile_size = tile_size if tile_size > 0 else getattr(vae_info.model, "tile_sample_min_height", 256)
+
         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()
+            # 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.
+            if tiled:
+                if tile_size > 0:
+                    vae.enable_tiling(tile_sample_min_height=tile_size, tile_sample_min_width=tile_size)
+                else:
+                    vae.enable_tiling()
+            else:
+                vae.disable_tiling()
 
             image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae.dtype)
             with torch.inference_mode():
@@ -91,7 +113,12 @@ 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.
+        if self.width and self.height and self.width > 0 and self.height > 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 +129,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

@@ -27,20 +27,34 @@
     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=FieldDescriptions.vae_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.
+        effective_tile_size = None
+        if tiled:
+            effective_tile_size = (
+                self.tile_size if self.tile_size > 0 else getattr(vae_info.model, "tile_sample_min_height", 256)
+            )
         # 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 +63,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,13 +77,15 @@ 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()
+                # 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 tiled:
+                    if self.tile_size > 0:
+                        vae.enable_tiling(tile_sample_min_height=self.tile_size, tile_sample_min_width=self.tile_size)
+                    else:
+                        vae.enable_tiling()
                 else:
                     vae.disable_tiling()
 

invokeai/backend/util/vae_working_memory.py

@@ -165,14 +165,22 @@ 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 full RGB image, which stays resident on the execution device either way.
     """
     latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1
 
@@ -210,7 +218,13 @@ 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.
+        working_memory += 3 * h * w * element_size
+    else:
+        working_memory = h * w * element_size * scaling_constant
 
     return int(working_memory)
 

invokeai/frontend/web/openapi.json

@@ -71710,6 +71710,27 @@
             "input": "connection",
             "orig_required": true
           },
+          "tiled": {
+            "default": false,
+            "description": "Processing using overlapping tiles (reduce memory consumption)",
+            "field_kind": "input",
+            "input": "any",
+            "orig_default": false,
+            "orig_required": false,
+            "title": "Tiled",
+            "type": "boolean"
+          },
+          "tile_size": {
+            "default": 0,
+            "description": "The tile size for VAE tiling in pixels (image space). If set to 0, the default tile size for the model will be used. Larger tile sizes generally produce better results at the cost of higher memory usage.",
+            "field_kind": "input",
+            "input": "any",
+            "multipleOf": 8,
+            "orig_default": 0,
+            "orig_required": false,
+            "title": "Tile Size",
+            "type": "integer"
+          },
           "width": {
             "anyOf": [
               {
@@ -71756,7 +71777,7 @@
         "tags": ["image", "latents", "vae", "i2l", "qwen_image"],
         "title": "Image to Latents - Qwen Image",
         "type": "object",
-        "version": "1.0.0",
+        "version": "1.1.0",
         "output": {
           "$ref": "#/components/schemas/LatentsOutput"
         }
@@ -71854,6 +71875,27 @@
             "input": "connection",
             "orig_required": true
           },
+          "tiled": {
+            "default": false,
+            "description": "Processing using overlapping tiles (reduce memory consumption)",
+            "field_kind": "input",
+            "input": "any",
+            "orig_default": false,
+            "orig_required": false,
+            "title": "Tiled",
+            "type": "boolean"
+          },
+          "tile_size": {
+            "default": 0,
+            "description": "The tile size for VAE tiling in pixels (image space). If set to 0, the default tile size for the model will be used. Larger tile sizes generally produce better results at the cost of higher memory usage.",
+            "field_kind": "input",
+            "input": "any",
+            "multipleOf": 8,
+            "orig_default": 0,
+            "orig_required": false,
+            "title": "Tile Size",
+            "type": "integer"
+          },
           "type": {
             "const": "qwen_image_l2i",
             "default": "qwen_image_l2i",
@@ -71866,7 +71908,7 @@
         "tags": ["latents", "image", "vae", "l2i", "qwen_image"],
         "title": "Latents to Image - Qwen Image",
         "type": "object",
-        "version": "1.0.0",
+        "version": "1.1.0",
         "output": {
           "$ref": "#/components/schemas/ImageOutput"
         }

invokeai/frontend/web/src/services/api/schema.ts

@@ -30307,6 +30307,18 @@ export type components = {
              * @default null
              */
             vae?: components["schemas"]["VAEField"] | null;
+            /**
+             * Tiled
+             * @description Processing using overlapping tiles (reduce memory consumption)
+             * @default false
+             */
+            tiled?: boolean;
+            /**
+             * Tile Size
+             * @description The tile size for VAE tiling in pixels (image space). If set to 0, the default tile size for the model will be used. Larger tile sizes generally produce better results at the cost of higher memory usage.
+             * @default 0
+             */
+            tile_size?: number;
             /**
              * Width
              * @description Resize the image to this width before encoding. If not set, encodes at the image's original size.
@@ -30368,6 +30380,18 @@ export type components = {
              * @default null
              */
             vae?: components["schemas"]["VAEField"] | null;
+            /**
+             * Tiled
+             * @description Processing using overlapping tiles (reduce memory consumption)
+             * @default false
+             */
+            tiled?: boolean;
+            /**
+             * Tile Size
+             * @description The tile size for VAE tiling in pixels (image space). If set to 0, the default tile size for the model will be used. Larger tile sizes generally produce better results at the cost of higher memory usage.
+             * @default 0
+             */
+            tile_size?: number;
             /**
              * type
              * @default qwen_image_l2i

tests/app/invocations/test_qwen_image_working_memory.py

@@ -178,3 +178,29 @@ def test_qwen_image_to_latents_requests_working_memory(self):
             mock_estimate.assert_called_once()
             assert mock_estimate.call_args.kwargs["operation"] == "encode"
             mock_vae_info.model_on_device.assert_called_once_with(working_mem_bytes=expected_memory)
+
+    def test_qwen_image_to_latents_passes_resolved_tile_size_to_the_estimate(self):
+        """Tiling only helps if the *estimate* shrinks with it.
+
+        The cache reserves whatever the estimator returns, so enabling tiling on the VAE without
+        telling the estimator would leave it reserving the full-frame figure (~11 GB at 2560x1440)
+        and evicting models to honour it: the encode would be bounded, but nothing else would fit.
+        tile_size=0 means "use the model default", which must be resolved before estimating.
+        """
+        _mock_vae, mock_vae_info = self._mock_vae_info()
+        mock_vae_info.model.tile_sample_min_height = 256
+        mock_image_tensor = torch.zeros(1, 3, 512, 512)
+
+        estimation_path = "invokeai.app.invocations.qwen_image_image_to_latents.estimate_vae_working_memory_qwen_image"
+
+        for tiled, tile_size, expected in ((False, 0, None), (True, 0, 256), (True, 512, 512)):
+            with patch(estimation_path) as mock_estimate:
+                mock_estimate.return_value = 1024
+                try:
+                    QwenImageImageToLatentsInvocation.vae_encode(
+                        mock_vae_info, mock_image_tensor, tiled=tiled, tile_size=tile_size
+                    )
+                except Exception:
+                    # Downstream encode math fails under mocking; only the estimate call matters.
+                    pass
+                assert mock_estimate.call_args.kwargs["tile_size"] == expected, f"tiled={tiled}, tile_size={tile_size}"