metadata
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 (NeurIPS 2026), Junyi Wu, 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
Usage
We can use the standard diffusers pipeline:
1-step SDXL UNet generation (DIDR)
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)
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
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
License
DIDR is released under Creative Commons Attribution-NonCommercial 4.0 International License.
Citation
If you find DIDR useful or relevant to your research, please kindly cite our paper:
@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 and our Z-Image generators from Z-Image-Turbo. We thank the authors for releasing their models.