feat(download): support ModelScope repositories in the URL downloader

ModelScope became a linkable source, but downloading from it was impossible:
the URL picker only recognised huggingface.co, the file listing hit a
huggingface-only endpoint, the resolve URL was hardcoded, and the default
path template always wrote into a `huggingface/` directory.

Move the download knowledge into the providers so the handlers stay generic:

- `ModelSource` gains `list_files()`, `file_download_url()`,
  `default_revision` and `default_subdir`. `HuggingFaceSource` keeps the Hub
  tree API (`/api/models/{id}/tree/{rev}`, LFS-aware sizes, `main`).
  `ModelScopeSource` uses `/api/v1/models/{id}/repo/files?Revision=master`
  — which reports real byte sizes for LFS files, so no HEAD probe is needed,
  and which only accepts `master` (an HF-imported repo still 404s on `main`)
  — and downloads through `/models/{id}/resolve/{rev}/{path}`. That URL
  redirects to a CDN target carrying a time-limited `auth_key`, so it is
  rebuilt on every request and never cached, which is also what keeps
  resumable Range requests working.
- `hf_handlers.py`/`HfHandler` become `model_source_handlers.py`/
  `ModelSourceHandler` with `list_model_source_files` and
  `download_model_source`. New routes `/api/lm/model-source-files` and
  `/api/lm/download-model-source`; the old `/api/lm/hf-repo-files` and
  `/api/lm/download-hf-model` paths stay as aliases, and a payload without
  `platform` still means Hugging Face, so existing callers are unaffected.
- A downloaded sidecar now records `source_platform` + `source_url` (with the
  `hf_url` alias only for Hugging Face) instead of always writing `hf_url`,
  and `use_default_paths` files ModelScope downloads under
  `modelscope/<owner>/<repo>`. The now-unused shared HF aiohttp session and
  its shutdown hook are gone; providers open short-lived sessions.
- Frontend: `detectUrlType` returns the platform-neutral
  `model-source-repo` / `model-source-file` plus an explicit `platform`, the
  DownloadManager's `hf*` state and methods are renamed to `source*`, every
  `source === 'huggingface'` check becomes `isExternalModelSource()`, and
  batch groups are keyed by `platform:repo` so the same `owner/name` on two
  sites renders as two groups. A bare `owner/name` still means Hugging Face.
- `is_valid_source_id()` centralises repo-id validation (exactly
  `owner/name`, no traversal, no leading dot). This also fixes the old HF
  download check that rejected any dot in the name, i.e. legitimate repos
  such as `black-forest-labs/FLUX.1-dev`.

Verified against the live APIs: the example repo lists 8 weight files with
correct sizes, and a ranged GET of the built resolve URL returns 206 after
following the redirect to the CDN. Backend 2853 passed; frontend 1143 JS +
91 Vue passed. The nine locales carry the refreshed download copy in the
next commit.
This commit is contained in:
Will Miao
2026-09-14 07:42:51 +08:00
parent b9bf006998
commit 38d4c59b4c
22 changed files with 1953 additions and 667 deletions
+128 -1
View File
@@ -20,12 +20,15 @@ HTTP handlers never need site-specific branching.
from __future__ import annotations
import logging
import os
import re
from dataclasses import dataclass
from typing import Any, Optional
from typing import Any, Iterable, Optional
import aiohttp
from ...utils.constants import MODEL_FILE_EXTENSIONS
logger = logging.getLogger(__name__)
#: Shared HTTP timeout for model-card fetches.
@@ -58,6 +61,36 @@ class SourceRef:
"""Canonical URL of the model page."""
class ModelSourceError(Exception):
"""Raised when a model source cannot satisfy a request.
Carries the HTTP status the API handler should answer with, so the
handlers stay free of per-site error mapping.
"""
def __init__(self, message: str, status: int = 502) -> None:
super().__init__(message)
self.status = status
#: Repository ids are always exactly ``owner/name``. Components may contain
#: dots (``black-forest-labs/FLUX.1-dev``) but must not be empty, ``.`` / ``..``,
#: or start with a dot - the id is used as a path segment on disk.
_SOURCE_ID_COMPONENT = re.compile(r"^[A-Za-z0-9_][A-Za-z0-9_.\-]*$")
def is_valid_source_id(source_id: str) -> bool:
"""Return ``True`` when *source_id* is a safe ``owner/name`` repository id."""
if not source_id or not isinstance(source_id, str) or source_id.count("/") != 1:
return False
owner, name = source_id.split("/", 1)
return all(
part and part not in (".", "..") and _SOURCE_ID_COMPONENT.match(part)
for part in (owner, name)
)
async def fetch_text(url: str, *, timeout: int = HTTP_TIMEOUT) -> str:
"""Fetch *url* and return its body as text, or ``""`` on any failure.
@@ -80,6 +113,34 @@ async def fetch_text(url: str, *, timeout: int = HTTP_TIMEOUT) -> str:
return ""
async def fetch_json(
url: str, *, timeout: int = HTTP_TIMEOUT
) -> tuple[int, Any]:
"""Fetch *url* and return ``(status, parsed_body)``.
Unlike :func:`fetch_text` this reports the status, because callers such as
the file-listing endpoints need to distinguish "repo not found" (404) from
a transport failure. ``parsed_body`` is ``None`` when the response is not
JSON or the request failed outright (status ``0``).
"""
try:
async with aiohttp.ClientSession(
headers={"User-Agent": USER_AGENT},
timeout=aiohttp.ClientTimeout(total=timeout),
) as session:
async with session.get(url) as resp:
if resp.status != 200:
return resp.status, None
try:
return resp.status, await resp.json(content_type=None)
except Exception:
return resp.status, None
except Exception as exc: # pragma: no cover - network dependent
logger.debug("Failed to fetch %s: %s", url, exc)
return 0, None
class ModelSource:
"""Description and I/O for one external model hosting site."""
@@ -95,6 +156,12 @@ class ModelSource:
#: Whether models can be downloaded directly from this site.
supports_download: bool = False
#: Branch used when the caller does not pass an explicit revision.
default_revision: str = ""
#: Sub-directory the "use default paths" template places downloads in.
default_subdir: str = ""
#: Lenient pattern used to recognise URLs already stored in metadata.
#: Captures the site-specific source id in group ``id``.
url_pattern: re.Pattern[str] | None = None
@@ -163,6 +230,45 @@ class ModelSource:
return ""
# ------------------------------------------------------------------
# Download support
# ------------------------------------------------------------------
async def list_files(
self, source_id: str, revision: str = ""
) -> list[dict[str, Any]]:
"""List downloadable weight files in *source_id*.
Returns ``[{"filename": <repo-relative path>, "size": <bytes>}]``,
largest first, filtered to :data:`MODEL_FILE_EXTENSIONS`. Sites
without download support return an empty list.
Raises :class:`ModelSourceError` when the repository cannot be read,
so the handler can surface "not found" separately from a transport
failure.
"""
return []
def file_download_url(
self, source_id: str, filename: str, revision: str = ""
) -> str:
"""Return the direct (redirecting) download URL for one file."""
raise ModelSourceError(
f"{self.label or self.platform} does not support downloads", status=400
)
def resolve_revision(self, revision: str = "") -> str:
"""Return *revision*, falling back to this site's default branch."""
return revision or self.default_revision
def page_url_for_file(self, source_id: str, filename: str) -> str:
"""Return the human-facing page for *filename* inside *source_id*."""
return self.canonical_url(source_id)
def __repr__(self) -> str: # pragma: no cover - debugging aid
return f"<ModelSource {self.platform}>"
@@ -175,12 +281,33 @@ def clean_source_url(url: Any) -> str:
return url.strip()
def filter_weight_files(entries: Iterable[tuple[str, int]]) -> list[dict[str, Any]]:
"""Keep model-weight files from ``(path, size)`` pairs, largest first.
Every site lists a lot more than weights (READMEs, configs, tokenizers,
…); the download picker only ever wants the files ComfyUI can load, which
is exactly :data:`MODEL_FILE_EXTENSIONS`.
"""
files = [
{"filename": path, "size": int(size or 0)}
for path, size in entries
if path and os.path.splitext(path)[1].lower() in MODEL_FILE_EXTENSIONS
]
files.sort(key=lambda entry: entry["size"], reverse=True)
return files
__all__ = [
"GROUP_PREFIXES",
"HTTP_TIMEOUT",
"ModelSource",
"ModelSourceError",
"SourceRef",
"USER_AGENT",
"clean_source_url",
"fetch_json",
"fetch_text",
"filter_weight_files",
"is_valid_source_id",
]