feat: generalize temporal tiling across video VAEs (#1926)

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 @@ protected:
                                        const sd::Tensor<float>& 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<float> _compute_temporal_tiled(const int n_threads,
+                                                      const sd::Tensor<float>& 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<float>& 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<float> compute_with_temporal_tiling(const int n_threads,
+                                                   const sd::Tensor<float>& 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<float>* tensor) {
         GGML_ASSERT(tensor != nullptr);
         for (int64_t i = 0; i < tensor->numel(); ++i) {
@@ -40,10 +122,15 @@ protected:
                                     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<float>& 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<float>();
@@ -86,6 +173,10 @@ public:
 
     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 @@ public:
                                    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 @@ public:
         int64_t t0              = ggml_time_ms();
         sd::Tensor<float> input = x;
         sd::Tensor<float> output;
-        set_tiling_params(tiling_params);
 
         if (tiling_params.enabled) {
             const int scale_factor = get_scale_factor();
@@ -201,10 +295,14 @@ public:
                 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 @@ public:
     virtual sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents)                           = 0;
     virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& 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 {