Commit 40b1788 for stable-diffusion.cpp
commit 40b1788476bbe8fe42fa538ddf7d32408a126b93
Author: losewayy <1134285612@qq.com>
Date: Sat Oct 24 01:09:05 2026 +0800
feat: allow grouped temporal decoding for Wan2.2 TI2V 5B (#2196)
diff --git a/docs/wan.md b/docs/wan.md
index 87a960f..2031de7 100644
--- a/docs/wan.md
+++ b/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
```
diff --git a/examples/common/common.cpp b/examples/common/common.cpp
index d7bc9ba..bd25f90 100644
--- a/examples/common/common.cpp
+++ b/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},
{"",
diff --git a/examples/server/api.md b/examples/server/api.md
index c8fa2fb..0d83214 100644
--- a/examples/server/api.md
+++ b/examples/server/api.md
@@ -873,7 +873,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:
diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp
index 64c1533..1099f4b 100644
--- a/src/core/ggml_extend.cpp
+++ b/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};
}
diff --git a/src/model/vae/vae_tiling.hpp b/src/model/vae/vae_tiling.hpp
index 50cf083..acad046 100644
--- a/src/model/vae/vae_tiling.hpp
+++ b/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);
}
diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp
index 9dd30fd..848427c 100644
--- a/src/model/vae/wan_vae.hpp
+++ b/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),