ControlGenAI/ImageReFL_HPS_SD

ImageReFL

Recent advances in diffusion models have led to impressive image generation capabilities, but aligning these models with human preferences remains challenging. Reward-based fine-tuning using models trained on human feedback improves alignment but often harms diversity, producing less varied outputs. In this work, we address this trade-off with two contributions. First, we introduce \textit{combined generation}, a novel sampling strategy that applies a reward-tuned diffusion model only in the later stages of the generation process, while preserving the base model for earlier steps. This approach mitigates early-stage overfitting and helps retain global structure and diversity. Second, we propose \textit{ImageReFL}, a fine-tuning method that improves image diversity with minimal loss in quality by training on real images and incorporating multiple regularizers, including diffusion and ReFL losses. Our approach outperforms conventional reward tuning methods on standard quality and diversity metrics. A user study further confirms that our method better balances human preference alignment and visual diversity.

Model Details

This implementation is based on Stable Diffusion 1.5 and was trained using the reward model HPSv2.1 with the ImageReFL algorithm.
Inference uses the combined generation approach described in the ImageReFL paper.

Model Sources

• Repository
• Paper

How to Get Started with the Model

Model support classical Stable Diffusion inference, but with few addititonal paramters:

• original_unet_steps regulates the number of diffusion steps performed with the original U-Net model. The recommended number is 30 for models based on SD 1.5, and 35 for models based on SDXL.

Example of inference:

from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "ControlGenAI/ImageReFL_HPS_SD",
    trust_remote_code=True
).to(device)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt).images[0]

Citation

If you use this code or our findings for your research, please cite our paper:

@misc{sorokin2025imagereflbalancingqualitydiversity,
      title={ImageReFL: Balancing Quality and Diversity in Human-Aligned Diffusion Models}, 
      author={Dmitrii Sorokin and Maksim Nakhodnov and Andrey Kuznetsov and Aibek Alanov},
      year={2025},
      eprint={2505.22569},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2505.22569}, 
}