Make the auxilary text encoder patching more reliable with custom projector

Author: isidentical

src/diffusers/loaders.py

@@ -109,6 +109,20 @@ def text_encoder_attn_modules(text_encoder):
     return attn_modules
 
 
+def text_encoder_aux_modules(text_encoder):
+    aux_modules = []
+
+    if isinstance(text_encoder, CLIPTextModel):
+        for i, layer in enumerate(text_encoder.text_model.encoder.layers):
+            mlp_mod = layer.mlp
+            name = f"text_model.encoder.layers.{i}.mlp"
+            aux_modules.append((name, mlp_mod))
+    else:
+        raise ValueError(f"do not know how to get aux modules for: {text_encoder.__class__.__name__}")
+
+    return aux_modules
+
+
 def text_encoder_lora_state_dict(text_encoder):
     state_dict = {}
 
@@ -1079,6 +1093,7 @@ def load_lora_into_text_encoder(cls, state_dict, network_alpha, text_encoder, lo
             text_encoder_lora_state_dict = {
                 k.replace(f"{cls.text_encoder_name}.", ""): v for k, v in state_dict.items() if k in text_encoder_keys
             }
+            text_encoder_lora_state_dict = {**text_encoder_lora_state_dict, **state_dict_aux}
             if len(text_encoder_lora_state_dict) > 0:
                 logger.info(f"Loading {cls.text_encoder_name}.")
 
@@ -1119,13 +1134,26 @@ def load_lora_into_text_encoder(cls, state_dict, network_alpha, text_encoder, lo
                             f"{name}.out_proj.lora_linear_layer.down.weight"
                         ] = text_encoder_lora_state_dict.pop(f"{name}.to_out_lora.down.weight")
 
+                    for name, _ in text_encoder_aux_modules(text_encoder):
+                        text_encoder_lora_state_dict[
+                            f"{name}.fc1.lora_linear_layer.up.weight"
+                        ] = text_encoder_lora_state_dict.pop(f"{name}.fc1.lora.up.weight")
+                        text_encoder_lora_state_dict[
+                            f"{name}.fc2.lora_linear_layer.up.weight"
+                        ] = text_encoder_lora_state_dict.pop(f"{name}.fc2.lora.up.weight")
+
+                        text_encoder_lora_state_dict[
+                            f"{name}.fc1.lora_linear_layer.down.weight"
+                        ] = text_encoder_lora_state_dict.pop(f"{name}.fc1.lora.down.weight")
+                        text_encoder_lora_state_dict[
+                            f"{name}.fc2.lora_linear_layer.down.weight"
+                        ] = text_encoder_lora_state_dict.pop(f"{name}.fc2.lora.down.weight")
+
                 rank = text_encoder_lora_state_dict[
                     "text_model.encoder.layers.0.self_attn.out_proj.lora_linear_layer.up.weight"
                 ].shape[1]
 
                 cls._modify_text_encoder(text_encoder, lora_scale, network_alpha, rank=rank)
-                if state_dict_aux:
-                    cls._load_lora_aux_for_text_encoder(text_encoder, state_dict_aux, network_alpha=network_alpha)
 
                 # set correct dtype & device
                 text_encoder_lora_state_dict = {
@@ -1157,36 +1185,10 @@ def _remove_text_encoder_monkey_patch_classmethod(cls, text_encoder):
                 attn_module.v_proj = attn_module.v_proj.regular_linear_layer
                 attn_module.out_proj = attn_module.out_proj.regular_linear_layer
 
-    @classmethod
-    def _load_lora_aux_for_text_encoder(cls, text_encoder, state_dict, network_alpha=None):
-        lora_grouped_dict = defaultdict(dict)
-        for key, value in state_dict.items():
-            attn_processor_key, sub_key = ".".join(key.split(".")[:-3]), ".".join(key.split(".")[-3:])
-            lora_grouped_dict[attn_processor_key][sub_key] = value
-
-        for key, value_dict in lora_grouped_dict.items():
-            rank = value_dict["lora.down.weight"].shape[0]
-            target_modules = [module for name, module in text_encoder.named_modules() if name == key]
-            if len(target_modules) == 0:
-                logger.warning(f"Could not find module {key} in the model. Skipping.")
-                continue
-
-            target_module = target_modules[0]
-            value_dict = {k.replace("lora.", ""): v for k, v in value_dict.items()}
-            lora_layer = LoRALinearLayer(target_module.in_features, target_module.out_features, rank, network_alpha)
-            lora_layer.load_state_dict(value_dict)
-            lora_layer.to(device=text_encoder.device, dtype=text_encoder.dtype)
-
-            old_forward = target_module.forward
-
-            def make_new_forward(old_forward, lora_layer):
-                def new_forward(x):
-                    return old_forward(x) + lora_layer(x)
-
-                return new_forward
-
-            # Monkey-patch.
-            target_module.forward = make_new_forward(old_forward, lora_layer)
+        for _, aux_module in text_encoder_aux_modules(text_encoder):
+            if isinstance(aux_module.fc1, PatchedLoraProjection):
+                aux_module.fc1 = aux_module.fc1.regular_linear_layer
+                aux_module.fc2 = aux_module.fc2.regular_linear_layer
 
     @classmethod
     def _modify_text_encoder(cls, text_encoder, lora_scale=1, network_alpha=None, rank=4, dtype=None):
@@ -1220,6 +1222,13 @@ def _modify_text_encoder(cls, text_encoder, lora_scale=1, network_alpha=None, ra
             )
             lora_parameters.extend(attn_module.out_proj.lora_linear_layer.parameters())
 
+        for _, aux_module in text_encoder_aux_modules(text_encoder):
+            aux_module.fc1 = PatchedLoraProjection(aux_module.fc1, lora_scale, network_alpha, rank=rank, dtype=dtype)
+            lora_parameters.extend(aux_module.fc1.lora_linear_layer.parameters())
+
+            aux_module.fc2 = PatchedLoraProjection(aux_module.fc2, lora_scale, network_alpha, rank=rank, dtype=dtype)
+            lora_parameters.extend(aux_module.fc2.lora_linear_layer.parameters())
+
         return lora_parameters
 
     @classmethod

tests/models/test_layers_utils.py

@@ -22,6 +22,7 @@
 
 from diffusers.models.attention import GEGLU, AdaLayerNorm, ApproximateGELU
 from diffusers.models.embeddings import get_timestep_embedding
+from diffusers.models.lora import LinearWithLoRA
 from diffusers.models.resnet import Downsample2D, ResnetBlock2D, Upsample2D
 from diffusers.models.transformer_2d import Transformer2DModel
 from diffusers.utils import torch_device
@@ -482,7 +483,7 @@ def test_spatial_transformer_default_ff_layers(self):
 
         assert spatial_transformer_block.transformer_blocks[0].ff.net[0].__class__ == GEGLU
         assert spatial_transformer_block.transformer_blocks[0].ff.net[1].__class__ == nn.Dropout
-        assert spatial_transformer_block.transformer_blocks[0].ff.net[2].__class__ == nn.Linear
+        assert spatial_transformer_block.transformer_blocks[0].ff.net[2].__class__ == LinearWithLoRA
 
         dim = 32
         inner_dim = 128
@@ -506,7 +507,7 @@ def test_spatial_transformer_geglu_approx_ff_layers(self):
 
         assert spatial_transformer_block.transformer_blocks[0].ff.net[0].__class__ == ApproximateGELU
         assert spatial_transformer_block.transformer_blocks[0].ff.net[1].__class__ == nn.Dropout
-        assert spatial_transformer_block.transformer_blocks[0].ff.net[2].__class__ == nn.Linear
+        assert spatial_transformer_block.transformer_blocks[0].ff.net[2].__class__ == LinearWithLoRA
 
         dim = 32
         inner_dim = 128