{"id":18264554,"url":"https://github.com/modeltc/tfmq-dm","last_synced_at":"2025-04-04T21:30:41.485Z","repository":{"id":226733935,"uuid":"769502884","full_name":"ModelTC/TFMQ-DM","owner":"ModelTC","description":"[CVPR 2024 Highlight] This is the official PyTorch implementation of \"TFMQ-DM: Temporal Feature Maintenance Quantization for Diffusion Models\".","archived":false,"fork":false,"pushed_at":"2024-08-01T07:45:51.000Z","size":123417,"stargazers_count":61,"open_issues_count":0,"forks_count":4,"subscribers_count":8,"default_branch":"main","last_synced_at":"2025-03-20T19:15:18.965Z","etag":null,"topics":["cvpr","cvpr2024","ddim","diffusion-models","highlight","ldm","post-training-quantization","quantization","stable-diffusion"],"latest_commit_sha":null,"homepage":"https://modeltc.github.io/TFMQ-DM/","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/ModelTC.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-03-09T09:04:18.000Z","updated_at":"2025-03-16T02:56:35.000Z","dependencies_parsed_at":"2024-08-01T09:21:49.343Z","dependency_job_id":null,"html_url":"https://github.com/ModelTC/TFMQ-DM","commit_stats":null,"previous_names":["modeltc/tfmq-dm"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ModelTC%2FTFMQ-DM","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ModelTC%2FTFMQ-DM/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ModelTC%2FTFMQ-DM/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ModelTC%2FTFMQ-DM/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/ModelTC","download_url":"https://codeload.github.com/ModelTC/TFMQ-DM/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247251973,"owners_count":20908601,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["cvpr","cvpr2024","ddim","diffusion-models","highlight","ldm","post-training-quantization","quantization","stable-diffusion"],"created_at":"2024-11-05T11:15:05.403Z","updated_at":"2025-04-04T21:30:36.476Z","avatar_url":"https://github.com/ModelTC.png","language":"Jupyter Notebook","funding_links":[],"categories":[],"sub_categories":[],"readme":"# TFMQ-DM: Temporal Feature Maintenance Quantization for Diffusion Models\n[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) \n[![arXiv](https://img.shields.io/badge/TFMQ--DM-2311.16503-b31b1b)](https://arxiv.org/abs/2311.16503)\n[![GitHub Stars](https://img.shields.io/github/stars/ModelTC/TFMQ-DM.svg?style=social\u0026label=Star\u0026maxAge=60)](https://github.com/ModelTC/TFMQ-DM)\n\n[[paper](https://arxiv.org/abs/2311.16503)][[slides](assets/slides.pdf)][[poster](assets/poster.pdf)][[project page](https://modeltc.github.io/TFMQ-DM/)]\n\n[Yushi Huang*](https://github.com/Harahan), [Ruihao Gong*](https://xhplus.github.io/), [Jing Liu](https://jing-liu.com/), [Tianlong Chen](https://tianlong-chen.github.io/), [Xianglong Liu📧](https://xlliu-beihang.github.io/)\n\n(* denotes equal contribution, 📧 denotes corresponding author.)\n\nThis is the official implementation of our paper [TFMQ-DM](https://arxiv.org/abs/2311.16503), a  novel training-free framework that achieves a new state-of-the-art result in PTQ of diffusion models, especially under 4-bit weight quantization, and significantly accelerates quantization time. For some hard tasks, e.g., CelebA-HQ, our method reduces the FID score by 6.7.\n\u003cdiv align=center\u003e\n\t\u003cfigure class=\"second\"\u003e\n\t    \u003cimg src=\"./img/sd.png\" width=\"460\"/\u003e\u003cimg src=\"./img/ldm.png\" width=\"282\"/\u003e\n\t\u003c/figure\u003e\n\u003c/div\u003e\n\n(Left) Images generated by w4a8 quantized and full-precision Stable Diffusion with MS-COCO captions. (Right) Random samples from w4a8 quantized and full-precision LDM-4 on CelebA-HQ. These results show that our TFMQ-DM outperforms the previous PTQ methods.  More qualitative and quantitative results can be found in our paper.\n\n## News\n\n* **Apr 5, 2024**: 🌟 Our paper has been selected as a **Highlight Poster** at CVPR 2024 **(top 2.8%)**! 🎉 Cheers!\n\n* **Mar 31, 2024**: 🔥 We release our Python code for quantizing all the models presented in our paper. Have a try!\n\n* **Feb 27, 2024**: 🌟 Our paper has been accepted by CVPR 2024! 🎉 Cheers!\n\n## Overview\n\n![overview](./img/overview.png)\n\nThe Diffusion model, a prevalent framework for image generation, encounters significant challenges in terms of broad applicability due to its extended inference times and substantial memory requirements. Efficient Post-training Quantization (PTQ) is pivotal for addressing these issues in traditional models. Different from traditional models, diffusion models heavily depend on the time-step $t$ to achieve satisfactory multi-round denoising. Usually, $t$ from the finite set \\{ $1, \\ldots, T$ \\} is encoded to a temporal feature by a few modules irrespective of the sampling data. However, existing PTQ methods do not optimize these modules separately. They adopt inappropriate reconstruction targets and complex calibration methods, resulting in a severe disturbance of the temporal feature and denoising trajectory, as well as a low compression efficiency. To solve these, we propose a Temporal Feature Maintenance Quantization (TFMQ) framework building upon a Temporal Information Block which is just related to the time-step $t$ and unrelated to the sampling data. Powered by the pioneering block design, we devise temporal information aware reconstruction (TIAR) and finite set calibration (FSC) to align the full-precision temporal features in a limited time. Equipped with the framework, we can maintain the most temporal information and ensure the end-to-end generation quality. Extensive experiments on various datasets and diffusion models prove our state-of-the-art results. \n\n## Quick Started\n\nAfter cloning the repository, you can follow these steps to complete the model's quantization inference process.\n\n### Requirements\n\n```bash\nconda env create -f ./stable-diffusion/environment.yaml\nconda activate ldm\npip install -r requirements.txt\n```\n\nIf you encounter errors while installing the packages listed in requirements.txt, you can try installing each Python package individually using the pip command.\n\nBefore quantization, you need to download the pre-trained weights (preliminary checkpoints):\n\n```bash\n# ----------- DDIM -----------\n# The model execution script will automatically download the required checkpoints during runtime.\n\n# ----------- LDM ------------\ncd ./stable-diffusion/\nsh ./scripts/download_first_stages.sh\nsh ./scripts/download_models.sh\nmkdir -p models/ldm/cin256-v2/\nwget -O models/ldm/cin256-v2/model.ckpt https://ommer-lab.com/files/latent-diffusion/nitro/cin/model.ckpt\ncd ../\n\n# ----- Stable Diffusion -----\nwget https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt\nmkdir -p ./stable-diffusion/models/ldm/stable-diffusion-v1/\nmv sd-v1-4.ckpt ./stable-diffusion/models/ldm/stable-diffusion-v1/sd-v1-4.ckpt\n```\n\n### Quantization\n\nIn this part, we will first generate some calibration data before quantizing. Alternatively, you can generate calibration data separately by adding `torch.save` in the script, and quantize when needed by adding `torch.load` and commenting out the code for generating calibration data.\n\nIf you want to quantize your diffusion models on multiple GPUs, add `--multi_gpu` to the corresponding command, except for `DDIM`. Additionally, you can remove `--use_aq`，`--aq 8` to cancel activation quantization.\n\nAdditionally, before quantizing Stable Diffusion, you still need to prepare for some prompts. For example, you can download [MS-COCO](https://cocodataset.org/#download), and use the path of `captions_train*.json` as `\u003cPATH/TO/LOAD/DATA\u003e` in the following command. We would use 128 prompts within `captions_train*.json` as a part of calibration data.\n\n```bash\n# ----------- DDIM -----------\npython sample_diffusion_ddim.py --config ddim/configs/cifar10.yml --timesteps 100 --eta 0 --skip_type quad --wq \u003c4 OR 8\u003e --ptq --aq 8 -l \u003cPATH/TO/SAVE/LOG\u003e --cali --use_aq --cali_save_path \u003cPATH/TO/SAVE/QUANTIZED/MODEL\u003e --interval_length 5\n\n# ----------- LDM ------------\n# LSUN-Bedrooms\npython sample_diffusion_ldm.py -r ./stable-diffusion/models/ldm/lsun_beds256/model.ckpt -c 200 -e 1.0 --seed 40 --wq \u003c4 OR 8\u003e --ptq --aq 8 -l \u003cPATH/TO/SAVE/LOG\u003e --cali --use_aq --cali_save_path \u003cPATH/TO/SAVE/QUANTIZED/MODEL\u003e --interval_length 10\n# LSUN-Churches\npython sample_diffusion_ldm.py -r ./stable-diffusion/models/ldm/lsun_churches256/model.ckpt -c 400 -e 0.0 --seed 40 --wq \u003c4 OR 8\u003e --ptq --aq 8 -l \u003cPATH/TO/SAVE/LOG\u003e --cali --use_aq --cali_save_path \u003cPATH/TO/SAVE/QUANTIZED/MODEL\u003e --interval_length 25\n# CelebA-HQ\npython sample_diffusion_ldm.py -r ./stable-diffusion/models/ldm/celeba256/model.ckpt -c 200 -e 0.0 --seed 40 --wq \u003c4 OR 8\u003e --ptq --aq 8 -l \u003cPATH/TO/SAVE/LOG\u003e --cali --use_aq --cali_save_path \u003cPATH/TO/SAVE/QUANTIZED/MODEL\u003e --interval_length 10\n# FFHQ\npython sample_diffusion_ldm.py -r ./stable-diffusion/models/ldm/ffhq256/model.ckpt -c 200 -e 1.0 --seed 40 --wq \u003c4 OR 8\u003e --ptq --aq 8 -l \u003cPATH/TO/SAVE/LOG\u003e --cali --use_aq --cali_save_path \u003cPATH/TO/SAVE/QUANTIZED/MODEL\u003e --interval_length 10\n# ImageNet\npython latent_imagenet_diffusion.py -e 0.0 --ddim_steps 20 --scale 3.0 --seed 40 --wq \u003c4 OR 8\u003e --ptq --aq 8 --outdir \u003cPATH/TO/SAVE/LOG\u003e --cali --use_aq --cali_save_path \u003cPATH/TO/SAVE/QUANTIZED/MODEL\u003e\n\n# ----- Stable Diffusion -----\npython txt2img.py --plms --no_grad_ckpt --ddim_steps 50 --seed 40 --cond --wq \u003c4 OR 8\u003e --ptq --aq 8 --outdir \u003cPATH/TO/SAVE/LOG\u003e --cali --skip_grid --use_aq --ckpt ./stable-diffusion/models/ldm/stable-diffusion-v1/sd-v1-4.ckpt --config stable-diffusion/configs/stable-diffusion/v1-inference.yaml --data_path \u003cPATH/TO/LOAD/DATA\u003e --cali_save_path \u003cPATH/TO/SAVE/QUANTIZED/MODEL\u003e\n```\n\n### Inference\n\nAfter the quantization process, you can generate images you like. You can remove `--use_aq`，`--aq 8` to cancel activation quantization.\n\n```bash\n# ----------- DDIM -----------\npython sample_diffusion_ddim.py --config ddim/configs/cifar10.yml --timesteps 100 --eta 0 --skip_type quad --wq \u003c4 OR 8\u003e --ptq --aq 8 -l \u003cPATH/TO/SAVE/RESULT\u003e --use_aq --cali_ckpt \u003cPATH/TO/LOAD/QUANTIZED/MODEL\u003e --max_images 128\n\n# ----------- LDM ------------\n# LSUN-Bedrooms\npython sample_diffusion_ldm.py -r ./stable-diffusion/models/ldm/lsun_beds256/model.ckpt -c 200 -e 1.0 --seed 40 --wq \u003c4 OR 8\u003e --ptq --aq 8 -l \u003cPATH/TO/SAVE/RESULT\u003e --use_aq --cali_ckpt \u003cPATH/TO/LOAD/QUANTIZED/MODEL\u003e -n 5 --batch_size 5\n# LSUN-Churches\npython sample_diffusion_ldm.py -r ./stable-diffusion/models/ldm/lsun_churches256/model.ckpt -c 400 -e 0.0 --seed 40 --wq \u003c4 OR 8\u003e --ptq --aq 8 -l \u003cPATH/TO/SAVE/RESULT\u003e --use_aq --cali_ckpt \u003cPATH/TO/LOAD/QUANTIZED/MODEL\u003e -n 5 --batch_size 5\n# CelebA-HQ\npython sample_diffusion_ldm.py -r ./stable-diffusion/models/ldm/celeba256/model.ckpt -c 200 -e 0.0 --seed 40 --wq \u003c4 OR 8\u003e --ptq --aq 8 -l \u003cPATH/TO/SAVE/RESULT\u003e --use_aq --cali_ckpt \u003cPATH/TO/LOAD/QUANTIZED/MODEL\u003e -n 5 --batch_size 5\n# FFHQ\npython sample_diffusion_ldm.py -r ./stable-diffusion/models/ldm/ffhq256/model.ckpt -c 200 -e 1.0 --seed 40 --wq \u003c4 OR 8\u003e --ptq --aq 8 -l \u003cPATH/TO/SAVE/RESULT\u003e --use_aq --cali_ckpt \u003cPATH/TO/LOAD/QUANTIZED/MODEL\u003e -n 5 --batch_size 5\n# ImageNet\npython latent_imagenet_diffusion.py -e 0.0 --ddim_steps 20 --scale 3.0 --seed 40 --wq \u003c4 OR 8\u003e --ptq --aq 8 --outdir \u003cPATH/TO/SAVE/RESULT\u003e  --use_aq --cali_ckpt \u003cPATH/TO/LOAD/QUANTIZED/MODEL\u003e --n_sample_per_class 2 --classes \u003cCLASSES. e.g. 7,489,765\u003e\n\n# ----- Stable Diffusion -----\npython txt2img.py --prompt \u003cPROMPT. e.g. \"A white dog.\"\u003e --plms --no_grad_ckpt --ddim_steps 50 --seed 40 --cond --n_iter 1 --n_samples 1 --wq \u003c4 OR 8\u003e --ptq --aq 8 --skip_grid --outdir \u003cPATH/TO/SAVE/RESULT\u003e --use_aq --ckpt ./stable-diffusion/models/ldm/stable-diffusion-v1/sd-v1-4.ckpt --config stable-diffusion/configs/stable-diffusion/v1-inference.yaml --cali_ckpt \u003cPATH/TO/LOAD/QUANTIZED/MODEL\u003e\n```\n\n## Acknowledgments\n\nOur code was developed based on [ddim](https://github.com/ermongroup/ddim), [latent-diffusion](https://github.com/CompVis/latent-diffusion) and [stable-diffusion](https://github.com/CompVis/stable-diffusion). We referred to [BRECQ](https://github.com/yhhhli/BRECQ) and [Q-Diffusion](https://github.com/Xiuyu-Li/q-diffusion) for the blockwise calibration implementation.\n\nWe thank [OpenVINO](https://github.com/openvinotoolkit/openvino) for providing the framework to deploy our quantized model and measure acceleration. Moreover, we also thank [torch-fidelity](https://github.com/toshas/torch-fidelity), [guided-diffusion](https://github.com/openai/guided-diffusion), and [clip-score](https://github.com/Taited/clip-score) for IS, sFID, FID and CLIP score computation.\n\n## Citation\n\nIf you find our TFMQ-DM useful or relevant to your research, please kindly cite our paper:\n\n```\n@InProceedings{Huang_2024_CVPR,\n    author    = {Huang, Yushi and Gong, Ruihao and Liu, Jing and Chen, Tianlong and Liu, Xianglong},\n    title     = {TFMQ-DM: Temporal Feature Maintenance Quantization for Diffusion Models},\n    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},\n    month     = {June},\n    year      = {2024},\n    pages     = {7362-7371}\n}\n\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmodeltc%2Ftfmq-dm","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fmodeltc%2Ftfmq-dm","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmodeltc%2Ftfmq-dm/lists"}