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
+1 -1
View File
@@ -71,7 +71,7 @@ Enriches models linked to an external model site with metadata extracted by an L
| Platform | Link | AI enrichment | Direct download | | Platform | Link | AI enrichment | Direct download |
| --- | --- | --- | --- | | --- | --- | --- | --- |
| Hugging Face | yes | yes | yes | | Hugging Face | yes | yes | yes |
| ModelScope | yes | yes | no | | ModelScope | yes | yes | yes |
| TensorArt | yes | no (see below) | no | | TensorArt | yes | no (see below) | no |
TensorArt is link-only: `tensor.art` sits behind a Cloudflare managed challenge and its internal API requires session authorization, so the backend cannot read its model pages. Linking still stores the canonical page URL and the "View on TensorArt" link works. TensorArt is link-only: `tensor.art` sits behind a Cloudflare managed challenge and its internal API requires session authorization, so the backend cannot read its model pages. Linking still stores the canonical page URL and the "View on TensorArt" link works.
+4 -4
View File
@@ -1395,9 +1395,9 @@
"download": { "download": {
"title": "Download Model from URL", "title": "Download Model from URL",
"titleWithType": "Download {type} from URL", "titleWithType": "Download {type} from URL",
"civitaiUrl": "CivitAI URL(s):", "civitaiUrl": "Model URL(s):",
"placeholder": "https://civitai.com/models/...", "placeholder": "https://civitai.com/models/...",
"urlHint": "Enter one CivitAI, CivArchive, or Hugging Face URL per line. Supports multiple URLs for batch download.", "urlHint": "Enter one CivitAI, CivArchive, Hugging Face, or ModelScope URL per line. Supports multiple URLs for batch download.",
"selectHfFiles": "Select file(s) to download from this repository:", "selectHfFiles": "Select file(s) to download from this repository:",
"selectAll": "Select All", "selectAll": "Select All",
"fetchingRepoFiles": "Fetching repository files...", "fetchingRepoFiles": "Fetching repository files...",
@@ -1430,9 +1430,9 @@
"inLibrary": "In Library" "inLibrary": "In Library"
}, },
"errors": { "errors": {
"invalidUrl": "Invalid CivitAI URL format", "invalidUrl": "Invalid model URL format",
"noVersions": "No versions available for this model", "noVersions": "No versions available for this model",
"mixedSources": "Cannot mix CivitAI and Hugging Face URLs in the same batch.", "mixedSources": "Cannot mix CivitAI and Hugging Face / ModelScope URLs in the same batch.",
"noModelFiles": "No model files found in this repository." "noModelFiles": "No model files found in this repository."
}, },
"status": { "status": {
-7
View File
@@ -472,12 +472,5 @@ class LoraManager:
scanner.cancel_task() scanner.cancel_task()
logger.debug("LoRA Manager: Cancelled %s", name) logger.debug("LoRA Manager: Cancelled %s", name)
# Close shared aiohttp sessions to avoid "Unclosed client session" warnings
try:
from py.routes.handlers.hf_handlers import close_hf_api_session
await close_hf_api_session()
except Exception as exc:
logger.debug("Error closing HF API session: %s", exc)
except Exception as e: except Exception as e:
logger.error(f"Error during cleanup: {e}", exc_info=True) logger.error(f"Error during cleanup: {e}", exc_info=True)
+10 -7
View File
@@ -55,7 +55,7 @@ from ...utils.constants import (
VALID_LORA_TYPES, VALID_LORA_TYPES,
VALID_OTHER_CIVITAI_TYPES, VALID_OTHER_CIVITAI_TYPES,
) )
from .hf_handlers import HfHandler from .model_source_handlers import ModelSourceHandler
from .agent_handlers import AgentHandler from .agent_handlers import AgentHandler
from .download_routing_handlers import DownloadRoutingHandler from .download_routing_handlers import DownloadRoutingHandler
from .model_handlers import ModelCivitaiHandler from .model_handlers import ModelCivitaiHandler
@@ -4001,7 +4001,7 @@ class MiscHandlerSet:
doctor: DoctorHandler, doctor: DoctorHandler,
example_workflows: ExampleWorkflowsHandler, example_workflows: ExampleWorkflowsHandler,
base_model: BaseModelHandlerSet, base_model: BaseModelHandlerSet,
hf_handler: Any = None, model_source_handler: Any = None,
agent_handler: Any = None, agent_handler: Any = None,
download_routing: Any = None, download_routing: Any = None,
) -> None: ) -> None:
@@ -4022,7 +4022,7 @@ class MiscHandlerSet:
self.doctor = doctor self.doctor = doctor
self.example_workflows = example_workflows self.example_workflows = example_workflows
self.base_model = base_model self.base_model = base_model
self.hf_handler = hf_handler self.model_source_handler = model_source_handler
self.agent_handler = agent_handler self.agent_handler = agent_handler
self.download_routing = download_routing self.download_routing = download_routing
@@ -4076,10 +4076,13 @@ class MiscHandlerSet:
"get_example_workflows": self.example_workflows.get_example_workflows, "get_example_workflows": self.example_workflows.get_example_workflows,
"get_example_workflow": self.example_workflows.get_example_workflow, "get_example_workflow": self.example_workflows.get_example_workflow,
# Hugging Face handlers # Hugging Face handlers
"get_hf_repo_files": self.hf_handler.get_hf_repo_files, # External model sources (Hugging Face / ModelScope)
"download_hf_model": self.hf_handler.download_hf_model, "list_model_source_files": self.model_source_handler.list_model_source_files,
"set_hf_url": self.hf_handler.set_hf_url, "download_model_source": self.model_source_handler.download_model_source,
"get_model_sources": self.hf_handler.get_model_sources, "get_hf_repo_files": self.model_source_handler.list_model_source_files,
"download_hf_model": self.model_source_handler.download_model_source,
"set_hf_url": self.model_source_handler.set_hf_url,
"get_model_sources": self.model_source_handler.get_model_sources,
# Agent skill handlers # Agent skill handlers
"get_agent_skills": self.agent_handler.get_agent_skills, "get_agent_skills": self.agent_handler.get_agent_skills,
"execute_agent_skill": self.agent_handler.execute_agent_skill, "execute_agent_skill": self.agent_handler.execute_agent_skill,
@@ -1,8 +1,13 @@
"""Handlers for Hugging Face model listing and download. """Handlers for external model sources: linking, file listing and downloads.
Minimal MVP implementation uses direct HTTP to the HF API for file Covers every site registered in :mod:`py.services.model_sources`. The module
listing and the project's existing aiohttp-based Downloader for was Hugging Face only (``hf_handlers.py`` / ``HfHandler``) until ModelScope
downloading. No huggingface_hub dependency required. downloads were added; the per-site differences now live in the providers, so
this file has no platform branches beyond the capability lookups.
The historical route paths (``/api/lm/set-hf-url``, ``/api/lm/hf-repo-files``,
``/api/lm/download-hf-model``) are still registered as aliases of the generic
handlers, so existing callers keep working.
""" """
from __future__ import annotations from __future__ import annotations
@@ -10,10 +15,8 @@ from __future__ import annotations
import json import json
import logging import logging
import os import os
import re
from typing import Any from typing import Any
import aiohttp
from aiohttp import web from aiohttp import web
from ...config import config from ...config import config
@@ -23,14 +26,17 @@ from ...services.downloader import (
) )
from ...services.aria2_downloader import Aria2Downloader from ...services.aria2_downloader import Aria2Downloader
from ...services.model_sources import ( from ...services.model_sources import (
ModelSourceError,
SourceRef,
detect_source, detect_source,
get_download_source,
is_valid_source_id,
list_sources, list_sources,
normalize_metadata_source, normalize_metadata_source,
) )
from ...services.settings_manager import get_settings_manager from ...services.settings_manager import get_settings_manager
from ...services.service_registry import ServiceRegistry from ...services.service_registry import ServiceRegistry
from ...services.websocket_manager import ws_manager from ...services.websocket_manager import ws_manager
from ...utils.constants import MODEL_FILE_EXTENSIONS
from ...utils.metadata_manager import MetadataManager from ...utils.metadata_manager import MetadataManager
from ...utils.models import LoraMetadata, CheckpointMetadata, EmbeddingMetadata from ...utils.models import LoraMetadata, CheckpointMetadata, EmbeddingMetadata
@@ -39,28 +45,6 @@ logger = logging.getLogger(__name__)
_DEFAULT_MODEL_CLASS = LoraMetadata _DEFAULT_MODEL_CLASS = LoraMetadata
_DEFAULT_SCANNER_GETTER = "get_lora_scanner" _DEFAULT_SCANNER_GETTER = "get_lora_scanner"
# Shared aiohttp session for HF API calls (created on first use)
_hf_api_session: aiohttp.ClientSession | None = None
async def _get_hf_api_session() -> aiohttp.ClientSession:
"""Get or create the shared aiohttp session for HF API calls."""
global _hf_api_session # needed because we reassign the module-level name
if _hf_api_session is None or _hf_api_session.closed:
_hf_api_session = aiohttp.ClientSession(
headers={"User-Agent": "ComfyUI-LoRA-Manager/1.0"},
timeout=aiohttp.ClientTimeout(total=30),
)
return _hf_api_session
async def close_hf_api_session() -> None:
"""Close the shared HF API session, if it was ever created."""
global _hf_api_session
if _hf_api_session is not None and not _hf_api_session.closed:
await _hf_api_session.close()
_hf_api_session = None
def _infer_model_type(model_root: str) -> tuple[Any, str]: def _infer_model_type(model_root: str) -> tuple[Any, str]:
"""Determine model class and scanner by matching ``model_root`` against the """Determine model class and scanner by matching ``model_root`` against the
@@ -101,18 +85,19 @@ def _infer_model_type(model_root: str) -> tuple[Any, str]:
return _DEFAULT_MODEL_CLASS, _DEFAULT_SCANNER_GETTER return _DEFAULT_MODEL_CLASS, _DEFAULT_SCANNER_GETTER
async def _save_hf_metadata(dest_path: str, repo: str, model_root: str) -> None: async def _save_source_metadata(
dest_path: str, ref: SourceRef, model_root: str
) -> None:
"""Create a proper .metadata.json and add the model to the scanner cache. """Create a proper .metadata.json and add the model to the scanner cache.
Uses ``MetadataManager.create_default_metadata()`` which computes the Uses ``MetadataManager.create_default_metadata()`` which computes the
SHA256 hash, extracts safetensors header metadata (base_model), and SHA256 hash, extracts safetensors header metadata (base_model), and
produces a fully-populated ``LoraMetadata`` (or ``CheckpointMetadata`` / produces a fully-populated ``LoraMetadata`` (or ``CheckpointMetadata`` /
``EmbeddingMetadata``) object. We then overlay HF-specific fields and ``EmbeddingMetadata``) object. We then overlay the external-source fields
register the model in the in-memory scanner cache so it appears and register the model in the in-memory scanner cache so it appears
immediately without a full filesystem walk. immediately without a full filesystem walk.
""" """
try: try:
hf_url = f"https://huggingface.co/{repo}"
model_class, scanner_getter_name = _infer_model_type(model_root) model_class, scanner_getter_name = _infer_model_type(model_root)
# 1. Create proper metadata (computes SHA256, reads safetensors headers) # 1. Create proper metadata (computes SHA256, reads safetensors headers)
@@ -123,15 +108,21 @@ async def _save_hf_metadata(dest_path: str, repo: str, model_root: str) -> None:
logger.warning("create_default_metadata returned None for %s", dest_path) logger.warning("create_default_metadata returned None for %s", dest_path)
return return
# 2. Overlay HF-specific fields # 2. Overlay the external-source fields (`hf_url` is written by
metadata._unknown_fields["hf_url"] = hf_url # normalisation for Hugging Face only)
metadata._unknown_fields["source_url"] = hf_url fields = metadata._unknown_fields
metadata._unknown_fields["source_platform"] = "huggingface" fields["source_url"] = ref.url
metadata.from_civitai = False # HF models are not from CivitAI fields["source_platform"] = ref.platform
if ref.platform == "huggingface":
fields["hf_url"] = ref.url
metadata.from_civitai = False # externally-sourced models are not from CivitAI
# 3. Save metadata atomically # 3. Save metadata atomically
await MetadataManager.save_metadata(dest_path, metadata) await MetadataManager.save_metadata(dest_path, metadata)
logger.info("Saved HF metadata (with hf_url) for %s", dest_path) logger.info(
"Saved %s metadata (source=%s) for %s",
ref.platform, ref.url, dest_path,
)
# 4. Determine relative folder path for cache # 4. Determine relative folder path for cache
# model_root is an absolute path; dest_path is under it # model_root is an absolute path; dest_path is under it
@@ -145,13 +136,12 @@ async def _save_hf_metadata(dest_path: str, repo: str, model_root: str) -> None:
if scanner_getter is not None: if scanner_getter is not None:
scanner = await scanner_getter() scanner = await scanner_getter()
if scanner is not None: if scanner is not None:
metadata_dict = metadata.to_dict() metadata_dict = normalize_metadata_source(metadata.to_dict())
metadata_dict["hf_url"] = hf_url
await scanner.add_model_to_cache(metadata_dict, folder) await scanner.add_model_to_cache(metadata_dict, folder)
logger.info("Added %s to scanner cache (folder=%s)", dest_path, folder) logger.info("Added %s to scanner cache (folder=%s)", dest_path, folder)
except Exception as exc: except Exception as exc:
logger.warning("Failed to save HF metadata for %s: %s", dest_path, exc) logger.warning("Failed to save source metadata for %s: %s", dest_path, exc)
def _find_matching_root(dest_dir: str) -> str | None: def _find_matching_root(dest_dir: str) -> str | None:
@@ -193,14 +183,23 @@ async def _add_to_scanner_cache(dest_path: str, metadata: dict[str, Any]) -> Non
await scanner.update_single_model_cache(dest_path, dest_path, metadata) await scanner.update_single_model_cache(dest_path, dest_path, metadata)
class HfHandler: def _unsupported_platform_error(platform: str) -> web.Response:
"""Handle Hugging Face model browsing and download.""" supported = ", ".join(source.label for source in list_sources() if source.supports_download)
return web.json_response(
{"error": f"'{platform}' does not support downloads. Supported: {supported}"},
status=400,
)
class ModelSourceHandler:
"""Handle external model browsing, linking and downloads."""
async def get_model_sources(self, request: web.Request) -> web.Response: async def get_model_sources(self, request: web.Request) -> web.Response:
"""List the external model sites the UI can link a model to. """List the external model sites the UI can link a model to.
Used by the "Link Model" dialog to validate URLs client-side and to Used by the "Link Model" dialog to validate URLs client-side, to
explain which sites support AI metadata enrichment. explain which sites support AI metadata enrichment, and to pick the
right download endpoint/revision.
""" """
return web.json_response([ return web.json_response([
@@ -209,6 +208,7 @@ class HfHandler:
"label": source.label, "label": source.label,
"supports_enrichment": source.supports_enrichment, "supports_enrichment": source.supports_enrichment,
"supports_download": source.supports_download, "supports_download": source.supports_download,
"default_revision": source.default_revision,
"example_url": source.canonical_url( "example_url": source.canonical_url(
"user/repo" if source.platform != "tensorart" else "827823520299086029" "user/repo" if source.platform != "tensorart" else "827823520299086029"
), ),
@@ -219,10 +219,12 @@ class HfHandler:
async def set_hf_url(self, request: web.Request) -> web.Response: async def set_hf_url(self, request: web.Request) -> web.Response:
"""Link a model file to its page on an external model site. """Link a model file to its page on an external model site.
Accepts ``source_url`` (preferred) or the legacy ``hf_url`` / Accepts ``source_url`` (preferred) or the legacy ``hf_url`` / ``url``
``url`` payload key. Hugging Face, ModelScope, and TensorArt URLs payload key. Every registered site is recognised and the platform is
are recognised; the platform is stored alongside the canonical URL. stored alongside the canonical URL. TensorArt models can be linked and
TensorArt models can be linked and browsed, but not AI-enriched. browsed, but not AI-enriched.
The route path keeps its historical ``set-hf-url`` name.
""" """
try: try:
@@ -337,74 +339,60 @@ class HfHandler:
status=500, status=500,
) )
async def get_hf_repo_files(self, request: web.Request) -> web.Response: async def list_model_source_files(self, request: web.Request) -> web.Response:
"""List model-weight files from a HF repo with real file sizes. """List the downloadable weight files of an external repository.
Uses the HF tree API endpoint which returns accurate file sizes Query params: ``platform``, ``repo`` (``owner/name``), ``revision``
(including LFS-tracked files), unlike the model info endpoint. (optional; each site has its own default branch).
Returns a JSON array of ``{"filename", "size"}``, largest first
the same shape the Hugging Face endpoint has always returned.
""" """
repo = request.query.get("repo", "").strip()
if not repo or "/" not in repo: platform = (request.query.get("platform") or "").strip()
repo = (request.query.get("repo") or "").strip()
revision = (request.query.get("revision") or "").strip()
source = get_download_source(platform)
if source is None:
return _unsupported_platform_error(platform)
if not is_valid_source_id(repo):
return web.json_response( return web.json_response(
{"error": "Missing or invalid 'repo' parameter (expected user/repo)"}, {"error": "Missing or invalid 'repo' parameter (expected owner/name)"},
status=400, status=400,
) )
url = f"https://huggingface.co/api/models/{repo}/tree/main"
try: try:
session = await _get_hf_api_session() files = await source.list_files(repo, revision)
async with session.get(url) as resp: except ModelSourceError as exc:
if resp.status == 404: return web.json_response({"error": str(exc)}, status=exc.status)
return web.json_response(
{"error": f"Repo '{repo}' not found"}, status=404
)
if resp.status != 200:
text = await resp.text()
return web.json_response(
{"error": f"HF API error {resp.status}: {text[:200]}"},
status=resp.status,
)
tree: list[dict[str, Any]] = await resp.json()
except Exception as exc: except Exception as exc:
logger.error("Failed to fetch HF repo files: %s", exc) logger.error("Failed to list %s files in %s: %s", platform, repo, exc)
return web.json_response({"error": str(exc)}, status=502) return web.json_response({"error": str(exc)}, status=502)
files: list[dict[str, Any]] = []
for entry in tree:
path: str = entry.get("path", "")
ext = os.path.splitext(path)[1].lower()
if ext not in MODEL_FILE_EXTENSIONS:
continue
size = entry.get("size", 0) or 0
if size == 0 and "lfs" in entry:
size = entry["lfs"].get("size", 0) or 0
files.append({
"filename": path,
"size": size,
})
files.sort(key=lambda f: f["size"], reverse=True)
return web.json_response(files) return web.json_response(files)
async def download_hf_model(self, request: web.Request) -> web.Response: async def download_model_source(self, request: web.Request) -> web.Response:
"""Download a single file from Hugging Face into the model directory. """Download a single file from an external repository.
POST JSON body:: POST JSON body::
{ {
"repo": "dx8152/Flux2-Klein-9B-Consistency", "platform": "modelscope",
"filename": "Flux2-Klein-9B-consistency-V2.safetensors", "repo": "owner/name",
"revision": "main", "filename": "subdir/model.safetensors",
"revision": "master",
"model_root": "loras", "model_root": "loras",
"relative_path": "", "relative_path": "",
"use_default_paths": false, "use_default_paths": false,
"download_id": "optional-batch-id" "download_id": "optional-batch-id"
} }
``platform`` defaults to ``huggingface`` when omitted, which keeps the
legacy ``/api/lm/download-hf-model`` payload working unchanged.
If ``download_id`` is provided, real-time progress (bytes, speed, If ``download_id`` is provided, real-time progress (bytes, speed,
percentage) is broadcast via the WebSocket progress system, matching percentage) is broadcast via the WebSocket progress system.
the CivitAI download experience.
Respects the ``download_backend`` setting (``aria2`` or ``default``). Respects the ``download_backend`` setting (``aria2`` or ``default``).
""" """
@@ -413,30 +401,33 @@ class HfHandler:
except json.JSONDecodeError: except json.JSONDecodeError:
return web.json_response({"error": "Invalid JSON"}, status=400) return web.json_response({"error": "Invalid JSON"}, status=400)
platform = (payload.get("platform") or "huggingface").strip()
repo = (payload.get("repo") or "").strip() repo = (payload.get("repo") or "").strip()
filename = (payload.get("filename") or "").strip() filename = (payload.get("filename") or "").strip()
revision = (payload.get("revision") or "main").strip() revision = (payload.get("revision") or "").strip()
model_root = (payload.get("model_root") or "").strip() model_root = (payload.get("model_root") or "").strip()
relative_path = (payload.get("relative_path") or "").strip() relative_path = (payload.get("relative_path") or "").strip()
use_default_paths = bool(payload.get("use_default_paths", False)) use_default_paths = bool(payload.get("use_default_paths", False))
download_id: str | None = payload.get("download_id") download_id: str | None = payload.get("download_id")
logger.info( logger.info(
"download_hf_model: repo=%s file=%s root=%s download_id=%s", "download_model_source: platform=%s repo=%s file=%s root=%s download_id=%s",
repo, filename, model_root, download_id, platform, repo, filename, model_root, download_id,
) )
source = get_download_source(platform)
if source is None:
return _unsupported_platform_error(platform)
if not repo or not filename: if not repo or not filename:
return web.json_response( return web.json_response(
{"error": "Missing required fields: 'repo' and 'filename'"}, status=400 {"error": "Missing required fields: 'repo' and 'filename'"}, status=400
) )
# Validate repo format — must be user/repo_name # `owner/name` only; the components become path segments below.
if repo.count("/") != 1 or not re.match(r"^[a-zA-Z0-9_.-]+/[a-zA-Z0-9_.-]+$", repo): if not is_valid_source_id(repo):
return web.json_response({"error": f"Invalid repo format: {repo}"}, status=400)
author, repo_name = repo.split("/", 1)
if ".." in (author, repo_name) or "." in (author, repo_name):
return web.json_response({"error": f"Invalid repo format: {repo}"}, status=400) return web.json_response({"error": f"Invalid repo format: {repo}"}, status=400)
owner, repo_name = repo.split("/", 1)
# Validate filename — must not contain path traversal # Validate filename — must not contain path traversal
if ".." in filename: if ".." in filename:
@@ -455,21 +446,21 @@ class HfHandler:
# unnecessary when the frontend sends the path from its own dropdown # unnecessary when the frontend sends the path from its own dropdown
# (populated from scanner roots). Using the "business path" directly # (populated from scanner roots). Using the "business path" directly
# keeps dest_path consistent with scanner roots so that later folder # keeps dest_path consistent with scanner roots so that later folder
# derivation (in _save_hf_metadata) works correctly. # derivation (in _save_source_metadata) works correctly.
if os.path.isabs(model_root): if os.path.isabs(model_root):
base_dir = os.path.normpath(model_root) base_dir = os.path.normpath(model_root)
else: else:
base_dir = os.path.normpath(os.path.join(os.getcwd(), "models", model_root)) base_dir = os.path.normpath(os.path.join(os.getcwd(), "models", model_root))
if use_default_paths: if use_default_paths:
target_dir = os.path.join(base_dir, "huggingface", author, repo_name) target_dir = os.path.join(base_dir, source.default_subdir, owner, repo_name)
elif relative_path: elif relative_path:
target_dir = os.path.join(base_dir, relative_path) target_dir = os.path.join(base_dir, relative_path)
else: else:
target_dir = base_dir target_dir = base_dir
# Strip HF repo subdirectory — "diffusion_models/xxx.safetensors" # Strip the repository sub-directory — "diffusion_models/xxx.safetensors"
# is an HF repo convention, not meaningful for local storage. # is a repository convention, not meaningful for local storage.
file_base = os.path.basename(filename) file_base = os.path.basename(filename)
os.makedirs(target_dir, exist_ok=True) os.makedirs(target_dir, exist_ok=True)
@@ -477,16 +468,18 @@ class HfHandler:
# Check if already exists (simple skip) # Check if already exists (simple skip)
if os.path.exists(dest_path) and os.path.getsize(dest_path) > 0: if os.path.exists(dest_path) and os.path.getsize(dest_path) > 0:
logger.info("download_hf_model: file already exists, skipping — %s", dest_path) logger.info("download_model_source: file already exists, skipping — %s", dest_path)
return web.json_response({ return web.json_response({
"success": True, "success": True,
"message": f"File already exists: {dest_path}", "message": f"File already exists: {dest_path}",
"path": dest_path, "path": dest_path,
}) })
# Build HF resolve URL # Built per request: sites that redirect to a CDN hand out a
resolve_url = ( # time-limited token in the redirect, so the URL must never be cached.
f"https://huggingface.co/{repo}/resolve/{revision}/{filename}" resolve_url = source.file_download_url(repo, filename, revision)
ref = SourceRef(
platform=source.platform, source_id=repo, url=source.canonical_url(repo)
) )
# Set up progress callback if download_id is provided # Set up progress callback if download_id is provided
@@ -528,28 +521,27 @@ class HfHandler:
if download_backend == "aria2": if download_backend == "aria2":
aria2 = await Aria2Downloader.get_instance() aria2 = await Aria2Downloader.get_instance()
aid = download_id or f"hf_{repo}_{filename}" aid = download_id or f"{source.platform}_{repo}_{filename}"
try: try:
hf_success, hf_result = await aria2.download_file( ok, result = await aria2.download_file(
url=resolve_url, url=resolve_url,
save_path=dest_path, save_path=dest_path,
download_id=aid, download_id=aid,
progress_callback=progress_callback, progress_callback=progress_callback,
) )
if hf_success: if ok:
await _save_hf_metadata(dest_path, repo, model_root) await _save_source_metadata(dest_path, ref, model_root)
return web.json_response({ return web.json_response({
"success": True, "success": True,
"message": f"Downloaded to {dest_path}", "message": f"Downloaded to {dest_path}",
"path": dest_path, "path": dest_path,
}) })
else: return web.json_response(
return web.json_response( {"success": False, "error": result or "aria2 download failed"},
{"success": False, "error": hf_result or "aria2 download failed"}, status=500,
status=500, )
)
except Exception as exc: except Exception as exc:
logger.error("HF download (aria2) failed: %s", exc) logger.error("%s download (aria2) failed: %s", platform, exc)
return web.json_response( return web.json_response(
{"success": False, "error": str(exc)}, status=500 {"success": False, "error": str(exc)}, status=500
) )
@@ -565,19 +557,18 @@ class HfHandler:
progress_callback=progress_callback, progress_callback=progress_callback,
) )
if success: if success:
await _save_hf_metadata(dest_path, repo, model_root) await _save_source_metadata(dest_path, ref, model_root)
return web.json_response({ return web.json_response({
"success": True, "success": True,
"message": f"Downloaded to {result}", "message": f"Downloaded to {result}",
"path": result, "path": result,
}) })
else: return web.json_response(
return web.json_response( {"success": False, "error": result or "Download failed"},
{"success": False, "error": result or "Download failed"}, status=500,
status=500, )
)
except Exception as exc: except Exception as exc:
logger.error("HF download failed: %s", exc) logger.error("%s download failed: %s", platform, exc)
return web.json_response( return web.json_response(
{"success": False, "error": str(exc)}, status=500 {"success": False, "error": str(exc)}, status=500
) )
+8 -1
View File
@@ -99,7 +99,11 @@ MISC_ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
RouteDefinition( RouteDefinition(
"GET", "/api/lm/delete-model-version", "delete_model_version" "GET", "/api/lm/delete-model-version", "delete_model_version"
), ),
# Hugging Face model endpoints # External model source endpoints (Hugging Face / ModelScope).
# The hf-* paths are the historical names, kept as aliases.
RouteDefinition(
"GET", "/api/lm/model-source-files", "list_model_source_files"
),
RouteDefinition( RouteDefinition(
"GET", "/api/lm/hf-repo-files", "get_hf_repo_files" "GET", "/api/lm/hf-repo-files", "get_hf_repo_files"
), ),
@@ -107,6 +111,9 @@ MISC_ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
RouteDefinition( RouteDefinition(
"POST", "/api/lm/download/routing", "get_download_routing" "POST", "/api/lm/download/routing", "get_download_routing"
), ),
RouteDefinition(
"POST", "/api/lm/download-model-source", "download_model_source"
),
RouteDefinition( RouteDefinition(
"POST", "/api/lm/download-hf-model", "download_hf_model" "POST", "/api/lm/download-hf-model", "download_hf_model"
), ),
+3 -3
View File
@@ -39,7 +39,7 @@ from .handlers.misc_handlers import (
build_service_registry_adapter, build_service_registry_adapter,
) )
from .handlers.base_model_handlers import BaseModelHandlerSet from .handlers.base_model_handlers import BaseModelHandlerSet
from .handlers.hf_handlers import HfHandler from .handlers.model_source_handlers import ModelSourceHandler
from .handlers.agent_handlers import AgentHandler from .handlers.agent_handlers import AgentHandler
from .handlers.download_routing_handlers import DownloadRoutingHandler from .handlers.download_routing_handlers import DownloadRoutingHandler
from .misc_route_registrar import MiscRouteRegistrar from .misc_route_registrar import MiscRouteRegistrar
@@ -139,7 +139,7 @@ class MiscRoutes:
doctor = DoctorHandler(settings_service=self._settings) doctor = DoctorHandler(settings_service=self._settings)
example_workflows = ExampleWorkflowsHandler() example_workflows = ExampleWorkflowsHandler()
base_model = BaseModelHandlerSet() base_model = BaseModelHandlerSet()
hf_handler = HfHandler() model_source_handler = ModelSourceHandler()
agent_handler = AgentHandler() agent_handler = AgentHandler()
download_routing = DownloadRoutingHandler() download_routing = DownloadRoutingHandler()
@@ -161,7 +161,7 @@ class MiscRoutes:
doctor=doctor, doctor=doctor,
example_workflows=example_workflows, example_workflows=example_workflows,
base_model=base_model, base_model=base_model,
hf_handler=hf_handler, model_source_handler=model_source_handler,
agent_handler=agent_handler, agent_handler=agent_handler,
download_routing=download_routing, download_routing=download_routing,
) )
+12
View File
@@ -12,10 +12,14 @@ from .base import (
GROUP_PREFIXES, GROUP_PREFIXES,
HTTP_TIMEOUT, HTTP_TIMEOUT,
ModelSource, ModelSource,
ModelSourceError,
SourceRef, SourceRef,
USER_AGENT, USER_AGENT,
clean_source_url, clean_source_url,
fetch_json,
fetch_text, fetch_text,
filter_weight_files,
is_valid_source_id,
) )
from .huggingface import HuggingFaceSource from .huggingface import HuggingFaceSource
from .modelscope import ModelScopeSource from .modelscope import ModelScopeSource
@@ -24,6 +28,8 @@ from .registry import (
SOURCE_PLATFORM_FIELD, SOURCE_PLATFORM_FIELD,
SOURCE_URL_FIELD, SOURCE_URL_FIELD,
detect_source, detect_source,
downloadable_sources,
get_download_source,
get_source, get_source,
get_source_platform, get_source_platform,
has_external_source, has_external_source,
@@ -40,6 +46,7 @@ __all__ = [
"HTTP_TIMEOUT", "HTTP_TIMEOUT",
"LEGACY_HF_URL_FIELD", "LEGACY_HF_URL_FIELD",
"ModelSource", "ModelSource",
"ModelSourceError",
"HuggingFaceSource", "HuggingFaceSource",
"ModelScopeSource", "ModelScopeSource",
"SOURCE_PLATFORM_FIELD", "SOURCE_PLATFORM_FIELD",
@@ -49,10 +56,15 @@ __all__ = [
"USER_AGENT", "USER_AGENT",
"clean_source_url", "clean_source_url",
"detect_source", "detect_source",
"downloadable_sources",
"fetch_json",
"fetch_text", "fetch_text",
"filter_weight_files",
"get_download_source",
"get_source", "get_source",
"get_source_platform", "get_source_platform",
"has_external_source", "has_external_source",
"is_valid_source_id",
"list_sources", "list_sources",
"normalize_metadata_source", "normalize_metadata_source",
"resolve_source_ref", "resolve_source_ref",
+128 -1
View File
@@ -20,12 +20,15 @@ HTTP handlers never need site-specific branching.
from __future__ import annotations from __future__ import annotations
import logging import logging
import os
import re import re
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any, Optional from typing import Any, Iterable, Optional
import aiohttp import aiohttp
from ...utils.constants import MODEL_FILE_EXTENSIONS
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
#: Shared HTTP timeout for model-card fetches. #: Shared HTTP timeout for model-card fetches.
@@ -58,6 +61,36 @@ class SourceRef:
"""Canonical URL of the model page.""" """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: async def fetch_text(url: str, *, timeout: int = HTTP_TIMEOUT) -> str:
"""Fetch *url* and return its body as text, or ``""`` on any failure. """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 "" 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: class ModelSource:
"""Description and I/O for one external model hosting site.""" """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. #: Whether models can be downloaded directly from this site.
supports_download: bool = False 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. #: Lenient pattern used to recognise URLs already stored in metadata.
#: Captures the site-specific source id in group ``id``. #: Captures the site-specific source id in group ``id``.
url_pattern: re.Pattern[str] | None = None url_pattern: re.Pattern[str] | None = None
@@ -163,6 +230,45 @@ class ModelSource:
return "" 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 def __repr__(self) -> str: # pragma: no cover - debugging aid
return f"<ModelSource {self.platform}>" return f"<ModelSource {self.platform}>"
@@ -175,12 +281,33 @@ def clean_source_url(url: Any) -> str:
return url.strip() 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__ = [ __all__ = [
"GROUP_PREFIXES", "GROUP_PREFIXES",
"HTTP_TIMEOUT", "HTTP_TIMEOUT",
"ModelSource", "ModelSource",
"ModelSourceError",
"SourceRef", "SourceRef",
"USER_AGENT", "USER_AGENT",
"clean_source_url", "clean_source_url",
"fetch_json",
"fetch_text", "fetch_text",
"filter_weight_files",
"is_valid_source_id",
] ]
+59 -2
View File
@@ -2,9 +2,18 @@
from __future__ import annotations from __future__ import annotations
import logging
import re import re
from .base import ModelSource, fetch_text from .base import (
ModelSource,
ModelSourceError,
fetch_json,
fetch_text,
filter_weight_files,
)
logger = logging.getLogger(__name__)
#: Lenient — used to normalise URLs already stored in metadata; tolerates #: Lenient — used to normalise URLs already stored in metadata; tolerates
#: sub-paths such as ``/resolve/main/model.safetensors``. #: sub-paths such as ``/resolve/main/model.safetensors``.
@@ -25,6 +34,8 @@ class HuggingFaceSource(ModelSource):
label = "Hugging Face" label = "Hugging Face"
supports_enrichment = True supports_enrichment = True
supports_download = True supports_download = True
default_revision = "main"
default_subdir = "huggingface"
url_pattern = _URL_PATTERN url_pattern = _URL_PATTERN
strict_url_pattern = _STRICT_URL_PATTERN strict_url_pattern = _STRICT_URL_PATTERN
@@ -32,7 +43,7 @@ class HuggingFaceSource(ModelSource):
return f"https://huggingface.co/{source_id}" return f"https://huggingface.co/{source_id}"
def asset_base_url(self, source_id: str, revision: str = "") -> str: def asset_base_url(self, source_id: str, revision: str = "") -> str:
return f"https://huggingface.co/{source_id}/resolve/{revision or 'main'}" return f"https://huggingface.co/{source_id}/resolve/{self.resolve_revision(revision)}"
async def fetch_model_card(self, source_id: str) -> str: async def fetch_model_card(self, source_id: str) -> str:
"""Fetch ``README.md`` from Hugging Face (tries ``main``, then ``master``).""" """Fetch ``README.md`` from Hugging Face (tries ``main``, then ``master``)."""
@@ -45,5 +56,51 @@ class HuggingFaceSource(ModelSource):
return text return text
return "" return ""
async def list_files(
self, source_id: str, revision: str = ""
) -> list[dict]:
"""List weight files via the Hub tree API.
The tree endpoint (rather than the model-info endpoint) is used
because it reports accurate sizes for LFS-tracked files.
"""
revision = self.resolve_revision(revision)
status, payload = await fetch_json(
f"https://huggingface.co/api/models/{source_id}/tree/{revision}"
)
if status == 404:
raise ModelSourceError(f"Repository '{source_id}' not found", status=404)
if status != 200 or not isinstance(payload, list):
raise ModelSourceError(
f"Hugging Face API error while listing '{source_id}' (HTTP {status})"
)
entries = []
for entry in payload:
if not isinstance(entry, dict):
continue
path = entry.get("path", "")
size = entry.get("size", 0) or 0
if not size and isinstance(entry.get("lfs"), dict):
size = entry["lfs"].get("size", 0) or 0
entries.append((path, size))
return filter_weight_files(entries)
def file_download_url(
self, source_id: str, filename: str, revision: str = ""
) -> str:
return (
f"https://huggingface.co/{source_id}/resolve/"
f"{self.resolve_revision(revision)}/{filename}"
)
def page_url_for_file(self, source_id: str, filename: str) -> str:
return (
f"https://huggingface.co/{source_id}/blob/{self.default_revision}/{filename}"
)
__all__ = ["HuggingFaceSource"] __all__ = ["HuggingFaceSource"]
+73 -5
View File
@@ -2,20 +2,38 @@
ModelScope exposes the same "model card as README.md" convention as ModelScope exposes the same "model card as README.md" convention as
Hugging Face, including a YAML frontmatter block that often carries Hugging Face, including a YAML frontmatter block that often carries
``base_model:`` and ``trigger_words:``. Two public endpoints are used, ``base_model:`` and ``trigger_words:``. Three public endpoints are used,
neither of which requires an API key for public models: none of which requires an API key for public models:
* ``/models/{owner}/{name}/resolve/{revision}/README.md`` raw model card * ``/models/{owner}/{name}/resolve/{revision}/README.md`` raw model card
* ``/api/v1/models/{owner}/{name}/repo?Revision=..&FilePath=README.md`` * ``/api/v1/models/{owner}/{name}/repo?Revision=..&FilePath=README.md``
the same content through the API, used as a fallback when the resolve the same content through the API, used as a fallback when the resolve
URL is unavailable. URL is unavailable.
* ``/api/v1/models/{owner}/{name}/repo/files?Revision=..`` the file
listing backing the download picker. It reports real sizes for LFS
files (not the pointer size), so no extra HEAD request is needed.
Downloads go through ``/models/{owner}/{name}/resolve/{revision}/{path}``,
which redirects to a CDN URL carrying a time-limited ``auth_key``.
Requesting the resolve URL fresh on every attempt (which the shared
downloader does, including for resumable Range requests) keeps that key
valid; the CDN URL must never be cached.
""" """
from __future__ import annotations from __future__ import annotations
import logging
import re import re
from .base import ModelSource, fetch_text from .base import (
ModelSource,
ModelSourceError,
fetch_json,
fetch_text,
filter_weight_files,
)
logger = logging.getLogger(__name__)
_URL_PATTERN = re.compile( _URL_PATTERN = re.compile(
r"https?://(?:www\.)?modelscope\.(?:cn|com)/models/(?P<id>[^/?#\s]+/[^/?#\s]+)" r"https?://(?:www\.)?modelscope\.(?:cn|com)/models/(?P<id>[^/?#\s]+/[^/?#\s]+)"
@@ -41,7 +59,9 @@ class ModelScopeSource(ModelSource):
platform = "modelscope" platform = "modelscope"
label = "ModelScope" label = "ModelScope"
supports_enrichment = True supports_enrichment = True
supports_download = False supports_download = True
default_revision = "master"
default_subdir = "modelscope"
url_pattern = _URL_PATTERN url_pattern = _URL_PATTERN
strict_url_pattern = _STRICT_URL_PATTERN strict_url_pattern = _STRICT_URL_PATTERN
@@ -49,7 +69,10 @@ class ModelScopeSource(ModelSource):
return f"https://modelscope.cn/models/{source_id}" return f"https://modelscope.cn/models/{source_id}"
def asset_base_url(self, source_id: str, revision: str = "") -> str: def asset_base_url(self, source_id: str, revision: str = "") -> str:
return f"https://modelscope.cn/models/{source_id}/resolve/{revision or 'master'}" return (
f"https://modelscope.cn/models/{source_id}/resolve/"
f"{self.resolve_revision(revision)}"
)
async def fetch_model_card(self, source_id: str) -> str: async def fetch_model_card(self, source_id: str) -> str:
"""Fetch the model card, preferring the raw resolve URL.""" """Fetch the model card, preferring the raw resolve URL."""
@@ -72,5 +95,50 @@ class ModelScopeSource(ModelSource):
return text return text
return "" return ""
async def list_files(
self, source_id: str, revision: str = ""
) -> list[dict]:
"""List weight files via the repo files API.
``master`` is the only branch name the API accepts even repos
imported from Hugging Face are addressed as ``master`` (``main``
returns 404) so no fallback probing is done here.
"""
revision = self.resolve_revision(revision)
status, payload = await fetch_json(
"https://modelscope.cn/api/v1/models/"
f"{source_id}/repo/files?Revision={revision}"
)
if status == 404:
raise ModelSourceError(f"Repository '{source_id}' not found", status=404)
if status != 200 or not isinstance(payload, dict):
raise ModelSourceError(
f"ModelScope API error while listing '{source_id}' (HTTP {status})"
)
entries = []
for entry in (payload.get("Data") or {}).get("Files") or []:
if not isinstance(entry, dict) or entry.get("Type") != "blob":
continue
entries.append((entry.get("Path", ""), entry.get("Size", 0) or 0))
return filter_weight_files(entries)
def file_download_url(
self, source_id: str, filename: str, revision: str = ""
) -> str:
return (
f"https://modelscope.cn/models/{source_id}/resolve/"
f"{self.resolve_revision(revision)}/{filename}"
)
def page_url_for_file(self, source_id: str, filename: str) -> str:
return (
f"https://modelscope.cn/models/{source_id}/file/view/"
f"{self.default_revision}/{filename}"
)
__all__ = ["ModelScopeSource"] __all__ = ["ModelScopeSource"]
+17
View File
@@ -56,6 +56,21 @@ def source_label(platform: Optional[str], default: str = "") -> str:
return source.label if source else default return source.label if source else default
def downloadable_sources() -> list[ModelSource]:
"""Return the sources whose repositories can be downloaded directly."""
return [source for source in _SOURCES if source.supports_download]
def get_download_source(platform: Optional[str]) -> Optional[ModelSource]:
"""Return the source for *platform*, but only when it supports downloads."""
source = get_source(platform)
if source is None or not source.supports_download:
return None
return source
def detect_source(url: Optional[str], *, strict: bool = False) -> Optional[SourceRef]: def detect_source(url: Optional[str], *, strict: bool = False) -> Optional[SourceRef]:
"""Return the :class:`SourceRef` for *url*, or ``None`` if unsupported.""" """Return the :class:`SourceRef` for *url*, or ``None`` if unsupported."""
@@ -197,6 +212,8 @@ __all__ = [
"SOURCE_PLATFORM_FIELD", "SOURCE_PLATFORM_FIELD",
"SOURCE_URL_FIELD", "SOURCE_URL_FIELD",
"detect_source", "detect_source",
"downloadable_sources",
"get_download_source",
"get_source", "get_source",
"get_source_platform", "get_source_platform",
"has_external_source", "has_external_source",
+11 -3
View File
@@ -203,10 +203,18 @@ export const DOWNLOAD_ENDPOINTS = {
exampleImagesMissing: '/api/lm/download-example-images' // Download only missing example images exampleImagesMissing: '/api/lm/download-example-images' // Download only missing example images
}; };
// Hugging Face API endpoints // External model source endpoints (Hugging Face / ModelScope).
// The hf-* paths are the historical names, kept as server-side aliases.
export const MODEL_SOURCE_ENDPOINTS = {
repoFiles: '/api/lm/model-source-files',
download: '/api/lm/download-model-source',
sources: '/api/lm/model-sources',
};
/** @deprecated use MODEL_SOURCE_ENDPOINTS */
export const HF_ENDPOINTS = { export const HF_ENDPOINTS = {
repoFiles: '/api/lm/hf-repo-files', repoFiles: MODEL_SOURCE_ENDPOINTS.repoFiles,
download: '/api/lm/download-hf-model', download: MODEL_SOURCE_ENDPOINTS.download,
}; };
// WebSocket endpoints // WebSocket endpoints
+46 -9
View File
@@ -8,6 +8,7 @@ import {
isValidModelType, isValidModelType,
DOWNLOAD_ENDPOINTS, DOWNLOAD_ENDPOINTS,
HF_ENDPOINTS, HF_ENDPOINTS,
MODEL_SOURCE_ENDPOINTS,
WS_ENDPOINTS WS_ENDPOINTS
} from './apiConfig.js'; } from './apiConfig.js';
import { resetAndReload } from './modelApiFactory.js'; import { resetAndReload } from './modelApiFactory.js';
@@ -1367,30 +1368,52 @@ export class BaseModelApiClient {
} }
} }
async fetchHfRepoFiles(repo, revision = 'main') { /**
* List the downloadable weight files of an external repository.
* @param {string} repo - `owner/name`
* @param {string} [platform] - `huggingface` (default) or `modelscope`
* @param {string} [revision] - branch; each site has its own default
*/
async fetchModelSourceFiles(repo, platform = 'huggingface', revision = '') {
try { try {
const params = new URLSearchParams({ repo, revision }); const params = new URLSearchParams({ repo, platform });
const response = await fetch(`${HF_ENDPOINTS.repoFiles}?${params}`); if (revision) params.set('revision', revision);
const response = await fetch(`${MODEL_SOURCE_ENDPOINTS.repoFiles}?${params}`);
if (!response.ok) { if (!response.ok) {
const err = await response.json().catch(() => ({})); const err = await response.json().catch(() => ({}));
throw new Error(err.error || 'Failed to fetch HF repo files'); throw new Error(err.error || 'Failed to fetch repository files');
} }
return await response.json(); return await response.json();
} catch (error) { } catch (error) {
console.error('Error fetching HF repo files:', error); console.error('Error fetching repository files:', error);
throw error; throw error;
} }
} }
async downloadHfModel({ repo, filename, revision, modelRoot, relativePath, useDefaultPaths, download_id }) { /** Backwards-compatible Hugging Face wrapper. */
async fetchHfRepoFiles(repo, revision = 'main') {
return this.fetchModelSourceFiles(repo, 'huggingface', revision);
}
async downloadModelSource({
platform = 'huggingface',
repo,
filename,
revision,
modelRoot,
relativePath,
useDefaultPaths,
download_id,
}) {
try { try {
const response = await fetch(HF_ENDPOINTS.download, { const response = await fetch(MODEL_SOURCE_ENDPOINTS.download, {
method: 'POST', method: 'POST',
headers: { 'Content-Type': 'application/json' }, headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ body: JSON.stringify({
platform,
repo, repo,
filename, filename,
revision: revision || 'main', revision: revision || '',
model_root: modelRoot, model_root: modelRoot,
relative_path: relativePath || '', relative_path: relativePath || '',
use_default_paths: useDefaultPaths || false, use_default_paths: useDefaultPaths || false,
@@ -1404,11 +1427,25 @@ export class BaseModelApiClient {
return await response.json(); return await response.json();
} catch (error) { } catch (error) {
console.error('Error downloading HF model:', error); console.error('Error downloading model:', error);
throw error; throw error;
} }
} }
/** Backwards-compatible Hugging Face wrapper. */
async downloadHfModel({ repo, filename, revision, modelRoot, relativePath, useDefaultPaths, download_id }) {
return this.downloadModelSource({
platform: 'huggingface',
repo,
filename,
revision: revision || 'main',
modelRoot,
relativePath,
useDefaultPaths,
download_id,
});
}
_buildQueryParams(baseParams, pageState) { _buildQueryParams(baseParams, pageState) {
const params = new URLSearchParams(baseParams); const params = new URLSearchParams(baseParams);
const isExcludedView = pageState.viewMode === 'excluded'; const isExcludedView = pageState.viewMode === 'excluded';
+182 -169
View File
@@ -13,6 +13,13 @@ import { buildCivitaiUrl, extractCivitaiModelUrlParts, normalizeCivitaiPageHost
import { formatFileSize } from '../utils/formatters.js'; import { formatFileSize } from '../utils/formatters.js';
import { showDownloadBatchSummary } from '../components/DownloadBatchSummaryModal.js'; import { showDownloadBatchSummary } from '../components/DownloadBatchSummaryModal.js';
import { openOtherModelsSettings } from '../utils/otherModels.js'; import { openOtherModelsSettings } from '../utils/otherModels.js';
import {
buildModelSourceFilePage,
detectModelSourceDownloadUrl,
getModelSource,
isExternalModelSource,
isValidRepoId,
} from '../utils/modelSourceHelpers.js';
export class DownloadManager { export class DownloadManager {
constructor() { constructor() {
@@ -39,10 +46,11 @@ export class DownloadManager {
this.isBatchMode = false; this.isBatchMode = false;
this.editingBatchIndex = -1; this.editingBatchIndex = -1;
// HF download state // External repository download state (Hugging Face / ModelScope)
this.hfRepoId = null; this.sourcePlatform = 'huggingface';
this.hfSelectedFiles = []; this.sourceRepoId = null;
this.hfRepoCollapsed = {}; this.sourceSelectedFiles = [];
this.sourceRepoCollapsed = {};
this.loadingManager = new LoadingManager(); this.loadingManager = new LoadingManager();
this.folderTreeManager = new FolderTreeManager(); this.folderTreeManager = new FolderTreeManager();
@@ -186,10 +194,11 @@ export class DownloadManager {
// Reset default path toggle // Reset default path toggle
this.loadDefaultPathSetting(); this.loadDefaultPathSetting();
// Reset HF state // Reset external repository state
this.hfRepoId = null; this.sourcePlatform = 'huggingface';
this.hfSelectedFiles = []; this.sourceRepoId = null;
this.hfRepoCollapsed = {}; this.sourceSelectedFiles = [];
this.sourceRepoCollapsed = {};
} }
async retrieveVersionsForModel(modelId, source = null) { async retrieveVersionsForModel(modelId, source = null) {
@@ -212,10 +221,12 @@ export class DownloadManager {
// Detect URL types — all URLs must share the same source type // Detect URL types — all URLs must share the same source type
const urlTypes = urls.map(u => DownloadManager.detectUrlType(u)); const urlTypes = urls.map(u => DownloadManager.detectUrlType(u));
const isHf = urlTypes.every(t => t && (t.type === 'hf-resolve' || t.type === 'hf-repo')); const isExternalSource = urlTypes.every(
t => t && (t.type === 'model-source-repo' || t.type === 'model-source-file')
);
const isCivitai = urlTypes.every(t => t && t.type === 'civitai'); const isCivitai = urlTypes.every(t => t && t.type === 'civitai');
if (!isHf && !isCivitai) { if (!isExternalSource && !isCivitai) {
const allValid = urlTypes.every(t => t !== null); const allValid = urlTypes.every(t => t !== null);
if (!allValid) { if (!allValid) {
errorElement.textContent = translate('modals.download.errors.invalidUrl'); errorElement.textContent = translate('modals.download.errors.invalidUrl');
@@ -228,8 +239,8 @@ export class DownloadManager {
} }
} }
if (isHf) { if (isExternalSource) {
return this._validateAndFetchHf(urls, errorElement); return this._validateAndFetchExternalRepo(urls, errorElement);
} }
// --- Original CivitAI flow below --- // --- Original CivitAI flow below ---
@@ -327,45 +338,66 @@ export class DownloadManager {
this.showBatchPreviewStep(); this.showBatchPreviewStep();
} }
// ---- Hugging Face download flow ---- // ---- External repository download flow (Hugging Face / ModelScope) ----
async _validateAndFetchHf(urls, errorElement) { /** Rendering group key: the same repo on two sites is two groups. */
_externalGroupKey(item) {
return `${item.source}:${item.repo || 'unknown'}`;
}
_defaultRevisionFor(platform) {
const source = getModelSource(platform);
return (source && source.defaultRevision) || '';
}
_makeExternalItem(url, info, file) {
return {
url,
source: info.platform,
platform: info.platform,
repo: info.repo,
revision: file.revision || this._defaultRevisionFor(info.platform),
filename: file.filename,
displayName: file.filename,
fileSizeBytes: file.size,
selectedVersion: true,
versions: [],
checked: false,
error: null,
};
}
/** Fetch a repository's weight files as flat batch items. */
async _fetchExternalRepoItems(url, info) {
const revision = this._defaultRevisionFor(info.platform);
const files = await this.apiClient.fetchModelSourceFiles(
info.repo, info.platform, revision
);
if (!files || files.length === 0) {
throw new Error(translate('modals.download.errors.noModelFiles'));
}
return files.map(file => this._makeExternalItem(url, info, { ...file, revision }));
}
async _validateAndFetchExternalRepo(urls, errorElement) {
if (urls.length === 1) { if (urls.length === 1) {
const info = DownloadManager.detectUrlType(urls[0]); const info = DownloadManager.detectUrlType(urls[0]);
// Direct file resolve URL → skip file selection, go to location // Direct file URL → skip file selection, go to location
if (info.type === 'hf-resolve') { if (info.type === 'model-source-file') {
this.isBatchMode = false; this.isBatchMode = false;
this.hfRepoId = info.repo; this.sourcePlatform = info.platform;
this.hfSelectedFiles = [info.filename]; this.sourceRepoId = info.repo;
this.source = 'huggingface'; this.sourceSelectedFiles = [info.filename];
this.source = info.platform;
this.proceedToLocation(); this.proceedToLocation();
return; return;
} }
// Repo URL → fetch file list and convert to batch items // Repo URL → fetch file list and convert to batch items
try { try {
this.loadingManager.showSimpleLoading(translate('modals.download.fetchingRepoFiles')); this.loadingManager.showSimpleLoading(translate('modals.download.fetchingRepoFiles'));
const files = await this.apiClient.fetchHfRepoFiles(info.repo);
if (!files || files.length === 0) {
throw new Error(translate('modals.download.errors.noModelFiles'));
}
this.isBatchMode = true; this.isBatchMode = true;
this.batchModels = []; this.batchModels = await this._fetchExternalRepoItems(urls[0], info);
this.source = 'huggingface'; this.source = info.platform;
for (const file of files) {
this.batchModels.push({
url: urls[0],
source: 'huggingface',
repo: info.repo,
filename: file.filename,
revision: 'main',
displayName: file.filename,
fileSizeBytes: file.size,
selectedVersion: true,
versions: [],
checked: false,
error: null,
});
}
this.showBatchPreviewStep(); this.showBatchPreviewStep();
} catch (err) { } catch (err) {
errorElement.textContent = err.message; errorElement.textContent = err.message;
@@ -375,10 +407,9 @@ export class DownloadManager {
return; return;
} }
// Multiple HF URLs → batch mode: flatten all files from all repos // Multiple URLs → batch mode: flatten all files from all repos
this.isBatchMode = true; this.isBatchMode = true;
this.batchModels = []; this.batchModels = [];
this.source = 'huggingface';
this.loadingManager.showSimpleLoading(translate('modals.download.fetchingRepoFiles')); this.loadingManager.showSimpleLoading(translate('modals.download.fetchingRepoFiles'));
for (const url of urls) { for (const url of urls) {
@@ -387,42 +418,15 @@ export class DownloadManager {
this.batchModels.push({ url, error: 'Invalid URL', versions: [], selectedVersion: null }); this.batchModels.push({ url, error: 'Invalid URL', versions: [], selectedVersion: null });
continue; continue;
} }
if (info.type === 'hf-resolve') { this.source = info.platform;
this.batchModels.push({ if (info.type === 'model-source-file') {
url, this.batchModels.push(this._makeExternalItem(url, info, {
source: 'huggingface',
repo: info.repo,
filename: info.filename, filename: info.filename,
revision: info.revision || 'main', revision: info.revision,
displayName: info.filename, }));
selectedVersion: true, } else if (info.type === 'model-source-repo') {
versions: [],
checked: false,
error: null,
});
} else if (info.type === 'hf-repo') {
try { try {
const files = await this.apiClient.fetchHfRepoFiles(info.repo); this.batchModels.push(...await this._fetchExternalRepoItems(url, info));
if (!files || files.length === 0) {
this.batchModels.push({ url, error: 'No model files found', versions: [], selectedVersion: null });
continue;
}
// Flatten: create one batch item per file, all checked by default
for (const file of files) {
this.batchModels.push({
url,
source: 'huggingface',
repo: info.repo,
filename: file.filename,
revision: 'main',
displayName: file.filename,
fileSizeBytes: file.size,
selectedVersion: true,
versions: [],
checked: false,
error: null,
});
}
} catch (err) { } catch (err) {
this.batchModels.push({ url, error: err.message, versions: [], selectedVersion: null }); this.batchModels.push({ url, error: err.message, versions: [], selectedVersion: null });
} }
@@ -480,7 +484,8 @@ export class DownloadManager {
* Detect the source type of a download URL. * Detect the source type of a download URL.
* @param {string} url * @param {string} url
* @returns {{ type: string, repo?: string, filename?: string, revision?: string } | null} * @returns {{ type: string, repo?: string, filename?: string, revision?: string } | null}
* type: 'civitai' | 'civarchive' | 'hf-resolve' | 'hf-repo' | 'direct-http' * type: 'civitai' | 'civarchive' | 'model-source-file' | 'model-source-repo'
* | 'direct-http'
*/ */
static detectUrlType(url) { static detectUrlType(url) {
const trimmed = url.trim(); const trimmed = url.trim();
@@ -492,38 +497,27 @@ export class DownloadManager {
return { type: 'civitai' }; return { type: 'civitai' };
} }
// Hugging Face resolve/blob URL → direct file // External model sources (Hugging Face / ModelScope). Repository URLs
// "blob" is the web preview page; it maps 1:1 to the "resolve" download URL // list every weight file; resolve URLs point at one file. Both are
const hfResolveMatch = trimmed.match(/huggingface\.co\/([^/\s]+\/[^/\s]+)\/(?:resolve|blob)\/([^/\s]+)\/(.+)/i); // recognised through the shared registry, so adding a site is a
if (hfResolveMatch) { // registry change rather than a change here.
return { const sourceInfo = detectModelSourceDownloadUrl(trimmed);
type: 'hf-resolve', if (sourceInfo) {
repo: hfResolveMatch[1],
revision: hfResolveMatch[2],
filename: hfResolveMatch[3],
};
}
// Hugging Face repo URL (huggingface.co/user/repo or bare user/repo path)
// Require huggingface.co prefix for full URLs; bare user/repo only without ://
const hfRepoMatch = trimmed.match(
trimmed.includes('://')
? /^https?:\/\/huggingface\.co\/([a-zA-Z0-9_.-]+\/[a-zA-Z0-9_.-]+)(?:\/?$|$)/
: /^([a-zA-Z0-9_.-]+\/[a-zA-Z0-9_.-]+)$/
);
if (hfRepoMatch) {
// Reject path-traversal patterns like "../.." or "user/.." // Reject path-traversal patterns like "../.." or "user/.."
const parts = hfRepoMatch[1].split('/'); if (!isValidRepoId(sourceInfo.repo)) {
if (parts.some(p => p === '.' || p === '..')) {
return null; return null;
} }
return { return {
type: 'hf-repo', type: sourceInfo.kind === 'file' ? 'model-source-file' : 'model-source-repo',
repo: hfRepoMatch[1], platform: sourceInfo.platform,
repo: sourceInfo.repo,
...(sourceInfo.kind === 'file'
? { revision: sourceInfo.revision, filename: sourceInfo.filename }
: {}),
}; };
} }
// Direct HTTP(S) URL (non-HF) // Direct HTTP(S) URL (non model-source)
if (/^https?:\/\//i.test(trimmed)) { if (/^https?:\/\//i.test(trimmed)) {
return { type: 'direct-http' }; return { type: 'direct-http' };
} }
@@ -931,7 +925,7 @@ export class DownloadManager {
} }
// In single-URL mode, validate version selection (skip for HF) // In single-URL mode, validate version selection (skip for HF)
if (!this.isBatchMode && this.source !== 'huggingface') { if (!this.isBatchMode && !isExternalModelSource(this.source)) {
if (!this.currentVersion) { if (!this.currentVersion) {
showToast('toast.loras.pleaseSelectVersion', {}, 'error'); showToast('toast.loras.pleaseSelectVersion', {}, 'error');
return; return;
@@ -1164,12 +1158,15 @@ export class DownloadManager {
/** /**
* Synthesize a clickable URL for a single-download failure entry. * Synthesize a clickable URL for a single-download failure entry.
* Single downloads have no pasted URL, so the modal link is derived from * Single downloads have no pasted URL, so the modal link is derived from
* the model/version ids (CivitAI) or the HF repo/file (HuggingFace). * the model/version ids (CivitAI) or the external repo/file.
*/ */
_buildSingleItemUrl({ modelId, versionId, source, repo = null, filename = null }) { _buildSingleItemUrl({ modelId, versionId, source, repo = null, filename = null }) {
if (source === 'huggingface' && repo) { if (isExternalModelSource(source) && repo) {
const base = `https://huggingface.co/${encodeURI(repo)}`; return buildModelSourceFilePage({
return filename ? `${base}/blob/${encodeURI('main')}/${encodeURI(filename)}` : base; platform: source,
repo,
filename,
}) || getModelSource(source).canonical(repo);
} }
if (modelId) { if (modelId) {
return buildCivitaiUrl({ return buildCivitaiUrl({
@@ -1490,8 +1487,8 @@ export class DownloadManager {
* matched card-by-card via `_reconcileViewAfterDownload`; HF * matched card-by-card via `_reconcileViewAfterDownload`; HF
* downloads (no CivitAI identity to match) keep the legacy reload. * downloads (no CivitAI identity to match) keep the legacy reload.
*/ */
async _reconcileBatchViewAfterDownload(completedCivitaiItems = [], hfCompletedCount = 0) { async _reconcileBatchViewAfterDownload(completedCivitaiItems = [], externalCompletedCount = 0) {
if (hfCompletedCount > 0) { if (externalCompletedCount > 0) {
await resetAndReload(true); await resetAndReload(true);
return; return;
} }
@@ -1661,10 +1658,11 @@ export class DownloadManager {
return failedItems.length === 0; return failedItems.length === 0;
} }
async _downloadHfSingle({ modelRoot, targetFolder, useDefaultPaths, files = null }) { async _downloadExternalRepoFiles({ modelRoot, targetFolder, useDefaultPaths, files = null }) {
modalManager.closeModal('downloadModal'); modalManager.closeModal('downloadModal');
this.loadingManager.restoreProgressBar(); this.loadingManager.restoreProgressBar();
const filesToDownload = files || this.hfSelectedFiles; const platform = this.sourcePlatform;
const filesToDownload = files || this.sourceSelectedFiles;
const totalFiles = filesToDownload.length; const totalFiles = filesToDownload.length;
const updateProgress = this.loadingManager.showDownloadProgress(totalFiles); const updateProgress = this.loadingManager.showDownloadProgress(totalFiles);
@@ -1720,10 +1718,11 @@ export class DownloadManager {
} }
}; };
const response = await this.apiClient.downloadHfModel({ const response = await this.apiClient.downloadModelSource({
repo: this.hfRepoId, platform,
repo: this.sourceRepoId,
filename, filename,
revision: 'main', revision: this._defaultRevisionFor(platform),
modelRoot, modelRoot,
relativePath: targetFolder, relativePath: targetFolder,
useDefaultPaths, useDefaultPaths,
@@ -1738,10 +1737,10 @@ export class DownloadManager {
} else { } else {
failedFiles.push({ failedFiles.push({
item: { item: {
source: 'huggingface', source: platform,
repo: this.hfRepoId, repo: this.sourceRepoId,
filename, filename,
url: this._buildSingleItemUrl({ source: 'huggingface', repo: this.hfRepoId, filename }), url: this._buildSingleItemUrl({ source: platform, repo: this.sourceRepoId, filename }),
}, },
error: response?.error || 'Unknown error', error: response?.error || 'Unknown error',
name: filename, name: filename,
@@ -1749,13 +1748,13 @@ export class DownloadManager {
} }
} catch (err) { } catch (err) {
if (!cancelled) { if (!cancelled) {
console.error(`Failed to download HF file ${filename}:`, err); console.error(`Failed to download repo file ${filename}:`, err);
failedFiles.push({ failedFiles.push({
item: { item: {
source: 'huggingface', source: platform,
repo: this.hfRepoId, repo: this.sourceRepoId,
filename, filename,
url: this._buildSingleItemUrl({ source: 'huggingface', repo: this.hfRepoId, filename }), url: this._buildSingleItemUrl({ source: platform, repo: this.sourceRepoId, filename }),
}, },
error: err?.message || 'Unknown error', error: err?.message || 'Unknown error',
name: filename, name: filename,
@@ -1781,7 +1780,7 @@ export class DownloadManager {
total: totalFiles, total: totalFiles,
completed: completedDownloads, completed: completedDownloads,
failedItems: failedFiles, failedItems: failedFiles,
onRetry: () => this._downloadHfSingle({ onRetry: () => this._downloadExternalRepoFiles({
modelRoot, modelRoot,
targetFolder, targetFolder,
useDefaultPaths, useDefaultPaths,
@@ -1831,7 +1830,7 @@ export class DownloadManager {
const validCount = this.batchModels.filter(m => { const validCount = this.batchModels.filter(m => {
if (m.error) return false; if (m.error) return false;
if (m.source === 'huggingface') return m.checked !== false; if (isExternalModelSource(m.source)) return m.checked !== false;
return m.selectedVersion; return m.selectedVersion;
}).length; }).length;
document.getElementById('downloadModalTitle').textContent = document.getElementById('downloadModalTitle').textContent =
@@ -1839,7 +1838,9 @@ export class DownloadManager {
` (${validCount})`; ` (${validCount})`;
const list = document.getElementById('batchPreviewList'); const list = document.getElementById('batchPreviewList');
const hasHfItems = this.batchModels.some(m => m.source === 'huggingface' && !m.error); const hasExternalItems = this.batchModels.some(
m => isExternalModelSource(m.source) && !m.error
);
// Error items render flat, outside any group // Error items render flat, outside any group
const errorItemsHtml = this.batchModels.map((item, index) => { const errorItemsHtml = this.batchModels.map((item, index) => {
@@ -1863,7 +1864,7 @@ export class DownloadManager {
// CivitAI items render flat, outside any group (unchanged) // CivitAI items render flat, outside any group (unchanged)
const civitaiItemsHtml = this.batchModels.map((item, index) => { const civitaiItemsHtml = this.batchModels.map((item, index) => {
if (item.error) return null; if (item.error) return null;
if (item.source === 'huggingface') return null; if (isExternalModelSource(item.source)) return null;
const ver = item.selectedVersion; const ver = item.selectedVersion;
const firstImage = ver?.images?.find(img => !img.url.endsWith('.mp4')); const firstImage = ver?.images?.find(img => !img.url.endsWith('.mp4'));
const thumbnailUrl = firstImage ? firstImage.url : '/loras_static/images/no-preview.png'; const thumbnailUrl = firstImage ? firstImage.url : '/loras_static/images/no-preview.png';
@@ -1901,25 +1902,30 @@ export class DownloadManager {
`; `;
}).filter(Boolean).join(''); }).filter(Boolean).join('');
// Group HF items by repo (data model stays flat — only rendering groups) // Group external-repository items by platform + repo so that the same
const hfGroups = {}; // `owner/name` on two sites stays in two groups (data model stays flat
// — only rendering groups).
const externalGroups = {};
this.batchModels.forEach((item, index) => { this.batchModels.forEach((item, index) => {
if (item.error || item.source !== 'huggingface') return; if (item.error || !isExternalModelSource(item.source)) return;
const repo = item.repo || 'unknown'; const groupKey = this._externalGroupKey(item);
if (!hfGroups[repo]) hfGroups[repo] = []; if (!externalGroups[groupKey]) {
hfGroups[repo].push({ item, index }); externalGroups[groupKey] = { repo: item.repo || 'unknown', items: [] };
}
externalGroups[groupKey].items.push({ item, index });
}); });
const renderHfItem = ({ item, index }) => { const renderExternalItem = ({ item, index }) => {
const hfSize = item.fileSizeBytes ? formatFileSize(item.fileSizeBytes) : '?'; const fileSize = item.fileSizeBytes ? formatFileSize(item.fileSizeBytes) : '?';
const badge = getModelSource(item.source)?.label || item.source;
return ` return `
<div class="batch-preview-item" data-index="${index}"> <div class="batch-preview-item" data-index="${index}">
<input type="checkbox" class="batch-preview-checkbox" <input type="checkbox" class="batch-preview-checkbox"
data-index="${index}" ${item.checked !== false ? 'checked' : ''} /> data-index="${index}" ${item.checked !== false ? 'checked' : ''} />
<div class="batch-preview-info"> <div class="batch-preview-info">
<div class="batch-preview-name">${item.displayName || item.filename || `HF #${index}`} <span class="hf-badge">HF</span></div> <div class="batch-preview-name">${item.displayName || item.filename || `${badge} #${index}`} <span class="hf-badge">${badge}</span></div>
<div class="batch-preview-meta"> <div class="batch-preview-meta">
<span>${hfSize}</span> <span>${fileSize}</span>
<span>${item.repo || ''}</span> <span>${item.repo || ''}</span>
</div> </div>
</div> </div>
@@ -1930,32 +1936,32 @@ export class DownloadManager {
`; `;
}; };
const hfGroupsHtml = Object.keys(hfGroups).map(repo => { const externalGroupsHtml = Object.keys(externalGroups).map(groupKey => {
const items = hfGroups[repo]; const { repo, items } = externalGroups[groupKey];
const isCollapsed = this.hfRepoCollapsed[repo] === true; const isCollapsed = this.sourceRepoCollapsed[groupKey] === true;
const allChecked = items.every(({ item }) => item.checked !== false); const allChecked = items.every(({ item }) => item.checked !== false);
const fileCount = items.length; const fileCount = items.length;
return ` return `
<div class="batch-preview-group" data-repo="${repo}"> <div class="batch-preview-group" data-repo="${groupKey}">
<div class="batch-preview-group-header"> <div class="batch-preview-group-header">
<i class="fas fa-chevron-right batch-preview-group-toggle ${isCollapsed ? '' : 'expanded'}"></i> <i class="fas fa-chevron-right batch-preview-group-toggle ${isCollapsed ? '' : 'expanded'}"></i>
<span class="batch-preview-group-name">${repo}</span> <span class="batch-preview-group-name">${repo}</span>
<span class="batch-preview-group-count">${fileCount} ${translate('modals.download.fileSelection.files', {}, 'files')}</span> <span class="batch-preview-group-count">${fileCount} ${translate('modals.download.fileSelection.files', {}, 'files')}</span>
<input type="checkbox" class="batch-preview-group-select-all" data-repo="${repo}" ${allChecked ? 'checked' : ''} /> <input type="checkbox" class="batch-preview-group-select-all" data-repo="${groupKey}" ${allChecked ? 'checked' : ''} />
</div> </div>
<div class="batch-preview-group-body ${isCollapsed ? '' : 'expanded'}"> <div class="batch-preview-group-body ${isCollapsed ? '' : 'expanded'}">
${items.map(renderHfItem).join('')} ${items.map(renderExternalItem).join('')}
</div> </div>
</div> </div>
`; `;
}).join(''); }).join('');
let itemsHtml = errorItemsHtml + civitaiItemsHtml + hfGroupsHtml; let itemsHtml = errorItemsHtml + civitaiItemsHtml + externalGroupsHtml;
// Prepend select-all toolbar if there are HF items with checkboxes // Prepend select-all toolbar if there are external items with checkboxes
if (hasHfItems) { if (hasExternalItems) {
const allChecked = this.batchModels const allChecked = this.batchModels
.filter(m => m.source === 'huggingface' && !m.error) .filter(m => isExternalModelSource(m.source) && !m.error)
.every(m => m.checked !== false); .every(m => m.checked !== false);
itemsHtml = ` itemsHtml = `
<div class="batch-preview-select-all"> <div class="batch-preview-select-all">
@@ -1980,13 +1986,18 @@ export class DownloadManager {
// Global select-all // Global select-all
const selectAll = document.getElementById('batchSelectAll'); const selectAll = document.getElementById('batchSelectAll');
if (selectAll) { if (selectAll) {
const hfItems = this.batchModels.filter(m => m.source === 'huggingface' && !m.error); const externalItems = this.batchModels.filter(
selectAll.checked = hfItems.length > 0 && hfItems.every(m => m.checked !== false); m => isExternalModelSource(m.source) && !m.error
);
selectAll.checked = externalItems.length > 0
&& externalItems.every(m => m.checked !== false);
} }
// Per-group select-all // Per-group select-all
list.querySelectorAll('.batch-preview-group-select-all').forEach(gsa => { list.querySelectorAll('.batch-preview-group-select-all').forEach(gsa => {
const repo = gsa.dataset.repo; const repo = gsa.dataset.repo;
const repoItems = this.batchModels.filter(m => m.source === 'huggingface' && !m.error && m.repo === repo); const repoItems = this.batchModels.filter(
m => isExternalModelSource(m.source) && !m.error && this._externalGroupKey(m) === repo
);
gsa.checked = repoItems.length > 0 && repoItems.every(m => m.checked !== false); gsa.checked = repoItems.length > 0 && repoItems.every(m => m.checked !== false);
}); });
}; };
@@ -1998,7 +2009,7 @@ export class DownloadManager {
const repo = groupSelectAll.dataset.repo; const repo = groupSelectAll.dataset.repo;
const checked = groupSelectAll.checked; const checked = groupSelectAll.checked;
this.batchModels.forEach((m, idx) => { this.batchModels.forEach((m, idx) => {
if (m.source === 'huggingface' && !m.error && m.repo === repo) { if (isExternalModelSource(m.source) && !m.error && this._externalGroupKey(m) === repo) {
m.checked = checked; m.checked = checked;
const cb = list.querySelector(`.batch-preview-checkbox[data-index="${idx}"]`); const cb = list.querySelector(`.batch-preview-checkbox[data-index="${idx}"]`);
if (cb) cb.checked = checked; if (cb) cb.checked = checked;
@@ -2014,9 +2025,9 @@ export class DownloadManager {
const repo = group.dataset.repo; const repo = group.dataset.repo;
const body = group.querySelector('.batch-preview-group-body'); const body = group.querySelector('.batch-preview-group-body');
const toggle = group.querySelector('.batch-preview-group-toggle'); const toggle = group.querySelector('.batch-preview-group-toggle');
const isCollapsed = this.hfRepoCollapsed[repo]; const isCollapsed = this.sourceRepoCollapsed[repo];
if (isCollapsed) { if (isCollapsed) {
this.hfRepoCollapsed[repo] = false; this.sourceRepoCollapsed[repo] = false;
body.style.transition = ''; // restore in case collapse was interrupted body.style.transition = ''; // restore in case collapse was interrupted
body.classList.add('expanded'); body.classList.add('expanded');
toggle.classList.add('expanded'); toggle.classList.add('expanded');
@@ -2025,13 +2036,13 @@ export class DownloadManager {
body.style.maxHeight = body.scrollHeight + 'px'; body.style.maxHeight = body.scrollHeight + 'px';
const onEnd = (e) => { const onEnd = (e) => {
if (e.propertyName !== 'max-height') return; if (e.propertyName !== 'max-height') return;
if (this.hfRepoCollapsed[repo] !== false) return; if (this.sourceRepoCollapsed[repo] !== false) return;
body.style.maxHeight = ''; // fall back to .expanded's 9999px body.style.maxHeight = ''; // fall back to .expanded's 9999px
body.removeEventListener('transitionend', onEnd); body.removeEventListener('transitionend', onEnd);
}; };
body.addEventListener('transitionend', onEnd); body.addEventListener('transitionend', onEnd);
} else { } else {
this.hfRepoCollapsed[repo] = true; this.sourceRepoCollapsed[repo] = true;
body.style.maxHeight = body.scrollHeight + 'px'; body.style.maxHeight = body.scrollHeight + 'px';
requestAnimationFrame(() => { requestAnimationFrame(() => {
// animate only max-height; keep expanded so opacity stays 1 // animate only max-height; keep expanded so opacity stays 1
@@ -2040,7 +2051,7 @@ export class DownloadManager {
toggle.classList.remove('expanded'); toggle.classList.remove('expanded');
const onEnd = (e) => { const onEnd = (e) => {
if (e.propertyName !== 'max-height') return; if (e.propertyName !== 'max-height') return;
if (this.hfRepoCollapsed[repo] !== true) return; // state changed since if (this.sourceRepoCollapsed[repo] !== true) return; // state changed since
body.classList.remove('expanded'); body.classList.remove('expanded');
body.style.transition = ''; body.style.transition = '';
body.removeEventListener('transitionend', onEnd); body.removeEventListener('transitionend', onEnd);
@@ -2119,7 +2130,7 @@ export class DownloadManager {
// For HF items, respect the checked flag; for CivitAI items, use selectedVersion // For HF items, respect the checked flag; for CivitAI items, use selectedVersion
const validModels = this.batchModels.filter(m => { const validModels = this.batchModels.filter(m => {
if (m.error) return false; if (m.error) return false;
if (m.source === 'huggingface') return m.checked !== false; if (isExternalModelSource(m.source)) return m.checked !== false;
return m.selectedVersion; return m.selectedVersion;
}); });
if (validModels.length === 0) return; if (validModels.length === 0) return;
@@ -2172,8 +2183,8 @@ export class DownloadManager {
} }
if (!this.isBatchMode) { if (!this.isBatchMode) {
// Single-item download // Single-item download
if (this.source === 'huggingface') { if (isExternalModelSource(this.source)) {
return this._downloadHfSingle({ return this._downloadExternalRepoFiles({
modelRoot, modelRoot,
targetFolder, targetFolder,
useDefaultPaths, useDefaultPaths,
@@ -2228,7 +2239,7 @@ export class DownloadManager {
if (m.error) return false; if (m.error) return false;
if (!m.selectedVersion) return false; if (!m.selectedVersion) return false;
// HF items have selectedVersion as a boolean marker + checked flag // HF items have selectedVersion as a boolean marker + checked flag
if (m.source === 'huggingface') return m.checked !== false; if (isExternalModelSource(m.source)) return m.checked !== false;
return !m.selectedVersion.existsLocally; return !m.selectedVersion.existsLocally;
}); });
if (downloadItems.length === 0) { if (downloadItems.length === 0) {
@@ -2255,10 +2266,11 @@ export class DownloadManager {
let cancelled = false; let cancelled = false;
const failedItems = []; const failedItems = [];
// Successful CivitAI items are reconciled in place afterwards // Successful CivitAI items are reconciled in place afterwards
// (their cards can be matched by model id); HF items keep the // (their cards can be matched by model id); externally-sourced items
// legacy full reload because they have no CivitAI identity (#1078). // keep the legacy full reload because they have no CivitAI identity
// (#1078).
const completedCivitaiItems = []; const completedCivitaiItems = [];
let hfCompletedCount = 0; let externalCompletedCount = 0;
loadingManager.showCancelButton(async () => { loadingManager.showCancelButton(async () => {
if (cancelled) return; if (cancelled) return;
@@ -2301,15 +2313,15 @@ export class DownloadManager {
const item = downloadItems[i]; const item = downloadItems[i];
const name = item.displayName || item.filename || (item.selectedVersion?.name || `Model #${item.modelId}`); const name = item.displayName || item.filename || (item.selectedVersion?.name || `Model #${item.modelId}`);
const isHf = item.source === 'huggingface'; const isExternal = isExternalModelSource(item.source);
updateProgress(0, completedDownloads, name); updateProgress(0, completedDownloads, name);
loadingManager.setStatus(`${i + 1}/${downloadItems.length}: ${name}`); loadingManager.setStatus(`${i + 1}/${downloadItems.length}: ${name}`);
try { try {
let response; let response;
if (isHf) { if (isExternal) {
const downloadId = Date.now().toString() + '_hf_' + i; const downloadId = Date.now().toString() + '_src_' + i;
const wsHf = new WebSocket(`${wsProtocol}${window.location.host}/ws/download-progress?id=${downloadId}`); const wsHf = new WebSocket(`${wsProtocol}${window.location.host}/ws/download-progress?id=${downloadId}`);
try { try {
await new Promise((resolve, reject) => { await new Promise((resolve, reject) => {
@@ -2329,10 +2341,11 @@ export class DownloadManager {
} }
}; };
response = await this.apiClient.downloadHfModel({ response = await this.apiClient.downloadModelSource({
platform: item.platform || item.source,
repo: item.repo, repo: item.repo,
filename: item.filename, filename: item.filename,
revision: item.revision || 'main', revision: item.revision || this._defaultRevisionFor(item.platform || item.source),
modelRoot, modelRoot,
relativePath: targetFolder, relativePath: targetFolder,
useDefaultPaths, useDefaultPaths,
@@ -2363,8 +2376,8 @@ export class DownloadManager {
} else { } else {
completedDownloads++; completedDownloads++;
updateProgress(100, completedDownloads, ''); updateProgress(100, completedDownloads, '');
if (isHf) { if (isExternal) {
hfCompletedCount++; externalCompletedCount++;
} else { } else {
completedCivitaiItems.push(item); completedCivitaiItems.push(item);
} }
@@ -2398,7 +2411,7 @@ export class DownloadManager {
}); });
} }
await this._reconcileBatchViewAfterDownload(completedCivitaiItems, hfCompletedCount); await this._reconcileBatchViewAfterDownload(completedCivitaiItems, externalCompletedCount);
} }
async downloadVersionWithDefaults(modelType, modelId, versionId, { async downloadVersionWithDefaults(modelType, modelId, versionId, {
+119 -1
View File
@@ -19,21 +19,35 @@ export const MODEL_SOURCES = [
groupPrefix: 'hf', groupPrefix: 'hf',
supportsEnrichment: true, supportsEnrichment: true,
supportsDownload: true, supportsDownload: true,
defaultRevision: 'main',
defaultSubdir: 'huggingface',
exampleUrl: 'https://huggingface.co/user/repo', exampleUrl: 'https://huggingface.co/user/repo',
placeholder: 'https://huggingface.co/user/repo', placeholder: 'https://huggingface.co/user/repo',
pattern: /^https?:\/\/(?:www\.)?huggingface\.co\/([^/?#\s]+\/[^/?#\s]+)/i, pattern: /^https?:\/\/(?:www\.)?huggingface\.co\/([^/?#\s]+\/[^/?#\s]+)/i,
// `blob` is the web preview page; it maps 1:1 to the `resolve` download URL.
filePattern:
/^https?:\/\/(?:www\.)?huggingface\.co\/([^/?#\s]+\/[^/?#\s]+)\/(?:resolve|blob)\/([^/?#\s]+)\/(.+)$/i,
canonical: (id) => `https://huggingface.co/${id}`, canonical: (id) => `https://huggingface.co/${id}`,
filePage: (id, filename) => `https://huggingface.co/${id}/blob/main/${filename}`,
// Bare `user/repo` has always meant Hugging Face; keep that meaning.
bareRepoPattern: /^([A-Za-z0-9_.-]+\/[A-Za-z0-9_.-]+)$/,
}, },
{ {
platform: 'modelscope', platform: 'modelscope',
label: 'ModelScope', label: 'ModelScope',
groupPrefix: 'ms', groupPrefix: 'ms',
supportsEnrichment: true, supportsEnrichment: true,
supportsDownload: false, supportsDownload: true,
defaultRevision: 'master',
defaultSubdir: 'modelscope',
exampleUrl: 'https://modelscope.cn/models/user/repo', exampleUrl: 'https://modelscope.cn/models/user/repo',
placeholder: 'https://modelscope.cn/models/user/repo', placeholder: 'https://modelscope.cn/models/user/repo',
pattern: /^https?:\/\/(?:www\.)?modelscope\.(?:cn|com)\/models\/([^/?#\s]+\/[^/?#\s]+)/i, pattern: /^https?:\/\/(?:www\.)?modelscope\.(?:cn|com)\/models\/([^/?#\s]+\/[^/?#\s]+)/i,
filePattern:
/^https?:\/\/(?:www\.)?modelscope\.(?:cn|com)\/models\/([^/?#\s]+\/[^/?#\s]+)\/resolve\/([^/?#\s]+)\/(.+)$/i,
canonical: (id) => `https://modelscope.cn/models/${id}`, canonical: (id) => `https://modelscope.cn/models/${id}`,
filePage: (id, filename) =>
`https://modelscope.cn/models/${id}/file/view/master/${filename}`,
}, },
{ {
platform: 'tensorart', platform: 'tensorart',
@@ -41,10 +55,14 @@ export const MODEL_SOURCES = [
groupPrefix: 'ta', groupPrefix: 'ta',
supportsEnrichment: false, supportsEnrichment: false,
supportsDownload: false, supportsDownload: false,
defaultRevision: '',
defaultSubdir: '',
exampleUrl: 'https://tensor.art/models/827823520299086029', exampleUrl: 'https://tensor.art/models/827823520299086029',
placeholder: 'https://tensor.art/models/827823520299086029', placeholder: 'https://tensor.art/models/827823520299086029',
pattern: /^https?:\/\/(?:www\.)?(?:tensor\.art|tusi\.cn)\/models\/(\d+)/i, pattern: /^https?:\/\/(?:www\.)?(?:tensor\.art|tusi\.cn)\/models\/(\d+)/i,
filePattern: null,
canonical: (id) => `https://tensor.art/models/${id}`, canonical: (id) => `https://tensor.art/models/${id}`,
filePage: null,
}, },
]; ];
@@ -171,3 +189,103 @@ export function openModelSource(url) {
if (!url) return; if (!url) return;
window.open(url, '_blank', 'noopener,noreferrer'); window.open(url, '_blank', 'noopener,noreferrer');
} }
// ---------------------------------------------------------------------------
// Download support
// ---------------------------------------------------------------------------
/** Sources whose repositories the backend can download from. */
export const DOWNLOADABLE_SOURCES = MODEL_SOURCES.filter((s) => s.supportsDownload);
/**
* Whether a DownloadManager `source` value refers to an external repository
* download (as opposed to a CivitAI/CivArchive version or a direct link).
*/
export function isExternalModelSource(source) {
return DOWNLOADABLE_SOURCES.some((s) => s.platform === source);
}
/** Return the downloadable source descriptor for a platform, or null. */
export function getDownloadSource(platform) {
const source = getModelSource(platform);
return source && source.supportsDownload ? source : null;
}
/** Normalise a repository id: reject traversal, exactly one slash. */
export function isValidRepoId(repo) {
if (!repo || typeof repo !== 'string' || repo.split('/').length !== 2) return false;
return repo
.split('/')
.every((part) => part && part !== '.' && part !== '..' && /^[A-Za-z0-9_][\w.-]*$/.test(part));
}
/**
* Recognise a downloadable model-source URL.
*
* Handles both a repository page and a direct file (resolve) URL for every
* source that supports downloads, plus the historical bare `owner/name`
* shorthand, which only ever meant Hugging Face.
*
* @returns {{kind: 'repo'|'file', platform: string, label: string,
* repo: string, revision?: string, filename?: string}|null}
*/
export function detectModelSourceDownloadUrl(url) {
if (!url || typeof url !== 'string') return null;
const candidate = url.trim();
if (!candidate) return null;
// Direct file URLs first: the repo pattern would match their prefix and
// lose the revision/filename.
for (const source of DOWNLOADABLE_SOURCES) {
if (!source.filePattern) continue;
const match = candidate.match(source.filePattern);
if (match) {
return {
kind: 'file',
platform: source.platform,
label: source.label,
repo: match[1],
revision: match[2],
filename: match[3],
};
}
}
for (const source of DOWNLOADABLE_SOURCES) {
const match = candidate.match(source.pattern);
if (match) {
return {
kind: 'repo',
platform: source.platform,
label: source.label,
repo: match[1],
};
}
}
if (!candidate.includes('://')) {
for (const source of DOWNLOADABLE_SOURCES) {
if (!source.bareRepoPattern) continue;
const match = candidate.match(source.bareRepoPattern);
if (match && isValidRepoId(match[1])) {
return {
kind: 'repo',
platform: source.platform,
label: source.label,
repo: match[1],
};
}
}
}
return null;
}
/** Human-facing page for one file of an external repository. */
export function buildModelSourceFilePage({ platform, repo, filename }) {
const source = getModelSource(platform);
if (!source || !source.filePage || !filename) {
return source ? source.canonical(repo) : null;
}
return source.filePage(repo, filename);
}
@@ -26,6 +26,7 @@ const {
}, },
}, },
downloadModel: vi.fn(), downloadModel: vi.fn(),
downloadModelSource: vi.fn(),
downloadHfModel: vi.fn(), downloadHfModel: vi.fn(),
cancelDownload: vi.fn(), cancelDownload: vi.fn(),
getPageState: vi.fn(() => ({})), getPageState: vi.fn(() => ({})),
@@ -158,7 +159,7 @@ describe('DownloadManager batch download summary flow', () => {
// Reset the shared mocks so mockResolvedValueOnce queues and call // Reset the shared mocks so mockResolvedValueOnce queues and call
// history never leak between tests. // history never leak between tests.
mockApiClient.downloadModel.mockReset(); mockApiClient.downloadModel.mockReset();
mockApiClient.downloadHfModel.mockReset(); mockApiClient.downloadModelSource.mockReset();
mockApiClient.cancelDownload.mockReset(); mockApiClient.cancelDownload.mockReset();
showToastMock.mockClear(); showToastMock.mockClear();
showDownloadBatchSummaryMock.mockClear(); showDownloadBatchSummaryMock.mockClear();
@@ -406,14 +407,15 @@ describe('DownloadManager batch download summary flow', () => {
expect(showToastMock).toHaveBeenCalledWith('toast.loras.downloadCompleted', {}, 'success'); expect(showToastMock).toHaveBeenCalledWith('toast.loras.downloadCompleted', {}, 'success');
}); });
it('shows a summary for HF partial failure and retries only the failed files', async () => { it('shows a summary for external repo partial failure and retries only the failed files', async () => {
manager.hfRepoId = 'user/repo'; manager.sourcePlatform = 'huggingface';
manager.hfSelectedFiles = ['a.safetensors', 'b.safetensors']; manager.sourceRepoId = 'user/repo';
mockApiClient.downloadHfModel manager.sourceSelectedFiles = ['a.safetensors', 'b.safetensors'];
mockApiClient.downloadModelSource
.mockResolvedValueOnce({ success: true }) .mockResolvedValueOnce({ success: true })
.mockResolvedValueOnce({ success: false, error: 'denied' }); .mockResolvedValueOnce({ success: false, error: 'denied' });
const result = await manager._downloadHfSingle({ modelRoot: '/m', useDefaultPaths: true }); const result = await manager._downloadExternalRepoFiles({ modelRoot: '/m', useDefaultPaths: true });
expect(result).toBe(false); expect(result).toBe(false);
expect(showDownloadBatchSummaryMock).toHaveBeenCalledTimes(1); expect(showDownloadBatchSummaryMock).toHaveBeenCalledTimes(1);
@@ -428,7 +430,9 @@ describe('DownloadManager batch download summary flow', () => {
await summary.onRetry(); await summary.onRetry();
expect(mockApiClient.downloadHfModel).toHaveBeenCalledTimes(3); expect(mockApiClient.downloadModelSource).toHaveBeenCalledTimes(3);
expect(mockApiClient.downloadHfModel.mock.calls[2][0].filename).toBe('b.safetensors'); const retryArgs = mockApiClient.downloadModelSource.mock.calls[2][0];
expect(retryArgs.filename).toBe('b.safetensors');
expect(retryArgs.platform).toBe('huggingface');
}); });
}); });
@@ -0,0 +1,269 @@
import { beforeEach, describe, expect, it, vi } from 'vitest';
const {
DOWNLOAD_MANAGER_MODULE,
MODAL_MANAGER_MODULE,
UI_HELPERS_MODULE,
STATE_MODULE,
LOADING_MANAGER_MODULE,
API_FACTORY_MODULE,
STORAGE_HELPERS_MODULE,
FOLDER_TREE_MANAGER_MODULE,
I18N_HELPERS_MODULE,
SUMMARY_MODULE,
mockApiClient,
mockLoadingManager,
showDownloadBatchSummaryMock,
} = vi.hoisted(() => {
const mockApiClient = {
apiConfig: { config: { displayName: 'LoRA', singularName: 'lora' } },
downloadModel: vi.fn(),
downloadModelSource: vi.fn(),
fetchModelSourceFiles: vi.fn(),
cancelDownload: vi.fn(),
getPageState: vi.fn(() => ({})),
};
const mockLoadingManager = {
showSimpleLoading: vi.fn(),
hide: vi.fn(),
restoreProgressBar: vi.fn(),
showDownloadProgress: vi.fn(() => vi.fn()),
setStatus: vi.fn(),
showCancelButton: vi.fn(),
};
return {
DOWNLOAD_MANAGER_MODULE: new URL('../../../static/js/managers/DownloadManager.js', import.meta.url).pathname,
MODAL_MANAGER_MODULE: new URL('../../../static/js/managers/ModalManager.js', import.meta.url).pathname,
UI_HELPERS_MODULE: new URL('../../../static/js/utils/uiHelpers.js', import.meta.url).pathname,
STATE_MODULE: new URL('../../../static/js/state/index.js', import.meta.url).pathname,
LOADING_MANAGER_MODULE: new URL('../../../static/js/managers/LoadingManager.js', import.meta.url).pathname,
API_FACTORY_MODULE: new URL('../../../static/js/api/modelApiFactory.js', import.meta.url).pathname,
STORAGE_HELPERS_MODULE: new URL('../../../static/js/utils/storageHelpers.js', import.meta.url).pathname,
FOLDER_TREE_MANAGER_MODULE: new URL('../../../static/js/components/FolderTreeManager.js', import.meta.url).pathname,
I18N_HELPERS_MODULE: new URL('../../../static/js/utils/i18nHelpers.js', import.meta.url).pathname,
SUMMARY_MODULE: new URL('../../../static/js/components/DownloadBatchSummaryModal.js', import.meta.url).pathname,
mockApiClient,
mockLoadingManager,
showDownloadBatchSummaryMock: vi.fn(),
};
});
vi.mock(MODAL_MANAGER_MODULE, () => ({
modalManager: { showModal: vi.fn(), closeModal: vi.fn() },
}));
vi.mock(UI_HELPERS_MODULE, () => ({ showToast: vi.fn() }));
vi.mock(STATE_MODULE, () => ({
state: { global: { settings: {} }, loadingManager: mockLoadingManager },
}));
vi.mock(LOADING_MANAGER_MODULE, () => ({
LoadingManager: vi.fn(() => mockLoadingManager),
}));
vi.mock(API_FACTORY_MODULE, () => ({
getModelApiClient: vi.fn(() => mockApiClient),
resetAndReload: vi.fn(),
}));
vi.mock(STORAGE_HELPERS_MODULE, () => ({
getStorageItem: vi.fn((_key, defaultValue) => defaultValue),
setStorageItem: vi.fn(),
}));
vi.mock(FOLDER_TREE_MANAGER_MODULE, () => ({
FolderTreeManager: vi.fn(() => ({ clearSelection: vi.fn(), init: vi.fn() })),
}));
vi.mock(I18N_HELPERS_MODULE, () => ({
translate: vi.fn((_, __, fallback) => fallback ?? ''),
}));
vi.mock(SUMMARY_MODULE, () => ({
showDownloadBatchSummary: showDownloadBatchSummaryMock,
}));
class FakeWebSocket {
constructor(url) {
this.url = url;
this.onopen = null;
this.onmessage = null;
this.onerror = null;
this.close = vi.fn();
queueMicrotask(() => {
if (this.onopen) this.onopen();
});
}
}
const MS_REPO_URL = 'https://modelscope.cn/models/jj3550945163/Krea-2-LORA';
describe('DownloadManager external model source downloads', () => {
let DownloadManager;
let manager;
beforeEach(async () => {
document.body.innerHTML = '';
vi.stubGlobal('WebSocket', FakeWebSocket);
mockApiClient.downloadModelSource.mockReset();
mockApiClient.fetchModelSourceFiles.mockReset();
mockLoadingManager.showSimpleLoading.mockReset();
({ DownloadManager } = await import(DOWNLOAD_MANAGER_MODULE));
manager = new DownloadManager();
manager.apiClient = mockApiClient;
manager.showBatchPreviewStep = vi.fn();
manager.proceedToLocation = vi.fn();
});
it('loads a ModelScope repo as batch items on the master revision', async () => {
mockApiClient.fetchModelSourceFiles.mockResolvedValue([
{ filename: 'a.safetensors', size: 10 },
{ filename: 'sub/b.safetensors', size: 20 },
]);
const errorElement = { textContent: '' };
await manager._validateAndFetchExternalRepo([MS_REPO_URL], errorElement);
expect(mockApiClient.fetchModelSourceFiles).toHaveBeenCalledWith(
'jj3550945163/Krea-2-LORA',
'modelscope',
'master'
);
expect(errorElement.textContent).toBe('');
expect(manager.source).toBe('modelscope');
expect(manager.isBatchMode).toBe(true);
expect(manager.batchModels).toHaveLength(2);
expect(manager.batchModels[0]).toMatchObject({
source: 'modelscope',
platform: 'modelscope',
repo: 'jj3550945163/Krea-2-LORA',
revision: 'master',
filename: 'a.safetensors',
fileSizeBytes: 10,
displayName: 'a.safetensors',
});
expect(manager.showBatchPreviewStep).toHaveBeenCalled();
});
it('keeps Hugging Face on its own revision', async () => {
mockApiClient.fetchModelSourceFiles.mockResolvedValue([
{ filename: 'a.safetensors', size: 10 },
]);
await manager._validateAndFetchExternalRepo(
['https://huggingface.co/user/repo'],
{ textContent: '' }
);
expect(mockApiClient.fetchModelSourceFiles).toHaveBeenCalledWith(
'user/repo',
'huggingface',
'main'
);
expect(manager.batchModels[0].revision).toBe('main');
});
it('surfaces a listing failure on the URL field', async () => {
mockApiClient.fetchModelSourceFiles.mockRejectedValue(new Error('Repository not found'));
const errorElement = { textContent: '' };
await manager._validateAndFetchExternalRepo([MS_REPO_URL], errorElement);
expect(errorElement.textContent).toBe('Repository not found');
expect(manager.showBatchPreviewStep).not.toHaveBeenCalled();
});
it('skips file selection for a direct ModelScope file URL', async () => {
await manager._validateAndFetchExternalRepo(
[`${MS_REPO_URL}/resolve/master/Krea-2-LORA_c1-st1000.safetensors`],
{ textContent: '' }
);
expect(manager.isBatchMode).toBe(false);
expect(manager.sourcePlatform).toBe('modelscope');
expect(manager.sourceRepoId).toBe('jj3550945163/Krea-2-LORA');
expect(manager.sourceSelectedFiles).toEqual(['Krea-2-LORA_c1-st1000.safetensors']);
expect(manager.proceedToLocation).toHaveBeenCalled();
});
it('downloads a single ModelScope file through the generic endpoint', async () => {
mockApiClient.downloadModelSource.mockResolvedValue({ success: true });
manager.sourcePlatform = 'modelscope';
manager.sourceRepoId = 'jj3550945163/Krea-2-LORA';
manager.sourceSelectedFiles = ['Krea-2-LORA_c1-st1000.safetensors'];
await manager._downloadExternalRepoFiles({
modelRoot: '/models/loras',
targetFolder: '',
useDefaultPaths: true,
});
expect(mockApiClient.downloadModelSource).toHaveBeenCalledTimes(1);
expect(mockApiClient.downloadModelSource.mock.calls[0][0]).toMatchObject({
platform: 'modelscope',
repo: 'jj3550945163/Krea-2-LORA',
filename: 'Krea-2-LORA_c1-st1000.safetensors',
revision: 'master',
});
});
it('carries the platform through a batch download', async () => {
mockApiClient.downloadModelSource.mockResolvedValue({ success: true });
manager.showBatchPreviewStep = vi.fn();
await manager.executeBatchDownload(
[
{
source: 'modelscope',
platform: 'modelscope',
repo: 'u/r',
filename: 'f.safetensors',
revision: 'master',
displayName: 'f.safetensors',
checked: true,
},
],
{ modelRoot: '/models/loras', targetFolder: '', useDefaultPaths: true }
);
expect(mockApiClient.downloadModelSource).toHaveBeenCalledTimes(1);
expect(mockApiClient.downloadModelSource.mock.calls[0][0]).toMatchObject({
platform: 'modelscope',
repo: 'u/r',
filename: 'f.safetensors',
revision: 'master',
});
expect(mockApiClient.downloadModel).not.toHaveBeenCalled();
});
it('links failures to the ModelScope file page', async () => {
expect(
manager._buildSingleItemUrl({
source: 'modelscope',
repo: 'u/r',
filename: 'sub/f.safetensors',
})
).toBe('https://modelscope.cn/models/u/r/file/view/master/sub/f.safetensors');
expect(
manager._buildSingleItemUrl({
source: 'huggingface',
repo: 'u/r',
filename: 'f.safetensors',
})
).toBe('https://huggingface.co/u/r/blob/main/f.safetensors');
});
it('groups the same repo name on two platforms separately', () => {
const hf = { source: 'huggingface', repo: 'u/r' };
const ms = { source: 'modelscope', repo: 'u/r' };
expect(manager._externalGroupKey(hf)).toBe('huggingface:u/r');
expect(manager._externalGroupKey(ms)).toBe('modelscope:u/r');
expect(manager._externalGroupKey(hf)).not.toBe(manager._externalGroupKey(ms));
});
});
@@ -1,14 +1,15 @@
import { describe, it, expect } from 'vitest'; import { describe, it, expect } from 'vitest';
import { DownloadManager } from '../../../static/js/managers/DownloadManager.js'; import { DownloadManager } from '../../../static/js/managers/DownloadManager.js';
describe('DownloadManager.detectUrlType — HF URL detection', () => { describe('DownloadManager.detectUrlType — external model source URLs', () => {
it('detects HF resolve URL with file', () => { it('detects HF resolve URL with file', () => {
const result = DownloadManager.detectUrlType( const result = DownloadManager.detectUrlType(
'https://huggingface.co/dx8152/Flux2-Klein-9B-Consistency/resolve/main/Flux2-Klein-9B-consistency-V2.safetensors' 'https://huggingface.co/dx8152/Flux2-Klein-9B-Consistency/resolve/main/Flux2-Klein-9B-consistency-V2.safetensors'
); );
expect(result).toEqual({ expect(result).toEqual({
type: 'hf-resolve', type: 'model-source-file',
platform: 'huggingface',
repo: 'dx8152/Flux2-Klein-9B-Consistency', repo: 'dx8152/Flux2-Klein-9B-Consistency',
revision: 'main', revision: 'main',
filename: 'Flux2-Klein-9B-consistency-V2.safetensors', filename: 'Flux2-Klein-9B-consistency-V2.safetensors',
@@ -20,7 +21,8 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
'https://huggingface.co/user/repo/resolve/main/subdir/model.safetensors' 'https://huggingface.co/user/repo/resolve/main/subdir/model.safetensors'
); );
expect(result).toEqual({ expect(result).toEqual({
type: 'hf-resolve', type: 'model-source-file',
platform: 'huggingface',
repo: 'user/repo', repo: 'user/repo',
revision: 'main', revision: 'main',
filename: 'subdir/model.safetensors', filename: 'subdir/model.safetensors',
@@ -32,7 +34,8 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
'https://huggingface.co/dx8152/Flux2-Klein-9B-Consistency' 'https://huggingface.co/dx8152/Flux2-Klein-9B-Consistency'
); );
expect(result).toEqual({ expect(result).toEqual({
type: 'hf-repo', type: 'model-source-repo',
platform: 'huggingface',
repo: 'dx8152/Flux2-Klein-9B-Consistency', repo: 'dx8152/Flux2-Klein-9B-Consistency',
}); });
}); });
@@ -40,7 +43,8 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
it('detects HF repo URL (bare user/repo)', () => { it('detects HF repo URL (bare user/repo)', () => {
const result = DownloadManager.detectUrlType('dx8152/Flux2-Klein-9B-Consistency'); const result = DownloadManager.detectUrlType('dx8152/Flux2-Klein-9B-Consistency');
expect(result).toEqual({ expect(result).toEqual({
type: 'hf-repo', type: 'model-source-repo',
platform: 'huggingface',
repo: 'dx8152/Flux2-Klein-9B-Consistency', repo: 'dx8152/Flux2-Klein-9B-Consistency',
}); });
}); });
@@ -50,7 +54,8 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
'https://huggingface.co/user/repo/' 'https://huggingface.co/user/repo/'
); );
expect(result).toEqual({ expect(result).toEqual({
type: 'hf-repo', type: 'model-source-repo',
platform: 'huggingface',
repo: 'user/repo', repo: 'user/repo',
}); });
}); });
@@ -60,7 +65,8 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
'https://huggingface.co/Comfy-Org/z_image_turbo/blob/main/split_files/diffusion_models/z_image_turbo_bf16.safetensors' 'https://huggingface.co/Comfy-Org/z_image_turbo/blob/main/split_files/diffusion_models/z_image_turbo_bf16.safetensors'
); );
expect(result).toEqual({ expect(result).toEqual({
type: 'hf-resolve', type: 'model-source-file',
platform: 'huggingface',
repo: 'Comfy-Org/z_image_turbo', repo: 'Comfy-Org/z_image_turbo',
revision: 'main', revision: 'main',
filename: 'split_files/diffusion_models/z_image_turbo_bf16.safetensors', filename: 'split_files/diffusion_models/z_image_turbo_bf16.safetensors',
@@ -115,7 +121,7 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
const result = DownloadManager.detectUrlType( const result = DownloadManager.detectUrlType(
'https://huggingface.co/user/repo/resolve/main/file.safetensors' 'https://huggingface.co/user/repo/resolve/main/file.safetensors'
); );
expect(result?.type).toBe('hf-resolve'); expect(result?.type).toBe('model-source-file');
}); });
it('prefers CivitAI over HF when both match', () => { it('prefers CivitAI over HF when both match', () => {
@@ -126,4 +132,51 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
); );
expect(result?.type).toBe('civitai'); expect(result?.type).toBe('civitai');
}); });
it('detects a ModelScope repo URL', () => {
const result = DownloadManager.detectUrlType(
'https://modelscope.cn/models/jj3550945163/Krea-2-LORA'
);
expect(result).toEqual({
type: 'model-source-repo',
platform: 'modelscope',
repo: 'jj3550945163/Krea-2-LORA',
});
});
it('detects a ModelScope file URL with revision and subdirectory', () => {
const result = DownloadManager.detectUrlType(
'https://modelscope.cn/models/AI-ModelScope/stable-diffusion-v1-5/resolve/master/vae/diffusion_pytorch_model.bin'
);
expect(result).toEqual({
type: 'model-source-file',
platform: 'modelscope',
repo: 'AI-ModelScope/stable-diffusion-v1-5',
revision: 'master',
filename: 'vae/diffusion_pytorch_model.bin',
});
});
it('detects a ModelScope view sub-page as a repo URL', () => {
const result = DownloadManager.detectUrlType(
'https://www.modelscope.cn/models/user/repo/summary'
);
expect(result).toEqual({
type: 'model-source-repo',
platform: 'modelscope',
repo: 'user/repo',
});
});
it('does not treat a bare owner/name as ModelScope', () => {
// The shorthand has always meant Hugging Face; ModelScope needs its host.
const result = DownloadManager.detectUrlType('user/repo');
expect(result.platform).toBe('huggingface');
});
it('rejects path traversal in either platform', () => {
expect(
DownloadManager.detectUrlType('https://modelscope.cn/models/../etc/passwd')
).toBeNull();
});
}); });
-308
View File
@@ -1,308 +0,0 @@
"""Tests for the HuggingFace link handler (``set_hf_url``).
Regression coverage for issue #1094: linking a model to HuggingFace must not
clear its CivitAI provenance or metadata, so both "View on CivitAI" and
"View on Hugging Face" can coexist.
"""
from __future__ import annotations
import json
import os
from typing import Any
from unittest.mock import AsyncMock
import pytest
from py.routes.handlers import hf_handlers
from py.routes.handlers.hf_handlers import HfHandler
from py.utils.metadata_manager import MetadataManager
def _json_payload(response) -> dict[str, Any]:
assert response.text is not None
return json.loads(response.text)
class FakeRequest:
def __init__(self, *, json_data=None):
self._json_data = json_data or {}
async def json(self):
return self._json_data
def _sidecar_path(model_path) -> str:
return f"{os.path.splitext(str(model_path))[0]}.metadata.json"
@pytest.fixture
def hf_env(tmp_path, monkeypatch):
"""Point HF linking at *tmp_path* and stub the scanner cache write."""
monkeypatch.setattr(hf_handlers, "_find_matching_root", lambda _dir: str(tmp_path))
cache_write = AsyncMock()
monkeypatch.setattr(hf_handlers, "_add_to_scanner_cache", cache_write)
return {"root": tmp_path, "cache_write": cache_write}
async def _write_model(model_path, payload: dict[str, Any]) -> None:
model_path.write_bytes(b"x" * 32)
await MetadataManager.save_metadata(str(model_path), payload)
@pytest.mark.asyncio
async def test_set_hf_url_keeps_civitai_metadata_and_provenance(tmp_path, hf_env):
model_path = tmp_path / "civitai_model.safetensors"
await _write_model(
model_path,
{
"file_name": "civitai_model",
"model_name": "CivitAI Model",
"file_path": str(model_path),
"size": 32,
"modified": 1.0,
"sha256": "a" * 64,
"base_model": "SDXL 1.0",
"preview_url": "",
"from_civitai": True,
"civitai": {"id": 111, "modelId": 222, "name": "v1", "trainedWords": []},
},
)
response = await HfHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"hf_url": "https://huggingface.co/user/repo",
}
)
)
assert response.status == 200
assert _json_payload(response)["success"] is True
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["hf_url"] == "https://huggingface.co/user/repo"
# Linking HF must not erase the model's CivitAI provenance or data.
assert saved["from_civitai"] is True
assert saved["civitai"]["modelId"] == 222
assert saved["civitai"]["id"] == 111
hf_env["cache_write"].assert_awaited_once()
cached_metadata = hf_env["cache_write"].await_args.args[1]
assert cached_metadata["hf_url"] == "https://huggingface.co/user/repo"
assert cached_metadata["from_civitai"] is True
assert cached_metadata["civitai"]["modelId"] == 222
@pytest.mark.asyncio
async def test_set_hf_url_does_not_force_from_civitai_false(tmp_path, hf_env):
"""A model without CivitAI data keeps its existing provenance flag."""
model_path = tmp_path / "hf_only.safetensors"
await _write_model(
model_path,
{
"file_name": "hf_only",
"model_name": "HF Only",
"file_path": str(model_path),
"size": 32,
"modified": 1.0,
"sha256": "b" * 64,
"base_model": "Unknown",
"preview_url": "",
"from_civitai": True,
},
)
response = await HfHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"hf_url": "https://huggingface.co/user/repo",
}
)
)
assert response.status == 200
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["hf_url"] == "https://huggingface.co/user/repo"
assert saved["from_civitai"] is True
@pytest.mark.asyncio
async def test_set_hf_url_rejects_non_repo_url(tmp_path, hf_env):
model_path = tmp_path / "model.safetensors"
await _write_model(
model_path,
{
"file_name": "model",
"model_name": "model",
"file_path": str(model_path),
"size": 32,
"modified": 1.0,
"sha256": "c" * 64,
"base_model": "Unknown",
"preview_url": "",
},
)
response = await HfHandler().set_hf_url(
FakeRequest(json_data={"file_path": str(model_path), "hf_url": "https://example.com/x"})
)
assert response.status == 400
payload = _json_payload(response)
assert payload["success"] is False
hf_env["cache_write"].assert_not_awaited()
# ---------------------------------------------------------------------------
# Multi-source linking (ModelScope / TensorArt)
# ---------------------------------------------------------------------------
async def _write_plain_model(model_path, sha: str = "d" * 64) -> None:
await _write_model(
model_path,
{
"file_name": "model",
"model_name": "model",
"file_path": str(model_path),
"size": 32,
"modified": 1.0,
"sha256": sha,
"base_model": "Unknown",
"preview_url": "",
},
)
@pytest.mark.asyncio
async def test_set_hf_url_accepts_modelscope_and_stores_source_fields(tmp_path, hf_env):
model_path = tmp_path / "ms_model.safetensors"
await _write_plain_model(model_path)
response = await HfHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://modelscope.cn/models/jj3550945163/Krea-2-LORA",
}
)
)
assert response.status == 200
payload = _json_payload(response)
assert payload["source_platform"] == "modelscope"
assert payload["source_url"] == "https://modelscope.cn/models/jj3550945163/Krea-2-LORA"
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["source_platform"] == "modelscope"
assert saved["source_url"] == "https://modelscope.cn/models/jj3550945163/Krea-2-LORA"
# No stale Hugging Face alias for a ModelScope model.
assert saved.get("hf_url", "") == ""
cached_metadata = hf_env["cache_write"].await_args.args[1]
assert cached_metadata["source_platform"] == "modelscope"
@pytest.mark.asyncio
async def test_set_hf_url_accepts_tensorart_url(tmp_path, hf_env):
model_path = tmp_path / "ta_model.safetensors"
await _write_plain_model(model_path, sha="e" * 64)
response = await HfHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": (
"https://tensor.art/models/827823520299086029/"
"Vivid-Impressions-Storybook-Sstyle-V1.0"
),
}
)
)
assert response.status == 200
payload = _json_payload(response)
assert payload["source_platform"] == "tensorart"
# The canonical page URL is stored, without the slug.
assert payload["source_url"] == "https://tensor.art/models/827823520299086029"
@pytest.mark.asyncio
async def test_set_hf_url_canonicalises_modelscope_subpage(tmp_path, hf_env):
model_path = tmp_path / "ms_sub.safetensors"
await _write_plain_model(model_path, sha="f" * 64)
response = await HfHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://modelscope.cn/models/user/repo/summary",
}
)
)
assert response.status == 200
assert _json_payload(response)["source_url"] == "https://modelscope.cn/models/user/repo"
@pytest.mark.asyncio
async def test_set_hf_url_is_idempotent_for_modelscope(tmp_path, hf_env):
model_path = tmp_path / "ms_twice.safetensors"
await _write_plain_model(model_path, sha="1" * 64)
request = FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://modelscope.cn/models/user/repo",
}
)
await HfHandler().set_hf_url(request)
await HfHandler().set_hf_url(request)
# The second call short-circuits without rewriting the cache entry.
assert hf_env["cache_write"].await_count == 1
@pytest.mark.asyncio
async def test_set_hf_url_switching_source_clears_hf_alias(tmp_path, hf_env):
model_path = tmp_path / "switch.safetensors"
await _write_plain_model(model_path, sha="2" * 64)
await HfHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://huggingface.co/user/repo",
}
)
)
await HfHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://modelscope.cn/models/user/repo",
}
)
)
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["source_platform"] == "modelscope"
assert saved.get("hf_url", "") == ""
@pytest.mark.asyncio
async def test_get_model_sources_lists_capabilities():
response = await HfHandler().get_model_sources(FakeRequest())
sources = _json_payload(response)
by_platform = {s["platform"]: s for s in sources}
assert set(by_platform) == {"huggingface", "modelscope", "tensorart"}
assert by_platform["huggingface"]["supports_enrichment"] is True
assert by_platform["modelscope"]["supports_enrichment"] is True
# TensorArt is link-only: no accessible model card for the backend.
assert by_platform["tensorart"]["supports_enrichment"] is False
assert by_platform["modelscope"]["supports_download"] is False
assert all(s["example_url"] for s in sources)
+635
View File
@@ -0,0 +1,635 @@
"""Tests for the external model-source handlers.
Covers linking (``set_hf_url``), file listing and downloads across the
registered platforms (Hugging Face / ModelScope).
Regression coverage for issue #1094: linking a model to HuggingFace must not
clear its CivitAI provenance or metadata, so both "View on CivitAI" and
"View on Hugging Face" can coexist.
"""
from __future__ import annotations
import json
import os
from types import SimpleNamespace
from typing import Any
from unittest.mock import AsyncMock
import pytest
from py.routes.handlers import model_source_handlers
from py.routes.handlers.model_source_handlers import ModelSourceHandler
from py.services.model_sources import ModelSourceError, SourceRef
from py.services.service_registry import ServiceRegistry
from py.utils.models import LoraMetadata
from py.utils.metadata_manager import MetadataManager
def _json_payload(response) -> dict[str, Any]:
assert response.text is not None
return json.loads(response.text)
class FakeRequest:
def __init__(self, *, json_data=None, query=None):
self._json_data = json_data or {}
self.query = query or {}
async def json(self):
return self._json_data
def _sidecar_path(model_path) -> str:
return f"{os.path.splitext(str(model_path))[0]}.metadata.json"
@pytest.fixture
def source_env(tmp_path, monkeypatch):
"""Point HF linking at *tmp_path* and stub the scanner cache write."""
monkeypatch.setattr(model_source_handlers, "_find_matching_root", lambda _dir: str(tmp_path))
cache_write = AsyncMock()
monkeypatch.setattr(model_source_handlers, "_add_to_scanner_cache", cache_write)
return {"root": tmp_path, "cache_write": cache_write}
async def _write_model(model_path, payload: dict[str, Any]) -> None:
model_path.write_bytes(b"x" * 32)
await MetadataManager.save_metadata(str(model_path), payload)
@pytest.mark.asyncio
async def test_set_hf_url_keeps_civitai_metadata_and_provenance(tmp_path, source_env):
model_path = tmp_path / "civitai_model.safetensors"
await _write_model(
model_path,
{
"file_name": "civitai_model",
"model_name": "CivitAI Model",
"file_path": str(model_path),
"size": 32,
"modified": 1.0,
"sha256": "a" * 64,
"base_model": "SDXL 1.0",
"preview_url": "",
"from_civitai": True,
"civitai": {"id": 111, "modelId": 222, "name": "v1", "trainedWords": []},
},
)
response = await ModelSourceHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"hf_url": "https://huggingface.co/user/repo",
}
)
)
assert response.status == 200
assert _json_payload(response)["success"] is True
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["hf_url"] == "https://huggingface.co/user/repo"
# Linking HF must not erase the model's CivitAI provenance or data.
assert saved["from_civitai"] is True
assert saved["civitai"]["modelId"] == 222
assert saved["civitai"]["id"] == 111
source_env["cache_write"].assert_awaited_once()
cached_metadata = source_env["cache_write"].await_args.args[1]
assert cached_metadata["hf_url"] == "https://huggingface.co/user/repo"
assert cached_metadata["from_civitai"] is True
assert cached_metadata["civitai"]["modelId"] == 222
@pytest.mark.asyncio
async def test_set_hf_url_does_not_force_from_civitai_false(tmp_path, source_env):
"""A model without CivitAI data keeps its existing provenance flag."""
model_path = tmp_path / "hf_only.safetensors"
await _write_model(
model_path,
{
"file_name": "hf_only",
"model_name": "HF Only",
"file_path": str(model_path),
"size": 32,
"modified": 1.0,
"sha256": "b" * 64,
"base_model": "Unknown",
"preview_url": "",
"from_civitai": True,
},
)
response = await ModelSourceHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"hf_url": "https://huggingface.co/user/repo",
}
)
)
assert response.status == 200
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["hf_url"] == "https://huggingface.co/user/repo"
assert saved["from_civitai"] is True
@pytest.mark.asyncio
async def test_set_hf_url_rejects_non_repo_url(tmp_path, source_env):
model_path = tmp_path / "model.safetensors"
await _write_model(
model_path,
{
"file_name": "model",
"model_name": "model",
"file_path": str(model_path),
"size": 32,
"modified": 1.0,
"sha256": "c" * 64,
"base_model": "Unknown",
"preview_url": "",
},
)
response = await ModelSourceHandler().set_hf_url(
FakeRequest(json_data={"file_path": str(model_path), "hf_url": "https://example.com/x"})
)
assert response.status == 400
payload = _json_payload(response)
assert payload["success"] is False
source_env["cache_write"].assert_not_awaited()
# ---------------------------------------------------------------------------
# Multi-source linking (ModelScope / TensorArt)
# ---------------------------------------------------------------------------
async def _write_plain_model(model_path, sha: str = "d" * 64) -> None:
await _write_model(
model_path,
{
"file_name": "model",
"model_name": "model",
"file_path": str(model_path),
"size": 32,
"modified": 1.0,
"sha256": sha,
"base_model": "Unknown",
"preview_url": "",
},
)
@pytest.mark.asyncio
async def test_set_hf_url_accepts_modelscope_and_stores_source_fields(tmp_path, source_env):
model_path = tmp_path / "ms_model.safetensors"
await _write_plain_model(model_path)
response = await ModelSourceHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://modelscope.cn/models/jj3550945163/Krea-2-LORA",
}
)
)
assert response.status == 200
payload = _json_payload(response)
assert payload["source_platform"] == "modelscope"
assert payload["source_url"] == "https://modelscope.cn/models/jj3550945163/Krea-2-LORA"
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["source_platform"] == "modelscope"
assert saved["source_url"] == "https://modelscope.cn/models/jj3550945163/Krea-2-LORA"
# No stale Hugging Face alias for a ModelScope model.
assert saved.get("hf_url", "") == ""
cached_metadata = source_env["cache_write"].await_args.args[1]
assert cached_metadata["source_platform"] == "modelscope"
@pytest.mark.asyncio
async def test_set_hf_url_accepts_tensorart_url(tmp_path, source_env):
model_path = tmp_path / "ta_model.safetensors"
await _write_plain_model(model_path, sha="e" * 64)
response = await ModelSourceHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": (
"https://tensor.art/models/827823520299086029/"
"Vivid-Impressions-Storybook-Sstyle-V1.0"
),
}
)
)
assert response.status == 200
payload = _json_payload(response)
assert payload["source_platform"] == "tensorart"
# The canonical page URL is stored, without the slug.
assert payload["source_url"] == "https://tensor.art/models/827823520299086029"
@pytest.mark.asyncio
async def test_set_hf_url_canonicalises_modelscope_subpage(tmp_path, source_env):
model_path = tmp_path / "ms_sub.safetensors"
await _write_plain_model(model_path, sha="f" * 64)
response = await ModelSourceHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://modelscope.cn/models/user/repo/summary",
}
)
)
assert response.status == 200
assert _json_payload(response)["source_url"] == "https://modelscope.cn/models/user/repo"
@pytest.mark.asyncio
async def test_set_hf_url_is_idempotent_for_modelscope(tmp_path, source_env):
model_path = tmp_path / "ms_twice.safetensors"
await _write_plain_model(model_path, sha="1" * 64)
request = FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://modelscope.cn/models/user/repo",
}
)
await ModelSourceHandler().set_hf_url(request)
await ModelSourceHandler().set_hf_url(request)
# The second call short-circuits without rewriting the cache entry.
assert source_env["cache_write"].await_count == 1
@pytest.mark.asyncio
async def test_set_hf_url_switching_source_clears_hf_alias(tmp_path, source_env):
model_path = tmp_path / "switch.safetensors"
await _write_plain_model(model_path, sha="2" * 64)
await ModelSourceHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://huggingface.co/user/repo",
}
)
)
await ModelSourceHandler().set_hf_url(
FakeRequest(
json_data={
"file_path": str(model_path),
"source_url": "https://modelscope.cn/models/user/repo",
}
)
)
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["source_platform"] == "modelscope"
assert saved.get("hf_url", "") == ""
@pytest.mark.asyncio
async def test_get_model_sources_lists_capabilities():
response = await ModelSourceHandler().get_model_sources(FakeRequest())
sources = _json_payload(response)
by_platform = {s["platform"]: s for s in sources}
assert set(by_platform) == {"huggingface", "modelscope", "tensorart"}
assert by_platform["huggingface"]["supports_enrichment"] is True
assert by_platform["modelscope"]["supports_enrichment"] is True
# TensorArt is link-only: no accessible model card for the backend.
assert by_platform["tensorart"]["supports_enrichment"] is False
assert by_platform["modelscope"]["supports_download"] is True
assert by_platform["modelscope"]["default_revision"] == "master"
assert by_platform["tensorart"]["supports_download"] is False
assert all(s["example_url"] for s in sources)
# ---------------------------------------------------------------------------
# File listing
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_list_model_source_files_returns_provider_result(monkeypatch):
captured: dict = {}
async def fake_list_files(self, source_id, revision=""):
captured["source_id"] = source_id
captured["revision"] = revision
return [{"filename": "a.safetensors", "size": 10}]
monkeypatch.setattr(
"py.services.model_sources.modelscope.ModelScopeSource.list_files",
fake_list_files,
)
response = await ModelSourceHandler().list_model_source_files(
FakeRequest(
query={
"platform": "modelscope",
"repo": "jj3550945163/Krea-2-LORA",
"revision": "v1",
}
)
)
assert response.status == 200
assert _json_payload(response) == [{"filename": "a.safetensors", "size": 10}]
assert captured == {"source_id": "jj3550945163/Krea-2-LORA", "revision": "v1"}
@pytest.mark.asyncio
async def test_list_model_source_files_rejects_link_only_platform():
response = await ModelSourceHandler().list_model_source_files(
FakeRequest(query={"platform": "tensorart", "repo": "u/r"})
)
assert response.status == 400
assert "does not support downloads" in _json_payload(response)["error"]
@pytest.mark.asyncio
async def test_list_model_source_files_rejects_unsafe_repo():
for repo in ("noslash", "../etc/passwd", "u/.."):
response = await ModelSourceHandler().list_model_source_files(
FakeRequest(query={"platform": "modelscope", "repo": repo})
)
assert response.status == 400, repo
assert "repo" in _json_payload(response)["error"]
@pytest.mark.asyncio
async def test_list_model_source_files_maps_missing_repo_to_404(monkeypatch):
async def fake_list_files(self, source_id, revision=""):
raise ModelSourceError(f"Repository '{source_id}' not found", status=404)
monkeypatch.setattr(
"py.services.model_sources.modelscope.ModelScopeSource.list_files",
fake_list_files,
)
response = await ModelSourceHandler().list_model_source_files(
FakeRequest(query={"platform": "modelscope", "repo": "u/r"})
)
assert response.status == 404
assert "not found" in _json_payload(response)["error"]
@pytest.mark.asyncio
async def test_list_model_source_files_maps_transport_failure_to_502(monkeypatch):
async def fake_list_files(self, source_id, revision=""):
raise ModelSourceError("upstream exploded", status=502)
monkeypatch.setattr(
"py.services.model_sources.modelscope.ModelScopeSource.list_files",
fake_list_files,
)
response = await ModelSourceHandler().list_model_source_files(
FakeRequest(query={"platform": "modelscope", "repo": "u/r"})
)
assert response.status == 502
# ---------------------------------------------------------------------------
# Downloads
# ---------------------------------------------------------------------------
def _stub_download_backend(monkeypatch) -> dict:
"""Replace the downloader/settings plumbing with a recording stub."""
captured: dict = {}
async def fake_download_file(**kwargs):
captured.update(kwargs)
return True, kwargs["save_path"]
class _Downloader:
download_file = staticmethod(fake_download_file)
async def fake_get_downloader():
return _Downloader()
class _Settings:
def get(self, key, default=None):
return default
monkeypatch.setattr(model_source_handlers, "get_downloader", fake_get_downloader)
monkeypatch.setattr(
model_source_handlers, "get_settings_manager", lambda: _Settings()
)
return captured
@pytest.mark.asyncio
async def test_download_model_source_modelscope_uses_resolve_url(tmp_path, monkeypatch):
captured = _stub_download_backend(monkeypatch)
saved = AsyncMock()
monkeypatch.setattr(model_source_handlers, "_save_source_metadata", saved)
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "modelscope",
"repo": "jj3550945163/Krea-2-LORA",
"filename": "Krea-2-LORA_c1-st1000.safetensors",
"model_root": str(tmp_path),
}
)
)
assert response.status == 200
assert captured["url"] == (
"https://modelscope.cn/models/jj3550945163/Krea-2-LORA/resolve/master/"
"Krea-2-LORA_c1-st1000.safetensors"
)
assert captured["save_path"] == str(tmp_path / "Krea-2-LORA_c1-st1000.safetensors")
ref = saved.await_args.args[1]
assert ref.platform == "modelscope"
assert ref.source_id == "jj3550945163/Krea-2-LORA"
assert ref.url == "https://modelscope.cn/models/jj3550945163/Krea-2-LORA"
@pytest.mark.asyncio
async def test_download_model_source_modelscope_default_paths(tmp_path, monkeypatch):
captured = _stub_download_backend(monkeypatch)
saved = AsyncMock()
monkeypatch.setattr(model_source_handlers, "_save_source_metadata", saved)
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "modelscope",
"repo": "owner/name",
"filename": "nested/model.safetensors",
"model_root": str(tmp_path),
"use_default_paths": True,
}
)
)
assert response.status == 200
# The site gets its own sub-directory, mirroring `huggingface/<owner>/<repo>`.
assert captured["save_path"] == str(
tmp_path / "modelscope" / "owner" / "name" / "model.safetensors"
)
@pytest.mark.asyncio
async def test_download_model_source_defaults_to_huggingface(tmp_path, monkeypatch):
"""The legacy /api/lm/download-hf-model payload has no `platform` key."""
captured = _stub_download_backend(monkeypatch)
monkeypatch.setattr(model_source_handlers, "_save_source_metadata", AsyncMock())
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"repo": "user/repo",
"filename": "f.safetensors",
"revision": "main",
"model_root": str(tmp_path),
}
)
)
assert response.status == 200
assert captured["url"] == (
"https://huggingface.co/user/repo/resolve/main/f.safetensors"
)
@pytest.mark.asyncio
async def test_download_model_source_rejects_link_only_platform(tmp_path):
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "tensorart",
"repo": "u/r",
"filename": "f.safetensors",
"model_root": str(tmp_path),
}
)
)
assert response.status == 400
assert "does not support downloads" in _json_payload(response)["error"]
@pytest.mark.asyncio
async def test_download_model_source_rejects_unsafe_input(tmp_path, monkeypatch):
_stub_download_backend(monkeypatch)
cases = [
({"repo": "noslash", "filename": "f.safetensors"}, "repo format"),
({"repo": "u/r", "filename": "../../etc/passwd"}, "Invalid filename"),
(
{"repo": "u/r", "filename": "f.safetensors", "relative_path": "/abs"},
"relative_path must not be absolute",
),
(
{"repo": "u/r", "filename": "f.safetensors", "relative_path": "../up"},
"Invalid relative_path",
),
]
for extra, expected in cases:
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "modelscope",
"model_root": str(tmp_path),
**extra,
}
)
)
assert response.status == 400, extra
assert expected in _json_payload(response)["error"], extra
@pytest.mark.asyncio
async def test_download_model_source_skips_existing_file(tmp_path, monkeypatch):
captured = _stub_download_backend(monkeypatch)
monkeypatch.setattr(model_source_handlers, "_save_source_metadata", AsyncMock())
(tmp_path / "f.safetensors").write_bytes(b"already here")
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "modelscope",
"repo": "u/r",
"filename": "f.safetensors",
"model_root": str(tmp_path),
}
)
)
assert response.status == 200
assert "already exists" in _json_payload(response)["message"]
assert captured == {}
@pytest.mark.asyncio
@pytest.mark.parametrize(
("platform", "url", "expect_hf_alias"),
[
("modelscope", "https://modelscope.cn/models/u/r", False),
("huggingface", "https://huggingface.co/u/r", True),
],
)
async def test_save_source_metadata_writes_platform_fields(
tmp_path, monkeypatch, platform, url, expect_hf_alias
):
"""A download's sidecar must record its own platform (and no stale HF alias)."""
model_path = tmp_path / "downloaded.safetensors"
model_path.write_bytes(b"x" * 32)
metadata = LoraMetadata(
file_name="downloaded",
model_name="Downloaded",
file_path=str(model_path),
size=32,
modified=1.0,
sha256="a" * 64,
base_model="SDXL 1.0",
preview_url="",
)
monkeypatch.setattr(
model_source_handlers.MetadataManager,
"create_default_metadata",
AsyncMock(return_value=metadata),
)
scanner = SimpleNamespace(add_model_to_cache=AsyncMock())
monkeypatch.setattr(
ServiceRegistry, "get_lora_scanner", AsyncMock(return_value=scanner)
)
monkeypatch.setattr(
model_source_handlers, "_infer_model_type", lambda _root: (LoraMetadata, "get_lora_scanner")
)
ref = SourceRef(platform=platform, source_id="u/r", url=url)
await model_source_handlers._save_source_metadata(str(model_path), ref, str(tmp_path))
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["source_platform"] == platform
assert saved["source_url"] == url
assert bool(saved.get("hf_url", "")) is expect_hf_alias
cached = scanner.add_model_to_cache.await_args.args[0]
assert cached["source_platform"] == platform
assert cached["source_url"] == url
+184 -2
View File
@@ -14,10 +14,14 @@ from py.services.model_sources import (
ModelScopeSource, ModelScopeSource,
TensorArtSource, TensorArtSource,
detect_source, detect_source,
downloadable_sources,
get_download_source,
get_source, get_source,
get_source_platform, get_source_platform,
has_external_source, has_external_source,
is_valid_source_id,
list_sources, list_sources,
ModelSourceError,
normalize_metadata_source, normalize_metadata_source,
resolve_source_ref, resolve_source_ref,
source_group_key, source_group_key,
@@ -127,11 +131,15 @@ class TestCapabilities:
source = get_source("huggingface") source = get_source("huggingface")
assert source.supports_enrichment is True assert source.supports_enrichment is True
assert source.supports_download is True assert source.supports_download is True
assert source.default_revision == "main"
assert source.default_subdir == "huggingface"
def test_modelscope_supports_enrichment_but_not_download(self): def test_modelscope_supports_enrichment_and_download(self):
source = get_source("modelscope") source = get_source("modelscope")
assert source.supports_enrichment is True assert source.supports_enrichment is True
assert source.supports_download is False assert source.supports_download is True
assert source.default_revision == "master"
assert source.default_subdir == "modelscope"
def test_tensorart_is_link_only(self): def test_tensorart_is_link_only(self):
source = get_source("tensorart") source = get_source("tensorart")
@@ -322,3 +330,177 @@ class TestAssetBaseUrl:
ModelScopeSource().asset_base_url("u/r") ModelScopeSource().asset_base_url("u/r")
== "https://modelscope.cn/models/u/r/resolve/master" == "https://modelscope.cn/models/u/r/resolve/master"
) )
# ---------------------------------------------------------------------------
# Download support
# ---------------------------------------------------------------------------
class TestListFiles:
@pytest.mark.asyncio
async def test_huggingface_reads_tree_api_with_lfs_sizes(self, monkeypatch):
captured: dict = {}
async def fake_fetch_json(url, **_kwargs):
captured["url"] = url
return 200, [
{"path": "README.md", "size": 120},
{"path": "a/model.safetensors", "size": 300},
{"path": "b.safetensors", "size": 0, "lfs": {"size": 200}},
]
monkeypatch.setattr(
"py.services.model_sources.huggingface.fetch_json", fake_fetch_json
)
files = await HuggingFaceSource().list_files("u/r")
assert captured["url"] == "https://huggingface.co/api/models/u/r/tree/main"
assert files == [
{"filename": "a/model.safetensors", "size": 300},
{"filename": "b.safetensors", "size": 200},
]
@pytest.mark.asyncio
async def test_huggingface_honours_explicit_revision(self, monkeypatch):
captured: dict = {}
async def fake_fetch_json(url, **_kwargs):
captured["url"] = url
return 200, []
monkeypatch.setattr(
"py.services.model_sources.huggingface.fetch_json", fake_fetch_json
)
await HuggingFaceSource().list_files("u/r", "v2.0")
assert captured["url"].endswith("/tree/v2.0")
@pytest.mark.asyncio
async def test_modelscope_reads_repo_files_api(self, monkeypatch):
captured: dict = {}
async def fake_fetch_json(url, **_kwargs):
captured["url"] = url
return 200, {
"Data": {
"Files": [
# directories are listed too and must be dropped
{"Type": "tree", "Path": "vae", "Size": 0},
{"Type": "blob", "Path": "README.md", "Size": 100},
{"Type": "blob", "Path": "sub/model.safetensors", "Size": 500},
{"Type": "blob", "Path": "model.ckpt", "Size": 200},
]
}
}
monkeypatch.setattr(
"py.services.model_sources.modelscope.fetch_json", fake_fetch_json
)
files = await ModelScopeSource().list_files("u/r")
assert captured["url"] == (
"https://modelscope.cn/api/v1/models/u/r/repo/files?Revision=master"
)
assert files == [
{"filename": "sub/model.safetensors", "size": 500},
{"filename": "model.ckpt", "size": 200},
]
@pytest.mark.asyncio
async def test_missing_repo_is_reported_as_not_found(self, monkeypatch):
async def fake_fetch_json(url, **_kwargs):
return 404, None
monkeypatch.setattr(
"py.services.model_sources.modelscope.fetch_json", fake_fetch_json
)
with pytest.raises(ModelSourceError) as excinfo:
await ModelScopeSource().list_files("u/r")
assert excinfo.value.status == 404
assert "not found" in str(excinfo.value)
@pytest.mark.asyncio
async def test_transport_failure_is_reported_as_bad_gateway(self, monkeypatch):
async def fake_fetch_json(url, **_kwargs):
return 0, None
monkeypatch.setattr(
"py.services.model_sources.huggingface.fetch_json", fake_fetch_json
)
with pytest.raises(ModelSourceError) as excinfo:
await HuggingFaceSource().list_files("u/r")
assert excinfo.value.status == 502
class TestDownloadUrls:
def test_huggingface_resolve_url(self):
assert HuggingFaceSource().file_download_url("u/r", "sub/f.safetensors") == (
"https://huggingface.co/u/r/resolve/main/sub/f.safetensors"
)
def test_modelscope_resolve_url_defaults_to_master(self):
assert ModelScopeSource().file_download_url("u/r", "sub/f.safetensors") == (
"https://modelscope.cn/models/u/r/resolve/master/sub/f.safetensors"
)
def test_explicit_revision_wins(self):
assert ModelScopeSource().file_download_url("u/r", "f.bin", "v1") == (
"https://modelscope.cn/models/u/r/resolve/v1/f.bin"
)
def test_tensorart_refuses_to_build_a_download_url(self):
source = TensorArtSource()
assert source.supports_download is False
with pytest.raises(ModelSourceError):
source.file_download_url("123", "f.safetensors")
@pytest.mark.asyncio
async def test_tensorart_lists_nothing(self):
assert await TensorArtSource().list_files("123") == []
class TestSourceIdValidation:
@pytest.mark.parametrize(
"source_id",
["u/r", "black-forest-labs/FLUX.1-dev", "AI-ModelScope/stable-diffusion-v1-5"],
)
def test_accepts_repo_ids(self, source_id):
assert is_valid_source_id(source_id) is True
@pytest.mark.parametrize(
"source_id",
[
"",
"noslash",
"a/b/c",
"../etc/passwd",
"u/..",
"u/.",
".hidden/r",
"u/r with space",
"/r",
"u/",
],
)
def test_rejects_unsafe_ids(self, source_id):
assert is_valid_source_id(source_id) is False
class TestDownloadSourceRegistry:
def test_downloadable_sources_excludes_link_only_sites(self):
platforms = {source.platform for source in downloadable_sources()}
assert platforms == {"huggingface", "modelscope"}
def test_get_download_source_rejects_link_only_platform(self):
assert get_download_source("tensorart") is None
assert get_download_source("nope") is None
assert get_download_source("modelscope").platform == "modelscope"
assert get_download_source("huggingface").platform == "huggingface"