Fix lora inject, added weight self apply lora (#39)

* Develop (#34) * Add parameter to control rank of decomposition (#28) * ENH: allow controlling rank of approximation * Training script accepts lora_rank * feat : statefully monkeypatch different loras + example ipynb + readme Co-authored-by: brian6091 <brian6091@gmail.com> * release : version 0.0.4, now able to tune rank, now add loras dynamically * readme : add brain6091's discussions * fix:inject lora in to_out module list * feat: added weight self apply lora * chore: add import copy * fix: readded r Co-authored-by: Simo Ryu <35953539+cloneofsimo@users.noreply.github.com> Co-authored-by: brian6091 <brian6091@gmail.com> Co-authored-by: SimoRyu <cloneofsimo@korea.ac.kr>
cloneofsimo · Dec 15, 2022 · fececf3 · fececf3
1 parent 9f31bd0
commit fececf3
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 19 deletions.
diff --git a/README.md b/README.md
@@ -32,19 +32,21 @@
 
 > "style of sks, superman", with pop-art style LoRA model.
 
-# Web Demo
-
-Integrated into [Huggingface Spaces 🤗](https://huggingface.co/spaces) using [Gradio](https://github.com/gradio-app/gradio). Try out the Web Demo [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/ysharma/Low-rank-Adaptation)
-
 ## Main Features
 
 - Fine-tune Stable diffusion models twice as faster than dreambooth method, by Low-rank Adaptation
-- Get insanely small end result, easy to share and download.
-- Easy to use, compatible with diffusers
-- Sometimes even better performance than full fine-tuning (but left as future work for extensive comparisons)
-- Merge checkpoints by merging LoRA
+- Get insanely small end result (3MB for just unet, 6MB for both unet + clip), easy to share and download.
+- Easy to use, compatible with `diffusers`
+- Sometimes _even better performance_ than full fine-tuning (but left as future work for extensive comparisons)
+- Merge checkpoints + Build recipes by merging LoRAs together
 - Fine-tune both CLIP & Unet to gain better results.
 
+# Web Demo
+
+- Integrated into [Huggingface Spaces 🤗](https://huggingface.co/spaces) using [Gradio](https://github.com/gradio-app/gradio). Try out the Web Demo [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/ysharma/Low-rank-Adaptation)
+
+- Easy [colab running example](https://colab.research.google.com/drive/1iSFDpRBKEWr2HLlz243rbym3J2X95kcy?usp=sharing) of Dreambooth by @pedrogengo
+
 # UPDATES & Notes
 
 - **You can now fine-tune text_encoder as well! Enabled with simple `--train_text_encoder`**
@@ -236,6 +238,15 @@ $$
 
 # Tips and Discussions
 
+## **Training tips in general**
+
+I'm curating a list of tips and discussions here. Feel free to add your own tips and discussions with a PR!
+
+- Discussion by @nitrosocke, can be found [here](https://github.com/cloneofsimo/lora/issues/19#issuecomment-1347149627)
+- Configurations by @xsteenbrugge, Using Clip-interrogator to get a decent prompt seems to work well for him, https://twitter.com/xsteenbrugge/status/1602799180698763264
+- Super easy [colab running example](https://colab.research.google.com/drive/1iSFDpRBKEWr2HLlz243rbym3J2X95kcy?usp=sharing) of Dreambooth by @pedrogengo
+- [Amazing in-depth analysis](https://github.com/cloneofsimo/lora/discussions/37) on the effect of rank, $\alpha_{unet}$, $\alpha_{clip}$, and training configurations from brian6091!
+
 ### **How long should you train?**
 
 Effect of fine tuning (both Unet + CLIP) can be seen in the following image, where each image is another 500 steps.
@@ -255,10 +266,6 @@ You can see that with 2500 steps, you already get somewhat good results.
 
 People using dreambooth are used to using lr around `1e-6`, but this is way too small for training LoRAs. **I've tried using 1e-4, and it is OK**. I think these values should be more explored statistically.
 
-### **Training tips in general**
-
-- Discussion by @nitrosocke, can be found here. https://github.com/cloneofsimo/lora/issues/19
-
 ### **What happens to Text Encoder LoRA and Unet LoRA?**
 
 Let's see: the following is only using Unet LoRA:
@@ -289,6 +296,17 @@ With LoRA Text Encoder, Unet, all the schedulers, guidance scale, negative promp
 
 > Left with tuned $\alpha_{unet} = 0.6$, $\alpha_{text} = 0.9$, right with $\alpha_{unet} = 1.0$, $\alpha_{text} = 1.0$.
 
+Here is an extensive visualization on the effect of $\alpha_{unet}$, $\alpha_{text}$, by @brian6091 from [his analysis
+](https://github.com/cloneofsimo/lora/discussions/37)
+
+<!-- #region -->
+<p align="center">
+<img  src="contents/comp_scale_clip_unet.jpg">
+</p>
+<!-- #endregion -->
+
+> "a photo of (S\*)", trained with 21 images, with rank 16 LoRA. More details can be found [here](https://github.com/cloneofsimo/lora/discussions/37)
+
 ---
 
 TODOS

diff --git a/contents/comp_scale_clip_unet.jpg b/contents/comp_scale_clip_unet.jpg
diff --git a/lora_diffusion/lora.py b/lora_diffusion/lora.py
@@ -3,6 +3,7 @@
 
 import numpy as np
 import PIL
+import copy
 import torch
 import torch.nn.functional as F
 
@@ -54,24 +55,27 @@ def inject_trainable_lora(
                         _child_module.in_features,
                         _child_module.out_features,
                         _child_module.bias is not None,
-                        r,
+                        r
                     )
                     _tmp.linear.weight = weight
                     if bias is not None:
                         _tmp.linear.bias = bias
 
                     # switch the module
-                    _module._modules[name] = _tmp
+                    if name == 'to_out.0':
+                        _module._modules['to_out']._modules['0'] = _tmp
+                    else:
+                        _module._modules[name] = _tmp
 
                     require_grad_params.append(
-                        _module._modules[name].lora_up.parameters()
+                        _tmp.lora_up.parameters()
                     )
                     require_grad_params.append(
-                        _module._modules[name].lora_down.parameters()
+                        _tmp.lora_down.parameters()
                     )
 
-                    _module._modules[name].lora_up.weight.requires_grad = True
-                    _module._modules[name].lora_down.weight.requires_grad = True
+                    _tmp.lora_up.weight.requires_grad = True
+                    _tmp.lora_down.weight.requires_grad = True
                     names.append(name)
 
     return require_grad_params, names
@@ -136,6 +140,30 @@ def weight_apply_lora(
                     )
                     _child_module.weight = nn.Parameter(weight)
 
+def weight_self_apply_lora(
+    lora_model, target_replace_module=["CrossAttention", "Attention"], alpha=1.0
+):  
+    # Self apply weights to a lora-injected model, and return normal model
+    for _module in lora_model.modules():
+        if _module.__class__.__name__ in target_replace_module:
+            for name, _child_module in _module.named_modules():
+                if _child_module.__class__.__name__ == "LoraInjectedLinear":
+
+                    weight = _child_module.linear.weight
+                    up_weight = _child_module.lora_up.weight.to(weight.device)
+                    down_weight = _child_module.lora_down.weight.to(weight.device)
+
+                    linear = copy.deepcopy(_child_module.linear)
+                    # W <- W + U * D
+                    linear.weight = torch.nn.Parameter(weight + alpha * (up_weight @ down_weight).type(
+                        weight.dtype
+                    ))
+
+                    if name == 'to_out.0':
+                        _module._modules['to_out']._modules['0'] = linear
+                    else:
+                        _module._modules[name] = linear
+    return lora_model
 
 def monkeypatch_lora(
     model, loras, target_replace_module=["CrossAttention", "Attention"]

diff --git a/setup.py b/setup.py
@@ -6,7 +6,7 @@
 setup(
     name="lora_diffusion",
     py_modules=["lora_diffusion"],
-    version="0.0.3",
+    version="0.0.4",
     description="Low Rank Adaptation for Diffusion Models. Works with Stable Diffusion out-of-the-box.",
     author="Simo Ryu",
     packages=find_packages(),