diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 35812157c..d46be91a4 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -1013,7 +1013,7 @@ ArgOptions SDGenerationParams::get_options() { &extra_sample_args}, {"", "--extra-tiling-args", - "extra VAE tiling args, key=value list. LTX video VAE supports temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)", + "extra VAE tiling args, key=value list. Supported video VAEs accept temporal_tile_frames/temporal_tile_size (default: 4), temporal_tile_overlap (default: 1)", (int)',', &extra_tiling_args}, {"", @@ -1230,7 +1230,7 @@ ArgOptions SDGenerationParams::get_options() { &vae_tiling_params.enabled}, {"", "--temporal-tiling", - "enable temporal tiling for LTX video VAE decode", + "enable temporal tiling for supported video VAE decode", true, &vae_tiling_params.temporal_tiling}, {"", diff --git a/examples/server/api.md b/examples/server/api.md index 29d65877e..3ec9a1da8 100644 --- a/examples/server/api.md +++ b/examples/server/api.md @@ -518,7 +518,8 @@ Shared default fields used by both `img_gen` and `vid_gen`: | `output_format` | `string` | | `output_compression` | `integer` | -`vae_tiling_params.extra_tiling_args` accepts a key=value list. For LTX video VAE temporal tiling, `temporal_tile_frames` defaults to `4` and `temporal_tile_overlap` defaults to `1`. +`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. `img_gen`-specific default fields: diff --git a/src/core/backend_fit.cpp b/src/core/backend_fit.cpp index 460194740..0ba2df14f 100644 --- a/src/core/backend_fit.cpp +++ b/src/core/backend_fit.cpp @@ -364,15 +364,11 @@ namespace sd::backend_fit { } bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) { - if (prefer_temporal_tiling) { - if (tiling_params.temporal_tiling) { - return false; - } + const char* retry_mode = nullptr; + if (prefer_temporal_tiling && !tiling_params.temporal_tiling) { tiling_params.temporal_tiling = true; - } else { - if (tiling_params.enabled) { - return false; - } + retry_mode = tiling_params.enabled ? "spatial+temporal" : "temporal"; + } else if (!tiling_params.enabled) { tiling_params.enabled = true; if (tiling_params.tile_size_x <= 0) { tiling_params.tile_size_x = 256; @@ -380,10 +376,13 @@ namespace sd::backend_fit { 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; } LOG_WARN("auto-fit: VAE decode failed (likely out of memory); retrying with %s tiling", - tiling_params.temporal_tiling ? "temporal" : "spatial"); + retry_mode); return true; } diff --git a/src/model/vae/hunyuan_vae.hpp b/src/model/vae/hunyuan_vae.hpp index b16126181..938f5ffe8 100644 --- a/src/model/vae/hunyuan_vae.hpp +++ b/src/model/vae/hunyuan_vae.hpp @@ -758,6 +758,15 @@ namespace Hunyuan { return "hunyuan_video_vae"; } + bool supports_temporal_tiling(VAETemporalDirection direction) const override { + return direction == VAETemporalDirection::DECODE; + } + + int get_temporal_tile_output_scale(VAETemporalDirection direction) const override { + SD_UNUSED(direction); + return 4; + } + void get_param_tensors(std::map& tensors) override { if (!decode_only) { encoder.get_param_tensors(tensors, weight_prefix + ".encoder"); diff --git a/src/model/vae/ltx_vae.hpp b/src/model/vae/ltx_vae.hpp index a19ce820d..5629b939b 100644 --- a/src/model/vae/ltx_vae.hpp +++ b/src/model/vae/ltx_vae.hpp @@ -1213,9 +1213,6 @@ struct LTXVideoVAE : public VAE { static constexpr int DEFAULT_TEMPORAL_TILE_OVERLAP = 1; bool decode_only; - bool temporal_tiling_enabled = false; - int temporal_tile_frames = DEFAULT_TEMPORAL_TILE_FRAMES; - int temporal_tile_overlap = DEFAULT_TEMPORAL_TILE_OVERLAP; int ltx_vae_version; bool timestep_conditioning; int patch_size; @@ -1248,62 +1245,22 @@ struct LTXVideoVAE : public VAE { return "ltx_video_vae"; } - void set_temporal_tiling_enabled(bool enabled) override { - temporal_tiling_enabled = enabled; + bool supports_temporal_tiling(VAETemporalDirection direction) const override { + return direction == VAETemporalDirection::DECODE; } - void set_tiling_params(const sd_tiling_params_t& params) override { - temporal_tiling_enabled = params.temporal_tiling; - temporal_tile_frames = DEFAULT_TEMPORAL_TILE_FRAMES; - temporal_tile_overlap = DEFAULT_TEMPORAL_TILE_OVERLAP; - - for (const auto& [key, value] : parse_key_value_args(params.extra_tiling_args, "LTX VAE extra tiling arg")) { - int parsed = 0; - if (!parse_strict_int(value, parsed)) { - LOG_WARN("ignoring invalid LTX VAE extra tiling arg '%s=%s'", key.c_str(), value.c_str()); - } else if (key == "temporal_tile_frames") { - temporal_tile_frames = std::max(1, parsed); - } else if (key == "temporal_tile_overlap") { - temporal_tile_overlap = std::max(0, parsed); - } else { - LOG_WARN("ignoring unknown LTX VAE extra tiling arg '%s'", key.c_str()); - } - } + int get_default_temporal_tile_frames(VAETemporalDirection direction) const override { + SD_UNUSED(direction); + return DEFAULT_TEMPORAL_TILE_FRAMES; } - void get_param_tensors(std::map& tensors) override { - vae.get_param_tensors(tensors, weight_prefix); + int get_default_temporal_tile_overlap(VAETemporalDirection direction) const override { + SD_UNUSED(direction); + return DEFAULT_TEMPORAL_TILE_OVERLAP; } - struct TemporalTilePlan { - int frames = 1; - int overlap = 0; - int stride = 1; - int num_tiles = 1; - }; - - TemporalTilePlan resolve_temporal_tile_plan(int64_t total_frames) const { - TemporalTilePlan plan; - plan.frames = std::max(1, temporal_tile_frames); - plan.overlap = std::max(0, temporal_tile_overlap); - - if (plan.overlap >= plan.frames) { - LOG_WARN("temporal_tile_overlap (%d) is greater than or equal to temporal_tile_frames (%d), adjusting values to avoid empty decode windows", - plan.overlap, - plan.frames); - plan.overlap = plan.frames - 1; - } - if (total_frames > 1 && plan.overlap >= total_frames) { - LOG_WARN("temporal_tile_overlap (%d) is greater than or equal to total latent frames (%lld), adjusting values to decode at least one tile", - plan.overlap, - (long long)total_frames); - plan.overlap = static_cast(total_frames - 1); - } - - plan.stride = std::max(1, plan.frames - plan.overlap); - int64_t tiled_frames = std::max(1, total_frames - plan.overlap); - plan.num_tiles = total_frames > 0 ? static_cast((tiled_frames + plan.stride - 1) / plan.stride) : 0; - return plan; + void get_param_tensors(std::map& tensors) override { + vae.get_param_tensors(tensors, weight_prefix); } std::string temporal_feat_cache_name(size_t feat_idx) const { @@ -1365,52 +1322,53 @@ struct LTXVideoVAE : public VAE { sd::Tensor decode_temporal_tiled_streaming(const int n_threads, const sd::Tensor& input, - size_t expected_dim) { + size_t expected_dim, + const VAETemporalTilingConfig& config) { const int64_t total_frames = input.shape()[2]; - TemporalTilePlan plan = resolve_temporal_tile_plan(total_frames); + auto plan = make_vae_temporal_tile_plan(total_frames, config); LOG_DEBUG("Using streaming temporal tiling: temporal_tile_frames=%d, temporal_tile_overlap=%d, total latent frames=%lld, resulting in %d tiles", - plan.frames, + plan.tile_frames, plan.overlap, (long long)total_frames, - plan.num_tiles); + (int)plan.tiles.size()); free_cache_ctx_and_buffer(); cache_tensor_map.clear(); - sd::Tensor output; - for (int64_t start = 0; start < total_frames - plan.overlap; start += plan.stride) { - const int64_t end = std::min(total_frames, start + plan.frames); - const int chunk_overlap = end < total_frames ? plan.overlap : 0; - auto z_chunk = sd::ops::slice(input, 2, start, end); - + auto output = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor& z_chunk, const VAETemporalTile& tile) { LOG_DEBUG("LTX VAE temporal tile %lld/%d: latent frames [%lld, %lld), overlap=%d", - (long long)(start / plan.stride + 1), - plan.num_tiles, - (long long)start, - (long long)end, - chunk_overlap); + (long long)tile.index + 1, + (int)plan.tiles.size(), + (long long)tile.start, + (long long)tile.end, + tile.overlap); auto get_graph = [&]() -> ggml_cgraph* { return build_temporal_tile_graph(z_chunk, - static_cast(start), - chunk_overlap); + static_cast(tile.start), + tile.overlap); }; - auto chunk = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true, true, true), - expected_dim); - if (chunk.empty()) { - free_cache_ctx_and_buffer(); - cache_tensor_map.clear(); - return {}; - } - output = output.empty() ? std::move(chunk) : sd::ops::concat(output, chunk, 2); - } + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true, true, true), + expected_dim); + }); free_cache_ctx_and_buffer(); cache_tensor_map.clear(); return output; } + sd::Tensor _compute_temporal_tiled(const int n_threads, + const sd::Tensor& input, + VAETemporalDirection direction, + const VAETemporalTilingConfig& config) override { + GGML_ASSERT(direction == VAETemporalDirection::DECODE); + return decode_temporal_tiled_streaming(n_threads, + input, + static_cast(input.dim()), + config); + } + ggml_cgraph* build_latent_statistics_graph(const sd::Tensor& z_tensor, bool normalize) { ggml_cgraph* gf = new_graph_custom(1024); ggml_tensor* z = make_input(z_tensor); @@ -1446,9 +1404,6 @@ struct LTXVideoVAE : public VAE { input = sd::ops::slice(input, 2, 0, cropped_t); } } - if (decode_graph && temporal_tiling_enabled && input.dim() == 5 && input.shape()[2] > 1) { - return decode_temporal_tiled_streaming(n_threads, input, expected_dim); - } auto get_graph = [&]() -> ggml_cgraph* { return build_graph(input, decode_graph); }; diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index 163a9447a..a8aa6b5f1 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -558,10 +558,11 @@ namespace MiniMaxH3VAE { } static sd_tiling_params_t h3_tiling(sd_tiling_params_t params) { - params.enabled = true; - params.tile_size_x = 16; - params.tile_size_y = 16; - params.target_overlap = 0.25f; + params.enabled = true; + params.temporal_tiling = false; + params.tile_size_x = 16; + params.tile_size_y = 16; + params.target_overlap = 0.25f; return params; } @@ -624,15 +625,13 @@ namespace MiniMaxH3VAE { if (pad > 0) { input = repeat_last_frame(input, pad); } - sd::Tensor result; - for (int64_t start = 0; start < input.shape()[2]; start += 17) { - auto chunk = sd::ops::slice(input, 2, start, start + 17); - auto encoded = VAE::encode(n_threads, chunk, tiling, circular_x, circular_y); - if (encoded.empty()) { - return {}; - } - result = result.empty() ? std::move(encoded) - : sd::ops::concat(result, encoded, 2); + auto plan = make_vae_temporal_tile_plan(input.shape()[2], {17, 0}); + auto result = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor& chunk, const VAETemporalTile& tile) { + SD_UNUSED(tile); + return VAE::encode(n_threads, chunk, tiling, circular_x, circular_y); + }); + if (result.empty()) { + return {}; } if (result.shape()[2] > 3) { result = sd::ops::slice(result, 2, 0, result.shape()[2] - 3); @@ -685,22 +684,21 @@ namespace MiniMaxH3VAE { input = repeat_last_frame(input, pad_tokens); } - sd::Tensor result; sd::Tensor overlap; - for (int64_t i = 0; i < num_chunks; ++i) { - int64_t start = i * tokens_per_chunk; - int64_t end = std::min(start + tokens_per_chunk + token_overlap, - input.shape()[2]); - auto chunk = sd::ops::slice(input, 2, start, end); - auto decoded = VAE::decode(n_threads, - chunk, - tiling, - true, - circular_x, - circular_y, - silent); + auto plan = make_vae_temporal_tile_plan( + input.shape()[2], + {static_cast(tokens_per_chunk + token_overlap), static_cast(token_overlap)}); + GGML_ASSERT(plan.tiles.size() == static_cast(num_chunks)); + auto result = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor& chunk, const VAETemporalTile& tile) { + auto decoded = VAE::decode(n_threads, + chunk, + tiling, + true, + circular_x, + circular_y, + silent); if (decoded.empty()) { - return {}; + return sd::Tensor(); } int64_t first_end = std::min(frames_per_chunk, decoded.shape()[2]); @@ -712,8 +710,6 @@ namespace MiniMaxH3VAE { first = blend_temporal(overlap, first, frame_overlap); overlap = {}; } - result = result.empty() ? std::move(first) - : sd::ops::concat(result, first, 2); if (decoded.shape()[2] > frames_per_chunk + frame_pre_padding) { overlap = sd::ops::slice(decoded, @@ -721,10 +717,14 @@ namespace MiniMaxH3VAE { frames_per_chunk + frame_pre_padding, decoded.shape()[2]); } - if (i == num_chunks - 1 && !overlap.empty()) { - result = sd::ops::concat(result, overlap, 2); + if (tile.last && !overlap.empty()) { + first = sd::ops::concat(first, overlap, 2); overlap = {}; } + return first; + }); + if (result.empty()) { + return {}; } int64_t expected_frames = input.shape()[2] <= 1 ? 1 : ((x.shape()[2] - 2) / 5) * 17 + 5; diff --git a/src/model/vae/tae.hpp b/src/model/vae/tae.hpp index a03e5d21a..d291bb785 100644 --- a/src/model/vae/tae.hpp +++ b/src/model/vae/tae.hpp @@ -819,6 +819,21 @@ struct TinyVideoAutoEncoder : public VAE { return "taehv"; } + bool supports_temporal_tiling(VAETemporalDirection direction) const override { + return direction == VAETemporalDirection::DECODE && !sd_version_is_minimax_h3(version); + } + + int get_temporal_tile_output_scale(VAETemporalDirection direction) const override { + SD_UNUSED(direction); + int scale = 1; + for (bool upscale : taehv.time_upscale) { + if (upscale) { + scale *= 2; + } + } + return scale; + } + void get_param_tensors(std::map& tensors) override { taehv.get_param_tensors(tensors, weight_prefix); } diff --git a/src/model/vae/vae.hpp b/src/model/vae/vae.hpp index 62a6dd066..f3adb0ccc 100644 --- a/src/model/vae/vae.hpp +++ b/src/model/vae/vae.hpp @@ -3,6 +3,7 @@ #include "core/tensor_ggml.hpp" #include "model/common/block.hpp" +#include "model/vae/vae_tiling.hpp" #include "model_manager.h" struct VAE : public GGMLRunner { @@ -14,6 +15,87 @@ struct VAE : public GGMLRunner { const sd::Tensor& z, bool decode_graph) = 0; + virtual bool supports_temporal_tiling(VAETemporalDirection direction) const { + SD_UNUSED(direction); + return false; + } + + virtual int get_default_temporal_tile_frames(VAETemporalDirection direction) const { + SD_UNUSED(direction); + return 4; + } + + virtual int get_default_temporal_tile_overlap(VAETemporalDirection direction) const { + SD_UNUSED(direction); + return 1; + } + + virtual int get_temporal_tile_output_scale(VAETemporalDirection direction) const { + SD_UNUSED(direction); + return 1; + } + + virtual sd::Tensor _compute_temporal_tiled(const int n_threads, + const sd::Tensor& input, + VAETemporalDirection direction, + const VAETemporalTilingConfig& config) { + if (direction != VAETemporalDirection::DECODE) { + return _compute(n_threads, input, false); + } + + VAETemporalTilingConfig resolved_config = config; + const int output_scale = get_temporal_tile_output_scale(direction); + if (output_scale > 1 && + resolved_config.overlap == 0 && + input.shape()[2] > resolved_config.tile_frames) { + LOG_WARN("%s temporal decode requires at least one overlapping latent frame; using overlap=1", + get_desc().c_str()); + resolved_config.overlap = 1; + } + + auto plan = make_vae_temporal_tile_plan(input.shape()[2], resolved_config); + LOG_DEBUG("%s temporal tiling: tile_frames=%d, overlap=%d, total_frames=%lld, tiles=%d", + get_desc().c_str(), + plan.tile_frames, + plan.overlap, + (long long)input.shape()[2], + (int)plan.tiles.size()); + return process_vae_temporal_tiles_blended( + input, + plan, + output_scale, + [&](const sd::Tensor& input_tile, const VAETemporalTile& tile) { + LOG_DEBUG("%s temporal tile %d/%d: input frames [%lld, %lld)", + get_desc().c_str(), + tile.index + 1, + (int)plan.tiles.size(), + (long long)tile.start, + (long long)tile.end); + return _compute(n_threads, input_tile, true); + }); + } + + sd::Tensor compute_with_temporal_tiling(const int n_threads, + const sd::Tensor& input, + VAETemporalDirection direction, + const sd_tiling_params_t& tiling_params) { + if (!tiling_params.temporal_tiling || input.dim() != 5 || input.shape()[2] <= 1) { + return _compute(n_threads, input, direction == VAETemporalDirection::DECODE); + } + if (!supports_temporal_tiling(direction)) { + LOG_WARN("%s does not support temporal tiling for %s; processing the full temporal dimension", + get_desc().c_str(), + direction == VAETemporalDirection::DECODE ? "decode" : "encode"); + return _compute(n_threads, input, direction == VAETemporalDirection::DECODE); + } + + auto config = resolve_vae_temporal_tiling_config( + tiling_params, + get_default_temporal_tile_frames(direction), + get_default_temporal_tile_overlap(direction)); + return _compute_temporal_tiled(n_threads, input, direction, config); + } + static inline void scale_tensor_to_minus1_1(sd::Tensor* tensor) { GGML_ASSERT(tensor != nullptr); for (int64_t i = 0; i < tensor->numel(); ++i) { @@ -40,10 +122,15 @@ struct VAE : public GGMLRunner { bool circular_x, bool circular_y, bool decode_graph, + const sd_tiling_params_t& tiling_params, const char* error_message, bool silent = false) { auto on_processing = [&](const sd::Tensor& input_tile) { - auto output_tile = _compute(n_threads, input_tile, decode_graph); + auto output_tile = compute_with_temporal_tiling( + n_threads, + input_tile, + decode_graph ? VAETemporalDirection::DECODE : VAETemporalDirection::ENCODE, + tiling_params); if (output_tile.empty()) { LOG_ERROR("%s", error_message); return sd::Tensor(); @@ -86,6 +173,10 @@ struct VAE : public GGMLRunner { virtual int get_encoder_output_channels(int input_channels) = 0; + bool can_temporal_tile_decode() const { + return supports_temporal_tiling(VAETemporalDirection::DECODE); + } + void get_tile_sizes(int& tile_size_x, int& tile_size_y, float& tile_overlap, @@ -151,9 +242,13 @@ struct VAE : public GGMLRunner { circular_x, circular_y, false, + tiling_params, "vae encode compute failed while processing a tile"); } else { - output = _compute(n_threads, input, false); + output = compute_with_temporal_tiling(n_threads, + input, + VAETemporalDirection::ENCODE, + tiling_params); } runner_done(); @@ -177,7 +272,6 @@ struct VAE : public GGMLRunner { int64_t t0 = ggml_time_ms(); sd::Tensor input = x; sd::Tensor output; - set_tiling_params(tiling_params); if (tiling_params.enabled) { const int scale_factor = get_scale_factor(); @@ -201,10 +295,14 @@ struct VAE : public GGMLRunner { circular_x, circular_y, true, + tiling_params, "vae decode compute failed while processing a tile", silent); } else { - output = _compute(n_threads, input, true); + output = compute_with_temporal_tiling(n_threads, + input, + VAETemporalDirection::DECODE, + tiling_params); } runner_done(); @@ -226,10 +324,6 @@ struct VAE : public GGMLRunner { virtual sd::Tensor vae_to_diffusion_latents(const sd::Tensor& latents) = 0; virtual void get_param_tensors(std::map& tensors) = 0; virtual void set_conv2d_scale(float scale) { SD_UNUSED(scale); }; - virtual void set_temporal_tiling_enabled(bool enabled) { SD_UNUSED(enabled); }; - virtual void set_tiling_params(const sd_tiling_params_t& params) { - set_temporal_tiling_enabled(params.temporal_tiling); - }; }; struct FakeVAE : public VAE { diff --git a/src/model/vae/vae_tiling.hpp b/src/model/vae/vae_tiling.hpp new file mode 100644 index 000000000..50cf08355 --- /dev/null +++ b/src/model/vae/vae_tiling.hpp @@ -0,0 +1,213 @@ +#ifndef __SD_MODEL_VAE_VAE_TILING_HPP__ +#define __SD_MODEL_VAE_VAE_TILING_HPP__ + +#include +#include +#include +#include + +#include "core/tensor.hpp" +#include "core/util.h" + +enum class VAETemporalDirection { + ENCODE, + DECODE, +}; + +struct VAETemporalTilingConfig { + int tile_frames = 1; + int overlap = 0; +}; + +struct VAETemporalTile { + int index = 0; + int64_t start = 0; + int64_t end = 0; + int overlap = 0; + bool first = false; + bool last = false; +}; + +struct VAETemporalTilePlan { + int tile_frames = 1; + int overlap = 0; + int stride = 1; + std::vector tiles; +}; + +inline VAETemporalTilingConfig resolve_vae_temporal_tiling_config(const sd_tiling_params_t& params, + int default_tile_frames, + int default_overlap) { + VAETemporalTilingConfig config; + config.tile_frames = std::max(1, default_tile_frames); + config.overlap = std::max(0, default_overlap); + + for (const auto& [key, value] : parse_key_value_args(params.extra_tiling_args, "VAE extra tiling arg")) { + if (key != "temporal_tile_frames" && key != "temporal_tile_size" && key != "temporal_tile_overlap") { + continue; + } + + int parsed = 0; + if (!parse_strict_int(value, parsed)) { + LOG_WARN("ignoring invalid VAE extra tiling arg '%s=%s'", key.c_str(), value.c_str()); + } else if (key == "temporal_tile_overlap") { + config.overlap = std::max(0, parsed); + } else { + config.tile_frames = std::max(1, parsed); + } + } + return config; +} + +inline VAETemporalTilePlan make_vae_temporal_tile_plan(int64_t total_frames, + const VAETemporalTilingConfig& config) { + VAETemporalTilePlan plan; + plan.tile_frames = std::max(1, config.tile_frames); + plan.overlap = std::max(0, config.overlap); + if (total_frames <= 1) { + plan.overlap = 0; + } + + if (plan.overlap >= plan.tile_frames) { + LOG_WARN("temporal_tile_overlap (%d) is greater than or equal to temporal_tile_frames (%d), adjusting values to avoid empty decode windows", + plan.overlap, + plan.tile_frames); + plan.overlap = plan.tile_frames - 1; + } + if (total_frames > 1 && plan.overlap >= total_frames) { + LOG_WARN("temporal_tile_overlap (%d) is greater than or equal to total frames (%lld), adjusting values to process at least one tile", + plan.overlap, + (long long)total_frames); + plan.overlap = static_cast(total_frames - 1); + } + + plan.stride = std::max(1, plan.tile_frames - plan.overlap); + for (int64_t start = 0; start < total_frames - plan.overlap; start += plan.stride) { + VAETemporalTile tile; + tile.index = static_cast(plan.tiles.size()); + tile.start = start; + tile.end = std::min(total_frames, start + plan.tile_frames); + tile.overlap = tile.end < total_frames ? plan.overlap : 0; + tile.first = start == 0; + tile.last = tile.end == total_frames; + plan.tiles.push_back(tile); + } + return plan; +} + +template +inline sd::Tensor process_vae_temporal_tiles(const sd::Tensor& input, + const VAETemporalTilePlan& plan, + Fn&& on_processing) { + sd::Tensor output; + for (const auto& tile : plan.tiles) { + auto input_tile = sd::ops::slice(input, 2, tile.start, tile.end); + auto output_tile = on_processing(input_tile, tile); + if (output_tile.empty()) { + return {}; + } + output = output.empty() ? std::move(output_tile) + : sd::ops::concat(output, output_tile, 2); + } + return output; +} + +template +inline sd::Tensor process_vae_temporal_tiles_blended(const sd::Tensor& input, + const VAETemporalTilePlan& plan, + int output_scale, + Fn&& on_processing) { + GGML_ASSERT(output_scale >= 1); + const int64_t output_frames = 1 + (input.shape()[2] - 1) * output_scale; + const int overlap_frames = plan.overlap > 0 ? 1 + (plan.overlap - 1) * output_scale : 0; + std::vector weights(static_cast(output_frames), 0.f); + sd::Tensor output; + + auto smootherstep = [](float value) { + return value * value * value * (value * (value * 6.f - 15.f) + 10.f); + }; + + for (const auto& tile : plan.tiles) { + auto input_tile = sd::ops::slice(input, 2, tile.start, tile.end); + auto output_tile = on_processing(input_tile, tile); + if (output_tile.empty()) { + return {}; + } + + const int64_t expected_tile_frames = 1 + (input_tile.shape()[2] - 1) * output_scale; + if (output_tile.dim() < 3 || output_tile.shape()[2] != expected_tile_frames) { + LOG_ERROR("unexpected temporal tile output shape: expected %lld frames, got %lld", + (long long)expected_tile_frames, + output_tile.dim() < 3 ? -1LL : (long long)output_tile.shape()[2]); + return {}; + } + + if (output.empty()) { + auto output_shape = output_tile.shape(); + output_shape[2] = output_frames; + output = sd::Tensor::zeros(std::move(output_shape)); + } else { + if (output.dim() != output_tile.dim()) { + LOG_ERROR("temporal tile output rank mismatch: expected %lld, got %lld", + (long long)output.dim(), + (long long)output_tile.dim()); + return {}; + } + for (size_t dim = 0; dim < static_cast(output.dim()); ++dim) { + if (dim != 2 && output.shape()[dim] != output_tile.shape()[dim]) { + LOG_ERROR("temporal tile output shape mismatch at dimension %zu", dim); + return {}; + } + } + } + + const int64_t output_start = tile.start * output_scale; + const int64_t inner = output.shape()[0] * output.shape()[1]; + const int64_t outer = output.numel() / (inner * output.shape()[2]); + const int64_t tile_frames = output_tile.shape()[2]; + for (int64_t frame = 0; frame < tile_frames; ++frame) { + float weight = 1.f; + if (!tile.first && overlap_frames > 0 && frame < overlap_frames) { + weight *= smootherstep(static_cast(frame + 1) / + static_cast(overlap_frames + 1)); + } + if (!tile.last && overlap_frames > 0 && frame >= tile_frames - overlap_frames) { + weight *= smootherstep(static_cast(tile_frames - frame) / + static_cast(overlap_frames + 1)); + } + + const int64_t output_frame = output_start + frame; + GGML_ASSERT(output_frame >= 0 && output_frame < output_frames); + weights[static_cast(output_frame)] += weight; + for (int64_t outer_index = 0; outer_index < outer; ++outer_index) { + const int64_t src_offset = (outer_index * tile_frames + frame) * inner; + const int64_t dst_offset = (outer_index * output_frames + output_frame) * inner; + for (int64_t inner_index = 0; inner_index < inner; ++inner_index) { + output[dst_offset + inner_index] += output_tile[src_offset + inner_index] * weight; + } + } + } + } + + if (output.empty()) { + return {}; + } + const int64_t inner = output.shape()[0] * output.shape()[1]; + const int64_t outer = output.numel() / (inner * output.shape()[2]); + for (int64_t frame = 0; frame < output_frames; ++frame) { + const float weight = weights[static_cast(frame)]; + if (weight <= 0.f) { + LOG_ERROR("temporal tiling left output frame %lld uncovered", (long long)frame); + return {}; + } + for (int64_t outer_index = 0; outer_index < outer; ++outer_index) { + const int64_t offset = (outer_index * output_frames + frame) * inner; + for (int64_t inner_index = 0; inner_index < inner; ++inner_index) { + output[offset + inner_index] /= weight; + } + } + } + return output; +} + +#endif // __SD_MODEL_VAE_VAE_TILING_HPP__ diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index 949562096..422a57823 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -1219,24 +1219,40 @@ namespace WAN { return out; } - ggml_tensor* decode_partial(GGMLRunnerContext* ctx, - ggml_tensor* z, - int i, - int64_t b = 1) { + ggml_tensor* decode_tiled_chunk(GGMLRunnerContext* ctx, + ggml_tensor* z, + int chunk_idx, + int64_t b = 1) { // z: [b*c, t, h, w] GGML_ASSERT(b == 1); auto decoder = std::dynamic_pointer_cast(blocks["decoder"]); auto conv2 = std::dynamic_pointer_cast(blocks["conv2"]); - auto x = conv2->forward(ctx, z); - // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decode_partial.prelude", "x"); - auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1); // [b*c, 1, h, w] - _conv_idx = 0; - auto out = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i); - out = unpatchify(ctx->ggml_ctx, out, patch_size, b); - // sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode_partial.final", "out"); - return out; + ggml_tensor* x; + if (is_2D) { + auto conv2_2d = std::dynamic_pointer_cast(blocks["conv2"]); + x = conv2_2d->forward(ctx, z); + } else { + x = conv2->forward(ctx, z); + } + + ggml_tensor* out = nullptr; + for (int64_t frame = 0; frame < x->ne[2]; ++frame) { + const int global_frame = chunk_idx + static_cast(frame); + auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, frame, frame + 1); + _conv_idx = 0; + auto out_frame = decoder->forward(ctx, in, b, _feat_map, _conv_idx, global_frame); + if (is_2D && global_frame > 0) { + auto repeated = out_frame; + for (int repeat = 1; repeat < 4; ++repeat) { + repeated = ggml_concat(ctx->ggml_ctx, repeated, out_frame, 2); + } + out_frame = repeated; + } + out = out == nullptr ? out_frame : ggml_concat(ctx->ggml_ctx, out, out_frame, 2); + } + return unpatchify(ctx->ggml_ctx, out, patch_size, b); } }; @@ -1272,6 +1288,15 @@ namespace WAN { return "wan_vae"; } + bool supports_temporal_tiling(VAETemporalDirection direction) const override { + return direction == VAETemporalDirection::DECODE; + } + + int get_temporal_tile_output_scale(VAETemporalDirection direction) const override { + SD_UNUSED(direction); + return 4; + } + void get_param_tensors(std::map& tensors) override { ae.get_param_tensors(tensors, weight_prefix); } @@ -1346,8 +1371,8 @@ namespace WAN { return gf; } - ggml_cgraph* build_graph_partial(const sd::Tensor& z_tensor, bool decode_graph, int i) { - ggml_cgraph* gf = new_graph_custom(20480); + ggml_cgraph* build_temporal_tile_graph(const sd::Tensor& z_tensor, int chunk_idx) { + ggml_cgraph* gf = new_graph_custom(std::max(20480, 10240 * z_tensor.shape()[2])); ae.clear_cache(); @@ -1360,7 +1385,7 @@ namespace WAN { auto runner_ctx = get_context(); - ggml_tensor* out = decode_graph ? ae.decode_partial(&runner_ctx, z, i) : ae.encode(&runner_ctx, z); + ggml_tensor* out = ae.decode_tiled_chunk(&runner_ctx, z, chunk_idx); for (size_t feat_idx = 0; feat_idx < ae._feat_map.size(); feat_idx++) { ggml_tensor* feat_cache = ae._feat_map[feat_idx]; @@ -1375,58 +1400,60 @@ namespace WAN { return gf; } - sd::Tensor _compute(const int n_threads, - const sd::Tensor& z, - bool decode_graph) override { - if (true) { - sd::Tensor input; - if (z.dim() == 4) { - input = z.unsqueeze(2); - } - auto get_graph = [&]() -> ggml_cgraph* { - if (input.empty()) { - return build_graph(z, decode_graph); - } else { - return build_graph(input, decode_graph); - } - }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true, true, true), - input.empty() ? z.dim() : input.dim()); - if (!result.empty() && z.dim() == 4) { - result.squeeze_(2); - } - return result; - } else { // chunk 1 result is weird - ae.clear_cache(); - int64_t t = z.shape()[2]; - int i = 0; + sd::Tensor _compute_temporal_tiled(const int n_threads, + const sd::Tensor& input, + VAETemporalDirection direction, + const VAETemporalTilingConfig& config) override { + GGML_ASSERT(direction == VAETemporalDirection::DECODE); + VAETemporalTilingConfig stateful_config = config; + stateful_config.overlap = 0; + auto plan = make_vae_temporal_tile_plan(input.shape()[2], stateful_config); + + LOG_DEBUG("Wan VAE stateful temporal tiling: tile_frames=%d, total latent frames=%lld, tiles=%d", + plan.tile_frames, + (long long)input.shape()[2], + (int)plan.tiles.size()); + + free_cache_ctx_and_buffer(); + cache_tensor_map.clear(); + ae.clear_cache(); + + auto output = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor& input_tile, const VAETemporalTile& tile) { + LOG_DEBUG("Wan VAE temporal tile %d/%d: latent frames [%lld, %lld)", + tile.index + 1, + (int)plan.tiles.size(), + (long long)tile.start, + (long long)tile.end); auto get_graph = [&]() -> ggml_cgraph* { - return build_graph_partial(z, decode_graph, i); + return build_temporal_tile_graph(input_tile, static_cast(tile.start)); }; - auto out_opt = GGMLRunner::compute(get_graph, n_threads, true, true, true); - if (!out_opt.has_value()) { - return {}; - } - sd::Tensor out = std::move(*out_opt); - ae.clear_cache(); - if (t == 1) { - return out; - } + return restore_trailing_singleton_dims( + GGMLRunner::compute(get_graph, n_threads, true, true, true), + static_cast(input.dim())); + }); - sd::Tensor output = std::move(out); + free_cache_ctx_and_buffer(); + cache_tensor_map.clear(); + ae.clear_cache(); + return output; + } - for (i = 1; i < t; i++) { - auto chunk_opt = GGMLRunner::compute(get_graph, n_threads, true, true, true); - if (!chunk_opt.has_value()) { - return {}; - } - out = std::move(*chunk_opt); - ae.clear_cache(); - output = sd::ops::concat(output, out, 2); - } - free_cache_ctx_and_buffer(); - return output; + sd::Tensor _compute(const int n_threads, + const sd::Tensor& z, + bool decode_graph) override { + sd::Tensor input; + if (z.dim() == 4) { + input = z.unsqueeze(2); + } + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(input.empty() ? z : input, decode_graph); + }; + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true, true, true), + input.empty() ? z.dim() : input.dim()); + if (!result.empty() && z.dim() == 4) { + result.squeeze_(2); } + return result; } void test() { diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 54ff82928..d2193e14c 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -2378,11 +2378,9 @@ class StableDiffusionGGML { sd::Tensor vae_latents; sd::Tensor decoded; if (preview_vae) { - preview_vae->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling); vae_latents = preview_vae->diffusion_to_vae_latents(_latents); decoded = preview_vae->decode(n_threads, vae_latents, vae_tiling_params, is_video, circular_x, circular_y, true); } else { - first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling); vae_latents = first_stage_model->diffusion_to_vae_latents(_latents); decoded = first_stage_model->decode(n_threads, vae_latents, vae_tiling_params, is_video, circular_x, circular_y, true); } @@ -3084,16 +3082,14 @@ class StableDiffusionGGML { if (sd_version_is_pid(version) || sd_version_is_minit2i(version)) { return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f); } - auto latents = first_stage_model->diffusion_to_vae_latents(x); - first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling); - auto decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); - if (decoded.empty() && auto_fit_enabled) { - bool prefer_temporal_tiling = decode_video && std::dynamic_pointer_cast(first_stage_model) != nullptr; - if (sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) { - first_stage_model->free_compute_buffer(); - first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling); - decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); - } + auto latents = first_stage_model->diffusion_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() && + auto_fit_enabled && + sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) { + first_stage_model->free_compute_buffer(); + decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); } return decoded; }