TeaCache/TeaCache4Cosmos/README.md

TeaCache4Cosmos

TeaCache can speedup Cosmos 2.0x without much visual quality degradation, in a training-free manner. The following video shows the results generated by TeaCache-Cosmos with various rel_l1_thresh values: 0 (original), 0.3 (1.4x speedup), and 0.4(2.0x speedup).

https://github.com/user-attachments/assets/28570179-0f22-42ee-8958-88bb48d209b4

https://github.com/user-attachments/assets/21341bd4-c0d5-4b5a-8b7d-cb2103897d2c

📈 Inference Latency Comparisons on a Single H800 GPU

Cosmos-t2v	TeaCache (0.3)	TeaCache (0.4)
~449 s	~327 s	~227 s

Cosmos-i2v	TeaCache (0.3)	TeaCache (0.4)
~453 s	~330 s	~229 s

Usage

Follow Cosmos to clone the repo and finish the installation, then you can modify the rel_l1_thresh to obtain your desired trade-off between latency and visul quality, and change the checkpoint_dir, prompt, input_image_or_video_path to customize your video.

You need to copy the running script to the Cosmos project folder:

cd TeaCache4Cosmos
cp *.py /path/to/Cosmos/

For T2V inference, you can use the following command:

cd /path/to/Cosmos/

python3 teacache_sample_video_t2v.py \
    --rel_l1_thresh 0.3 \
    --checkpoint_dir checkpoints \
    --diffusion_transformer_dir Cosmos-1.0-Diffusion-7B-Text2World \
    --prompt Inside the cozy ambiance of a bustling coffee house, a young woman with wavy chestnut hair and wearing a rust-colored cozy sweater stands amid the chatter and clinking of cups. She smiles warmly at the camera, her green eyes glinting with joy and subtle hints of laughter. The camera frames her elegantly, emphasizing the soft glow of the lighting on her smooth, clear skin and the detailed textures of her woolen attire. Her genuine smile is the centerpiece of the shot, showcasing her enjoyment in the quaint café setting, with steaming mugs and blurred patrons in the background. \
    --disable_prompt_upsampler \
    --offload_prompt_upsampler

For I2V inference, you can use the following command:

cd /path/to/Cosmos/

python3 teacache_sample_video_i2v.py \
    --rel_l1_thresh 0.4 \
    --checkpoint_dir checkpoints \
    --diffusion_transformer_dir Cosmos-1.0-Diffusion-7B-Video2World \
    --prompt "A girl gazed at the camera and smiled, her hair drifting in the wind." \
    --input_image_or_video_path image/path \
    --num_input_frames 1 \
    --disable_prompt_upsampler \
    --offload_prompt_upsampler

Citation

If you find TeaCache is useful in your research or applications, please consider giving us a star 🌟 and citing it by the following BibTeX entry.

@article{liu2024timestep,
  title={Timestep Embedding Tells: It's Time to Cache for Video Diffusion Model},
  author={Liu, Feng and Zhang, Shiwei and Wang, Xiaofeng and Wei, Yujie and Qiu, Haonan and Zhao, Yuzhong and Zhang, Yingya and Ye, Qixiang and Wan, Fang},
  journal={arXiv preprint arXiv:2411.19108},
  year={2024}
}

Acknowledgements

We would like to thank the contributors to the Cosmos.