Update app.py

app.py

@@ -57,55 +57,14 @@ class calculateDuration:
         
 @spaces.GPU(duration=120)
 @torch.inference_mode()
-def generate_image(prompt, lora_strings_json,  steps, seed, cfg_scale, width, height, progress):
-
-    lora_configs = None
-    adapter_names = []
-    if lora_strings_json:
-        try:
-            lora_configs = json.loads(lora_strings_json)
-        except:
-            gr.Warning("Parse lora config json failed")
-            print("parse lora config json failed")
-            
-        if lora_configs:
-            with calculateDuration("Loading LoRA weights"):
-                adapter_weights = []
-                for lora_info in lora_configs:
-                    lora_repo = lora_info.get("repo")
-                    weights = lora_info.get("weights")
-                    adapter_name = lora_info.get("adapter_name")
-                    adapter_weight = lora_info.get("adapter_weight")
-                    
-                    if lora_repo and weights and adapter_name:
-                        retry_count = 3
-                        for attempt in range(retry_count):
-                            try:
-                                pipe.load_lora_weights(lora_repo, weight_name=weights, adapter_name=adapter_name)
-                                adapter_names.append(adapter_name)
-                                adapter_weights.append(adapter_weight)
-                                break  # Load successful, exit retry loop
-                            except ValueError as e:
-                                print(f"Attempt {attempt+1}/{retry_count} failed to load LoRA adapter: {e}")
-                                if attempt == retry_count - 1:
-                                    print(f"Error loading LoRA adapter: {adapter_name} after {retry_count} attempts")
-                                else:
-                                    time.sleep(1)  # Wait before retrying
-                        
-                # set lora weights
-                if len(adapter_names) > 0:
-                    pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
+def generate_image(prompt, adapter_names,  steps, seed, cfg_scale, width, height, progress):
 
     
     gr.Info("Start to generate images ...")
-
     with calculateDuration(f"Make a new generator:{seed}"):
         pipe.to(device)
         generator = torch.Generator(device=device).manual_seed(seed)
     
-    if len(adapter_names) > 0:
-        pipe.fuse_lora(adapter_names=adapter_names)
-        
     with calculateDuration("Generating image"):
         # Generate image
         generated_image = pipe(
@@ -119,8 +78,6 @@ def generate_image(prompt, lora_strings_json,  steps, seed, cfg_scale, width, he
         ).images[0]
     
     progress(99, "Generate image success!")
-    if len(adapter_names) > 0:
-        pipe.unfuse_lora()
     return generated_image
 
 
@@ -173,12 +130,49 @@ def run_lora(prompt, lora_strings_json,  cfg_scale, steps, randomize_seed, seed,
 
     # Load LoRA weights
     gr.Info("Start to load loras ...")
+    lora_configs = None
+    adapter_names = []
+    if lora_strings_json:
+        try:
+            lora_configs = json.loads(lora_strings_json)
+        except:
+            gr.Warning("Parse lora config json failed")
+            print("parse lora config json failed")
+            
+        if lora_configs:
+            with calculateDuration("Loading LoRA weights"):
+                adapter_weights = []
+                for lora_info in lora_configs:
+                    lora_repo = lora_info.get("repo")
+                    weights = lora_info.get("weights")
+                    adapter_name = lora_info.get("adapter_name")
+                    adapter_weight = lora_info.get("adapter_weight")
+                    
+                    if lora_repo and weights and adapter_name:
+                        retry_count = 3
+                        for attempt in range(retry_count):
+                            try:
+                                pipe.load_lora_weights(lora_repo, weight_name=weights, adapter_name=adapter_name)
+                                adapter_names.append(adapter_name)
+                                adapter_weights.append(adapter_weight)
+                                break  # Load successful, exit retry loop
+                            except ValueError as e:
+                                print(f"Attempt {attempt+1}/{retry_count} failed to load LoRA adapter: {e}")
+                                if attempt == retry_count - 1:
+                                    print(f"Error loading LoRA adapter: {adapter_name} after {retry_count} attempts")
+                                else:
+                                    time.sleep(1)  # Wait before retrying
+                        
+                # set lora weights
+                if len(adapter_names) > 0:
+                    pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
+
     
     # Generate image
     error_message = ""
     try:
         print("Start applying for zeroGPU resources")
-        final_image = generate_image(prompt, lora_strings_json, steps, seed, cfg_scale, width, height, progress)
+        final_image = generate_image(prompt, adapter_names, steps, seed, cfg_scale, width, height, progress)
     except Exception as e:
         error_message =  str(e)
         gr.Error(error_message)