diff --git a/src/core/backend_fit.cpp b/src/core/backend_fit.cpp index 85cfc727c3..0eddb568ad 100644 --- a/src/core/backend_fit.cpp +++ b/src/core/backend_fit.cpp @@ -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; + 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; } diff --git a/src/core/backend_fit.h b/src/core/backend_fit.h index 9ef298b3b4..6db5315847 100644 --- a/src/core/backend_fit.h +++ b/src/core/backend_fit.h @@ -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__ diff --git a/src/model/vae/vae.hpp b/src/model/vae/vae.hpp index 4d8d13517b..8383f7cf56 100644 --- a/src/model/vae/vae.hpp +++ b/src/model/vae/vae.hpp @@ -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; @@ -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(latent_x), tile_size_x, tile_overlap); + } + if (!circular_y) { + tile_size_y = sd_tiling_seam_safe_tile_size(static_cast(latent_y), tile_size_y, tile_overlap); + } } virtual sd::Tensor encode(int n_threads, @@ -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, @@ -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); } diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index e41c1151a8..c30112715a 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -2740,6 +2740,13 @@ int StableDiffusionGGML::align_video_frames(int frames) { sd::Tensor StableDiffusionGGML::encode_to_vae_latents(const sd::Tensor& 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); + } if (latents.empty()) { return {}; } diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp index a48d04d13d..8c6ffffa4d 100644 --- a/src/pipeline/image.cpp +++ b/src/pipeline/image.cpp @@ -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); diff --git a/src/runtime/tiling.cpp b/src/runtime/tiling.cpp index f0a9dbc73b..daa2e0eb5e 100644 --- a/src/runtime/tiling.cpp +++ b/src/runtime/tiling.cpp @@ -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); + const int overlap = static_cast(size * overlap_factor); + const int min_overlap = std::max(2, static_cast(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& tensor) { GGML_ASSERT(tensor.dim() >= 2); return tensor.shape()[0] * tensor.shape()[1]; diff --git a/src/runtime/tiling.h b/src/runtime/tiling.h index a832a644e3..d97949ce59 100644 --- a/src/runtime/tiling.h +++ b/src/runtime/tiling.h @@ -19,4 +19,8 @@ sd::Tensor process_tiles_2d(const sd::Tensor& 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__