Update app.py

app.py

@@ -15,7 +15,7 @@ import time
 import requests
 import pandas as pd
 
-# Load prompts for randomization
+#Load prompts for randomization
 df = pd.read_csv('prompts.csv', header=None)
 prompt_values = df.values.flatten()
 
@@ -88,19 +88,26 @@ def update_selection(evt: gr.SelectData, selected_indices, loras_state, width, h
     if selected_index in selected_indices:
         selected_indices.remove(selected_index)
     else:
-        selected_indices.append(selected_index)
-
-    selected_infos = []
-    lora_scales = []
-    lora_images = []
-    for idx, sel_idx in enumerate(selected_indices):
-        lora = loras_state[sel_idx]
-        selected_info = f"### LoRA {idx+1} Selected: [{lora['title']}](https://huggingface.co/{lora['repo']}) ✨"
-        lora_scale = 1.15  # Default value
-        lora_image = lora['image']
-        selected_infos.append(selected_info)
-        lora_scales.append(lora_scale)
-        lora_images.append(lora_image)
+        if len(selected_indices) < 2:
+            selected_indices.append(selected_index)
+        else:
+            gr.Warning("You can select up to 2 LoRAs, remove one to select a new one.")
+            return gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), width, height, gr.update(), gr.update()
+
+    selected_info_1 = "Select a LoRA 1"
+    selected_info_2 = "Select a LoRA 2"
+    lora_scale_1 = 1.15
+    lora_scale_2 = 1.15
+    lora_image_1 = None
+    lora_image_2 = None
+    if len(selected_indices) >= 1:
+        lora1 = loras_state[selected_indices[0]]
+        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨"
+        lora_image_1 = lora1['image']
+    if len(selected_indices) >= 2:
+        lora2 = loras_state[selected_indices[1]]
+        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨"
+        lora_image_2 = lora2['image']
 
     if selected_indices:
         last_selected_lora = loras_state[selected_indices[-1]]
@@ -108,81 +115,60 @@ def update_selection(evt: gr.SelectData, selected_indices, loras_state, width, h
     else:
         new_placeholder = "Type a prompt after selecting a LoRA"
 
-    # Update the number of visible LoRA slots based on the number of selected indices
-    num_selected = len(selected_indices)
-    lora_slot_updates = []
-    for i in range(num_selected):
-        lora_slot_updates.append(gr.update(visible=True))
-    for i in range(num_selected, 5):  # Assuming max 5 LoRAs
-        lora_slot_updates.append(gr.update(visible=False))
-
-    return (
-        gr.update(placeholder=new_placeholder),
-        selected_infos,
-        selected_indices,
-        lora_scales,
-        width,
-        height,
-        lora_images,
-        lora_slot_updates
-    )
-
-def add_lora_slot(num_loras_visible):
-    if num_loras_visible < 5:  # Assuming max 5 LoRAs
-        num_loras_visible += 1
-    else:
-        gr.Warning("Maximum number of LoRAs reached.")
-    lora_slot_updates = []
-    for i in range(num_loras_visible):
-        lora_slot_updates.append(gr.update(visible=True))
-    for i in range(num_loras_visible, 5):
-        lora_slot_updates.append(gr.update(visible=False))
-    return num_loras_visible, lora_slot_updates
-
-def remove_lora_slot(idx_to_remove, selected_indices):
-    if idx_to_remove < len(selected_indices):
-        selected_indices.pop(idx_to_remove)
-    selected_infos = []
-    lora_scales = []
-    lora_images = []
-    for idx, sel_idx in enumerate(selected_indices):
-        lora = loras_state[sel_idx]
-        selected_info = f"### LoRA {idx+1} Selected: [{lora['title']}]({lora['repo']}) ✨"
-        lora_scale = 1.15  # Default value
-        lora_image = lora['image']
-        selected_infos.append(selected_info)
-        lora_scales.append(lora_scale)
-        lora_images.append(lora_image)
-    num_loras_visible = len(selected_indices)
-    lora_slot_updates = []
-    for i in range(num_loras_visible):
-        lora_slot_updates.append(gr.update(visible=True))
-    for i in range(num_loras_visible, 5):
-        lora_slot_updates.append(gr.update(visible=False))
-    return selected_infos, selected_indices, lora_scales, lora_images, num_loras_visible, lora_slot_updates
+    return gr.update(placeholder=new_placeholder), selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, width, height, lora_image_1, lora_image_2
+
+def remove_lora_1(selected_indices, loras_state):
+    if len(selected_indices) >= 1:
+        selected_indices.pop(0)
+    selected_info_1 = "Select a LoRA 1"
+    selected_info_2 = "Select a LoRA 2"
+    lora_scale_1 = 1.15
+    lora_scale_2 = 1.15
+    lora_image_1 = None
+    lora_image_2 = None
+    if len(selected_indices) >= 1:
+        lora1 = loras_state[selected_indices[0]]
+        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨"
+        lora_image_1 = lora1['image']
+    if len(selected_indices) >= 2:
+        lora2 = loras_state[selected_indices[1]]
+        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨"
+        lora_image_2 = lora2['image']
+    return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2
+
+def remove_lora_2(selected_indices, loras_state):
+    if len(selected_indices) >= 2:
+        selected_indices.pop(1)
+    selected_info_1 = "Select LoRA 1"
+    selected_info_2 = "Select LoRA 2"
+    lora_scale_1 = 1.15
+    lora_scale_2 = 1.15
+    lora_image_1 = None
+    lora_image_2 = None
+    if len(selected_indices) >= 1:
+        lora1 = loras_state[selected_indices[0]]
+        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨"
+        lora_image_1 = lora1['image']
+    if len(selected_indices) >= 2:
+        lora2 = loras_state[selected_indices[1]]
+        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨"
+        lora_image_2 = lora2['image']
+    return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2
 
 def randomize_loras(selected_indices, loras_state):
-    num_to_select = random.randint(1, 5)
-    selected_indices = random.sample(range(len(loras_state)), num_to_select)
-    selected_infos = []
-    lora_scales = []
-    lora_images = []
-    for idx, sel_idx in enumerate(selected_indices):
-        lora = loras_state[sel_idx]
-        selected_info = f"### LoRA {idx+1} Selected: [{lora['title']}](https://huggingface.co/{lora['repo']}) ✨"
-        lora_scale = 1.15  # Default value
-        lora_image = lora['image']
-        selected_infos.append(selected_info)
-        lora_scales.append(lora_scale)
-        lora_images.append(lora_image)
+    if len(loras_state) < 2:
+        raise gr.Error("Not enough LoRAs to randomize.")
+    selected_indices = random.sample(range(len(loras_state)), 2)
+    lora1 = loras_state[selected_indices[0]]
+    lora2 = loras_state[selected_indices[1]]
+    selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨"
+    selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨"
+    lora_scale_1 = 1.15
+    lora_scale_2 = 1.15
+    lora_image_1 = lora1['image']
+    lora_image_2 = lora2['image']
     random_prompt = random.choice(prompt_values)
-    num_loras_visible = len(selected_indices)
-    lora_slot_updates = []
-    for i in range(num_loras_visible):
-        lora_slot_updates.append(gr.update(visible=True))
-    for i in range(num_loras_visible, 5):
-        lora_slot_updates.append(gr.update(visible=False))
-    return selected_infos, selected_indices, lora_scales, lora_images, random_prompt, num_loras_visible, lora_slot_updates
+    return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2, random_prompt
 
 def add_custom_lora(custom_lora, selected_indices, current_loras, gallery):
     if custom_lora:
@@ -206,44 +192,83 @@ def add_custom_lora(custom_lora, selected_indices, current_loras, gallery):
             
             # Update gallery
             gallery_items = [(item["image"], item["title"]) for item in current_loras]
-            # Update selected_indices
-            selected_indices.append(existing_item_index)
+            # Update selected_indices if there's room
+            if len(selected_indices) < 2:
+                selected_indices.append(existing_item_index)
+            else:
+                gr.Warning("You can select up to 2 LoRAs, remove one to select a new one.")
 
             # Update selected_info and images
-            selected_infos = []
-            lora_scales = []
-            lora_images = []
-            for idx, sel_idx in enumerate(selected_indices):
-                lora = current_loras[sel_idx]
-                selected_info = f"### LoRA {idx+1} Selected: {lora['title']} ✨"
-                lora_scale = 1.15  # Default value
-                lora_image = lora['image'] if lora['image'] else None
-                selected_infos.append(selected_info)
-                lora_scales.append(lora_scale)
-                lora_images.append(lora_image)
-            num_loras_visible = len(selected_indices)
-            lora_slot_updates = []
-            for i in range(num_loras_visible):
-                lora_slot_updates.append(gr.update(visible=True))
-            for i in range(num_loras_visible, 5):
-                lora_slot_updates.append(gr.update(visible=False))
+            selected_info_1 = "Select a LoRA 1"
+            selected_info_2 = "Select a LoRA 2"
+            lora_scale_1 = 1.15
+            lora_scale_2 = 1.15
+            lora_image_1 = None
+            lora_image_2 = None
+            if len(selected_indices) >= 1:
+                lora1 = current_loras[selected_indices[0]]
+                selected_info_1 = f"### LoRA 1 Selected: {lora1['title']} ✨"
+                lora_image_1 = lora1['image'] if lora1['image'] else None
+            if len(selected_indices) >= 2:
+                lora2 = current_loras[selected_indices[1]]
+                selected_info_2 = f"### LoRA 2 Selected: {lora2['title']} ✨"
+                lora_image_2 = lora2['image'] if lora2['image'] else None
             print("Finished adding custom LoRA")
             return (
                 current_loras,
                 gr.update(value=gallery_items),
-                selected_infos, 
+                selected_info_1, 
+                selected_info_2,
                 selected_indices,
-                lora_scales,
-                lora_images,
-                num_loras_visible,
-                lora_slot_updates
+                lora_scale_1,
+                lora_scale_2,
+                lora_image_1,
+                lora_image_2
             )
         except Exception as e:
             print(e)
             gr.Warning(str(e))
-            return current_loras, gr.update(), selected_infos, selected_indices, lora_scales, lora_images, 0, []
+            return current_loras, gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update()
     else:
-        return current_loras, gr.update(), [], selected_indices, [], [], 0, []
+        return current_loras, gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update()
+
+def remove_custom_lora(selected_indices, current_loras, gallery):
+    if current_loras:
+        custom_lora_repo = current_loras[-1]['repo']
+        # Remove from loras list
+        current_loras = current_loras[:-1]
+        # Remove from selected_indices if selected
+        custom_lora_index = len(current_loras)
+        if custom_lora_index in selected_indices:
+            selected_indices.remove(custom_lora_index)
+    # Update gallery
+    gallery_items = [(item["image"], item["title"]) for item in current_loras]
+    # Update selected_info and images
+    selected_info_1 = "Select a LoRA 1"
+    selected_info_2 = "Select a LoRA 2"
+    lora_scale_1 = 1.15
+    lora_scale_2 = 1.15
+    lora_image_1 = None
+    lora_image_2 = None
+    if len(selected_indices) >= 1:
+        lora1 = current_loras[selected_indices[0]]
+        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨"
+        lora_image_1 = lora1['image']
+    if len(selected_indices) >= 2:
+        lora2 = current_loras[selected_indices[1]]
+        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨"
+        lora_image_2 = lora2['image']
+    return (
+        current_loras,
+        gr.update(value=gallery_items),
+        selected_info_1,
+        selected_info_2,
+        selected_indices,
+        lora_scale_1,
+        lora_scale_2,
+        lora_image_1,
+        lora_image_2
+    )
 
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress):
     print("Generating image...")
@@ -283,7 +308,7 @@ def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps
     return final_image
 
 @spaces.GPU(duration=75)
-def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_indices, lora_scales, randomize_seed, seed, width, height, loras_state, progress=gr.Progress(track_tqdm=True)):
+def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state, progress=gr.Progress(track_tqdm=True)):
     if not selected_indices:
         raise gr.Error("You must select at least one LoRA before proceeding.")
 
@@ -315,7 +340,7 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
             lora_name = f"lora_{idx}"
             lora_names.append(lora_name)
             print(f"Lora Name: {lora_name}")
-            lora_weights.append(lora_scales[idx] if idx < len(lora_scales) else 1.15)
+            lora_weights.append(lora_scale_1 if idx == 0 else lora_scale_2)
             lora_path = lora['repo']
             weight_name = lora.get("weights")
             print(f"Lora Path: {lora_path}")
@@ -326,6 +351,8 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
                 low_cpu_mem_usage=True,
                 adapter_name=lora_name
             )
+#            if image_input is not None: pipe_i2i = pipe_to_use
+#            else: pipe = pipe_to_use
         print("Loaded LoRAs:", lora_names)
         print("Adapter weights:", lora_weights)
         if image_input is not None:
@@ -446,34 +473,34 @@ with gr.Blocks(css=css, delete_cache=(60, 60)) as app:
     )
     loras_state = gr.State(loras)
     selected_indices = gr.State([])
-    num_loras_visible = gr.State(2)
     with gr.Row():
         with gr.Column(scale=3):
             prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Type a prompt after selecting a LoRA")
         with gr.Column(scale=1):
             generate_button = gr.Button("Generate", variant="primary", elem_classes=["button_total"])
     with gr.Row(elem_id="loaded_loras"):
-        add_lora_button = gr.Button("Add LoRA", variant="primary")
-        randomize_button = gr.Button("🎲", variant="secondary", scale=1, elem_id="random_btn")
-        lora_slots = []
-        for i in range(5):  # Assuming max 5 LoRAs
-            with gr.Column(visible=(i < 2), elem_id=f"lora_slot_{i}"):
-                with gr.Row():
-                    with gr.Column(scale=0, min_width=50):
-                        lora_image = gr.Image(label=f"LoRA {i+1} Image", interactive=False, min_width=50, width=50, show_label=False, show_share_button=False, show_download_button=False, show_fullscreen_button=False, height=50)
-                    with gr.Column(scale=3, min_width=100):
-                        selected_info = gr.Markdown(f"Select a LoRA {i+1}")
-                    with gr.Column(scale=5, min_width=50):
-                        lora_scale = gr.Slider(label=f"LoRA {i+1} Scale", minimum=0, maximum=3, step=0.01, value=1.15)
-                with gr.Row():
-                    remove_button = gr.Button("Remove", size="sm")
-                lora_slots.append({
-                    'lora_image': lora_image,
-                    'selected_info': selected_info,
-                    'lora_scale': lora_scale,
-                    'remove_button': remove_button,
-                    'visible': (i < 2),
-                })
+        with gr.Column(scale=1, min_width=25):
+            randomize_button = gr.Button("🎲", variant="secondary", scale=1, elem_id="random_btn")
+        with gr.Column(scale=8):
+            with gr.Row():
+                with gr.Column(scale=0, min_width=50):
+                    lora_image_1 = gr.Image(label="LoRA 1 Image", interactive=False, min_width=50, width=50, show_label=False, show_share_button=False, show_download_button=False, show_fullscreen_button=False, height=50)
+                with gr.Column(scale=3, min_width=100):
+                    selected_info_1 = gr.Markdown("Select a LoRA 1")
+                with gr.Column(scale=5, min_width=50):
+                    lora_scale_1 = gr.Slider(label="LoRA 1 Scale", minimum=0, maximum=3, step=0.01, value=1.15)
+            with gr.Row():
+                remove_button_1 = gr.Button("Remove", size="sm")
+        with gr.Column(scale=8):
+            with gr.Row():
+                with gr.Column(scale=0, min_width=50):
+                    lora_image_2 = gr.Image(label="LoRA 2 Image", interactive=False, min_width=50, width=50, show_label=False, show_share_button=False, show_download_button=False, show_fullscreen_button=False, height=50)
+                with gr.Column(scale=3, min_width=100):
+                    selected_info_2 = gr.Markdown("Select a LoRA 2")
+                with gr.Column(scale=5, min_width=50):
+                    lora_scale_2 = gr.Slider(label="LoRA 2 Scale", minimum=0, maximum=3, step=0.01, value=1.15)
+            with gr.Row():
+                remove_button_2 = gr.Button("Remove", size="sm")
     with gr.Row():
         with gr.Column():
             with gr.Group():
@@ -518,33 +545,36 @@ with gr.Blocks(css=css, delete_cache=(60, 60)) as app:
     gallery.select(
         update_selection,
         inputs=[selected_indices, loras_state, width, height],
-        outputs=[prompt, [slot['selected_info'] for slot in lora_slots], selected_indices, [slot['lora_scale'] for slot in lora_slots], width, height, [slot['lora_image'] for slot in lora_slots], [slot['lora_image'] for slot in lora_slots]]
+        outputs=[prompt, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, width, height, lora_image_1, lora_image_2])
+    remove_button_1.click(
+        remove_lora_1,
+        inputs=[selected_indices, loras_state],
+        outputs=[selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
     )
-    add_lora_button.click(
-        add_lora_slot,
-        inputs=[num_loras_visible],
-        outputs=[num_loras_visible, [gr.update(visible=True) for slot in lora_slots]]
+    remove_button_2.click(
+        remove_lora_2,
+        inputs=[selected_indices, loras_state],
+        outputs=[selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
     )
-    for idx, slot in enumerate(lora_slots):
-        slot['remove_button'].click(
-            remove_lora_slot,
-            inputs=[idx, selected_indices],
-            outputs=[[s['selected_info'] for s in lora_slots], selected_indices, [s['lora_scale'] for s in lora_slots], [s['lora_image'] for s in lora_slots], num_loras_visible, [gr.update(visible=s['visible']) for s in lora_slots]]
-        )
     randomize_button.click(
         randomize_loras,
         inputs=[selected_indices, loras_state],
-        outputs=[[slot['selected_info'] for slot in lora_slots], selected_indices, [slot['lora_scale'] for slot in lora_slots], [slot['lora_image'] for slot in lora_slots], prompt, num_loras_visible, [gr.update(visible=slot['visible']) for slot in lora_slots]]
+        outputs=[selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2, prompt]
     )
     add_custom_lora_button.click(
         add_custom_lora,
         inputs=[custom_lora, selected_indices, loras_state, gallery],
-        outputs=[loras_state, gallery, [slot['selected_info'] for slot in lora_slots], selected_indices, [slot['lora_scale'] for slot in lora_slots], [slot['lora_image'] for slot in lora_slots], num_loras_visible, [gr.update(visible=slot['visible']) for slot in lora_slots]]
+        outputs=[loras_state, gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
+    )
+    remove_custom_lora_button.click(
+        remove_custom_lora,
+        inputs=[selected_indices, loras_state, gallery],
+        outputs=[loras_state, gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
     )
     gr.on(
         triggers=[generate_button.click, prompt.submit],
         fn=run_lora,
-        inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_indices, [slot['lora_scale'] for slot in lora_slots], randomize_seed, seed, width, height, loras_state],
+        inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state],
         outputs=[result, seed, progress_bar]
     ).then(
         fn=lambda x, history: update_history(x, history),
@@ -553,4 +583,4 @@ with gr.Blocks(css=css, delete_cache=(60, 60)) as app:
     )
 
 app.queue()
-app.launch()
+app.launch()