<feat> unify hacked_lora_forward

app.py

@@ -99,79 +99,72 @@ def init_basemodel():
         image_processor=image_processor,
     )
 
+    # insert LoRA
+    lora_config = LoraConfig(
+        r=16,
+        lora_alpha=16,
+        init_lora_weights="gaussian",
+        target_modules=[
+            'attn.to_k', 'attn.to_q', 'attn.to_v', 'attn.to_out.0',
+            'attn.add_k_proj', 'attn.add_q_proj', 'attn.add_v_proj', 'attn.to_add_out',
+            'ff.net.0.proj', 'ff.net.2',
+            'ff_context.net.0.proj', 'ff_context.net.2',
+            'norm1_context.linear', 'norm1.linear',
+            'norm.linear', 'proj_mlp', 'proj_out',
+        ]
+    )
+    transformer.add_adapter(lora_config)
+
+    # hack LoRA forward
+    def create_hacked_forward(module):
+        if not hasattr(module, 'original_forward'):
+            module.original_forward = module.forward
+        lora_forward = module.forward
+        non_lora_forward = module.base_layer.forward
+        img_sequence_length = int((512 / 8 / 2) ** 2)
+        encoder_sequence_length = 144 + 252 # encoder sequence: 144 img 252 txt
+        num_imgs = 4
+        num_generated_imgs = 3
+
+        def hacked_lora_forward(self, x, *args, **kwargs):
+            if x.shape[1] == img_sequence_length * num_imgs and len(x.shape) > 2:
+                return torch.cat((
+                    lora_forward(x[:, :-img_sequence_length*num_generated_imgs], *args, **kwargs),
+                    non_lora_forward(x[:, -img_sequence_length*num_generated_imgs:], *args, **kwargs)
+                ), dim=1)
+            elif x.shape[1] == encoder_sequence_length * 2 or x.shape[1] == encoder_sequence_length:
+                return lora_forward(x, *args, **kwargs)
+            elif x.shape[1] == img_sequence_length * num_imgs + encoder_sequence_length:
+                return torch.cat((
+                    lora_forward(x[:, :(num_imgs - num_generated_imgs)*img_sequence_length], *args, **kwargs),
+                    non_lora_forward(x[:, (num_imgs - num_generated_imgs)*img_sequence_length:-encoder_sequence_length], *args, **kwargs),
+                    lora_forward(x[:, -encoder_sequence_length:], *args, **kwargs)
+                ), dim=1)
+            elif x.shape[1] == img_sequence_length * num_imgs + encoder_sequence_length * 2:
+                return torch.cat((
+                    lora_forward(x[:, :(num_imgs - num_generated_imgs)*img_sequence_length], *args, **kwargs),
+                    non_lora_forward(x[:, (num_imgs - num_generated_imgs)*img_sequence_length:-2*encoder_sequence_length], *args, **kwargs),
+                    lora_forward(x[:, -2*encoder_sequence_length:], *args, **kwargs)
+                ), dim=1)
+            elif x.shape[1] == 3072:
+                return non_lora_forward(x, *args, **kwargs)
+            else:
+                raise ValueError(
+                    f"hacked_lora_forward receives unexpected sequence length: {x.shape[1]}, input shape: {x.shape}!"
+                )
+
+        return hacked_lora_forward.__get__(module, type(module))
+
+    for n, m in transformer.named_modules():
+        if isinstance(m, peft.tuners.lora.layer.Linear):
+            m.forward = create_hacked_forward(m)
+
 
 @spaces.GPU
 def process_image_and_text(condition_image, target_prompt, condition_image_prompt, task, random_seed, num_steps, inpainting, fill_x1, fill_x2, fill_y1, fill_y2):
     # set up the model
     global pipe, current_task, transformer
     if current_task != task:
-        if current_task is None:
-            # insert LoRA
-            lora_config = LoraConfig(
-                r=16,
-                lora_alpha=16,
-                init_lora_weights="gaussian",
-                target_modules=[
-                    'attn.to_k', 'attn.to_q', 'attn.to_v', 'attn.to_out.0',
-                    'attn.add_k_proj', 'attn.add_q_proj', 'attn.add_v_proj', 'attn.to_add_out',
-                    'ff.net.0.proj', 'ff.net.2',
-                    'ff_context.net.0.proj', 'ff_context.net.2',
-                    'norm1_context.linear', 'norm1.linear',
-                    'norm.linear', 'proj_mlp', 'proj_out',
-                ]
-            )
-            transformer.add_adapter(lora_config)
-        else:
-            def restore_forward(module):
-                def restored_forward(self, x, *args, **kwargs):
-                    return module.original_forward(x, *args, **kwargs)
-                return restored_forward.__get__(module, type(module))
-            
-            for n, m in transformer.named_modules():
-                if isinstance(m, peft.tuners.lora.layer.Linear):
-                    m.forward = restore_forward(m)
-
-        current_task = task
-
-        # hack LoRA forward
-        def create_hacked_forward(module):
-            if not hasattr(module, 'original_forward'):
-                module.original_forward = module.forward
-            lora_forward = module.forward
-            non_lora_forward = module.base_layer.forward
-            img_sequence_length = int((512 / 8 / 2) ** 2)
-            encoder_sequence_length = 144 + 252 # encoder sequence: 144 img 252 txt
-            num_imgs = 4
-            num_generated_imgs = 3
-            num_encoder_sequences = 2 if task in ['subject_driven', 'style_transfer'] else 1
-
-            def hacked_lora_forward(self, x, *args, **kwargs):
-                if x.shape[1] == img_sequence_length * num_imgs and len(x.shape) > 2:
-                    return torch.cat((
-                        lora_forward(x[:, :-img_sequence_length*num_generated_imgs], *args, **kwargs),
-                        non_lora_forward(x[:, -img_sequence_length*num_generated_imgs:], *args, **kwargs)
-                    ), dim=1)
-                elif x.shape[1] == encoder_sequence_length * num_encoder_sequences or x.shape[1] == encoder_sequence_length:
-                    return lora_forward(x, *args, **kwargs)
-                elif x.shape[1] == img_sequence_length * num_imgs + encoder_sequence_length * num_encoder_sequences:
-                    return torch.cat((
-                        lora_forward(x[:, :(num_imgs - num_generated_imgs)*img_sequence_length], *args, **kwargs),
-                        non_lora_forward(x[:, (num_imgs - num_generated_imgs)*img_sequence_length:-num_encoder_sequences*encoder_sequence_length], *args, **kwargs),
-                        lora_forward(x[:, -num_encoder_sequences*encoder_sequence_length:], *args, **kwargs)
-                    ), dim=1)
-                elif x.shape[1] == 3072:
-                    return non_lora_forward(x, *args, **kwargs)
-                else:
-                    raise ValueError(
-                        f"hacked_lora_forward receives unexpected sequence length: {x.shape[1]}, input shape: {x.shape}!"
-                    )
-
-            return hacked_lora_forward.__get__(module, type(module))
-
-        for n, m in transformer.named_modules():
-            if isinstance(m, peft.tuners.lora.layer.Linear):
-                m.forward = create_hacked_forward(m)
-
         # load LoRA weights
         model_root = hf_hub_download(
             repo_id="Kunbyte/DRA-Ctrl",