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

[Pfannkuchensack]

Summary

Feature. InvokeAI's FLUX.1 autoencoder (invokeai/backend/flux/modules/autoencoder.py, our own port of the BFL reference) has no tiling — grep -c "tiling\|tile" returned 0. Peak decode memory therefore grows with the square of the output, unbounded: measured on real weights, 4.35 GiB at 1024², 9.50 at 1536², 16.72 at 2048². What happens when that does not fit depends on the platform, and neither outcome is acceptable:

• Where the allocator is the limit, the decode raises OutOfMemoryError and the generation is lost.
• On Windows with recent NVIDIA drivers, it does not fail — the driver falls back to system memory and the decode crawls over PCIe instead, taking the rest of the machine down with it.

This PR bounds the peak at 1.23 GiB flat, at any resolution, so neither path is reached.

Nine call sites share that autoencoder: flux_vae_decode / flux_vae_encode, z_image_latents_to_image / z_image_image_to_latents, anima_latents_to_image / anima_image_to_latents and the PiD nodes. Z-Image is where it surfaces in practice — both its GGUF starter bundles declare flux_vae as a dependency, so small-VRAM users land on exactly this combination — but FLUX.1 itself has the same exposure.

Why tiling went on the class, not into one node. The alternative considered was swapping the loader to a diffusers AutoencoderKL, which brings enable_tiling() with it. An audit of the call sites rejected it:

• flux_vae_encode.py:45 asserts isinstance(vae_info.model, AutoEncoder) — there is no diffusers path.
• pid_upscale.py:124 explicitly rejects a diffusers VAE with a written rationale, then undoes FLUX scaling by hand (normalised_latent / ae.scale_factor + ae.shift_factor) — arithmetic that only holds for our class's encode convention.
• anima_latents_to_image.py:143-155 dispatches over Union[AutoencoderKLWan, FluxAutoEncoder]; if the FLUX VAE became an AutoencoderKL it would fall through both branches and raise TypeError.

Re-deriving PiD's scaling and rewriting Anima's class dispatch in the same PR that introduces tiling is not a trade worth making. Implementing tiling on AutoEncoder instead covers all nine call sites, changes no class, and gives the nodes the same enable_tiling() / disable_tiling() spelling they already use for the diffusers VAEs.

How. 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 then cannot be sliced out of the latent. Overlaps scale with the coordinates because they are nothing but coordinate differences. Each finished tile is moved to the host before the next is decoded — that, not the tile size alone, is what makes the peak flat rather than merely smaller.

Also in this PR:

• estimate_vae_working_memory_flux takes an optional tile_size, defaulting to None so the six existing call sites are untouched. tile_size=0 resolves through getattr, not a direct attribute read — the Z-Image nodes also hand a diffusers AutoencoderKL to this estimator, and the SD1/SDXL sibling's vae.tile_sample_min_size would raise on it.
• The Z-Image decode node gains 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, no schema change.
• _is_oom_error moves out of the Anima node into invokeai/backend/util/oom.py and is imported by both. A second copy is how the next backend's spelling gets added to only one of them.
• The three stray print() calls in vae_working_memory.py (lines 74, 96, 355 — cogview4, flux, sd3) are removed. They fired on every decode.

Encode stays untiled, deliberately. It peaks at roughly half of decode and would need its own tiled encode. No half-wired field was added for it.

One behaviour change for existing users, worth stating explicitly. force_tiled_decode was previously ignored by the Z-Image decode node — that is part of what this fixes. Anyone who already has it set in invokeai.yaml therefore gets tiled decoding where they used to get untiled, and a tiled decode of this architecture is not bit-identical (see below). "Off by default, so nothing changes unless you enable it" is not quite true for them. Nothing changes for anyone who has not set that flag.

Reviewed against upstream invoke-ai#9427

The Qwen-Image tiling PR solved this same problem for a different VAE and collected a detailed
review. Every finding there was checked against this branch. Three of them landed:

• The cached VAE was being mutated permanently. 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. A tiled Z-Image decode therefore left use_tiling=True behind, and the nodes that share
  the FLUX autoencoder 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 have silently decoded and encoded tiled. Now scoped through scoped_vae_tiling,
  which restores every tiling attribute in a finally. Same shape as SD's patch_vae_tiling_params
  and Qwen's patch_qwen_image_vae_tiling, neither of which fits these classes.
• The estimator read the leaked value back. getattr(vae, "tile_sample_min_size", ...) returns
  whatever the previous invocation set, not the default this node asked for. tile_size <= 0 now
  resolves against a module-level constant.
• No lower bound on the field. tile_size is multiple_of=8 with no floor, so a small value
  produced an enormous tile count, and a negative one raised ValueError in the middle of a
  generation. Both now resolve through resolve_tile_size, with a 128px cost floor.

Two findings do not apply, checked rather than assumed:

• Silent truncation when the tile is smaller than the stride cannot occur here. That bug comes
  from diffusers stepping the tile loop by one quantity and slicing by another; this implementation
  preallocates the destination at the exact output size and merges into it. Asserted anyway, over
  every value the field can carry against seven deliberately awkward shapes — 49 combinations, all
  exact.
• Compute growing quadratically with tile size is the opposite way round here: larger tiles mean
  fewer tiles, and are measurably more accurate. FieldDescriptions.vae_tile_size's existing
  wording — larger tiles cost memory and give better results — is correct for this VAE, where for
  the Qwen one it had to be corrected.

One more surfaced from running the checks rather than from the review: the new module was initially
named vae_tiling.py, colliding 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 entire suite. Renamed to vae_tiling_scope.

The tile geometry is exact

A tiled decode of this architecture can never match a single-pass one in general: the decoder's GroupNorms normalise over the whole spatial extent and its mid-block attention is global, so both see different input when the image arrives in tiles. That is inherent, and equally true of diffusers' tiling. With those two operators removed the decoder is purely convolutional and the two must agree — they do, to 1.0e-07, seams included. That test is what actually verifies the slicing, the coordinate scale-up and the blending; everything else would only be measuring the architecture.

Smaller tiles are less accurate, not more

This inverts the intuition that seams are answered with more overlap. On the test fixture at 96×96 latents: 512px tiles land at 1.0e-07, while 256px tiles at the same 128px overlap drift to 4.4e-03 — four orders of magnitude worse. Halving the tile multiplies the seams, and the blend bands then sit closer to each tile's own zero-padded border. Prefer a large tile that fits over a small one that fits comfortably.

Related Issues / Discussions

Design notes and the measurements this was built from live in the local .ideas/prompt-z-image-tiled-decode.md and .ideas/vae-decode-oom-on-rocm.md.

QA Instructions

Measured on RTX 4090, torch 2.7.1+cu128, Windows, against the real ae.safetensors FLUX.1 VAE weights in bf16 — not a fixture.

Peak memory and decode time

Peak reserved (which is what InvokeAI's working-memory estimator targets), single decode, counters reset per measurement:

px	untiled	tiled	saved	untiled	tiled	slower
1024	4.35 GiB	1.23 GiB	3.5×	0.42 s	0.45 s	1.09×
1536	9.50 GiB	1.23 GiB	7.7×	0.38 s	0.76 s	1.97×
2048	16.72 GiB	1.23 GiB	13.6×	0.74 s	1.18 s	1.60×

The tiled peak is flat because finished tiles leave the device immediately. On a 16 GB card, 2048² goes from "does not fit" to routine.

Untiled fails, tiled succeeds — demonstrated, not argued

With PyTorch's own allocator capped at 3 GiB (its cap, not the driver's — so it raises before the driver is asked and Windows' sysmem fallback never engages):

1536px untiled: OutOfMemoryError   is_oom_error=True
1536px tiled  : COMPLETED  peak allocated 0.70 GiB  reserved 1.23 GiB

is_oom_error=True is the same predicate the node's retry uses, so this is the exact error class that now turns into a slower success instead of a lost generation.

Seam check on real images

Six generated images, resized to 1536², encoded with this VAE's own encoder and decoded both ways — in-distribution content with real interior seams (four tiles per axis at the shipped 512/128 geometry):

	maxdiff	mean abs	worst column / median column
range over 6 images	0.078 – 0.295	0.007 – 0.016	1.20× – 2.18×
mean	0.186	0.0122	1.71×

A visible seam is a localised spike — a column orders of magnitude above the median. There is none: the worst column is at most 2.2× the median. Amplifying the difference 20× shows it tracking image content (fur, edges, texture), which is the GroupNorm/attention effect, with faint per-tile level shifts in flat regions and no boundary line. The decoded image is clean across a large flat pavement area, which is the case seams show up in first.

Stated plainly because it is a real gap: the mean error (0.0122) is above the diffusers reference for tiling this same VAE (0.0022). It is not visible, but it is not equal either. The most likely cause is blend width — this blends over a fixed 128px inside an overlap that calc_tiles_min_overlap actually spreads to 168–176px, so part of the available overlap goes unused. Widening the blend to the available overlap is a plausible follow-up; it is not done here.

The untiled path is unchanged

Two consecutive untiled decodes are bit-identical, and an untiled decode after a tiled one is bit-identical to one before it — the cached, shared VAE instance does not leak tiling state between invocations.

Test suite

• tests/backend/flux/modules/test_autoencoder_tiling.py — 16 tests: tiling state and its restoration, geometry exactness, four tile layouts (evenly divided, not evenly divided, odd on both axes, tiling on one axis only), an image smaller than one tile, batch independence, and that finished tiles leave the decode device.
• tests/backend/util/test_vae_tiling_scope.py — 66 tests: the sentinel and the cost floor, the 49-combination shape sweep, state restored on the normal / untiled / exception paths, no geometry leaking between two scopes, no tiling flag left on the shared VAE, and a diffusers AutoencoderKL handled without arguments.
• tests/app/invocations/test_z_image_tiled_decode.py — 21 tests: estimator behaviour including a regression guard that tile_size=None reproduces today's value exactly, plus node wiring, the restored state, and the OOM retry.
• Mutation-checked, each caught: ignoring force_tiled_decode (1 failure), removing the retry (4), never passing tile_size to the estimator (1), not setting the tiling state (5), dropping the finally restore (6), resolving the sentinel off the module again (12), removing the cost floor (4).
• tests/app/invocations tests/backend/flux tests/backend/tiles tests/backend/util tests/backend/stable_diffusion — 1283 passed, 1 skipped, 6 xfailed. pytest --collect-only clean. ruff check / format --check clean.
• openapi.json regenerated as CI does it (generate, then pnpm prettier --write): 22 insertions, 1 deletion — the two new fields and the version bump, nothing else. pnpm typegen adds 12 lines to schema.ts.

Not run

A full generation through a live server with force_tiled_decode set in invokeai.yaml. The wiring is covered by mocked tests and the mutation check above, but the config switch has not been exercised end to end.

Merge Plan

Nothing special. No DB schema, no redux slice, no dependency change. The openapi.json change is additive — two optional fields with defaults on one prototype node — so existing workflows deserialise unchanged.

Independent of #172, which is in flight in parallel: the two touch disjoint files and can merge in either order.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable) — 103 tests across three files; seven separate mutations verified as caught
• ❗Changes to a redux slice have a corresponding migration — n/a, no redux changes
• Documentation added / updated (if applicable) — n/a; the new node fields carry the shared FieldDescriptions text
• Updated What's New copy (if doing a release after this PR) — worth a line: large images now decode in bounded memory instead of failing or falling back to system RAM

🤖 Generated with Claude Code