feat(vae): allow grouped temporal decoding for Wan2.2 TI2V 5B

docs/wan.md

@@ -184,6 +184,8 @@ Notes:
 
 <video src=../assets/wan/Wan2.2_5B_i2v.mp4 controls="controls" muted="muted" type="video/mp4"></video>
 
+TI2V 5B can decode multiple latent frames together when temporal tiling is enabled. For example, `--temporal-tiling --vae-tiling --vae-tile-size 256x256 --extra-tiling-args "temporal_tile_frames=4,temporal_decode_frames=2"` keeps up to four latent frames in each temporal tile and decodes up to two at once, with spatial tiles of 256×256 pixels. `temporal_decode_frames` defaults to `1`, preserving frame-by-frame decoding; other model versions ignore it. Increasing it requires more memory. Adjust the spatial tile size to fit the available memory.
+
 ### Wan2.1 FLF2V 14B
 
 ```

examples/common/common.cpp

@@ -1194,7 +1194,7 @@ ArgOptions SDGenerationParams::get_options() {
          &extra_sample_args},
         {"",
          "--extra-tiling-args",
-         "extra VAE tiling args, key=value list. Supported video VAEs accept temporal_tile_frames/temporal_tile_size (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); Wan2.2 TI2V 5B accepts temporal_decode_frames (default: 1)",
          (int)',',
          &extra_tiling_args},
         {"",

examples/server/api.md

@@ -877,7 +877,7 @@ Shared default fields used by both `img_gen` and `vid_gen`:
 
 `vae_tiling_params.tile_size_w` and `tile_size_h` are in **image pixels**, with `0` selecting the 256-pixel default. Both encode and decode use these sizes without an encoding multiplier. Positive `rel_size_w`/`rel_size_h` values override the corresponding absolute size: values up to 1 are dimension fractions, and values greater than 1 are target tile counts. Set `enabled` to use spatial tiling. Sizes are aligned down to the VAE scale factor and capped at the input dimensions; explicit sizes below the minimum supported tile size are rejected. These fields previously used latent units; see [VAE tiling](../../docs/performance.md#use-vae-tiling-to-reduce-encode-and-decode-memory-usage) for migration and OOM retry behavior.
 
-`vae_tiling_params.extra_tiling_args` accepts a key=value list. Supported video VAEs accept `temporal_tile_frames` (alias `temporal_tile_size`, default `4`) and `temporal_tile_overlap` (default `1`).
+`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`). Wan2.2 TI2V 5B uses `temporal_decode_frames` (default `1`) to decode multiple latent frames per forward pass, trading additional memory for throughput; other model versions ignore this option.
 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/ggml_extend.cpp

@@ -554,9 +554,9 @@ std::vector<ggml_tensor*> split_image_qkv(ggml_context* ctx,
     qkv         = ggml_cont(ctx, ggml_ext_torch_permute(ctx, qkv, 0, 1, 3, 2));  // [3, N, C, H*W]
 
     int64_t offset = qkv->nb[2] * qkv->ne[2];
-    auto q         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 0);  // [N, C, H, W]
-    auto k         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 1);  // [N, C, H, W]
-    auto v         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 2);  // [N, C, H, W]
+    auto q         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[2], offset * 0);  // [N, C, H, W]
+    auto k         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[2], offset * 1);  // [N, C, H, W]
+    auto v         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[2], offset * 2);  // [N, C, H, W]
     return {q, k, v};
 }
 

src/model/vae/vae_tiling.hpp

@@ -15,8 +15,9 @@ enum class VAETemporalDirection {
 };
 
 struct VAETemporalTilingConfig {
-    int tile_frames = 1;
-    int overlap     = 0;
+    int tile_frames   = 1;
+    int overlap       = 0;
+    int decode_frames = 1;
 };
 
 struct VAETemporalTile {
@@ -43,7 +44,8 @@ inline VAETemporalTilingConfig resolve_vae_temporal_tiling_config(const sd_tilin
     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") {
+        if (key != "temporal_tile_frames" && key != "temporal_tile_size" && key != "temporal_tile_overlap" &&
+            key != "temporal_decode_frames") {
             continue;
         }
 
@@ -52,6 +54,8 @@ inline VAETemporalTilingConfig resolve_vae_temporal_tiling_config(const sd_tilin
             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 if (key == "temporal_decode_frames") {
+            config.decode_frames = std::max(1, parsed);
         } else {
             config.tile_frames = std::max(1, parsed);
         }

src/model/vae/wan_vae.hpp

@@ -1287,9 +1287,11 @@ namespace WAN {
         ggml_tensor* decode_tiled_chunk(GGMLRunnerContext* ctx,
                                         ggml_tensor* z,
                                         int chunk_idx,
-                                        int64_t b = 1) {
+                                        int64_t b         = 1,
+                                        int decode_frames = 1) {
             // z: [b*c, t, h, w]
             GGML_ASSERT(b == 1);
+            GGML_ASSERT(decode_frames >= 1);
 
             auto decoder = std::dynamic_pointer_cast<Decoder3d>(blocks["decoder"]);
             auto conv2   = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv2"]);
@@ -1303,11 +1305,13 @@ namespace WAN {
             }
 
             ggml_tensor* out = nullptr;
-            for (int64_t frame = 0; frame < x->ne[2]; ++frame) {
+            for (int64_t frame = 0; frame < x->ne[2];) {
                 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);
+                // The first latent frame bypasses temporal upsampling and primes the causal caches.
+                const int64_t count = !is_2D && global_frame > 0 ? std::min<int64_t>(decode_frames, x->ne[2] - frame) : 1;
+                auto in             = ggml_ext_slice(ctx->ggml_ctx, x, 2, frame, frame + count);
+                _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) {
@@ -1316,6 +1320,7 @@ namespace WAN {
                     out_frame = repeated;
                 }
                 out = out == nullptr ? out_frame : ggml_concat(ctx->ggml_ctx, out, out_frame, 2);
+                frame += count;
             }
             return unpatchify(ctx->ggml_ctx, out, patch_size, b);
         }
@@ -1455,7 +1460,7 @@ namespace WAN {
             return gf;
         }
 
-        ggml_cgraph* build_temporal_tile_graph(const sd::Tensor<float>& z_tensor, int chunk_idx) {
+        ggml_cgraph* build_temporal_tile_graph(const sd::Tensor<float>& z_tensor, int chunk_idx, int decode_frames) {
             ggml_cgraph* gf = new_graph_custom(std::max<size_t>(20480, 10240 * z_tensor.shape()[2]));
 
             ae.clear_cache();
@@ -1469,7 +1474,8 @@ namespace WAN {
 
             auto runner_ctx = get_context(gf);
 
-            ggml_tensor* out = ae.decode_tiled_chunk(&runner_ctx, z, chunk_idx);
+            ggml_tensor* out = ae.decode_tiled_chunk(&runner_ctx, z, chunk_idx, 1,
+                                                     version == VERSION_WAN2_2_TI2V ? decode_frames : 1);
 
             for (size_t feat_idx = 0; feat_idx < ae._feat_map.size(); feat_idx++) {
                 ggml_tensor* feat_cache = ae._feat_map[feat_idx];
@@ -1492,8 +1498,9 @@ namespace WAN {
             stateful_config.overlap                 = 0;
             auto plan                               = make_vae_temporal_tile_plan(input.shape()[2], stateful_config);
 
-            LOG_VERBOSE("Wan VAE stateful temporal tiling: tile_frames=%d, total latent frames=%lld, tiles=%d",
+            LOG_VERBOSE("Wan VAE stateful temporal tiling: tile_frames=%d, decode_frames=%d, total latent frames=%lld, tiles=%d",
                         plan.tile_frames,
+                        version == VERSION_WAN2_2_TI2V && !ae.is_2D ? config.decode_frames : 1,
                         (long long)input.shape()[2],
                         (int)plan.tiles.size());
 
@@ -1507,7 +1514,7 @@ namespace WAN {
                             (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 build_temporal_tile_graph(input_tile, static_cast<int>(tile.start), config.decode_frames);
                 };
                 return restore_trailing_singleton_dims(
                     GGMLRunner::compute(get_graph, n_threads, false),