Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 35 additions & 4 deletions src/core/backend_fit.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -479,26 +479,57 @@ namespace sd::backend_fit {
}

bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) {
const char* retry_mode = nullptr;
return prepare_vae_retry_tiling(tiling_params, prefer_temporal_tiling, false);
}

bool prepare_vae_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling, bool encode) {
// Spatial retry tiles are a fraction of the latent: half of it for decode. Image VAE encode tiles are
// scaled up 2x by get_tile_sizes, so encode starts a step lower to actually split the image.
// get_tile_sizes keeps the tiles overlapping (sd_tiling_seam_safe_tile_size), so a fractional tile never
// collapses into two tiles that meet edge to edge. If a retry still fails, the tiles are halved again
// down to an eighth of the latent before giving up.
const float first_rel_size = encode ? 0.25f : 0.5f;
const float min_rel_size = 0.125f;
Comment on lines +491 to +492

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Scale the encode retry floor with its tile factor

For ordinary image VAEs, get_tile_sizes multiplies relative encode tile sizes by encoding_factor = 2.0, but this shared floor stops retrying once rel_size_* reaches 0.125. Consequently encode retries use effective axis fractions of 0.5 and 0.25 and then give up, never trying the advertised one-eighth-sized tile; an encode that still OOMs at one quarter but would fit at one eighth therefore fails unnecessarily. The encode floor must account for the encoding factor, just as first_rel_size already does.

Useful? React with 👍 / 👎.

const char* stage = encode ? "encode" : "decode";
const char* retry_mode = nullptr;
if (prefer_temporal_tiling && !tiling_params.temporal_tiling) {
tiling_params.temporal_tiling = true;
retry_mode = tiling_params.enabled ? "spatial+temporal" : "temporal";
} else if (!tiling_params.enabled) {
tiling_params.enabled = true;
tiling_params.rel_size_x = 0.5f;
tiling_params.rel_size_y = 0.5f;
tiling_params.rel_size_x = first_rel_size;
tiling_params.rel_size_y = first_rel_size;
if (tiling_params.tile_size_x <= 0) {
tiling_params.tile_size_x = 256;
}
if (tiling_params.tile_size_y <= 0) {
tiling_params.tile_size_y = 256;
}
if (tiling_params.target_overlap <= 0.f) {
tiling_params.target_overlap = 0.5f;
}
retry_mode = tiling_params.temporal_tiling ? "spatial+temporal" : "spatial";
} else if (tiling_params.rel_size_x > 0.f && tiling_params.rel_size_x <= 1.f &&
tiling_params.rel_size_y > 0.f && tiling_params.rel_size_y <= 1.f &&
std::max(tiling_params.rel_size_x, tiling_params.rel_size_y) > min_rel_size) {
// An axis already below the floor keeps its size: clamping it up would grow the tile that just failed.
auto shrink = [&](float rel) { return rel > min_rel_size ? std::max(rel * 0.5f, min_rel_size) : rel; };
tiling_params.rel_size_x = shrink(tiling_params.rel_size_x);
tiling_params.rel_size_y = shrink(tiling_params.rel_size_y);
if (tiling_params.target_overlap <= 0.f) {
tiling_params.target_overlap = 0.5f;
}
LOG_WARN("VAE %s failed (likely out of memory); retrying with smaller tiles (%.3g x %.3g of the latent)",
stage,
tiling_params.rel_size_x,
tiling_params.rel_size_y);
return true;
} else {
return false;
}

LOG_WARN("VAE decode failed (likely out of memory); retrying with %s tiling",
LOG_WARN("VAE %s failed (likely out of memory); retrying with %s tiling",
stage,
retry_mode);
return true;
}
Expand Down
6 changes: 6 additions & 0 deletions src/core/backend_fit.h
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,12 @@ namespace sd::backend_fit {
bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params,
bool prefer_temporal_tiling);

// Turns on (or tightens) VAE tiling after an encode/decode failed, most likely out of memory.
// Returns false once there is nothing left to try.
bool prepare_vae_retry_tiling(sd_tiling_params_t& tiling_params,
bool prefer_temporal_tiling,
bool encode);

} // namespace sd::backend_fit

#endif // __SD_BACKEND_FIT_H__
15 changes: 12 additions & 3 deletions src/model/vae/vae.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -184,7 +184,9 @@ struct VAE : public GGMLRunner {
const sd_tiling_params_t& params,
int64_t latent_x,
int64_t latent_y,
float encoding_factor = 1.0f) {
float encoding_factor = 1.0f,
bool circular_x = false,
bool circular_y = false) {
tile_overlap = std::max(std::min(params.target_overlap, 0.5f), 0.0f);
auto get_tile_size = [&](int requested_size, float factor, int64_t latent_size) {
const int default_tile_size = 32;
Expand All @@ -205,6 +207,13 @@ struct VAE : public GGMLRunner {

tile_size_x = get_tile_size(params.tile_size_x, params.rel_size_x, latent_x);
tile_size_y = get_tile_size(params.tile_size_y, params.rel_size_y, latent_y);
// Circular axes wrap their tiles instead of pinning the last one to the edge, so they keep the requested size.
if (!circular_x) {
tile_size_x = sd_tiling_seam_safe_tile_size(static_cast<int>(latent_x), tile_size_x, tile_overlap);
}
if (!circular_y) {
tile_size_y = sd_tiling_seam_safe_tile_size(static_cast<int>(latent_y), tile_size_y, tile_overlap);
}
}

virtual sd::Tensor<float> encode(int n_threads,
Expand All @@ -230,7 +239,7 @@ struct VAE : public GGMLRunner {
// tiles to 64 latent pixels so a 512px SD image is encoded as one tile.
const float encode_tile_factor = sd_version_is_minimax_h3(version) ? 1.f : (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) ? 1.30539f
: 2.0f;
get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, encode_tile_factor);
get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, encode_tile_factor, circular_x, circular_y);
LOG_VERBOSE("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
output = tiled_compute(input,
n_threads,
Expand Down Expand Up @@ -280,7 +289,7 @@ struct VAE : public GGMLRunner {
int64_t H = input.shape()[1] * scale_factor;
float tile_overlap;
int tile_size_x, tile_size_y;
get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, input.shape()[0], input.shape()[1]);
get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, input.shape()[0], input.shape()[1], 1.0f, circular_x, circular_y);
if (!silent) {
LOG_VERBOSE("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
}
Expand Down
7 changes: 7 additions & 0 deletions src/pipeline/diffusion_engine.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2740,6 +2740,13 @@ int StableDiffusionGGML::align_video_frames(int frames) {

sd::Tensor<float> StableDiffusionGGML::encode_to_vae_latents(const sd::Tensor<float>& x) {
auto latents = first_stage_model->encode(n_threads, x, vae_tiling_params, circular_x, circular_y);
// Same fallback as decode_first_stage, on a copy: an encode that ran out of memory is retried with
// overlapping tiles for this call only.
sd_tiling_params_t retry_tiling_params = vae_tiling_params;
while (latents.empty() &&
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);
Comment on lines +2746 to +2748

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Reconfigure circular padding before tiled encode retries

When image VAE tiling starts disabled and an encode OOM triggers this retry under circular_x or circular_y, configure_image_vae_axes has already enabled circular padding inside the VAE graph. This newly tiled call also passes the circular flags to process_tiles_2d, so every tile wraps both at the tiler level and internally within its own convolution graph. The explicit-tiling path avoids this by disabling the model's circular axis whenever an axis is split; the retry path needs the same reconfiguration or it produces artifacts at internal tile boundaries for circular img2img/control/reference encodes.

Useful? React with 👍 / 👎.

}
if (latents.empty()) {
return {};
}
Expand Down
5 changes: 4 additions & 1 deletion src/pipeline/image.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,10 @@ namespace sd::pipeline {
overlap,
sd_img_gen_params->vae_tiling_params,
latent_size_x,
latent_size_y);
latent_size_y,
1.0f,
sd->circular_x,
sd->circular_y);

sd->circular_x = sd->circular_x && (tile_size_x >= latent_size_x);
sd->circular_y = sd->circular_y && (tile_size_y >= latent_size_y);
Expand Down
31 changes: 31 additions & 0 deletions src/runtime/tiling.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,37 @@ static void sd_tiling_calc_tiles(int& num_tiles_dim,
}
}

int sd_tiling_seam_safe_tile_size(int small_dim, int tile_size, float tile_overlap_factor) {
// sd_tiling_calc_tiles lets the real overlap fall well below the target for some sizes. The worst case is a
// tile a little over half of the axis (e.g. round(dim / 2)): it falls back to two tiles overlapping by
// 2 * tile_size - small_dim, which is 0 or 1 latent, so the cross-fade has nothing to blend over and the tile
// edge shows as a line. Pick the nearest tile size (smaller first; a larger one only up to twice the request,
// so an explicitly small tile cannot balloon) whose tiles overlap by at least half the target and by 2+ latents.
const int min_tile_size = 4;
if (tile_overlap_factor <= 0.f || tile_size >= small_dim) {
return tile_size;
}
auto overlaps_enough = [&](int size) {
int num_tiles;
float overlap_factor;
sd_tiling_calc_tiles(num_tiles, overlap_factor, small_dim, size, tile_overlap_factor, false);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Skip non-circular sizing for circular VAE axes

When --circularx or --circulary is combined with VAE tiling, this helper still selects the replacement size using non-circular geometry. For example, a latent axis of 57 with the default tile size 32 is changed to 28: the previous circular layout visits offsets 0, 19, and 38 and has unit aggregate blend weight, while size 28 makes process_tiles_2d visit 0, 14, 28, 42, and 56 even though the calculation reports four tiles, causing 26 positions to receive up to approximately twice the intended weight. This corrupts the tileable output with a bright band, so circular state needs to be included in the sizing decision or this adjustment should be bypassed for circular axes.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in ee5829c: get_tile_sizes now takes the circular flags and skips sd_tiling_seam_safe_tile_size on a circular axis, so those axes keep the tile size master used. Confirmed with a constant-input process_tiles_2d check: on a circular 57-latent axis the default 32 tile sums to unit weight, while the adjusted 28 tile does not.

const int overlap = static_cast<int>(size * overlap_factor);
const int min_overlap = std::max(2, static_cast<int>(size * tile_overlap_factor * 0.5f));
return num_tiles < 2 || overlap >= min_overlap;
};
for (int size = tile_size; size >= min_tile_size; --size) {
if (overlaps_enough(size)) {
return size;
}
}
for (int size = tile_size + 1; size < small_dim && size <= 2 * tile_size; ++size) {
if (overlaps_enough(size)) {
return size;
}
}
return tile_size;
}

static int64_t sd_tensor_plane_size(const sd::Tensor<float>& tensor) {
GGML_ASSERT(tensor.dim() >= 2);
return tensor.shape()[0] * tensor.shape()[1];
Expand Down
4 changes: 4 additions & 0 deletions src/runtime/tiling.h
Original file line number Diff line number Diff line change
Expand Up @@ -19,4 +19,8 @@ sd::Tensor<float> process_tiles_2d(const sd::Tensor<float>& input,
const TileProcessCallback& on_processing,
bool silent = false);

// Returns tile_size, adjusted if needed so a non-circular axis of small_dim is never split into tiles that
// barely overlap (see the definition for why that shows as a seam).
int sd_tiling_seam_safe_tile_size(int small_dim, int tile_size, float tile_overlap_factor);

#endif // __SD_RUNTIME_TILING_H__