| --- |
| license: cc-by-nc-4.0 |
| library_name: diffusers |
| pipeline_tag: text-to-image |
| base_model: |
| - stabilityai/stable-diffusion-xl-base-1.0 |
| - Tongyi-MAI/Z-Image-Turbo |
| tags: |
| - text-to-image |
| - stable-diffusion-xl |
| - z-image |
| - diffusion distillation |
| - one-step generation |
| - rlhf |
| --- |
| |
| # DIDR Model Card |
|
|
|  |
|
|
| > [**Diff-Instruct with Diffused Reward: Towards Principled One-step Generator RL**](https://arxiv.org/abs/2605.24001) (NeurIPS 2026), |
| > [Junyi Wu](https://junyi-w.github.io), Weijian Luo, Haoyang Zheng, Ruizhe Zhang, Guang Lin |
|
|
| ## Contact |
|
|
| Feel free to contact us if you have any questions about the paper! |
|
|
| Junyi Wu [wu2393@purdue.edu](mailto:wu2393@purdue.edu) |
|
|
| ## Usage |
|
|
| We can use the standard diffusers pipeline: |
|
|
| #### 1-step SDXL UNet generation (DIDR) |
|
|
| ```python |
| import torch |
| from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler |
| from huggingface_hub import hf_hub_download |
| from safetensors.torch import load_file |
| base_model_id = "stabilityai/stable-diffusion-xl-base-1.0" |
| repo_name = "Junyi-W/DIDR" |
| ckpt_name = "didr_sdxl_1step_unet_fp16.safetensors" |
| # Load model. |
| unet = UNet2DConditionModel.from_config(UNet2DConditionModel.load_config(base_model_id, subfolder="unet")).to("cuda", torch.float16) |
| unet.load_state_dict(load_file(hf_hub_download(repo_name, ckpt_name), device="cuda")) |
| pipe = DiffusionPipeline.from_pretrained(base_model_id, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda") |
| pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config) |
| prompt="a photo of a cat" |
| image=pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0, timesteps=[399]).images[0] |
| ``` |
|
|
| #### 1-step SDXL UNet generation (DIDR-longer) |
|
|
| ```python |
| import torch |
| from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler |
| from huggingface_hub import hf_hub_download |
| from safetensors.torch import load_file |
| base_model_id = "stabilityai/stable-diffusion-xl-base-1.0" |
| repo_name = "Junyi-W/DIDR" |
| ckpt_name = "didr_longer_sdxl_1step_unet_fp16.safetensors" |
| # Load model. |
| unet = UNet2DConditionModel.from_config(UNet2DConditionModel.load_config(base_model_id, subfolder="unet")).to("cuda", torch.float16) |
| unet.load_state_dict(load_file(hf_hub_download(repo_name, ckpt_name), device="cuda")) |
| pipe = DiffusionPipeline.from_pretrained(base_model_id, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda") |
| pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config) |
| prompt="a photo of a cat" |
| image=pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0, timesteps=[399]).images[0] |
| ``` |
|
|
| #### 1-step Z-Image generation |
|
|
| ```python |
| import torch |
| from diffusers import ZImagePipeline, ZImageTransformer2DModel |
| base_model_id = "Tongyi-MAI/Z-Image-Turbo" |
| repo_name = "Junyi-W/DIDR" |
| # Load model. |
| transformer = ZImageTransformer2DModel.from_pretrained(repo_name, subfolder="zimage_transformer", torch_dtype=torch.bfloat16) |
| pipe = ZImagePipeline.from_pretrained(base_model_id, transformer=transformer, torch_dtype=torch.bfloat16).to("cuda") |
| prompt="A photorealistic portrait of a young woman in traditional Chinese red Hanfu, intricate golden embroidery, dramatic lighting, ultra high definition" |
| image=pipe(prompt=prompt, height=1024, width=1024, num_inference_steps=1, guidance_scale=0.0).images[0] |
| ``` |
|
|
| For more information, please refer to the [code repository](https://github.com/Junyi-W/DIDR) |
|
|
|
|
| ## License |
|
|
| DIDR is released under [Creative Commons Attribution-NonCommercial 4.0 International License](https://creativecommons.org/licenses/by-nc/4.0/deed.en). |
|
|
|
|
| ## Citation |
|
|
| If you find DIDR useful or relevant to your research, please kindly cite our paper: |
|
|
| ```bib |
| @inproceedings{wu2026didr, |
| title={Diff-Instruct with Diffused Reward: Towards Principled One-step Generator RL}, |
| author={Wu, Junyi and Luo, Weijian and Zheng, Haoyang and Zhang, Ruizhe and Lin, Guang}, |
| booktitle={Advances in Neural Information Processing Systems (NeurIPS)}, |
| year={2026} |
| } |
| ``` |
|
|
|
|
| ## Acknowledgments |
|
|
| Our SDXL generators are initialized from [DMD2](https://huggingface.co/tianweiy/DMD2) and our Z-Image generators from [Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo). We thank the authors for releasing their models. |
|
|
