VAE tiling: decode 2x2 tiles again at 456-504 px while keeping the seam fix

docs/performance.md

@@ -50,9 +50,13 @@ overlap. For example, `0.5x0.5` uses half the width and height in both encode an
 decode. The target overlap is clamped to 0 through 0.5 and the actual overlap is
 adjusted to fit the image. On an axis that does not wrap, a tile size that would leave
 neighboring tiles overlapping by less than half the target (or by fewer than 2 latent
-pixels) is adjusted to the nearest size that does not: smaller first, larger only up to
-twice the requested size. Without it, a tile a little over half of the axis gives two
-tiles that barely overlap and the tile edge shows as a line. Size and overlap options
+pixels) is replaced. The replacement is the size that decodes the fewest latent pixels
+along the axis (tiles times size) among sizes whose overlaps reach the full target overlap,
+capped at 8 latent pixels, preferring sizes within 25% above the request. It never uses
+more tiles than the nearest size that overlaps by half the target (smaller first, larger
+only up to twice the request), and falls back to that size when nothing better qualifies.
+Without it, a tile a little over half of the axis gives two tiles that barely overlap and
+the tile edge shows as a line. Size and overlap options
 require `--vae-tiling`.
 
 **Migration:** `--vae-tile-size` and the C/JSON fields `tile_size_w` and

src/model/vae/vae.hpp

@@ -166,6 +166,10 @@ struct VAE : public GGMLRunner {
     }
 
 public:
+    // OOM-retry cap on latent tile sizes (0 = none).
+    int retry_tile_limit_w = 0;
+    int retry_tile_limit_h = 0;
+
     VAE(SDVersion version,
         ggml_backend_t backend,
         const std::string& weight_prefix                    = "",
@@ -249,10 +253,16 @@ struct VAE : public GGMLRunner {
         }
         // Circular axes wrap their tiles instead of pinning the last one to the edge, so they keep the requested size.
         if (!circular_x && !circular_x_enabled) {
-            tile_size_w = sd_tiling_seam_safe_tile_size(static_cast<int>(latent_w), tile_size_w, tile_overlap);
+            tile_size_w = sd_tiling_seam_safe_tile_size(static_cast<int>(latent_w), tile_size_w, tile_overlap,
+                                                        retry_tile_limit_w);
+        } else if (retry_tile_limit_w > 0) {
+            tile_size_w = std::min(tile_size_w, std::max(4, retry_tile_limit_w));
         }
         if (!circular_y && !circular_y_enabled) {
-            tile_size_h = sd_tiling_seam_safe_tile_size(static_cast<int>(latent_h), tile_size_h, tile_overlap);
+            tile_size_h = sd_tiling_seam_safe_tile_size(static_cast<int>(latent_h), tile_size_h, tile_overlap,
+                                                        retry_tile_limit_h);
+        } else if (retry_tile_limit_h > 0) {
+            tile_size_h = std::min(tile_size_h, std::max(4, retry_tile_limit_h));
         }
         return true;
     }

src/pipeline/diffusion_engine.cpp

@@ -2895,6 +2895,38 @@ int StableDiffusionGGML::align_video_frames(int frames) {
     return latent_frames_to_video_frames(video_frames_to_latent_frames(frames));
 }
 
+// After an OOM (last_w/last_h > 0) cap below the failed size: get_tile_sizes may widen a halved request back to it.
+static bool sd_vae_tile_sizes(VAE& vae,
+                              const sd_tiling_params_t& tiling_params,
+                              int64_t latent_w,
+                              int64_t latent_h,
+                              bool circular_x,
+                              bool circular_y,
+                              int last_w,
+                              int last_h,
+                              int& tile_size_w,
+                              int& tile_size_h) {
+    vae.retry_tile_limit_w = last_w > 0 ? std::max(4, last_w - 1) : 0;
+    vae.retry_tile_limit_h = last_h > 0 ? std::max(4, last_h - 1) : 0;
+    tile_size_w            = static_cast<int>(latent_w);
+    tile_size_h            = static_cast<int>(latent_h);
+    float tile_overlap;
+    if (tiling_params.enabled &&
+        !vae.get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling_params,
+                            latent_w, latent_h, circular_x, circular_y)) {
+        return false;
+    }
+    return last_w <= 0 || !tiling_params.enabled || tile_size_w < last_w || tile_size_h < last_h;
+}
+
+struct VAERetryTileLimitGuard {
+    VAE& vae;
+    ~VAERetryTileLimitGuard() {
+        vae.retry_tile_limit_w = 0;
+        vae.retry_tile_limit_h = 0;
+    }
+};
+
 sd::Tensor<float> StableDiffusionGGML::encode_to_vae_latents(const sd::Tensor<float>& x) {
     // Same out-of-memory fallback as decode_first_stage, on a copy of the tiling parameters, so the
     // smaller tiles apply to this call only.
@@ -2903,17 +2935,12 @@ sd::Tensor<float> StableDiffusionGGML::encode_to_vae_latents(const sd::Tensor<fl
     const int64_t latent_w  = x.shape()[0] / scale_factor;
     const int64_t latent_h  = x.shape()[1] / scale_factor;
     sd::Tensor<float> latents;
+    VAERetryTileLimitGuard limit_guard{*first_stage_model};
     for (int last_w = 0, last_h = 0;;) {
-        int tile_size_w = static_cast<int>(latent_w);
-        int tile_size_h = static_cast<int>(latent_h);
-        float tile_overlap;
-        if (tiling_params.enabled &&
-            !first_stage_model->get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling_params,
-                                               latent_w, latent_h, circular_x, circular_y)) {
-            return {};
-        }
-        // get_tile_sizes can widen a retry tile to keep the overlap, so stop once a retry no longer shrinks it.
-        if (tile_size_w == last_w && tile_size_h == last_h) {
+        int tile_size_w;
+        int tile_size_h;
+        if (!sd_vae_tile_sizes(*first_stage_model, tiling_params, latent_w, latent_h, circular_x, circular_y,
+                               last_w, last_h, tile_size_w, tile_size_h)) {
             return {};
         }
         latents = first_stage_model->encode(n_threads, x, tiling_params, circular_x, circular_y);
@@ -2951,17 +2978,12 @@ sd::Tensor<float> StableDiffusionGGML::decode_first_stage(const sd::Tensor<float
     auto tiling_params                = first_stage_model->resolve_tiling_params(vae_tiling_params);
     const bool prefer_temporal_tiling = decode_video && latents.dim() == 5 && latents.shape()[2] > 1 &&
                                         first_stage_model->can_temporal_tile_decode();
+    VAERetryTileLimitGuard limit_guard{*first_stage_model};
     for (int last_w = 0, last_h = 0;;) {
-        int tile_size_w = static_cast<int>(latents.shape()[0]);
-        int tile_size_h = static_cast<int>(latents.shape()[1]);
-        float tile_overlap;
-        if (tiling_params.enabled &&
-            !first_stage_model->get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling_params,
-                                               latents.shape()[0], latents.shape()[1], circular_x, circular_y)) {
-            return {};
-        }
-        // get_tile_sizes can widen a retry tile, so stop once a spatial retry no longer shrinks it.
-        if (tile_size_w == last_w && tile_size_h == last_h) {
+        int tile_size_w;
+        int tile_size_h;
+        if (!sd_vae_tile_sizes(*first_stage_model, tiling_params, latents.shape()[0], latents.shape()[1],
+                               circular_x, circular_y, last_w, last_h, tile_size_w, tile_size_h)) {
             return {};
         }
         auto decoded = first_stage_model->decode(n_threads, latents, tiling_params, decode_video, circular_x, circular_y);

src/runtime/tiling.cpp

@@ -1,6 +1,7 @@
 #include "runtime/tiling.h"
 
 #include <algorithm>
+#include <climits>
 #include <cstdint>
 #include <cstdlib>
 #include <exception>
@@ -60,35 +61,84 @@ static std::vector<TileSpan> sd_tiling_plan_axis(int dimension,
     return tiles;
 }
 
-int sd_tiling_seam_safe_tile_size(int small_dim, int tile_size, float tile_overlap_factor) {
-    // A tile a little over half the axis (32 of 62 latents) gets two tiles overlapping by 1-2 latents, which shows
-    // as a seam. Pick the nearest size (smaller first, larger only up to 2x) whose overlaps are >= half the target
-    // and >= 2 latents.
+int sd_tiling_seam_safe_tile_size(int small_dim, int tile_size, float tile_overlap_factor, int max_tile_size) {
+    // Overlaps of 1-2 latents (32-latent tiles on a 62 axis) show as a seam; require >= half the target and >= 2.
     const int min_tile_size = 4;
+    const int limit         = max_tile_size > 0 ? std::max(max_tile_size, min_tile_size) : INT_MAX;
+    tile_size               = std::min(tile_size, limit);
     if (tile_overlap_factor <= 0.f || tile_size >= small_dim) {
         return tile_size;
     }
-    auto overlaps_enough = [&](int size) {
-        const int min_overlap = std::max(2, static_cast<int>(size * tile_overlap_factor * 0.5f));
-        const auto tiles      = sd_tiling_plan_axis(small_dim, size, tile_overlap_factor, false);
+    auto overlaps_enough = [&](int size, bool replacement) {
+        int min_overlap = std::max(2, static_cast<int>(size * tile_overlap_factor * 0.5f));
+        if (replacement) {
+            // Half the target still leaves a line on 16x VAEs (Qwen-Image-2.1 at 480 px); 8 latents does not.
+            min_overlap = std::max(min_overlap, std::min(8, static_cast<int>(size * tile_overlap_factor)));
+        }
+        const auto tiles = sd_tiling_plan_axis(small_dim, size, tile_overlap_factor, false);
         for (size_t i = 0; i + 1 < tiles.size(); ++i) {
             if (tiles[i].overlap_after < min_overlap) {
                 return false;
             }
         }
         return true;
     };
-    for (int size = tile_size; size >= min_tile_size; --size) {
-        if (overlaps_enough(size)) {
-            return size;
+    if (overlaps_enough(tile_size, false)) {
+        return tile_size;
+    }
+    auto tile_count = [&](int size) {
+        return static_cast<int>(sd_tiling_plan_axis(small_dim, size, tile_overlap_factor, false).size());
+    };
+    // The nearest passing size (smaller first, then larger up to 2x) caps the tile count so tiny tiles never win.
+    int nearest = 0;
+    for (int size = tile_size - 1; size >= min_tile_size && nearest == 0; --size) {
+        if (overlaps_enough(size, false)) {
+            nearest = size;
+        }
+    }
+    for (int size = tile_size + 1; nearest == 0 && size < small_dim && size <= std::min(2 * tile_size, limit); ++size) {
+        if (overlaps_enough(size, false)) {
+            nearest = size;
         }
     }
-    for (int size = tile_size + 1; size < small_dim && size <= 2 * tile_size; ++size) {
-        if (overlaps_enough(size)) {
+    if (nearest == 0) {
+        return tile_size;
+    }
+    const int max_tiles = tile_count(nearest);
+    auto pick = [&](int lowest, int highest, bool replacement) {
+        int best_size     = 0;
+        int64_t best_cost = 0;
+        int best_tiles    = 0;
+        for (int size = lowest; size <= highest; ++size) {
+            if (!overlaps_enough(size, replacement)) {
+                continue;
+            }
+            const int tiles = tile_count(size);
+            if (tiles > max_tiles) {
+                continue;
+            }
+            const int64_t cost = static_cast<int64_t>(tiles) * size;
+            const bool better = best_size == 0 || cost < best_cost || (cost == best_cost && tiles < best_tiles);
+            if (better) {
+                best_size  = size;
+                best_cost  = cost;
+                best_tiles = tiles;
+            }
+        }
+        return best_size;
+    };
+    const int capped = std::min({small_dim - 1, tile_size + tile_size / 4, limit});
+    const int widest = std::min({small_dim - 1, 2 * tile_size, limit});
+    for (bool replacement : {true, false}) {
+        int size = pick(min_tile_size, capped, replacement);
+        if (size == 0) {
+            size = pick(capped + 1, widest, replacement);
+        }
+        if (size != 0) {
             return size;
         }
     }
-    return tile_size;
+    return nearest;
 }
 
 static int64_t sd_tensor_plane_size(const sd::Tensor<float>& tensor) {

src/runtime/tiling.h

@@ -44,8 +44,7 @@ sd::Tensor<float> process_tiles_2d_batched(const sd::Tensor<float>& input,
                                            const TileBatchProcessCallback& 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);
+// tile_size adjusted so tiles never barely overlap; max_tile_size > 0 (OOM retry) caps it, never below 4.
+int sd_tiling_seam_safe_tile_size(int small_dim, int tile_size, float tile_overlap_factor, int max_tile_size = 0);
 
 #endif  // __SD_RUNTIME_TILING_H__