DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-11 11:26:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
losewayy authored and GitHub committed 2026-10-24 01:09:05 +08:00
1 parent cd7f986058
commit 40b1788476
6 files changed
+30 -17
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/model/vae/vae_tiling.hpp
+7 -3
View File
@@ -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
+16 -9
View File
@@ -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),