Update app.py

app.py

@@ -55,52 +55,73 @@ class calculateDuration:
         else:
             print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
 
-def update_selection(evt: gr.SelectData, width, height, selected_lora1, selected_lora2):
-    selected_lora = loras[evt.index]
-    new_placeholder = f"Type a prompt for {selected_lora['title']}"
+def update_selection(evt: gr.SelectData, selected_indices, width, height):
+    selected_index = evt.index
+    selected_indices = selected_indices or []
+    if selected_index in selected_indices:
+        # LoRA is already selected, remove it
+        selected_indices.remove(selected_index)
+    else:
+        if len(selected_indices) < 2:
+            selected_indices.append(selected_index)
+        else:
+            raise gr.Error("You can select up to 2 LoRAs only.")
 
     # Initialize outputs
-    outputs = []
+    selected_info_1 = ""
+    selected_info_2 = ""
+    if len(selected_indices) >= 1:
+        lora1 = loras[selected_indices[0]]
+        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨"
+    if len(selected_indices) >= 2:
+        lora2 = loras[selected_indices[1]]
+        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨"
 
-    if selected_lora1 is None:
-        selected_lora1 = selected_lora
-        selected_lora1_info = f"### LoRA 1 Selected: [{selected_lora1['title']}](https://huggingface.co/{selected_lora1['repo']}) ✨"
-        lora_scale1_visible = True
-        remove_lora1_visible = True
-    elif selected_lora2 is None:
-        selected_lora2 = selected_lora
-        selected_lora2_info = f"### LoRA 2 Selected: [{selected_lora2['title']}](https://huggingface.co/{selected_lora2['repo']}) ✨"
-        lora_scale2_visible = True
-        remove_lora2_visible = True
+    # Update prompt placeholder based on last selected LoRA
+    if selected_indices:
+        last_selected_lora = loras[selected_indices[-1]]
+        new_placeholder = f"Type a prompt for {last_selected_lora['title']}"
     else:
-        raise gr.Error("You can only select up to two LoRAs. Please remove one before selecting another.")
-
-    # Update placeholder
-    placeholder_update = gr.update(placeholder=new_placeholder)
-
-    # For width and height adjustment
-    if "aspect" in selected_lora:
-        if selected_lora["aspect"] == "portrait":
-            width = 768
-            height = 1024
-        elif selected_lora["aspect"] == "landscape":
-            width = 1024
-            height = 768
-        else:
-            width = 1024
-            height = 1024
+        new_placeholder = "Type a prompt after selecting a LoRA"
 
-    return placeholder_update, selected_lora1, selected_lora2, selected_lora1_info, selected_lora2_info, gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), width, height
+    return (
+        gr.update(placeholder=new_placeholder),
+        selected_info_1,
+        selected_info_2,
+        selected_indices,
+        width,
+        height,
+    )
 
-def remove_selected_lora1(selected_lora1, selected_lora1_info):
-    selected_lora1 = None
-    selected_lora1_info = ""
-    return selected_lora1, selected_lora1_info, gr.update(visible=False), gr.update(visible=False)
+def remove_lora_1(selected_indices):
+    selected_indices = selected_indices or []
+    if len(selected_indices) >= 1:
+        selected_indices.pop(0)
+    # Update selected_info_1 and selected_info_2
+    selected_info_1 = ""
+    selected_info_2 = ""
+    if len(selected_indices) >= 1:
+        lora1 = loras[selected_indices[0]]
+        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨"
+    if len(selected_indices) >= 2:
+        lora2 = loras[selected_indices[1]]
+        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨"
+    return selected_info_1, selected_info_2, selected_indices
 
-def remove_selected_lora2(selected_lora2, selected_lora2_info):
-    selected_lora2 = None
-    selected_lora2_info = ""
-    return selected_lora2, selected_lora2_info, gr.update(visible=False), gr.update(visible=False)
+def remove_lora_2(selected_indices):
+    selected_indices = selected_indices or []
+    if len(selected_indices) >= 2:
+        selected_indices.pop(1)
+    # Update selected_info_1 and selected_info_2
+    selected_info_1 = ""
+    selected_info_2 = ""
+    if len(selected_indices) >= 1:
+        lora1 = loras[selected_indices[0]]
+        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨"
+    if len(selected_indices) >= 2:
+        lora2 = loras[selected_indices[1]]
+        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨"
+    return selected_info_1, selected_info_2, selected_indices
 
 @spaces.GPU(duration=70)
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress):
@@ -115,6 +136,7 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
             width=width,
             height=height,
             generator=generator,
+            joint_attention_kwargs={"scale": 1.0},
             output_type="pil",
             good_vae=good_vae,
         ):
@@ -134,59 +156,54 @@ def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps
         width=width,
         height=height,
         generator=generator,
+        joint_attention_kwargs={"scale": 1.0},
         output_type="pil",
     ).images[0]
     return final_image 
-
-def run_lora(prompt, image_input, image_strength, cfg_scale, steps, randomize_seed, seed, width, height, selected_lora1, selected_lora2, lora_scale1, lora_scale2, progress=gr.Progress(track_tqdm=True)):
-    if selected_lora1 is None and selected_lora2 is None:
-        raise gr.Error("You must select at least one LoRA before proceeding.")
     
-    # Build the prompt mash
-    prompt_mash = prompt
+def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, progress=gr.Progress(track_tqdm=True)):
+    if not selected_indices:
+        raise gr.Error("You must select at least one LoRA before proceeding.")
 
-    # Handle trigger words and positions
-    trigger_words = []
-    if selected_lora1 is not None:
-        trigger_word1 = selected_lora1.get("trigger_word", "")
-        if trigger_word1:
-            if selected_lora1.get("trigger_position") == "prepend":
-                trigger_words.insert(0, trigger_word1)
-            else:
-                trigger_words.append(trigger_word1)
-    if selected_lora2 is not None:
-        trigger_word2 = selected_lora2.get("trigger_word", "")
-        if trigger_word2:
-            if selected_lora2.get("trigger_position") == "prepend":
-                trigger_words.insert(0, trigger_word2)
+    selected_loras = [loras[idx] for idx in selected_indices]
+
+    # Build the prompt with trigger words
+    prompt_mash = prompt
+    for lora in selected_loras:
+        trigger_word = lora.get('trigger_word', '')
+        if trigger_word:
+            if lora.get("trigger_position") == "prepend":
+                prompt_mash = f"{trigger_word} {prompt_mash}"
             else:
-                trigger_words.append(trigger_word2)
-    # Combine trigger words with the prompt
-    if trigger_words:
-        prompt_mash = f"{' '.join(trigger_words)} {prompt}"
+                prompt_mash = f"{prompt_mash} {trigger_word}"
 
-    with calculateDuration("Unloading LoRAs"):
+    # Unload previous LoRA weights
+    with calculateDuration("Unloading LoRA"):
         pipe.unload_lora_weights()
         pipe_i2i.unload_lora_weights()
-        
+
     # Load LoRA weights with respective scales
     with calculateDuration("Loading LoRA weights"):
-        if image_input is not None:
-            if selected_lora1 is not None:
-                pipe_i2i.load_lora_weights(selected_lora1['repo'], weight_name=selected_lora1.get('weights'), scale=lora_scale1)
-            if selected_lora2 is not None:
-                pipe_i2i.load_lora_weights(selected_lora2['repo'], weight_name=selected_lora2.get('weights'), scale=lora_scale2)
-        else:
-            if selected_lora1 is not None:
-                pipe.load_lora_weights(selected_lora1['repo'], weight_name=selected_lora1.get('weights'), scale=lora_scale1)
-            if selected_lora2 is not None:
-                pipe.load_lora_weights(selected_lora2['repo'], weight_name=selected_lora2.get('weights'), scale=lora_scale2)
-                
+        for idx, lora in enumerate(selected_loras):
+            lora_path = lora['repo']
+            scale = lora_scale_1 if idx == 0 else lora_scale_2
+            if image_input is not None:
+                if "weights" in lora:
+                    pipe_i2i.load_lora_weights(lora_path, weight_name=lora["weights"], multiplier=scale)
+                else:
+                    pipe_i2i.load_lora_weights(lora_path, multiplier=scale)
+            else:
+                if "weights" in lora:
+                    pipe.load_lora_weights(lora_path, weight_name=lora["weights"], multiplier=scale)
+                else:
+                    pipe.load_lora_weights(lora_path, multiplier=scale)
+
     # Set random seed for reproducibility
     with calculateDuration("Randomizing seed"):
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
-            
+
+    # Generate image
     if image_input is not None:
         final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, seed)
         yield final_image, seed, gr.update(visible=False)
@@ -196,37 +213,37 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, randomize_se
         final_image = None
         step_counter = 0
         for image in image_generator:
-            step_counter += 1
+            step_counter+=1
             final_image = image
             progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
             yield image, seed, gr.update(value=progress_bar, visible=True)
         yield final_image, seed, gr.update(value=progress_bar, visible=False)
-        
+
 def get_huggingface_safetensors(link):
-    split_link = link.split("/")
-    if(len(split_link) == 2):
-        model_card = ModelCard.load(link)
-        base_model = model_card.data.get("base_model")
-        print(base_model)
-        if((base_model != "black-forest-labs/FLUX.1-dev") and (base_model != "black-forest-labs/FLUX.1-schnell")):
-            raise Exception("Not a FLUX LoRA!")
-        image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
-        trigger_word = model_card.data.get("instance_prompt", "")
-        image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
-        fs = HfFileSystem()
-        try:
-            list_of_files = fs.ls(link, detail=False)
-            for file in list_of_files:
-                if(file.endswith(".safetensors")):
-                    safetensors_name = file.split("/")[-1]
-                if (not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))):
-                    image_elements = file.split("/")
-                    image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
-        except Exception as e:
-            print(e)
-            gr.Warning(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
-            raise Exception(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
-        return split_link[1], link, safetensors_name, trigger_word, image_url
+  split_link = link.split("/")
+  if(len(split_link) == 2):
+            model_card = ModelCard.load(link)
+            base_model = model_card.data.get("base_model")
+            print(base_model)
+            if((base_model != "black-forest-labs/FLUX.1-dev") and (base_model != "black-forest-labs/FLUX.1-schnell")):
+                raise Exception("Not a FLUX LoRA!")
+            image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
+            trigger_word = model_card.data.get("instance_prompt", "")
+            image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
+            fs = HfFileSystem()
+            try:
+                list_of_files = fs.ls(link, detail=False)
+                for file in list_of_files:
+                    if(file.endswith(".safetensors")):
+                        safetensors_name = file.split("/")[-1]
+                    if (not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))):
+                      image_elements = file.split("/")
+                      image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
+            except Exception as e:
+              print(e)
+              gr.Warning(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+              raise Exception(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+            return split_link[1], link, safetensors_name, trigger_word, image_url
 
 def check_custom_model(link):
     if(link.startswith("https://")):
@@ -286,8 +303,8 @@ css = '''
 #title img{width: 100px; margin-right: 0.5em}
 #gallery .grid-wrap{height: 10vh}
 #lora_list{background: var(--block-background-fill);padding: 0 1em .3em; font-size: 90%}
-.custom_lora_card .card_internal{display: flex;height: 100px;margin-top: .5em}
-.custom_lora_card .card_internal img{margin-right: 1em}
+.card_internal{display: flex;height: 100px;margin-top: .5em}
+.card_internal img{margin-right: 1em}
 .styler{--form-gap-width: 0px !important}
 #progress{height:30px}
 #progress .generating{display:none}
@@ -299,11 +316,10 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 3600)) as app:
         """<h1><img src="https://huggingface.co/spaces/multimodalart/flux-lora-the-explorer/resolve/main/flux_lora.png" alt="LoRA"> LoRA Lab</h1>""",
         elem_id="title",
     )
-    selected_lora1 = gr.State(None)
-    selected_lora2 = gr.State(None)
+    selected_indices = gr.State([])
     with gr.Row():
         with gr.Column(scale=3):
-            prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Type a prompt after selecting LoRAs")
+            prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Type a prompt after selecting a LoRA")
         with gr.Column(scale=1, elem_id="gen_column"):
             generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
     with gr.Row():
@@ -320,21 +336,18 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 3600)) as app:
                 gr.Markdown("[Check the list of FLUX LoRas](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.1-dev)", elem_id="lora_list")
             custom_lora_info = gr.HTML(visible=False)
             custom_lora_button = gr.Button("Remove custom LoRA", visible=False)
-            # Selected LoRAs section
-            gr.Markdown("### Selected LoRAs")
-            with gr.Row():
-                with gr.Column():
-                    selected_lora1_info = gr.Markdown("", visible=False)
-                    lora_scale1 = gr.Slider(label="LoRA 1 Scale", minimum=0, maximum=3, step=0.01, value=0.95, visible=False)
-                    remove_lora1_button = gr.Button("Remove LoRA 1", visible=False)
-                with gr.Column():
-                    selected_lora2_info = gr.Markdown("", visible=False)
-                    lora_scale2 = gr.Slider(label="LoRA 2 Scale", minimum=0, maximum=3, step=0.01, value=0.95, visible=False)
-                    remove_lora2_button = gr.Button("Remove LoRA 2", visible=False)
         with gr.Column():
             progress_bar = gr.Markdown(elem_id="progress",visible=False)
             result = gr.Image(label="Generated Image")
-
+    with gr.Row():
+        with gr.Column():
+            selected_info_1 = gr.Markdown("")
+            lora_scale_1 = gr.Slider(label="LoRA 1 Scale", minimum=0, maximum=3, step=0.01, value=0.95)
+            remove_button_1 = gr.Button("Remove LoRA 1")
+        with gr.Column():
+            selected_info_2 = gr.Markdown("")
+            lora_scale_2 = gr.Slider(label="LoRA 2 Scale", minimum=0, maximum=3, step=0.01, value=0.95)
+            remove_button_2 = gr.Button("Remove LoRA 2")
     with gr.Row():
         with gr.Accordion("Advanced Settings", open=False):
             with gr.Row():
@@ -352,35 +365,35 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 3600)) as app:
                 with gr.Row():
                     randomize_seed = gr.Checkbox(True, label="Randomize seed")
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
-
+    
     gallery.select(
         update_selection,
-        inputs=[width, height, selected_lora1, selected_lora2],
-        outputs=[prompt, selected_lora1, selected_lora2, selected_lora1_info, selected_lora2_info, lora_scale1, remove_lora1_button, lora_scale2, remove_lora2_button, width, height]
+        inputs=[selected_indices, width, height],
+        outputs=[prompt, selected_info_1, selected_info_2, selected_indices, width, height]
     )
-    remove_lora1_button.click(
-        remove_selected_lora1,
-        inputs=[selected_lora1, selected_lora1_info],
-        outputs=[selected_lora1, selected_lora1_info, lora_scale1, remove_lora1_button]
+    remove_button_1.click(
+        remove_lora_1,
+        inputs=[selected_indices],
+        outputs=[selected_info_1, selected_info_2, selected_indices]
     )
-    remove_lora2_button.click(
-        remove_selected_lora2,
-        inputs=[selected_lora2, selected_lora2_info],
-        outputs=[selected_lora2, selected_lora2_info, lora_scale2, remove_lora2_button]
+    remove_button_2.click(
+        remove_lora_2,
+        inputs=[selected_indices],
+        outputs=[selected_info_1, selected_info_2, selected_indices]
     )
     custom_lora.input(
         add_custom_lora,
         inputs=[custom_lora],
-        outputs=[custom_lora_info, custom_lora_button, gallery, selected_lora1_info, selected_lora2_info, prompt]
+        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info_1, selected_indices, prompt]
     )
     custom_lora_button.click(
         remove_custom_lora,
-        outputs=[custom_lora_info, custom_lora_button, gallery, selected_lora1_info, selected_lora2_info, custom_lora]
+        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info_1, selected_indices, custom_lora]
     )
     gr.on(
         triggers=[generate_button.click, prompt.submit],
         fn=run_lora,
-        inputs=[prompt, input_image, image_strength, cfg_scale, steps, randomize_seed, seed, width, height, selected_lora1, selected_lora2, lora_scale1, lora_scale2],
+        inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height],
         outputs=[result, seed, progress_bar]
     )