Handle downloading of R8 weights before inferring model

Author: fofr

cog-safe-push-configs/standalone.yaml

@@ -5,7 +5,7 @@ test_model: fofr/staging
 # Tests
 predict:
   compare_outputs: false
-  predict_timeout: 600
+  predict_timeout: 1800
   test_cases:
     - inputs:
         prompt: "a MNCRFTMOV cat"

cog-safe-push-configs/trained.yaml

@@ -8,5 +8,10 @@ predict:
   predict_timeout: 600
   test_cases:
     - inputs:
-        prompt: "a LAEZEL portrait painting"
+        prompt: "a 14b LAEZEL is laughing"
+        # url contains lora.safetensors, which defaults to 14b
         replicate_weights: "https://replicate.delivery/xezq/begAR055rwyeb0sjY60xM01L9i2fJ8T11ofpHGfL6PCEYf2FF/trained_model.tar"
+    - inputs:
+        prompt: "a 1.3b HOMER is laughing"
+        # url contains 1.3b-lora.safetensors, which defaults to 1.3b
+        replicate_weights: "https://replicate.delivery/xezq/7kXMcHfE8GwHZ6CU4TR2iTod20dPX8SerMCMrRn16efou1hRB/trained_model.tar"

predict.py

@@ -13,6 +13,7 @@
 OUTPUT_DIR = "/tmp/outputs"
 INPUT_DIR = "/tmp/inputs"
 COMFYUI_TEMP_OUTPUT_DIR = "ComfyUI/temp"
+COMFYUI_LORAS_DIR = "ComfyUI/models/loras"
 ALL_DIRECTORIES = [OUTPUT_DIR, INPUT_DIR, COMFYUI_TEMP_OUTPUT_DIR]
 
 mimetypes.add_type("image/webp", ".webp")
@@ -94,8 +95,6 @@ def setup(self):
         self.comfyUI.handle_weights(
             workflow,
             weights_to_download=[
-                "wan2.1_t2v_1.3B_bf16.safetensors",
-                "wan2.1_t2v_14B_bf16.safetensors",
                 "wan_2.1_vae.safetensors",
                 "umt5_xxl_fp16.safetensors",
             ],
@@ -186,23 +185,7 @@ def update_workflow(self, workflow, **kwargs):
             if kwargs["lora_filename"]:
                 lora_loader["lora_name"] = kwargs["lora_filename"]
             elif kwargs["lora_url"]:
-                url = kwargs["lora_url"]
-                if m := re.match(
-                    r"^(?:https?://replicate.com/)?([^/]+)/([^/]+)/?$", url
-                ):
-                    owner, model_name = m.groups()
-                    lora_filename = download_replicate_weights(
-                        f"https://replicate.com/{owner}/{model_name}/_weights",
-                        "ComfyUI/models/loras",
-                    )
-                    lora_loader["lora_name"] = lora_filename
-                elif url.startswith("https://replicate.delivery"):
-                    lora_filename = download_replicate_weights(
-                        url, "ComfyUI/models/loras"
-                    )
-                    lora_loader["lora_name"] = lora_filename
-                else:
-                    lora_loader["lora_name"] = url
+                lora_loader["lora_name"] = kwargs["lora_url"]
 
             lora_loader["strength_model"] = kwargs["lora_strength_model"]
             lora_loader["strength_clip"] = kwargs["lora_strength_clip"]
@@ -233,11 +216,31 @@ def generate(
         seed = seed_helper.generate(seed)
 
         lora_filename = None
+        inferred_model_type = None
         if replicate_weights:
-            lora_filename, model_type = download_replicate_weights(
-                replicate_weights, "ComfyUI/models/loras"
+            lora_filename, inferred_model_type = download_replicate_weights(
+                replicate_weights, COMFYUI_LORAS_DIR
             )
-            model = model_type
+            model = inferred_model_type
+        elif lora_url:
+            if m := re.match(
+                r"^(?:https?://replicate.com/)?([^/]+)/([^/]+)/?$", lora_url
+            ):
+                owner, model_name = m.groups()
+                lora_filename, inferred_model_type = download_replicate_weights(
+                    f"https://replicate.com/{owner}/{model_name}/_weights",
+                    COMFYUI_LORAS_DIR,
+                )
+            elif lora_url.startswith("https://replicate.delivery"):
+                lora_filename, inferred_model_type = download_replicate_weights(
+                    lora_url, COMFYUI_LORAS_DIR
+                )
+
+            if inferred_model_type and inferred_model_type != model:
+                print(
+                    f"Warning: Model type mismatch between requested model ({model}) and inferred model type ({inferred_model_type}). Using {inferred_model_type}."
+                )
+                model = inferred_model_type
 
         if resolution == "720p" and model == "1.3b":
             print("Warning: 720p is not supported for 1.3b, using 480p instead")