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

[Pfannkuchensack]

Summary

Both Qwen-Image VAE nodes reserve working memory for a full-frame operation, which at high resolutions is more than a 24 GB card has — so the cache evicts everything else to honour it. On CUDA at 2560x1440:

Node	reserved, untiled	reserved, tiled (256px)
qwen_image_l2i (decode)	19.91 GiB	0.55 GiB
qwen_image_i2l (encode)	10.99 GiB	0.26 GiB

Tiling is the intended escape hatch, but today it does not work on either node:

• i2l hardcoded vae.disable_tiling() — there was no way to enable it at all.
• l2i honours the global force_tiled_decode, but computes its working-memory estimate before and independently of that flag. So tiling bounds the VAE while the cache still reserves the full-frame figure — the memory is never actually freed for anything else. Effectively inert.

This PR adds tiled / tile_size input fields to both nodes, following the SD/SDXL i2l/l2i nodes, OR'd with the global force_tiled_decode. Off by default, so behaviour is unchanged unless enabled.

estimate_vae_working_memory_qwen_image gains a matching tile_size parameter, and both nodes resolve tile_size=0 to the default (256px) before estimating. Tiled, it budgets one tile plus 25% overlap plus the pixel-space buffers, mirroring estimate_vae_working_memory_wan. Without this the change would be cosmetic on i2l and remain inert on l2i.

Measured end-to-end through the i2l node at 2560x1440: 10.99 → 0.26 GiB reserved, 9.26 → 0.17 GiB actual peak, identical latent shape. Tiled latents differ by ~1.4% relative L2 on noise input (worst case for tile blending; real images blend far better), which is why this stays opt-in.

Applying the tile geometry correctly

AutoencoderKLQwenImage carries a tile size and a tile stride (stock: 256px / 192px). tiled_encode / tiled_decode step the tile loops by stride but slice each accumulated tile to min, and enable_tiling inherits the module's current value for any argument left out. Passing only tile_sample_min_* therefore breaks in both directions:

• Below the inherited 192px stride, whole bands of the image are never processed and the output is silently smaller than requested — a 128px tile turns a 512x512 decode into 384x384, with no exception raised. On i2l the undersized latent flows onward as a perfectly valid-looking tensor.
• Above it, the tile count is fixed by the stride, so a larger tile_size grows every tile without eliminating any: at 2560x1440 the loops emit 112 tiles regardless, making compute scale with tile_size² (8.0x a full frame at 512px, 17.9x at 768px).

So both nodes pass all four parameters, keeping the stock 3/4 stride ratio, as anima_latents_to_image already does. The stride is additionally rounded down to a multiple of the VAE's 8x spatial compression: the loops step in one space and slice in the other, so the pixel stride must be exactly 8x the latent stride or the two disagree. With a proportional stride the processed area stays flat at ~2.0x a full frame at any tile size.

tile_size values between 1 and 64 are clamped to 64 as a cost floor (1620 tiles at 64px vs 57,600 at 8px on a 2560×1440 frame). The field has to accept the 0 "use the default" sentinel and so cannot carry a pydantic lower bound.

Scoping the tiling state

enable_tiling writes the geometry straight onto the module and disable_tiling only clears use_tiling — it does not restore the sizes. That module is the model cache's own instance (LoadedModel.model and model_on_device both return it, and as_qwen_image_vae deliberately returns the same object to keep partial-loading hooks intact), so a tile_size set once would stick for the lifetime of the cache entry and leak into later invocations.

It also crosses model families: a native-layout qwen_image_vae single file is classified with the Anima base, so anima_latents_to_image — which sets 512/384 and never restores — can feed the same instance in a single workflow.

Both nodes now apply the geometry through a context manager that restores all five attributes in a finally, and resolve the tile_size=0 sentinel against a constant rather than reading the module's current tile_sample_min_height. Together that makes the nodes both non-leaking and immune to geometry left behind by anything else.

Related Issues / Discussions

None filed — found while debugging VRAM exhaustion in a latents → image → upscale → image → latents workflow at 2560x1440 on a 24 GB card, where a ~12 GB transformer stays resident across the VAE round-trip.

QA Instructions

Reproducing the limit (no code needed): run a Qwen-Image / Krea-2 img2img round-trip at ~2560x1440 while a large transformer is resident. The VAE nodes request ~20 GB (decode) and ~11 GB (encode) of working memory, forcing the cache to evict the transformer; on a 24 GB card this surfaces as Loading 0.0 MB into VRAM, but only -N MB were requested and models loading at <100%.

With this PR:

1. Enable tiled on the Latents to Image and/or Image to Latents - Qwen Image nodes (or set force_tiled_decode: true in invokeai.yaml for both).
2. Re-run the workflow. The nodes should complete without evicting the transformer, and the VRAM warning should not appear.
3. Compare against a tiled: false run. Images should be visually equivalent — tile seams are the failure mode to look for; none were observed on photographic content.
4. Leave tiled off and confirm behaviour is identical to main. This is the default path and the most important check.
5. tile_size: 0 uses the default (256px). Larger values mean fewer, larger tiles: reserved memory scales with tile_size² (256 → ~0.55 GiB decode, 512 → ~1.9 GiB, at any resolution) while the processed area stays ~2.0x a full frame either way. Values 1–64 are clamped to 64.
6. Chain an Anima decode and a Qwen-Image decode against the same VAE in one workflow, in either order, and confirm both produce correctly sized output — the tiling state must not carry between them.

Automated: pytest tests/app/invocations/test_qwen_image_working_memory.py — covers the resolved tile size reaching the estimator, the arguments actually passed to enable_tiling on both nodes, the estimator's tiled arithmetic, and a class of tests against a real (tiny, randomly initialised) AutoencoderKLQwenImage asserting shape preservation across tile sizes, state restoration, and immunity to geometry left by another node.

Not measured: the decode's runtime peak was not benchmarked separately — only its reservation, and the encode end-to-end. The decode constants themselves are unchanged from the existing calibration.

Behaviour changes worth calling out

• force_tiled_decode now affects encode. ORing the global flag into i2l matches image_to_latents.py, so the precedent is established — but users who already have force_tiled_decode: true in invokeai.yaml will get a ~1.4%-different latent with no node or workflow change on their side. "Off by default" holds for everyone else.
• width/height zero-handling on i2l. The workflow UI cannot represent None in a number input and sends 0 for "unset", which reached resize((0, 0)) and raised "height and width must be > 0". Non-positive values are now treated as unset — which also means a half-filled pair (e.g. width=1024, height=0) encodes at the original size rather than erroring.

Merge Plan

Nothing special. Both node versions bumped 1.0.0 → 1.1.0; all new fields have defaults, so saved workflows load unchanged and keep current (untiled) behaviour. The new fields change the generated OpenAPI schema, so schema.ts needs regenerating as usual.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration — n/a, backend only
• Documentation added / updated (if applicable) — n/a, field descriptions are inline
• Updated What's New copy (if doing a release after this PR) — worth a line for force_tiled_decode users, see above

[lstein]

Thanks for this — the diagnosis is right and well argued. The observation that l2i's tiling was inert because the estimate was computed independently of the flag is a genuine find, and threading tile_size into estimate_vae_working_memory_qwen_image is the correct fix for it. The untiled path is byte-identical (the estimator's else branch is untouched and both nodes pass tile_size=None when tiled is false), schema/openapi are regenerated correctly including both version bumps, and CI is fully green.

The problem is in how tiling is turned on. Both nodes call

vae.enable_tiling(tile_sample_min_height=tile_size, tile_sample_min_width=tile_size)

and never touch tile_sample_stride_height / tile_sample_stride_width. That leads to two independent ways to get silently corrupted output.

---

Blocker 1 — tile_size below 192 silently truncates the image / latent

AutoencoderKLQwenImage.__init__ sets tile_sample_min_* = 256 and tile_sample_stride_* = 192. tiled_encode / tiled_decode step the tile loops by stride but slice each accumulated tile to min:

for i in range(0, height, tile_latent_stride_height):        # advances by stride
    ...
result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width])

So whenever min < stride, whole bands of the image are never processed and the output is smaller than requested. tile_size is declared multiple_of=8 with no lower bound, so every value from 8 to 184 hits this.

Measured against a real AutoencoderKLQwenImage (tiny random weights, CPU, diffusers 0.39.0):

encode 512x512 image   tile_size=192 -> (1,16,1,64,64)    ok
                       tile_size=128 -> (1,16,1,48,48)    *** WRONG ***
                       tile_size=64  -> (1,16,1,24,24)    *** WRONG ***

decode 64x64 latent    tile_size=192 -> (1,3,1,512,512)   ok
                       tile_size=128 -> (1,3,1,384,384)   *** WRONG ***
                       tile_size=64  -> (1,3,1,192,192)   *** WRONG ***

No exception is raised. On i2l the undersized latent flows onward as a perfectly valid-looking tensor.

The same omission has a second, opposite-direction effect above 256: the number of tiles is fixed by the stride, so raising tile_size grows every tile without eliminating any. At 2560x1440 the loops emit ceil(1440/192) x ceil(2560/192) = 8 x 14 = 112 tiles regardless of tile_size:

tile_size	tiles	pixels processed	vs full-frame
256	112	7.3 Mpx	2.0x
384	112	16.5 Mpx	4.5x
512	112	29.4 Mpx	8.0x
768	112	66.1 Mpx	17.9x

FieldDescriptions.vae_tile_size says "larger tile sizes generally produce better results at the cost of higher memory usage", and QA step 5 recommends going larger — but on this VAE the real cost of a larger tile is quadratic compute, not just memory. With a proportional stride, tile_size=512 would be ~1.8x, not 8x.

Suggested fix: pass all four parameters, keeping the stock 256/192 ratio, e.g.

vae.enable_tiling(
    tile_sample_min_height=tile_size,
    tile_sample_min_width=tile_size,
    tile_sample_stride_height=tile_size * 3 // 4,
    tile_sample_stride_width=tile_size * 3 // 4,
)

This is exactly what anima_latents_to_image.py:144-149 already does. A ge= floor on the field would be belt-and-braces.

---

Blocker 2 — enable_tiling permanently mutates the cached VAE

self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height

enable_tiling writes through to the module, and disable_tiling() only clears use_tiling — it does not restore the sizes. That module is the model cache's own instance: LoadedModel.model and model_on_device both return self._cache_record.cached_model.model (load_base.py:103,110), and as_qwen_image_vae deliberately returns the same object. So a tile_size set once sticks for the lifetime of the cache entry.

Replaying the node's exact call sequence twice against one instance:

run A  tile_size=512  ->  model min now 512
run B  tile_size=0    ->  estimate=512, tiles=512     (documented as "the VAE default, 256")

run A  tile_size=128  ->  (1,3,1,384,384)
run B  tile_size=0    ->  (1,3,1,384,384)             *** still corrupt; the user set nothing ***

then   i2l tiled=True tile_size=0  ->  latent (1,16,1,48,48)   *** the leak crosses nodes ***

It also crosses model families. anima_latents_to_image sets min=512, stride=384 (ANIMA_VAE_TILE_SIZE / ANIMA_VAE_TILE_STRIDE) on the same Wan-classified VAE instance these nodes reinterpret — a native-layout qwen_image_vae single file is classified with the Anima base, as the node comments note, so one VAE loader can feed both nodes in a single workflow. Afterwards the documented-safe value corrupts, because the leaked stride is now larger than it:

after Anima l2i:                min=512  stride=384
qwen l2i with tile_size=256 ->  (1,3,1,384,384)      expected (1,3,1,512,512)

The SD l2i this PR follows avoids precisely this with patch_vae_tiling_params — a context manager that restores the original values in a finally. Note qwen_image_latents_to_image.py:92 still carries the now-vestigial tiling_context = nullcontext(); that is the hook the SD node uses for it.

Fixing blocker 1 by always passing all four parameters also closes most of this, since nothing is then inherited. The residual piece is getattr(vae_info.model, "tile_sample_min_height", 256) in the estimate, which reads whatever the previous run leaked rather than a known default.

---

Repro script

Both blockers above, self-contained (no weights needed):

import torch
from diffusers.models.autoencoders.autoencoder_kl_qwenimage import AutoencoderKLQwenImage

vae = AutoencoderKLQwenImage(
    base_dim=4, z_dim=16, dim_mult=[1, 1, 1, 1], num_res_blocks=1,
    attn_scales=[], temperal_downsample=[False, True, True],
).eval()
z = torch.randn(1, 16, 1, 64, 64)  # -> 512x512

# Blocker 1
vae.enable_tiling(tile_sample_min_height=128, tile_sample_min_width=128)
with torch.inference_mode():
    print(tuple(vae.decode(z, return_dict=False)[0].shape))   # (1,3,1,384,384)

# Blocker 2 -- the node's tile_size=0 path, run afterwards
vae.enable_tiling()                                            # no args == "model default"
with torch.inference_mode():
    print(tuple(vae.decode(z, return_dict=False)[0].shape))   # (1,3,1,384,384), still

---

Non-blockers

The tiled estimate under-counts the accumulated tile buffers. working_memory += 3 * h * w * element_size budgets one RGB frame, but tiled_decode holds every decoded tile in rows (about (min/stride)^2 ~ 1.78x a full frame at the defaults) simultaneously with result_rows (~1x) and the final torch.cat plus its slice (~2x). At 2560x1440 fp16 that is roughly 116 MB actual against 22 MB budgeted — comfortably absorbed by the 25% tile slack within the 497 MB total. But the shortfall grows linearly with output area while the tile term stays constant, so it degrades in exactly the regime tiling exists for. estimate_vae_working_memory_wan, which the new docstring says this mirrors, uses clip_copies = 2 on decode; this one effectively uses 1.

force_tiled_decode now changes encode results. ORing it into i2l matches image_to_latents.py:158, so the precedent is fine — but users who already have it set in invokeai.yaml get a ~1.4%-different latent with no node or workflow change on their side. "Off by default, so behaviour is unchanged unless enabled" is not quite true for them; worth a line in the PR body / What's New rather than a code change.

The width/height zero-handling fix is an unrelated drive-by, untested and unmentioned in the summary. In self.width and self.height and self.width > 0 and self.height > 0 the first two clauses are redundant with the last two. Also note width=1024, height=0 now silently encodes at the original size rather than erroring — reasonable, but a behaviour change worth stating.

Test coverage. The new test patches the estimator, so nothing exercises the estimator's tiled arithmetic, the arguments actually passed to enable_tiling (where both blockers live), or the l2i node at all — and l2i is where the inert-tiling bug being fixed actually was. The try / except Exception: pass would also let the test pass if vae_encode blew up immediately after the estimate call.

---

Attacks that came back clean

• Untiled path is byte-identical to main; the estimator's non-tiled branch is untouched and both nodes pass tile_size=None when tiled is false.
• Negative tile_size values fall through to the > 0 guard and are treated as "default" — no crash.
• Images at or below the tile size over-reserve slightly rather than under-reserve (diffusers skips tiling entirely there).
• QwenImageImageToLatentsInvocation.vae_encode has no other callers, so the new keyword-only defaults break nothing.
• element_size derivation and the ROCm/CUDA constant selection are unchanged.
• schema.ts and openapi.json are correctly regenerated, including both 1.0.0 -> 1.1.0 bumps and the multipleOf: 8 constraint; typegen-checks and openapi-checks both pass.
• ruff check clean; all 8 tests in test_qwen_image_working_memory.py pass.

[lstein]

Both blockers are fixed properly — not worked around — and the fix is better than what I suggested: pulling the geometry into patch_qwen_image_vae_tiling rather than open-coding four arguments at each call site means the pixel/latent stride relationship is stated once, next to the reason it matters. Approving.

I re-attacked both findings against 4ca19dfcc5 rather than reading the diff for plausibility.

Blocker 1 — tile stride

Swept every value the field can carry — tile_size=0 plus every multiple of 8 up to 1024 — across 8 shapes chosen for pathological leftovers (16x16, 80x80, 400x400, 272x144, 1024x576, 144x272, 1200x80, 64x64), asserting both the decoded pixel shape and the encoded latent shape:

total combos: 1032   failures: 0

The min > stride direction is fixed too — processed area at 2560x1440 is now flat instead of quadratic:

tile_size	192	256	384	512	768
before	1.4x	2.0x	4.5x	8.0x	17.9x
after	1.80x	1.99x	1.80x	1.99x	2.40x

Blocker 2 — cached-module mutation

Confirmed against diffusers 0.39.0 that enable_tiling on AutoencoderKLQwenImage and AutoencoderKLWan mutates exactly the five attributes the context manager saves — no sixth. Restore verified on the normal, untiled and exception paths, and the Anima scenario now holds in both directions: leaked 512/384 cannot change a Qwen decode, and is still intact for whoever set it afterwards.

I also went after the hazard the context manager newly opens up — two invocations interleaving save/restore on one module. It isn't reachable: one session worker per device, per-device caches, and _adopt_shared_cpu_weights deep-copies a meta shell (load_default.py:537, "Independent module per device") so two workers never hold the same VAE object.

Estimator

image_copies = 5 checks out. Measured peak allocated during tiled_decode on GPU, in multiples of one RGB frame:

1440x2560  tile=256 -> 3.78    tile=512 -> 5.21
2048x2048  tile=256 -> 3.73    tile=512 -> 5.09
3072x3072  tile=256 -> 3.73    tile=512 -> 3.93
encode     tile=256 -> 0.62 .. 1.12

The one overshoot (5.21 at 1440x2560/512) is ~4.6 MB against a ~1.9 GB tile term, so the 25% overlap slack absorbs it many times over.

Pushed to this branch (3952b42)

One residual, which I've fixed rather than sending back:

The multiple-of-8 stride rounding was load-bearing but untested. I mutation-tested each element of the fix. Removing the stride arguments failed 20 tests, dropping the finally restore failed 6, removing the clamp failed 2, regressing the sentinel to a module read failed 3 — but deleting // 8 * 8 from _tile_stride_for left all 35 tests green. It is not cosmetic: tile_size is multiple_of=8, so 72 and 80 are reachable from the UI and their raw 3/4 stride (54, 60) is not, and a 512x512 encode at min=72, stride=54 returns a 57x57 latent instead of 64x64 — silently, exactly like the original bug. Above 136 it raises outright. Now covered at the argument level (72 -> 48, 80 -> 56), end-to-end through the real tiny VAE, and with an explicit pin on the failure mode; the mutation now fails 5 tests.

The same commit corrects the rationale on QWEN_IMAGE_VAE_MIN_TILE_SIZE. _tile_stride_for already floors the stride at 8, so the derived latent step can never collapse to 0, and 8 / 24 / 32 all round to clean multiples of 8 and produce correctly sized output — 64 is not the smallest valid tile. It's a cost floor: the tile count grows with the inverse square of the stride, 1620 tiles at 64px against 57,600 at 8px on a 2560x1440 frame. The clamp is right; only the reason was wrong, and it would have stopped someone from lowering it later.

The PR body now covers every non-blocker from the first round — the force_tiled_decode encode change, the width/height drive-by, and the corrected guidance on larger tiles. ruff check and ruff format clean, 801 passed in tests/app/invocations/.