Update README.md

README.md

@@ -6,21 +6,27 @@ This repository contains offset versions of https://huggingface.co/mhdang/dpo-sd
 These can be added directly to any initialized UNet to inject DPO training into it. See the code below for usage (diffusers only.)
 
 ```py
-def inject_dpo(unet: UNet2DConditionModel, dpo_path: str, strict: bool = False) -> None:
+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_path (`str`)
+      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="cpu") as f:
+  with safe_open(dpo_offset_path, framework="pt", device=device) as f:
       for key in f.keys():
           key_parts = key.split(".")
           current_layer = unet
@@ -35,9 +41,11 @@ def inject_dpo(unet: UNet2DConditionModel, dpo_path: str, strict: bool = False)
               layer_param = getattr(current_layer, key_parts[-1], None)
               if layer_param is None:
                   if strict:
-                      raise IOError(f"Couldn't get weighht parameter for key {key}")
+                      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
@@ -46,36 +54,68 @@ 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)
-
-# download DPO offsets
-dpo_path = huggingface_hub.hf_hub_download("painebenjamin/sd-dpo-offsets", "sd_v15_unet_dpo_offset.safetensors")
-# inject
-inject_dpo(pipe.unet, dpo_path)
+pipe.to(device)
 
 # make image
 prompt = "Two cats playing chess on a tree branch"
-image = pipe(prompt, guidance_scale=7.5).images[0] 
+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 sdxl pipeline
+# load sdv15 pipeline
+device = "cuda"
 model_id = "Lykon/dreamshaper-xl-1-0"
-pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16")
+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_path = huggingface_hub.hf_hub_download("painebenjamin/sd-dpo-offsets", "sd_xl_unet_dpo_offset.safetensors")
+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_path)
+inject_dpo(pipe.unet, dpo_offset_path, device)
 
-# make image
-prompt = "Two cats playing chess on a tree branch"
-image = pipe(prompt, guidance_scale=7.5).images[0]
-image.save("cats_playing_chess.png")
-```
+# 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)