Make sure to always load the highest trained safetensors file for all cases (#36)

- Make sure to always load the highest trained safetensors file for all cases (c576d4ca850474f1cbb84f4ca2ff8ff449ea1f68)


Co-authored-by: Sylvain Filoni <[email protected]>

app.py

@@ -12,6 +12,7 @@ from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_d
 import copy
 import random
 import time
+import re
 
 # Load LoRAs from JSON file
 with open('loras.json', 'r') as f:
@@ -172,30 +173,73 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
         yield final_image, seed, gr.update(value=progress_bar, visible=False)
         
 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
-            fs = HfFileSystem()
-            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)
-              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
+    split_link = link.split("/")
+    if len(split_link) != 2:
+        raise Exception("Invalid Hugging Face repository link format.")
+
+    # Load model card
+    model_card = ModelCard.load(link)
+    base_model = model_card.data.get("base_model")
+    print(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
+
+    # Initialize Hugging Face file system
+    fs = HfFileSystem()
+    try:
+        list_of_files = fs.ls(link, detail=False)
+        
+        # Initialize variables for safetensors selection
+        safetensors_name = None
+        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:
+            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.startswith("https://")):