diff --git a/docs/wan.md b/docs/wan.md index 87a960f3..2031de7d 100644 --- a/docs/wan.md +++ b/docs/wan.md @@ -184,6 +184,8 @@ Notes: +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 d7bc9bab..bd25f905 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 c8fa2fbc..0d83214e 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 64c15332..1099f4b0 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -554,9 +554,9 @@ std::vector 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 50cf0835..acad0462 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 9dd30fdf..848427ca 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(blocks["decoder"]); auto conv2 = std::dynamic_pointer_cast(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(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(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& z_tensor, int chunk_idx) { + ggml_cgraph* build_temporal_tile_graph(const sd::Tensor& z_tensor, int chunk_idx, int decode_frames) { ggml_cgraph* gf = new_graph_custom(std::max(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(tile.start)); + return build_temporal_tile_graph(input_tile, static_cast(tile.start), config.decode_frames); }; return restore_trailing_singleton_dims( GGMLRunner::compute(get_graph, n_threads, false),