<!-- ## **TeaCache4Cosmos** -->
# TeaCache4Cosmos

[TeaCache](https://github.com/LiewFeng/TeaCache) can speedup [Cosmos](https://github.com/NVIDIA/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](https://github.com/NVIDIA/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:

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

For T2V inference, you can use the following command:

```bash
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:

```bash
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](https://github.com/NVIDIA/Cosmos).
