Server-Side Model Downloads

1. Overview & scope

ComfyUI workflows declare model dependencies inline via properties.models entries on loader nodes — each one carries a filename, a directory (e.g. loras, checkpoints), and a URL to fetch the file from. Until this feature, when a workflow loaded with a missing model, the frontend offered the user a download button that triggered a plain browser download via a synthesized <a download> click. Files landed in the user's Downloads folder; users then had to manually move them into ComfyUI/models/<directory>/. Gated HuggingFace models couldn't be downloaded at all without manual huggingface-cli login + hf_hub_download out-of-band.

This change moves the fetch to the server, lands files in the correct on-disk location, and adds authenticated HuggingFace support so gated models can be downloaded after a one-click OAuth flow.

2. Architecture at a glance

Key idea: every concern lives in app/model_downloader/ as a self-contained subsystem. Wiring into the rest of ComfyUI is two lines in server.py (register_routes(self.app)) and one feature-flag entry in comfy_api/feature_flags.py.

3. Download mechanism

3.1 The singleton

DOWNLOAD_SERVER (in app/model_downloader/download_server.py) is the process-wide registry of in-flight downloads. It exists so that:

• Multiple polling tabs see a single coherent view of "what's downloading right now."
• At most one download per model_id can run at a time — preventing two concurrent writers to the same destination path.
• Cancellation is just removal from the registry; the worker discovers cancellation on its next chunk-boundary check.

3.2 DownloadSession state

@dataclass
class DownloadSession:
    model_id: str             # e.g. "loras/my_lora.safetensors"
    url: str                  # the URL we're fetching from
    progress: Optional[float] # fraction in [0,1]; None until total known
    bytes_downloaded: int
    total_bytes: Optional[int]
    epoch: int                # see "atomicity" below

The registry is a plain dict[str, DownloadSession] guarded by a threading.Lock (callable from both the asyncio event-loop thread and the download-worker tasks).

3.3 Download lifecycle

3.4 Atomicity

Three independent atomicity guarantees, each addressing a different race:

1. File atomicity: downloads write to <final>.tmp and use os.replace for the promotion. A crashed/cancelled download leaves only the .tmp, never a partial-but-named-correct file that a loader would happily load and silently produce garbage outputs.
2. Registry atomicity: try_register holds the lock and inserts iff no entry exists. A second concurrent request for the same model_id returns None; the route then 409s and rolls back any sessions it had already registered in this batch.
3. Epoch-based cancellation: each DownloadSession carries an epoch counter assigned on registration. If a user cancels and then immediately re-triggers the download (same model_id), a new session with a new epoch is registered. The old worker, still running on the cancelled session, observes is_active(session) as False (epoch mismatch), rolls back its own .tmp, and exits without affecting the new session. Prevents the old worker's late finish() from accidentally evicting the new session.

3.5 Orphan cleanup

When the server restarts mid-download, any .tmp file is by definition orphaned. DOWNLOAD_SERVER.sweep_orphan_tmp_files() walks every registered model folder and removes *.tmp files. Idempotent; runs on the first download request rather than module import to keep the import path I/O-free.

3.6 Auth headers (HuggingFace)

When session.url is on huggingface.co and a token is stored, stream_to_disk attaches Authorization: Bearer <access_token> to the GET. Non-HF URLs receive no auth header (avoids token leakage to other hosts). This is HF's documented way to access gated repos with a personal access token — no reliance on huggingface_hub's download API.

4. HuggingFace OAuth mechanism

4.1 Why OAuth at all

Some HF repos are gated — the user has to accept a license / be approved before they can download. The bearer token from a logged-in HF account passes that gate. Rather than asking the user to paste a personal access token (security-awful UX), we run a proper OAuth 2.0 Authorization Code flow with PKCE, identical pattern to what the huggingface-cli login command does internally.

4.2 Where the token lives

Token shape (mirrors what HF returns from the token endpoint):

{
  "access_token":  "hf_oauth_…",
  "refresh_token": "…",         // null if not granted
  "expires_at":    1739895432.0, // absolute epoch seconds
  "scope":         "openid profile read-repos"
}

4.3 Token lifecycle

4.4 PKCE + state protection

Standard OAuth 2.0 PKCE (RFC 7636) with the SHA-256 method:

• Random 64-byte URL-safe verifier never leaves the server process.
• SHA-256 hash of the verifier sent as the code_challenge in the authorize URL.
• Random 32-byte state validated on callback; mismatches return 400 and the token exchange is skipped (CSRF defence).

5. API reference & frontend usage

All routes live under /api/, use kebab-case paths, and POST for input-bearing operations even when they're "read-only" — keeps semantics uniform and avoids URL-length limits when payloads grow.

{
  "models": {
    "loras/foo.safetensors": "https://huggingface.co/org/repo/resolve/main/foo.safetensors",
    "checkpoints/bar.safetensors": "https://huggingface.co/.../bar.safetensors"
  }
}

{
  "models": {
    "loras/foo.safetensors": {
      "state": "downloading",                     // "available" | "missing" | "downloading"
      "progress": {
        "bytes_downloaded": 1024000,
        "total_bytes":      29145431166,
        "progress":         0.000035                 // null until total known
      },
      "file_size":          29145431166,             // bytes; null if not probed
      "is_hf_downloadable": true                     // null for non-HF / probe failure
    },
    "checkpoints/bar.safetensors": {
      "state": "missing",
      "progress": null,
      "file_size": 1234567890,
      "is_hf_downloadable": false                    // gated, no access
    }
  },
  "hf_auth": {
    "token_available": true,
    "eligible":         true
  }
}

{
  "models": {
    "loras/foo.safetensors": "https://huggingface.co/.../foo.safetensors"
  }
}

HTTP 202 Accepted
{
  "accepted":  true,
  "scheduled": ["loras/foo.safetensors"]
}

HTTP 400 / 409
{
  "error": {
    "code":    "MODEL_NOT_DOWNLOADABLE",  // INVALID_MODEL_ID / URL_NOT_ALLOWED /
                                          // ALREADY_AVAILABLE / ALREADY_DOWNLOADING /
                                          // MODEL_NOT_DOWNLOADABLE / EMPTY_REQUEST
    "message": "…human-readable…",
    "details": { "model_id": "loras/foo.safetensors", "url": "https://…" }
  }
}

{ "model_id": "loras/foo.safetensors" }

{ "cancelled": true }    // or HTTP 404 with NOT_DOWNLOADING if no active session

{
  "token_available": true,
  "username":        "ogluzman"   // resolved via HfApi.whoami(); null if token invalid
}

{
  "authorize_url": "https://huggingface.co/oauth/authorize?client_id=…&state=…&code_challenge=…"
}

{ "logged_out": true }

6. Loopback eligibility gate

6.1 The rule

# app/model_downloader/hf_auth/eligibility.py

def is_hf_auth_eligible() -> bool:
    return _is_loopback(args.listen) and not args.multi_user

HF auth surfaces — both the login flow and the settings panel — appear iff this returns True.

6.2 Why it exists

Core ComfyUI has no authentication. Any HF token the server holds is implicitly shared by anyone who can reach the server. In a single-user local install that's fine — the OS user is the boundary, the loopback bind keeps remote actors out. In any other deployment it would be a credential-leak by misconfiguration:

• Non-loopback: anyone on the network who can reach the port could trigger downloads using the operator's HF account.
• --multi-user mode: multiple declared users (via the unauthenticated comfy-user header) would all share one HF token implicitly — Alice's prompts would silently fetch gated content as Bob.

Both cases are real credential leakage that the operator probably didn't realize they were enabling. The gate disables the feature instead of shipping a footgun.

6.3 What's gated

6.4 Implementation note

We had to inline a copy of is_loopback in eligibility.py (rather than importing from server.py) because comfy_api/feature_flags.py evaluates its registry at module-import time — earlier than server.py defines the helper. The inlined version is ~20 lines, mirrors server.is_loopback exactly, and is the kind of thing worth flagging if anyone ever does a "shared util" cleanup pass.

7. Probe caching strategy

The polling endpoint runs probe_url(url) for every model on every tick. To keep that cheap (HuggingFace round-trip per probe is >100ms), the probe layer caches what's safe to cache and recomputes what isn't:

Single-flight protection: a per-URL asyncio.Lock dedupes concurrent probes for the same URL — when many polls land in the same tick, exactly one of them runs the HF call and the others await the same result. Failures aren't cached (they're transient by nature; retry next call).

8. Frontend ↔ backend separation

8.1 Local dev workflow

# Backend  (one terminal)
cd ComfyUI
python main.py --listen 127.0.0.1 --port 8189 --cpu

# Frontend (another terminal)
cd ComfyUI_frontend
DEV_SERVER_COMFYUI_URL=http://127.0.0.1:8189 pnpm dev
# Vite serves at http://localhost:5173 and proxies /api/* to the backend

Open http://localhost:5173 in a browser — you get the Vite dev server with HMR, talking to your local backend.

8.2 Frontend integration points

• Feature-flag gate. MissingModelCard.vue renders the new MissingModelCardServerSide.vue when isServerSideDownloadsAvailable() returns true (the server_side_model_downloads server feature flag). Old servers silently fall through to the legacy in-browser download path.
• Eligibility-flag gate. The HF settings panel and login banner only render when hf_auth_eligible is true. Read once at startup.
• Single store of truth. useServerSideDownloadsStore (Pinia) holds the entire view of the polling response. Components read; only the store mutates.
• Three gated-row variants. The per-row gated message changes based on (eligible, logged-in):
  • Not eligible: "open the repository page to accept the license, then place the file in models/<dir>/ manually."
  • Eligible, not logged in: "log in with HuggingFace above to enable the download."
  • Eligible, logged in: "visit the repository page to accept the license, then come back to download" (license acceptance does the rest via the 1 Hz poll).

9. Tests & OpenAPI spec

9.1 Test coverage

~70 unit tests in two files under tests-unit/app_test/:

• model_downloader_test.py — allowlist, path validation, registry lifecycle (including epoch race semantics), orphan .tmp cleanup, precondition gating on all four model routes, atomic batch behavior.
• hf_auth_test.py — token store (save / load / chmod / corruption / refresh), eligibility under (listen, multi_user) matrix, URL parsing, probe caching (intrinsic + size + skip-when-not-downloadable), all three HF auth routes, PKCE primitives + authorize URL shape.

$ pytest tests-unit/app_test/model_downloader_test.py tests-unit/app_test/hf_auth_test.py -q
71 passed in 0.23s

9.2 OpenAPI spec

All six routes are documented in openapi.yaml with request/response schemas. The spec is hand-maintained — there's no codegen between handler signatures and the YAML. §10 flags this as a long-term tech-debt item.

Lint is enforced in CI via Spectral (.github/workflows/openapi-lint.yml); local run:

npx -y @stoplight/spectral-cli@6 lint openapi.yaml --ruleset .spectral.yaml --fail-severity=error

10. Open follow-ups & gotchas

10.1 Things deliberately not done

• Per-user HF tokens. Requires a real auth layer in core ComfyUI (sessions, login, identity). Out of scope; the loopback gate is the substitute.
• Hash verification of downloaded files. Some properties.models[*].hash entries carry a SHA. We don't verify; trust the source. Easy to add if needed (one method on stream_to_disk).
• Resumable downloads. A failed download starts from zero next time. Range requests + offset tracking would add it.
• Codegen for the OpenAPI spec. Spec is hand-maintained and lint-checked, not derived from handler signatures. Long-term direction is probably pydantic-driven schema export, but that's a project unto itself.

10.2 Convention summary

• Route paths: kebab-case for all new routes (/api/models-availability-status, etc.). Older endpoints use snake_case; newer assets endpoints use kebab; we picked kebab to match the newer direction.
• Error envelope: {"error": {"code": "MACHINE_READABLE", "message": "human", "details": {...}}}. Matches the pattern in app/assets/api/routes.py.
• Pydantic at the boundary: request schemas in schemas_in.py, response schemas in schemas_out.py, validated via Schema.model_validate(payload) in handlers.
• Logging prefix: all logs use [model_downloader] or [hf_auth] prefixes for grep-ability.

10.3 Useful greps

# Find every backend file touched by this feature
ls app/model_downloader app/model_downloader/api app/model_downloader/hf_auth

# Find every place is_loopback is consulted (3 callers)
grep -rn "is_loopback" --include="*.py" app/ server.py

# Confirm the HF OAuth callback port and redirect URI
grep -n "CALLBACK_PORT\|REDIRECT_URI" app/model_downloader/hf_auth/oauth.py

# Run the test suite for just this feature
.venv/bin/python -m pytest tests-unit/app_test/model_downloader_test.py \
                            tests-unit/app_test/hf_auth_test.py -q

11. HuggingFace OAuth app setup

Step-by-step walkthrough for creating the OAuth app whose client_id goes into HF_CLIENT_ID. Reflects what the HuggingFace settings UI looked like at the time this feature was developed; HF occasionally moves things around but the fields themselves are stable.

11.1 Navigation

1. Sign in at huggingface.co with the Comfy-Org-controlled account that should own the app.
2. Open user settings (avatar menu → Settings).
3. In the left sidebar, click Connected Apps. (Not Access Tokens — that's for personal access tokens, a different concept.)
4. Click Create app (or similar — the button label has varied).

11.2 Fields to fill in

11.3 Scopes — exactly which boxes to tick

HF groups scopes into sections. The bare minimum for this feature is three checkboxes total. Leave everything else off.

11.4 Public app + PKCE

After creation, HF will label the app a Public app and explicitly note: "No client secret. Use PKCE or device code flow for authentication." This is expected and correct — we use PKCE (see §4). Do not click Add client secret; we don't need it and having one without using it would be a future footgun.

11.5 Wire the client_id into the code

The Credentials section of the new app shows a Client ID in the form of a UUID (e.g. a8189e14-9246-4f19-bd6a-a307bdcb9276). Copy that value and paste it verbatim into:

# app/model_downloader/hf_auth/oauth.py  (around line 49)
HF_CLIENT_ID = "paste-the-uuid-here"

That's the only code change required. Restart ComfyUI; POST /api/hf-auth-login-start should now produce an authorize_url that huggingface.co accepts.

11.6 Test the round-trip

1. Start ComfyUI on loopback: python main.py --listen 127.0.0.1 --port 8189
2. Confirm eligibility:

   curl -s http://127.0.0.1:8189/api/features | grep hf_auth_eligible
   # expect: "hf_auth_eligible": true

3. Trigger the login flow:

   curl -s -X POST http://127.0.0.1:8189/api/hf-auth-login-start | python3 -m json.tool
   # expect: {"authorize_url": "https://huggingface.co/oauth/authorize?client_id=<your-uuid>&..."}

4. Open authorize_url in a browser. The consent screen should display the Application Name you chose and list the three scopes (openid, profile, gated-repos). Click Authorize.
5. HF redirects to http://127.0.0.1:41954/api/auth/huggingface/callback?code=...&state=.... Our local callback server completes the token exchange and renders a "Login complete" page.
6. Confirm token is held:

   curl -s http://127.0.0.1:8189/api/hf-auth-token-status | python3 -m json.tool
   # expect: {"token_available": true, "username": "your-hf-username"}

Once that round-trip works, the missing-models card will use the token automatically for every subsequent gated probe and download.

11.7 If you need to change the callback port

The port 41954 is arbitrary — chosen to be high and unlikely to collide. If you ever need to change it, three things must move together:

• CALLBACK_PORT in app/model_downloader/hf_auth/oauth.py.
• The Redirect URL registered on the HuggingFace app (must match exactly, including port).
• The redirect-URI constant in any test fixtures (search for the port number in tests-unit/app_test/hf_auth_test.py).

If they drift out of sync, HF will reject the redirect with a redirect_uri_mismatch error and the callback never lands.

Generated as a feature handover. Living document — keep it updated as the feature evolves, or replace with a proper docs site entry once one exists.