Update app.py

app.py

@@ -182,17 +182,23 @@ def get_huggingface_safetensors(link):
     # Load model card
     model_card = ModelCard.load(link)
     base_model = model_card.data.get("base_model")
-    print(base_model)
+    print(f"Base model: {base_model}")
 
     # Validate model type
     if base_model not in {"black-forest-labs/FLUX.1-dev", "black-forest-labs/FLUX.1-schnell"}:
         raise Exception("Not a FLUX LoRA!")
 
     # Extract image and trigger word
-    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
-
+    try:
+        image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
+        print(f"Image path {image_path}")
+        trigger_word = model_card.data.get("instance_prompt", "")
+        print(f"Image path {trigger_word}")
+        image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
+        print(f"Image URL {image_url}")
+    except Exception as e:
+        raise gr.Error(e)
+        
     # Initialize Hugging Face file system
     fs = HfFileSystem()
     try:
@@ -239,7 +245,7 @@ def get_huggingface_safetensors(link):
     except Exception as e:
         print(e)
         raise Exception("You didn't include a valid Hugging Face repository with a *.safetensors LoRA")
-
+    
     return split_link[1], link, safetensors_name, trigger_word, image_url