SYZhang0805/DisBack

Distribution Backtracking Builds A Faster Convergence Trajectory for One-step Diffusion Distillation,
Shengyuan Zhang¹, Ling Yang², Zejian Li*¹, An Zhao¹, Chenye Meng¹, Changyuan Yang³, Guang Yang³, Zhiyuan Yang³, Lingyun Sun^1
1Zhejiang University ²Peking University ³Alibaba Group*

Contact

Feel free to contact us if you have any questions about the paper!

Shengyuan Zhang zhangshengyuan@zju.edu.cn

Usage

For one-step text-to-image generation, DisBack can use the standard diffuser pipeline:

import torch
from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler
from huggingface_hub import hf_hub_download

base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
repo_name = "SYZhang0805/DisBack"
ckpt_name = "SDXL_DisBack.bin"

unet = UNet2DConditionModel.from_config(base_model_id, subfolder="unet").to("cuda", torch.float16)
unet.load_state_dict(torch.load(hf_hub_download(repo_name, ckpt_name), map_location="cuda"))

pipe = DiffusionPipeline.from_pretrained(base_model_id, unet=unet, torch_dtype=torch.float16, use_safetensors=True, variant="fp16").to("cuda")
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
prompt="A photo of a dog." 
image=pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0, timesteps=[399], height=1024, width=1024).images[0]
image.save('output.png', 'PNG')

For more details, please refer to our github repository

License

DisBack is released under MIT license

Citation

If you find our paper useful or relevant to your research, please kindly cite our papers:

@article{zhang2024distributionbacktrackingbuildsfaster,
    title={Distribution Backtracking Builds A Faster Convergence Trajectory for One-step Diffusion Distillation}, 
    author={Shengyuan Zhang and Ling Yang and Zejian Li and An Zhao and Chenye Meng and Changyuan Yang and Guang Yang and Zhiyuan Yang and Lingyun Sun},
    journal={arXiv 2408.15991},
    year={2024}
}

Credits

DisBack is highly built on the following amazing open-source projects:

DMD2: Improved Distribution Matching Distillation for Fast Image Synthesis

Diff-Instruct: Diff-Instruct: A Universal Approach for Transferring Knowledge From Pre-trained Diffusion Models

ScoreGAN: Unifying GANs and Score-Based Diffusion as Generative Particle Models

Thanks to the maintainers of these projects for their contribution to this project!