From cc76eb1ffe42ee74408f4d95e9d02620133e891e Mon Sep 17 00:00:00 2001 From: olyasir Date: Fri, 17 Jul 2026 10:48:46 +0300 Subject: [PATCH 1/2] fix: match reference first-chunk semantics in Wan VAE temporal upsample In Resample::forward (upsample3d), the first latent-frame chunk skipped time_conv entirely and never seeded the temporal feat cache, so the second chunk ran time_conv against zero padding instead of chunk-0 context. The reference implementation (Wan vae2_2.py, "Rep" cache semantics) runs time_conv with causal zero padding on chunk 0, doubles the frames, trims the first one, and stores the input tail in the cache for the next chunk. Effect on Wan2.2 VAE decode vs the PyTorch reference (5-frame 480x832, CPU backend, f16 weights): cosine 0.9959 / 27.2 dB PSNR before, cosine 1.000000 / 79.0 dB after. Output was visually near-identical, so this only shows up in numeric parity testing. Encode and TAEHV paths are unaffected. --- src/wan.hpp | 39 ++++++++++++++++++++------------------- 1 file changed, 20 insertions(+), 19 deletions(-) diff --git a/src/wan.hpp b/src/wan.hpp index 68f020e25..5f69b6994 100644 --- a/src/wan.hpp +++ b/src/wan.hpp @@ -166,32 +166,33 @@ namespace WAN { if (feat_cache.size() > 0) { int idx = feat_idx; feat_idx += 1; - if (chunk_idx == 0) { - // feat_cache[idx] == nullptr, pass + auto time_conv = std::dynamic_pointer_cast(blocks["time_conv"]); + + // mirrors the reference (vae2_2.py Resample.forward): + // cache_x = x[:, :, -CACHE_T:].clone() + // if feat_cache[idx] == "Rep": x = time_conv(x) # chunk 0: causal zero pad + // else: ; x = time_conv(x, feat_cache[idx]) + // feat_cache[idx] = cache_x + auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]); + if (chunk_idx == 0) { // "Rep" + x = time_conv->forward(ctx, x); } else { - auto time_conv = std::dynamic_pointer_cast(blocks["time_conv"]); - - auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]); - if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) { // chunk_idx >= 2 + if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) { // cache last frame of last two chunk cache_x = ggml_concat(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]), cache_x, 2); } - if (chunk_idx == 1 && cache_x->ne[2] < 2) { // Rep - cache_x = ggml_pad_ext(ctx->ggml_ctx, cache_x, 0, 0, 0, 0, (int)cache_x->ne[2], 0, 0, 0); - // aka cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device),cache_x],dim=2) - } - if (chunk_idx == 1) { - x = time_conv->forward(ctx, x); - } else { - x = time_conv->forward(ctx, x, feat_cache[idx]); - } - feat_cache[idx] = cache_x; - x = ggml_reshape_4d(ctx->ggml_ctx, x, w * h, t, c, 2); // (2, c, t, h*w) - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 3, 1, 2)); // (c, t, 2, h*w) - x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, 2 * t, c); // (c, t*2, h, w) + x = time_conv->forward(ctx, x, feat_cache[idx]); + } + feat_cache[idx] = cache_x; + x = ggml_reshape_4d(ctx->ggml_ctx, x, w * h, t, c, 2); // (2, c, t, h*w) + x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 3, 1, 2)); // (c, t, 2, h*w) + x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, 2 * t, c); // (c, t*2, h, w) + if (chunk_idx == 0) { + // reference: if first_chunk: x = x[:, :, 1:] + x = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]); } } } From e17495a6d5af1b8026af8c4c2b2c18b6b245bb3e Mon Sep 17 00:00:00 2001 From: olyasir Date: Fri, 17 Jul 2026 13:49:07 +0300 Subject: [PATCH 2/2] fix: assert time_conv block exists before use in upsample3d path --- src/wan.hpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/wan.hpp b/src/wan.hpp index 5f69b6994..a49244533 100644 --- a/src/wan.hpp +++ b/src/wan.hpp @@ -166,7 +166,10 @@ namespace WAN { if (feat_cache.size() > 0) { int idx = feat_idx; feat_idx += 1; - auto time_conv = std::dynamic_pointer_cast(blocks["time_conv"]); + auto time_conv_it = blocks.find("time_conv"); + GGML_ASSERT(time_conv_it != blocks.end()); + auto time_conv = std::dynamic_pointer_cast(time_conv_it->second); + GGML_ASSERT(time_conv != nullptr); // mirrors the reference (vae2_2.py Resample.forward): // cache_x = x[:, :, -CACHE_T:].clone()