Update app.py

app.py

@@ -385,35 +385,82 @@ run_lora.zerogpu = True
 
 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(f"Base model: {base_model}")
-        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
-        fs = HfFileSystem()
+    if len(split_link) != 2:
+        raise Exception("Invalid Hugging Face repository link format.")
+
+    print(f"Repository attempted: {split_link}")
+    
+    # Load model card
+    model_card = ModelCard.load(link)
+    base_model = model_card.data.get("base_model")
+    print(f"Base model: {base_model}")
+
+    # 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!")
+        
+    # Extract image and trigger word
+    print("Before trying to get image")
+    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}")
+
+        
+    # Initialize Hugging Face file system
+    fs = HfFileSystem()
+    try:
+        list_of_files = fs.ls(link, detail=False)
+        
+        # Initialize variables for safetensors selection
         safetensors_name = None
-        try:
-            list_of_files = fs.ls(link, detail=False)
-            for file in list_of_files:
-                if file.endswith(".safetensors"):
-                    safetensors_name = file.split("/")[-1]
-                if not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp")):
-                    image_elements = file.split("/")
-                    image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
-        except Exception as e:
-            print(e)
-            raise gr.Error("Invalid Hugging Face repository with a *.safetensors LoRA")
+        highest_trained_file = None
+        highest_steps = -1
+        last_safetensors_file = None
+        step_pattern = re.compile(r"_0{3,}\d+")  # Detects step count `_000...`
+
+        for file in list_of_files:
+            filename = file.split("/")[-1]
+
+            # Select safetensors file
+            if filename.endswith(".safetensors"):
+                last_safetensors_file = filename  # Track last encountered file
+                
+                match = step_pattern.search(filename)
+                if not match:
+                    # Found a full model without step numbers, return immediately
+                    safetensors_name = filename
+                    break
+                else:
+                    # Extract step count and track highest
+                    steps = int(match.group().lstrip("_"))
+                    if steps > highest_steps:
+                        highest_trained_file = filename
+                        highest_steps = steps
+
+            # Select an image file if not found in model card
+            if not image_url and filename.lower().endswith((".jpg", ".jpeg", ".png", ".webp")):
+                image_url = f"https://huggingface.co/{link}/resolve/main/{filename}"
+
+        # If no full model found, fall back to the most trained safetensors file
         if not safetensors_name:
-            raise gr.Error("No *.safetensors file found in the repository")
-        return split_link[1], link, safetensors_name, trigger_word, image_url
-    else:
-        raise gr.Error("Invalid Hugging Face repository link")
+            safetensors_name = highest_trained_file if highest_trained_file else last_safetensors_file
+
+        # If still no safetensors file found, raise an exception
+        if not safetensors_name:
+            raise Exception("No valid *.safetensors file found in the repository.")
+
+    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
 
 def check_custom_model(link):
     if link.endswith(".safetensors"):