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

Co-authored-by: leejet <leejet714@gmail.com>

examples/cli/main.cpp

@@ -81,7 +81,7 @@ struct SDCliParams {
              &metadata_format},
             {"",
              "--preview-path",
-             "path to write preview image to (default: ./preview.png). Multi-frame previews support .avi, .webm, and animated .webp",
+             "path to write preview image to (default: ./preview.png). For image generation, the filename can have %03d placeholder for sequential numbering. Multi-frame previews support .avi, .webm, and animated .webp",
              0,
              &preview_path},
             {"",
@@ -94,7 +94,7 @@ struct SDCliParams {
         options.int_options = {
             {"",
              "--preview-interval",
-             "interval in denoising steps between consecutive updates of the image preview file (default is 1, meaning updating at every step)",
+             "preview interval: in each sampling pass, positive N updates every Nth denoiser step and -N previews only completed logical step N; 0 previews the final completed step of the first pass (base-resolution or high-noise). Default: 1",
              &preview_interval},
             {"",
              "--output-begin-idx",
@@ -377,29 +377,6 @@ bool load_images_from_dir(const std::string dir,
     return true;
 }
 
-void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy, void* data) {
-    (void)step;
-    (void)is_noisy;
-    SDCliParams* cli_params = (SDCliParams*)data;
-    // is_noisy is set to true if the preview corresponds to noisy latents, false if it's denoised latents
-    // unused in this app, it will either be always noisy or always denoised here
-    if (frame_count == 1) {
-        if (!write_image_to_file(cli_params->preview_path,
-                                 image->data,
-                                 image->width,
-                                 image->height,
-                                 image->channel,
-                                 "",
-                                 cli_params->compression_quality)) {
-            LOG_ERROR("save preview image to '%s' failed", cli_params->preview_path.c_str());
-        }
-    } else {
-        if (create_video_from_sd_images(cli_params->preview_path.c_str(), image, frame_count, cli_params->preview_fps, cli_params->compression_quality) != 0) {
-            LOG_ERROR("save preview video to '%s' failed", cli_params->preview_path.c_str());
-        }
-    }
-}
-
 std::string format_frame_idx(std::string pattern, int frame_idx) {
     std::smatch match;
     std::string result = pattern;
@@ -419,6 +396,36 @@ std::string format_frame_idx(std::string pattern, int frame_idx) {
     return result;
 }
 
+int continuous_preview_counter = 0;
+
+void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy, void* data) {
+    (void)step;
+    (void)is_noisy;
+    SDCliParams* cli_params = (SDCliParams*)data;
+    // is_noisy is set to true if the preview corresponds to noisy latents, false if it's denoised latents
+    // unused in this app, it will either be always noisy or always denoised here
+    if (frame_count == 1) {
+        fs::path path = cli_params->preview_path;
+        if (encoded_image_format_from_path(path.string()) == EncodedImageFormat::UNKNOWN)
+            path += ".png";
+        if (std::regex_search(path.string(), format_specifier_regex))
+            path = fs::path(format_frame_idx(path.string(), continuous_preview_counter++));
+        if (!write_image_to_file(path.string(),
+                                 image->data,
+                                 image->width,
+                                 image->height,
+                                 image->channel,
+                                 "",
+                                 cli_params->compression_quality)) {
+            LOG_ERROR("save preview image to '%s' failed", path.string().c_str());
+        }
+    } else {
+        if (create_video_from_sd_images(cli_params->preview_path.c_str(), image, frame_count, cli_params->preview_fps, cli_params->compression_quality) != 0) {
+            LOG_ERROR("save preview video to '%s' failed", cli_params->preview_path.c_str());
+        }
+    }
+}
+
 static fs::path get_video_audio_sidecar_path(const SDCliParams& cli_params) {
     fs::path out_path     = cli_params.output_path;
     fs::path base_path    = out_path;

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:
 

include/stable-diffusion.h

@@ -446,6 +446,9 @@ typedef bool (*sd_graph_eval_callback_t)(struct ggml_tensor* t, bool ask, void*
 
 SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data);
 SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data);
+// In each sampling pass, a positive interval previews every Nth denoiser step, while a
+// negative interval previews only completed logical step -interval. Zero previews the final
+// completed step of the first sampling pass (base-resolution or high-noise).
 SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy, void* data);
 SD_API void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data);
 SD_API int32_t sd_get_num_physical_cores();

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;
     }
 

src/model/diffusion/minimax_h3.hpp

@@ -123,25 +123,6 @@ namespace MiniMaxH3 {
         return to_shift * base / (1.f + (to_shift - 1.f) * base);
     }
 
-    static float time_shift_slope(float sigma, float from_shift, float to_shift) {
-        float base = sigma / (from_shift + sigma * (1.f - from_shift));
-        float a    = 1.f + (from_shift - 1.f) * base;
-        float b    = 1.f + (to_shift - 1.f) * base;
-        return to_shift * a * a / (from_shift * b * b);
-    }
-
-    static float time_shift_step_scale(float sigma,
-                                       float next_sigma,
-                                       float from_shift,
-                                       float to_shift) {
-        if (!std::isfinite(next_sigma) || next_sigma < 0.f || next_sigma == sigma) {
-            return time_shift_slope(sigma, from_shift, to_shift);
-        }
-        float shifted_sigma      = time_shift_sigma(sigma, from_shift, to_shift);
-        float shifted_next_sigma = time_shift_sigma(next_sigma, from_shift, to_shift);
-        return (shifted_sigma - shifted_next_sigma) / (sigma - next_sigma);
-    }
-
     struct TimeEmbedder : public GGMLBlock {
         TimeEmbedder(int64_t input_dim, int64_t hidden_dim, int64_t output_dim) {
             blocks["proj_in"]  = std::make_shared<Linear>(input_dim, hidden_dim, true, true);
@@ -606,8 +587,7 @@ namespace MiniMaxH3 {
                                                       const std::vector<TokenModulationSpan>& segments,
                                                       const std::vector<SequenceSegment>& sequence_segments,
                                                       const TokenModulationSpan& video_segment,
-                                                      const TokenModulationSpan& audio_segment,
-                                                      float audio_slope) {
+                                                      const TokenModulationSpan& audio_segment) {
             auto video_proj = std::dynamic_pointer_cast<Linear>(blocks["video_patch_proj"]);
             auto audio_proj = std::dynamic_pointer_cast<Linear>(blocks["audio_patch_proj"]);
 
@@ -727,7 +707,7 @@ namespace MiniMaxH3 {
                                                audio->ne[2]);
             audio_out        = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, audio_out, 1, 2, 0, 3));
             video_out        = ggml_ext_scale(ctx->ggml_ctx, video_out, -1.f);
-            audio_out        = ggml_ext_scale(ctx->ggml_ctx, audio_out, -audio_slope);
+            audio_out        = ggml_ext_scale(ctx->ggml_ctx, audio_out, -1.f);
             return {video_out, audio_out};
         }
     };
@@ -1045,17 +1025,16 @@ namespace MiniMaxH3 {
                                  const std::vector<MiniMaxH3ReferenceBlock>& reference_blocks,
                                  int audio_length,
                                  float video_shift,
-                                 float audio_shift,
-                                 float next_video_sigma) {
+                                 float audio_shift) {
             auto split        = split_av_latents(packed, audio_length);
             video_input_cache = std::move(split.first);
             audio_input_cache = std::move(split.second);
             GGML_ASSERT(!audio_input_cache.empty());
             GGML_ASSERT(!context_tensor.empty());
 
-            auto video   = make_input(video_input_cache);
-            auto audio   = make_input(audio_input_cache);
-            auto context = make_input(context_tensor);
+            auto video         = make_input(video_input_cache);
+            auto audio_carrier = make_input(audio_input_cache);
+            auto context       = make_input(context_tensor);
             std::vector<ggml_tensor*> condition_inputs;
             condition_inputs.reserve(condition_videos.size());
             for (const auto& condition : condition_videos) {
@@ -1067,21 +1046,26 @@ namespace MiniMaxH3 {
                 audio_condition_inputs.push_back(make_input(condition));
             }
 
-            float sigma_v = std::clamp(timestep[0] / 1000.f, 1e-6f, 1.f);
-            float t_v     = 1.f - sigma_v;
-            float t_a     = 1.f - time_shift_sigma(sigma_v, video_shift, audio_shift);
-            auto layout   = build_layout(context_tensor.shape()[1],
-                                         video_input_cache.shape()[2],
-                                         video_input_cache.shape()[1],
-                                         video_input_cache.shape()[0],
-                                         audio_length,
-                                         condition_videos,
-                                         condition_audios,
-                                         keyframe_indices,
-                                         reference_blocks,
-                                         text_tags,
-                                         t_v,
-                                         t_a);
+            float sigma_v     = std::clamp(timestep[0] / 1000.f, 1e-6f, 1.f);
+            float sigma_a     = time_shift_sigma(sigma_v, video_shift, audio_shift);
+            float audio_scale = video_shift / audio_shift;
+            float t_v         = 1.f - sigma_v;
+            float t_a         = 1.f - sigma_a;
+            // The sampler carries c_a = (sigma_v / sigma_a) * x_a so the packed
+            // latent follows one sigma schedule. Restore x_a for the H3 network.
+            auto audio  = ggml_ext_scale(compute_ctx, audio_carrier, sigma_a / sigma_v);
+            auto layout = build_layout(context_tensor.shape()[1],
+                                       video_input_cache.shape()[2],
+                                       video_input_cache.shape()[1],
+                                       video_input_cache.shape()[0],
+                                       audio_length,
+                                       condition_videos,
+                                       condition_audios,
+                                       keyframe_indices,
+                                       reference_blocks,
+                                       text_tags,
+                                       t_v,
+                                       t_a);
 
             position_input_cache = sd::Tensor<float>(
                 {3, static_cast<int64_t>(layout.positions.size() / 3)},
@@ -1142,19 +1126,15 @@ namespace MiniMaxH3 {
                                             layout.segments,
                                             layout.sequence_segments,
                                             layout.video_segment,
-                                            layout.audio_segment,
-                                            // The generic Euler sampler advances the packed tensor by
-                                            // `next_video_sigma - sigma_v`. For that sampler, scale H3's
-                                            // audio velocity by the exact ratio of the independent audio
-                                            // step. The derivative approximation substantially oversteps
-                                            // at low step counts (the Turbo use case). Retain the local
-                                            // slope for samplers that make extra/intermediate evaluations.
-                                            time_shift_step_scale(sigma_v,
-                                                                  next_video_sigma,
-                                                                  video_shift,
-                                                                  audio_shift));
-            auto merged     = merge_av_latents(compute_ctx, output.first, output.second);
-            auto graph      = new_graph_custom(H3_GRAPH_SIZE);
+                                            layout.audio_segment);
+            // Convert the model's audio velocity to d(c_a) / d(sigma_v).
+            output.second = ggml_add(compute_ctx,
+                                     ggml_ext_scale(compute_ctx, audio, 1.f - audio_scale),
+                                     ggml_ext_scale(compute_ctx,
+                                                    output.second,
+                                                    1.f + (audio_scale - 1.f) * sigma_a));
+            auto merged   = merge_av_latents(compute_ctx, output.first, output.second);
+            auto graph    = new_graph_custom(H3_GRAPH_SIZE);
             ggml_build_forward_expand(graph, merged);
             return graph;
         }
@@ -1184,8 +1164,7 @@ namespace MiniMaxH3 {
                                    reference_blocks,
                                    extra->audio_length,
                                    extra->video_sigma_shift,
-                                   extra->audio_sigma_shift,
-                                   extra->next_video_sigma);
+                                   extra->audio_sigma_shift);
             };
             return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph,
                                                                               n_threads,

src/model/diffusion/model.hpp

@@ -108,8 +108,6 @@ struct MiniMaxH3DiffusionExtra {
     int audio_length                                              = 0;
     float video_sigma_shift                                       = 12.f;
     float audio_sigma_shift                                       = 3.f;
-    // Negative when the outer sampler is not a single-evaluation Euler step.
-    float next_video_sigma = -1.f;
 };
 
 struct MiniT2IDiffusionExtra {

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,64 +1245,24 @@ 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;
+    int get_default_temporal_tile_frames(VAETemporalDirection direction) const override {
+        SD_UNUSED(direction);
+        return DEFAULT_TEMPORAL_TILE_FRAMES;
+    }
 
-        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_overlap(VAETemporalDirection direction) const override {
+        SD_UNUSED(direction);
+        return DEFAULT_TEMPORAL_TILE_OVERLAP;
     }
 
     void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         vae.get_param_tensors(tensors, weight_prefix);
     }
 
-    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;
-    }
-
     std::string temporal_feat_cache_name(size_t feat_idx) const {
         return "ltx_vae_temporal_feat:" + std::to_string(feat_idx);
     }
@@ -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,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;

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);
     }

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 {

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 <algorithm>
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#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<VAETemporalTile> 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<int>(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<int>(plan.tiles.size());
+        tile.start   = start;
+        tile.end     = std::min<int64_t>(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 <typename Fn>
+inline sd::Tensor<float> process_vae_temporal_tiles(const sd::Tensor<float>& input,
+                                                    const VAETemporalTilePlan& plan,
+                                                    Fn&& on_processing) {
+    sd::Tensor<float> 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 <typename Fn>
+inline sd::Tensor<float> process_vae_temporal_tiles_blended(const sd::Tensor<float>& 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<float> weights(static_cast<size_t>(output_frames), 0.f);
+    sd::Tensor<float> 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<float>::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<size_t>(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<float>(frame + 1) /
+                                       static_cast<float>(overlap_frames + 1));
+            }
+            if (!tile.last && overlap_frames > 0 && frame >= tile_frames - overlap_frames) {
+                weight *= smootherstep(static_cast<float>(tile_frames - frame) /
+                                       static_cast<float>(overlap_frames + 1));
+            }
+
+            const int64_t output_frame = output_start + frame;
+            GGML_ASSERT(output_frame >= 0 && output_frame < output_frames);
+            weights[static_cast<size_t>(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<size_t>(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__

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<Decoder3d>(blocks["decoder"]);
             auto conv2   = std::dynamic_pointer_cast<CausalConv3d>(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<Conv2dBut3d>(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<int>(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<std::string, ggml_tensor*>& 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<float>& z_tensor, bool decode_graph, int i) {
-            ggml_cgraph* gf = new_graph_custom(20480);
+        ggml_cgraph* build_temporal_tile_graph(const sd::Tensor<float>& z_tensor, int chunk_idx) {
+            ggml_cgraph* gf = new_graph_custom(std::max<size_t>(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<float> _compute_temporal_tiled(const int n_threads,
+                                                  const sd::Tensor<float>& 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<float>& 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_temporal_tile_graph(input_tile, static_cast<int>(tile.start));
+                };
+                return restore_trailing_singleton_dims(
+                    GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
+                    static_cast<size_t>(input.dim()));
+            });
+
+            free_cache_ctx_and_buffer();
+            cache_tensor_map.clear();
+            ae.clear_cache();
+            return output;
+        }
+
         sd::Tensor<float> _compute(const int n_threads,
                                    const sd::Tensor<float>& z,
                                    bool decode_graph) override {
-            if (true) {
-                sd::Tensor<float> 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<float>(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;
-                auto get_graph = [&]() -> ggml_cgraph* {
-                    return build_graph_partial(z, decode_graph, i);
-                };
-                auto out_opt = GGMLRunner::compute<float>(get_graph, n_threads, true, true, true);
-                if (!out_opt.has_value()) {
-                    return {};
-                }
-                sd::Tensor<float> out = std::move(*out_opt);
-                ae.clear_cache();
-                if (t == 1) {
-                    return out;
-                }
-
-                sd::Tensor<float> output = std::move(out);
-
-                for (i = 1; i < t; i++) {
-                    auto chunk_opt = GGMLRunner::compute<float>(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<float> 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<float>(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() {

src/runtime/denoiser.hpp

@@ -1043,6 +1043,16 @@ struct Denoiser {
                                                     const sd::Tensor<float>& latent) = 0;
     virtual float noise_level_to_sigma(float noise_level)                            = 0;
 
+    virtual sd::Tensor<float> process_latent_in(const sd::Tensor<float>& latent) {
+        // An empty result means the original latent can be used unchanged.
+        SD_UNUSED(latent);
+        return {};
+    }
+
+    virtual sd::Tensor<float> process_latent_out(sd::Tensor<float> latent) {
+        return latent;
+    }
+
     virtual std::vector<float> get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler_type, SDVersion version, const char* extra_sample_args = nullptr) {
         auto bound_t_to_sigma = std::bind(&Denoiser::t_to_sigma, this, std::placeholders::_1);
         std::shared_ptr<SigmaScheduler> scheduler;
@@ -1286,6 +1296,40 @@ struct DiscreteFlowDenoiser : public Denoiser {
     }
 };
 
+struct H3AVFlowDenoiser : public DiscreteFlowDenoiser {
+    int64_t video_channels;
+    float audio_shift;
+
+    H3AVFlowDenoiser(float shift, float audio_shift, int64_t video_channels)
+        : DiscreteFlowDenoiser(shift),
+          video_channels(video_channels),
+          audio_shift(audio_shift) {
+        GGML_ASSERT(shift > 0.f && audio_shift > 0.f && video_channels > 0);
+    }
+
+    sd::Tensor<float> process_latent_in(const sd::Tensor<float>& latent) override {
+        return scale_audio(latent, shift / audio_shift);
+    }
+
+    sd::Tensor<float> process_latent_out(sd::Tensor<float> latent) override {
+        auto transformed = scale_audio(latent, audio_shift / shift);
+        if (transformed.empty()) {
+            return latent;
+        }
+        return transformed;
+    }
+
+private:
+    sd::Tensor<float> scale_audio(const sd::Tensor<float>& latent, float scale) const {
+        if (scale == 1.f || latent.dim() < 4 || latent.shape()[3] <= video_channels) {
+            return {};
+        }
+        auto video = sd::ops::slice(latent, 3, 0, video_channels);
+        auto audio = sd::ops::slice(latent, 3, video_channels, latent.shape()[3]) * scale;
+        return sd::ops::concat(video, audio, 3);
+    }
+};
+
 struct FluxFlowDenoiser : public DiscreteFlowDenoiser {
     FluxFlowDenoiser() = default;
 

src/runtime/preview_interval.h

@@ -0,0 +1,45 @@
+#ifndef __SD_RUNTIME_PREVIEW_INTERVAL_H__
+#define __SD_RUNTIME_PREVIEW_INTERVAL_H__
+
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+
+namespace sd::preview {
+
+    constexpr std::uint64_t logical_sample_step(int step) {
+        return step < 0 ? static_cast<std::uint64_t>(-static_cast<std::int64_t>(step))
+                        : static_cast<std::uint64_t>(step);
+    }
+
+    constexpr bool sample_step_is_complete(int step,
+                                           std::size_t total_steps,
+                                           bool terminal_sigma_is_zero) {
+        return step > 0 ||
+               (terminal_sigma_is_zero &&
+                step < 0 &&
+                logical_sample_step(step) == static_cast<std::uint64_t>(total_steps));
+    }
+
+    constexpr bool should_preview_sample_step(int step,
+                                              std::size_t total_steps,
+                                              bool terminal_sigma_is_zero,
+                                              int interval,
+                                              bool preview_final_step) {
+        if (interval > 0) {
+            return step % interval == 0;
+        }
+        if (!sample_step_is_complete(step, total_steps, terminal_sigma_is_zero)) {
+            return false;
+        }
+
+        std::uint64_t logical_step = logical_sample_step(step);
+        if (interval < 0) {
+            std::uint64_t requested_step = static_cast<std::uint64_t>(-static_cast<std::int64_t>(interval));
+            return logical_step == requested_step;
+        }
+        return preview_final_step && logical_step == static_cast<std::uint64_t>(total_steps);
+    }
+}  // namespace sd::preview
+
+#endif  // __SD_RUNTIME_PREVIEW_INTERVAL_H__

src/stable-diffusion.cpp

@@ -61,6 +61,7 @@
 #include "model/vae/wan_vae.hpp"
 #include "runtime/denoiser.hpp"
 #include "runtime/guidance.h"
+#include "runtime/preview_interval.h"
 #include "runtime/sample-cache.h"
 #include "upscaler.h"
 
@@ -1860,6 +1861,9 @@ public:
                     if (sd_version_is_ltxav(version)) {
                         LOG_INFO("running in LTXAV FLOW mode");
                         denoiser = std::make_shared<FluxFlowDenoiser>();
+                    } else if (sd_version_is_minimax_h3(version)) {
+                        LOG_INFO("running in MiniMax H3 AV FLOW mode");
+                        denoiser = std::make_shared<H3AVFlowDenoiser>(default_flow_shift, 3.f, get_latent_channel());
                     } else {
                         LOG_INFO("running in FLOW mode");
                         denoiser = std::make_shared<DiscreteFlowDenoiser>();
@@ -2374,11 +2378,9 @@ public:
             sd::Tensor<float> vae_latents;
             sd::Tensor<float> 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);
             }
@@ -2467,8 +2469,11 @@ public:
                                     sd_get_preview_mode()};
     }
 
-    void report_sample_progress(int step, size_t total_steps, int64_t* last_progress_us) {
-        if (step > 0 || step == -(int)total_steps) {
+    void report_sample_progress(int step,
+                                size_t total_steps,
+                                bool terminal_sigma_is_zero,
+                                int64_t* last_progress_us) {
+        if (sd::preview::sample_step_is_complete(step, total_steps, terminal_sigma_is_zero)) {
             int64_t now        = ggml_time_us();
             int showstep       = std::abs(step);
             float step_seconds = last_progress_us != nullptr && *last_progress_us > 0
@@ -2530,6 +2535,7 @@ public:
                              int audio_length,
                              float frame_rate,
                              const sd_cache_params_t* cache_params,
+                             bool preview_final_step,
                              const sd::Tensor<float>& video_positions = {}) {
         struct RunnerDoneOnExit {
             GGMLRunner* runner = nullptr;
@@ -2589,8 +2595,9 @@ public:
             }
         }
 
-        size_t steps       = sigmas.size() - 1;
-        bool has_skiplayer = (slg_scale != 0.0f || slg_uncond) && !skip_layers.empty();
+        size_t steps                = sigmas.size() - 1;
+        bool terminal_sigma_is_zero = sigmas.back() == 0.f;
+        bool has_skiplayer          = (slg_scale != 0.0f || slg_uncond) && !skip_layers.empty();
         if (has_skiplayer && !sd_version_is_dit(version)) {
             has_skiplayer = false;
             LOG_WARN("SLG is incompatible with this model type");
@@ -2617,10 +2624,14 @@ public:
         int64_t last_progress_us     = ggml_time_us();
         SamplePreviewContext preview = prepare_sample_preview_context();
 
-        sd::Tensor<float> x_t      = !noise.empty()
-                                         ? denoiser->noise_scaling(sigmas[0], noise, init_latent)
-                                         : init_latent;
-        sd::Tensor<float> denoised = x_t;
+        sd::Tensor<float> processed_init_latent       = denoiser->process_latent_in(init_latent);
+        const sd::Tensor<float>& sampling_init_latent = processed_init_latent.empty()
+                                                            ? init_latent
+                                                            : processed_init_latent;
+        sd::Tensor<float> x_t                         = !noise.empty()
+                                                            ? denoiser->noise_scaling(sigmas[0], noise, sampling_init_latent)
+                                                            : sampling_init_latent;
+        sd::Tensor<float> denoised                    = x_t;
 
         auto denoise = [&](const sd::Tensor<float>& x, float sigma, int step) -> sd::guidance::GuiderOutput {
             if (get_cancel_flag() == SD_CANCEL_ALL) {
@@ -2639,14 +2650,21 @@ public:
             float c_out  = scaling[1];
             float c_in   = scaling[2];
 
+            bool preview_needed = preview.callback != nullptr &&
+                                  sd::preview::should_preview_sample_step(step,
+                                                                          steps,
+                                                                          terminal_sigma_is_zero,
+                                                                          sd_get_preview_interval(),
+                                                                          preview_final_step);
+
             std::vector<float> base_timesteps_vec = prepare_sample_timesteps(sigma, shifted_timestep);
             std::vector<float> timesteps_vec      = base_timesteps_vec;
             sd::Tensor<float> audio_timesteps_tensor;
             if (sd_version_is_ltxav(version) && !denoise_mask.empty()) {
-                timesteps_vec          = process_ltxav_video_timesteps(base_timesteps_vec, init_latent, denoise_mask);
+                timesteps_vec          = process_ltxav_video_timesteps(base_timesteps_vec, sampling_init_latent, denoise_mask);
                 audio_timesteps_tensor = sd::Tensor<float>({static_cast<int64_t>(base_timesteps_vec.size())}, base_timesteps_vec);
             } else {
-                timesteps_vec = process_timesteps(timesteps_vec, init_latent, denoise_mask, step);
+                timesteps_vec = process_timesteps(timesteps_vec, sampling_init_latent, denoise_mask, step);
             }
             const std::vector<float>& scaling_timesteps_vec = (sd_version_is_ltxav(version) && !denoise_mask.empty())
                                                                   ? base_timesteps_vec
@@ -2661,29 +2679,25 @@ public:
             }
             sd::Tensor<float> noised_input = x * c_in;
             if (!denoise_mask.empty() && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version) || sd_version_is_lingbot_video(version))) {
-                noised_input = noised_input * denoise_mask + init_latent * (1.0f - denoise_mask);
+                noised_input = noised_input * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
             }
 
             if (cache_runtime.spectrum_enabled && cache_runtime.spectrum.should_predict()) {
                 cache_runtime.spectrum.predict(&denoised);
                 if (!denoise_mask.empty()) {
-                    denoised = denoised * denoise_mask + init_latent * (1.0f - denoise_mask);
+                    denoised = denoised * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
                 }
-                if (sd_should_preview_denoised() && preview.callback != nullptr) {
-                    if (step % sd_get_preview_interval() == 0) {
-                        preview_image(step, denoised, version, preview.mode, preview.callback, preview.data, false);
-                    }
+                if (preview_needed && sd_should_preview_denoised()) {
+                    preview_image(step, denoised, version, preview.mode, preview.callback, preview.data, false);
                 }
-                report_sample_progress(step, steps, &last_progress_us);
+                report_sample_progress(step, steps, terminal_sigma_is_zero, &last_progress_us);
                 sd::guidance::GuiderOutput output;
                 output.pred = denoised;
                 return output;
             }
 
-            if (sd_should_preview_noisy() && preview.callback != nullptr) {
-                if (step % sd_get_preview_interval() == 0) {
-                    preview_image(step, noised_input, version, preview.mode, preview.callback, preview.data, true);
-                }
+            if (preview_needed && sd_should_preview_noisy()) {
+                preview_image(step, noised_input, version, preview.mode, preview.callback, preview.data, true);
             }
 
             sd::Tensor<float> cond_out;
@@ -2765,11 +2779,7 @@ public:
                         condition.c_reference_blocks.empty() ? nullptr : &condition.c_reference_blocks,
                         audio_length,
                         std::isfinite(active_flow_shift) ? active_flow_shift : 12.f,
-                        3.f,
-                        method == EULER_SAMPLE_METHOD && step > 0 &&
-                                static_cast<size_t>(step) < sigmas.size()
-                            ? sigmas[step]
-                            : -1.f};
+                        3.f};
                 } else if (sd_version_is_ltxav(version)) {
                     diffusion_params.extra = LTXAVDiffusionExtra{
                         nullptr,
@@ -2894,14 +2904,12 @@ public:
                 cache_runtime.spectrum.update(denoised);
             }
             if (!denoise_mask.empty()) {
-                denoised = denoised * denoise_mask + init_latent * (1.0f - denoise_mask);
+                denoised = denoised * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
             }
-            if (sd_should_preview_denoised() && preview.callback != nullptr) {
-                if (step % sd_get_preview_interval() == 0) {
-                    preview_image(step, denoised, version, preview.mode, preview.callback, preview.data, false);
-                }
+            if (preview_needed && sd_should_preview_denoised()) {
+                preview_image(step, denoised, version, preview.mode, preview.callback, preview.data, false);
             }
-            report_sample_progress(step, steps, &last_progress_us);
+            report_sample_progress(step, steps, terminal_sigma_is_zero, &last_progress_us);
             output.pred = denoised;
             return output;
         };
@@ -2924,6 +2932,7 @@ public:
         if (inverse_noise_scaling) {
             x0 = denoiser->inverse_noise_scaling(sigmas[sigmas.size() - 1], x0);
         }
+        x0 = denoiser->process_latent_out(std::move(x0));
 
         if (control_net) {
             control_net->free_control_ctx();
@@ -3073,16 +3082,14 @@ public:
         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<LTXVideoVAE>(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;
     }
@@ -5719,7 +5726,8 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
                                                    1.f,
                                                    0,
                                                    static_cast<float>(request.fps),
-                                                   request.cache_params);
+                                                   request.cache_params,
+                                                   true);
         int64_t sampling_end  = ggml_time_ms();
         if (!x_0.empty()) {
             LOG_INFO("sampling completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000);
@@ -5840,7 +5848,8 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
                                                             1.f,
                                                             0,
                                                             static_cast<float>(request.fps),
-                                                            request.cache_params);
+                                                            request.cache_params,
+                                                            false);
             int64_t hires_sample_end   = ggml_time_ms();
             if (!x_0.empty()) {
                 LOG_INFO("hires sampling %d/%d completed, taking %.2fs",
@@ -6974,6 +6983,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
                                                            latents.audio_length,
                                                            static_cast<float>(request.fps),
                                                            request.cache_params,
+                                                           true,
                                                            latents.video_positions);
         int64_t sampling_end          = ggml_time_ms();
         if (x_t_sampled.empty()) {
@@ -7016,6 +7026,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
                                                         latents.audio_length,
                                                         static_cast<float>(request.fps),
                                                         request.cache_params,
+                                                        plan.high_noise_sample_steps <= 0,
                                                         latents.video_positions);
 
     int64_t sampling_end = ggml_time_ms();
@@ -7154,6 +7165,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
                                             latents.audio_length,
                                             static_cast<float>(hires_request.fps),
                                             hires_request.cache_params,
+                                            false,
                                             hires_video_positions);
         sampling_end   = ggml_time_ms();
         if (final_latent.empty()) {