Determine model automatically from r8 lora filename

Author: fofr

.github/workflows/ci.yaml

@@ -4,7 +4,7 @@ on:
   workflow_dispatch:
     inputs:
       models:
-        description: 'Comma-separated list of models (standalone, trained1.3b, trained14b) to push or "all"'
+        description: 'Comma-separated list of models (standalone, trained) to push or "all"'
         type: string
         default: 'all'
 
@@ -44,7 +44,7 @@ jobs:
         REPLICATE_API_TOKEN: ${{ secrets.REPLICATE_API_TOKEN }}
       run: |
         if [ "${{ inputs.models }}" = "all" ]; then
-          models="standalone, trained1.3b, trained14b"
+          models="standalone, trained"
         else
           models="${{ inputs.models }}"
         fi

cog-safe-push-configs/trained14b.yaml cog-safe-push-configs/trained.yaml

@@ -217,7 +217,7 @@ def generate(
         negative_prompt: str | None = None,
         aspect_ratio: str = "16:9",
         frames: int = 81,
-        model: str = "14b",
+        model: str | None = None,
         lora_url: str | None = None,
         lora_strength_model: float = 1.0,
         lora_strength_clip: float = 1.0,
@@ -232,15 +232,16 @@ def generate(
         self.comfyUI.cleanup(ALL_DIRECTORIES)
         seed = seed_helper.generate(seed)
 
-        if resolution == "720p" and model == "1.3b":
-            print("Warning: 720p is not supported for 1.3b, using 480p instead")
-            resolution = "480p"
-
         lora_filename = None
         if replicate_weights:
-            lora_filename = download_replicate_weights(
+            lora_filename, model_type = download_replicate_weights(
                 replicate_weights, "ComfyUI/models/loras"
             )
+            model = model_type
+
+        if resolution == "720p" and model == "1.3b":
+            print("Warning: 720p is not supported for 1.3b, using 480p instead")
+            resolution = "480p"
 
         with open(api_json_file, "r") as file:
             workflow = json.loads(file.read())
@@ -308,9 +309,7 @@ def predict(
         )
 
 
-class Trained14BLoraPredictor(Predictor):
-    model = "14b"
-
+class TrainedLoraPredictor(Predictor):
     def predict(
         self,
         prompt: str = Inputs.prompt,
@@ -332,7 +331,7 @@ def predict(
             negative_prompt=negative_prompt,
             aspect_ratio=aspect_ratio,
             frames=frames,
-            model=self.model,
+            model=None,
             resolution=resolution,
             lora_url=None,
             lora_strength_model=lora_strength_model,
@@ -344,7 +343,3 @@ def predict(
             seed=seed,
             replicate_weights=replicate_weights,
         )
-
-
-class Trained1_3BLoraPredictor(Trained14BLoraPredictor):
-    model = "1.3b"

cog-safe-push-configs/trained1.3b.yaml

@@ -5,19 +5,30 @@
 import hashlib
 from pathlib import Path
 
-def get_filename_from_url(url: str, extension: str) -> str:
-    """Generate a unique filename from a URL using MD5 hash"""
-    filename = hashlib.md5(url.encode()).hexdigest()
-    return f"{filename}.{extension}"
+
+def url_hash(url: str) -> str:
+    return hashlib.md5(url.encode()).hexdigest()
+
+
+def get_filename_from_url(url: str, model_type: str) -> str:
+    return f"{model_type}_{url_hash(url)}.safetensors"
+
 
 def download_replicate_weights(url: str, lora_dir: str):
     """Downloads weights from a Replicate tarball URL and extracts the safetensors file"""
-    unique_filename = get_filename_from_url(url, "safetensors")
-    target_path = Path(lora_dir) / unique_filename
+    hash = url_hash(url)
+
+    # Check if either version already exists
+    lora_dir_path = Path(lora_dir)
+    existing_14b = lora_dir_path / f"14b_{hash}.safetensors"
+    existing_1_3b = lora_dir_path / f"1.3b_{hash}.safetensors"
 
-    if target_path.exists():
-        print(f"✅ {unique_filename} already cached")
-        return unique_filename
+    if existing_14b.exists():
+        print(f"✅ {existing_14b.name} already cached")
+        return existing_14b.name, "14b"
+    if existing_1_3b.exists():
+        print(f"✅ {existing_1_3b.name} already cached")
+        return existing_1_3b.name, "1.3b"
 
     with tempfile.TemporaryDirectory() as temp_dir:
         temp_dir = Path(temp_dir)
@@ -41,6 +52,9 @@ def download_replicate_weights(url: str, lora_dir: str):
         if len(safetensors_paths) > 1:
             raise ValueError("Multiple .safetensors files found in tarball")
 
+        model_type = "1.3b" if "1.3b" in safetensors_paths[0].name.lower() else "14b"
+        unique_filename = get_filename_from_url(url, model_type)
+        target_path = Path(lora_dir) / unique_filename
         shutil.move(safetensors_paths[0], target_path)
 
-    return unique_filename
+        return unique_filename, model_type

predict.py

@@ -34,10 +34,8 @@ MODEL=$1
 # Set the predictor name
 if [[ "$MODEL" == "standalone" ]]; then
     export PREDICTOR=StandaloneLoraPredictor
-elif [[ "$MODEL" == "trained14b" ]]; then
-    export PREDICTOR=Trained14BLoraPredictor
-elif [[ "$MODEL" == "trained1.3b" ]]; then
-    export PREDICTOR=Trained1_3BLoraPredictor
+elif [[ "$MODEL" == "trained" ]]; then
+    export PREDICTOR=TrainedLoraPredictor
 else
     echo "Unknown model: $MODEL. Valid models are:"
     list_models