feat: generalize temporal tiling across video VAEs

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},
         {"",

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:
 

src/core/backend_fit.cpp

@@ -364,26 +364,25 @@ 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;
             }
             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;
     }
 

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<std::string, ggml_tensor*>& tensors) override {
             if (!decode_only) {
                 encoder.get_param_tensors(tensors, weight_prefix + ".encoder");

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<std::string, ggml_tensor*>& 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<int>(total_frames - 1);
-        }
-
-        plan.stride          = std::max(1, plan.frames - plan.overlap);
-        int64_t tiled_frames = std::max<int64_t>(1, total_frames - plan.overlap);
-        plan.num_tiles       = total_frames > 0 ? static_cast<int>((tiled_frames + plan.stride - 1) / plan.stride) : 0;
-        return plan;
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& 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<float> decode_temporal_tiled_streaming(const int n_threads,
                                                       const sd::Tensor<float>& 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<float> output;
-        for (int64_t start = 0; start < total_frames - plan.overlap; start += plan.stride) {
-            const int64_t end       = std::min<int64_t>(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<float>& 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<int>(start),
-                                                 chunk_overlap);
+                                                 static_cast<int>(tile.start),
+                                                 tile.overlap);
             };
-            auto chunk = restore_trailing_singleton_dims(GGMLRunner::compute<float>(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<float>(get_graph, n_threads, true, true, true),
+                                                   expected_dim);
+        });
 
         free_cache_ctx_and_buffer();
         cache_tensor_map.clear();
         return output;
     }
 
+    sd::Tensor<float> _compute_temporal_tiled(const int n_threads,
+                                              const sd::Tensor<float>& input,
+                                              VAETemporalDirection direction,
+                                              const VAETemporalTilingConfig& config) override {
+        GGML_ASSERT(direction == VAETemporalDirection::DECODE);
+        return decode_temporal_tiled_streaming(n_threads,
+                                               input,
+                                               static_cast<size_t>(input.dim()),
+                                               config);
+    }
+
     ggml_cgraph* build_latent_statistics_graph(const sd::Tensor<float>& 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);
         };

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<float> 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<float>& 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<float> result;
             sd::Tensor<float> 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<int>(tokens_per_chunk + token_overlap), static_cast<int>(token_overlap)});
+            GGML_ASSERT(plan.tiles.size() == static_cast<size_t>(num_chunks));
+            auto result = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& 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<float>();
                 }
 
                 int64_t first_end = std::min<int64_t>(frames_per_chunk, decoded.shape()[2]);
@@ -712,19 +710,21 @@ 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,
                                              2,
                                              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;

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<std::string, ggml_tensor*>& tensors) override {
         taehv.get_param_tensors(tensors, weight_prefix);
     }