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