From 904b775b228746dd4448bd8364abe6c6f726a896 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Fri, 28 Aug 2026 05:51:48 +0200 Subject: [PATCH 1/2] feat(vae): tiled decode for the FLUX.1 autoencoder InvokeAI's own FLUX.1 AutoEncoder had no tiling, so a decode that did not fit in VRAM failed instead of degrading. Nine call sites share that VAE -- FLUX.1, Z-Image, Anima and the PiD nodes -- and Z-Image is where it surfaces, because its GGUF bundles depend on this VAE and small-VRAM users reach for exactly that combination. Tiling goes on the class rather than into one node, so all nine are covered and the nodes can set the tiling state in the same enable_tiling() / disable_tiling() spelling they already use for the diffusers VAEs. 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 and hands back tile edges that are not multiples of the compression factor, which cannot be sliced out of the latent. Measured: with the decoder's two globally-scoped operators removed (the mid-block attention and the GroupNorms, both of which make a tiled decode differ by construction), a tiled decode reproduces the single-pass one to 1.0e-07 -- seams included. Accuracy degrades with smaller tiles, not with smaller overlaps: at 96x96 latents, 512px tiles land at 1.0e-07 while 256px tiles at the same 128px overlap drift to 4.4e-03. Prefer a large tile that fits over a small one that fits comfortably. Also here: - estimate_vae_working_memory_flux takes an optional tile_size, defaulting to None so the six existing call sites are unchanged. tile_size=0 resolves via getattr, because the Z-Image nodes also pass a diffusers AutoencoderKL to this estimator. - The Z-Image decode node gets tiled/tile_size fields, honours force_tiled_decode, sets the shared VAE's tiling state explicitly on every run, and retries once tiled on OOM. Node version 1.1.0 -> 1.2.0. - flux_vae_decode gets the OOM retry only -- no new fields, no version bump. - _is_oom_error moves out of the Anima node into backend/util/oom.py, so the next backend's spelling cannot be added to one copy only. - The three stray debug prints in vae_working_memory.py are gone. Encode stays untiled deliberately: it peaks at roughly half of decode and would need its own tiled encode. --- .../app/invocations/anima_latents_to_image.py | 31 +-- invokeai/app/invocations/flux_vae_decode.py | 26 ++- .../invocations/z_image_latents_to_image.py | 64 ++++-- invokeai/backend/flux/modules/autoencoder.py | 138 ++++++++++++ invokeai/backend/util/oom.py | 34 +++ invokeai/backend/util/vae_working_memory.py | 41 ++-- invokeai/frontend/web/openapi.json | 23 +- .../frontend/web/src/services/api/schema.ts | 12 + .../invocations/test_z_image_tiled_decode.py | 179 +++++++++++++++ .../flux/modules/test_autoencoder_tiling.py | 208 ++++++++++++++++++ 10 files changed, 691 insertions(+), 65 deletions(-) create mode 100644 invokeai/backend/util/oom.py create mode 100644 tests/app/invocations/test_z_image_tiled_decode.py create mode 100644 tests/backend/flux/modules/test_autoencoder_tiling.py diff --git a/invokeai/app/invocations/anima_latents_to_image.py b/invokeai/app/invocations/anima_latents_to_image.py index 18b5db028d5..132f5a6c5e2 100644 --- a/invokeai/app/invocations/anima_latents_to_image.py +++ b/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. diff --git a/invokeai/app/invocations/flux_vae_decode.py b/invokeai/app/invocations/flux_vae_decode.py index 8caa3728471..70fc7fbb782 100644 --- a/invokeai/app/invocations/flux_vae_decode.py +++ b/invokeai/app/invocations/flux_vae_decode.py @@ -18,6 +18,7 @@ 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_working_memory import estimate_vae_working_memory_flux @@ -59,10 +60,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 +73,25 @@ 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 the tiling state of the shared, cached VAE + # instance is set explicitly rather than inherited from whoever decoded last. + vae.disable_tiling() + try: + 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() + vae.enable_tiling() + img = decode() img = img.clamp(-1, 1) img = rearrange(img[0], "c h w -> h w c") # noqa: F821 diff --git a/invokeai/app/invocations/z_image_latents_to_image.py b/invokeai/app/invocations/z_image_latents_to_image.py index 3ad07de4b7d..88223a50ce1 100644 --- a/invokeai/app/invocations/z_image_latents_to_image.py +++ b/invokeai/app/invocations/z_image_latents_to_image.py @@ -21,6 +21,7 @@ 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_working_memory import estimate_vae_working_memory_flux # Z-Image can use either the Diffusers AutoencoderKL or the FLUX AutoEncoder @@ -32,7 +33,7 @@ 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): @@ -40,6 +41,11 @@ class ZImageLatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard): 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 +59,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 +88,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() + # The VAE instance is cached and shared across invocations, so the tiling state is always + # set explicitly -- otherwise one tiled run would leave every later run tiled. + self._set_tiling(vae, enabled=use_tiling) # Clear memory as VAE decode can request a lot TorchDevice.empty_cache() - with torch.inference_mode(): - 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) + 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 - latents = latents / scaling_factor - if shift_factor is not None: - latents = latents + shift_factor + def decode() -> torch.Tensor: + if isinstance(vae, FluxAutoEncoder): + return vae.decode(latents) + return vae.decode(latents, return_dict=False)[0] - img = vae.decode(latents, return_dict=False)[0] + with torch.inference_mode(): + try: + 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() + self._set_tiling(vae, enabled=True) + img = decode() img = img.clamp(-1, 1) img = rearrange(img[0], "c h w -> h w c") @@ -112,3 +133,12 @@ def invoke(self, context: InvocationContext) -> ImageOutput: image_dto = context.images.save(image=img_pil) return ImageOutput.build(image_dto) + + def _set_tiling(self, vae: ZImageVAE, enabled: bool) -> None: + """Set the VAE's tiling state explicitly, in whichever class's spelling applies.""" + if not enabled: + vae.disable_tiling() + elif isinstance(vae, FluxAutoEncoder) and self.tile_size: + vae.enable_tiling(tile_sample_min_size=self.tile_size) + else: + vae.enable_tiling() diff --git a/invokeai/backend/flux/modules/autoencoder.py b/invokeai/backend/flux/modules/autoencoder.py index 6b072a82f63..01b093fd7cf 100644 --- a/invokeai/backend/flux/modules/autoencoder.py +++ b/invokeai/backend/flux/modules/autoencoder.py @@ -2,10 +2,20 @@ 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 + @dataclass class AutoEncoderParams: @@ -298,6 +308,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 +377,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)) diff --git a/invokeai/backend/util/oom.py b/invokeai/backend/util/oom.py new file mode 100644 index 00000000000..0352f0fdb30 --- /dev/null +++ b/invokeai/backend/util/oom.py @@ -0,0 +1,34 @@ +"""Recognising an out-of-memory failure across backends.""" + +import torch + + +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 a 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. + + This lives in one place on purpose: a second copy is how the next backend's spelling gets added + to one of them only. + """ + 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 + ) diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 9970c6b5d96..81d002c61ba 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -7,7 +7,7 @@ from diffusers.models.autoencoders.autoencoder_tiny import AutoencoderTiny from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR -from invokeai.backend.flux.modules.autoencoder import AutoEncoder +from invokeai.backend.flux.modules.autoencoder import DEFAULT_TILE_SAMPLE_MIN_SIZE, AutoEncoder _WAN_VAE_SINGLE_FRAME_DECODE_SCALING_CONSTANT = 2900 _WAN_VAE_VIDEO_DECODE_SCALING_CONSTANT_A14B = 6500 @@ -71,29 +71,46 @@ def estimate_vae_working_memory_cogview4( scaling_constant = 2200 if operation == "decode" else 1100 working_memory = h * w * element_size * scaling_constant - print(f"estimate_vae_working_memory_cogview4: {int(working_memory)}") - return int(working_memory) def estimate_vae_working_memory_flux( - operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoEncoder + operation: Literal["encode", "decode"], + image_tensor: torch.Tensor, + vae: AutoEncoder, + tile_size: int | None = None, ) -> int: - """Estimate the working memory required by the invocation in bytes.""" + """Estimate the working memory required by the invocation in bytes. - latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 + `tile_size` is in output pixels and defaults to None, i.e. a single-pass decode -- the six + existing call sites depend on that signature. When set, the estimate is bounded by one tile + instead of the whole image, because a tiled decode never holds more than that. `tile_size=0` + means "whatever the VAE's own default is"; a VAE that has no such default falls back to the + autoencoder's tile size. + """ - out_h = latent_scale_factor_for_operation * image_tensor.shape[-2] - out_w = latent_scale_factor_for_operation * image_tensor.shape[-1] + latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 element_size = next(vae.parameters()).element_size() # This constant is determined experimentally and takes into consideration both allocated and reserved memory. See #8414 # Encoding uses ~45% the working memory as decoding. scaling_constant = 2200 if operation == "decode" else 1100 - working_memory = out_h * out_w * element_size * scaling_constant - - print(f"estimate_vae_working_memory_flux: {int(working_memory)}") + if tile_size is not None: + if tile_size == 0: + # Not every VAE reaching this estimator is an InvokeAI AutoEncoder -- the Z-Image nodes + # also pass a diffusers AutoencoderKL -- so this cannot dereference the attribute the way + # estimate_vae_working_memory_sd15_sdxl does. + tile_size = getattr(vae, "tile_sample_min_size", DEFAULT_TILE_SAMPLE_MIN_SIZE) + assert isinstance(tile_size, int) + out_h = tile_size + out_w = tile_size + # A 25% margin for tile overlap and the number of tiles, mirroring the SD1/SDXL estimator. + working_memory = out_h * out_w * element_size * scaling_constant * 1.25 + else: + out_h = latent_scale_factor_for_operation * image_tensor.shape[-2] + out_w = latent_scale_factor_for_operation * image_tensor.shape[-1] + working_memory = out_h * out_w * element_size * scaling_constant return int(working_memory) @@ -352,6 +369,4 @@ def estimate_vae_working_memory_sd3( working_memory = h * w * element_size * scaling_constant - print(f"estimate_vae_working_memory_sd3: {int(working_memory)}") - return int(working_memory) diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 2e75b0f0569..1e71f2e893e 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -98979,6 +98979,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": "z_image_l2i", "default": "z_image_l2i", @@ -98991,7 +99012,7 @@ "tags": ["latents", "image", "vae", "l2i", "z-image"], "title": "Latents to Image - Z-Image", "type": "object", - "version": "1.1.0", + "version": "1.2.0", "output": { "$ref": "#/components/schemas/ImageOutput" } diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index ed1d063ce18..d610549dbc0 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -44633,6 +44633,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 z_image_l2i diff --git a/tests/app/invocations/test_z_image_tiled_decode.py b/tests/app/invocations/test_z_image_tiled_decode.py new file mode 100644 index 00000000000..9a3939ef913 --- /dev/null +++ b/tests/app/invocations/test_z_image_tiled_decode.py @@ -0,0 +1,179 @@ +"""Tiling controls, the OOM fallback, and the tile-aware working-memory estimate for Z-Image.""" + +from unittest.mock import MagicMock, patch + +import pytest +import torch +from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL + +from invokeai.app.invocations.z_image_latents_to_image import ZImageLatentsToImageInvocation +from invokeai.backend.flux.modules.autoencoder import DEFAULT_TILE_SAMPLE_MIN_SIZE +from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEncoder +from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux + + +def _mock_flux_vae(element_size_bytes: int = 2) -> MagicMock: + vae = MagicMock(spec=FluxAutoEncoder) + dtype = torch.float16 if element_size_bytes == 2 else torch.float32 + # A fresh iterator per call: the decode path reads `parameters()` after the estimator already + # has, and a single stored iterator would be exhausted by then. + vae.parameters.side_effect = lambda: iter([torch.zeros(1, dtype=dtype)]) + return vae + + +class TestFluxWorkingMemoryEstimate: + def test_the_default_reproduces_the_untiled_estimate(self): + """Regression guard for the six call sites that pass no tile_size at all.""" + latents = torch.zeros(1, 16, 128, 128) + # 1024x1024 output px * 2 bytes * 2200 + expected = 1024 * 1024 * 2 * 2200 + actual = estimate_vae_working_memory_flux(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) + assert actual == expected + + @pytest.mark.parametrize("latent_hw", [(128, 128), (192, 192), (256, 256)]) + def test_a_tiled_estimate_does_not_grow_with_the_image(self, latent_hw): + """The point of the bound: peak memory is one tile, whatever the resolution.""" + h, w = latent_hw + estimate = estimate_vae_working_memory_flux( + operation="decode", image_tensor=torch.zeros(1, 16, h, w), vae=_mock_flux_vae(), tile_size=512 + ) + assert estimate == int(512 * 512 * 2 * 2200 * 1.25) + + def test_a_tiled_estimate_is_smaller_than_the_untiled_one_where_it_matters(self): + latents = torch.zeros(1, 16, 192, 192) # 1536px, the resolution that OOMs on 16GB + untiled = estimate_vae_working_memory_flux(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) + tiled = estimate_vae_working_memory_flux( + operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=512 + ) + assert tiled < untiled + + def test_tile_size_zero_resolves_against_the_vae(self): + vae = _mock_flux_vae() + vae.tile_sample_min_size = 384 + estimate = estimate_vae_working_memory_flux( + operation="decode", image_tensor=torch.zeros(1, 16, 192, 192), vae=vae, tile_size=0 + ) + assert estimate == int(384 * 384 * 2 * 2200 * 1.25) + + def test_tile_size_zero_on_a_vae_without_a_default_does_not_raise(self): + # The Z-Image nodes also hand a diffusers AutoencoderKL to this estimator; the SD1/SDXL + # sibling dereferences `vae.tile_sample_min_size` directly and would raise here. + vae = MagicMock(spec=AutoencoderKL) + vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.float16)]) + del vae.tile_sample_min_size + estimate = estimate_vae_working_memory_flux( + operation="decode", image_tensor=torch.zeros(1, 16, 192, 192), vae=vae, tile_size=0 + ) + assert estimate == int(DEFAULT_TILE_SAMPLE_MIN_SIZE**2 * 2 * 2200 * 1.25) + + +def _build_decode_mocks(latents: torch.Tensor, decoded: torch.Tensor, force_tiled_decode: bool = False): + """Wire ZImageLatentsToImageInvocation.invoke to run end-to-end on CPU against a mocked FLUX VAE.""" + vae = _mock_flux_vae() + vae.decode.return_value = decoded + + vae_info = MagicMock() + vae_info.model = vae + vae_info.compute_device = torch.device("cpu") + cm = MagicMock() + cm.__enter__ = MagicMock(return_value=(None, vae)) + cm.__exit__ = MagicMock(return_value=None) + vae_info.model_on_device.return_value = cm + + context = MagicMock() + context.models.load.return_value = vae_info + context.tensors.load.return_value = latents + # A bare MagicMock config would read as truthy and silently tile everything. + context.config.get.return_value.force_tiled_decode = force_tiled_decode + image_dto = MagicMock() + image_dto.image_name = "test.png" + image_dto.width = decoded.shape[-1] + image_dto.height = decoded.shape[-2] + context.images.save.return_value = image_dto + return vae, vae_info, context + + +def _build_invocation(tiled: bool = False, tile_size: int = 0) -> ZImageLatentsToImageInvocation: + return ZImageLatentsToImageInvocation.model_construct( + latents=MagicMock(latents_name="test_latents"), + vae=MagicMock(vae=MagicMock()), + tiled=tiled, + tile_size=tile_size, + ) + + +class TestTilingIsWired: + def test_the_default_decodes_untiled(self): + vae, _, context = _build_decode_mocks(torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512)) + _build_invocation().invoke(context) + vae.disable_tiling.assert_called_once() + vae.enable_tiling.assert_not_called() + + def test_the_node_field_reaches_the_tiled_path(self): + vae, _, context = _build_decode_mocks(torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512)) + _build_invocation(tiled=True).invoke(context) + vae.enable_tiling.assert_called_once_with() + vae.disable_tiling.assert_not_called() + + def test_force_tiled_decode_reaches_the_tiled_path(self): + """The config switch a small-VRAM user actually has; the node field is not the only way in.""" + vae, _, context = _build_decode_mocks( + torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512), force_tiled_decode=True + ) + _build_invocation().invoke(context) + vae.enable_tiling.assert_called_once_with() + + def test_a_requested_tile_size_is_passed_through(self): + vae, _, context = _build_decode_mocks(torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512)) + _build_invocation(tiled=True, tile_size=384).invoke(context) + vae.enable_tiling.assert_called_once_with(tile_sample_min_size=384) + + @pytest.mark.parametrize("tiled,expected_tile_size", [(False, None), (True, 0)]) + def test_the_estimate_is_tile_bounded_only_when_tiling(self, tiled, expected_tile_size): + path = "invokeai.app.invocations.z_image_latents_to_image.estimate_vae_working_memory_flux" + _, _, context = _build_decode_mocks(torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512)) + with patch(path, return_value=1024) as estimate: + _build_invocation(tiled=tiled).invoke(context) + assert estimate.call_args.kwargs["tile_size"] == expected_tile_size + + +class TestOomFallback: + @pytest.mark.parametrize( + "oom_error", + [ + torch.cuda.OutOfMemoryError("CUDA out of memory. Tried to allocate 5.9 GiB"), + RuntimeError("CUDA error: out of memory"), + RuntimeError("cuDNN error: CUDNN_STATUS_ALLOC_FAILED"), + RuntimeError("Native API failed. Native API returns: UR_RESULT_ERROR_OUT_OF_DEVICE_MEMORY"), + ], + ) + def test_an_untiled_oom_retries_once_tiled(self, oom_error): + decoded = torch.zeros(1, 3, 512, 512) + vae, _, context = _build_decode_mocks(torch.zeros(1, 16, 64, 64), decoded) + vae.decode.side_effect = [oom_error, decoded] + + result = _build_invocation().invoke(context) + + assert vae.decode.call_count == 2 + vae.enable_tiling.assert_called_once_with() + assert result.width == 512 + + def test_a_non_oom_error_propagates_without_a_retry(self): + vae, _, context = _build_decode_mocks(torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512)) + vae.decode.side_effect = RuntimeError("Input type (float) and weight type (half) should be the same") + + with pytest.raises(RuntimeError, match="weight type"): + _build_invocation().invoke(context) + + assert vae.decode.call_count == 1 + vae.enable_tiling.assert_not_called() + + def test_an_oom_while_already_tiled_reraises(self): + vae, _, context = _build_decode_mocks(torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512)) + vae.decode.side_effect = torch.cuda.OutOfMemoryError("CUDA out of memory") + + with pytest.raises(torch.cuda.OutOfMemoryError): + _build_invocation(tiled=True).invoke(context) + + assert vae.decode.call_count == 1 + vae.enable_tiling.assert_called_once() diff --git a/tests/backend/flux/modules/test_autoencoder_tiling.py b/tests/backend/flux/modules/test_autoencoder_tiling.py new file mode 100644 index 00000000000..fc21697c54b --- /dev/null +++ b/tests/backend/flux/modules/test_autoencoder_tiling.py @@ -0,0 +1,208 @@ +"""Tiled decode for InvokeAI's FLUX.1 autoencoder. + +The reference numbers for tiled-vs-untiled agreement come from measurement, not taste: diffusers' +own tiling of this VAE gives max 0.082 / mean 0.0022 per pixel at 1536px on a +/-1 image. A tiled +decode of this architecture can never be exact, because the decoder's GroupNorms normalise over the +whole spatial extent and its mid-block attention is global -- both see a different input when the +image arrives in tiles. What *is* exact is the tile geometry, and `test_geometry_is_exact_without_ +the_global_operators` pins it by removing those two operators and demanding equality. +""" + +import numpy as np +import pytest +import torch +from torch import nn + +from invokeai.backend.flux.modules.autoencoder import ( + DEFAULT_TILE_OVERLAP, + DEFAULT_TILE_SAMPLE_MIN_SIZE, + AutoEncoder, + AutoEncoderParams, +) + + +def _build_autoencoder() -> AutoEncoder: + """A structurally faithful but tiny FLUX.1 autoencoder: same block layout and the same + 8x spatial compression, with the channel counts cut down so the test runs on CPU.""" + params = AutoEncoderParams( + resolution=256, + in_channels=3, + ch=32, # GroupNorm(32) means this cannot go lower + out_ch=3, + ch_mult=[1, 2, 4, 4], + num_res_blocks=1, + z_channels=16, + scale_factor=0.3611, + shift_factor=0.1159, + ) + torch.manual_seed(0) + return AutoEncoder(params).eval() + + +def _make_purely_convolutional(ae: AutoEncoder) -> None: + """Strip the two operators with an unbounded receptive field from the decoder. + + Both make a tiled decode differ from an untiled one by construction, everywhere in the image + rather than only near the seams, which is why they have to go before the geometry can be + asserted at all. + """ + ae.decoder.mid.attn_1 = nn.Identity() + for name, module in list(ae.decoder.named_modules()): + if not isinstance(module, nn.GroupNorm): + continue + parent: nn.Module = ae.decoder + *path, attr = name.split(".") + for step in path: + parent = parent[int(step)] if step.isdigit() else getattr(parent, step) + setattr(parent, attr, nn.Identity()) + + +class TestTilingState: + def test_tiling_is_off_by_default(self): + ae = _build_autoencoder() + assert ae.use_tiling is False + assert ae.tile_sample_min_size == DEFAULT_TILE_SAMPLE_MIN_SIZE + assert ae.tile_overlap == DEFAULT_TILE_OVERLAP + + def test_spatial_compression_follows_ch_mult(self): + assert _build_autoencoder().spatial_compression == 8 + + def test_disable_tiling_restores_the_untiled_result(self): + ae = _build_autoencoder() + z = torch.randn(1, 16, 96, 96) + with torch.no_grad(): + before = ae.decode(z) + ae.enable_tiling() + ae.decode(z) + ae.disable_tiling() + after = ae.decode(z) + # The VAE instance is cached and shared across invocations, so a tiled run must not leave + # the next one tiled. + assert torch.equal(before, after) + + @pytest.mark.parametrize( + "kwargs,message", + [ + ({"tile_sample_min_size": 500}, "divisible by 8"), + ({"tile_overlap": 100}, "divisible by 8"), + ({"tile_sample_min_size": 128, "tile_overlap": 128}, "must be smaller than"), + ], + ) + def test_geometry_that_cannot_be_sliced_is_rejected(self, kwargs, message): + # A tile edge that is not a multiple of the compression factor has no exact latent slice, + # and an overlap at least as large as the tile makes the layout degenerate. + with pytest.raises(ValueError, match=message): + _build_autoencoder().enable_tiling(**kwargs) + + +class TestTiledDecode: + def test_geometry_is_exact_without_the_global_operators(self): + ae = _build_autoencoder() + _make_purely_convolutional(ae) + z = torch.randn(1, 16, 96, 96) + with torch.no_grad(): + untiled = ae.decode(z) + ae.enable_tiling() + tiled = ae.decode(z) + # Purely convolutional: the slicing, the coordinate scale-up and the linear blending must + # reproduce the single-pass decode to float32 precision, seams included. + assert torch.allclose(untiled, tiled, atol=1e-6) + + @pytest.mark.parametrize( + "latent_hw", + [ + (96, 96), # tiles do not divide the image evenly + (128, 128), # 2x2 tiles, evenly divided + (100, 77), # odd on both axes + (64, 160), # tiling on one axis only + ], + ) + def test_shape_and_dtype_survive_every_layout(self, latent_hw): + ae = _build_autoencoder() + h, w = latent_hw + z = torch.randn(1, 16, h, w) + with torch.no_grad(): + untiled = ae.decode(z) + ae.enable_tiling() + tiled = ae.decode(z) + assert tiled.shape == untiled.shape == (1, 3, h * 8, w * 8) + assert tiled.dtype == untiled.dtype == z.dtype + assert tiled.device == z.device + + def test_an_image_smaller_than_one_tile_is_decoded_in_a_single_pass(self): + ae = _build_autoencoder() + z = torch.randn(1, 16, 32, 32) # 256px, well under the 512px tile + with torch.no_grad(): + untiled = ae.decode(z) + ae.enable_tiling() + tiled = ae.decode(z) + # Not merely close: there is nothing to tile, so it must be the same computation. + assert torch.equal(untiled, tiled) + + def test_batches_are_decoded_independently(self): + ae = _build_autoencoder() + _make_purely_convolutional(ae) + z = torch.randn(2, 16, 96, 96) + with torch.no_grad(): + ae.enable_tiling() + batched = ae.decode(z) + singles = torch.cat([ae.decode(z[i : i + 1]) for i in range(2)]) + assert batched.shape == (2, 3, 768, 768) + assert torch.allclose(batched, singles, atol=1e-6) + + def test_the_shipped_geometry_is_the_accurate_one(self): + """Smaller tiles are less accurate, not more -- measured, and the reason for the default. + + Halving the tile at a fixed image size multiplies the seams, and the blend bands then sit + closer to each tile's own zero-padded border. On this fixture, purely convolutional, at + 96x96 latents: 512px tiles reproduce the single-pass decode to 1.0e-07, while 256px tiles + with the same 128px overlap drift to 4.4e-03 -- four orders of magnitude worse. + """ + ae = _build_autoencoder() + _make_purely_convolutional(ae) + z = torch.randn(1, 16, 96, 96) + with torch.no_grad(): + untiled = ae.decode(z) + ae.enable_tiling(tile_sample_min_size=DEFAULT_TILE_SAMPLE_MIN_SIZE) + shipped = ae.decode(z) + ae.enable_tiling(tile_sample_min_size=256) + smaller = ae.decode(z) + + assert shipped.shape == smaller.shape == untiled.shape + assert (untiled - shipped).abs().max() < 1e-6 + assert (untiled - shipped).abs().max() < (untiled - smaller).abs().max() + + def test_a_tile_too_small_for_the_default_overlap_shrinks_it_instead_of_raising(self): + # The workflow UI lets a user type any multiple of 8 into `tile_size`; a value under the + # default 128px overlap must not turn into a failed generation. + ae = _build_autoencoder() + ae.enable_tiling(tile_sample_min_size=128) + assert ae.tile_overlap == 64 + assert ae.tile_overlap % ae.spatial_compression == 0 + with torch.no_grad(): + assert ae.decode(torch.randn(1, 16, 96, 96)).shape == (1, 3, 768, 768) + + def test_finished_tiles_do_not_stay_on_the_decode_device(self): + # Bounding the peak is the entire point: a tile is moved off the device as soon as it is + # decoded, so what the merge sees is numpy on the host. + ae = _build_autoencoder() + seen: list[type] = [] + real_merge = None + + import invokeai.backend.flux.modules.autoencoder as autoencoder_module + + real_merge = autoencoder_module.merge_tiles_with_linear_blending + + def spy(dst_image, tiles, tile_images, blend_amount): + seen.extend(type(t) for t in tile_images) + return real_merge(dst_image, tiles, tile_images, blend_amount) + + autoencoder_module.merge_tiles_with_linear_blending = spy + try: + ae.enable_tiling() + with torch.no_grad(): + ae.decode(torch.randn(1, 16, 96, 96)) + finally: + autoencoder_module.merge_tiles_with_linear_blending = real_merge + + assert seen and all(t is np.ndarray for t in seen) From 80e340d483f5e45b9fdbd07445475a251eb3c53f Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Fri, 28 Aug 2026 07:24:48 +0200 Subject: [PATCH 2/2] fix(vae): scope the FLUX autoencoder's tiling state to one decode Reviewing this against upstream #9427 -- the Qwen-Image tiling PR, which hit the same class of problem -- turned up two silent bugs here. enable_tiling() writes the geometry onto the module and disable_tiling() restores only the flag, and that module is the model cache's own instance. So a tiled Z-Image decode, or a single OOM retry, left use_tiling=True behind on the shared FLUX autoencoder. The nodes that reach that same instance but never touch the flag -- flux_vae_encode, pid_upscale, flux_pid_decode, and Anima's FLUX branch, whose disable_tiling() sits in its diffusers branch only -- would then have decoded and encoded tiled without asking. No error, no log line, just different output. scoped_vae_tiling sets the state for one block and restores every tiling attribute in a finally, on the normal, untiled and exception paths alike. It is the same shape as SD's patch_vae_tiling_params and Qwen's patch_qwen_image_vae_tiling; neither fits these classes, since the SD one is typed to AutoencoderKL/AutoencoderTiny, patches three attributes the FLUX autoencoder does not have, and leaves use_tiling to the caller. Second, the estimator resolved tile_size=0 with getattr(vae, "tile_sample_min_size", ...) -- which returns whatever the previous invocation left on the cached module rather than the default the node is asking for. It now resolves against a module-level constant, the same correction #9427 needed. Third, the node field had no lower bound: small values produced an enormous tile count, and a negative one raised ValueError in the middle of a generation. resolve_tile_size owns both the sentinel and a 128px cost floor -- a cost floor, not a validity one: the geometry stays correct all the way down, but the tile count grows with the inverse square of the tile size, and small tiles are measurably less accurate. The module was first named vae_tiling.py, which collided with the existing stable_diffusion/vae_tiling.py, and its test collided by basename with tests/backend/stable_diffusion/test_vae_tiling.py -- a collection error that aborts the whole suite. Renamed to vae_tiling_scope. Two findings from that review do not apply and were checked rather than assumed: truncation when the tile is smaller than the stride cannot happen here (the destination is preallocated at the exact output size and tiles merge into it), and compute goes the other way round -- larger tiles mean fewer tiles and are more accurate, so the existing field description is correct for this VAE. Mutations verified as caught: dropping the finally restore (6 tests), resolving the sentinel off the module again (12), removing the cost floor (4). --- invokeai/app/invocations/flux_vae_decode.py | 14 +- .../invocations/z_image_latents_to_image.py | 24 +-- invokeai/backend/flux/modules/autoencoder.py | 21 +++ invokeai/backend/util/vae_tiling_scope.py | 62 ++++++++ invokeai/backend/util/vae_working_memory.py | 19 +-- .../invocations/test_z_image_tiled_decode.py | 37 +++-- tests/backend/util/test_vae_tiling_scope.py | 147 ++++++++++++++++++ 7 files changed, 281 insertions(+), 43 deletions(-) create mode 100644 invokeai/backend/util/vae_tiling_scope.py create mode 100644 tests/backend/util/test_vae_tiling_scope.py diff --git a/invokeai/app/invocations/flux_vae_decode.py b/invokeai/app/invocations/flux_vae_decode.py index 70fc7fbb782..6748750e8a2 100644 --- a/invokeai/app/invocations/flux_vae_decode.py +++ b/invokeai/app/invocations/flux_vae_decode.py @@ -19,6 +19,7 @@ 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 @@ -79,19 +80,20 @@ def decode() -> torch.Tensor: return vae.decode(latents) return vae.decode(latents, return_dict=False)[0] - # This node has no tiling controls, so the tiling state of the shared, cached VAE - # instance is set explicitly rather than inherited from whoever decoded last. - vae.disable_tiling() + # 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: - img = decode() + 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() - vae.enable_tiling() - img = decode() + 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 diff --git a/invokeai/app/invocations/z_image_latents_to_image.py b/invokeai/app/invocations/z_image_latents_to_image.py index 88223a50ce1..f2c53cb0c83 100644 --- a/invokeai/app/invocations/z_image_latents_to_image.py +++ b/invokeai/app/invocations/z_image_latents_to_image.py @@ -22,6 +22,7 @@ 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 @@ -88,10 +89,6 @@ 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) - # The VAE instance is cached and shared across invocations, so the tiling state is always - # set explicitly -- otherwise one tiled run would leave every later run tiled. - self._set_tiling(vae, enabled=use_tiling) - # Clear memory as VAE decode can request a lot TorchDevice.empty_cache() @@ -111,9 +108,13 @@ def decode() -> torch.Tensor: 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: - img = decode() + 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 @@ -121,8 +122,8 @@ def decode() -> torch.Tensor: # tiling, which caps the peak allocation regardless of resolution. context.util.signal_progress("VAE decode ran out of memory, retrying tiled") TorchDevice.empty_cache() - self._set_tiling(vae, enabled=True) - img = decode() + 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") @@ -133,12 +134,3 @@ def decode() -> torch.Tensor: image_dto = context.images.save(image=img_pil) return ImageOutput.build(image_dto) - - def _set_tiling(self, vae: ZImageVAE, enabled: bool) -> None: - """Set the VAE's tiling state explicitly, in whichever class's spelling applies.""" - if not enabled: - vae.disable_tiling() - elif isinstance(vae, FluxAutoEncoder) and self.tile_size: - vae.enable_tiling(tile_sample_min_size=self.tile_size) - else: - vae.enable_tiling() diff --git a/invokeai/backend/flux/modules/autoencoder.py b/invokeai/backend/flux/modules/autoencoder.py index 01b093fd7cf..86cb7b78f66 100644 --- a/invokeai/backend/flux/modules/autoencoder.py +++ b/invokeai/backend/flux/modules/autoencoder.py @@ -16,6 +16,27 @@ 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: diff --git a/invokeai/backend/util/vae_tiling_scope.py b/invokeai/backend/util/vae_tiling_scope.py new file mode 100644 index 00000000000..6f76f2dade2 --- /dev/null +++ b/invokeai/backend/util/vae_tiling_scope.py @@ -0,0 +1,62 @@ +"""Scoped VAE tiling state. + +The VAE instances these helpers take belong to the model cache and are shared across invocations and +across *nodes*: the FLUX.1 autoencoder is reached by nine call sites, and a diffusers AutoencoderKL +loaded for Z-Image is reached by several more. Tiling is a property of one decode, not of the model, +so it has to be restored rather than merely turned off -- `disable_tiling()` clears the flag but +leaves the geometry, and most of the consumers never touch the flag at all. + +Two siblings solve the same problem for their own classes: `patch_qwen_image_vae_tiling` for the +Qwen-Image VAE, and `stable_diffusion.vae_tiling.patch_vae_tiling_params` for SD's. Neither fits +here -- the SD one is typed to AutoencoderKL/AutoencoderTiny, patches three diffusers-specific +attributes the FLUX autoencoder does not have, and leaves `use_tiling` to the caller. +""" + +from contextlib import contextmanager +from typing import Any, Iterator + +from invokeai.backend.flux.modules.autoencoder import AutoEncoder, resolve_tile_size + +# Attributes that carry tiling state, across the VAE classes that reach these nodes: InvokeAI's FLUX +# AutoEncoder and diffusers' AutoencoderKL. Read defensively -- a class that has none of them simply +# has nothing to restore. +_TILING_ATTRS = ( + "use_tiling", + "tile_sample_min_size", + "tile_overlap", + "tile_latent_min_size", + "tile_overlap_factor", +) + +_MISSING = object() + + +@contextmanager +def scoped_vae_tiling(vae: Any, tile_size: int | None) -> Iterator[None]: + """Set the VAE's tiling state for the duration of the block, then restore exactly what was there. + + `tile_size=None` decodes in a single pass; `0` means "the VAE's own default"; any other value is + the tile size in output pixels. + """ + original = {name: getattr(vae, name, _MISSING) for name in _TILING_ATTRS} + try: + if tile_size is None: + vae.disable_tiling() + elif _accepts_tile_size(vae): + # resolve_tile_size owns the 0/negative sentinel and the cost floor, so every caller + # gets the same answer for the same field value. + vae.enable_tiling(tile_sample_min_size=resolve_tile_size(tile_size)) + else: + # This class does not expose a settable size -- diffusers' AutoencoderKL.enable_tiling() + # takes no arguments and uses its own geometry. + vae.enable_tiling() + yield + finally: + for name, value in original.items(): + if value is not _MISSING: + setattr(vae, name, value) + + +def _accepts_tile_size(vae: Any) -> bool: + """True if `enable_tiling` takes a tile size. Diffusers' AutoencoderKL takes no arguments.""" + return isinstance(vae, AutoEncoder) diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 81d002c61ba..513aaf34fa8 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -7,7 +7,7 @@ from diffusers.models.autoencoders.autoencoder_tiny import AutoencoderTiny from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR -from invokeai.backend.flux.modules.autoencoder import DEFAULT_TILE_SAMPLE_MIN_SIZE, AutoEncoder +from invokeai.backend.flux.modules.autoencoder import AutoEncoder, resolve_tile_size _WAN_VAE_SINGLE_FRAME_DECODE_SCALING_CONSTANT = 2900 _WAN_VAE_VIDEO_DECODE_SCALING_CONSTANT_A14B = 6500 @@ -84,9 +84,8 @@ def estimate_vae_working_memory_flux( `tile_size` is in output pixels and defaults to None, i.e. a single-pass decode -- the six existing call sites depend on that signature. When set, the estimate is bounded by one tile - instead of the whole image, because a tiled decode never holds more than that. `tile_size=0` - means "whatever the VAE's own default is"; a VAE that has no such default falls back to the - autoencoder's tile size. + instead of the whole image, because a tiled decode never holds more than that. `tile_size <= 0` + is the nodes' "use the default" sentinel; see `resolve_tile_size`. """ latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 @@ -97,14 +96,10 @@ def estimate_vae_working_memory_flux( scaling_constant = 2200 if operation == "decode" else 1100 if tile_size is not None: - if tile_size == 0: - # Not every VAE reaching this estimator is an InvokeAI AutoEncoder -- the Z-Image nodes - # also pass a diffusers AutoencoderKL -- so this cannot dereference the attribute the way - # estimate_vae_working_memory_sd15_sdxl does. - tile_size = getattr(vae, "tile_sample_min_size", DEFAULT_TILE_SAMPLE_MIN_SIZE) - assert isinstance(tile_size, int) - out_h = tile_size - out_w = tile_size + # Resolved against a module-level constant, never by reading `vae.tile_sample_min_size`: the + # VAE belongs to the model cache, so that attribute reflects whatever the previous + # invocation set rather than the default this node is asking for. + out_h = out_w = resolve_tile_size(tile_size) # A 25% margin for tile overlap and the number of tiles, mirroring the SD1/SDXL estimator. working_memory = out_h * out_w * element_size * scaling_constant * 1.25 else: diff --git a/tests/app/invocations/test_z_image_tiled_decode.py b/tests/app/invocations/test_z_image_tiled_decode.py index 9a3939ef913..255cc88ac3b 100644 --- a/tests/app/invocations/test_z_image_tiled_decode.py +++ b/tests/app/invocations/test_z_image_tiled_decode.py @@ -7,7 +7,7 @@ from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL from invokeai.app.invocations.z_image_latents_to_image import ZImageLatentsToImageInvocation -from invokeai.backend.flux.modules.autoencoder import DEFAULT_TILE_SAMPLE_MIN_SIZE +from invokeai.backend.flux.modules.autoencoder import DEFAULT_TILE_SAMPLE_MIN_SIZE, MIN_TILE_SAMPLE_SIZE from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEncoder from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux @@ -47,25 +47,34 @@ def test_a_tiled_estimate_is_smaller_than_the_untiled_one_where_it_matters(self) ) assert tiled < untiled - def test_tile_size_zero_resolves_against_the_vae(self): + def test_the_sentinel_does_not_read_the_size_off_the_vae(self): + """Upstream #9427 found this exact shape of bug in the Qwen estimator: reading + `vae.tile_sample_min_size` returns whatever the *previous* invocation left on the cached + module, not the default this node is asking for.""" vae = _mock_flux_vae() - vae.tile_sample_min_size = 384 + vae.tile_sample_min_size = 384 # as if a previous run had set it estimate = estimate_vae_working_memory_flux( operation="decode", image_tensor=torch.zeros(1, 16, 192, 192), vae=vae, tile_size=0 ) - assert estimate == int(384 * 384 * 2 * 2200 * 1.25) + assert estimate == int(DEFAULT_TILE_SAMPLE_MIN_SIZE**2 * 2 * 2200 * 1.25) - def test_tile_size_zero_on_a_vae_without_a_default_does_not_raise(self): + def test_the_sentinel_works_on_a_vae_without_that_attribute_at_all(self): # The Z-Image nodes also hand a diffusers AutoencoderKL to this estimator; the SD1/SDXL # sibling dereferences `vae.tile_sample_min_size` directly and would raise here. vae = MagicMock(spec=AutoencoderKL) - vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.float16)]) + vae.parameters.side_effect = lambda: iter([torch.zeros(1, dtype=torch.float16)]) del vae.tile_sample_min_size estimate = estimate_vae_working_memory_flux( operation="decode", image_tensor=torch.zeros(1, 16, 192, 192), vae=vae, tile_size=0 ) assert estimate == int(DEFAULT_TILE_SAMPLE_MIN_SIZE**2 * 2 * 2200 * 1.25) + def test_a_tile_below_the_cost_floor_is_estimated_at_the_floor(self): + estimate = estimate_vae_working_memory_flux( + operation="decode", image_tensor=torch.zeros(1, 16, 192, 192), vae=_mock_flux_vae(), tile_size=8 + ) + assert estimate == int(MIN_TILE_SAMPLE_SIZE**2 * 2 * 2200 * 1.25) + def _build_decode_mocks(latents: torch.Tensor, decoded: torch.Tensor, force_tiled_decode: bool = False): """Wire ZImageLatentsToImageInvocation.invoke to run end-to-end on CPU against a mocked FLUX VAE.""" @@ -109,10 +118,20 @@ def test_the_default_decodes_untiled(self): vae.disable_tiling.assert_called_once() vae.enable_tiling.assert_not_called() + def test_the_tiling_state_is_restored_after_the_invocation(self): + """The VAE belongs to the model cache and is shared with nodes that never touch the flag -- + FLUX.1 encode, PiD, Anima's FLUX branch. A tiled run must not leave them tiled.""" + vae, _, context = _build_decode_mocks(torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512)) + vae.use_tiling = False + vae.tile_sample_min_size = DEFAULT_TILE_SAMPLE_MIN_SIZE + _build_invocation(tiled=True, tile_size=256).invoke(context) + assert vae.use_tiling is False + assert vae.tile_sample_min_size == DEFAULT_TILE_SAMPLE_MIN_SIZE + def test_the_node_field_reaches_the_tiled_path(self): vae, _, context = _build_decode_mocks(torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512)) _build_invocation(tiled=True).invoke(context) - vae.enable_tiling.assert_called_once_with() + vae.enable_tiling.assert_called_once_with(tile_sample_min_size=DEFAULT_TILE_SAMPLE_MIN_SIZE) vae.disable_tiling.assert_not_called() def test_force_tiled_decode_reaches_the_tiled_path(self): @@ -121,7 +140,7 @@ def test_force_tiled_decode_reaches_the_tiled_path(self): torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512), force_tiled_decode=True ) _build_invocation().invoke(context) - vae.enable_tiling.assert_called_once_with() + vae.enable_tiling.assert_called_once_with(tile_sample_min_size=DEFAULT_TILE_SAMPLE_MIN_SIZE) def test_a_requested_tile_size_is_passed_through(self): vae, _, context = _build_decode_mocks(torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512)) @@ -155,7 +174,7 @@ def test_an_untiled_oom_retries_once_tiled(self, oom_error): result = _build_invocation().invoke(context) assert vae.decode.call_count == 2 - vae.enable_tiling.assert_called_once_with() + vae.enable_tiling.assert_called_once_with(tile_sample_min_size=DEFAULT_TILE_SAMPLE_MIN_SIZE) assert result.width == 512 def test_a_non_oom_error_propagates_without_a_retry(self): diff --git a/tests/backend/util/test_vae_tiling_scope.py b/tests/backend/util/test_vae_tiling_scope.py new file mode 100644 index 00000000000..ff39f3f9636 --- /dev/null +++ b/tests/backend/util/test_vae_tiling_scope.py @@ -0,0 +1,147 @@ +"""Scoped tiling state, and the tile-size sentinel. + +Both are here because of what the Qwen-Image tiling PR (upstream #9427) found the hard way: an +`enable_tiling()` call writes through to the model cache's own module, `disable_tiling()` restores +the flag but not the geometry, and an estimator that reads the size back off the module gets +whatever the previous invocation left rather than the default it asked for. +""" + +import pytest +import torch +from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL + +from invokeai.backend.flux.modules.autoencoder import ( + DEFAULT_TILE_OVERLAP, + DEFAULT_TILE_SAMPLE_MIN_SIZE, + MIN_TILE_SAMPLE_SIZE, + AutoEncoder, + AutoEncoderParams, + resolve_tile_size, +) +from invokeai.backend.util.vae_tiling_scope import scoped_vae_tiling + + +def _build_autoencoder() -> AutoEncoder: + params = AutoEncoderParams( + resolution=256, + in_channels=3, + ch=32, + out_ch=3, + ch_mult=[1, 2, 4, 4], + num_res_blocks=1, + z_channels=16, + scale_factor=0.3611, + shift_factor=0.1159, + ) + torch.manual_seed(0) + return AutoEncoder(params).eval() + + +class TestResolveTileSize: + @pytest.mark.parametrize("sentinel", [0, -8, -1024]) + def test_the_sentinel_resolves_to_the_module_default(self, sentinel): + # The workflow UI cannot send None, so 0 is "use the default". A negative value is not worth + # failing a generation over -- upstream #9427 confirmed the same for the Qwen nodes. + assert resolve_tile_size(sentinel) == DEFAULT_TILE_SAMPLE_MIN_SIZE + + @pytest.mark.parametrize("small", [8, 64, MIN_TILE_SAMPLE_SIZE - 8]) + def test_values_below_the_cost_floor_are_clamped(self, small): + # A cost floor, not a correctness one: the tile count grows with the inverse square of the + # tile size, and small tiles are also measurably less accurate. + assert resolve_tile_size(small) == MIN_TILE_SAMPLE_SIZE + + @pytest.mark.parametrize("size", [128, 256, 384, 512, 1024]) + def test_usable_values_pass_through(self, size): + assert resolve_tile_size(size) == size + + +class TestEveryFieldValueProducesTheRightShape: + """The sweep upstream #9427 used to catch silent truncation. + + That bug cannot occur here -- the destination is preallocated at the exact output size and tiles + are merged into it, rather than a loop stepping by one quantity and slicing by another -- but the + guarantee is worth asserting rather than reasoning about, across the shapes most likely to leave + an awkward remainder. + """ + + @pytest.mark.parametrize("latent_hw", [(2, 2), (10, 10), (50, 50), (34, 18), (128, 72), (18, 34), (150, 10)]) + @pytest.mark.parametrize("tile_size", [0, 8, 128, 256, 384, 512, 1024]) + def test_output_shape_is_exact(self, latent_hw, tile_size): + ae = _build_autoencoder() + h, w = latent_hw + z = torch.randn(1, 16, h, w) + with torch.no_grad(), scoped_vae_tiling(ae, tile_size): + out = ae.decode(z) + assert out.shape == (1, 3, h * 8, w * 8) + + +class TestStateIsRestored: + def test_the_normal_path_restores_everything(self): + ae = _build_autoencoder() + before = (ae.use_tiling, ae.tile_sample_min_size, ae.tile_overlap) + with scoped_vae_tiling(ae, 256): + assert ae.use_tiling is True + assert ae.tile_sample_min_size == 256 + assert (ae.use_tiling, ae.tile_sample_min_size, ae.tile_overlap) == before + + def test_the_untiled_path_restores_everything(self): + # Entering with tiling already on: the block must decode untiled and hand the state back. + ae = _build_autoencoder() + ae.enable_tiling(tile_sample_min_size=256) + before = (ae.use_tiling, ae.tile_sample_min_size, ae.tile_overlap) + with scoped_vae_tiling(ae, None): + assert ae.use_tiling is False + assert (ae.use_tiling, ae.tile_sample_min_size, ae.tile_overlap) == before + + def test_an_exception_still_restores(self): + # The OOM retry path raises through this context manager, so `finally` is load-bearing. + ae = _build_autoencoder() + before = (ae.use_tiling, ae.tile_sample_min_size, ae.tile_overlap) + with pytest.raises(RuntimeError, match="boom"): + with scoped_vae_tiling(ae, 256): + raise RuntimeError("boom") + assert (ae.use_tiling, ae.tile_sample_min_size, ae.tile_overlap) == before + + def test_geometry_does_not_leak_between_two_scopes(self): + """The bug this exists for: `disable_tiling()` clears the flag but keeps the geometry, so a + size set once would otherwise silently become the default for everyone afterwards.""" + ae = _build_autoencoder() + with scoped_vae_tiling(ae, 256): + pass + assert ae.tile_sample_min_size == DEFAULT_TILE_SAMPLE_MIN_SIZE + assert ae.tile_overlap == DEFAULT_TILE_OVERLAP + with scoped_vae_tiling(ae, 0): + assert ae.tile_sample_min_size == DEFAULT_TILE_SAMPLE_MIN_SIZE + + def test_a_tiled_decode_does_not_leave_the_shared_vae_tiled(self): + """Nine nodes reach this class and most never touch the tiling flag -- FLUX.1 encode, PiD, + and Anima's FLUX branch among them. A leaked flag would silently tile their work.""" + ae = _build_autoencoder() + z = torch.randn(1, 16, 96, 96) + with torch.no_grad(): + with scoped_vae_tiling(ae, 0): + ae.decode(z) + assert ae.use_tiling is False + # What an unguarded consumer would get next, decoded with no tiling call of its own. + after = ae.decode(z) + ae.disable_tiling() + expected = ae.decode(z) + assert torch.equal(after, expected) + + +class TestDiffusersVaesAreHandledToo: + def test_a_class_without_a_settable_size_is_enabled_without_arguments(self): + # diffusers' AutoencoderKL.enable_tiling() takes no parameters; passing one would raise. + vae = AutoencoderKL( + in_channels=3, + out_channels=3, + latent_channels=4, + block_out_channels=(32,), + down_block_types=("DownEncoderBlock2D",), + up_block_types=("UpDecoderBlock2D",), + layers_per_block=1, + norm_num_groups=32, + ) + with scoped_vae_tiling(vae, 384): + assert vae.use_tiling is True + assert vae.use_tiling is False