---
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

![image/jpeg](assets/teaser.jpg)

> [**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.