feat: add --vae_cpu flag for improved VRAM optimization on consumer GPUsMovie elves #535
Conversation
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
Added environment.yml for project dependencies and wok/go.sh script. Updated .gitignore to exclude symlinks in wok directory. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
Add --vae_cpu argument to enable VAE offloading for consumer GPUs with limited VRAM. When enabled, VAE initializes on CPU and moves to GPU only when needed for encoding/decoding operations. Key changes: - Add --vae_cpu argument to generate.py (mirrors --t5_cpu pattern) - Update all 4 pipelines (T2V, I2V, FLF2V, VACE) with conditional VAE offloading - Fix DiT offloading to free VRAM before T5 loading when offload_model=True - Handle VAE scale tensors (mean/std) during device transfers Benefits: - Saves ~100-200MB VRAM without performance degradation - Enables T2V-1.3B on more consumer GPUs (tested on 11.49GB GPU) - Backward compatible (default=False) - Consistent with existing --t5_cpu flag Test results on 11.49 GiB VRAM GPU: - Baseline: OOM (needed 80MB, only 85MB free) - With --vae_cpu: Success - With --t5_cpu: Success - With both flags: Success (maximum VRAM savings) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
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.
Description
Problem
Users with consumer-grade GPUs (like RTX 4090 with 11.49 GB VRAM) encounter OOM errors when running the T2V-1.3B model even with existing optimization flags (
--offload_model True --t5_cpu). The OOM occurs because the VAE remains on GPU throughout the entire generation pipeline despite only being needed briefly for encoding/decoding.Solution
This PR adds a
--vae_cpuflag that works similarly to the existing--t5_cpuflag. When enabled:Implementation Details
--vae_cpuargument togenerate.py(mirrors--t5_cpupattern)offload_model=Trueandt5_cpu=False, DiT now offloads before T5 loads to prevent OOMmeanandstdtensors move with the model