Update app.py

app.py

@@ -106,6 +106,10 @@ def get_huggingface_safetensors(link):
   split_link = link.split("/")
   if(len(split_link) == 2):
             model_card = ModelCard.load(link)
+            base_model = model_card.data.get("base_model")
+            print(base_model)
+            if((base_model != "black-forest-labs/FLUX.1-dev") and (base_model != "black-forest-labs/FLUX.1-schnell")):
+                raise Exception("Not a FLUX LoRA!")
             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
@@ -120,8 +124,8 @@ def get_huggingface_safetensors(link):
                       image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
             except Exception as e:
               print(e)
-              gr.Warning(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA {e}")
-              raise Exception(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA {e}")
+              gr.Warning(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+              raise Exception(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
             return split_link[1], link, safetensors_name, trigger_word, image_url
 
 def check_custom_model(link):