Studio: size seam-free VAE tiles from measured peaks (no slower than stock, no low-VRAM OOM) - #12764

Follows #12736 (merged). #12736 made every image VAE's tiled decode seam-free; this makes those wide tiles no slower than the stock tiles they replace, and closes the low-VRAM OOM corners, while keeping the seam fix.

Problem

Measured after #12736, against main before it, on real weights:

• Qwen-Image at low free VRAM: the wide tiles are sized from one per-pixel figure that is 2 to 5x the real decode peak of most VAEs, so at about 0.6 GiB free Qwen-Image decodes at the 32-latent floor and is 1.4 to 1.8x slower than the stock 32 / 8 tiles.
• FLUX.1 / SDXL / FLUX.2 at an edge sliver (1600 px): the 128-latent wide tiles decode about 6 ms (9%) slower than stock.
• HunyuanImage-2.1: the budget figure is 5x its measured peak, so 2048 px rarely gets tiles above the floor, and below about 1.4 GiB free the 1.3 GiB floor tile does not fit at all and OOMs where the stock tiles would not.
• Qwen-Image-2.1: at about 0.35 GiB free the floor tile (inherited from Studio: seam-free VAE tiles for Qwen-Image-2.1 on low-VRAM tiers (fixes thin horizontal / vertical lines) #12696) OOMs, where the stock tiles decode.

Change

diffusion_vae_tiling.py (and the fused VAE module's batched loop, so the fallback keeps it):

• The decode budget comes from measured peaks. Each figure is bf16 MiB per latent of tile area, unfused / fused, the worst over tile sides 8 to 128, measured on the real weights:
  • Qwen-Image: 0.27 / 0.18.
  • Qwen-Image-2.1: 1.7 / 0.43 (Studio: seam-free VAE tiles for Qwen-Image-2.1 on low-VRAM tiers (fixes thin horizontal / vertical lines) #12696's figure; measured 1.64).
  • HunyuanImage-2.1: 1.3 (no fused path).
  • AutoencoderKL and FLUX.2: 0.17 / 0.09.
  • A VAE not measured takes the worst of its ratio.
  • Tiles are sized on the fused figure where Studio's fused VAE kernels are installed.
  • The old per-pixel figure was 5x HunyuanImage's real peak and 2.9x AutoencoderKL's, so those VAEs rarely got tiles above the floor.
• At a sliver, a VAE whose stock tile is the floor tile (AutoencoderKL, FLUX.2) takes the fewest decoder calls, then the widest tile, that decode no more latents per side than the stock loop. At 1600 px that is two 120-latent tiles: 240 decoded latents, as stock, in 4 calls instead of 9. Those VAEs grow their tiles past the floor only into memory the device has free, not the allocator's cached blocks.
• When even the floor tile cannot fit the free VRAM, the decode uses the stock tiles and logs why (once per VAE). This only applies to a VAE whose stock tile is smaller than its floor tile: HunyuanImage (1.3 GiB against 0.2) and Qwen-Image-2.1 (0.43 GiB fused). An OOM in the wide tiles retries at the floor tile, then in the stock tiles. Losing quality is better than an OOM.
• When the wide path falls back to stock tiles on the Wan family, it keeps diffusion_vae_fused's batched tile loop, as fast as before Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736.
• The geometry is read once at install, and the blend weights are cached per VAE.

Performance

Arms: pre-#12736 is main before #12736 (df4eb29579). #12696 is 27916a55fb. #12736 is acc589fbe5, the code main carries now (the merged commit only tightened comments). Head is this PR (a4a70e3a74, the decode path is unchanged by the merges since).

Decode alone, emulated free VRAM

Each VAE's tiled decode of a real latent (the untiled encode of DIV2K 0010), with Studio's fused VAE kernels installed in load order. Exclusive B200. mem_get_info is capped at the given free memory with no allocator slack. The cap is not enforced, so a #12736 figure above the free memory would OOM on a real card. Arms are interleaved in one process. Median of 9 decodes after one warm decode (15 for the AutoencoderKL VAEs), in seconds. PSNR and the worst 64 px window are against the untiled decode of the same latent.

• Qwen-Image: the fused figure (0.18 MiB per latent) buys 40 to 56-latent tiles at 0.6 GiB free, and the head decodes 2 to 2.3x faster than pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736 there, where Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736 was 1.4 to 1.8x slower. At 2 GiB free it is 5.5 to 8x faster.
• Qwen-Image-2.1: at the floor tile the head is torch.equal to Studio: seam-free VAE tiles for Qwen-Image-2.1 on low-VRAM tiers (fixes thin horizontal / vertical lines) #12696. With 1 GiB or more free it takes larger tiles: 1.3 to 2.2x faster than Studio: seam-free VAE tiles for Qwen-Image-2.1 on low-VRAM tiers (fixes thin horizontal / vertical lines) #12696, and closer to untiled (53 to 61 dB against 50 to 54). At 0.35 GiB free the floor tile (about 0.45 GiB with the fused kernels) cannot fit, so it decodes in the stock tiles. Studio: seam-free VAE tiles for Qwen-Image-2.1 on low-VRAM tiers (fixes thin horizontal / vertical lines) #12696 OOMs there (see the low-free table below).
• HunyuanImage-2.1: a 1024 px canvas is one 32-latent tile, decoded untiled whenever it fits. At 4 GiB free, the 1.3 MiB figure gives 2048 px 64x32 tiles where the old figure kept the floor tile: 2x faster than pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736 and 1.5x faster than Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736, at 32.8 dB against 28.3. Below about 1.4 GiB free the floor tile (1.3 GiB) cannot fit, so it decodes in the stock tiles, as pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736 does. Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736 OOMs there.
• AutoencoderKL / FLUX.2 at the 1600 px sliver: the 120-latent tiles decode as many latents as stock in 4 calls instead of 9. Over 25 interleaved runs at the floor, SDXL took 60.6 ms against pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736's 66.0 ms, and FLUX.1 60.9 ms against 65.8 to 71.2 ms. Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736 was 68 to 76 ms. The floor peak drops from 1,408 to 1,270 MiB. Line scores (the worst row or column mean error above its 128 px neighbourhood, in levels) match the 128-latent tiles: 0.32 to 1.40 against 0.32 to 1.25, and stock's 2.17 to 6.22, on DIV2K 0010. PSNR sits between stock and Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736. 2048 px keeps the stock path, bit-identical.
• The per-decode overhead of the rule is the tile decodes themselves: four bare 120-latent tile decodes take 60.4 to 61.3 ms, and the whole wide decode takes 60.4 to 61.5 ms.

End to end through Studio's DiffusionBackend

In-process DiffusionBackend.load_pipeline / generate on an exclusive B200, with transformer_quant=auto and every other setting at its default. Studio's 16, 12 and 8 GB tiers are simulated: mem_get_info reports that total, and the allocator is capped with set_per_process_memory_fraction. Each process is fresh, with fresh compile caches. The host was over-subscribed for part of the run (load average up to 780 on 192 cores), so every arm of a tier is loaded at once and the renders run one at a time, with the arm order rotating each repeat. Host load then lands on every arm alike. Pre-#12736 runs twice per tier (a and b) as the A/A. Per size: 1 cold render, then 5 warm. The table is the median of warm renders 2 to 5. Steps: 6 (Qwen-Image, Qwen-Image-2.1, HunyuanImage-2.1, FLUX.1-dev), 8 (SDXL), 4 (FLUX.2-klein). Peak alloc is torch.cuda.max_memory_allocated over the image. Device is this process's NVML memory. RSS is the host resident peak. PSNR and window are measured on the output PNG against the untiled decode of the arm's own latent. The latent is identical in every arm. Full card only for HunyuanImage, SDXL and FLUX.2-klein: no family tiles there, so every arm runs the same code.

• Every tiled Qwen-Image-2.1, Qwen-Image and HunyuanImage cell decodes 1.5 to 28x faster on the head, except one noisy cell (below). Totals are 0.2 to 0.7 s faster where the decode was a visible share. Head images are at 43.8 to 59.8 dB against the untiled decode, mostly 54 to 60 (the level a Studio untiled decode itself reaches here). Pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736's tiled images are at 27.7 to 53.0 dB.
• The decode peak is higher where the head decodes in larger tiles or untiled: up to +2.7 GB (Qwen-Image-2.1 at 1664 px on 16 GB) and +1.9 GB (HunyuanImage at 2048 px). That is by design. A tile is only that large when three quarters of the free VRAM holds it, and no cell OOMed.
• HunyuanImage at 12 GB, 2048 px: the head's median decode is 3.15 s because 2 of its 5 renders stalled for 5.7 to 6.9 s in streaming offload. Pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736 had the same stall (7.0 s in 1 of 5). The other three head decodes took 0.48 to 0.84 s, against pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736's 0.96 to 1.36 s.
• FLUX.1 at 2048 px, SDXL and FLUX.2-klein decode through the same stock path as pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736, and the images are bit-identical. Their total-time differences are within the A/A spread (SDXL's renders on CPU offload swing from 1.3 to 37 s in every arm).
• The final commit fixed FLUX.1 at 8 GB, 1600 px. Before it, the head decoded untiled out of the allocator's cache and stalled on cache flushes in 3 of 5 renders (0.214 s median decode against pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736's 0.082). It now decodes in 120-latent tiles: 0.065 s against 0.068, a total of 1.83 against 1.93 s (A/A 1.82 / 2.00), and line scores of 0.67 to 0.79 against pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736's 2.4 to 3.6 at the sliver edge. The worst window is 3.53 against 3.12, at the same interior texture window in both, not a seam.

Very low free VRAM at decode time

The same e2e processes at the 12 GB tier. Tiling is forced on, and a filler tensor holds all but 0.6 or 0.35 GiB of the simulated free memory while the VAE decodes the render's latent. Median of 3 after a warm decode, in seconds:

• In this case the head never OOMs, and Studio: seam-free VAE tiles for Qwen-Image-2.1 on low-VRAM tiers (fixes thin horizontal / vertical lines) #12696 does at 0.35 GiB. Where even the floor tile cannot fit, Qwen-Image-2.1 and HunyuanImage take the same stock tiles as pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736 (the same images), and the log says why. The HunyuanImage time differences there are the same code path, within pre a / pre b's own spread.
• FLUX.1, SDXL and FLUX.2-klein OOM at these levels in every arm. Their stock tile is a 1,408 MiB 1024 px tile. The head never asks for more than stock there.
• In the decode-alone runs above, the one case still slower than pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736 is Qwen-Image at 0.35 GiB free and 1664 to 2048 px: 0.494 and 0.952 s against 0.406 and 0.601 s. The widest tile that fits is 38 to 40 latents with a 16-latent overlap, which decodes more latents than stock's 32-latent tiles with 8-latent blends. At 0.6 GiB free it is already 2x faster than pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736. The Studio tiers above never reached this case: Qwen-Image had 1.7 to 10.8 GiB free at decode time in every cell.

Real low-VRAM cards

Colab L4 (22.5 GB), Kaggle T4 (15 GB) and Colab T4 High-RAM, plus Colab G4 (RTX PRO 6000) for HunyuanImage, which needs about 77 GB of host RAM at load, more than an L4 VM has. Studio is installed by its own install.sh. The 12 and 16 GB tiers are a real VRAM ballast plus a reported total. Pre-#12736 and head are interleaved, one after the other twice, with a fresh process each. Warm totals and decodes are medians of 4 renders, at 6 steps.

• The pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736 A/A spread was 0.0 to 0.9 s in total and at most 0.03 s in decode. The decode gains are well outside it. Total-time differences under 1 s on L4 and G4 are not.
• These cards had 2 to 10 GB free at decode time, so the head decoded untiled in every cell and matched untiled (54 to 69 dB). Pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736's stock tiles were at 28 to 54 dB.
• The cloud cells ran head 00dbb42e78. The two later commits change only the AutoencoderKL / FLUX.2 sliver path, so the Qwen-Image-2.1, Qwen-Image and HunyuanImage cells are unchanged at a4a70e3a74. The L4 FLUX.1 cell was rerun on a4a70e3a74 (last row). Once model CPU offload has parked the denoiser, the device has room for one untiled 1600 px decode there, so it stays untiled, as at 00dbb42e78. On the T4, with 8 GB free at decode time, FLUX.1's 1600 px decode is one tile on either head.

Tests

• tests/test_diffusion_vae_tiling_general.py (2,388 tests, CPU, tiny diffusers VAEs of each class):
  • geometry and ratio read off AutoencoderKLQwenImage (8x), AutoencoderKLQwenImage21 (16x), AutoencoderKLHunyuanImage (32x), AutoencoderKL and AutoencoderKLFlux2 (8x);
  • the stock-layout floor per ratio, including the sliver sides;
  • wide tiles full and edge-aligned, with weights summing to 1, at every length and ratio;
  • seam reduction on Qwen-Image and HunyuanImage;
  • AutoencoderKL bit-identical to stock where its tiles meet the floor, and on the wide path for a sliver;
  • which encodes are replaced, and the Qwen-Image encode closer to untiled;
  • the kill switch per class, KEEP_STOCK, and an unreadable geometry keeping stock;
  • every image family's VAE class (read from the diffusers pipeline's vae annotation) is one the rule reads, or is listed in KEEP_STOCK.
• Perf tests added:
  • the per-class figures bound the measured peaks and stay within 10% of them; the fused figure is used only while the fused kernels are installed and not fallen back;
  • at every sliver length from 129 to 419 latents and three budgets, an AutoencoderKL layout makes fewer decoder calls than stock, decodes no more latents per side (up to rounding), and keeps every overlap at 32 or more;
  • the 32-latent floor layout (Qwen-Image, Qwen-Image-2.1, HunyuanImage) is unchanged at every length;
  • a floor that cannot fit decodes in the stock tiles and logs once; Qwen-Image and AutoencoderKL never fall back that way, since their stock tiles need as much;
  • an OOM retries at the floor, then in the stock tiles, and other errors propagate;
  • the stock fallback keeps the fused batched loop, which takes tiled_decode back on uninstall;
  • only AutoencoderKL / FLUX.2 size past the floor without the allocator's cache;
  • geometry and blend weights are computed once per VAE.
• On acc589fbe5 (Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736 as merged, before these commits), 298 of the perf tests fail. On the head all 2,089 of them pass.
• tests/test_diffusion_vae_tiling.py (Studio: seam-free VAE tiles for Qwen-Image-2.1 on low-VRAM tiers (fixes thin horizontal / vertical lines) #12696, 580 tests) passes unchanged.
• The related suites pass, 4,222 tests: diffusion backend, memory, calibrated activation, speed, fused VAE, single-frame VAE, upscale gate, VAE prebuild, channels-last, Qwen-Image-2.1 VAE layout.
• Both tiling suites (2,968 tests) pass on torch 2.6.0 (CPU) and 2.12.1. The module uses no torchao.
• Qwen-Image-2.1 on real weights, Studio: seam-free VAE tiles for Qwen-Image-2.1 on low-VRAM tiers (fixes thin horizontal / vertical lines) #12696's module against this one, at 1024x1024, 1344x1344, 2048x2048, 2752x1536 and 1184x864, with 0.6 to 8 GiB free:
  • without the fused kernels, tiled_decode is torch.equal wherever the floor tile fits;
  • with them, it is torch.equal at the floor tile;
  • the installed vae._encode is torch.equal in every case.

Not covered

• Real 8 and 16 GB consumer cards were not run. The real-card evidence is an L4 and a T4 with VRAM ballast, and a G4 for HunyuanImage. The 8 / 12 / 16 GB end-to-end tiers are simulated on a B200.
• HunyuanImage was not run on an L4 (the VM is short of host RAM) or a T4 (fp16, no GGUF). Qwen-Image FP8 did not fit an L4 VM's host RAM, so the GGUF ran instead.
• Qwen-Image at 0.35 GiB free and 1664 to 2048 px still decodes 1.2 to 1.6x slower than pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736 (decode alone, emulated). No Studio tier here got near that free memory.
• The very-low-free cases are emulated: a filler tensor under the simulated cap, or a capped mem_get_info. On Windows (WDDM), an oversized tile spills to shared memory instead of raising an OOM, so only the floor-fit check protects there, not the OOM retry.
• Part of the end-to-end run had the host over-subscribed. The interleaving puts that load on every arm alike, but SDXL's CPU-offload renders swung 1.3 to 37 s in every arm, so its totals say nothing. Its decode runs the same code as pre-Studio: seam-free tiled decode for every image VAE whose stock tiles fall under the floor (HunyuanImage-2.1 grid, Qwen-Image, slivers) #12736.
• The cloud cells ran 00dbb42e78 (and FLUX.1 on an L4 at a4a70e3a74). The FLUX.1 T4 cell predates the last two commits, which only change the sliver path; with 8 GB free at decode time it decodes one tile on either.
• The per-class figures were measured on a B200 with diffusers 0.41.0.dev0 and torch 2.12.1. Other GPUs and allocator settings, such as expandable_segments, may peak differently. The 25% headroom and the OOM retry cover that, but it was not measured.
• FLUX.2 (and, less so, FLUX.1 and SDXL) shift colour per tile above 1024 px even where the stock tiles meet the floor: 27.3 dB at 2048 px for FLUX.2, a greener bottom-right tile. It is a per-tile offset, not a line, and those sizes keep the stock path here.
• The boundary-gradient seam score does not separate stock from wide on these photos for any VAE, Qwen-Image-2.1 included. The PSNR, the worst 64 px window and the line score carry the result.
• ComfyUI's own OOM fallback for the Qwen-Image VAE uses 32-latent tiles with 8-latent overlaps, as diffusers does. This PR follows its 16-latent 2D decode overlap instead, which measured slightly closer to untiled.
• AMD, Windows and Apple Silicon were not run.