fix LoRA swap on Zero GPU

app.py

@@ -39,25 +39,6 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 
-def check_and_load_lora_user(user_lora_selector, user_lora_weight, gr_lora_loaded):
-    flash_sdxl_id = "jasperai/flash-sdxl"
-
-    if user_lora_selector == "" or user_lora_selector == "":
-        raise gr.Error("Please select a LoRA before running the inference.")
-
-    if gr_lora_loaded != user_lora_selector:
-        gr.Info("Loading LoRA")
-        pipe.unload_lora_weights()
-        pipe.load_lora_weights(flash_sdxl_id, adapter_name="lora")
-        pipe.load_lora_weights(user_lora_selector, adapter_name="user")
-        pipe.set_adapters(["lora", "user"], adapter_weights=[1.0, user_lora_weight])
-        gr.Info("LoRA Loaded")
-
-        gr_lora_loaded = user_lora_selector
-
-    return gr_lora_loaded
-
-
 def rescale_lora(user_lora_weight):
 
     global pipe
@@ -86,22 +67,26 @@ def infer(
     guidance_scale,
     user_lora_selector,
     user_lora_weight,
-    gr_lora_loaded,
+    progress=gr.Progress(track_tqdm=True)
 ):
     flash_sdxl_id = "jasperai/flash-sdxl"
 
-    if user_lora_selector == "" or user_lora_selector == "":
-        raise gr.Error("Please select a LoRA before running the inference.")
+    gr.Info("Checking LoRA")
 
-    if gr_lora_loaded != user_lora_selector:
-        gr.Info("Loading LoRA")
+    new_adapter_id = user_lora_selector.replace("/", "_")
+    loaded_adapters = pipe.get_list_adapters()
+
+    print(loaded_adapters["unet"])
+    print(new_adapter_id)
+
+    if new_adapter_id not in loaded_adapters["unet"]:
+        gr.Info("Swapping LoRA")
         pipe.unload_lora_weights()
         pipe.load_lora_weights(flash_sdxl_id, adapter_name="lora")
-        pipe.load_lora_weights(user_lora_selector, adapter_name="user")
-        pipe.set_adapters(["lora", "user"], adapter_weights=[1.0, user_lora_weight])
-        gr.Info("LoRA Loaded")
+        pipe.load_lora_weights(user_lora_selector, adapter_name=new_adapter_id)
 
-        gr_lora_loaded = user_lora_selector
+    pipe.set_adapters(["lora", new_adapter_id], adapter_weights=[1.0, user_lora_weight])
+    gr.Info("LoRA setup done")
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -111,6 +96,8 @@ def infer(
     if pre_prompt != "":
         prompt = f"{pre_prompt} {prompt}"
 
+    gr.Info("Generation Stage")
+
     image = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
@@ -119,7 +106,7 @@ def infer(
         generator=generator,
     ).images[0]
 
-    return image, gr_lora_loaded
+    return image 
 
 
 css = """
@@ -160,7 +147,6 @@ with gr.Blocks(css=css) as demo:
     # Index of selected LoRA
     gr_sdxl_loras = gr.State(value=sdxl_loras_raw)
     # Serve as memory for currently loaded lora in pipe
-    gr_lora_loaded = gr.State(value="")
     gr_lora_id = gr.State(value="")
 
     with gr.Row():
@@ -285,11 +271,10 @@ with gr.Blocks(css=css) as demo:
             negative_prompt,
             guidance_scale,
             user_lora_selector,
-            user_lora_weight,
-            gr_lora_loaded,
+            user_lora_weight
         ],
-        outputs=[result, gr_lora_loaded],
-        show_progress="minimal",
+        outputs=[result],
+        # show_progress="full",
     )
 
     user_lora_weight.change(