Update app.py

app.py

@@ -7,6 +7,7 @@ import spaces
 import torch
 from diffusers import DiffusionPipeline, LCMScheduler
 
+# Load the JSON data
 with open("sdxl_lora.json", "r") as file:
     data = json.load(file)
     sdxl_loras_raw = [
@@ -31,24 +32,19 @@ model_id = "stabilityai/stable-diffusion-xl-base-1.0"
 
 pipe = DiffusionPipeline.from_pretrained(model_id, variant="fp16")
 pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-pipe.load_lora_weights("alvdansen/flash-lora-araminta-k-styles", adapter_name="lora")
 pipe.to(device=DEVICE, dtype=torch.float16)
 
-
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-
-def update_selection(
-    selected_state: gr.SelectData,
-    gr_sdxl_loras,
-):
-
+def update_selection(selected_state: gr.SelectData, gr_sdxl_loras):
     lora_id = gr_sdxl_loras[selected_state.index]["repo"]
     trigger_word = gr_sdxl_loras[selected_state.index]["trigger_word"]
-
     return lora_id, trigger_word
 
+def load_lora_for_style(style_repo):
+    pipe.unload_lora_weights()  # Unload any previously loaded weights
+    pipe.load_lora_weights(style_repo, adapter_name="lora")
 
 @spaces.GPU
 def infer(
@@ -63,19 +59,8 @@ def infer(
     user_lora_weight,
     progress=gr.Progress(track_tqdm=True),
 ):
-    flash_sdxl_id = "jasperai/flash-sdxl"
-
-    new_adapter_id = user_lora_selector.replace("/", "_")
-    loaded_adapters = pipe.get_list_adapters()
-
-    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=new_adapter_id)
-
-    pipe.set_adapters(["lora", new_adapter_id], adapter_weights=[1.0, user_lora_weight])
-    gr.Info("LoRA setup done")
+    # Load the appropriate LoRA weights
+    load_lora_for_style(user_lora_selector)
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -95,19 +80,15 @@ def infer(
 
     return image
 
-
 css = """
-
 h1 {
     text-align: center;
     display:block;
 }
-
 p {
     text-align: justify;
     display:block;
 }
-
 """
 
 if torch.cuda.is_available():
@@ -116,11 +97,9 @@ else:
     power_device = "CPU"
 
 with gr.Blocks(css=css) as demo:
-
     gr.Markdown(
         f"""
     # ⚡ FlashDiffusion: FlashLoRA ⚡
-
     This is an interactive demo of [Flash Diffusion](https://gojasper.github.io/flash-diffusion-project/) **on top of** existing LoRAs.
     
     The distillation method proposed in [Flash Diffusion: Accelerating Any Conditional Diffusion Model for Few Steps Image Generation](http://arxiv.org/abs/2406.02347) *by Clément Chadebec, Onur Tasar, Eyal Benaroche and Benjamin Aubin* from Jasper Research.
@@ -137,11 +116,8 @@ with gr.Blocks(css=css) as demo:
     gr_lora_id = gr.State(value="")
 
     with gr.Row():
-
         with gr.Blocks():
-
             with gr.Column():
-
                 user_lora_selector = gr.Textbox(
                     label="Current Selected LoRA",
                     max_lines=1,
@@ -166,9 +142,7 @@ with gr.Blocks(css=css) as demo:
                 )
 
         with gr.Column():
-
             with gr.Row():
-
                 prompt = gr.Text(
                     label="Prompt",
                     show_label=False,
@@ -183,7 +157,6 @@ with gr.Blocks(css=css) as demo:
             result = gr.Image(label="Result", show_label=False)
 
             with gr.Accordion("Advanced Settings", open=False):
-
                 pre_prompt = gr.Text(
                     label="Pre-Prompt",
                     show_label=True,
@@ -204,7 +177,6 @@ with gr.Blocks(css=css) as demo:
                 randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
                 with gr.Row():
-
                     num_inference_steps = gr.Slider(
                         label="Number of inference steps",
                         minimum=4,
@@ -214,7 +186,6 @@ with gr.Blocks(css=css) as demo:
                     )
 
                 with gr.Row():
-
                     guidance_scale = gr.Slider(
                         label="Guidance Scale",
                         minimum=1,
@@ -242,7 +213,6 @@ with gr.Blocks(css=css) as demo:
             run_button.click,
             seed.change,
             randomize_seed.change,
-            # prompt.change,
             prompt.submit,
             negative_prompt.change,
             negative_prompt.submit,
@@ -261,7 +231,6 @@ with gr.Blocks(css=css) as demo:
             user_lora_weight,
         ],
         outputs=[result],
-        # show_progress="full",
     )
 
     gallery.select(
@@ -279,5 +248,4 @@ with gr.Blocks(css=css) as demo:
         "This demo is only for research purpose. Users are solely responsible for any content they create, and it is their obligation to ensure that it adheres to appropriate and ethical standards."
     )
 
-
-demo.queue().launch()
+demo.queue().launch()