Skip to content

fix: stream Wan VAE decode output to CPU#2042

Open
dexhunter wants to merge 1 commit into
kijai:mainfrom
dexhunter:dex/stream-vae-decode-cpu
Open

fix: stream Wan VAE decode output to CPU#2042
dexhunter wants to merge 1 commit into
kijai:mainfrom
dexhunter:dex/stream-vae-decode-cpu

Conversation

@dexhunter

Copy link
Copy Markdown

Summary

  • stream standard non-tiled Wan 2.1 VAE output into a preallocated CPU float32 tensor
  • remove the growing GPU temporal concatenation used by the ComfyUI decode node
  • preserve existing GPU-return, tiled, end-frame, 38-channel, and external VAE paths

Fixes #1826.

Why

VideoVAE_.decode() currently grows the complete decoded tensor with
torch.cat() at every latent timestep. The node then immediately transfers the
completed result to CPU float32. For long videos, the repeated GPU accumulation
can exhaust VRAM after sampling has already completed.

This change allocates the final host tensor once and copies each decoded temporal
chunk directly into its final slice. The decoder cache sequence and default
GPU-return behavior are unchanged.

Validation

Autoresearch trace

The source-hashed evaluator uses current main
088128b224242e110d3906c6750e9a3a348a659b, verifies the downloaded
Wan2_1_VAE_bf16.safetensors SHA-256, and runs the public
WanVideoVAE.decode() path with fixed-seed latents.

Output frames Current main peak CUDA allocation Candidate Temporal cat output Wall time
81 3,389,870,592 B 3,240,571,392 B 1,585,152,000 B -> 0 B 4,569 ms -> 3,774 ms
321 3,860,316,672 B 3,268,649,472 B 24,035,328,000 B -> 0 B 15,793 ms -> 14,905 ms

For both comparison sizes:

  • streamed output is byte-identical to current main (max_abs_error = 0)
  • the candidate's default GPU-return path is also byte-identical and retains its existing behavior
  • node dispatch preserves legacy/external decode signatures

The candidate also completed the issue-sized 3,497-frame decode with shape
[1, 3, 3497, 480, 640], zero temporal cats, and a 3,633,374,208-byte peak CUDA
allocation delta. The final CPU output is 12.89 GB and process RSS reached 19.00
GB, so this addresses VRAM accumulation rather than total host-memory usage.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

generating 2 minute video infinitetalk wan 2.1 - crash at wanvideo decode node wan_video_vae.py OOM error

1 participant