From 19bbbca1c736bbb9538679fc0ae690cb2b46b492 Mon Sep 17 00:00:00 2001 From: leejet Date: Fri, 25 Sep 2026 18:20:32 +0800 Subject: [PATCH] refactor: define VAE tile dimensions in image pixels (#2059) --- docs/performance.md | 48 +++++++++++ examples/common/common.cpp | 58 ++++++------- examples/server/api.md | 42 +++++----- examples/server/frontend | 2 +- examples/server/routes_sdcpp.cpp | 8 +- include/stable-diffusion.h | 10 ++- src/core/backend_fit.cpp | 43 +++++++--- src/core/backend_fit.h | 5 +- src/model/vae/minimax_h3_vae.hpp | 18 ++-- src/model/vae/vae.hpp | 132 +++++++++++++++++++----------- src/pipeline/diffusion_engine.cpp | 26 ++++-- src/pipeline/image.cpp | 28 ++++--- src/runtime/tiling.cpp | 48 +++++------ src/runtime/tiling.h | 4 +- 14 files changed, 298 insertions(+), 174 deletions(-) diff --git a/docs/performance.md b/docs/performance.md index 8e292081..6797ad4b 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -21,6 +21,54 @@ CPU fallback. It excludes weights and cache buffers. Within a runner lifecycle, the summary is printed only on the first graph or when backend capacities or the segment count change. +## Use VAE tiling to reduce encode and decode memory usage. + +`--vae-tiling` enables spatial tiling for both VAE encoding and decoding. The +default tile size is 256x256 **image pixels**, independent of the VAE scale factor: + +```shell +--vae-tiling --vae-tile-size 256x256 --vae-tile-overlap 0.5 +``` + +`--vae-tile-size` accepts one size or `WIDTHxHEIGHT`. A zero dimension uses the +256-pixel default. Sizes are rounded down to a multiple of the VAE scale factor +and capped at the current input dimensions. Explicit sizes below four latent +pixels per axis (or the full axis when it is smaller) are rejected. Encoding and +decoding use the same spatial sizes, without an additional encoding multiplier. +Inputs that fit within a tile are processed as one tile. + +For a 512x512 image with the default 50% overlap, both encoding and decoding use +3x3 tiles. A 256-pixel tile corresponds to 32 latent pixels for an 8x VAE, 16 for +a 16x VAE, and 8 for a 32x VAE. Smaller tiles reduce each graph's memory demand, +but overlapping work can increase processing time and tiling can affect image +quality, especially during encoding. Use larger tiles when more context is needed. + +`--vae-relative-tile-size` overrides the absolute size on each axis with a positive +value. Values up to and including 1 specify a fraction of the current input size; +values greater than 1 specify a target number of tiles per axis, accounting for +overlap. For example, `0.5x0.5` uses half the width and height in both encode and +decode. The target overlap is clamped to 0 through 0.5 and the actual overlap is +adjusted to fit the image. Size and overlap options require `--vae-tiling`. + +**Migration:** `--vae-tile-size` and the C/JSON fields `tile_size_w` and +`tile_size_h` now use image pixels instead of latent units. The C/JSON fields +`tile_size_x/y` have been renamed to `tile_size_w/h`, and `rel_size_x/y` to +`rel_size_w/h`. The command-line option names are unchanged. For example, an old +decode tile size of 32 corresponds to 256 pixels for an 8x VAE or 512 pixels for a +16x VAE. Encoding no longer enlarges explicit or relative tile sizes. + +The main VAE decode path retries allocation failures with smaller tiles, even +without `--vae-tiling`. Supported video VAEs first try temporal tiling; spatial +retries use at most 256-pixel tiles initially and then halve the effective tile +dimensions down to the minimum size. Each spatial retry must reduce the effective +tile size. These runtime adjustments do not change the caller's parameters. +Execution failures are not retried, and encoding has no automatic OOM retry. + +`--temporal-tiling` remains independent of spatial tiling. MiniMax H3 always uses +spatial tiling (256x256 pixels and 25% overlap by default) and its own temporal +windows. With `--vae-tiling`, its overlap follows `--vae-tile-overlap`; explicit +spatial sizes are honored. + ## Offload weights to the CPU to save VRAM without reducing generation speed. Using `--offload-to-cpu` allows you to offload weights to the CPU, saving VRAM without reducing generation speed. diff --git a/examples/common/common.cpp b/examples/common/common.cpp index fa7354bc..dc15b4e2 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -1341,7 +1341,7 @@ ArgOptions SDGenerationParams::get_options() { &embed_image_metadata}, {"", "--vae-tiling", - "process vae in tiles to reduce memory usage", + "process vae encode and decode in spatial tiles to reduce memory usage (default: 256x256 image pixels)", true, &vae_tiling_params.enabled}, {"", @@ -1605,12 +1605,12 @@ ArgOptions SDGenerationParams::get_options() { size_t x_pos = tile_size_str.find('x'); try { if (x_pos != std::string::npos) { - std::string tile_x_str = tile_size_str.substr(0, x_pos); - std::string tile_y_str = tile_size_str.substr(x_pos + 1); - vae_tiling_params.tile_size_x = std::stoi(tile_x_str); - vae_tiling_params.tile_size_y = std::stoi(tile_y_str); + std::string tile_w_str = tile_size_str.substr(0, x_pos); + std::string tile_h_str = tile_size_str.substr(x_pos + 1); + vae_tiling_params.tile_size_w = std::stoi(tile_w_str); + vae_tiling_params.tile_size_h = std::stoi(tile_h_str); } else { - vae_tiling_params.tile_size_x = vae_tiling_params.tile_size_y = std::stoi(tile_size_str); + vae_tiling_params.tile_size_w = vae_tiling_params.tile_size_h = std::stoi(tile_size_str); } } catch (const std::invalid_argument&) { return -1; @@ -1628,12 +1628,12 @@ ArgOptions SDGenerationParams::get_options() { size_t x_pos = rel_size_str.find('x'); try { if (x_pos != std::string::npos) { - std::string rel_x_str = rel_size_str.substr(0, x_pos); - std::string rel_y_str = rel_size_str.substr(x_pos + 1); - vae_tiling_params.rel_size_x = std::stof(rel_x_str); - vae_tiling_params.rel_size_y = std::stof(rel_y_str); + std::string rel_w_str = rel_size_str.substr(0, x_pos); + std::string rel_h_str = rel_size_str.substr(x_pos + 1); + vae_tiling_params.rel_size_w = std::stof(rel_w_str); + vae_tiling_params.rel_size_h = std::stof(rel_h_str); } else { - vae_tiling_params.rel_size_x = vae_tiling_params.rel_size_y = std::stof(rel_size_str); + vae_tiling_params.rel_size_w = vae_tiling_params.rel_size_h = std::stof(rel_size_str); } } catch (const std::invalid_argument&) { return -1; @@ -1763,11 +1763,11 @@ ArgOptions SDGenerationParams::get_options() { on_scm_policy_arg}, {"", "--vae-tile-size", - "tile size for vae tiling in latent units, not image pixels, format [X]x[Y] (default: 32x32)", + "tile size for vae encode and decode in image pixels, format [W]x[H] or [S] (default: 256x256; requires --vae-tiling)", on_tile_size_arg}, {"", "--vae-relative-tile-size", - "relative tile size for vae tiling, format [X]x[Y], in fraction of image size if < 1, in number of tiles per dim if >=1 (overrides --vae-tile-size)", + "relative tile size for vae encode and decode, format [W]x[H] or [S]: <=1 is a dimension fraction, >1 a target tile count (overrides --vae-tile-size; requires --vae-tiling)", on_relative_tile_size_arg}, {"", "--prompt-file", @@ -2224,20 +2224,20 @@ bool SDGenerationParams::from_json_str( if (tiling_json.contains("temporal_tiling") && tiling_json["temporal_tiling"].is_boolean()) { vae_tiling_params.temporal_tiling = tiling_json["temporal_tiling"]; } - if (tiling_json.contains("tile_size_x") && tiling_json["tile_size_x"].is_number_integer()) { - vae_tiling_params.tile_size_x = tiling_json["tile_size_x"]; + if (tiling_json.contains("tile_size_w") && tiling_json["tile_size_w"].is_number_integer()) { + vae_tiling_params.tile_size_w = tiling_json["tile_size_w"]; } - if (tiling_json.contains("tile_size_y") && tiling_json["tile_size_y"].is_number_integer()) { - vae_tiling_params.tile_size_y = tiling_json["tile_size_y"]; + if (tiling_json.contains("tile_size_h") && tiling_json["tile_size_h"].is_number_integer()) { + vae_tiling_params.tile_size_h = tiling_json["tile_size_h"]; } if (tiling_json.contains("target_overlap") && tiling_json["target_overlap"].is_number()) { vae_tiling_params.target_overlap = tiling_json["target_overlap"]; } - if (tiling_json.contains("rel_size_x") && tiling_json["rel_size_x"].is_number()) { - vae_tiling_params.rel_size_x = tiling_json["rel_size_x"]; + if (tiling_json.contains("rel_size_w") && tiling_json["rel_size_w"].is_number()) { + vae_tiling_params.rel_size_w = tiling_json["rel_size_w"]; } - if (tiling_json.contains("rel_size_y") && tiling_json["rel_size_y"].is_number()) { - vae_tiling_params.rel_size_y = tiling_json["rel_size_y"]; + if (tiling_json.contains("rel_size_h") && tiling_json["rel_size_h"].is_number()) { + vae_tiling_params.rel_size_h = tiling_json["rel_size_h"]; } if (tiling_json.contains("extra_tiling_args") && tiling_json["extra_tiling_args"].is_string()) { extra_tiling_args = tiling_json["extra_tiling_args"].get(); @@ -2934,11 +2934,11 @@ std::string SDGenerationParams::to_string() const { << " vae_tiling_params: { " << vae_tiling_params.enabled << ", " << vae_tiling_params.temporal_tiling << ", " - << vae_tiling_params.tile_size_x << ", " - << vae_tiling_params.tile_size_y << ", " + << vae_tiling_params.tile_size_w << ", " + << vae_tiling_params.tile_size_h << ", " << vae_tiling_params.target_overlap << ", " - << vae_tiling_params.rel_size_x << ", " - << vae_tiling_params.rel_size_y << ", " + << vae_tiling_params.rel_size_w << ", " + << vae_tiling_params.rel_size_h << ", " << "\"" << extra_tiling_args << "\" },\n" << "}"; return oss.str(); @@ -3140,11 +3140,11 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params, root["vae_tiling"] = { {"enabled", gen_params.vae_tiling_params.enabled}, {"temporal_tiling", gen_params.vae_tiling_params.temporal_tiling}, - {"tile_size_x", gen_params.vae_tiling_params.tile_size_x}, - {"tile_size_y", gen_params.vae_tiling_params.tile_size_y}, + {"tile_size_w", gen_params.vae_tiling_params.tile_size_w}, + {"tile_size_h", gen_params.vae_tiling_params.tile_size_h}, {"target_overlap", gen_params.vae_tiling_params.target_overlap}, - {"rel_size_x", gen_params.vae_tiling_params.rel_size_x}, - {"rel_size_y", gen_params.vae_tiling_params.rel_size_y}, + {"rel_size_w", gen_params.vae_tiling_params.rel_size_w}, + {"rel_size_h", gen_params.vae_tiling_params.rel_size_h}, {"extra_tiling_args", gen_params.extra_tiling_args}, }; } diff --git a/examples/server/api.md b/examples/server/api.md index 9bcdc9f0..541e260a 100644 --- a/examples/server/api.md +++ b/examples/server/api.md @@ -524,11 +524,11 @@ Shared default fields used by both `img_gen` and `vid_gen`: | `vae_tiling_params` | `object` | | `vae_tiling_params.enabled` | `boolean` | | `vae_tiling_params.temporal_tiling` | `boolean` | -| `vae_tiling_params.tile_size_x` | `integer` | -| `vae_tiling_params.tile_size_y` | `integer` | +| `vae_tiling_params.tile_size_w` | `integer` | +| `vae_tiling_params.tile_size_h` | `integer` | | `vae_tiling_params.target_overlap` | `number` | -| `vae_tiling_params.rel_size_x` | `number` | -| `vae_tiling_params.rel_size_y` | `number` | +| `vae_tiling_params.rel_size_w` | `number` | +| `vae_tiling_params.rel_size_h` | `number` | | `vae_tiling_params.extra_tiling_args` | `string` | | `cache_mode` | `string` | | `cache_option` | `string` | @@ -537,6 +537,8 @@ Shared default fields used by both `img_gen` and `vid_gen`: | `output_format` | `string` | | `output_compression` | `integer` | +`vae_tiling_params.tile_size_w` and `tile_size_h` are in **image pixels**, with `0` selecting the 256-pixel default. Both encode and decode use these sizes without an encoding multiplier. Positive `rel_size_w`/`rel_size_h` values override the corresponding absolute size: values up to 1 are dimension fractions, and values greater than 1 are target tile counts. Set `enabled` to use spatial tiling. Sizes are aligned down to the VAE scale factor and capped at the input dimensions; explicit sizes below the minimum supported tile size are rejected. These fields previously used latent units; see [VAE tiling](../../docs/performance.md#use-vae-tiling-to-reduce-encode-and-decode-memory-usage) for migration and OOM retry behavior. + `vae_tiling_params.extra_tiling_args` accepts a key=value list. Supported video VAEs accept `temporal_tile_frames` (alias `temporal_tile_size`, default `4`) and `temporal_tile_overlap` (default `1`). LTX and Wan preserve causal state between temporal tiles. Hunyuan Video and TAEHV use overlap blending. MiniMax H3 keeps its model-specific fixed temporal windows because its latent-to-frame mapping is non-linear. @@ -767,11 +769,11 @@ Example: "vae_tiling_params": { "enabled": false, "temporal_tiling": false, - "tile_size_x": 0, - "tile_size_y": 0, + "tile_size_w": 0, + "tile_size_h": 0, "target_overlap": 0.5, - "rel_size_x": 0.0, - "rel_size_y": 0.0, + "rel_size_w": 0.0, + "rel_size_h": 0.0, "extra_tiling_args": "" }, @@ -900,11 +902,11 @@ Other native fields: | `vae_tiling_params` | `object` | | `vae_tiling_params.enabled` | `boolean` | | `vae_tiling_params.temporal_tiling` | `boolean` | -| `vae_tiling_params.tile_size_x` | `integer` | -| `vae_tiling_params.tile_size_y` | `integer` | +| `vae_tiling_params.tile_size_w` | `integer` | +| `vae_tiling_params.tile_size_h` | `integer` | | `vae_tiling_params.target_overlap` | `number` | -| `vae_tiling_params.rel_size_x` | `number` | -| `vae_tiling_params.rel_size_y` | `number` | +| `vae_tiling_params.rel_size_w` | `number` | +| `vae_tiling_params.rel_size_h` | `number` | | `vae_tiling_params.extra_tiling_args` | `string` | | `cache_mode` | `string` | | `cache_option` | `string` | @@ -1115,11 +1117,11 @@ Example: "vae_tiling_params": { "enabled": false, "temporal_tiling": false, - "tile_size_x": 0, - "tile_size_y": 0, + "tile_size_w": 0, + "tile_size_h": 0, "target_overlap": 0.5, - "rel_size_x": 0.0, - "rel_size_y": 0.0, + "rel_size_w": 0.0, + "rel_size_h": 0.0, "extra_tiling_args": "" }, @@ -1240,11 +1242,11 @@ Other native fields: | `vae_tiling_params` | `object` | | `vae_tiling_params.enabled` | `boolean` | | `vae_tiling_params.temporal_tiling` | `boolean` | -| `vae_tiling_params.tile_size_x` | `integer` | -| `vae_tiling_params.tile_size_y` | `integer` | +| `vae_tiling_params.tile_size_w` | `integer` | +| `vae_tiling_params.tile_size_h` | `integer` | | `vae_tiling_params.target_overlap` | `number` | -| `vae_tiling_params.rel_size_x` | `number` | -| `vae_tiling_params.rel_size_y` | `number` | +| `vae_tiling_params.rel_size_w` | `number` | +| `vae_tiling_params.rel_size_h` | `number` | | `vae_tiling_params.extra_tiling_args` | `string` | | `cache_mode` | `string` | | `cache_option` | `string` | diff --git a/examples/server/frontend b/examples/server/frontend index c4bce3d6..dd74a8e8 160000 --- a/examples/server/frontend +++ b/examples/server/frontend @@ -1 +1 @@ -Subproject commit c4bce3d6b3f236614cca21014f076083b7270ba8 +Subproject commit dd74a8e808aaa8b26124217424b23058935184de diff --git a/examples/server/routes_sdcpp.cpp b/examples/server/routes_sdcpp.cpp index 60e68ca2..b8767b62 100644 --- a/examples/server/routes_sdcpp.cpp +++ b/examples/server/routes_sdcpp.cpp @@ -78,11 +78,11 @@ static json make_vae_tiling_json(const sd_tiling_params_t& params) { return { {"enabled", params.enabled}, {"temporal_tiling", params.temporal_tiling}, - {"tile_size_x", params.tile_size_x}, - {"tile_size_y", params.tile_size_y}, + {"tile_size_w", params.tile_size_w}, + {"tile_size_h", params.tile_size_h}, {"target_overlap", params.target_overlap}, - {"rel_size_x", params.rel_size_x}, - {"rel_size_y", params.rel_size_y}, + {"rel_size_w", params.rel_size_w}, + {"rel_size_h", params.rel_size_h}, {"extra_tiling_args", params.extra_tiling_args ? params.extra_tiling_args : ""}, }; } diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 616bc377..8e5577d5 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -173,11 +173,13 @@ enum lora_apply_mode_t { typedef struct { bool enabled; bool temporal_tiling; - int tile_size_x; - int tile_size_y; + // Spatial tile dimensions in image pixels for both encode and decode; 0 uses 256. + int tile_size_w; + int tile_size_h; float target_overlap; - float rel_size_x; - float rel_size_y; + // Positive values override tile_size: <= 1 is a dimension fraction, > 1 a target tile count. + float rel_size_w; + float rel_size_h; const char* extra_tiling_args; } sd_tiling_params_t; diff --git a/src/core/backend_fit.cpp b/src/core/backend_fit.cpp index 2ffc1cbb..6a7ab433 100644 --- a/src/core/backend_fit.cpp +++ b/src/core/backend_fit.cpp @@ -478,7 +478,12 @@ namespace sd::backend_fit { return true; } - bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling, ggml_status status) { + bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, + bool prefer_temporal_tiling, + ggml_status status, + int latent_tile_size_w, + int latent_tile_size_h, + int scale_factor) { // Execution failures can leave the device unusable; tiling only helps with allocation failures. if (status != GGML_STATUS_ALLOC_FAILED) { return false; @@ -487,19 +492,31 @@ namespace sd::backend_fit { 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; - 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; - } - retry_mode = tiling_params.temporal_tiling ? "spatial+temporal" : "spatial"; } else { - return false; + if (latent_tile_size_w <= 0 || latent_tile_size_h <= 0 || scale_factor <= 0) { + return false; + } + auto smaller_tile = [&](int size) { + int next_size = size / 2; + if (!tiling_params.enabled) { + next_size = std::min(next_size, 256 / scale_factor); + } + return std::min(size, std::max(4, next_size)); + }; + const int tile_size_w = smaller_tile(latent_tile_size_w); + const int tile_size_h = smaller_tile(latent_tile_size_h); + if (tile_size_w == latent_tile_size_w && tile_size_h == latent_tile_size_h) { + return false; + } + tiling_params.enabled = true; + tiling_params.rel_size_w = 0.0f; + tiling_params.rel_size_h = 0.0f; + tiling_params.tile_size_w = tile_size_w * scale_factor; + tiling_params.tile_size_h = tile_size_h * scale_factor; + retry_mode = tiling_params.temporal_tiling ? "spatial+temporal" : "spatial"; + LOG_WARN("Reducing VAE decode tiles from %dx%d to %dx%d image pixels", + latent_tile_size_w * scale_factor, latent_tile_size_h * scale_factor, + tiling_params.tile_size_w, tiling_params.tile_size_h); } LOG_WARN("VAE decode ran out of memory; retrying with %s tiling", diff --git a/src/core/backend_fit.h b/src/core/backend_fit.h index 0c682a40..ad1b60c6 100644 --- a/src/core/backend_fit.h +++ b/src/core/backend_fit.h @@ -17,7 +17,10 @@ namespace sd::backend_fit { bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling, - ggml_status status); + ggml_status status, + int latent_tile_size_w, + int latent_tile_size_h, + int scale_factor); } // namespace sd::backend_fit diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index 77756643..f6c8eb49 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -556,12 +556,18 @@ namespace MiniMaxH3VAE { tensor.shape()[3]}); } - static sd_tiling_params_t h3_tiling(sd_tiling_params_t params) { + sd_tiling_params_t resolve_tiling_params(sd_tiling_params_t params) const override { + if (!params.enabled) { + params.target_overlap = 0.25f; + } + if (params.tile_size_w == 0 && params.rel_size_w == 0.f) { + params.tile_size_w = 256; + } + if (params.tile_size_h == 0 && params.rel_size_h == 0.f) { + params.tile_size_h = 256; + } params.enabled = true; params.temporal_tiling = false; - params.tile_size_x = 16; - params.tile_size_y = 16; - params.target_overlap = 0.25f; return params; } @@ -605,7 +611,7 @@ namespace MiniMaxH3VAE { bool circular_x = false, bool circular_y = false) override { auto input = ensure_video_shape(x); - auto tiling = h3_tiling(tiling_params); + auto tiling = resolve_tiling_params(tiling_params); if (input.shape()[2] == 1) { auto encoded = VAE::encode(n_threads, input, tiling, circular_x, circular_y); if (!encoded.empty() && encoded.shape()[2] > 1) { @@ -646,7 +652,7 @@ namespace MiniMaxH3VAE { bool circular_y = false, bool silent = false) override { auto input = ensure_video_shape(x); - auto tiling = h3_tiling(tiling_params); + auto tiling = resolve_tiling_params(tiling_params); if (input.shape()[2] == 1) { auto decoded = VAE::decode(n_threads, input, diff --git a/src/model/vae/vae.hpp b/src/model/vae/vae.hpp index 4d8d1351..157b5164 100644 --- a/src/model/vae/vae.hpp +++ b/src/model/vae/vae.hpp @@ -1,6 +1,9 @@ #ifndef __SD_MODEL_VAE_VAE_HPP__ #define __SD_MODEL_VAE_VAE_HPP__ +#include +#include + #include "core/tensor_ggml.hpp" #include "model/common/block.hpp" #include "model/vae/vae_tiling.hpp" @@ -117,8 +120,8 @@ protected: int output_width, int output_height, int scale, - int p_tile_size_x, - int p_tile_size_y, + int p_tile_size_w, + int p_tile_size_h, float tile_overlap_factor, bool circular_x, bool circular_y, @@ -138,17 +141,28 @@ protected: } return output_tile; }; - return ::process_tiles_2d(input, - output_width, - output_height, - scale, - p_tile_size_x, - p_tile_size_y, - tile_overlap_factor, - circular_x, - circular_y, - on_processing, - silent); + const bool original_circular_x = circular_x_enabled; + const bool original_circular_y = circular_y_enabled; + const int64_t latent_width = decode_graph ? input.shape()[0] : output_width; + const int64_t latent_height = decode_graph ? input.shape()[1] : output_height; + circular_x = circular_x || original_circular_x; + circular_y = circular_y || original_circular_y; + // Full-width axes wrap in convolutions; split axes wrap between tiles. + set_circular_axes(circular_x && p_tile_size_w >= latent_width, + circular_y && p_tile_size_h >= latent_height); + auto output = ::process_tiles_2d(input, + output_width, + output_height, + scale, + p_tile_size_w, + p_tile_size_h, + tile_overlap_factor, + circular_x && p_tile_size_w < latent_width, + circular_y && p_tile_size_h < latent_height, + on_processing, + silent); + set_circular_axes(original_circular_x, original_circular_y); + return output; } public: @@ -178,33 +192,48 @@ public: return supports_temporal_tiling(VAETemporalDirection::DECODE); } - void get_tile_sizes(int& tile_size_x, - int& tile_size_y, + virtual sd_tiling_params_t resolve_tiling_params(sd_tiling_params_t params) const { + return params; + } + + bool get_tile_sizes(int& tile_size_w, + int& tile_size_h, float& tile_overlap, const sd_tiling_params_t& params, - int64_t latent_x, - int64_t latent_y, - float encoding_factor = 1.0f) { - 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; - const int min_tile_dimension = 4; - int tile_size = default_tile_size; - // factor <= 1 means simple fraction of the latent dimension - // factor > 1 means number of tiles across that dimension - if (factor > 0.f) { - if (factor > 1.0) - factor = 1 / (factor - factor * tile_overlap + tile_overlap); - tile_size = static_cast(std::round(latent_size * factor)); - } else if (requested_size >= min_tile_dimension) { - tile_size = requested_size; + int64_t latent_w, + int64_t latent_h) { + const auto tiling = resolve_tiling_params(params); + if (latent_w <= 0 || latent_h <= 0 || + latent_w > std::numeric_limits::max() || latent_h > std::numeric_limits::max() || + !std::isfinite(tiling.target_overlap)) { + LOG_ERROR("invalid VAE tiling dimensions or overlap"); + return false; + } + const int scale_factor = get_scale_factor(); + tile_overlap = std::max(std::min(tiling.target_overlap, 0.5f), 0.0f); + auto get_tile_size = [&](int requested_size, double factor, int64_t latent_size, int& tile_size) { + if (requested_size < 0 || !std::isfinite(factor) || factor < 0.0) { + LOG_ERROR("VAE tile sizes and relative sizes must be finite and non-negative"); + return false; } - tile_size = static_cast(tile_size * encoding_factor); - return std::max(std::min(tile_size, static_cast(latent_size)), min_tile_dimension); + const int min_tile_dimension = std::min(4, static_cast(latent_size)); + double size = (requested_size > 0 ? requested_size : 256) / scale_factor; + if (factor > 0.0) { + if (factor > 1.0) { + factor = 1.0 / (factor * (1.0 - tile_overlap) + tile_overlap); + } + size = std::floor(static_cast(latent_size) * factor); + } + if (size < min_tile_dimension && (requested_size > 0 || factor > 0.0)) { + LOG_ERROR("VAE tile size must be at least %d image pixels on this axis", min_tile_dimension * scale_factor); + return false; + } + tile_size = static_cast(std::min(static_cast(latent_size), std::max(min_tile_dimension, size))); + return true; }; - 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); + return get_tile_size(tiling.tile_size_w, tiling.rel_size_w, latent_w, tile_size_w) && + get_tile_size(tiling.tile_size_h, tiling.rel_size_h, latent_h, tile_size_h); } virtual sd::Tensor encode(int n_threads, @@ -213,6 +242,7 @@ public: bool circular_x = false, bool circular_y = false) { int64_t t0 = ggml_time_ms(); + tiling_params = resolve_tiling_params(tiling_params); sd::Tensor input = x; sd::Tensor output; if (scale_input) { @@ -224,21 +254,19 @@ public: int64_t W = input.shape()[0] / scale_factor; int64_t H = input.shape()[1] / scale_factor; float tile_overlap; - int tile_size_x, tile_size_y; - // Image VAE encode is more sensitive to tile boundary context than decode. - // Keep the smaller legacy factor for video VAEs, but default image encode - // 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); - LOG_VERBOSE("VAE Tile size: %dx%d", tile_size_x, tile_size_y); + int tile_size_w, tile_size_h; + if (!get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling_params, W, H)) { + return {}; + } + LOG_VERBOSE("VAE encode tile size: %dx%d pixels (%dx%d latent)", + tile_size_w * scale_factor, tile_size_h * scale_factor, tile_size_w, tile_size_h); output = tiled_compute(input, n_threads, static_cast(W), static_cast(H), scale_factor, - tile_size_x, - tile_size_y, + tile_size_w, + tile_size_h, tile_overlap, circular_x, circular_y, @@ -271,6 +299,7 @@ public: bool circular_y = false, bool silent = false) { int64_t t0 = ggml_time_ms(); + tiling_params = resolve_tiling_params(tiling_params); sd::Tensor input = x; sd::Tensor output; @@ -279,10 +308,13 @@ public: int64_t W = input.shape()[0] * scale_factor; 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]); + int tile_size_w, tile_size_h; + if (!get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling_params, input.shape()[0], input.shape()[1])) { + return {}; + } if (!silent) { - LOG_VERBOSE("VAE Tile size: %dx%d", tile_size_x, tile_size_y); + LOG_VERBOSE("VAE decode tile size: %dx%d pixels (%dx%d latent)", + tile_size_w * scale_factor, tile_size_h * scale_factor, tile_size_w, tile_size_h); } output = tiled_compute( input, @@ -290,8 +322,8 @@ public: static_cast(W), static_cast(H), scale_factor, - tile_size_x, - tile_size_y, + tile_size_w, + tile_size_h, tile_overlap, circular_x, circular_y, diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index 3eba5d71..aacbec99 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -2883,14 +2883,26 @@ sd::Tensor StableDiffusionGGML::decode_first_stage(const sd::Tensordiffusion_to_vae_latents(x); - auto decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); - const bool prefer_temporal_tiling = decode_video && first_stage_model->can_temporal_tile_decode(); - while (decoded.empty() && - sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling, - first_stage_model->last_compute_status())) { - decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); + 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(); + for (;;) { + int tile_size_w = static_cast(latents.shape()[0]); + int tile_size_h = static_cast(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])) { + return {}; + } + auto decoded = first_stage_model->decode(n_threads, latents, tiling_params, decode_video, circular_x, circular_y); + if (!decoded.empty() || + !sd::backend_fit::prepare_vae_decode_retry_tiling(tiling_params, prefer_temporal_tiling, + first_stage_model->last_compute_status(), + tile_size_w, tile_size_h, first_stage_model->get_scale_factor())) { + return decoded; + } } - return decoded; } sd::Tensor StableDiffusionGGML::normalize_ltx_video_latents(const sd::Tensor& x) { diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp index 38b8f086..951e8dc3 100644 --- a/src/pipeline/image.cpp +++ b/src/pipeline/image.cpp @@ -35,19 +35,21 @@ namespace sd::pipeline { return original_axes; } - int tile_size_x, tile_size_y; + int tile_size_w, tile_size_h; float overlap; - int latent_size_x = request.width / request.vae_scale_factor; - int latent_size_y = request.height / request.vae_scale_factor; - sd->first_stage_model->get_tile_sizes(tile_size_x, - tile_size_y, - overlap, - sd_img_gen_params->vae_tiling_params, - latent_size_x, - latent_size_y); + int latent_size_w = request.width / request.vae_scale_factor; + int latent_size_h = request.height / request.vae_scale_factor; + if (!sd->first_stage_model->get_tile_sizes(tile_size_w, + tile_size_h, + overlap, + sd_img_gen_params->vae_tiling_params, + latent_size_w, + latent_size_h)) { + return original_axes; + } - sd->circular_x = sd->circular_x && (tile_size_x >= latent_size_x); - sd->circular_y = sd->circular_y && (tile_size_y >= latent_size_y); + sd->circular_x = sd->circular_x && (tile_size_w >= latent_size_w); + sd->circular_y = sd->circular_y && (tile_size_h >= latent_size_h); if (sd->first_stage_model) { sd->first_stage_model->set_circular_axes(sd->circular_x, sd->circular_y); @@ -56,8 +58,8 @@ namespace sd::pipeline { sd->preview_vae->set_circular_axes(sd->circular_x, sd->circular_y); } - sd->circular_x = original_axes.circular_x && (tile_size_x < latent_size_x); - sd->circular_y = original_axes.circular_y && (tile_size_y < latent_size_y); + sd->circular_x = original_axes.circular_x && (tile_size_w < latent_size_w); + sd->circular_y = original_axes.circular_y && (tile_size_h < latent_size_h); return original_axes; } diff --git a/src/runtime/tiling.cpp b/src/runtime/tiling.cpp index f0a9dbc7..38ea4c24 100644 --- a/src/runtime/tiling.cpp +++ b/src/runtime/tiling.cpp @@ -142,8 +142,8 @@ sd::Tensor process_tiles_2d(const sd::Tensor& input, int output_width, int output_height, int scale, - int p_tile_size_x, - int p_tile_size_y, + int p_tile_size_w, + int p_tile_size_h, float tile_overlap_factor, bool circular_x, bool circular_y, @@ -168,28 +168,28 @@ sd::Tensor process_tiles_2d(const sd::Tensor& input, int num_tiles_x; float tile_overlap_factor_x; - sd_tiling_calc_tiles(num_tiles_x, tile_overlap_factor_x, small_width, p_tile_size_x, tile_overlap_factor, circular_x); + sd_tiling_calc_tiles(num_tiles_x, tile_overlap_factor_x, small_width, p_tile_size_w, tile_overlap_factor, circular_x); int num_tiles_y; float tile_overlap_factor_y; - sd_tiling_calc_tiles(num_tiles_y, tile_overlap_factor_y, small_height, p_tile_size_y, tile_overlap_factor, circular_y); + sd_tiling_calc_tiles(num_tiles_y, tile_overlap_factor_y, small_height, p_tile_size_h, tile_overlap_factor, circular_y); - int tile_overlap_x = static_cast(p_tile_size_x * tile_overlap_factor_x); - int non_tile_overlap_x = p_tile_size_x - tile_overlap_x; - int tile_overlap_y = static_cast(p_tile_size_y * tile_overlap_factor_y); - int non_tile_overlap_y = p_tile_size_y - tile_overlap_y; - int tile_size_x = p_tile_size_x < small_width ? p_tile_size_x : small_width; - int tile_size_y = p_tile_size_y < small_height ? p_tile_size_y : small_height; - int input_tile_size_x = tile_size_x; - int input_tile_size_y = tile_size_y; - int output_tile_size_x = tile_size_x; - int output_tile_size_y = tile_size_y; + int tile_overlap_x = static_cast(p_tile_size_w * tile_overlap_factor_x); + int non_tile_overlap_x = p_tile_size_w - tile_overlap_x; + int tile_overlap_y = static_cast(p_tile_size_h * tile_overlap_factor_y); + int non_tile_overlap_y = p_tile_size_h - tile_overlap_y; + int tile_size_w = p_tile_size_w < small_width ? p_tile_size_w : small_width; + int tile_size_h = p_tile_size_h < small_height ? p_tile_size_h : small_height; + int input_tile_size_w = tile_size_w; + int input_tile_size_h = tile_size_h; + int output_tile_size_w = tile_size_w; + int output_tile_size_h = tile_size_h; if (decode) { - output_tile_size_x *= scale; - output_tile_size_y *= scale; + output_tile_size_w *= scale; + output_tile_size_h *= scale; } else { - input_tile_size_x *= scale; - input_tile_size_y *= scale; + input_tile_size_w *= scale; + input_tile_size_h *= scale; } int num_tiles = num_tiles_x * num_tiles_y; @@ -205,9 +205,9 @@ sd::Tensor process_tiles_2d(const sd::Tensor& input, } for (int y = 0; y < small_height && !last_y; y += non_tile_overlap_y) { int dy = 0; - if (!circular_y && y + tile_size_y >= small_height) { + if (!circular_y && y + tile_size_h >= small_height) { int original_y = y; - y = small_height - tile_size_y; + y = small_height - tile_size_h; dy = original_y - y; if (decode) { dy *= scale; @@ -216,9 +216,9 @@ sd::Tensor process_tiles_2d(const sd::Tensor& input, } for (int x = 0; x < small_width && !last_x; x += non_tile_overlap_x) { int dx = 0; - if (!circular_x && x + tile_size_x >= small_width) { + if (!circular_x && x + tile_size_w >= small_width) { int original_x = x; - x = small_width - tile_size_x; + x = small_width - tile_size_w; dx = original_x - x; if (decode) { dx *= scale; @@ -235,12 +235,12 @@ sd::Tensor process_tiles_2d(const sd::Tensor& input, int overlap_y_out = decode ? tile_overlap_y * scale : tile_overlap_y; int64_t t1 = ggml_time_ms(); - auto input_tile = sd_tensor_split_2d(input, input_tile_size_x, input_tile_size_y, x_in, y_in); + auto input_tile = sd_tensor_split_2d(input, input_tile_size_w, input_tile_size_h, x_in, y_in); auto output_tile = on_processing(input_tile); if (output_tile.empty()) { return {}; } - GGML_ASSERT(output_tile.shape()[0] == output_tile_size_x && output_tile.shape()[1] == output_tile_size_y); + GGML_ASSERT(output_tile.shape()[0] == output_tile_size_w && output_tile.shape()[1] == output_tile_size_h); if (output.empty()) { std::vector output_shape = output_tile.shape(); output_shape[0] = output_width; diff --git a/src/runtime/tiling.h b/src/runtime/tiling.h index a832a644..b203c542 100644 --- a/src/runtime/tiling.h +++ b/src/runtime/tiling.h @@ -11,8 +11,8 @@ sd::Tensor process_tiles_2d(const sd::Tensor& input, int output_width, int output_height, int scale, - int p_tile_size_x, - int p_tile_size_y, + int p_tile_size_w, + int p_tile_size_h, float tile_overlap_factor, bool circular_x, bool circular_y,