Update app.py

app.py

@@ -19,21 +19,25 @@ MAX_SEED = np.iinfo(np.int32).max
 pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
 
 # Load LoRA data (you'll need to create this JSON file or modify to load your LoRAs)
-
-with open("flux_loras.json", "r") as file:
-    data = json.load(file)
-    flux_loras_raw = [
-        {
-            "image": item["image"],
-            "title": item["title"],
-            "repo": item["repo"],
-            "trigger_word": item.get("trigger_word", ""),
-            "trigger_position": item.get("trigger_position", "prepend"),
-            "weights": item.get("weights", "pytorch_lora_weights.safetensors"),
-        }
-        for item in data
-    ]
-print(f"Loaded {len(flux_loras_raw)} LoRAs from JSON")
+try:
+    with open("flux_loras.json", "r") as file:
+        data = json.load(file)
+        flux_loras_raw = [
+            {
+                "image": item["image"],
+                "title": item["title"],
+                "repo": item["repo"],
+                "trigger_word": item.get("trigger_word", ""),
+                "trigger_position": item.get("trigger_position", "prepend"),
+                "weights": item.get("weights", "pytorch_lora_weights.safetensors"),
+                "likes": item.get("likes", 0),
+            }
+            for item in data
+        ]
+    print(f"Successfully loaded {len(flux_loras_raw)} LoRAs from JSON")
+except Exception as e:
+    print(f"Error loading flux_loras.json: {e}")
+    flux_loras_raw = []
 # Global variables for LoRA management
 current_lora = None
 lora_cache = {}
@@ -123,8 +127,16 @@ def remove_custom_lora():
 
 def classify_gallery(flux_loras):
     """Sort gallery by likes"""
-    sorted_gallery = sorted(flux_loras, key=lambda x: x.get("likes", 0), reverse=True)
-    return [(item["image"], item["title"]) for item in sorted_gallery], sorted_gallery
+    try:
+        sorted_gallery = sorted(flux_loras, key=lambda x: x.get("likes", 0), reverse=True)
+        gallery_items = []
+        for item in sorted_gallery:
+            if "image" in item and "title" in item:
+                gallery_items.append((item["image"], item["title"]))
+        return gallery_items, sorted_gallery
+    except Exception as e:
+        print(f"Error loading gallery: {e}")
+        return [], []
 
 def infer_with_lora_wrapper(input_image, prompt, selected_index, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.75, flux_loras=None, progress=gr.Progress(track_tqdm=True)):
     """Wrapper function to handle state serialization"""
@@ -149,29 +161,44 @@ def infer_with_lora(input_image, prompt, selected_index, custom_lora, seed=42, r
         lora_to_use = custom_lora
     elif selected_index is not None and flux_loras and selected_index < len(flux_loras):
         lora_to_use = flux_loras[selected_index]
-    print(f"Loaded {len(flux_loras)} LoRAs from JSON")
     # Load LoRA if needed
     if lora_to_use and lora_to_use != current_lora:
         try:
             # Unload current LoRA
             if current_lora:
                 pipe.unload_lora_weights()
+                print(f"Unloaded previous LoRA")
             
             # Load new LoRA
-            lora_path = load_lora_weights(lora_to_use["repo"], lora_to_use["weights"])
+            repo_id = lora_to_use.get("repo", "unknown")
+            weights_file = lora_to_use.get("weights", "pytorch_lora_weights.safetensors")
+            print(f"Loading LoRA: {repo_id} with weights: {weights_file}")
+            
+            lora_path = load_lora_weights(repo_id, weights_file)
             if lora_path:
                 pipe.load_lora_weights(lora_path, adapter_name="selected_lora")
                 pipe.set_adapters(["selected_lora"], adapter_weights=[lora_scale])
-                print(f"loaded: {lora_path} with scale {lora_scale}")
+                print(f"Successfully loaded: {lora_path} with scale {lora_scale}")
                 current_lora = lora_to_use
+            else:
+                print(f"Failed to load LoRA from {repo_id}")
+                gr.Warning(f"Failed to load LoRA style. Please try a different one.")
+                return None, seed, gr.update(visible=False)
             
         except Exception as e:
             print(f"Error loading LoRA: {e}")
             # Continue without LoRA
     else:
-        print(f"using already loaded lora: {lora_to_use}")
+        if lora_to_use:
+            print(f"Using already loaded LoRA: {lora_to_use.get('repo', 'unknown')}")
     
-    input_image = input_image.convert("RGB")
+    try:
+        # Convert image to RGB
+        input_image = input_image.convert("RGB")
+    except Exception as e:
+        print(f"Error processing image: {e}")
+        gr.Warning("Error processing the uploaded image. Please try a different image.")
+        return None, seed, gr.update(visible=False)
     
     # Check if LoRA is selected
     if lora_to_use is None: