feat(vae): tiled decode for the FLUX.1 autoencoder

invokeai/app/invocations/anima_latents_to_image.py

@@ -29,6 +29,7 @@
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEncoder
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.oom import is_oom_error
 from invokeai.backend.util.vae_working_memory import (
     estimate_vae_working_memory_anima,
     estimate_vae_working_memory_flux,
@@ -41,34 +42,6 @@
 ANIMA_VAE_TILE_STRIDE = 384
 
 
-def _is_oom_error(e: RuntimeError) -> bool:
-    """Return True if the error indicates an out-of-memory condition.
-
-    The caching allocator raises torch.cuda.OutOfMemoryError, but an OOM surfaced from inside a
-    cuDNN/cuBLAS kernel (e.g. workspace allocation in the Wan VAE's convolutions) arrives as a
-    plain RuntimeError, which must be matched by message. XPU exhaustion likewise arrives as a
-    plain RuntimeError, naming the Level Zero/UR result code (`..._OUT_OF_DEVICE_MEMORY`) rather
-    than the words "out of memory" -- so it needs its own spelling to be matched here.
-
-    `out_of_host_memory` is knowingly over-broad: Level Zero returns it for driver-side resource
-    failures generally (kernel compilation, handle exhaustion), not only host allocation. Matching
-    it means a genuinely broken decode costs one wasted tiled retry before the error re-raises
-    unchanged. That is preferred over the alternative -- a real host-memory exhaustion that skips
-    the retry -- because the retry is bounded and non-destructive, while a missed OOM fails a
-    generation that would have succeeded tiled.
-    """
-    if isinstance(e, torch.cuda.OutOfMemoryError):
-        return True
-    msg = str(e).lower()
-    return (
-        "out of memory" in msg
-        or "out_of_device_memory" in msg
-        or "out_of_host_memory" in msg
-        or "cudnn_status_alloc_failed" in msg
-        or "cublas_status_alloc_failed" in msg
-    )
-
-
 @invocation(
     "anima_l2i",
     title="Latents to Image - Anima",
@@ -181,7 +154,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
                     try:
                         decoded = vae.decode(latents, return_dict=False)[0]
                     except RuntimeError as e:
-                        if use_tiling or not _is_oom_error(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.

invokeai/app/invocations/flux_vae_decode.py

@@ -18,6 +18,8 @@
 from invokeai.backend.flux.modules.autoencoder import AutoEncoder
 from invokeai.backend.model_manager.load.load_base import LoadedModel
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.oom import is_oom_error
+from invokeai.backend.util.vae_tiling_scope import scoped_vae_tiling
 from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux
 
 
@@ -59,10 +61,7 @@ def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Ima
             # 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)
 
-            if isinstance(vae, AutoEncoder):
-                # BFL AutoEncoder returns tensor directly
-                img = vae.decode(latents)
-            else:
+            if not isinstance(vae, AutoEncoder):
                 # Diffusers AutoencoderKL returns DecoderOutput with .sample attribute
                 # Scale latents for diffusers VAE (FLUX uses shift_factor and scale_factor).
                 # `shift_factor` is optional on AutoencoderKL: the FLUX VAE sets one, but a plain
@@ -75,7 +74,26 @@ def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Ima
                 if shift_factor is not None:
                     latents = latents + shift_factor
 
-                img = vae.decode(latents, return_dict=False)[0]
+            def decode() -> torch.Tensor:
+                if isinstance(vae, AutoEncoder):
+                    # BFL AutoEncoder returns tensor directly
+                    return vae.decode(latents)
+                return vae.decode(latents, return_dict=False)[0]
+
+            # This node has no tiling controls, so it decodes untiled -- but says so explicitly
+            # rather than inheriting whatever the last node to touch this shared, cached VAE left
+            # behind, and restores that state afterwards.
+            try:
+                with scoped_vae_tiling(vae, None):
+                    img = decode()
+            except RuntimeError as e:
+                if 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.
+                TorchDevice.empty_cache()
+                with scoped_vae_tiling(vae, 0):
+                    img = decode()
 
         img = img.clamp(-1, 1)
         img = rearrange(img[0], "c h w -> h w c")  # noqa: F821

invokeai/app/invocations/z_image_latents_to_image.py

@@ -21,6 +21,8 @@
 from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEncoder
 from invokeai.backend.stable_diffusion.extensions.seamless import SeamlessExt
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.oom import is_oom_error
+from invokeai.backend.util.vae_tiling_scope import scoped_vae_tiling
 from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux
 
 # Z-Image can use either the Diffusers AutoencoderKL or the FLUX AutoEncoder
@@ -32,14 +34,19 @@
     title="Latents to Image - Z-Image",
     tags=["latents", "image", "vae", "l2i", "z-image"],
     category="latents",
-    version="1.1.0",
+    version="1.2.0",
     classification=Classification.Prototype,
 )
 class ZImageLatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Generates an image from latents using Z-Image VAE (supports both Diffusers and FLUX 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. We use this rather than `int | None`, because the workflow UI does not
+    # offer a way to directly set None values. The size applies to InvokeAI's FLUX AutoEncoder; a diffusers
+    # AutoencoderKL tiles with its own geometry, which it does not expose as a single settable size.
+    tile_size: int = InputField(default=0, multiple_of=8, description=FieldDescriptions.vae_tile_size)
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ImageOutput:
@@ -53,12 +60,14 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
             )
 
         is_flux_vae = isinstance(vae_info.model, FluxAutoEncoder)
+        use_tiling = self.tiled or context.config.get().force_tiled_decode
 
         # Estimate working memory needed for VAE decode
         estimated_working_memory = estimate_vae_working_memory_flux(
             operation="decode",
             image_tensor=latents,
             vae=vae_info.model,
+            tile_size=self.tile_size if use_tiling else None,
         )
 
         # FLUX VAE doesn't support seamless, so only apply for AutoencoderKL
@@ -80,28 +89,41 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
             # 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)
 
-            # Disable tiling for AutoencoderKL
-            if isinstance(vae, AutoencoderKL):
-                vae.disable_tiling()
-
             # Clear memory as VAE decode can request a lot
             TorchDevice.empty_cache()
 
-            with torch.inference_mode():
+            if not isinstance(vae, FluxAutoEncoder):
+                # AutoencoderKL - Apply scaling_factor and shift_factor from VAE config
+                # Z-Image uses: latents = latents / scaling_factor + shift_factor
+                # (the FLUX VAE handles scaling internally)
+                scaling_factor = vae.config.scaling_factor
+                shift_factor = getattr(vae.config, "shift_factor", None)
+
+                latents = latents / scaling_factor
+                if shift_factor is not None:
+                    latents = latents + shift_factor
+
+            def decode() -> torch.Tensor:
                 if isinstance(vae, FluxAutoEncoder):
-                    # FLUX VAE handles scaling internally
-                    img = vae.decode(latents)
-                else:
-                    # AutoencoderKL - Apply scaling_factor and shift_factor from VAE config
-                    # Z-Image uses: latents = latents / scaling_factor + shift_factor
-                    scaling_factor = vae.config.scaling_factor
-                    shift_factor = getattr(vae.config, "shift_factor", None)
-
-                    latents = latents / scaling_factor
-                    if shift_factor is not None:
-                        latents = latents + shift_factor
-
-                    img = vae.decode(latents, return_dict=False)[0]
+                    return vae.decode(latents)
+                return vae.decode(latents, return_dict=False)[0]
+
+            # The VAE belongs to the model cache and is shared with every other node that reaches
+            # this class -- FLUX.1 decode and encode, Anima, PiD. Tiling is a property of this one
+            # decode, not of the model, so the state is scoped and restored rather than left behind.
+            with torch.inference_mode():
+                try:
+                    with scoped_vae_tiling(vae, self.tile_size if use_tiling else None):
+                        img = 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")
+                    TorchDevice.empty_cache()
+                    with scoped_vae_tiling(vae, self.tile_size):
+                        img = decode()
 
             img = img.clamp(-1, 1)
             img = rearrange(img[0], "c h w -> h w c")

invokeai/backend/flux/modules/autoencoder.py

@@ -2,10 +2,41 @@
 
 from dataclasses import dataclass
 
+import numpy as np
 import torch
 from einops import rearrange
 from torch import Tensor, nn
 
+from invokeai.backend.tiles.tiles import calc_tiles_min_overlap, merge_tiles_with_linear_blending
+from invokeai.backend.tiles.utils import TBLR, Tile
+
+# Tile geometry for tiled decode, in output-pixel units. 512px tiles with a 128px minimum overlap is
+# the geometry the diffusers VAEs and the Anima node use. Both values must be divisible by the
+# autoencoder's spatial compression factor so that a pixel-space tile maps onto an exact latent slice.
+DEFAULT_TILE_SAMPLE_MIN_SIZE = 512
+DEFAULT_TILE_OVERLAP = 128
+
+# A cost floor, not a correctness one: the geometry stays valid all the way down, but the tile count
+# grows with the inverse square of the tile size. At 2048x2048 a 128px tile already emits 289 tiles;
+# a 16px tile would emit ~16k, and the per-tile kernel-launch overhead dominates long before that.
+# Small tiles are also measurably less accurate (see `enable_tiling`), so the low end of the node
+# field is clamped rather than honoured literally.
+MIN_TILE_SAMPLE_SIZE = 128
+
+
+def resolve_tile_size(tile_size: int) -> int:
+    """Resolve a node's ``tile_size`` field to the size the autoencoder 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, and a negative value is not worth failing a generation
+    over. It resolves to the module-level default rather than to whatever is currently set on the
+    VAE: the instance belongs to the model cache, so reading it back would return whatever the
+    previous invocation left there.
+    """
+    if tile_size <= 0:
+        return DEFAULT_TILE_SAMPLE_MIN_SIZE
+    return max(tile_size, MIN_TILE_SAMPLE_SIZE)
+
 
 @dataclass
 class AutoEncoderParams:
@@ -298,6 +329,55 @@ def __init__(self, params: AutoEncoderParams):
         self.scale_factor = params.scale_factor
         self.shift_factor = params.shift_factor
 
+        # Each level of `ch_mult` past the first halves the spatial resolution, so this is the ratio
+        # between output pixels and latent elements along one axis (8 for the FLUX.1 autoencoder).
+        self.spatial_compression = 2 ** (len(params.ch_mult) - 1)
+
+        self.use_tiling = False
+        self.tile_sample_min_size = DEFAULT_TILE_SAMPLE_MIN_SIZE
+        self.tile_overlap = DEFAULT_TILE_OVERLAP
+
+    def enable_tiling(
+        self,
+        tile_sample_min_size: int = DEFAULT_TILE_SAMPLE_MIN_SIZE,
+        tile_overlap: int | None = None,
+    ) -> None:
+        """Decode in overlapping tiles, bounding peak memory at the cost of some decode time.
+
+        Mirrors the `enable_tiling()` / `disable_tiling()` pair on the diffusers autoencoders so that
+        callers can set the tiling state the same way regardless of which VAE class they hold. Sizes
+        are in output pixels.
+
+        `tile_overlap` defaults to DEFAULT_TILE_OVERLAP, shrunk to half the tile if the caller asked
+        for a tile that small. The alternative -- raising -- would turn a tile size the workflow UI
+        lets a user type into a failed generation.
+
+        Note on accuracy: at the default 512/128 geometry a tiled decode reproduces the single-pass
+        one exactly (float32 epsilon, measured). It degrades as tiles get small relative to the
+        image, because more tiles mean the blend bands sit closer to the tiles' own zero-padded
+        borders. Prefer a large tile that fits over a small one that fits comfortably.
+        """
+        if tile_overlap is None:
+            tile_overlap = min(DEFAULT_TILE_OVERLAP, tile_sample_min_size // 2)
+            tile_overlap -= tile_overlap % self.spatial_compression
+        if tile_sample_min_size % self.spatial_compression != 0:
+            raise ValueError(
+                f"tile_sample_min_size must be divisible by {self.spatial_compression}, got {tile_sample_min_size}."
+            )
+        if tile_overlap % self.spatial_compression != 0:
+            raise ValueError(f"tile_overlap must be divisible by {self.spatial_compression}, got {tile_overlap}.")
+        if tile_overlap >= tile_sample_min_size:
+            raise ValueError(
+                f"tile_overlap ({tile_overlap}) must be smaller than tile_sample_min_size ({tile_sample_min_size})."
+            )
+        self.use_tiling = True
+        self.tile_sample_min_size = tile_sample_min_size
+        self.tile_overlap = tile_overlap
+
+    def disable_tiling(self) -> None:
+        """Decode in a single pass. The inverse of `enable_tiling()`."""
+        self.use_tiling = False
+
     def encode(self, x: Tensor, sample: bool = True, generator: torch.Generator | None = None) -> Tensor:
         """Run VAE encoding on input tensor x.
 
@@ -318,7 +398,86 @@ def encode(self, x: Tensor, sample: bool = True, generator: torch.Generator | No
 
     def decode(self, z: Tensor) -> Tensor:
         z = z / self.scale_factor + self.shift_factor
+        if self.use_tiling:
+            return self._tiled_decode(z)
         return self.decoder(z)
 
+    def _tiled_decode(self, z: Tensor) -> Tensor:
+        """Decode `z` as overlapping tiles, blended back together linearly.
+
+        `z` is expected to already be denormalised, i.e. this consumes what `decode()` hands to
+        `self.decoder`. Peak memory is bounded by one tile plus the destination image, because each
+        finished tile is moved to the CPU before the next one is decoded.
+
+        The tile layout is computed in *latent* space and scaled up afterwards. Computing it in pixel
+        space would be wrong: `calc_tiles_min_overlap` distributes the leftover with integer
+        division, so it hands back tile edges that are not multiples of `spatial_compression` and
+        therefore cannot be sliced out of `z`. Overlaps scale with the coordinates because they are
+        nothing but coordinate differences.
+        """
+        scale = self.spatial_compression
+        latent_tile_size = self.tile_sample_min_size // scale
+        latent_overlap = self.tile_overlap // scale
+        latent_height, latent_width = z.shape[-2], z.shape[-1]
+
+        # Nothing to gain from tiling something that already fits in a single tile.
+        if latent_height <= latent_tile_size and latent_width <= latent_tile_size:
+            return self.decoder(z)
+
+        latent_tiles = calc_tiles_min_overlap(
+            image_height=latent_height,
+            image_width=latent_width,
+            tile_height=latent_tile_size,
+            tile_width=latent_tile_size,
+            min_overlap=latent_overlap,
+        )
+        pixel_tiles = [
+            Tile(
+                coords=TBLR(
+                    top=t.coords.top * scale,
+                    bottom=t.coords.bottom * scale,
+                    left=t.coords.left * scale,
+                    right=t.coords.right * scale,
+                ),
+                overlap=TBLR(
+                    top=t.overlap.top * scale,
+                    bottom=t.overlap.bottom * scale,
+                    left=t.overlap.left * scale,
+                    right=t.overlap.right * scale,
+                ),
+            )
+            for t in latent_tiles
+        ]
+
+        out_channels = self.decoder.conv_out.out_channels
+        batch_images: list[Tensor] = []
+        for batch_idx in range(z.shape[0]):
+            tile_images: list[np.ndarray] = []
+            for latent_tile in latent_tiles:
+                latent_slice = z[
+                    batch_idx : batch_idx + 1,
+                    :,
+                    latent_tile.coords.top : latent_tile.coords.bottom,
+                    latent_tile.coords.left : latent_tile.coords.right,
+                ]
+                decoded_tile = self.decoder(latent_slice)
+                # Off the GPU immediately -- holding the finished tiles on the device is the thing
+                # tiling exists to avoid.
+                tile_images.append(decoded_tile[0].permute(1, 2, 0).float().cpu().numpy())
+
+            merged = np.zeros(
+                (latent_height * scale, latent_width * scale, out_channels),
+                dtype=tile_images[0].dtype,
+            )
+            merge_tiles_with_linear_blending(
+                dst_image=merged,
+                tiles=pixel_tiles,
+                tile_images=tile_images,
+                blend_amount=self.tile_overlap,
+            )
+            batch_images.append(torch.from_numpy(merged).permute(2, 0, 1))
+
+        return torch.stack(batch_images).to(device=z.device, dtype=z.dtype)
+
     def forward(self, x: Tensor) -> Tensor:
         return self.decode(self.encode(x))