minimax worker: lora_from_url.py

worker/lora_from_url.py

@@ -0,0 +1,65 @@
+"""ComfyUI custom node: load a LoRA fetched from a URL at request time.
+
+Lets serverless workflows reference LoRAs by URL (HF, civitai, anywhere)
+instead of requiring the file to be pre-baked into the image. Downloads are
+cached on the worker's disk keyed by URL hash. Supports HF_TOKEN and
+CIVITAI_TOKEN env vars for authenticated downloads.
+"""
+
+import hashlib
+import os
+import urllib.request
+
+import comfy.sd
+import comfy.utils
+
+CACHE_DIR = "/lora-cache"
+os.makedirs(CACHE_DIR, exist_ok=True)
+
+
+def _fetch(url: str) -> str:
+    path = os.path.join(CACHE_DIR, hashlib.sha1(url.encode()).hexdigest() + ".safetensors")
+    if os.path.exists(path):
+        return path
+    headers = {}
+    hf_token = os.environ.get("HF_TOKEN")
+    if hf_token and "huggingface.co" in url:
+        headers["Authorization"] = f"Bearer {hf_token}"
+    civitai_token = os.environ.get("CIVITAI_TOKEN")
+    if civitai_token and "civitai" in url:
+        url += ("&" if "?" in url else "?") + "token=" + civitai_token
+    req = urllib.request.Request(url, headers=headers)
+    tmp = path + ".part"
+    with urllib.request.urlopen(req, timeout=600) as r, open(tmp, "wb") as f:
+        while chunk := r.read(1 << 20):
+            f.write(chunk)
+    os.rename(tmp, path)
+    return path
+
+
+class LoraFromURLModelOnly:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "url": ("STRING", {"default": ""}),
+                "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
+            }
+        }
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "load"
+    CATEGORY = "loaders"
+
+    def load(self, model, url, strength):
+        url = (url or "").strip()
+        if not url or strength == 0:
+            return (model,)
+        lora_sd = comfy.utils.load_torch_file(_fetch(url), safe_load=True)
+        model_lora, _ = comfy.sd.load_lora_for_models(model, None, lora_sd, strength, 0)
+        return (model_lora,)
+
+
+NODE_CLASS_MAPPINGS = {"LoraFromURLModelOnly": LoraFromURLModelOnly}
+NODE_DISPLAY_NAME_MAPPINGS = {"LoraFromURLModelOnly": "Load LoRA from URL (model only)"}