ComfyUI/app/model_downloader/http_client.py
Alex 351119eb05 feat(model_downloader): add server-side model downloads with gated-repo support
Lets ComfyUI fetch the models a workflow needs directly on the server,
so users no longer have to locate each file and drop it into the correct
folder by hand.

Crucially it supports gated HuggingFace repositories: the user logs in
once via HuggingFace, after which the server can download models that
require license acceptance or authentication — previously a manual,
error-prone step. The frontend can surface per-model availability and
download progress through the accompanying API.
2026-06-25 15:59:41 +03:00

64 lines
2.0 KiB
Python

"""Lazy module-level aiohttp ClientSession.
A single shared session means TLS handshakes are reused across HEAD probes
and the subsequent GETs to the same host (HuggingFace is the dominant
case), which is a noticeable speedup on cold connections.
We deliberately don't close the session at process exit — aiohttp's
warning about unclosed sessions is benign at shutdown, and adding atexit
plumbing buys nothing because the OS reclaims the sockets anyway. The
session lifetime is the lifetime of the Python process.
"""
from __future__ import annotations
import asyncio
import ssl
from typing import Optional
import aiohttp
import certifi
# Larger per-host pool than aiohttp's default (=100 total / =0 per host)
# so concurrent gated probes + a download to the same host don't queue.
_CONNECTOR_LIMIT_PER_HOST = 8
_session: Optional[aiohttp.ClientSession] = None
_lock = asyncio.Lock()
def ssl_context() -> ssl.SSLContext:
"""TLS context pinned to certifi's CA bundle.
aiohttp's default context uses the OS trust store, which isn't wired up
on some Python installs (python.org macOS, slim containers) — there TLS
to huggingface.co fails with CERTIFICATE_VERIFY_FAILED.
"""
return ssl.create_default_context(cafile=certifi.where())
async def get_session() -> aiohttp.ClientSession:
"""Return the shared session, creating it on first call."""
global _session
if _session is not None and not _session.closed:
return _session
async with _lock:
if _session is None or _session.closed:
connector = aiohttp.TCPConnector(
limit_per_host=_CONNECTOR_LIMIT_PER_HOST,
ssl=ssl_context(),
)
_session = aiohttp.ClientSession(connector=connector)
return _session
def parse_content_length(value: Optional[str]) -> Optional[int]:
"""Parse a byte-count header value, or None if absent/malformed/negative."""
if not value:
return None
try:
n = int(value)
except ValueError:
return None
return n if n >= 0 else None