fix(wan2.1): causal VAE decode — T latents must yield 1+(T−1)·4 frames, not 4·T - #38
Open
xocialize wants to merge 1 commit into
Open
fix(wan2.1): causal VAE decode — T latents must yield 1+(T−1)·4 frames, not 4·T#38xocialize wants to merge 1 commit into
xocialize wants to merge 1 commit into
Conversation
WanVAE.decode ran the decoder over the whole sequence, where upsample3d
temporally doubles every frame: T latents -> 4*T output frames. The
reference Wan2.1 decoder is causal-chunked: the first latent frame is not
temporally upsampled ('Rep' first-chunk sentinel skips time_conv), giving
1+(T-1)*4 frames. Beyond the wrong count, the whole-sequence path produces
different boundary frames (up to 0.21 abs in [-1,1] terms, decaying over
~8 frames) and a +3-frame phase shift versus reference output.
- Resample: add cached upsample3d branch with the 'Rep' sentinel,
mirroring the existing cached downsample3d branch
- Decoder3d: accept feat_cache/feat_idx (mirror of Encoder3d)
- WanVAE.decode: chunked per-latent-frame decode with
_count_decoder_cache_slots (mirror of the encode path); memory is now
flat in T as a side effect
- decode_tiled: per-tile chunked decode; causal_temporal=True (the tiling
module already documents this mapping)
- tests: frame-count and encode/decode roundtrip coverage
Verified against the reference PyTorch Wan2.1 VAE (Wan2.1_VAE.pth, fp32
CPU): max_abs 2.6e-6 (T=1) / 6.3e-6 (T=5) on every output frame.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Bug
WanVAE.decode(wan_2) runsDecoder3dover the whole latent sequence, and theupsample3dpath temporally doubles every frame — soTlatent frames decode to4·Toutput frames. The reference Wan2.1 decoder (Wan-AI/Wan2.1vae.py, also diffusersAutoencoderKLWan) decodes causally one latent frame at a time: the first chunk sets the'Rep'sentinel and skipstime_conv, so the first latent frame is not temporally upsampled —T → 1 + (T−1)·4frames (65 for 17 latents, not 68).Beyond the wrong frame count, the whole-sequence path produces different pixels at the start: comparing against the reference decoder on identical latents, boundary frames differ by up to 0.21 (abs, in [-1,1] range), decaying over ~8 frames, with a +3-frame phase shift across the rest (~190 ms at 16 fps — breaks A/V sync for downstream use).
Interestingly
tiling.pyalready documents both mappings and hascausal_temporalsupport —decode_tiledjust passescausal_temporal=Falsefor Wan2.1, which contradicts the reference.Fix
Mirrors the existing encode path (which already does cached chunked processing correctly):
Resample: cachedupsample3dbranch with the'Rep'first-chunk sentinel, structured like the existing cacheddownsample3dbranchDecoder3d.__call__: acceptsfeat_cache/feat_idx(mirror ofEncoder3d)WanVAE.decode: per-latent-frame chunked decode +_count_decoder_cache_slots(mirror of_count_encoder_cache_slots); peak memory is now flat in T as a side effectdecode_tiled: per-tile chunked decode,causal_temporal=TrueT=1→1, 2→5, 5→17) and encode/decode roundtripWhole-sequence behavior is preserved when
feat_cache=Noneis passed to the decoder directly.Verification
Against the reference PyTorch Wan2.1 VAE (
Wan2.1_VAE.pth, fp32, CPU), identical latents:tests/pass (test_wan_vae.py80/80 including the new cases; full suite 387 passed — the one failure,test_wan_tiling.py::TestNonCausalTemporal::test_causal_vs_noncausal_output_size, fails identically on currentmainand is unrelated).Found while porting SCAIL-2 (Wan2.1-I2V fork) on top of the wan_2 substrate — happy to adjust scope/style.
🤖 Generated with Claude Code