Repository navigation
Keep VAE tiles overlapping on the out-of-memory retry (no more seam lines), and retry the encode too - #7
Keep VAE tiles overlapping on the out-of-memory retry (no more seam lines), and retry the encode too#7oobabooga wants to merge 4 commits into
Conversation
When an untiled VAE decode fails, the retry tiled at half the latent per axis. For most sizes that is a tile one latent over half, which sd_tiling_calc_tiles turns into two tiles overlapping by 0 or 1 latent: the cross-fade has nothing to blend over and the tile edge shows as a hard line through the image (Qwen-Image-2.1 1184x864: 110 levels). get_tile_sizes now picks the nearest tile size whose tiles overlap by at least half the target (never under 2 latents), preferring smaller tiles so memory never grows. This also fixes the --vae-tiling path at the few sizes where it hit the same two-tile case (57..63 latents decode, 113..127 encode). The retry also halves the tiles again (down to 1/8 of the latent) when the first tiled attempt still does not fit, and an untiled VAE encode that fails is now retried with tiles the same way instead of failing the job.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
Codex Review: Didn't find any major issues. More of your lovely PRs please. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ee5829cb00
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| sd::backend_fit::prepare_vae_retry_tiling(retry_tiling_params, false, true)) { | ||
| latents = first_stage_model->encode(n_threads, x, retry_tiling_params, circular_x, circular_y); |
There was a problem hiding this comment.
When a request enables a circular axis without upfront VAE tiling, configure_image_vae_axes leaves circular padding enabled inside the VAE graph. If this new fallback then enables spatial tiling after an encode OOM, the retry also passes circular_x/circular_y to the tile merger without switching the VAE graph to non-circular padding, so every individual tile wraps its own convolution boundaries while the tile grid wraps again. This produces incorrect interior tile boundaries for circular img2img/reference-image requests; the retry must apply the same mutually exclusive graph-padding versus tile-wrapping configuration used for upfront tiling.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Measured, no visible defect: Z-Image-Turbo 1008x1008 with --circular, decode retry under a real 5 GiB VRAM cap, against the uncapped untiled circular decode: 44.88 dB, seam score 2.24 as is, and 44.46 dB, 2.29 with the graph's circular padding switched off for the retry. Only each tile's outer pixels see the doubled wrap, and the smootherstep blend weights those near zero. Same retry path as master's decode retry; not changing it here.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
|
Codex Review: Didn't find any major issues. Keep them coming! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
2 participants
Problem
When
--vae-tilingis not passed and the untiled VAE decode runs out of memory,decode_first_stageretries with spatial tiling atrel_size 0.5, half the latent per axis (backend_fit.cpp,prepare_vae_decode_retry_tiling). For most sizes that is a tile one latent over half of the axis (round(dim / 2)), andsd_tiling_calc_tilesturns it into two tiles overlapping by2 * tile - dim= 0 or 1 latent. The smootherstep cross-fade then has nothing to blend over, so the tile edge shows as a hard line through the image. With Qwen-Image-2.1 (16x VAE) the edges also carry the VAE's border artifacts, so it is a broken band rather than a thin line.Examples of the retry geometry on master:
The same two-tile case is reachable from
--vae-tilingtoo, at the few sizes where the default 32-latent tile is just over half the axis: 57 to 63 latents on decode (Qwen-Image-2.1 992x992 overlaps by 2 latents) and 113 to 127 latents on encode (64-latent tiles).The encode side had no retry at all: an untiled VAE encode that runs out of memory fails the whole job (
vae encode compute failed).Change
sd_tiling_seam_safe_tile_size(runtime/tiling.cpp), applied per axis at the end ofVAE::get_tile_sizes: if the tile size the caller asked for would leave tiles overlapping by less than half the target overlap (or under 2 latents), use the nearest tile size that does not, trying smaller tiles first and never more than twice the requested size, so an explicitly small tile stays small. Sizes that already overlap enough are untouched, so the default--vae-tilinggeometry only changes at the sizes listed above.rel_size 0.5for its first attempt (now 36x26 instead of 37x27 at 1184x864, 3x3 tiles with about 0.47 overlap). If that still fails, it halves the tiles again, down to 1/8 of the latent, before giving up (an axis already below 1/8 keeps its size). Before, a failed first retry ended the job.encode_to_vae_latentsretries a failed encode the same way, on a copy of the tiling params so later calls are unaffected. Encode starts atrel_size 0.25because image VAE encode tiles are scaled up 2x inget_tile_sizes.Results
B200, one exclusive card per run, master (
1d02858, the build Studio pins; master only adds a CI commit) against this branch, same seed and prompt. The VRAM cap is real: a second process holds the rest of the card so the untiled decode's allocation fails and the retry runs. Seam score = the largest mean absolute difference along any row or column in a 64 px band, against the untiled decode of the same latent from the same binary (0 to 255 levels).--vae-tiling--vae-tiling--vae-tilingvae encode compute failed. This branch retries the encode with 64x64 tiles (3x3), then retries the decode twice, and finishes.Decode time (
decode_first_stage, seconds, including the failed untiled attempt):--vae-tilingpassed up frontThe retry is about 1 s slower than master's because it decodes 9 tiles instead of 4. It only runs after an out-of-memory failure.
Before / after
Out-of-memory retry under a VRAM cap. Left: untiled decode. Middle: master's retry (2x2 tiles, no overlap). Right: this branch (3x3 overlapping tiles). Top: Qwen-Image-2.1 Q4_K_M 1184x864. Bottom: Z-Image-Turbo Q4_K_M 1008x1008.
Tests
process_tiles_2dwith an identity tile function that adds a different constant to each tile, for every axis length from 8 to 256 latents and each geometry used here (retry at 0.5, 0.25 and 0.125 of the latent, encode retry, default 32-latent decode and 64-latent encode tiles). It checks that a constant input comes back unchanged (blend weights sum to 1) and that no step between neighbouring outputs is as large as the per-tile offset. On master it fails for 186 of 249 lengths with the 0.5 retry, 67 with the encode retry, and at 63 latents (default decode) and 127 latents (encode). With this PR it passes for all of them.sd-cli, no new warnings in the touched files.Validation on other hardware
vae encode compute failed); this branch retries the encode and the decode and finisheshipMallocfailure (6.3 GiB free)--vae-tilingZ-Image 496x496 (62 latents) on CUDA, HIP, Vulkan (Linux and Windows), MetalMerged with the MiniMax-H3 stack (unslothai#17 to unslothai#20), the only conflict is
src/runtime/tiling.h(both add a declaration; keep both). The combined tree gives the same retry and tiling output as this branch alone.Not covered
get_tile_sizes, so they only change at the thin-overlap sizes.binbcast.cu(GGML_ASSERT(s02 <= UINT32_MAX)) on master and on this branch alike. That is not an out-of-memory failure, so the retry does not run. It is not fixed here.