Update README.md

README.md

@@ -1,78 +1,91 @@
----
-license: apache-2.0
-datasets:
-- yuvalkirstain/pickapic_v1
-language:
-- en
-pipeline_tag: text-to-image
----
-# Step-aware Preference Optimization: Aligning Preference with Denoising Performance at Each Step
-
-<a href="https://arxiv.org/abs/2406.04314"><img src="https://img.shields.io/badge/Paper-arXiv-red?style=for-the-badge" height=22.5></a>
-<a href="https://github.com/RockeyCoss/SPO"><img src="https://img.shields.io/badge/Gihub-Code-succees?style=for-the-badge&logo=GitHub" height=22.5></a>
-<a href="https://rockeycoss.github.io/spo.github.io/"><img src="https://img.shields.io/badge/Project-Page-blue?style=for-the-badge" height=22.5></a>
-
-<table>
-  <tr>
-    <td><img src="assets/1.png" alt="teaser example 0" width="200"/></td>
-    <td><img src="assets/2.png" alt="teaser example 1" width="200"/></td>
-    <td><img src="assets/3.png" alt="teaser example 2" width="200"/></td>
-    <td><img src="assets/4.png" alt="teaser example 3" width="200"/></td>
-  </tr>
-</table>
-
-## Abstract
-<p>
-Recently, Direct Preference Optimization (DPO) has extended its success from aligning large language models (LLMs) to aligning text-to-image diffusion models with human preferences.
-Unlike most existing DPO methods that assume all diffusion steps share a consistent preference order with the final generated images, we argue that this assumption neglects step-specific denoising performance and that preference labels should be tailored to each step's contribution.
-</p>
-<p> 
-To address this limitation, we propose Step-aware Preference Optimization (SPO), a novel post-training approach that independently evaluates and adjusts the denoising performance at each step, using a <em>step-aware preference model</em> and a <em>step-wise resampler</em> to ensure accurate step-aware supervision.
-Specifically, at each denoising step, we sample a pool of images, find a suitable win-lose pair, and, most importantly, randomly select a single image from the pool to initialize the next denoising step. This step-wise resampler process ensures the next win-lose image pair comes from the same image, making the win-lose comparison independent of the previous step. To assess the preferences at each step, we train a separate step-aware preference model that can be applied to both noisy and clean images.
-</p> 
-<p>
-Our experiments with Stable Diffusion v1.5 and SDXL demonstrate that SPO significantly outperforms the latest Diffusion-DPO in aligning generated images with complex, detailed prompts and enhancing aesthetics, while also achieving more than 20&times; times faster in training efficiency. Code and model: <a ref="https://rockeycoss.github.io/spo.github.io/">https://rockeycoss.github.io/spo.github.io/</a>
-</p>
-
-## Model Description
-
-This model is fine-tuned from [runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5). It has been trained on 4,000 prompts for 10 epochs. This checkpoint is a LoRA checkpoint. We also provide a LoRA checkpoint compatible with [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui), which can be accessed [here](https://civitai.com/models/526379/spo-sd-v1-54k-p10eplorawebui)
-
-If you want to access the merged checkpoint that combines the LoRA checkpoint with the base model [runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5), please visit [SPO-SD-v1-5_4k-p_10ep](https://huggingface.co/SPO-Diffusion-Models/SPO-SD-v1-5_4k-p_10ep).
-
-
-
-## A quick example
-```python
-from diffusers import StableDiffusionPipeline
-import torch
-
-# load pipeline
-inference_dtype = torch.float16
-pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    torch_dtype=inference_dtype,
-)
-pipe.load_lora_weights("SPO-Diffusion-Models/SPO-SD-v1-5_4k-p_10ep_LoRA")
-pipe.to('cuda')
-
-generator=torch.Generator(device='cuda').manual_seed(42)
-image = pipe(
-    prompt='an image of a beautiful lake',
-    generator=generator,
-    guidance_scale=7.5,
-    output_type='pil',
-).images[0]
-image.save('lake.png')
-```
-
-## Citation
-If you find our work or codebase useful, please consider giving us a star and citing our work.
-```
-@article{liang2024step,
-  title={Step-aware Preference Optimization: Aligning Preference with Denoising Performance at Each Step},
-  author={Liang, Zhanhao and Yuan, Yuhui and Gu, Shuyang and Chen, Bohan and Hang, Tiankai and Li, Ji and Zheng, Liang},
-  journal={arXiv preprint arXiv:2406.04314},
-  year={2024}
-}
+---
+license: apache-2.0
+datasets:
+- yuvalkirstain/pickapic_v1
+language:
+- en
+pipeline_tag: text-to-image
+---
+# Aesthetic Post-Training Diffusion Models from Generic Preferences with Step-by-step Preference
+
+<a href="https://arxiv.org/abs/2406.04314"><img src="https://img.shields.io/badge/Paper-arXiv-red?style=for-the-badge" height=22.5></a>
+<a href="https://github.com/RockeyCoss/SPO"><img src="https://img.shields.io/badge/Gihub-Code-succees?style=for-the-badge&logo=GitHub" height=22.5></a>
+<a href="https://rockeycoss.github.io/spo.github.io/"><img src="https://img.shields.io/badge/Project-Page-blue?style=for-the-badge" height=22.5></a>
+
+<table>
+  <tr>
+    <td><img src="assets/1.png" alt="teaser example 0" width="200"/></td>
+    <td><img src="assets/2.png" alt="teaser example 1" width="200"/></td>
+    <td><img src="assets/3.png" alt="teaser example 2" width="200"/></td>
+    <td><img src="assets/4.png" alt="teaser example 3" width="200"/></td>
+  </tr>
+</table>
+
+## Abstract
+<p>
+    Generating visually appealing images is fundamental to modern text-to-image generation models. 
+    A potential solution to better aesthetics is direct preference optimization (DPO), 
+    which has been applied to diffusion models to improve general image quality including prompt alignment and aesthetics. 
+    Popular DPO methods propagate preference labels from clean image pairs to all the intermediate steps along the two generation trajectories. 
+    However, preference labels provided in existing datasets are blended with layout and aesthetic opinions, which would disagree with aesthetic preference. 
+    Even if aesthetic labels were provided (at substantial cost), it would be hard for the two-trajectory methods to capture nuanced visual differences at different steps.
+</p>
+<p>
+    To improve aesthetics economically, this paper uses existing generic preference data and introduces step-by-step preference optimization 
+    (SPO) that discards the propagation strategy and allows fine-grained image details to be assessed. Specifically, 
+    at each denoising step, we 1) sample a pool of candidates by denoising from a shared noise latent, 
+    2) use a step-aware preference model to find a suitable win-lose pair to supervise the diffusion model, and 
+    3) randomly select one from the pool to initialize the next denoising step. 
+    This strategy ensures that diffusion models focus on the subtle, fine-grained visual differences 
+    instead of layout aspect. We find that aesthetic can be significantly enhanced by accumulating these 
+    improved minor differences.
+</p>
+<p>
+    When fine-tuning Stable Diffusion v1.5 and SDXL, SPO yields significant 
+    improvements in aesthetics compared with existing DPO methods while not sacrificing image-text alignment 
+    compared with vanilla models. Moreover, SPO converges much faster than DPO methods due to the step-by-step 
+    alignment of fine-grained visual details.
+</p>
+
+## Model Description
+
+This model is fine-tuned from [runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5). It has been trained on 4,000 prompts for 10 epochs. This checkpoint is a LoRA checkpoint. We also provide a LoRA checkpoint compatible with [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui), which can be accessed [here](https://civitai.com/models/526379/spo-sd-v1-54k-p10eplorawebui)
+
+If you want to access the merged checkpoint that combines the LoRA checkpoint with the base model [runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5), please visit [SPO-SD-v1-5_4k-p_10ep](https://huggingface.co/SPO-Diffusion-Models/SPO-SD-v1-5_4k-p_10ep).
+
+
+
+## A quick example
+```python
+from diffusers import StableDiffusionPipeline
+import torch
+
+# load pipeline
+inference_dtype = torch.float16
+pipe = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=inference_dtype,
+)
+pipe.load_lora_weights("SPO-Diffusion-Models/SPO-SD-v1-5_4k-p_10ep_LoRA")
+pipe.to('cuda')
+
+generator=torch.Generator(device='cuda').manual_seed(42)
+image = pipe(
+    prompt='an image of a beautiful lake',
+    generator=generator,
+    guidance_scale=7.5,
+    output_type='pil',
+).images[0]
+image.save('lake.png')
+```
+
+## Citation
+If you find our work or codebase useful, please consider giving us a star and citing our work.
+```
+@article{liang2024step,
+  title={Aesthetic Post-Training Diffusion Models from Generic Preferences with Step-by-step Preference Optimization},
+  author={Liang, Zhanhao and Yuan, Yuhui and Gu, Shuyang and Chen, Bohan and Hang, Tiankai and Cheng, Mingxi and Li, Ji and Zheng, Liang},
+  journal={arXiv preprint arXiv:2406.04314},
+  year={2024}
+}
 ```