rockeycoss commited on
Commit
d85e1da
1 Parent(s): 55f8ed8
README.md CHANGED
@@ -1,3 +1,78 @@
1
  ---
2
  license: apache-2.0
 
 
 
 
 
3
  ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
  license: apache-2.0
3
+ datasets:
4
+ - yuvalkirstain/pickapic_v1
5
+ language:
6
+ - en
7
+ pipeline_tag: text-to-image
8
  ---
9
+ # Step-aware Preference Optimization: Aligning Preference with Denoising Performance at Each Step
10
+
11
+ <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>
12
+ <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>
13
+ <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>
14
+
15
+ <table>
16
+ <tr>
17
+ <td><img src="assets/1.png" alt="teaser example 0" width="200"/></td>
18
+ <td><img src="assets/2.png" alt="teaser example 1" width="200"/></td>
19
+ <td><img src="assets/3.png" alt="teaser example 2" width="200"/></td>
20
+ <td><img src="assets/4.png" alt="teaser example 3" width="200"/></td>
21
+ </tr>
22
+ </table>
23
+
24
+ ## Abstract
25
+ <p>
26
+ 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.
27
+ 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.
28
+ </p>
29
+ <p>
30
+ 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.
31
+ 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.
32
+ </p>
33
+ <p>
34
+ 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>
35
+ </p>
36
+
37
+ ## Model Description
38
+
39
+ 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)
40
+
41
+ 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).
42
+
43
+
44
+
45
+ ## A quick example
46
+ ```python
47
+ from diffusers import StableDiffusionPipeline
48
+ import torch
49
+
50
+ # load pipeline
51
+ inference_dtype = torch.float16
52
+ pipe = StableDiffusionPipeline.from_pretrained(
53
+ "runwayml/stable-diffusion-v1-5",
54
+ torch_dtype=inference_dtype,
55
+ )
56
+ pipe.load_lora_weights("SPO-Diffusion-Models/SPO-SD-v1-5_4k-p_10ep_LoRA")
57
+ pipe.to('cuda')
58
+
59
+ generator=torch.Generator(device='cuda').manual_seed(42)
60
+ image = pipe(
61
+ prompt='an image of a beautiful lake',
62
+ generator=generator,
63
+ guidance_scale=7.5,
64
+ output_type='pil',
65
+ ).images[0]
66
+ image.save('lake.png')
67
+ ```
68
+
69
+ ## Citation
70
+ If you find our work or codebase useful, please consider giving us a star and citing our work.
71
+ ```
72
+ @article{liang2024step,
73
+ title={Step-aware Preference Optimization: Aligning Preference with Denoising Performance at Each Step},
74
+ author={Liang, Zhanhao and Yuan, Yuhui and Gu, Shuyang and Chen, Bohan and Hang, Tiankai and Li, Ji and Zheng, Liang},
75
+ journal={arXiv preprint arXiv:2406.04314},
76
+ year={2024}
77
+ }
78
+ ```
assets/1.png ADDED
assets/2.png ADDED
assets/3.png ADDED
assets/4.png ADDED
spo-sd-v1-5_4k-p_10ep_lora_diffusers.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1be130c5be2de0beacadd3bf0bafe3bedd7e7a380729932a1e369fb29efa86f4
3
+ size 3226184