[Feature] Support Peft

.github/workflows/build.yml

@@ -43,6 +43,8 @@ jobs:
           pip install setuptools
       - name: Build and install
         run: rm -rf .eggs && pip install -e .[dev]
+      - name: Install diffusers from main branch
+        run: pip install git+https://github.com/huggingface/diffusers.git
       - name: Run unittests and generate coverage report
         run: |
           coverage run --branch --source=diffengine -m pytest tests/

README.md

@@ -193,6 +193,8 @@ For detailed user guides and advanced guides, please refer to our [Documentation
           <li><a href="configs/esd/README.md">Erasing Concepts from Diffusion Models (2023)</a></li>
           <li><a href="configs/ssd_1b/README.md">SSD-1B (2023)</a></li>
           <li><a href="configs/instruct_pix2pix/README.md">InstructPix2Pix (2022)</a></li>
+          <li><a href="configs/stable_diffusion_xl_loha/README.md">LoHa (ICLR'2022)</a></li>
+          <li><a href="configs/stable_diffusion_xl_lokr/README.md">LoKr (2022)</a></li>
         </ul>
       </td>
       <td>

configs/_base_/datasets/dog_dreambooth.py

@@ -29,5 +29,5 @@
         prompt=["A photo of sks dog in a bucket"] * 4,
         by_epoch=False,
         interval=100),
-    dict(type="LoRASaveHook"),
+    dict(type="PeftSaveHook"),
 ]

configs/_base_/datasets/dog_dreambooth_if.py

@@ -28,5 +28,5 @@
         prompt=["A photo of sks dog in a bucket"] * 4,
         by_epoch=False,
         interval=100),
-    dict(type="LoRASaveHook"),
+    dict(type="PeftSaveHook"),
 ]

configs/_base_/datasets/dog_dreambooth_xl.py

@@ -33,5 +33,5 @@
         interval=100,
         height=1024,
         width=1024),
-    dict(type="LoRASaveHook"),
+    dict(type="PeftSaveHook"),
 ]

configs/_base_/datasets/google_backpack_dreambooth_xl.py

@@ -35,5 +35,5 @@
         interval=100,
         height=1024,
         width=1024),
-    dict(type="LoRASaveHook"),
+    dict(type="PeftSaveHook"),
 ]

configs/_base_/datasets/keramer_face_dreambooth.py

@@ -29,5 +29,5 @@
         prompt=["a photo of sks person in suits"] * 4,
         by_epoch=False,
         interval=100),
-    dict(type="LoRASaveHook"),
+    dict(type="PeftSaveHook"),
 ]

configs/_base_/datasets/keramer_face_dreambooth_xl.py

@@ -33,5 +33,5 @@
         interval=100,
         height=1024,
         width=1024),
-    dict(type="LoRASaveHook"),
+    dict(type="PeftSaveHook"),
 ]

configs/_base_/datasets/potatohead_dreambooth_xl.py

@@ -33,5 +33,5 @@
         interval=100,
         height=1024,
         width=1024),
-    dict(type="LoRASaveHook"),
+    dict(type="PeftSaveHook"),
 ]

configs/_base_/datasets/starbucks_dreambooth_xl.py

@@ -33,5 +33,5 @@
         interval=100,
         height=1024,
         width=1024),
-    dict(type="LoRASaveHook"),
+    dict(type="PeftSaveHook"),
 ]

configs/_base_/models/deepfloyd_if_xl_lora.py

@@ -1,5 +1,9 @@
 model = dict(
     type="DeepFloydIF",
     model="DeepFloyd/IF-I-XL-v1.0",
-    lora_config=dict(rank=8),
+    unet_lora_config=dict(
+        type="LoRA",
+        r=8,
+        lora_alpha=1,
+        target_modules=["to_q", "to_v", "to_k", "to_out.0"]),
     gradient_checkpointing=True)

configs/_base_/models/small_sd_lora.py

@@ -1,2 +1,8 @@
 model = dict(
-    type="StableDiffusion", model="segmind/small-sd", lora_config=dict(rank=8))
+    type="StableDiffusion",
+    model="segmind/small-sd",
+    unet_lora_config=dict(
+        type="LoRA",
+        r=8,
+        lora_alpha=1,
+        target_modules=["to_q", "to_v", "to_k", "to_out.0"]))

configs/_base_/models/ssd_1b_lora.py

@@ -2,4 +2,8 @@
     type="StableDiffusionXL",
     model="segmind/SSD-1B",
     vae_model="madebyollin/sdxl-vae-fp16-fix",
-    lora_config=dict(rank=8))
+    unet_lora_config=dict(
+        type="LoRA",
+        r=8,
+        lora_alpha=1,
+        target_modules=["to_q", "to_v", "to_k", "to_out.0"]))

configs/_base_/models/stable_diffusion_v15_lora.py

@@ -1,4 +1,8 @@
 model = dict(
     type="StableDiffusion",
     model="runwayml/stable-diffusion-v1-5",
-    lora_config=dict(rank=8))
+    unet_lora_config=dict(
+        type="LoRA",
+        r=8,
+        lora_alpha=1,
+        target_modules=["to_q", "to_v", "to_k", "to_out.0"]))

configs/_base_/models/stable_diffusion_v21_lora.py

@@ -1,4 +1,8 @@
 model = dict(
     type="StableDiffusion",
     model="stabilityai/stable-diffusion-2-1",
-    lora_config=dict(rank=8))
+    unet_lora_config=dict(
+        type="LoRA",
+        r=8,
+        lora_alpha=1,
+        target_modules=["to_q", "to_v", "to_k", "to_out.0"]))

configs/_base_/models/stable_diffusion_xl_loha.py

@@ -0,0 +1,9 @@
+model = dict(
+    type="StableDiffusionXL",
+    model="stabilityai/stable-diffusion-xl-base-1.0",
+    vae_model="madebyollin/sdxl-vae-fp16-fix",
+    unet_lora_config=dict(
+        type="LoHa",
+        r=8,
+        alpha=1,
+        target_modules=["to_q", "to_v", "to_k", "to_out.0"]))

configs/_base_/models/stable_diffusion_xl_lokr.py

@@ -0,0 +1,9 @@
+model = dict(
+    type="StableDiffusionXL",
+    model="stabilityai/stable-diffusion-xl-base-1.0",
+    vae_model="madebyollin/sdxl-vae-fp16-fix",
+    unet_lora_config=dict(
+        type="LoKr",
+        r=8,
+        alpha=1,
+        target_modules=["to_q", "to_v", "to_k", "to_out.0"]))

configs/_base_/models/stable_diffusion_xl_lora.py

@@ -2,4 +2,8 @@
     type="StableDiffusionXL",
     model="stabilityai/stable-diffusion-xl-base-1.0",
     vae_model="madebyollin/sdxl-vae-fp16-fix",
-    lora_config=dict(rank=8))
+    unet_lora_config=dict(
+        type="LoRA",
+        r=8,
+        lora_alpha=1,
+        target_modules=["to_q", "to_v", "to_k", "to_out.0"]))

configs/_base_/models/tiny_sd_lora.py

@@ -1,2 +1,8 @@
 model = dict(
-    type="StableDiffusion", model="segmind/tiny-sd", lora_config=dict(rank=8))
+    type="StableDiffusion",
+    model="segmind/tiny-sd",
+    unet_lora_config=dict(
+        type="LoRA",
+        r=8,
+        lora_alpha=1,
+        target_modules=["to_q", "to_v", "to_k", "to_out.0"]))

configs/_base_/models/wuerstchen_prior_lora.py

@@ -2,4 +2,8 @@
     type="WuerstchenPriorModel",
     decoder_model="warp-ai/wuerstchen",
     prior_model="warp-ai/wuerstchen-prior",
-    lora_config=dict(rank=8))
+    prior_lora_config=dict(
+        type="LoRA",
+        r=8,
+        lora_alpha=1,
+        target_modules=["to_q", "to_v", "to_k", "to_out.0"]))

configs/deepfloyd_if_dreambooth/README.md

@@ -40,19 +40,29 @@ $ mim train diffengine configs/deepfloyd_if_dreambooth/deepfloyd_if_xl_dreamboot
 Once you have trained a model, specify the path to the saved model and utilize it for inference using the `diffusers.pipeline` module.
 
 ```py
+from pathlib import Path
+
 import torch
 from diffusers import DiffusionPipeline
+from peft import PeftModel
 
-checkpoint = 'work_dirs/deepfloyd_if_xl_dreambooth_lora_dog/step999'
+checkpoint = Path('work_dirs/deepfloyd_if_xl_dreambooth_lora_dog/step999')
 prompt = 'A photo of sks dog in a bucket'
 
-pipe = DiffusionPipeline.from_pretrained('DeepFloyd/IF-I-XL-v1.0')
+pipe = DiffusionPipeline.from_pretrained(
+    'runwayml/stable-diffusion-v1-5', torch_dtype=torch.float16)
 pipe.to('cuda')
-pipe.load_lora_weights(checkpoint)
+pipe.unet = PeftModel.from_pretrained(pipe.unet, checkpoint / "unet", adapter_name="default")
+if (checkpoint / "text_encoder").exists():
+    pipe.text_encoder = PeftModel.from_pretrained(
+        pipe.text_encoder, checkpoint / "text_encoder", adapter_name="default"
+    )
 
 image = pipe(
     prompt,
     num_inference_steps=50,
+    height=1024,
+    width=1024,
 ).images[0]
 image.save('demo.png')
 ```

configs/distill_sd_dreambooth/README.md

@@ -59,20 +59,29 @@ Settings:
 Once you have trained a model, specify the path to where the model is saved, and use it for inference with the `diffusers`.
 
 ```py
+from pathlib import Path
+
 import torch
 from diffusers import DiffusionPipeline
+from peft import PeftModel
 
-checkpoint = 'work_dirs/small_sd_dreambooth_lora_dog/step999'
+checkpoint = Path('work_dirs/small_sd_dreambooth_lora_dog/step999')
 prompt = 'A photo of sks dog in a bucket'
 
 pipe = DiffusionPipeline.from_pretrained(
-    'segmind/small-sd', torch_dtype=torch.float16)
+    'runwayml/stable-diffusion-v1-5', torch_dtype=torch.float16)
 pipe.to('cuda')
-pipe.load_lora_weights(checkpoint)
+pipe.unet = PeftModel.from_pretrained(pipe.unet, checkpoint / "unet", adapter_name="default")
+if (checkpoint / "text_encoder").exists():
+    pipe.text_encoder = PeftModel.from_pretrained(
+        pipe.text_encoder, checkpoint / "text_encoder", adapter_name="default"
+    )
 
 image = pipe(
     prompt,
     num_inference_steps=50,
+    height=1024,
+    width=1024,
 ).images[0]
 image.save('demo.png')
 ```

configs/distill_sd_dreambooth/small_sd_dreambooth_lora_keramer_face.py

@@ -19,5 +19,5 @@
         prompt=["Portrait photo of a sks person in suits"] * 4,
         by_epoch=False,
         interval=100),
-    dict(type="LoRASaveHook"),
+    dict(type="PeftSaveHook"),
 ]

configs/distill_sd_dreambooth/tiny_sd_dreambooth_lora_keramer_face.py

@@ -17,5 +17,5 @@
         prompt=["Portrait photo of a sks person in suits"] * 4,
         by_epoch=False,
         interval=100),
-    dict(type="LoRASaveHook"),
+    dict(type="PeftSaveHook"),
 ]

configs/ssd_1b/README.md

@@ -15,14 +15,6 @@ Today, Segmind is thrilled to announce the open sourcing of our new foundational
 ```
 ```
 
-## Dependencies
-
-Note that install diffusers from source to use SSD-1B.
-
-```
-pip install -U git+https://github.com/huggingface/diffusers.git
-```
-
 ## Run Training
 
 Run Training

configs/ssd_1b_dreambooth/README.md

@@ -22,14 +22,6 @@ Large text-to-image models achieved a remarkable leap in the evolution of AI, en
 }
 ```
 
-## Dependencies
-
-Note that install diffusers from source to use SSD-1B.
-
-```
-pip install -U git+https://github.com/huggingface/diffusers.git
-```
-
 ## Run Training
 
 Run Training
@@ -66,10 +58,13 @@ Settings:
 Once you have trained a model, specify the path to where the model is saved, and use it for inference with the `diffusers`.
 
 ```py
+from pathlib import Path
+
 import torch
 from diffusers import DiffusionPipeline, AutoencoderKL
+from peft import PeftModel
 
-checkpoint = 'work_dirs/ssd_1b_dreambooth_lora_dog/step499'
+checkpoint = Path('work_dirs/ssd_1b_dreambooth_lora_dog/step499')
 prompt = 'A photo of sks dog in a bucket'
 
 vae = AutoencoderKL.from_pretrained(
@@ -79,13 +74,21 @@ vae = AutoencoderKL.from_pretrained(
 pipe = DiffusionPipeline.from_pretrained(
     'segmind/SSD-1B', vae=vae, torch_dtype=torch.float16)
 pipe.to('cuda')
-pipe.load_lora_weights(checkpoint)
+pipe.unet = PeftModel.from_pretrained(pipe.unet, checkpoint / "unet", adapter_name="default")
+if (checkpoint / "text_encoder_one").exists():
+    pipe.text_encoder_one = PeftModel.from_pretrained(
+        pipe.text_encoder_one, checkpoint / "text_encoder_one", adapter_name="default"
+    )
+if (checkpoint / "text_encoder_two").exists():
+    pipe.text_encoder_one = PeftModel.from_pretrained(
+        pipe.text_encoder_two, checkpoint / "text_encoder_two", adapter_name="default"
+    )
 
 image = pipe(
     prompt,
     num_inference_steps=50,
-    width=1024,
     height=1024,
+    width=1024,
 ).images[0]
 image.save('demo.png')
 ```

configs/stable_diffusion_dreambooth/README.md

@@ -40,20 +40,29 @@ $ mim train diffengine configs/stable_diffusion_dreambooth/stable_diffusion_v15_
 Once you have trained a model, specify the path to the saved model and utilize it for inference using the `diffusers.pipeline` module.
 
 ```py
+from pathlib import Path
+
 import torch
 from diffusers import DiffusionPipeline
+from peft import PeftModel
 
-checkpoint = 'work_dirs/stable_diffusion_v15_dreambooth_lora_dog/step999'
+checkpoint = Path('work_dirs/stable_diffusion_v15_dreambooth_lora_dog/step999')
 prompt = 'A photo of sks dog in a bucket'
 
 pipe = DiffusionPipeline.from_pretrained(
     'runwayml/stable-diffusion-v1-5', torch_dtype=torch.float16)
 pipe.to('cuda')
-pipe.load_lora_weights(checkpoint)
+pipe.unet = PeftModel.from_pretrained(pipe.unet, checkpoint / "unet", adapter_name="default")
+if (checkpoint / "text_encoder").exists():
+    pipe.text_encoder = PeftModel.from_pretrained(
+        pipe.text_encoder, checkpoint / "text_encoder", adapter_name="default"
+    )
 
 image = pipe(
     prompt,
     num_inference_steps=50,
+    height=1024,
+    width=1024,
 ).images[0]
 image.save('demo.png')
 ```

configs/stable_diffusion_lora/README.md

@@ -42,20 +42,29 @@ $ mim train diffengine configs/stable_diffusion_lora/stable_diffusion_v15_lora_p
 Once you have trained a model, specify the path to the saved model and utilize it for inference using the `diffusers.pipeline` module.
 
 ```py
+from pathlib import Path
+
 import torch
 from diffusers import DiffusionPipeline
+from peft import PeftModel
 
-checkpoint = 'work_dirs/stable_diffusion_v15_lora_pokemon_blip/step10450'
+checkpoint = Path('work_dirs/stable_diffusion_v15_lora_pokemon_blip/step10450')
 prompt = 'yoda pokemon'
 
 pipe = DiffusionPipeline.from_pretrained(
     'runwayml/stable-diffusion-v1-5', torch_dtype=torch.float16)
 pipe.to('cuda')
-pipe.load_lora_weights(checkpoint)
+pipe.unet = PeftModel.from_pretrained(pipe.unet, checkpoint / "unet", adapter_name="default")
+if (checkpoint / "text_encoder").exists():
+    pipe.text_encoder = PeftModel.from_pretrained(
+        pipe.text_encoder, checkpoint / "text_encoder", adapter_name="default"
+    )
 
 image = pipe(
     prompt,
     num_inference_steps=50,
+    height=1024,
+    width=1024,
 ).images[0]
 image.save('demo.png')
 ```

configs/stable_diffusion_lora/stable_diffusion_v15_lora_pokemon_blip.py

@@ -9,5 +9,5 @@
 
 custom_hooks = [
     dict(type="VisualizationHook", prompt=["yoda pokemon"] * 4),
-    dict(type="LoRASaveHook"),
+    dict(type="PeftSaveHook"),
 ]