Update app.py

app.py

@@ -179,67 +179,66 @@ def randomize_loras(selected_indices, loras_state):
     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):
+    """Add a custom LoRA to the current loras and update the gallery."""
     if custom_lora:
         try:
             title, repo, path, trigger_word, image = check_custom_model(custom_lora)
             print(f"Loaded custom LoRA: {repo}")
-            existing_item_index = next((index for (index, item) in enumerate(current_loras) if item['repo'] == repo), None)
+            existing_item_index = next(
+                (index for index, item in enumerate(current_loras) if item['repo'] == repo), 
+                None
+            )
             if existing_item_index is None:
+                # Download and save LoRA if necessary
                 if repo.endswith(".safetensors") and repo.startswith("http"):
                     repo = download_file(repo)
+                # Add the new LoRA to the current list
                 new_item = {
-                    "image": image if image else "/valid/path/to/default/image.png",  # Ensure this path is valid
+                    "image": image or "/path/to/valid/default_image.png",  # Update this path
                     "title": title,
                     "repo": repo,
                     "weights": path,
-                    "trigger_word": trigger_word
+                    "trigger_word": trigger_word,
                 }
-                print(f"New LoRA: {new_item}")
-                existing_item_index = len(current_loras)
                 current_loras.append(new_item)
-            
-            # Update gallery
+                existing_item_index = len(current_loras) - 1
+
+            # Update gallery and indices
             gallery_items = [(item["image"], item["title"]) for item in current_loras]
-            # 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_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_info_1, 
-                selected_info_2,
-                selected_indices,
-                lora_scale_1,
-                lora_scale_2,
-                lora_image_1,
-                lora_image_2
+                *update_selected_info(selected_indices, current_loras),
             )
         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()
+            print(f"Error: {e}")
+            raise gr.Error(f"Failed to load custom LoRA: {str(e)}")
     else:
         return current_loras, gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update()
 
+def update_selected_info(selected_indices, loras_state):
+    """Update selection info and images for the UI."""
+    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 = "/path/to/valid/default_image.png"
+    lora_image_2 = "/path/to/valid/default_image.png"
+
+    if len(selected_indices) >= 1:
+        lora1 = loras_state[selected_indices[0]]
+        selected_info_1 = f"### LoRA 1 Selected: {lora1['title']} ✨"
+        lora_image_1 = lora1.get("image", lora_image_1)
+    if len(selected_indices) >= 2:
+        lora2 = loras_state[selected_indices[1]]
+        selected_info_2 = f"### LoRA 2 Selected: {lora2['title']} ✨"
+        lora_image_2 = lora2.get("image", lora_image_2)
+
+    return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2
+
 def remove_custom_lora(selected_indices, current_loras, gallery):
     if current_loras:
         custom_lora_repo = current_loras[-1]['repo']