feat(flux2): tile the FLUX.2 VAE encode and decode

docs/src/content/docs/Users Guide/03.Models/Local Models/flux-2.mdx

@@ -1,7 +1,7 @@
 ---
 title: FLUX.2
 description: Generate and edit images with FLUX.2 Klein (4B and 9B, distilled and Base) and FLUX.2 [dev], including multi-reference image editing.
-lastUpdated: 2026-09-30
+lastUpdated: 2026-10-16
 sidebar:
   order: 5
 ---
@@ -169,3 +169,14 @@ Decode](/users-guide/image-generation/advanced-topics/pid-decode/).
 
 **FLUX2 Denoise** also has a **CFG Scale** and a negative conditioning input, which the Generate tab leaves
 off (CFG 1.0).
+
+:::note[Tiled VAE decode]
+**Latents to Image - FLUX2** and **Image to Latents - FLUX2** have **Tiled** and **Tile Size** fields.
+Without them, Invoke still tiles a decode that is too large for the GPU (`auto_tiled_decode`) or that ran
+out of memory, and `force_tiled_decode` tiles every decode and encode; see [Tiled VAE decode and
+encode](/configuration/optimization/low-vram-mode/#tiled-vae-decode-and-encode). A tiled result is not
+pixel-identical to an untiled one: the difference is fine, texture-level noise spread over the whole image,
+not seams, so the same seed does not reproduce an untiled image exactly. On a 1024×1024 FLUX.2 Klein image,
+a tiled decode differed from the untiled one by 5.6 of 255 per pixel value on average (about 2%), with no
+visible seams; at 2048×2048 by 3.3 of 255.
+:::

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

@@ -202,19 +202,19 @@ clear_vram_after_session: true
 
 Without tiling, the memory a VAE needs grows with the square of the image size. For FLUX.1, an untiled 3072² encode peaks at about 18.8 GiB on an RTX 4090, and a tiled one at 0.8 GiB, while tiled memory stays roughly flat as the image grows. Tiling costs some speed on smaller images.
 
-Invoke tiles large decodes by itself: when an untiled decode would claim most of the memory the GPU keeps resident, it decodes in tiles instead, for the FLUX.1, Z-Image, Qwen-Image, Qwen-Image-2.1, Krea-2, Anima and Wan decodes, and for Qwen-Image-2.1's encodes as well. Because a tiled decode is not pixel-identical to an untiled one, this only kicks in for images too large for the card: for Z-Image on a 16GB GPU, for example, from about 1850x1850 on Nvidia and from about 1500x1500 on AMD, whose convolution library needs more memory. On a ROCm build under Windows the comparison uses the video-memory budget Windows grants Invoke, which is lower than the card's total and lower again while another program uses the GPU. Anima tiles a little earlier, because its untiled decode pushes the Anima model itself out of VRAM on 8GB cards. Wan videos are tiled from 832x480 on 8GB cards and from 720p on 16GB cards, where their untiled decodes need about as much memory as the card holds, or more.
+Invoke tiles large decodes by itself: when an untiled decode would claim most of the memory the GPU keeps resident, it decodes in tiles instead, for the FLUX.1, FLUX.2, Z-Image, Qwen-Image, Qwen-Image-2.1, Krea-2, Anima and Wan decodes, and for Qwen-Image-2.1's encodes as well. Because a tiled decode is not pixel-identical to an untiled one, this only kicks in for images too large for the card: for Z-Image on a 16GB GPU, for example, from about 1850x1850 on Nvidia and from about 1500x1500 on AMD, whose convolution library needs more memory. On a ROCm build under Windows the comparison uses the video-memory budget Windows grants Invoke, which is lower than the card's total and lower again while another program uses the GPU. Anima tiles a little earlier, because its untiled decode pushes the Anima model itself out of VRAM on 8GB cards. Wan videos are tiled from 832x480 on 8GB cards and from 720p on 16GB cards, where their untiled decodes need about as much memory as the card holds, or more.
 
-To always decode untiled unless you ask for tiling, turn this off. A decode too large for the card then runs out of memory: the FLUX.1 and Z-Image decodes retry in tiles, the others fail, and on Windows the driver may move the decode into system memory and run it very slowly instead. Anima is not affected by this setting: there, tiling a large decode is what keeps the Anima model in VRAM, and turning it off would only make those decodes slower.
+To always decode untiled unless you ask for tiling, turn this off. A decode too large for the card then runs out of memory: the FLUX.1, FLUX.2 and Z-Image decodes retry in tiles, the others fail, and on Windows the driver may move the decode into system memory and run it very slowly instead. Anima is not affected by this setting: there, tiling a large decode is what keeps the Anima model in VRAM, and turning it off would only make those decodes slower.
 
 ```yaml
 auto_tiled_decode: false
 ```
 
-- **Always tile:** set `force_tiled_decode: true` in `invokeai.yaml`. It covers the VAE decode and encode for SD1.5/SDXL, FLUX.1, Z-Image, Qwen-Image and Qwen-Image-2.1, the Anima encode, the control image of a Z-Image ControlNet, and the Krea-2 style reference. Other model families (for example SD3, FLUX.2, Wan and LTX-2) ignore it.
-- **Per node:** the SD, FLUX.1, Z-Image, Qwen-Image and Qwen-Image-2.1 **Latents to Image** and **Image to Latents** nodes, and **Image to Latents - Anima**, have **Tiled** and **Tile Size** fields for workflows. So does the **Z-Image ControlNet** node, for the control image the denoise encodes. The [Upscale panel](/users-guide/image-generation/upscale/) turns tiling on by itself.
-- **Automatic retry:** if a FLUX.1, Z-Image or Anima decode, or a FLUX.1 encode, runs out of memory, InvokeAI retries it once with tiling and logs that it did.
+- **Always tile:** set `force_tiled_decode: true` in `invokeai.yaml`. It covers the VAE decode and encode for SD1.5/SDXL, FLUX.1, FLUX.2, Z-Image, Qwen-Image and Qwen-Image-2.1, the Anima encode, the control image of a Z-Image ControlNet, and the Krea-2 style reference. Other model families (for example SD3, Wan and LTX-2) ignore it.
+- **Per node:** the SD, FLUX.1, FLUX.2, Z-Image, Qwen-Image and Qwen-Image-2.1 **Latents to Image** and **Image to Latents** nodes, and **Image to Latents - Anima**, have **Tiled** and **Tile Size** fields for workflows. So does the **Z-Image ControlNet** node, for the control image the denoise encodes. The [Upscale panel](/users-guide/image-generation/upscale/) turns tiling on by itself.
+- **Automatic retry:** if a FLUX.1, FLUX.2, Z-Image or Anima decode, or a FLUX.1 or FLUX.2 encode, runs out of memory, InvokeAI retries it once with tiling and logs that it did.
 
-A tiled decode is not pixel-identical to a single-pass one. The VAE looks at the whole image at once, so the difference is spread over the image rather than limited to the tile seams; it is small, but the same seed will not reproduce an untiled result exactly.
+A tiled decode is not pixel-identical to a single-pass one. The VAE looks at the whole image at once, so the difference is spread over the image rather than limited to the tile seams; it is small (for a 1024x1024 FLUX.2 Klein image, 5.6 of 255 per pixel value on average, with no visible seams), but the same seed will not reproduce an untiled result exactly.
 
 ### PiD decode optimization
 

docs/src/generated/settings.json

@@ -815,7 +815,7 @@
     {
       "category": "GENERATION",
       "default": false,
-      "description": "Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). A tiled decode is not pixel-identical to a single-pass one: a VAE decoder normalises and attends over the whole image, so the difference is spread across it rather than confined to the tile seams. Also tiles the encode. Read by the SD1.5/SDXL, FLUX.1, Z-Image and Qwen-Image VAE nodes, the Anima encode, the Z-Image control encode and the Krea-2 style reference; other model families ignore it.",
+      "description": "Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). A tiled decode is not pixel-identical to a single-pass one: a VAE decoder normalises and attends over the whole image, so the difference is spread across it rather than confined to the tile seams. Also tiles the encode. Read by the SD1.5/SDXL, FLUX.1, FLUX.2, Z-Image, Qwen-Image and Qwen-Image-2.1 VAE nodes, the Anima encode, the Z-Image control encode and the Krea-2 style reference; other model families ignore it.",
       "env_var": "INVOKEAI_FORCE_TILED_DECODE",
       "literal_values": [],
       "name": "force_tiled_decode",
@@ -826,7 +826,7 @@
     {
       "category": "GENERATION",
       "default": true,
-      "description": "Decode large images in tiles when an untiled decode's estimated working memory would take most of the GPU's memory (FLUX.1, Z-Image, Qwen-Image, Krea-2 and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1 and Z-Image decodes then still retry tiled after running out of memory. Anima keeps tiling a decode too large for its GPU either way, because tiling is what makes that case fast there.",
+      "description": "Decode large images in tiles when an untiled decode's estimated working memory would take most of the GPU's memory (FLUX.1, FLUX.2, Z-Image, Qwen-Image, Qwen-Image-2.1, Krea-2 and Wan decodes, and Qwen-Image-2.1 encodes). Turn off to decode untiled unless tiling is requested; the FLUX.1, FLUX.2 and Z-Image decodes then still retry tiled after running out of memory. Anima keeps tiling a decode too large for its GPU either way, because tiling is what makes that case fast there.",
       "env_var": "INVOKEAI_AUTO_TILED_DECODE",
       "literal_values": [],
       "name": "auto_tiled_decode",

invokeai/app/invocations/vae/flux2_vae_decode.py

@@ -21,15 +21,17 @@
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.model_manager.load.load_base import LoadedModel
 from invokeai.backend.util.devices import TorchDevice
-from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2
+from invokeai.backend.util.oom import is_oom_error
+from invokeai.backend.util.vae_tiling_scope import MIN_TILE_SAMPLE_SIZE, scoped_vae_tiling
+from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2, should_pretile_vae_decode
 
 
 @invocation(
     "flux2_vae_decode",
     title="Latents to Image - FLUX2",
     tags=["latents", "image", "vae", "l2i", "flux2", "klein"],
     category="latents",
-    version="1.0.0",
+    version="1.1.0",
     classification=Classification.Prototype,
 )
 class Flux2VaeDecodeInvocation(BaseInvocation, WithMetadata, WithBoard):
@@ -43,30 +45,79 @@ class Flux2VaeDecodeInvocation(BaseInvocation, WithMetadata, WithBoard):
         description=FieldDescriptions.vae,
         input=Input.Connection,
     )
+    tiled: bool = InputField(default=False, description=FieldDescriptions.tiled)
+    # 0 is the "use the model's default" sentinel shared with the other VAE nodes: the workflow UI
+    # cannot represent None in a number field.
+    tile_size: int = InputField(
+        default=0,
+        multiple_of=8,
+        description=f"{FieldDescriptions.vae_tile_size} Values between 1 and "
+        f"{MIN_TILE_SAMPLE_SIZE} are raised to {MIN_TILE_SAMPLE_SIZE}.",
+    )
 
-    def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Image:
+    def _vae_decode(self, context: InvocationContext, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Image:
         """Decode latents to image using FLUX.2 VAE.
 
         Input latents should already be in the correct space after BN denormalization
         was applied in the denoiser. The VAE expects (B, 32, H, W) format.
         """
+        config = context.config.get()
+        use_tiling = self.tiled or config.force_tiled_decode
+
         # Decoding at FLUX.2 resolutions costs multiple GB of activations (~4.3GB at 1024x1024),
         # far above the default working memory the cache would otherwise reserve. Tell it up front so
         # it offloads enough of the (possibly still resident) transformer to leave room.
-        estimated_working_memory = estimate_vae_working_memory_flux2(
-            operation="decode", image_tensor=latents, vae=vae_info.model, device=vae_info.compute_device
-        )
+        def estimate(tiled: bool) -> int:
+            return estimate_vae_working_memory_flux2(
+                operation="decode",
+                image_tensor=latents,
+                vae=vae_info.model,
+                tile_size=self.tile_size if tiled else None,
+                device=vae_info.compute_device,
+            )
+
+        estimated_working_memory = estimate(use_tiling)
+        # Where the untiled decode would claim most of what the device keeps resident, tile before it runs: on Windows
+        # an allocation that does not fit is paged rather than refused, so the retry below never fires.
+        if (
+            not use_tiling
+            and config.auto_tiled_decode
+            and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory)
+        ):
+            use_tiling = True
+            estimated_working_memory = estimate(use_tiling)
 
         with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
             vae_dtype = next(iter(vae.parameters())).dtype
             # Use the VAE's intended compute device (CUDA/MPS, or CPU if configured cpu_only). Do NOT infer it from
             # current param residency: partial loading may have temporarily offloaded all weights to RAM, which would
             # wrongly place the latents (and thus the whole decode) on the CPU (see #9373).
-            device = vae_info.compute_device
-            latents = latents.to(device=device, dtype=vae_dtype)
-
-            # Decode using diffusers API
-            decoded = vae.decode(latents, return_dict=False)[0]
+            latents = latents.to(device=vae_info.compute_device, dtype=vae_dtype)
+
+            def decode() -> torch.Tensor:
+                return vae.decode(latents, return_dict=False)[0]
+
+            # The VAE is the cache's shared instance, so its tiling state is set for this decode only and restored.
+            try:
+                with scoped_vae_tiling(vae, self.tile_size if use_tiling else None):
+                    decoded = decode()
+            except RuntimeError as e:
+                if use_tiling or not is_oom_error(e):
+                    raise
+                # The working-memory estimate was insufficient on this system. Retry once with
+                # tiling, which caps the peak allocation regardless of resolution.
+                context.util.signal_progress("VAE decode ran out of memory, retrying tiled")
+                context.logger.warning(
+                    "VAE decode ran out of memory; retrying with tiling. The tiled result is not identical to an "
+                    "untiled decode -- the decoder's normalisation and attention are global, so the difference is "
+                    "spread over the image rather than confined to the seams."
+                )
+                # The traceback pins the failed attempt's activations; see the FLUX.1 decode node.
+                e.__traceback__ = None
+                TorchDevice.empty_cache()
+                # 0, not `self.tile_size`: the field may hold a size left over from a run where `tiled` was off.
+                with scoped_vae_tiling(vae, 0):
+                    decoded = decode()
 
         # Convert from [-1, 1] to [0, 1] then to [0, 255] PIL image
         img = (decoded / 2 + 0.5).clamp(0, 1)
@@ -96,7 +147,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
 
         vae_info = context.models.load(self.vae.vae)
         context.util.signal_progress("Running VAE")
-        image = self._vae_decode(vae_info=vae_info, latents=latents)
+        image = self._vae_decode(context=context, vae_info=vae_info, latents=latents)
 
         TorchDevice.empty_cache()
         image_dto = context.images.save(image=image)