File size: 4,369 Bytes
b00eff9
 
 
c817021
 
 
 
 
6334eed
 
 
 
 
 
 
c817021
 
 
 
 
 
6334eed
c817021
 
 
 
 
 
6334eed
c817021
 
 
 
 
 
 
 
 
 
 
 
 
 
6334eed
 
c817021
 
6334eed
c817021
 
 
 
 
 
 
 
6334eed
c817021
 
6334eed
c817021
 
 
6334eed
 
 
c817021
6334eed
 
 
 
 
 
 
 
 
 
c817021
 
6334eed
 
 
 
 
 
c817021
 
 
 
6334eed
 
c817021
6334eed
 
c817021
6334eed
 
 
 
 
 
 
 
 
c817021
 
1669423
c817021
6334eed
c817021
6334eed
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
---
license: openrail++
---
This repository contains offset versions of https://huggingface.co/mhdang/dpo-sdxl-text2image-v1 and https://huggingface.co/mhdang/dpo-sd1.5-text2image-v1.

These can be added directly to any initialized UNet to inject DPO training into it. See the code below for usage (diffusers only.)

```py
from __future__ import annotations

from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from diffusers.models import UNet2DConditionModel

def inject_dpo(unet: UNet2DConditionModel, dpo_offset_path: str, device: str, strict: bool = False) -> None:
  """
  Injects DPO weights directly into your UNet.

  Args:
      unet (`UNet2DConditionModel`)
          The initialized UNet from your pipeline.
      dpo_offset_path (`str`)
          The path to the `.safetensors` file downloaded from https://huggingface.co/benjamin-paine/sd-dpo-offsets/.
          Make sure you're using the right file for the right base model.
      strict (`bool`, *optional*)
          Whether or not to raise errors when a weight cannot be applied. Defaults to false.
  """
  from safetensors import safe_open
  with safe_open(dpo_offset_path, framework="pt", device=device) as f:
      for key in f.keys():
          key_parts = key.split(".")
          current_layer = unet
          for key_part in key_parts[:-1]:
              current_layer = getattr(current_layer, key_part, None)
              if current_layer is None:
                  break
              if current_layer is None:
                  if strict:
                      raise IOError(f"Couldn't find a layer to inject key {key} in.")
                  continue
              layer_param = getattr(current_layer, key_parts[-1], None)
              if layer_param is None:
                  if strict:
                      raise IOError(f"Couldn't get weight parameter for key {key}")
                  continue
              layer_param.data += f.get_tensor(key)
```

Now you can use this function like so:

```py
from diffusers import StableDiffusionPipeline
import huggingface_hub
import torch

# load sdv15 pipeline
device = "cuda"
model_id = "Lykon/dreamshaper-8"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to(device)

# make image
prompt = "Two cats playing chess on a tree branch"
generator = torch.Generator(device=device)
generator.manual_seed(123456789)
image = pipe(prompt, guidance_scale=7.5, generator=generator).images[0] 
image.save("cats_playing_chess.png")

# download DPO offsets
dpo_offset_path = huggingface_hub.hf_hub_download("benjamin-paine/sd-dpo-offsets", "sd_v15_unet_dpo_offset.safetensors")
# inject
inject_dpo(pipe.unet, dpo_offset_path, device)

# make image again
generator.manual_seed(123456789)
image = pipe(prompt, guidance_scale=7.5, generator=generator).images[0] 
image.save("cats_playing_chess_dpo.png")
```

`cats_playing_chess.png`
![image/png](https://cdn-uploads.huggingface.co/production/uploads/64429aaf7feb866811b12f73/KqAohfKMXKVGTDpuBhhx6.png)

`cats_playing_chess_dpo.png`
![image/png](https://cdn-uploads.huggingface.co/production/uploads/64429aaf7feb866811b12f73/fY9j1q8ZazyNP4JbD0TTU.png)

Or for XL:

```py
from diffusers import StableDiffusionXLPipeline
import huggingface_hub
import torch

# load sdv15 pipeline
device = "cuda"
model_id = "Lykon/dreamshaper-xl-1-0"
pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to(device)

# make image
prompt = "Two cats playing chess on a tree branch"
generator = torch.Generator(device=device)
generator.manual_seed(123456789)
image = pipe(prompt, guidance_scale=7.5, generator=generator).images[0] 
image.save("cats_playing_chess_xl.png")

# download DPO offsets
dpo_offset_path = huggingface_hub.hf_hub_download("benjamin-paine/sd-dpo-offsets", "sd_xl_unet_dpo_offset.safetensors")
# inject
inject_dpo(pipe.unet, dpo_offset_path, device)

# make image again
generator.manual_seed(123456789)
image = pipe(prompt, guidance_scale=7.5, generator=generator).images[0] 
image.save("cats_playing_chess_xl_dpo.png")
```

`cats_playing_chess_xl.png`
![image/png](https://cdn-uploads.huggingface.co/production/uploads/64429aaf7feb866811b12f73/BufmVzFBsoYX_jipzErIo.png)

`cats_playing_chess_xl_dpo.png`
![image/png](https://cdn-uploads.huggingface.co/production/uploads/64429aaf7feb866811b12f73/Rj9FXI-vmrMwvepMSLMe7.png)