refactor to support patching LoRA into T5

williamberman · williamberman · commit fde2bf8807b5 · 2023-06-13T14:22:14.000-07:00
diff --git a/examples/dreambooth/train_dreambooth_lora.py b/examples/dreambooth/train_dreambooth_lora.py
@@ -38,7 +38,7 @@
 from torch.utils.data import Dataset
 from torchvision import transforms
 from tqdm.auto import tqdm
-from transformers import AutoTokenizer, PretrainedConfig
+from transformers import AutoTokenizer, CLIPTextModel, PretrainedConfig, T5EncoderModel
 
 import diffusers
 from diffusers import (
@@ -49,7 +49,7 @@
     StableDiffusionPipeline,
     UNet2DConditionModel,
 )
-from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
+from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin, text_encoder_attn_modules
 from diffusers.models.attention_processor import (
     AttnAddedKVProcessor,
     AttnAddedKVProcessor2_0,
@@ -59,7 +59,7 @@
     SlicedAttnAddedKVProcessor,
 )
 from diffusers.optimization import get_scheduler
-from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, check_min_version, is_wandb_available
+from diffusers.utils import check_min_version, is_wandb_available
 from diffusers.utils.import_utils import is_xformers_available
 
 
@@ -857,23 +857,25 @@ def main(args):
     unet_lora_layers = AttnProcsLayers(unet.attn_processors)
 
     # The text encoder comes from 🤗 transformers, so we cannot directly modify it.
-    # So, instead, we monkey-patch the forward calls of its attention-blocks. For this,
-    # we first load a dummy pipeline with the text encoder and then do the monkey-patching.
+    # So, instead, we monkey-patch the forward calls of its attention-blocks.
     text_encoder_lora_layers = None
     if args.train_text_encoder:
         text_lora_attn_procs = {}
-        for name, module in text_encoder.named_modules():
-            if name.endswith(TEXT_ENCODER_ATTN_MODULE):
-                text_lora_attn_procs[name] = LoRAAttnProcessor(
-                    hidden_size=module.out_proj.out_features, cross_attention_dim=None
-                )
+
+        for name, module in text_encoder_attn_modules(text_encoder):
+            if isinstance(text_encoder, CLIPTextModel):
+                hidden_size = module.out_proj.out_features
+                inner_dim = None
+            elif isinstance(text_encoder, T5EncoderModel):
+                hidden_size = module.d_model
+                inner_dim = module.inner_dim
+            else:
+                raise ValueError(f"{text_encoder.__class__.__name__} does not support LoRA training")
+
+            text_lora_attn_procs[name] = LoRAAttnProcessor(hidden_size=hidden_size, inner_dim=inner_dim)
+
         text_encoder_lora_layers = AttnProcsLayers(text_lora_attn_procs)
-        temp_pipeline = DiffusionPipeline.from_pretrained(
-            args.pretrained_model_name_or_path, text_encoder=text_encoder
-        )
-        temp_pipeline._modify_text_encoder(text_lora_attn_procs)
-        text_encoder = temp_pipeline.text_encoder
-        del temp_pipeline
+        LoraLoaderMixin._modify_text_encoder(text_lora_attn_procs, text_encoder)
 
     # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
     def save_model_hook(models, weights, output_dir):
diff --git a/src/diffusers/loaders.py b/src/diffusers/loaders.py
@@ -20,6 +20,7 @@
 import torch
 import torch.nn.functional as F
 from huggingface_hub import hf_hub_download
+from torch import nn
 
 from .models.attention_processor import (
     AttnAddedKVProcessor,
@@ -36,7 +37,6 @@
 from .utils import (
     DIFFUSERS_CACHE,
     HF_HUB_OFFLINE,
-    TEXT_ENCODER_ATTN_MODULE,
     _get_model_file,
     deprecate,
     is_safetensors_available,
@@ -49,7 +49,7 @@
     import safetensors
 
 if is_transformers_available():
-    from transformers import PreTrainedModel, PreTrainedTokenizer
+    from transformers import CLIPTextModel, PreTrainedModel, PreTrainedTokenizer, T5EncoderModel
 
 
 logger = logging.get_logger(__name__)
@@ -67,6 +67,36 @@
 CUSTOM_DIFFUSION_WEIGHT_NAME_SAFE = "pytorch_custom_diffusion_weights.safetensors"
 
 
+class PatchedLoraProjection(nn.Module):
+    def __init__(self, regular_linear_layer, lora_linear_layer, lora_scale=1):
+        super().__init__()
+        self.regular_linear_layer = regular_linear_layer
+        self.lora_linear_layer = lora_linear_layer
+        self.lora_scale = lora_scale
+
+    def forward(self, input):
+        return self.regular_linear_layer(input) + self.lora_scale * self.lora_linear_layer(input)
+
+
+def text_encoder_attn_modules(text_encoder):
+    attn_modules = []
+
+    if isinstance(text_encoder, CLIPTextModel):
+        for i, layer in enumerate(text_encoder.text_model.encoder.layers):
+            name = f"text_model.encoder.layers.{i}.self_attn"
+            mod = layer.self_attn
+            attn_modules.append((name, mod))
+    elif isinstance(text_encoder, T5EncoderModel):
+        for i, block in enumerate(text_encoder.encoder.block):
+            name = f"encoder.block.{i}.layer.0.SelfAttention"
+            mod = block.layer[0].SelfAttention
+            attn_modules.append((name, mod))
+    else:
+        raise ValueError(f"do not know how to get attention modules for: {text_encoder.__class__.__name__}")
+
+    return attn_modules
+
+
 class AttnProcsLayers(torch.nn.Module):
     def __init__(self, state_dict: Dict[str, torch.Tensor]):
         super().__init__()
@@ -942,7 +972,7 @@ def load_lora_weights(self, pretrained_model_name_or_path_or_dict: Union[str, Di
                 attn_procs_text_encoder = self._load_text_encoder_attn_procs(
                     text_encoder_lora_state_dict, network_alpha=network_alpha
                 )
-                self._modify_text_encoder(attn_procs_text_encoder)
+                self._modify_text_encoder(attn_procs_text_encoder, self.text_encoder, self.lora_scale)
 
                 # save lora attn procs of text encoder so that it can be easily retrieved
                 self._text_encoder_lora_attn_procs = attn_procs_text_encoder
@@ -968,20 +998,24 @@ def text_encoder_lora_attn_procs(self):
             return self._text_encoder_lora_attn_procs
         return
 
-    def _remove_text_encoder_monkey_patch(self):
-        # Loop over the CLIPAttention module of text_encoder
-        for name, attn_module in self.text_encoder.named_modules():
-            if name.endswith(TEXT_ENCODER_ATTN_MODULE):
-                # Loop over the LoRA layers
-                for _, text_encoder_attr in self._lora_attn_processor_attr_to_text_encoder_attr.items():
-                    # Retrieve the q/k/v/out projection of CLIPAttention
-                    module = attn_module.get_submodule(text_encoder_attr)
-                    if hasattr(module, "old_forward"):
-                        # restore original `forward` to remove monkey-patch
-                        module.forward = module.old_forward
-                        delattr(module, "old_forward")
-
-    def _modify_text_encoder(self, attn_processors: Dict[str, LoRAAttnProcessor]):
+    @classmethod
+    def _remove_text_encoder_monkey_patch(cls, text_encoder):
+        for _, attn_module in text_encoder_attn_modules(text_encoder):
+            if isinstance(text_encoder, CLIPTextModel):
+                attn_module.q_proj = attn_module.q_proj.regular_linear_layer
+                attn_module.k_proj = attn_module.k_proj.regular_linear_layer
+                attn_module.v_proj = attn_module.v_proj.regular_linear_layer
+                attn_module.out_proj = attn_module.out_proj.regular_linear_layer
+            elif isinstance(text_encoder, T5EncoderModel):
+                attn_module.q = attn_module.q.regular_linear_layer
+                attn_module.k = attn_module.k.regular_linear_layer
+                attn_module.v = attn_module.v.regular_linear_layer
+                attn_module.o = attn_module.o.regular_linear_layer
+            else:
+                raise ValueError(f"{text_encoder.__class__.__name__} does not support LoRA training")
+
+    @classmethod
+    def _modify_text_encoder(cls, attn_processors: Dict[str, LoRAAttnProcessor], text_encoder, lora_scale=1):
         r"""
         Monkey-patches the forward passes of attention modules of the text encoder.
 
@@ -991,40 +1025,29 @@ def _modify_text_encoder(self, attn_processors: Dict[str, LoRAAttnProcessor]):
         """
 
         # First, remove any monkey-patch that might have been applied before
-        self._remove_text_encoder_monkey_patch()
+        cls._remove_text_encoder_monkey_patch(text_encoder)
 
-        # Loop over the CLIPAttention module of text_encoder
-        for name, attn_module in self.text_encoder.named_modules():
-            if name.endswith(TEXT_ENCODER_ATTN_MODULE):
-                # Loop over the LoRA layers
-                for attn_proc_attr, text_encoder_attr in self._lora_attn_processor_attr_to_text_encoder_attr.items():
-                    # Retrieve the q/k/v/out projection of CLIPAttention and its corresponding LoRA layer.
-                    module = attn_module.get_submodule(text_encoder_attr)
-                    lora_layer = attn_processors[name].get_submodule(attn_proc_attr)
-
-                    # save old_forward to module that can be used to remove monkey-patch
-                    old_forward = module.old_forward = module.forward
-
-                    # create a new scope that locks in the old_forward, lora_layer value for each new_forward function
-                    # for more detail, see https://github.com/huggingface/diffusers/pull/3490#issuecomment-1555059060
-                    def make_new_forward(old_forward, lora_layer):
-                        def new_forward(x):
-                            result = old_forward(x) + self.lora_scale * lora_layer(x)
-                            return result
-
-                        return new_forward
-
-                    # Monkey-patch.
-                    module.forward = make_new_forward(old_forward, lora_layer)
-
-    @property
-    def _lora_attn_processor_attr_to_text_encoder_attr(self):
-        return {
-            "to_q_lora": "q_proj",
-            "to_k_lora": "k_proj",
-            "to_v_lora": "v_proj",
-            "to_out_lora": "out_proj",
-        }
+        for name, attn_module in text_encoder_attn_modules(text_encoder):
+            if isinstance(text_encoder, CLIPTextModel):
+                attn_module.q_proj = PatchedLoraProjection(
+                    attn_module.q_proj, attn_processors[name].to_q_lora, lora_scale
+                )
+                attn_module.k_proj = PatchedLoraProjection(
+                    attn_module.k_proj, attn_processors[name].to_k_lora, lora_scale
+                )
+                attn_module.v_proj = PatchedLoraProjection(
+                    attn_module.v_proj, attn_processors[name].to_v_lora, lora_scale
+                )
+                attn_module.out_proj = PatchedLoraProjection(
+                    attn_module.out_proj, attn_processors[name].to_out_lora, lora_scale
+                )
+            elif isinstance(text_encoder, T5EncoderModel):
+                attn_module.q = PatchedLoraProjection(attn_module.q, attn_processors[name].to_q_lora, lora_scale)
+                attn_module.k = PatchedLoraProjection(attn_module.k, attn_processors[name].to_k_lora, lora_scale)
+                attn_module.v = PatchedLoraProjection(attn_module.v, attn_processors[name].to_v_lora, lora_scale)
+                attn_module.o = PatchedLoraProjection(attn_module.o, attn_processors[name].to_out_lora, lora_scale)
+            else:
+                raise ValueError(f"{text_encoder.__class__.__name__} does not support LoRA training")
 
     def _load_text_encoder_attn_procs(
         self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **kwargs
diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py
@@ -549,17 +549,19 @@ class LoRAAttnProcessor(nn.Module):
             Equivalent to `alpha` but it's usage is specific to Kohya (A1111) style LoRAs.
     """
 
-    def __init__(self, hidden_size, cross_attention_dim=None, rank=4, network_alpha=None):
+    def __init__(self, hidden_size, cross_attention_dim=None, inner_dim=None, rank=4, network_alpha=None):
         super().__init__()
 
         self.hidden_size = hidden_size
         self.cross_attention_dim = cross_attention_dim
         self.rank = rank
+        if inner_dim is None:
+            inner_dim = hidden_size
 
-        self.to_q_lora = LoRALinearLayer(hidden_size, hidden_size, rank, network_alpha)
-        self.to_k_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank, network_alpha)
-        self.to_v_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank, network_alpha)
-        self.to_out_lora = LoRALinearLayer(hidden_size, hidden_size, rank, network_alpha)
+        self.to_q_lora = LoRALinearLayer(hidden_size, inner_dim, rank, network_alpha)
+        self.to_k_lora = LoRALinearLayer(cross_attention_dim or hidden_size, inner_dim, rank, network_alpha)
+        self.to_v_lora = LoRALinearLayer(cross_attention_dim or hidden_size, inner_dim, rank, network_alpha)
+        self.to_out_lora = LoRALinearLayer(inner_dim, hidden_size, rank, network_alpha)
 
     def __call__(
         self, attn: Attention, hidden_states, encoder_hidden_states=None, attention_mask=None, scale=1.0, temb=None
diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py
@@ -30,7 +30,6 @@
     ONNX_EXTERNAL_WEIGHTS_NAME,
     ONNX_WEIGHTS_NAME,
     SAFETENSORS_WEIGHTS_NAME,
-    TEXT_ENCODER_ATTN_MODULE,
     WEIGHTS_NAME,
 )
 from .deprecation_utils import deprecate
diff --git a/src/diffusers/utils/constants.py b/src/diffusers/utils/constants.py
@@ -30,4 +30,3 @@
 DIFFUSERS_DYNAMIC_MODULE_NAME = "diffusers_modules"
 HF_MODULES_CACHE = os.getenv("HF_MODULES_CACHE", os.path.join(hf_cache_home, "modules"))
 DEPRECATED_REVISION_ARGS = ["fp16", "non-ema"]
-TEXT_ENCODER_ATTN_MODULE = ".self_attn"
diff --git a/tests/models/test_lora_layers.py b/tests/models/test_lora_layers.py
@@ -23,7 +23,7 @@
 from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
 
 from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNet2DConditionModel
-from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
+from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin, text_encoder_attn_modules
 from diffusers.models.attention_processor import (
     Attention,
     AttnProcessor,
@@ -33,7 +33,7 @@
     LoRAXFormersAttnProcessor,
     XFormersAttnProcessor,
 )
-from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device
+from diffusers.utils import floats_tensor, torch_device
 
 
 def create_unet_lora_layers(unet: nn.Module):
@@ -63,11 +63,10 @@ def create_text_encoder_lora_attn_procs(text_encoder: nn.Module):
     lora_attn_processor_class = (
         LoRAAttnProcessor2_0 if hasattr(F, "scaled_dot_product_attention") else LoRAAttnProcessor
     )
-    for name, module in text_encoder.named_modules():
-        if name.endswith(TEXT_ENCODER_ATTN_MODULE):
-            text_lora_attn_procs[name] = lora_attn_processor_class(
-                hidden_size=module.out_proj.out_features, cross_attention_dim=None
-            )
+    for name, module in text_encoder_attn_modules(text_encoder):
+        text_lora_attn_procs[name] = lora_attn_processor_class(
+            hidden_size=module.out_proj.out_features, cross_attention_dim=None
+        )
     return text_lora_attn_procs
 
 
@@ -286,7 +285,7 @@ def test_text_encoder_lora_monkey_patch(self):
         set_lora_up_weights(text_attn_procs, randn_weight=False)
 
         # monkey patch
-        pipe._modify_text_encoder(text_attn_procs)
+        pipe._modify_text_encoder(text_attn_procs, pipe.text_encoder, pipe.lora_scale)
 
         # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
         del text_attn_procs
@@ -305,7 +304,7 @@ def test_text_encoder_lora_monkey_patch(self):
         set_lora_up_weights(text_attn_procs, randn_weight=True)
 
         # monkey patch
-        pipe._modify_text_encoder(text_attn_procs)
+        pipe._modify_text_encoder(text_attn_procs, pipe.text_encoder, pipe.lora_scale)
 
         # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
         del text_attn_procs
@@ -334,7 +333,7 @@ def test_text_encoder_lora_remove_monkey_patch(self):
         set_lora_up_weights(text_attn_procs, randn_weight=True)
 
         # monkey patch
-        pipe._modify_text_encoder(text_attn_procs)
+        pipe._modify_text_encoder(text_attn_procs, pipe.text_encoder, pipe.lora_scale)
 
         # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
         del text_attn_procs

Original file line number	Diff line number	Diff line change
`@@ -30,7 +30,6 @@`
`30`	`30`	`ONNX_EXTERNAL_WEIGHTS_NAME,`
`31`	`31`	`ONNX_WEIGHTS_NAME,`
`32`	`32`	`SAFETENSORS_WEIGHTS_NAME,`
`33`		`- TEXT_ENCODER_ATTN_MODULE,`
`34`	`33`	`WEIGHTS_NAME,`
`35`	`34`	`)`
`36`	`35`	`from .deprecation_utils import deprecate`