nodes/ComfyAssets-ComfyUI-KikoTools
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #114 from ComfyAssets/feat/model-downloader-tokens

Browse Source 
feat(model_downloader): server-side tokens and civitai.red/green links
This commit is contained in:
Vito
2026-11-07 04:23:33 -08:00
committed by GitHub
parent 5463011266 49fdc9642e
commit e7dd9f0a5a
11 changed files with 291 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+13
View File
@@ -10,6 +10,19 @@
pass it to a node that reads a whole folder. The folder must be inside a
saved path or ComfyUI's input/output, as for images.
### Model Downloader
- Civitai tokens and Hugging Face tokens are read on the ComfyUI server:
`CIVITAI_API_KEY` (also `CIVITAI_API_TOKEN`, `CIVITAI_TOKEN`) and `HF_TOKEN`
(also `HUGGING_FACE_HUB_TOKEN`, or the token `hf auth login` saved). Leave
`api_token` empty so shared workflows carry no token; a typed token still
wins. Each token goes only to its own site, and custom URLs only get a typed
token.
- civitai.green and civitai.red links download like civitai.com links; the
download goes through civitai.com.
- A Civitai download with a `filename` that already exists returns before any
API call.
## 1.3.1
### Startup
README.md
+25 -3
View File
@@ -442,7 +442,7 @@ Download models, LoRAs, and other assets directly from CivitAI, HuggingFace, and
- **Multi-Platform Support**: CivitAI, HuggingFace, and direct download URLs
- **Smart URL Detection**: Automatic detection of download source and file handling
- **API Token Support**: Optional authentication for private/gated models
- **API Tokens on the Server**: Civitai and Hugging Face tokens come from the ComfyUI server's environment, so workflows never carry them
- **Progress Reporting**: Real-time download progress with speed indicators
- **Resume Support**: Skip existing files or force re-download
- **Interrupt Handling**: Respects ComfyUI's "Cancel current run" button
@@ -450,9 +450,9 @@ Download models, LoRAs, and other assets directly from CivitAI, HuggingFace, and
- **Custom Filenames**: Override auto-detected filenames when needed
**Platform Features:**
- **CivitAI**: Model page URLs, version-specific downloads, API authentication
- **CivitAI**: Model page URLs, version-specific downloads, API authentication; civitai.com, civitai.green and civitai.red links
- **HuggingFace**: Blob and resolve URLs, branch/revision support, gated model access
- **Custom URLs**: Direct download links with bearer token authentication
- **Custom URLs**: Direct download links; only a token typed on the node is sent
**Use Cases:**
- Download models without leaving ComfyUI
@@ -461,6 +461,28 @@ Download models, LoRAs, and other assets directly from CivitAI, HuggingFace, and
- Build reproducible workflows with automatic model fetching
- Quickly test new models from the community
**Tokens:** many Civitai models and gated Hugging Face models need a token. Leave the node's `api_token` empty and set the tokens where ComfyUI starts; a token typed on the node is saved inside the workflow. Each token goes only to its own site.
| Site | Variable | Get a token |
|------|----------|-------------|
| Civitai | `CIVITAI_API_KEY` | civitai.com → Account settings → API Keys |
| Hugging Face | `HF_TOKEN` (or run `hf auth login` once) | huggingface.co → Settings → Access Tokens (read) |
Linux / macOS, in `~/.bashrc` or `~/.zshrc` (or the script that starts ComfyUI):
```bash
export CIVITAI_API_KEY="your-civitai-key"
export HF_TOKEN="hf_your-token"
```
Windows, once in a Command Prompt (then open a new window and start ComfyUI):
```bat
setx CIVITAI_API_KEY "your-civitai-key"
setx HF_TOKEN "hf_your-token"
```
For the portable build you can instead add `set CIVITAI_API_KEY=your-civitai-key` and `set HF_TOKEN=hf_your-token` above the `python` line in `run_nvidia_gpu.bat`.
Restart ComfyUI after setting them.
#### ⏱️ Workflow Timer
Real-time execution timer that displays workflow duration with millisecond precision.
kikotools/tools/model_downloader/base.py
+3 -2
View File
@@ -88,7 +88,7 @@ def safe_server_filename(name: Optional[str], default: str) -> str:
class _SafeRedirectHandler(urllib.request.HTTPRedirectHandler):
"""Only follow http(s) redirects and drop credentials on cross-host hops."""
"""Only follow http(s) redirects; drop credentials on cross-host or https-to-http hops."""
def redirect_request(self, req, fp, code, msg, headers, newurl):
if urlparse(newurl).scheme.lower() not in ALLOWED_SCHEMES:
@@ -96,8 +96,9 @@ class _SafeRedirectHandler(urllib.request.HTTPRedirectHandler):
newurl, code, "Redirect to disallowed URL scheme", headers, fp
)
new_request = super().redirect_request(req, fp, code, msg, headers, newurl)
old, new = urlparse(req.full_url), urlparse(newurl)
if new_request is not None and (
urlparse(newurl).netloc.lower() != urlparse(req.full_url).netloc.lower()
new.netloc.lower() != old.netloc.lower() or new.scheme.lower() != "https"
):
new_request.remove_header("Authorization")
return new_request
kikotools/tools/model_downloader/civitai.py
+19 -6
View File
@@ -9,6 +9,7 @@ import urllib.error
from typing import Optional, Dict, Any
from urllib.parse import urlparse, parse_qs, unquote, urljoin
from .detector import CIVITAI_HOSTS
from .base import (
USER_AGENT,
BaseDownloader,
@@ -23,7 +24,7 @@ API_BASE = "https://civitai.com/api/v1"
REDIRECT_STATUSES = (301, 302, 303, 307, 308)
HTTP_ERROR_HINTS = {
401: "Authentication required. Please provide a valid API token.",
401: "Authentication required. Set CIVITAI_API_KEY on the ComfyUI server (see the node's README).",
403: "Access forbidden. The model might be restricted or require special permissions.",
404: "Model not found. The URL might be incorrect or the model was removed.",
}
@@ -139,15 +140,27 @@ class CivitAIDownloader(BaseDownloader):
validate_filename(filename)
output_dir = self.resolve_output_dir(output_path)
# Validate that URL is from civitai.com domain
parsed_url = urlparse(url)
if parsed_url.netloc not in ("civitai.com", "www.civitai.com"):
if parsed_url.netloc not in CIVITAI_HOSTS:
raise ValueError(
f"Invalid URL: Only civitai.com URLs are supported, got {parsed_url.netloc}"
"Invalid URL: only civitai.com, civitai.green and civitai.red URLs "
f"are supported, got {parsed_url.netloc}"
)
# Convert web URL to API URL if needed
if "/api/download/models/" not in url:
# A file already there needs no API call when its name is known
if filename and not self.should_download(
os.path.join(output_dir, filename), force
):
logger.info("File already exists: %s", os.path.join(output_dir, filename))
return os.path.join(output_dir, filename)
# Downloads go through civitai.com whichever site the link came from
if "/api/download/models/" in url:
# https only: the request carries the token
url = urllib.parse.urlunparse(
parsed_url._replace(scheme="https", netloc="civitai.com")
)
else:
ids = self._parse_civitai_url(url)
# If we have a version ID, use it directly
kikotools/tools/model_downloader/detector.py
+19 -5
View File
@@ -5,9 +5,21 @@ from enum import Enum
from typing import TYPE_CHECKING, Optional
from urllib.parse import urlparse
from .tokens import resolve_token
if TYPE_CHECKING:
from .base import BaseDownloader
# Civitai serves models on civitai.com, civitai.green (SFW) and civitai.red (mature).
CIVITAI_HOSTS = (
"civitai.com",
"www.civitai.com",
"civitai.green",
"www.civitai.green",
"civitai.red",
"www.civitai.red",
)
class DownloaderType(Enum):
"""Types of supported downloaders"""
@@ -68,7 +80,7 @@ class URLDetector:
True if CivitAI URL
"""
# Validate exact domain match to prevent subdomain attacks
if parsed.netloc not in ("civitai.com", "www.civitai.com"):
if parsed.netloc not in CIVITAI_HOSTS:
return False
# Check for API download endpoint
@@ -111,7 +123,9 @@ class URLDetector:
Args:
url: URL to download from
api_token: Optional API token for authentication
api_token: Optional token typed on the node. When empty, Civitai and
Hugging Face use the server's token (see tokens.py); other hosts
get none.
Returns:
Appropriate downloader instance
@@ -124,14 +138,14 @@ class URLDetector:
if downloader_type == DownloaderType.CIVITAI:
from .civitai import CivitAIDownloader
return CivitAIDownloader(token=api_token)
return CivitAIDownloader(token=resolve_token("civitai", api_token))
elif downloader_type == DownloaderType.HUGGINGFACE:
from .huggingface import HuggingFaceDownloader
return HuggingFaceDownloader(token=api_token)
return HuggingFaceDownloader(token=resolve_token("huggingface", api_token))
else: # CUSTOM
from .custom import CustomDownloader
return CustomDownloader(token=api_token)
return CustomDownloader(token=resolve_token("custom", api_token))
kikotools/tools/model_downloader/huggingface.py
+4 -2
View File
@@ -5,7 +5,7 @@ import os
import urllib.request
import urllib.error
from typing import Optional
from urllib.parse import urlparse, quote
from urllib.parse import urlparse, urlunparse, quote
from .base import (
USER_AGENT,
@@ -18,7 +18,7 @@ from .base import (
logger = logging.getLogger(__name__)
HTTP_ERROR_HINTS = {
401: "Authentication required. Please provide a valid HuggingFace token.",
401: "Authentication required. Set HF_TOKEN on the ComfyUI server or run `hf auth login` (see the node's README).",
403: "Access forbidden. The model might be gated or require special permissions.",
404: "File not found. The URL might be incorrect or the file was removed.",
}
@@ -131,6 +131,8 @@ class HuggingFaceDownloader(BaseDownloader):
logger.info("Converted URL to: %s", download_url)
else:
download_url = url
# https only: the request may carry the server's token
download_url = urlunparse(urlparse(download_url)._replace(scheme="https"))
if not filename:
if url_info["filename"]:
kikotools/tools/model_downloader/node.py
+4 -1
View File
@@ -47,7 +47,10 @@ class ModelDownloaderNode(ComfyAssetsBaseNode):
{
"default": "",
"multiline": False,
"placeholder": "API token (CivitAI or HuggingFace)",
"placeholder": "Leave empty: uses CIVITAI_API_KEY / HF_TOKEN",
"tooltip": "Leave empty to use the server's token: CIVITAI_API_KEY for "
"Civitai, HF_TOKEN or `hf auth login` for Hugging Face. A token typed here "
"is saved inside the workflow, so anyone you share it with gets it.",
},
),
"force_download": (
kikotools/tools/model_downloader/tokens.py
+42
View File
@@ -0,0 +1,42 @@
"""API tokens for the downloaders, read on the server so shared workflows never carry them."""
import os
from pathlib import Path
from typing import Optional
# Environment variables checked per provider, first match wins.
ENV_VARS = {
"civitai": ("CIVITAI_API_KEY", "CIVITAI_API_TOKEN", "CIVITAI_TOKEN"),
"huggingface": ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"),
}
def _huggingface_login_token() -> Optional[str]:
"""The token `hf auth login` saved, if any."""
path = os.environ.get("HF_TOKEN_PATH")
if not path:
home = os.environ.get("HF_HOME") or os.path.join(
Path.home(), ".cache", "huggingface"
)
path = os.path.join(home, "token")
try:
return Path(path).read_text(encoding="utf-8").strip() or None
except OSError:
return None
def resolve_token(provider: str, override: Optional[str] = None) -> Optional[str]:
"""Token for provider: a token typed on the node, else the server's own.
Only Civitai and Hugging Face get server tokens, each its own, so a token
is never sent to another host.
"""
if override and override.strip():
return override.strip()
for name in ENV_VARS.get(provider, ()):
value = os.environ.get(name, "").strip()
if value:
return value
if provider == "huggingface":
return _huggingface_login_token()
return None
tests/unit/tools/model_downloader/test_downloads.py
+62
View File
@@ -282,6 +282,12 @@ class TestRedirectHandler:
)
assert new.get_header("Authorization") == "Bearer s"
def test_https_to_http_drops_authorization(self):
new = self.redirect(
"https://a.com/x", "http://a.com/y", {"Authorization": "Bearer s"}
)
assert not new.has_header("Authorization")
def test_disallowed_scheme_blocked(self):
with pytest.raises(urllib.error.HTTPError):
self.redirect("https://a.com/x", "file:///etc/passwd", {})
@@ -510,6 +516,14 @@ class TestCustom:
class TestHuggingFace:
def test_http_link_is_fetched_over_https(self, root, out, opener):
_, op = opener
op.open.return_value = FakeResponse(b"d", headers={"Content-Length": "1"})
make(HuggingFaceDownloader, root, token="hf_x").download(
"http://huggingface.co/u/r/resolve/main/m.safetensors", out
)
assert op.open.call_args.args[0].full_url.startswith("https://")
def test_download_blob_url_converted(self, root, out, opener):
_, op = opener
op.open.return_value = FakeResponse(b"d", headers={"Content-Length": "1"})
@@ -778,6 +792,54 @@ class TestCivitAI:
)
assert op.open.call_args.args[0].full_url.endswith("/api/download/models/9")
@pytest.mark.parametrize("host", ["civitai.red", "www.civitai.green"])
def test_red_and_green_links_download_through_civitai_com(
self, root, out, opener, host
):
_, op = opener
op.open.return_value = FakeResponse(b"d", status=200)
make(CivitAIDownloader, root).download(
f"https://{host}/api/download/models/77?type=Model",
out,
filename="m.safetensors",
)
assert (
op.open.call_args.args[0].full_url
== "https://civitai.com/api/download/models/77?type=Model"
)
def test_http_api_link_is_fetched_over_https(self, root, out, opener):
_, op = opener
op.open.return_value = FakeResponse(b"d", status=200)
make(CivitAIDownloader, root, token="t").download(
"http://civitai.red/api/download/models/77", out, filename="m.safetensors"
)
assert op.open.call_args.args[0].full_url == (
"https://civitai.com/api/download/models/77"
)
def test_red_model_page_resolves_version(self, root, out, opener):
_, op = opener
op.open.return_value = FakeResponse(b"d", status=200)
make(CivitAIDownloader, root).download(
"https://civitai.red/models/5?modelVersionId=77",
out,
filename="m.safetensors",
)
assert op.open.call_args.args[0].full_url == (
"https://civitai.com/api/download/models/77"
)
def test_existing_named_file_skips_the_network(self, root, out, opener):
_, op = opener
with open(os.path.join(out, "m.safetensors"), "wb") as f:
f.write(b"have")
path = make(CivitAIDownloader, root).download(
"https://civitai.com/models/5", out, filename="m.safetensors"
)
assert path == os.path.join(os.path.realpath(out), "m.safetensors")
op.open.assert_not_called()
def test_web_url_model_without_versions(self, root, out, opener):
_, op = opener
op.open.return_value = FakeResponse(b'{"modelVersions": []}')
tests/unit/tools/model_downloader/test_tokens.py
+83
View File
@@ -0,0 +1,83 @@
"""Token resolution: tokens come from the server, never from the workflow by default."""
import pytest
from kikotools.tools.model_downloader import tokens
from kikotools.tools.model_downloader.detector import URLDetector
TOKEN_VARS = (
"CIVITAI_API_KEY",
"CIVITAI_API_TOKEN",
"CIVITAI_TOKEN",
"HF_TOKEN",
"HUGGING_FACE_HUB_TOKEN",
"HF_TOKEN_PATH",
"HF_HOME",
)
@pytest.fixture(autouse=True)
def clean_env(monkeypatch, tmp_path):
for name in TOKEN_VARS:
monkeypatch.delenv(name, raising=False)
# No real `hf auth login` token leaks into the tests
monkeypatch.setenv("HF_HOME", str(tmp_path / "hf"))
class TestResolveToken:
def test_override_wins(self, monkeypatch):
monkeypatch.setenv("CIVITAI_API_KEY", "env")
assert tokens.resolve_token("civitai", " typed ") == "typed"
@pytest.mark.parametrize(
"name", ["CIVITAI_API_KEY", "CIVITAI_API_TOKEN", "CIVITAI_TOKEN"]
)
def test_civitai_env(self, monkeypatch, name):
monkeypatch.setenv(name, " civ ")
assert tokens.resolve_token("civitai", "") == "civ"
@pytest.mark.parametrize("name", ["HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"])
def test_huggingface_env(self, monkeypatch, name):
monkeypatch.setenv(name, "hf_x")
assert tokens.resolve_token("huggingface", None) == "hf_x"
def test_huggingface_login_file(self, tmp_path):
(tmp_path / "hf").mkdir()
(tmp_path / "hf" / "token").write_text("hf_saved\n")
assert tokens.resolve_token("huggingface") == "hf_saved"
def test_huggingface_token_path(self, monkeypatch, tmp_path):
path = tmp_path / "elsewhere"
path.write_text("hf_path")
monkeypatch.setenv("HF_TOKEN_PATH", str(path))
assert tokens.resolve_token("huggingface") == "hf_path"
def test_providers_do_not_share_tokens(self, monkeypatch):
monkeypatch.setenv("HF_TOKEN", "hf_x")
assert tokens.resolve_token("civitai") is None
monkeypatch.setenv("CIVITAI_API_KEY", "civ")
monkeypatch.delenv("HF_TOKEN")
assert tokens.resolve_token("huggingface") is None
def test_unknown_provider_only_uses_override(self, monkeypatch):
monkeypatch.setenv("HF_TOKEN", "hf_x")
assert tokens.resolve_token("custom") is None
assert tokens.resolve_token("custom", "t") == "t"
class TestDetectorTokens:
def test_civitai_gets_env_token(self, monkeypatch):
monkeypatch.setenv("CIVITAI_API_KEY", "civ")
d = URLDetector().get_downloader("https://civitai.com/models/1")
assert d.token == "civ"
def test_huggingface_gets_env_token(self, monkeypatch):
monkeypatch.setenv("HF_TOKEN", "hf_x")
d = URLDetector().get_downloader("https://huggingface.co/a/b/resolve/main/c")
assert d.token == "hf_x"
def test_custom_host_never_gets_env_tokens(self, monkeypatch):
monkeypatch.setenv("HF_TOKEN", "hf_x")
monkeypatch.setenv("CIVITAI_API_KEY", "civ")
d = URLDetector().get_downloader("https://example.com/model.safetensors")
assert d.token is None
tests/unit/tools/model_downloader/test_url_detection.py
+17
View File
@@ -28,6 +28,23 @@ class TestURLDetection:
result = detector.detect(url)
assert result == DownloaderType.CIVITAI
@pytest.mark.parametrize(
"url",
[
"https://civitai.red/models/123/name",
"https://www.civitai.red/api/download/models/456",
"https://civitai.green/models/123?modelVersionId=9",
"https://www.civitai.green/api/download/models/456",
],
)
def test_detect_civitai_red_and_green(self, url):
"""civitai.red and civitai.green links are Civitai links"""
assert URLDetector().detect(url) == DownloaderType.CIVITAI
def test_lookalike_civitai_domain_is_custom(self):
url = "https://civitai.red.evil.com/models/123"
assert URLDetector().detect(url) == DownloaderType.CUSTOM
def test_detect_huggingface_co_url(self):
"""Detect HuggingFace .co domain URL"""
url = "https://huggingface.co/username/repo-name/resolve/main/model.safetensors"