<feat> support changing models.

README.md

@@ -9,6 +9,7 @@ app_file: app.py
 pinned: false
 license: apache-2.0
 short_description: huggingface space for DRA-Ctrl.
+entrypoint: app.sh
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,8 +1,4 @@
 import os
-
-if 'SPACES_APP' in os.environ:
-    os.system("pip install flash-attn==2.7.3 --no-build-isolation")
-
 import sys
 import torch
 import diffusers
@@ -102,35 +98,40 @@ def init_basemodel():
 @spaces.GPU
 def process_image_and_text(condition_image, target_prompt, condition_image_prompt, task, random_seed, inpainting, fill_x1, fill_x2, fill_y1, fill_y2):
     # set up the model
-    global pipe, current_task
-    if pipe is None or current_task != task:
-        if current_task is not None:
-            transformer.delete_adapters('default')
+    global pipe, current_task, transformer
+    if current_task != task:
+        if current_task is None:
+            current_task = task
+
+            # 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):
-                    if hasattr(m, 'base_layer'):
-                        m.forward = m.base_layer.forward.__get__(m, type(m))
-
-        current_task = task
-
-        # 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)
+                    m.forward = restore_forward(m)
 
         # 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)

app.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+pip install torch==2.5.1
+pip install git+https://github.com/Dao-AILab/[email protected]#subdirectory=csrc
+python app.py