Upload app.py

app.py

@@ -349,6 +349,85 @@ def _hub_url_parts(url: str) -> tuple[str, str]:
     return "/".join(parts[:2]), "/".join(parts[4:])
 
 
+CIVITAI_HOSTS = ("civitai.com", "civitai.red", "civitai.green", "civitai.work")
+
+
+def _download_direct_lora(url: str) -> str:
+    """Fetch a `.safetensors` from a plain URL - CivitAI in particular - and return the local path.
+
+    The download happens on the Space's own machine, not in the visitor's browser, so a CivitAI
+    session in a browser tab has nothing to do with it: a gated model answers a server with an HTML
+    login page instead of weights. `CIVITAI_TOKEN` (Settings -> Variables and secrets) is appended
+    automatically when it is set, and the header of whatever comes back is checked so a login page
+    fails with a sentence that says what to do rather than a parse error deep inside safetensors.
+    """
+    import hashlib
+    from urllib.parse import urlparse, unquote
+
+    import requests
+
+    host = (urlparse(url).hostname or "").lower()
+    request_url = url
+    token = os.environ.get("CIVITAI_TOKEN", "").strip()
+    if token and any(host.endswith(known) for known in CIVITAI_HOSTS) and "token=" not in url:
+        request_url = url + ("&" if "?" in url else "?") + f"token={token}"
+
+    cache_dir = os.path.join(tempfile.gettempdir(), "url-loras")
+    os.makedirs(cache_dir, exist_ok=True)
+    cached = os.path.join(cache_dir, hashlib.sha256(url.encode()).hexdigest()[:16] + ".safetensors")
+    if os.path.exists(cached) and os.path.getsize(cached) > 1_000_000:
+        return cached
+
+    try:
+        response = requests.get(request_url, stream=True, timeout=120,
+                                headers={"User-Agent": "Mozilla/5.0"})
+        response.raise_for_status()
+    except Exception as error:
+        raise gr.Error(f"Could not download `{url}`: {error}")
+
+    content_type = (response.headers.get("content-type") or "").lower()
+    if "text/html" in content_type:
+        raise gr.Error(
+            "That link answered with a web page instead of a file. The model is gated, so the "
+            "Space needs its own key: add CIVITAI_TOKEN under Settings -> Variables and secrets."
+        )
+
+    disposition = response.headers.get("content-disposition", "")
+    name = unquote(re.findall(r'filename\*?=(?:UTF-8\'\'|")?([^";]+)', disposition)[0]) \
+        if "filename" in disposition else os.path.basename(urlparse(url).path)
+    if name and not name.lower().endswith(".safetensors") and "." in name:
+        print(f"[lora] {name} is not a .safetensors; trying it anyway")
+
+    written = 0
+    with open(cached, "wb") as handle:
+        for chunk in response.iter_content(chunk_size=1 << 20):
+            if chunk:
+                handle.write(chunk)
+                written += len(chunk)
+
+    if written < 1_000_000:
+        os.remove(cached)
+        raise gr.Error(
+            "That link returned only a few kilobytes - almost always a login or error page rather "
+            "than weights. Check the link, or add CIVITAI_TOKEN to the Space."
+        )
+
+    # safetensors starts with an 8-byte little-endian header length followed by that much JSON.
+    with open(cached, "rb") as handle:
+        header_len = int.from_bytes(handle.read(8), "little")
+        if not (0 < header_len < 100_000_000):
+            os.remove(cached)
+            raise gr.Error("The downloaded file is not a `.safetensors` (bad header).")
+        try:
+            json.loads(handle.read(header_len).decode("utf-8"))
+        except Exception:
+            os.remove(cached)
+            raise gr.Error("The downloaded file is not a `.safetensors` (unreadable header).")
+
+    print(f"[lora] downloaded {written / 1e6:.0f} MB from {host} -> {os.path.basename(cached)}")
+    return cached
+
+
 def resolve_lora(reference: str) -> str:
     """Turn what the user typed into a local `.safetensors` path.
 
@@ -363,8 +442,13 @@ def resolve_lora(reference: str) -> str:
     if os.path.exists(reference):
         return reference
     if reference.startswith(("http://", "https://")):
-        repo_id, filename = _hub_url_parts(reference)
-        return hf_hub_download(repo_id, filename)
+        from urllib.parse import urlparse
+
+        if (urlparse(reference).hostname or "").lower().endswith("huggingface.co"):
+            repo_id, filename = _hub_url_parts(reference)
+            return hf_hub_download(repo_id, filename)
+        # Anything else - CivitAI and any other direct link - is fetched as a plain file.
+        return _download_direct_lora(reference)
 
     parts = [part for part in reference.split("/") if part]
     if len(parts) > 2 and parts[-1].endswith(".safetensors"):
@@ -1213,15 +1297,20 @@ HERO = """
   <a href="https://huggingface.co/spaces/multimodalart/minimax-h3" target="_blank" rel="noopener">text / image to video</a></p>
   <div class="pills">
     <span>33B</span><span>joint video + audio</span><span>Turbo lora: 4&ndash;8 steps</span>
-    <span>ComfyUI lora accepted</span><span>up to 9 references</span><span>named profiles</span><span>GPU cost estimate</span><span>clip stitching with audio</span>
+    <span>ComfyUI lora accepted</span><span>up to 9 references</span><span>named profiles</span><span>CivitAI links</span><span>GPU cost estimate</span><span>clip stitching with audio</span>
   </div>
 </div>
 """
 
 LORA_HELP = """Each slot takes a Hugging Face repo (`owner/repo`), a file inside one
-(`owner/repo/name.safetensors`), a file URL, or a local path — or just drop the files below. A strength of `0`
-switches a slot off without clearing it. Adapters have to be trained against the `transformer_ref/` partition.
-ComfyUI-trained adapters are remapped to diffusers' module names on the fly, so no separate conversion step is needed.
+(`owner/repo/name.safetensors`), a **CivitAI download link**, any other direct `.safetensors` URL, or a local path —
+or just drop the files below. A strength of `0` switches a slot off without clearing it. Adapters have to be trained
+against the `transformer_ref/` partition. ComfyUI-trained adapters are remapped to diffusers' module names on the fly,
+so no separate conversion step is needed.
+
+CivitAI links are fetched by the Space itself, not by your browser, so being signed in there does not help a gated
+model — set `CIVITAI_TOKEN` under *Settings → Variables and secrets* and it is appended automatically. Downloads are
+cached, so a second run with the same link costs nothing.
 """
 
 TURBO_HELP = """A few-step distillation that renders joint video + soundtrack in **4–8 steps** instead of the usual