Views
No views yet
Distribution Backtracking Builds A Faster Convergence Trajectory for One-step Diffusion Distillation,
Shengyuan Zhang1, Ling Yang2, Zejian Li*1, An Zhao1, Chenye Meng1, Changyuan Yang3, Guang Yang3, Zhiyuan Yang3, Lingyun Sun1
1Zhejiang University 2Peking University 3Alibaba Group*
1import torch
2from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler
3from huggingface_hub import hf_hub_download
4
5base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
6repo_name = "SYZhang0805/DisBack"
7ckpt_name = "SDXL_DisBack.bin"
8
9unet = UNet2DConditionModel.from_config(base_model_id, subfolder="unet").to("cuda", torch.float16)
10unet.load_state_dict(torch.load(hf_hub_download(repo_name, ckpt_name), map_location="cuda"))
11
12pipe = DiffusionPipeline.from_pretrained(base_model_id, unet=unet, torch_dtype=torch.float16, use_safetensors=True, variant="fp16").to("cuda")
13pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
14prompt="A photo of a dog."
15image=pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0, timesteps=[399], height=1024, width=1024).images[0]
16image.save('output.png', 'PNG')1@article{zhang2024distributionbacktrackingbuildsfaster,
2 title={Distribution Backtracking Builds A Faster Convergence Trajectory for One-step Diffusion Distillation},
3 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},
4 journal={arXiv 2408.15991},
5 year={2024}
6}