Update app.py

app.py

@@ -389,8 +389,14 @@ def get_huggingface_safetensors(link):
         model_card = ModelCard.load(link)
         base_model = model_card.data.get("base_model")
         print(f"Base model: {base_model}")
-        if base_model not in ["black-forest-labs/FLUX.1-dev", "black-forest-labs/FLUX.1-schnell"]:
+       
+        # Validate model type
+        acceptable_models = {"black-forest-labs/FLUX.1-dev", "black-forest-labs/FLUX.1-schnell"}
+        models_to_check = base_model if isinstance(base_model, list) else [base_model]
+        if not any(model in acceptable_models for model in models_to_check):
             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