Update app.py

app.py

@@ -289,6 +289,31 @@ def get_huggingface_safetensors(link):
 
 def check_custom_model(link):
     print(f"Checking a custom model on: {link}")
+    
+    if link.endswith('.safetensors'):
+        if 'huggingface.co' in link:
+            parts = link.split('/')
+            try:
+                hf_index = parts.index('huggingface.co')
+                username = parts[hf_index + 1]
+                repo_name = parts[hf_index + 2]
+                repo = f"{username}/{repo_name}"
+                
+                safetensors_name = parts[-1]
+                
+                try:
+                    model_card = ModelCard.load(repo)
+                    trigger_word = model_card.data.get("instance_prompt", "")
+                    image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
+                    image_url = f"https://huggingface.co/{repo}/resolve/main/{image_path}" if image_path else None
+                except:
+                    trigger_word = ""
+                    image_url = None
+                
+                return repo_name, repo, safetensors_name, trigger_word, image_url
+            except:
+                raise Exception("Invalid safetensors URL format")
+    
     if link.startswith("https://"):
         if link.startswith("https://huggingface.co") or link.startswith("https://www.huggingface.co"):
             link_split = link.split("huggingface.co/")