{"id":13562911,"url":"https://github.com/tianweiy/DMD2","last_synced_at":"2025-04-03T19:31:51.767Z","repository":{"id":241297205,"uuid":"805164005","full_name":"tianweiy/DMD2","owner":"tianweiy","description":"(NeurIPS 2024 Oral 🔥) Improved Distribution Matching Distillation for Fast Image Synthesis","archived":false,"fork":false,"pushed_at":"2025-03-05T06:04:36.000Z","size":4541,"stargazers_count":703,"open_issues_count":31,"forks_count":39,"subscribers_count":6,"default_branch":"main","last_synced_at":"2025-03-13T22:34:29.343Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"other","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/tianweiy.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE.md","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-05-24T02:30:31.000Z","updated_at":"2025-03-13T12:48:52.000Z","dependencies_parsed_at":null,"dependency_job_id":"05ad9f99-b340-4d16-9b16-b79e7af7beda","html_url":"https://github.com/tianweiy/DMD2","commit_stats":null,"previous_names":["tianweiy/dmd2"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tianweiy%2FDMD2","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tianweiy%2FDMD2/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tianweiy%2FDMD2/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tianweiy%2FDMD2/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/tianweiy","download_url":"https://codeload.github.com/tianweiy/DMD2/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247065309,"owners_count":20877756,"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":[],"created_at":"2024-08-01T13:01:13.321Z","updated_at":"2025-04-03T19:31:51.752Z","avatar_url":"https://github.com/tianweiy.png","language":"Python","funding_links":[],"categories":["Accelerate","Python"],"sub_categories":[],"readme":"# Improved Distribution Matching Distillation for Fast Image Synthesis [[Huggingface Repo](https://huggingface.co/tianweiy/DMD2)][[ComfyUI](https://gist.github.com/comfyanonymous/fcce4ced378f74f4c46026b134faf27a)][[Colab](https://colab.research.google.com/drive/1iGk7IW2WosophOVYpdW_KZGIfYpATOm7?usp=sharing)]\n\nFew-step Text-to-Image Generation.\n\n![image/jpeg](docs/teaser.jpg)\n\n\u003e [**Improved Distribution Matching Distillation for Fast Image Synthesis**](https://tianweiy.github.io/dmd2/),            \n\u003e Tianwei Yin, Michaël Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Frédo Durand, William T. Freeman        \n\u003e *NeurIPS 2024 ([arXiv 2405.14867](https://arxiv.org/abs/2405.14867))*  \n\n## Contact \n\nFeel free to contact us if you have any questions about the paper!\n\nTianwei Yin [tianweiy@mit.edu](mailto:tianweiy@mit.edu)\n\n## Abstract\n\nRecent approaches have shown promises distilling diffusion models into\nefficient one-step generators. Among them, Distribution Matching Distillation\n(DMD) produces one-step generators that match their teacher in distribution,\nwithout enforcing a one-to-one correspondence with the sampling trajectories of\ntheir teachers. However, to ensure stable training, DMD requires an additional\nregression loss computed using a large set of noise-image pairs generated by\nthe teacher with many steps of a deterministic sampler. This is costly for\nlarge-scale text-to-image synthesis and limits the student's quality, tying it\ntoo closely to the teacher's original sampling paths. We introduce DMD2, a set\nof techniques that lift this limitation and improve DMD training. First, we\neliminate the regression loss and the need for expensive dataset construction.\nWe show that the resulting instability is due to the fake critic not estimating\nthe distribution of generated samples accurately and propose a two time-scale\nupdate rule as a remedy. Second, we integrate a GAN loss into the distillation\nprocedure, discriminating between generated samples and real images. This lets\nus train the student model on real data, mitigating the imperfect real score\nestimation from the teacher model, and enhancing quality. Lastly, we modify the\ntraining procedure to enable multi-step sampling. We identify and address the\ntraining-inference input mismatch problem in this setting, by simulating\ninference-time generator samples during training time. Taken together, our\nimprovements set new benchmarks in one-step image generation, with FID scores\nof 1.28 on ImageNet-64x64 and 8.35 on zero-shot COCO 2014, surpassing the\noriginal teacher despite a 500X reduction in inference cost. Further, we show\nour approach can generate megapixel images by distilling SDXL, demonstrating\nexceptional visual quality among few-step methods.\n\n## Environment Setup\n\n```.bash\n# In conda env \nconda create -n dmd2 python=3.8 -y \nconda activate dmd2 \n\npip install --upgrade anyio\npip install -r requirements.txt\npython setup.py  develop\n```\n\n## Inference Example\n\n#### ImageNet \n\n```.bash\npython -m demo.imagenet_example  --checkpoint_path IMAGENET_CKPT_PATH \n```\n\n#### Text-to-Image \n\n```.bash\n# Note: on the demo page, click ``Use Tiny VAE for faster decoding'' to enable much faster speed and lower memory consumption using a Tiny VAE from [madebyollin](https://huggingface.co/madebyollin/taesdxl)\n\n# 4 step (much higher quality than 1 step)\npython -m demo.text_to_image_sdxl --checkpoint_path SDXL_CKPT_PATH --precision float16\n\n# 1 step \npython -m demo.text_to_image_sdxl --num_step 1 --checkpoint_path SDXL_CKPT_PATH --precision float16 --conditioning_timestep 399\n```\n\nWe can also use the standard diffuser pipeline:\n\n#### 4-step UNet generation \n\n```python\nimport torch\nfrom diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler\nfrom huggingface_hub import hf_hub_download\nfrom safetensors.torch import load_file\nbase_model_id = \"stabilityai/stable-diffusion-xl-base-1.0\"\nrepo_name = \"tianweiy/DMD2\"\nckpt_name = \"dmd2_sdxl_4step_unet_fp16.bin\"\n\n# Load model.\nwith torch.device(\"meta\"):\n    unet = UNet2DConditionModel.from_config(base_model_id, subfolder=\"unet\").to(torch.float16)\nstate_dict_path = hf_hub_download(repo_name, ckpt_name)\nunet.load_state_dict(torch.load(state_dict_path), assign=True)\nunet.to(\"cuda\")\n\npipe = DiffusionPipeline.from_pretrained(base_model_id, unet=unet, torch_dtype=torch.float16, variant=\"fp16\").to(\"cuda\")\npipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)\nprompt=\"a photo of a cat\"\n\n# LCMScheduler's default timesteps are different from the one we used for training \nimage=pipe(prompt=prompt, num_inference_steps=4, guidance_scale=0, timesteps=[999, 749, 499, 249]).images[0]\n```\n\n#### 4-step LoRA generation \n\n```python\nimport torch\nfrom diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler\nfrom huggingface_hub import hf_hub_download\nfrom safetensors.torch import load_file\nbase_model_id = \"stabilityai/stable-diffusion-xl-base-1.0\"\nrepo_name = \"tianweiy/DMD2\"\nckpt_name = \"dmd2_sdxl_4step_lora_fp16.safetensors\"\n# Load model.\npipe = DiffusionPipeline.from_pretrained(base_model_id, torch_dtype=torch.float16, variant=\"fp16\").to(\"cuda\")\npipe.load_lora_weights(hf_hub_download(repo_name, ckpt_name))\npipe.fuse_lora(lora_scale=1.0)  # we might want to make the scale smaller for community models\n\npipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)\nprompt=\"a photo of a cat\"\n\n# LCMScheduler's default timesteps are different from the one we used for training \nimage=pipe(prompt=prompt, num_inference_steps=4, guidance_scale=0, timesteps=[999, 749, 499, 249]).images[0]\n```\n\n#### 1-step UNet generation \n\n```python\nimport torch\nfrom diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler\nfrom huggingface_hub import hf_hub_download\nfrom safetensors.torch import load_file\nbase_model_id = \"stabilityai/stable-diffusion-xl-base-1.0\"\nrepo_name = \"tianweiy/DMD2\"\nckpt_name = \"dmd2_sdxl_1step_unet_fp16.bin\"\n# Load model.\nunet = UNet2DConditionModel.from_config(base_model_id, subfolder=\"unet\").to(\"cuda\", torch.float16)\nunet.load_state_dict(torch.load(hf_hub_download(repo_name, ckpt_name), map_location=\"cuda\"))\npipe = DiffusionPipeline.from_pretrained(base_model_id, unet=unet, torch_dtype=torch.float16, variant=\"fp16\").to(\"cuda\")\npipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)\nprompt=\"a photo of a cat\"\nimage=pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0, timesteps=[399]).images[0]\n```\n\n#### 4-step T2I Adapter \n\n```python \nfrom diffusers import StableDiffusionXLAdapterPipeline, T2IAdapter, AutoencoderKL, UNet2DConditionModel, LCMScheduler\nfrom diffusers.utils import load_image, make_image_grid\nfrom controlnet_aux.canny import CannyDetector\nfrom huggingface_hub import hf_hub_download\nimport torch\n\n# load adapter\nadapter = T2IAdapter.from_pretrained(\"TencentARC/t2i-adapter-canny-sdxl-1.0\", torch_dtype=torch.float16, varient=\"fp16\").to(\"cuda\")\n\nvae=AutoencoderKL.from_pretrained(\"madebyollin/sdxl-vae-fp16-fix\", torch_dtype=torch.float16)\n\nbase_model_id = \"stabilityai/stable-diffusion-xl-base-1.0\"\nrepo_name = \"tianweiy/DMD2\"\nckpt_name = \"dmd2_sdxl_4step_unet_fp16.bin\"\n# Load model.\nunet = UNet2DConditionModel.from_config(base_model_id, subfolder=\"unet\").to(\"cuda\", torch.float16)\nunet.load_state_dict(torch.load(hf_hub_download(repo_name, ckpt_name), map_location=\"cuda\"))\n\npipe = StableDiffusionXLAdapterPipeline.from_pretrained(\n    base_model_id, unet=unet, vae=vae, adapter=adapter, torch_dtype=torch.float16, variant=\"fp16\", \n).to(\"cuda\")\npipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)\npipe.enable_xformers_memory_efficient_attention()\n\ncanny_detector = CannyDetector()\n\nurl = \"https://huggingface.co/Adapter/t2iadapter/resolve/main/figs_SDXLV1.0/org_canny.jpg\"\nimage = load_image(url)\n\n# Detect the canny map in low resolution to avoid high-frequency details\nimage = canny_detector(image, detect_resolution=384, image_resolution=1024)#.resize((1024, 1024))\n\nprompt = \"Mystical fairy in real, magic, 4k picture, high quality\"\n\ngen_images = pipe(\n  prompt=prompt,\n  image=image,\n  num_inference_steps=4,\n  guidance_scale=0, \n  adapter_conditioning_scale=0.8, \n  adapter_conditioning_factor=0.5,\n  timesteps=[999, 749, 499, 249]\n).images[0]\ngen_images.save('out_canny.png')\n```\n\nPretrained models can be found in [ImageNet](experiments/imagenet/README.md) and [SDXL](experiments/sdxl/README.md). \n\n## Training and Evaluation \n\n### ImageNet-64x64 \n\nPlease refer to [ImageNet-64x64](experiments/imagenet/README.md) for details.\n\n### SDXL\n\nPlease refer to [SDXL](experiments/sdxl/README.md) for details.\n\n### SDv1.5 \n\nPlease refer to [SDv1.5](experiments/sdv1.5/README.md) for details.\n\n## License\n\nImproved Distribution Matching Distillation is released under [Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License](LICENSE.md).\n\n## Known Issues \n\n- [ ] Current FSDP for SDXL training is really slow; help is greatly appreciated!\n- [ ] Current LORA training is actually slower than the full finetuning and takes the same amount of memory; help is greatly appreciated!\n\n\n## Citation \n\nIf you find DMD2 useful or relevant to your research, please kindly cite our papers:\n\n```bib\n@inproceedings{yin2024improved,\n    title={Improved Distribution Matching Distillation for Fast Image Synthesis},\n    author={Yin, Tianwei and Gharbi, Micha{\\\"e}l and Park, Taesung and Zhang, Richard and Shechtman, Eli and Durand, Fredo and Freeman, William T},\n    booktitle={NeurIPS},\n    year={2024}\n}\n\n@inproceedings{yin2024onestep,\n    title={One-step Diffusion with Distribution Matching Distillation},\n    author={Yin, Tianwei and Gharbi, Micha{\\\"e}l and Zhang, Richard and Shechtman, Eli and Durand, Fr{\\'e}do and Freeman, William T and Park, Taesung},\n    booktitle={CVPR},\n    year={2024}\n}\n```\n\n## Third-part Code\n\n[EDM](https://github.com/NVlabs/edm/tree/main) for [dnnlib](dnnlib), [torch_utils](torch_utils) and [edm](third_party/edm) folders.\n\n## Acknowledgments \n\nThis work was done while Tianwei Yin was a full-time student at MIT. It was developed based on our reimplementation of the original DMD paper. This work was supported by the National Science Foundation under Cooperative Agreement PHY-2019786 (The NSF AI Institute for Artificial Intelligence and Fundamental Interactions, http://iaifi.org/), by NSF Grant 2105819, by NSF CISE award 1955864, and by funding from Google, GIST, Amazon, and Quanta Computer.\n\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ftianweiy%2FDMD2","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Ftianweiy%2FDMD2","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ftianweiy%2FDMD2/lists"}