Update README.md

README.md

@@ -1,3 +1,81 @@
 ---
 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
+def inject_dpo(unet: UNet2DConditionModel, dpo_path: str, strict: bool = False) -> None:
+  """
+  Injects DPO weights directly into your UNet.
+
+  Args:
+      unet (`UNet2DConditionModel`)
+          The initialized UNet from your pipeline.
+      dpo_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:
+      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 weighht parameter for key {key}")
+              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
+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)
+
+# 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")
+```
+
+Or for XL:
+
+```py
+from diffusers import StableDiffusionXLPipeline
+
+# load sdxl pipeline
+model_id = "Lykon/dreamshaper-xl-1-0"
+pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16")
+
+# download DPO offsets
+dpo_path = huggingface_hub.hf_hub_download("painebenjamin/sd-dpo-offsets", "sd_xl_unet_dpo_offset.safetensors")
+# inject
+inject_dpo(pipe.unet, dpo_path)
+
+# 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")
+```