Update app.py

app.py

@@ -125,23 +125,22 @@ def update_selection(evt: gr.SelectData, selected_indices, loras_state, width, h
     if selected_index in selected_indices:
         selected_indices.remove(selected_index)
     else:
-        if len(selected_indices) < 3:  # 변경: 2에서 3으로
+        if len(selected_indices) < 3:
             selected_indices.append(selected_index)
         else:
             gr.Warning("You can select up to 3 LoRAs, remove one to select a new one.")
             return gr.update(), gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), width, height, gr.update(), gr.update(), gr.update()
 
-    selected_info_1 = gr.Markdown("LoRA 1 선택")
-    selected_info_2 = gr.Markdown("LoRA 2 선택")
-    selected_info_3 = gr.Markdown("LoRA 3 선택")
-
+    selected_info_1 = "Select LoRA 1"
+    selected_info_2 = "Select LoRA 2"
+    selected_info_3 = "Select LoRA 3"
     
     lora_scale_1 = 1.15
     lora_scale_2 = 1.15
     lora_scale_3 = 1.15    
     lora_image_1 = None
     lora_image_2 = None
-    lora_image_2 = None
+    lora_image_3 = None
     
     if len(selected_indices) >= 1:
         lora1 = loras_state[selected_indices[0]]
@@ -164,63 +163,65 @@ def update_selection(evt: gr.SelectData, selected_indices, loras_state, width, h
 
     return gr.update(placeholder=new_placeholder), selected_info_1, selected_info_2, selected_info_3, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, width, height, lora_image_1, lora_image_2, lora_image_3
 
-
-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)
+def remove_lora(selected_indices, loras_state, index_to_remove):
+    if len(selected_indices) > index_to_remove:
+        selected_indices.pop(index_to_remove)
+    
     selected_info_1 = "Select LoRA 1"
     selected_info_2 = "Select LoRA 2"
+    selected_info_3 = "Select LoRA 3"
     lora_scale_1 = 1.15
     lora_scale_2 = 1.15
+    lora_scale_3 = 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
+    lora_image_3 = None
+    
+    for i, idx in enumerate(selected_indices):
+        lora = loras_state[idx]
+        if i == 0:
+            selected_info_1 = f"### LoRA 1 Selected: [{lora['title']}]({lora['repo']}) ✨"
+            lora_image_1 = lora['image']
+        elif i == 1:
+            selected_info_2 = f"### LoRA 2 Selected: [{lora['title']}]({lora['repo']}) ✨"
+            lora_image_2 = lora['image']
+        elif i == 2:
+            selected_info_3 = f"### LoRA 3 Selected: [{lora['title']}]({lora['repo']}) ✨"
+            lora_image_3 = lora['image']
+    
+    return selected_info_1, selected_info_2, selected_info_3, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_image_1, lora_image_2, lora_image_3
+
+def remove_lora_1(selected_indices, loras_state):
+    return remove_lora(selected_indices, loras_state, 0)
+
+def remove_lora_2(selected_indices, loras_state):
+    return remove_lora(selected_indices, loras_state, 1)
+
+def remove_lora_3(selected_indices, loras_state):
+    return remove_lora(selected_indices, loras_state, 2)
 
 def randomize_loras(selected_indices, loras_state):
     try:
-        if len(loras_state) < 2:
+        if len(loras_state) < 3:
             raise gr.Error("Not enough LoRAs to randomize.")
-        selected_indices = random.sample(range(len(loras_state)), 2)
+        selected_indices = random.sample(range(len(loras_state)), 3)
         lora1 = loras_state[selected_indices[0]]
         lora2 = loras_state[selected_indices[1]]
+        lora3 = loras_state[selected_indices[2]]
         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']}) ✨"
+        selected_info_3 = f"### LoRA 3 Selected: [{lora3['title']}](https://huggingface.co/{lora3['repo']}) ✨"
         lora_scale_1 = 1.15
         lora_scale_2 = 1.15
+        lora_scale_3 = 1.15
         lora_image_1 = lora1['image']
         lora_image_2 = lora2['image']
+        lora_image_3 = lora3['image']
         random_prompt = random.choice(prompt_values)
-        return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2, random_prompt
+        return selected_info_1, selected_info_2, selected_info_3, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_image_1, lora_image_2, lora_image_3, random_prompt
     except Exception as e:
         print(f"Error in randomize_loras: {str(e)}")
-        return "Error", "Error", [], 1.15, 1.15, None, None, ""
+        return "Error", "Error", "Error", [], 1.15, 1.15, 1.15, None, None, None, ""
 
 def add_custom_lora(custom_lora, selected_indices, current_loras):
     if custom_lora:
@@ -245,18 +246,21 @@ def add_custom_lora(custom_lora, selected_indices, current_loras):
             # Update gallery
             gallery_items = [(item["image"], item["title"]) for item in current_loras]
             # Update selected_indices if there's room
-            if len(selected_indices) < 2:
+            if len(selected_indices) < 3:
                 selected_indices.append(existing_item_index)
             else:
-                gr.Warning("You can select up to 2 LoRAs, remove one to select a new one.")
+                gr.Warning("You can select up to 3 LoRAs, remove one to select a new one.")
 
             # Update selected_info and images
             selected_info_1 = "Select a LoRA 1"
             selected_info_2 = "Select a LoRA 2"
+            selected_info_3 = "Select a LoRA 3"
             lora_scale_1 = 1.15
             lora_scale_2 = 1.15
+            lora_scale_3 = 1.15
             lora_image_1 = None
             lora_image_2 = None
+            lora_image_3 = None
             if len(selected_indices) >= 1:
                 lora1 = current_loras[selected_indices[0]]
                 selected_info_1 = f"### LoRA 1 Selected: {lora1['title']} ✨"
@@ -265,24 +269,31 @@ def add_custom_lora(custom_lora, selected_indices, current_loras):
                 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
+            if len(selected_indices) >= 3:
+                lora3 = current_loras[selected_indices[2]]
+                selected_info_3 = f"### LoRA 3 Selected: {lora3['title']} ✨"
+                lora_image_3 = lora3['image'] if lora3['image'] else None
             print("Finished adding custom LoRA")
             return (
                 current_loras,
                 gr.update(value=gallery_items),
                 selected_info_1, 
                 selected_info_2,
+                selected_info_3,
                 selected_indices,
                 lora_scale_1,
                 lora_scale_2,
+                lora_scale_3,
                 lora_image_1,
-                lora_image_2
+                lora_image_2,
+                lora_image_3
             )
         except Exception as e:
             print(e)
             gr.Warning(str(e))
-            return current_loras, gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update()
+            return current_loras, gr.update(), gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
     else:
-        return current_loras, gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update()
+        return current_loras, gr.update(), gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
 
 def remove_custom_lora(selected_indices, current_loras):
     if current_loras:
@@ -298,10 +309,13 @@ def remove_custom_lora(selected_indices, current_loras):
     # Update selected_info and images
     selected_info_1 = "Select a LoRA 1"
     selected_info_2 = "Select a LoRA 2"
+    selected_info_3 = "Select a LoRA 3"
     lora_scale_1 = 1.15
     lora_scale_2 = 1.15
+    lora_scale_3 = 1.15
     lora_image_1 = None
     lora_image_2 = None
+    lora_image_3 = None
     if len(selected_indices) >= 1:
         lora1 = current_loras[selected_indices[0]]
         selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨"
@@ -310,16 +324,23 @@ def remove_custom_lora(selected_indices, current_loras):
         lora2 = current_loras[selected_indices[1]]
         selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨"
         lora_image_2 = lora2['image']
+    if len(selected_indices) >= 3:
+        lora3 = current_loras[selected_indices[2]]
+        selected_info_3 = f"### LoRA 3 Selected: [{lora3['title']}]({lora3['repo']}) ✨"
+        lora_image_3 = lora3['image']
     return (
         current_loras,
         gr.update(value=gallery_items),
         selected_info_1,
         selected_info_2,
+        selected_info_3,
         selected_indices,
         lora_scale_1,
         lora_scale_2,
+        lora_scale_3,
         lora_image_1,
-        lora_image_2
+        lora_image_2,
+        lora_image_3
     )
 
 @spaces.GPU(duration=75)
@@ -775,19 +796,21 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as a
         inputs=[selected_indices, loras_state],
         outputs=[selected_info_1, selected_info_2, selected_info_3, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_image_1, lora_image_2, lora_image_3]
     )
-
+    
     remove_button_2.click(
         remove_lora_2,
         inputs=[selected_indices, loras_state],
         outputs=[selected_info_1, selected_info_2, selected_info_3, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_image_1, lora_image_2, lora_image_3]
     )
-
+    
     remove_button_3.click(
         remove_lora_3,
         inputs=[selected_indices, loras_state],
         outputs=[selected_info_1, selected_info_2, selected_info_3, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_image_1, lora_image_2, lora_image_3]
     )
 
+ 
+
     randomize_button.click(
         randomize_loras,
         inputs=[selected_indices, loras_state],