mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-20 18:51:26 -03:00
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:
@@ -71,7 +71,7 @@ Enriches models linked to an external model site with metadata extracted by an L
|
||||
| Platform | Link | AI enrichment | Direct download |
|
||||
| --- | --- | --- | --- |
|
||||
| Hugging Face | yes | yes | yes |
|
||||
| ModelScope | yes | yes | no |
|
||||
| ModelScope | yes | yes | yes |
|
||||
| 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.
|
||||
|
||||
+4
-4
@@ -1395,9 +1395,9 @@
|
||||
"download": {
|
||||
"title": "Download Model from URL",
|
||||
"titleWithType": "Download {type} from URL",
|
||||
"civitaiUrl": "CivitAI URL(s):",
|
||||
"civitaiUrl": "Model URL(s):",
|
||||
"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:",
|
||||
"selectAll": "Select All",
|
||||
"fetchingRepoFiles": "Fetching repository files...",
|
||||
@@ -1430,9 +1430,9 @@
|
||||
"inLibrary": "In Library"
|
||||
},
|
||||
"errors": {
|
||||
"invalidUrl": "Invalid CivitAI URL format",
|
||||
"invalidUrl": "Invalid model URL format",
|
||||
"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."
|
||||
},
|
||||
"status": {
|
||||
|
||||
@@ -472,12 +472,5 @@ class LoraManager:
|
||||
scanner.cancel_task()
|
||||
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:
|
||||
logger.error(f"Error during cleanup: {e}", exc_info=True)
|
||||
|
||||
@@ -55,7 +55,7 @@ from ...utils.constants import (
|
||||
VALID_LORA_TYPES,
|
||||
VALID_OTHER_CIVITAI_TYPES,
|
||||
)
|
||||
from .hf_handlers import HfHandler
|
||||
from .model_source_handlers import ModelSourceHandler
|
||||
from .agent_handlers import AgentHandler
|
||||
from .download_routing_handlers import DownloadRoutingHandler
|
||||
from .model_handlers import ModelCivitaiHandler
|
||||
@@ -4001,7 +4001,7 @@ class MiscHandlerSet:
|
||||
doctor: DoctorHandler,
|
||||
example_workflows: ExampleWorkflowsHandler,
|
||||
base_model: BaseModelHandlerSet,
|
||||
hf_handler: Any = None,
|
||||
model_source_handler: Any = None,
|
||||
agent_handler: Any = None,
|
||||
download_routing: Any = None,
|
||||
) -> None:
|
||||
@@ -4022,7 +4022,7 @@ class MiscHandlerSet:
|
||||
self.doctor = doctor
|
||||
self.example_workflows = example_workflows
|
||||
self.base_model = base_model
|
||||
self.hf_handler = hf_handler
|
||||
self.model_source_handler = model_source_handler
|
||||
self.agent_handler = agent_handler
|
||||
self.download_routing = download_routing
|
||||
|
||||
@@ -4076,10 +4076,13 @@ class MiscHandlerSet:
|
||||
"get_example_workflows": self.example_workflows.get_example_workflows,
|
||||
"get_example_workflow": self.example_workflows.get_example_workflow,
|
||||
# Hugging Face handlers
|
||||
"get_hf_repo_files": self.hf_handler.get_hf_repo_files,
|
||||
"download_hf_model": self.hf_handler.download_hf_model,
|
||||
"set_hf_url": self.hf_handler.set_hf_url,
|
||||
"get_model_sources": self.hf_handler.get_model_sources,
|
||||
# External model sources (Hugging Face / ModelScope)
|
||||
"list_model_source_files": self.model_source_handler.list_model_source_files,
|
||||
"download_model_source": self.model_source_handler.download_model_source,
|
||||
"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
|
||||
"get_agent_skills": self.agent_handler.get_agent_skills,
|
||||
"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
|
||||
listing and the project's existing aiohttp-based Downloader for
|
||||
downloading. No huggingface_hub dependency required.
|
||||
Covers every site registered in :mod:`py.services.model_sources`. The module
|
||||
was Hugging Face only (``hf_handlers.py`` / ``HfHandler``) until ModelScope
|
||||
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
|
||||
@@ -10,10 +15,8 @@ from __future__ import annotations
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
|
||||
from ...config import config
|
||||
@@ -23,14 +26,17 @@ from ...services.downloader import (
|
||||
)
|
||||
from ...services.aria2_downloader import Aria2Downloader
|
||||
from ...services.model_sources import (
|
||||
ModelSourceError,
|
||||
SourceRef,
|
||||
detect_source,
|
||||
get_download_source,
|
||||
is_valid_source_id,
|
||||
list_sources,
|
||||
normalize_metadata_source,
|
||||
)
|
||||
from ...services.settings_manager import get_settings_manager
|
||||
from ...services.service_registry import ServiceRegistry
|
||||
from ...services.websocket_manager import ws_manager
|
||||
from ...utils.constants import MODEL_FILE_EXTENSIONS
|
||||
from ...utils.metadata_manager import MetadataManager
|
||||
from ...utils.models import LoraMetadata, CheckpointMetadata, EmbeddingMetadata
|
||||
|
||||
@@ -39,28 +45,6 @@ logger = logging.getLogger(__name__)
|
||||
_DEFAULT_MODEL_CLASS = LoraMetadata
|
||||
_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]:
|
||||
"""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
|
||||
|
||||
|
||||
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.
|
||||
|
||||
Uses ``MetadataManager.create_default_metadata()`` which computes the
|
||||
SHA256 hash, extracts safetensors header metadata (base_model), and
|
||||
produces a fully-populated ``LoraMetadata`` (or ``CheckpointMetadata`` /
|
||||
``EmbeddingMetadata``) object. We then overlay HF-specific fields and
|
||||
register the model in the in-memory scanner cache so it appears
|
||||
``EmbeddingMetadata``) object. We then overlay the external-source fields
|
||||
and register the model in the in-memory scanner cache so it appears
|
||||
immediately without a full filesystem walk.
|
||||
"""
|
||||
try:
|
||||
hf_url = f"https://huggingface.co/{repo}"
|
||||
model_class, scanner_getter_name = _infer_model_type(model_root)
|
||||
|
||||
# 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)
|
||||
return
|
||||
|
||||
# 2. Overlay HF-specific fields
|
||||
metadata._unknown_fields["hf_url"] = hf_url
|
||||
metadata._unknown_fields["source_url"] = hf_url
|
||||
metadata._unknown_fields["source_platform"] = "huggingface"
|
||||
metadata.from_civitai = False # HF models are not from CivitAI
|
||||
# 2. Overlay the external-source fields (`hf_url` is written by
|
||||
# normalisation for Hugging Face only)
|
||||
fields = metadata._unknown_fields
|
||||
fields["source_url"] = ref.url
|
||||
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
|
||||
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
|
||||
# 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:
|
||||
scanner = await scanner_getter()
|
||||
if scanner is not None:
|
||||
metadata_dict = metadata.to_dict()
|
||||
metadata_dict["hf_url"] = hf_url
|
||||
metadata_dict = normalize_metadata_source(metadata.to_dict())
|
||||
await scanner.add_model_to_cache(metadata_dict, folder)
|
||||
logger.info("Added %s to scanner cache (folder=%s)", dest_path, folder)
|
||||
|
||||
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:
|
||||
@@ -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)
|
||||
|
||||
|
||||
class HfHandler:
|
||||
"""Handle Hugging Face model browsing and download."""
|
||||
def _unsupported_platform_error(platform: str) -> web.Response:
|
||||
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:
|
||||
"""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
|
||||
explain which sites support AI metadata enrichment.
|
||||
Used by the "Link Model" dialog to validate URLs client-side, to
|
||||
explain which sites support AI metadata enrichment, and to pick the
|
||||
right download endpoint/revision.
|
||||
"""
|
||||
|
||||
return web.json_response([
|
||||
@@ -209,6 +208,7 @@ class HfHandler:
|
||||
"label": source.label,
|
||||
"supports_enrichment": source.supports_enrichment,
|
||||
"supports_download": source.supports_download,
|
||||
"default_revision": source.default_revision,
|
||||
"example_url": source.canonical_url(
|
||||
"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:
|
||||
"""Link a model file to its page on an external model site.
|
||||
|
||||
Accepts ``source_url`` (preferred) or the legacy ``hf_url`` /
|
||||
``url`` payload key. Hugging Face, ModelScope, and TensorArt URLs
|
||||
are recognised; the platform is stored alongside the canonical URL.
|
||||
TensorArt models can be linked and browsed, but not AI-enriched.
|
||||
Accepts ``source_url`` (preferred) or the legacy ``hf_url`` / ``url``
|
||||
payload key. Every registered site is recognised and the platform is
|
||||
stored alongside the canonical URL. TensorArt models can be linked and
|
||||
browsed, but not AI-enriched.
|
||||
|
||||
The route path keeps its historical ``set-hf-url`` name.
|
||||
"""
|
||||
|
||||
try:
|
||||
@@ -337,74 +339,60 @@ class HfHandler:
|
||||
status=500,
|
||||
)
|
||||
|
||||
async def get_hf_repo_files(self, request: web.Request) -> web.Response:
|
||||
"""List model-weight files from a HF repo with real file sizes.
|
||||
async def list_model_source_files(self, request: web.Request) -> web.Response:
|
||||
"""List the downloadable weight files of an external repository.
|
||||
|
||||
Uses the HF tree API endpoint which returns accurate file sizes
|
||||
(including LFS-tracked files), unlike the model info endpoint.
|
||||
Query params: ``platform``, ``repo`` (``owner/name``), ``revision``
|
||||
(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(
|
||||
{"error": "Missing or invalid 'repo' parameter (expected user/repo)"},
|
||||
{"error": "Missing or invalid 'repo' parameter (expected owner/name)"},
|
||||
status=400,
|
||||
)
|
||||
|
||||
url = f"https://huggingface.co/api/models/{repo}/tree/main"
|
||||
|
||||
try:
|
||||
session = await _get_hf_api_session()
|
||||
async with session.get(url) as resp:
|
||||
if resp.status == 404:
|
||||
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()
|
||||
files = await source.list_files(repo, revision)
|
||||
except ModelSourceError as exc:
|
||||
return web.json_response({"error": str(exc)}, status=exc.status)
|
||||
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)
|
||||
|
||||
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)
|
||||
|
||||
async def download_hf_model(self, request: web.Request) -> web.Response:
|
||||
"""Download a single file from Hugging Face into the model directory.
|
||||
async def download_model_source(self, request: web.Request) -> web.Response:
|
||||
"""Download a single file from an external repository.
|
||||
|
||||
POST JSON body::
|
||||
|
||||
{
|
||||
"repo": "dx8152/Flux2-Klein-9B-Consistency",
|
||||
"filename": "Flux2-Klein-9B-consistency-V2.safetensors",
|
||||
"revision": "main",
|
||||
"platform": "modelscope",
|
||||
"repo": "owner/name",
|
||||
"filename": "subdir/model.safetensors",
|
||||
"revision": "master",
|
||||
"model_root": "loras",
|
||||
"relative_path": "",
|
||||
"use_default_paths": false,
|
||||
"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,
|
||||
percentage) is broadcast via the WebSocket progress system, matching
|
||||
the CivitAI download experience.
|
||||
percentage) is broadcast via the WebSocket progress system.
|
||||
|
||||
Respects the ``download_backend`` setting (``aria2`` or ``default``).
|
||||
"""
|
||||
@@ -413,30 +401,33 @@ class HfHandler:
|
||||
except json.JSONDecodeError:
|
||||
return web.json_response({"error": "Invalid JSON"}, status=400)
|
||||
|
||||
platform = (payload.get("platform") or "huggingface").strip()
|
||||
repo = (payload.get("repo") 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()
|
||||
relative_path = (payload.get("relative_path") or "").strip()
|
||||
use_default_paths = bool(payload.get("use_default_paths", False))
|
||||
download_id: str | None = payload.get("download_id")
|
||||
|
||||
logger.info(
|
||||
"download_hf_model: repo=%s file=%s root=%s download_id=%s",
|
||||
repo, filename, model_root, download_id,
|
||||
"download_model_source: platform=%s repo=%s file=%s root=%s download_id=%s",
|
||||
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:
|
||||
return web.json_response(
|
||||
{"error": "Missing required fields: 'repo' and 'filename'"}, status=400
|
||||
)
|
||||
|
||||
# Validate repo format — must be user/repo_name
|
||||
if repo.count("/") != 1 or not re.match(r"^[a-zA-Z0-9_.-]+/[a-zA-Z0-9_.-]+$", 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):
|
||||
# `owner/name` only; the components become path segments below.
|
||||
if not is_valid_source_id(repo):
|
||||
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
|
||||
if ".." in filename:
|
||||
@@ -455,21 +446,21 @@ class HfHandler:
|
||||
# unnecessary when the frontend sends the path from its own dropdown
|
||||
# (populated from scanner roots). Using the "business path" directly
|
||||
# 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):
|
||||
base_dir = os.path.normpath(model_root)
|
||||
else:
|
||||
base_dir = os.path.normpath(os.path.join(os.getcwd(), "models", model_root))
|
||||
|
||||
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:
|
||||
target_dir = os.path.join(base_dir, relative_path)
|
||||
else:
|
||||
target_dir = base_dir
|
||||
|
||||
# Strip HF repo subdirectory — "diffusion_models/xxx.safetensors"
|
||||
# is an HF repo convention, not meaningful for local storage.
|
||||
# Strip the repository sub-directory — "diffusion_models/xxx.safetensors"
|
||||
# is a repository convention, not meaningful for local storage.
|
||||
file_base = os.path.basename(filename)
|
||||
|
||||
os.makedirs(target_dir, exist_ok=True)
|
||||
@@ -477,16 +468,18 @@ class HfHandler:
|
||||
|
||||
# Check if already exists (simple skip)
|
||||
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({
|
||||
"success": True,
|
||||
"message": f"File already exists: {dest_path}",
|
||||
"path": dest_path,
|
||||
})
|
||||
|
||||
# Build HF resolve URL
|
||||
resolve_url = (
|
||||
f"https://huggingface.co/{repo}/resolve/{revision}/{filename}"
|
||||
# Built per request: sites that redirect to a CDN hand out a
|
||||
# time-limited token in the redirect, so the URL must never be cached.
|
||||
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
|
||||
@@ -528,28 +521,27 @@ class HfHandler:
|
||||
|
||||
if download_backend == "aria2":
|
||||
aria2 = await Aria2Downloader.get_instance()
|
||||
aid = download_id or f"hf_{repo}_{filename}"
|
||||
aid = download_id or f"{source.platform}_{repo}_{filename}"
|
||||
try:
|
||||
hf_success, hf_result = await aria2.download_file(
|
||||
ok, result = await aria2.download_file(
|
||||
url=resolve_url,
|
||||
save_path=dest_path,
|
||||
download_id=aid,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
if hf_success:
|
||||
await _save_hf_metadata(dest_path, repo, model_root)
|
||||
if ok:
|
||||
await _save_source_metadata(dest_path, ref, model_root)
|
||||
return web.json_response({
|
||||
"success": True,
|
||||
"message": f"Downloaded to {dest_path}",
|
||||
"path": dest_path,
|
||||
})
|
||||
else:
|
||||
return web.json_response(
|
||||
{"success": False, "error": hf_result or "aria2 download failed"},
|
||||
status=500,
|
||||
)
|
||||
return web.json_response(
|
||||
{"success": False, "error": result or "aria2 download failed"},
|
||||
status=500,
|
||||
)
|
||||
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(
|
||||
{"success": False, "error": str(exc)}, status=500
|
||||
)
|
||||
@@ -565,19 +557,18 @@ class HfHandler:
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
if success:
|
||||
await _save_hf_metadata(dest_path, repo, model_root)
|
||||
await _save_source_metadata(dest_path, ref, model_root)
|
||||
return web.json_response({
|
||||
"success": True,
|
||||
"message": f"Downloaded to {result}",
|
||||
"path": result,
|
||||
})
|
||||
else:
|
||||
return web.json_response(
|
||||
{"success": False, "error": result or "Download failed"},
|
||||
status=500,
|
||||
)
|
||||
return web.json_response(
|
||||
{"success": False, "error": result or "Download failed"},
|
||||
status=500,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error("HF download failed: %s", exc)
|
||||
logger.error("%s download failed: %s", platform, exc)
|
||||
return web.json_response(
|
||||
{"success": False, "error": str(exc)}, status=500
|
||||
)
|
||||
@@ -99,7 +99,11 @@ MISC_ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
|
||||
RouteDefinition(
|
||||
"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(
|
||||
"GET", "/api/lm/hf-repo-files", "get_hf_repo_files"
|
||||
),
|
||||
@@ -107,6 +111,9 @@ MISC_ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
|
||||
RouteDefinition(
|
||||
"POST", "/api/lm/download/routing", "get_download_routing"
|
||||
),
|
||||
RouteDefinition(
|
||||
"POST", "/api/lm/download-model-source", "download_model_source"
|
||||
),
|
||||
RouteDefinition(
|
||||
"POST", "/api/lm/download-hf-model", "download_hf_model"
|
||||
),
|
||||
|
||||
@@ -39,7 +39,7 @@ from .handlers.misc_handlers import (
|
||||
build_service_registry_adapter,
|
||||
)
|
||||
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.download_routing_handlers import DownloadRoutingHandler
|
||||
from .misc_route_registrar import MiscRouteRegistrar
|
||||
@@ -139,7 +139,7 @@ class MiscRoutes:
|
||||
doctor = DoctorHandler(settings_service=self._settings)
|
||||
example_workflows = ExampleWorkflowsHandler()
|
||||
base_model = BaseModelHandlerSet()
|
||||
hf_handler = HfHandler()
|
||||
model_source_handler = ModelSourceHandler()
|
||||
agent_handler = AgentHandler()
|
||||
download_routing = DownloadRoutingHandler()
|
||||
|
||||
@@ -161,7 +161,7 @@ class MiscRoutes:
|
||||
doctor=doctor,
|
||||
example_workflows=example_workflows,
|
||||
base_model=base_model,
|
||||
hf_handler=hf_handler,
|
||||
model_source_handler=model_source_handler,
|
||||
agent_handler=agent_handler,
|
||||
download_routing=download_routing,
|
||||
)
|
||||
|
||||
@@ -12,10 +12,14 @@ from .base import (
|
||||
GROUP_PREFIXES,
|
||||
HTTP_TIMEOUT,
|
||||
ModelSource,
|
||||
ModelSourceError,
|
||||
SourceRef,
|
||||
USER_AGENT,
|
||||
clean_source_url,
|
||||
fetch_json,
|
||||
fetch_text,
|
||||
filter_weight_files,
|
||||
is_valid_source_id,
|
||||
)
|
||||
from .huggingface import HuggingFaceSource
|
||||
from .modelscope import ModelScopeSource
|
||||
@@ -24,6 +28,8 @@ from .registry import (
|
||||
SOURCE_PLATFORM_FIELD,
|
||||
SOURCE_URL_FIELD,
|
||||
detect_source,
|
||||
downloadable_sources,
|
||||
get_download_source,
|
||||
get_source,
|
||||
get_source_platform,
|
||||
has_external_source,
|
||||
@@ -40,6 +46,7 @@ __all__ = [
|
||||
"HTTP_TIMEOUT",
|
||||
"LEGACY_HF_URL_FIELD",
|
||||
"ModelSource",
|
||||
"ModelSourceError",
|
||||
"HuggingFaceSource",
|
||||
"ModelScopeSource",
|
||||
"SOURCE_PLATFORM_FIELD",
|
||||
@@ -49,10 +56,15 @@ __all__ = [
|
||||
"USER_AGENT",
|
||||
"clean_source_url",
|
||||
"detect_source",
|
||||
"downloadable_sources",
|
||||
"fetch_json",
|
||||
"fetch_text",
|
||||
"filter_weight_files",
|
||||
"get_download_source",
|
||||
"get_source",
|
||||
"get_source_platform",
|
||||
"has_external_source",
|
||||
"is_valid_source_id",
|
||||
"list_sources",
|
||||
"normalize_metadata_source",
|
||||
"resolve_source_ref",
|
||||
|
||||
@@ -20,12 +20,15 @@ HTTP handlers never need site-specific branching.
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
from typing import Any, Iterable, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ...utils.constants import MODEL_FILE_EXTENSIONS
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#: Shared HTTP timeout for model-card fetches.
|
||||
@@ -58,6 +61,36 @@ class SourceRef:
|
||||
"""Canonical URL of the model page."""
|
||||
|
||||
|
||||
class ModelSourceError(Exception):
|
||||
"""Raised when a model source cannot satisfy a request.
|
||||
|
||||
Carries the HTTP status the API handler should answer with, so the
|
||||
handlers stay free of per-site error mapping.
|
||||
"""
|
||||
|
||||
def __init__(self, message: str, status: int = 502) -> None:
|
||||
super().__init__(message)
|
||||
self.status = status
|
||||
|
||||
|
||||
#: Repository ids are always exactly ``owner/name``. Components may contain
|
||||
#: dots (``black-forest-labs/FLUX.1-dev``) but must not be empty, ``.`` / ``..``,
|
||||
#: or start with a dot - the id is used as a path segment on disk.
|
||||
_SOURCE_ID_COMPONENT = re.compile(r"^[A-Za-z0-9_][A-Za-z0-9_.\-]*$")
|
||||
|
||||
|
||||
def is_valid_source_id(source_id: str) -> bool:
|
||||
"""Return ``True`` when *source_id* is a safe ``owner/name`` repository id."""
|
||||
|
||||
if not source_id or not isinstance(source_id, str) or source_id.count("/") != 1:
|
||||
return False
|
||||
owner, name = source_id.split("/", 1)
|
||||
return all(
|
||||
part and part not in (".", "..") and _SOURCE_ID_COMPONENT.match(part)
|
||||
for part in (owner, name)
|
||||
)
|
||||
|
||||
|
||||
async def fetch_text(url: str, *, timeout: int = HTTP_TIMEOUT) -> str:
|
||||
"""Fetch *url* and return its body as text, or ``""`` on any failure.
|
||||
|
||||
@@ -80,6 +113,34 @@ async def fetch_text(url: str, *, timeout: int = HTTP_TIMEOUT) -> str:
|
||||
return ""
|
||||
|
||||
|
||||
async def fetch_json(
|
||||
url: str, *, timeout: int = HTTP_TIMEOUT
|
||||
) -> tuple[int, Any]:
|
||||
"""Fetch *url* and return ``(status, parsed_body)``.
|
||||
|
||||
Unlike :func:`fetch_text` this reports the status, because callers such as
|
||||
the file-listing endpoints need to distinguish "repo not found" (404) from
|
||||
a transport failure. ``parsed_body`` is ``None`` when the response is not
|
||||
JSON or the request failed outright (status ``0``).
|
||||
"""
|
||||
|
||||
try:
|
||||
async with aiohttp.ClientSession(
|
||||
headers={"User-Agent": USER_AGENT},
|
||||
timeout=aiohttp.ClientTimeout(total=timeout),
|
||||
) as session:
|
||||
async with session.get(url) as resp:
|
||||
if resp.status != 200:
|
||||
return resp.status, None
|
||||
try:
|
||||
return resp.status, await resp.json(content_type=None)
|
||||
except Exception:
|
||||
return resp.status, None
|
||||
except Exception as exc: # pragma: no cover - network dependent
|
||||
logger.debug("Failed to fetch %s: %s", url, exc)
|
||||
return 0, None
|
||||
|
||||
|
||||
class ModelSource:
|
||||
"""Description and I/O for one external model hosting site."""
|
||||
|
||||
@@ -95,6 +156,12 @@ class ModelSource:
|
||||
#: Whether models can be downloaded directly from this site.
|
||||
supports_download: bool = False
|
||||
|
||||
#: Branch used when the caller does not pass an explicit revision.
|
||||
default_revision: str = ""
|
||||
|
||||
#: Sub-directory the "use default paths" template places downloads in.
|
||||
default_subdir: str = ""
|
||||
|
||||
#: Lenient pattern used to recognise URLs already stored in metadata.
|
||||
#: Captures the site-specific source id in group ``id``.
|
||||
url_pattern: re.Pattern[str] | None = None
|
||||
@@ -163,6 +230,45 @@ class ModelSource:
|
||||
|
||||
return ""
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Download support
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def list_files(
|
||||
self, source_id: str, revision: str = ""
|
||||
) -> list[dict[str, Any]]:
|
||||
"""List downloadable weight files in *source_id*.
|
||||
|
||||
Returns ``[{"filename": <repo-relative path>, "size": <bytes>}]``,
|
||||
largest first, filtered to :data:`MODEL_FILE_EXTENSIONS`. Sites
|
||||
without download support return an empty list.
|
||||
|
||||
Raises :class:`ModelSourceError` when the repository cannot be read,
|
||||
so the handler can surface "not found" separately from a transport
|
||||
failure.
|
||||
"""
|
||||
|
||||
return []
|
||||
|
||||
def file_download_url(
|
||||
self, source_id: str, filename: str, revision: str = ""
|
||||
) -> str:
|
||||
"""Return the direct (redirecting) download URL for one file."""
|
||||
|
||||
raise ModelSourceError(
|
||||
f"{self.label or self.platform} does not support downloads", status=400
|
||||
)
|
||||
|
||||
def resolve_revision(self, revision: str = "") -> str:
|
||||
"""Return *revision*, falling back to this site's default branch."""
|
||||
|
||||
return revision or self.default_revision
|
||||
|
||||
def page_url_for_file(self, source_id: str, filename: str) -> str:
|
||||
"""Return the human-facing page for *filename* inside *source_id*."""
|
||||
|
||||
return self.canonical_url(source_id)
|
||||
|
||||
def __repr__(self) -> str: # pragma: no cover - debugging aid
|
||||
return f"<ModelSource {self.platform}>"
|
||||
|
||||
@@ -175,12 +281,33 @@ def clean_source_url(url: Any) -> str:
|
||||
return url.strip()
|
||||
|
||||
|
||||
def filter_weight_files(entries: Iterable[tuple[str, int]]) -> list[dict[str, Any]]:
|
||||
"""Keep model-weight files from ``(path, size)`` pairs, largest first.
|
||||
|
||||
Every site lists a lot more than weights (READMEs, configs, tokenizers,
|
||||
…); the download picker only ever wants the files ComfyUI can load, which
|
||||
is exactly :data:`MODEL_FILE_EXTENSIONS`.
|
||||
"""
|
||||
|
||||
files = [
|
||||
{"filename": path, "size": int(size or 0)}
|
||||
for path, size in entries
|
||||
if path and os.path.splitext(path)[1].lower() in MODEL_FILE_EXTENSIONS
|
||||
]
|
||||
files.sort(key=lambda entry: entry["size"], reverse=True)
|
||||
return files
|
||||
|
||||
|
||||
__all__ = [
|
||||
"GROUP_PREFIXES",
|
||||
"HTTP_TIMEOUT",
|
||||
"ModelSource",
|
||||
"ModelSourceError",
|
||||
"SourceRef",
|
||||
"USER_AGENT",
|
||||
"clean_source_url",
|
||||
"fetch_json",
|
||||
"fetch_text",
|
||||
"filter_weight_files",
|
||||
"is_valid_source_id",
|
||||
]
|
||||
|
||||
@@ -2,9 +2,18 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
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
|
||||
#: sub-paths such as ``/resolve/main/model.safetensors``.
|
||||
@@ -25,6 +34,8 @@ class HuggingFaceSource(ModelSource):
|
||||
label = "Hugging Face"
|
||||
supports_enrichment = True
|
||||
supports_download = True
|
||||
default_revision = "main"
|
||||
default_subdir = "huggingface"
|
||||
url_pattern = _URL_PATTERN
|
||||
strict_url_pattern = _STRICT_URL_PATTERN
|
||||
|
||||
@@ -32,7 +43,7 @@ class HuggingFaceSource(ModelSource):
|
||||
return f"https://huggingface.co/{source_id}"
|
||||
|
||||
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:
|
||||
"""Fetch ``README.md`` from Hugging Face (tries ``main``, then ``master``)."""
|
||||
@@ -45,5 +56,51 @@ class HuggingFaceSource(ModelSource):
|
||||
return text
|
||||
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"]
|
||||
|
||||
@@ -2,20 +2,38 @@
|
||||
|
||||
ModelScope exposes the same "model card as README.md" convention as
|
||||
Hugging Face, including a YAML frontmatter block that often carries
|
||||
``base_model:`` and ``trigger_words:``. Two public endpoints are used,
|
||||
neither of which requires an API key for public models:
|
||||
``base_model:`` and ``trigger_words:``. Three public endpoints are used,
|
||||
none of which requires an API key for public models:
|
||||
|
||||
* ``/models/{owner}/{name}/resolve/{revision}/README.md`` — raw model card
|
||||
* ``/api/v1/models/{owner}/{name}/repo?Revision=..&FilePath=README.md`` —
|
||||
the same content through the API, used as a fallback when the resolve
|
||||
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
|
||||
|
||||
import logging
|
||||
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(
|
||||
r"https?://(?:www\.)?modelscope\.(?:cn|com)/models/(?P<id>[^/?#\s]+/[^/?#\s]+)"
|
||||
@@ -41,7 +59,9 @@ class ModelScopeSource(ModelSource):
|
||||
platform = "modelscope"
|
||||
label = "ModelScope"
|
||||
supports_enrichment = True
|
||||
supports_download = False
|
||||
supports_download = True
|
||||
default_revision = "master"
|
||||
default_subdir = "modelscope"
|
||||
url_pattern = _URL_PATTERN
|
||||
strict_url_pattern = _STRICT_URL_PATTERN
|
||||
|
||||
@@ -49,7 +69,10 @@ class ModelScopeSource(ModelSource):
|
||||
return f"https://modelscope.cn/models/{source_id}"
|
||||
|
||||
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:
|
||||
"""Fetch the model card, preferring the raw resolve URL."""
|
||||
@@ -72,5 +95,50 @@ class ModelScopeSource(ModelSource):
|
||||
return text
|
||||
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"]
|
||||
|
||||
@@ -56,6 +56,21 @@ def source_label(platform: Optional[str], default: str = "") -> str:
|
||||
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]:
|
||||
"""Return the :class:`SourceRef` for *url*, or ``None`` if unsupported."""
|
||||
|
||||
@@ -197,6 +212,8 @@ __all__ = [
|
||||
"SOURCE_PLATFORM_FIELD",
|
||||
"SOURCE_URL_FIELD",
|
||||
"detect_source",
|
||||
"downloadable_sources",
|
||||
"get_download_source",
|
||||
"get_source",
|
||||
"get_source_platform",
|
||||
"has_external_source",
|
||||
|
||||
@@ -203,10 +203,18 @@ export const DOWNLOAD_ENDPOINTS = {
|
||||
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 = {
|
||||
repoFiles: '/api/lm/hf-repo-files',
|
||||
download: '/api/lm/download-hf-model',
|
||||
repoFiles: MODEL_SOURCE_ENDPOINTS.repoFiles,
|
||||
download: MODEL_SOURCE_ENDPOINTS.download,
|
||||
};
|
||||
|
||||
// WebSocket endpoints
|
||||
|
||||
@@ -8,6 +8,7 @@ import {
|
||||
isValidModelType,
|
||||
DOWNLOAD_ENDPOINTS,
|
||||
HF_ENDPOINTS,
|
||||
MODEL_SOURCE_ENDPOINTS,
|
||||
WS_ENDPOINTS
|
||||
} from './apiConfig.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 {
|
||||
const params = new URLSearchParams({ repo, revision });
|
||||
const response = await fetch(`${HF_ENDPOINTS.repoFiles}?${params}`);
|
||||
const params = new URLSearchParams({ repo, platform });
|
||||
if (revision) params.set('revision', revision);
|
||||
const response = await fetch(`${MODEL_SOURCE_ENDPOINTS.repoFiles}?${params}`);
|
||||
if (!response.ok) {
|
||||
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();
|
||||
} catch (error) {
|
||||
console.error('Error fetching HF repo files:', error);
|
||||
console.error('Error fetching repository files:', 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 {
|
||||
const response = await fetch(HF_ENDPOINTS.download, {
|
||||
const response = await fetch(MODEL_SOURCE_ENDPOINTS.download, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
platform,
|
||||
repo,
|
||||
filename,
|
||||
revision: revision || 'main',
|
||||
revision: revision || '',
|
||||
model_root: modelRoot,
|
||||
relative_path: relativePath || '',
|
||||
use_default_paths: useDefaultPaths || false,
|
||||
@@ -1404,11 +1427,25 @@ export class BaseModelApiClient {
|
||||
|
||||
return await response.json();
|
||||
} catch (error) {
|
||||
console.error('Error downloading HF model:', error);
|
||||
console.error('Error downloading model:', 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) {
|
||||
const params = new URLSearchParams(baseParams);
|
||||
const isExcludedView = pageState.viewMode === 'excluded';
|
||||
|
||||
@@ -13,6 +13,13 @@ import { buildCivitaiUrl, extractCivitaiModelUrlParts, normalizeCivitaiPageHost
|
||||
import { formatFileSize } from '../utils/formatters.js';
|
||||
import { showDownloadBatchSummary } from '../components/DownloadBatchSummaryModal.js';
|
||||
import { openOtherModelsSettings } from '../utils/otherModels.js';
|
||||
import {
|
||||
buildModelSourceFilePage,
|
||||
detectModelSourceDownloadUrl,
|
||||
getModelSource,
|
||||
isExternalModelSource,
|
||||
isValidRepoId,
|
||||
} from '../utils/modelSourceHelpers.js';
|
||||
|
||||
export class DownloadManager {
|
||||
constructor() {
|
||||
@@ -39,10 +46,11 @@ export class DownloadManager {
|
||||
this.isBatchMode = false;
|
||||
this.editingBatchIndex = -1;
|
||||
|
||||
// HF download state
|
||||
this.hfRepoId = null;
|
||||
this.hfSelectedFiles = [];
|
||||
this.hfRepoCollapsed = {};
|
||||
// External repository download state (Hugging Face / ModelScope)
|
||||
this.sourcePlatform = 'huggingface';
|
||||
this.sourceRepoId = null;
|
||||
this.sourceSelectedFiles = [];
|
||||
this.sourceRepoCollapsed = {};
|
||||
|
||||
this.loadingManager = new LoadingManager();
|
||||
this.folderTreeManager = new FolderTreeManager();
|
||||
@@ -186,10 +194,11 @@ export class DownloadManager {
|
||||
// Reset default path toggle
|
||||
this.loadDefaultPathSetting();
|
||||
|
||||
// Reset HF state
|
||||
this.hfRepoId = null;
|
||||
this.hfSelectedFiles = [];
|
||||
this.hfRepoCollapsed = {};
|
||||
// Reset external repository state
|
||||
this.sourcePlatform = 'huggingface';
|
||||
this.sourceRepoId = null;
|
||||
this.sourceSelectedFiles = [];
|
||||
this.sourceRepoCollapsed = {};
|
||||
}
|
||||
|
||||
async retrieveVersionsForModel(modelId, source = null) {
|
||||
@@ -212,10 +221,12 @@ export class DownloadManager {
|
||||
|
||||
// Detect URL types — all URLs must share the same source type
|
||||
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');
|
||||
|
||||
if (!isHf && !isCivitai) {
|
||||
if (!isExternalSource && !isCivitai) {
|
||||
const allValid = urlTypes.every(t => t !== null);
|
||||
if (!allValid) {
|
||||
errorElement.textContent = translate('modals.download.errors.invalidUrl');
|
||||
@@ -228,8 +239,8 @@ export class DownloadManager {
|
||||
}
|
||||
}
|
||||
|
||||
if (isHf) {
|
||||
return this._validateAndFetchHf(urls, errorElement);
|
||||
if (isExternalSource) {
|
||||
return this._validateAndFetchExternalRepo(urls, errorElement);
|
||||
}
|
||||
|
||||
// --- Original CivitAI flow below ---
|
||||
@@ -327,45 +338,66 @@ export class DownloadManager {
|
||||
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) {
|
||||
const info = DownloadManager.detectUrlType(urls[0]);
|
||||
// Direct file resolve URL → skip file selection, go to location
|
||||
if (info.type === 'hf-resolve') {
|
||||
// Direct file URL → skip file selection, go to location
|
||||
if (info.type === 'model-source-file') {
|
||||
this.isBatchMode = false;
|
||||
this.hfRepoId = info.repo;
|
||||
this.hfSelectedFiles = [info.filename];
|
||||
this.source = 'huggingface';
|
||||
this.sourcePlatform = info.platform;
|
||||
this.sourceRepoId = info.repo;
|
||||
this.sourceSelectedFiles = [info.filename];
|
||||
this.source = info.platform;
|
||||
this.proceedToLocation();
|
||||
return;
|
||||
}
|
||||
// Repo URL → fetch file list and convert to batch items
|
||||
try {
|
||||
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.batchModels = [];
|
||||
this.source = 'huggingface';
|
||||
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.batchModels = await this._fetchExternalRepoItems(urls[0], info);
|
||||
this.source = info.platform;
|
||||
this.showBatchPreviewStep();
|
||||
} catch (err) {
|
||||
errorElement.textContent = err.message;
|
||||
@@ -375,10 +407,9 @@ export class DownloadManager {
|
||||
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.batchModels = [];
|
||||
this.source = 'huggingface';
|
||||
this.loadingManager.showSimpleLoading(translate('modals.download.fetchingRepoFiles'));
|
||||
|
||||
for (const url of urls) {
|
||||
@@ -387,42 +418,15 @@ export class DownloadManager {
|
||||
this.batchModels.push({ url, error: 'Invalid URL', versions: [], selectedVersion: null });
|
||||
continue;
|
||||
}
|
||||
if (info.type === 'hf-resolve') {
|
||||
this.batchModels.push({
|
||||
url,
|
||||
source: 'huggingface',
|
||||
repo: info.repo,
|
||||
this.source = info.platform;
|
||||
if (info.type === 'model-source-file') {
|
||||
this.batchModels.push(this._makeExternalItem(url, info, {
|
||||
filename: info.filename,
|
||||
revision: info.revision || 'main',
|
||||
displayName: info.filename,
|
||||
selectedVersion: true,
|
||||
versions: [],
|
||||
checked: false,
|
||||
error: null,
|
||||
});
|
||||
} else if (info.type === 'hf-repo') {
|
||||
revision: info.revision,
|
||||
}));
|
||||
} else if (info.type === 'model-source-repo') {
|
||||
try {
|
||||
const files = await this.apiClient.fetchHfRepoFiles(info.repo);
|
||||
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,
|
||||
});
|
||||
}
|
||||
this.batchModels.push(...await this._fetchExternalRepoItems(url, info));
|
||||
} catch (err) {
|
||||
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.
|
||||
* @param {string} url
|
||||
* @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) {
|
||||
const trimmed = url.trim();
|
||||
@@ -492,38 +497,27 @@ export class DownloadManager {
|
||||
return { type: 'civitai' };
|
||||
}
|
||||
|
||||
// Hugging Face resolve/blob URL → direct file
|
||||
// "blob" is the web preview page; it maps 1:1 to the "resolve" download URL
|
||||
const hfResolveMatch = trimmed.match(/huggingface\.co\/([^/\s]+\/[^/\s]+)\/(?:resolve|blob)\/([^/\s]+)\/(.+)/i);
|
||||
if (hfResolveMatch) {
|
||||
return {
|
||||
type: 'hf-resolve',
|
||||
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) {
|
||||
// External model sources (Hugging Face / ModelScope). Repository URLs
|
||||
// list every weight file; resolve URLs point at one file. Both are
|
||||
// recognised through the shared registry, so adding a site is a
|
||||
// registry change rather than a change here.
|
||||
const sourceInfo = detectModelSourceDownloadUrl(trimmed);
|
||||
if (sourceInfo) {
|
||||
// Reject path-traversal patterns like "../.." or "user/.."
|
||||
const parts = hfRepoMatch[1].split('/');
|
||||
if (parts.some(p => p === '.' || p === '..')) {
|
||||
if (!isValidRepoId(sourceInfo.repo)) {
|
||||
return null;
|
||||
}
|
||||
return {
|
||||
type: 'hf-repo',
|
||||
repo: hfRepoMatch[1],
|
||||
type: sourceInfo.kind === 'file' ? 'model-source-file' : 'model-source-repo',
|
||||
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)) {
|
||||
return { type: 'direct-http' };
|
||||
}
|
||||
@@ -931,7 +925,7 @@ export class DownloadManager {
|
||||
}
|
||||
|
||||
// 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) {
|
||||
showToast('toast.loras.pleaseSelectVersion', {}, 'error');
|
||||
return;
|
||||
@@ -1164,12 +1158,15 @@ export class DownloadManager {
|
||||
/**
|
||||
* Synthesize a clickable URL for a single-download failure entry.
|
||||
* 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 }) {
|
||||
if (source === 'huggingface' && repo) {
|
||||
const base = `https://huggingface.co/${encodeURI(repo)}`;
|
||||
return filename ? `${base}/blob/${encodeURI('main')}/${encodeURI(filename)}` : base;
|
||||
if (isExternalModelSource(source) && repo) {
|
||||
return buildModelSourceFilePage({
|
||||
platform: source,
|
||||
repo,
|
||||
filename,
|
||||
}) || getModelSource(source).canonical(repo);
|
||||
}
|
||||
if (modelId) {
|
||||
return buildCivitaiUrl({
|
||||
@@ -1490,8 +1487,8 @@ export class DownloadManager {
|
||||
* matched card-by-card via `_reconcileViewAfterDownload`; HF
|
||||
* downloads (no CivitAI identity to match) keep the legacy reload.
|
||||
*/
|
||||
async _reconcileBatchViewAfterDownload(completedCivitaiItems = [], hfCompletedCount = 0) {
|
||||
if (hfCompletedCount > 0) {
|
||||
async _reconcileBatchViewAfterDownload(completedCivitaiItems = [], externalCompletedCount = 0) {
|
||||
if (externalCompletedCount > 0) {
|
||||
await resetAndReload(true);
|
||||
return;
|
||||
}
|
||||
@@ -1661,10 +1658,11 @@ export class DownloadManager {
|
||||
return failedItems.length === 0;
|
||||
}
|
||||
|
||||
async _downloadHfSingle({ modelRoot, targetFolder, useDefaultPaths, files = null }) {
|
||||
async _downloadExternalRepoFiles({ modelRoot, targetFolder, useDefaultPaths, files = null }) {
|
||||
modalManager.closeModal('downloadModal');
|
||||
this.loadingManager.restoreProgressBar();
|
||||
const filesToDownload = files || this.hfSelectedFiles;
|
||||
const platform = this.sourcePlatform;
|
||||
const filesToDownload = files || this.sourceSelectedFiles;
|
||||
const totalFiles = filesToDownload.length;
|
||||
const updateProgress = this.loadingManager.showDownloadProgress(totalFiles);
|
||||
|
||||
@@ -1720,10 +1718,11 @@ export class DownloadManager {
|
||||
}
|
||||
};
|
||||
|
||||
const response = await this.apiClient.downloadHfModel({
|
||||
repo: this.hfRepoId,
|
||||
const response = await this.apiClient.downloadModelSource({
|
||||
platform,
|
||||
repo: this.sourceRepoId,
|
||||
filename,
|
||||
revision: 'main',
|
||||
revision: this._defaultRevisionFor(platform),
|
||||
modelRoot,
|
||||
relativePath: targetFolder,
|
||||
useDefaultPaths,
|
||||
@@ -1738,10 +1737,10 @@ export class DownloadManager {
|
||||
} else {
|
||||
failedFiles.push({
|
||||
item: {
|
||||
source: 'huggingface',
|
||||
repo: this.hfRepoId,
|
||||
source: platform,
|
||||
repo: this.sourceRepoId,
|
||||
filename,
|
||||
url: this._buildSingleItemUrl({ source: 'huggingface', repo: this.hfRepoId, filename }),
|
||||
url: this._buildSingleItemUrl({ source: platform, repo: this.sourceRepoId, filename }),
|
||||
},
|
||||
error: response?.error || 'Unknown error',
|
||||
name: filename,
|
||||
@@ -1749,13 +1748,13 @@ export class DownloadManager {
|
||||
}
|
||||
} catch (err) {
|
||||
if (!cancelled) {
|
||||
console.error(`Failed to download HF file ${filename}:`, err);
|
||||
console.error(`Failed to download repo file ${filename}:`, err);
|
||||
failedFiles.push({
|
||||
item: {
|
||||
source: 'huggingface',
|
||||
repo: this.hfRepoId,
|
||||
source: platform,
|
||||
repo: this.sourceRepoId,
|
||||
filename,
|
||||
url: this._buildSingleItemUrl({ source: 'huggingface', repo: this.hfRepoId, filename }),
|
||||
url: this._buildSingleItemUrl({ source: platform, repo: this.sourceRepoId, filename }),
|
||||
},
|
||||
error: err?.message || 'Unknown error',
|
||||
name: filename,
|
||||
@@ -1781,7 +1780,7 @@ export class DownloadManager {
|
||||
total: totalFiles,
|
||||
completed: completedDownloads,
|
||||
failedItems: failedFiles,
|
||||
onRetry: () => this._downloadHfSingle({
|
||||
onRetry: () => this._downloadExternalRepoFiles({
|
||||
modelRoot,
|
||||
targetFolder,
|
||||
useDefaultPaths,
|
||||
@@ -1831,7 +1830,7 @@ export class DownloadManager {
|
||||
|
||||
const validCount = this.batchModels.filter(m => {
|
||||
if (m.error) return false;
|
||||
if (m.source === 'huggingface') return m.checked !== false;
|
||||
if (isExternalModelSource(m.source)) return m.checked !== false;
|
||||
return m.selectedVersion;
|
||||
}).length;
|
||||
document.getElementById('downloadModalTitle').textContent =
|
||||
@@ -1839,7 +1838,9 @@ export class DownloadManager {
|
||||
` (${validCount})`;
|
||||
|
||||
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
|
||||
const errorItemsHtml = this.batchModels.map((item, index) => {
|
||||
@@ -1863,7 +1864,7 @@ export class DownloadManager {
|
||||
// CivitAI items render flat, outside any group (unchanged)
|
||||
const civitaiItemsHtml = this.batchModels.map((item, index) => {
|
||||
if (item.error) return null;
|
||||
if (item.source === 'huggingface') return null;
|
||||
if (isExternalModelSource(item.source)) return null;
|
||||
const ver = item.selectedVersion;
|
||||
const firstImage = ver?.images?.find(img => !img.url.endsWith('.mp4'));
|
||||
const thumbnailUrl = firstImage ? firstImage.url : '/loras_static/images/no-preview.png';
|
||||
@@ -1901,25 +1902,30 @@ export class DownloadManager {
|
||||
`;
|
||||
}).filter(Boolean).join('');
|
||||
|
||||
// Group HF items by repo (data model stays flat — only rendering groups)
|
||||
const hfGroups = {};
|
||||
// Group external-repository items by platform + repo so that the same
|
||||
// `owner/name` on two sites stays in two groups (data model stays flat
|
||||
// — only rendering groups).
|
||||
const externalGroups = {};
|
||||
this.batchModels.forEach((item, index) => {
|
||||
if (item.error || item.source !== 'huggingface') return;
|
||||
const repo = item.repo || 'unknown';
|
||||
if (!hfGroups[repo]) hfGroups[repo] = [];
|
||||
hfGroups[repo].push({ item, index });
|
||||
if (item.error || !isExternalModelSource(item.source)) return;
|
||||
const groupKey = this._externalGroupKey(item);
|
||||
if (!externalGroups[groupKey]) {
|
||||
externalGroups[groupKey] = { repo: item.repo || 'unknown', items: [] };
|
||||
}
|
||||
externalGroups[groupKey].items.push({ item, index });
|
||||
});
|
||||
|
||||
const renderHfItem = ({ item, index }) => {
|
||||
const hfSize = item.fileSizeBytes ? formatFileSize(item.fileSizeBytes) : '?';
|
||||
const renderExternalItem = ({ item, index }) => {
|
||||
const fileSize = item.fileSizeBytes ? formatFileSize(item.fileSizeBytes) : '?';
|
||||
const badge = getModelSource(item.source)?.label || item.source;
|
||||
return `
|
||||
<div class="batch-preview-item" data-index="${index}">
|
||||
<input type="checkbox" class="batch-preview-checkbox"
|
||||
data-index="${index}" ${item.checked !== false ? 'checked' : ''} />
|
||||
<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">
|
||||
<span>${hfSize}</span>
|
||||
<span>${fileSize}</span>
|
||||
<span>${item.repo || ''}</span>
|
||||
</div>
|
||||
</div>
|
||||
@@ -1930,32 +1936,32 @@ export class DownloadManager {
|
||||
`;
|
||||
};
|
||||
|
||||
const hfGroupsHtml = Object.keys(hfGroups).map(repo => {
|
||||
const items = hfGroups[repo];
|
||||
const isCollapsed = this.hfRepoCollapsed[repo] === true;
|
||||
const externalGroupsHtml = Object.keys(externalGroups).map(groupKey => {
|
||||
const { repo, items } = externalGroups[groupKey];
|
||||
const isCollapsed = this.sourceRepoCollapsed[groupKey] === true;
|
||||
const allChecked = items.every(({ item }) => item.checked !== false);
|
||||
const fileCount = items.length;
|
||||
return `
|
||||
<div class="batch-preview-group" data-repo="${repo}">
|
||||
<div class="batch-preview-group" data-repo="${groupKey}">
|
||||
<div class="batch-preview-group-header">
|
||||
<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-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 class="batch-preview-group-body ${isCollapsed ? '' : 'expanded'}">
|
||||
${items.map(renderHfItem).join('')}
|
||||
${items.map(renderExternalItem).join('')}
|
||||
</div>
|
||||
</div>
|
||||
`;
|
||||
}).join('');
|
||||
|
||||
let itemsHtml = errorItemsHtml + civitaiItemsHtml + hfGroupsHtml;
|
||||
let itemsHtml = errorItemsHtml + civitaiItemsHtml + externalGroupsHtml;
|
||||
|
||||
// Prepend select-all toolbar if there are HF items with checkboxes
|
||||
if (hasHfItems) {
|
||||
// Prepend select-all toolbar if there are external items with checkboxes
|
||||
if (hasExternalItems) {
|
||||
const allChecked = this.batchModels
|
||||
.filter(m => m.source === 'huggingface' && !m.error)
|
||||
.filter(m => isExternalModelSource(m.source) && !m.error)
|
||||
.every(m => m.checked !== false);
|
||||
itemsHtml = `
|
||||
<div class="batch-preview-select-all">
|
||||
@@ -1980,13 +1986,18 @@ export class DownloadManager {
|
||||
// Global select-all
|
||||
const selectAll = document.getElementById('batchSelectAll');
|
||||
if (selectAll) {
|
||||
const hfItems = this.batchModels.filter(m => m.source === 'huggingface' && !m.error);
|
||||
selectAll.checked = hfItems.length > 0 && hfItems.every(m => m.checked !== false);
|
||||
const externalItems = this.batchModels.filter(
|
||||
m => isExternalModelSource(m.source) && !m.error
|
||||
);
|
||||
selectAll.checked = externalItems.length > 0
|
||||
&& externalItems.every(m => m.checked !== false);
|
||||
}
|
||||
// Per-group select-all
|
||||
list.querySelectorAll('.batch-preview-group-select-all').forEach(gsa => {
|
||||
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);
|
||||
});
|
||||
};
|
||||
@@ -1998,7 +2009,7 @@ export class DownloadManager {
|
||||
const repo = groupSelectAll.dataset.repo;
|
||||
const checked = groupSelectAll.checked;
|
||||
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;
|
||||
const cb = list.querySelector(`.batch-preview-checkbox[data-index="${idx}"]`);
|
||||
if (cb) cb.checked = checked;
|
||||
@@ -2014,9 +2025,9 @@ export class DownloadManager {
|
||||
const repo = group.dataset.repo;
|
||||
const body = group.querySelector('.batch-preview-group-body');
|
||||
const toggle = group.querySelector('.batch-preview-group-toggle');
|
||||
const isCollapsed = this.hfRepoCollapsed[repo];
|
||||
const isCollapsed = this.sourceRepoCollapsed[repo];
|
||||
if (isCollapsed) {
|
||||
this.hfRepoCollapsed[repo] = false;
|
||||
this.sourceRepoCollapsed[repo] = false;
|
||||
body.style.transition = ''; // restore in case collapse was interrupted
|
||||
body.classList.add('expanded');
|
||||
toggle.classList.add('expanded');
|
||||
@@ -2025,13 +2036,13 @@ export class DownloadManager {
|
||||
body.style.maxHeight = body.scrollHeight + 'px';
|
||||
const onEnd = (e) => {
|
||||
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.removeEventListener('transitionend', onEnd);
|
||||
};
|
||||
body.addEventListener('transitionend', onEnd);
|
||||
} else {
|
||||
this.hfRepoCollapsed[repo] = true;
|
||||
this.sourceRepoCollapsed[repo] = true;
|
||||
body.style.maxHeight = body.scrollHeight + 'px';
|
||||
requestAnimationFrame(() => {
|
||||
// animate only max-height; keep expanded so opacity stays 1
|
||||
@@ -2040,7 +2051,7 @@ export class DownloadManager {
|
||||
toggle.classList.remove('expanded');
|
||||
const onEnd = (e) => {
|
||||
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.style.transition = '';
|
||||
body.removeEventListener('transitionend', onEnd);
|
||||
@@ -2119,7 +2130,7 @@ export class DownloadManager {
|
||||
// For HF items, respect the checked flag; for CivitAI items, use selectedVersion
|
||||
const validModels = this.batchModels.filter(m => {
|
||||
if (m.error) return false;
|
||||
if (m.source === 'huggingface') return m.checked !== false;
|
||||
if (isExternalModelSource(m.source)) return m.checked !== false;
|
||||
return m.selectedVersion;
|
||||
});
|
||||
if (validModels.length === 0) return;
|
||||
@@ -2172,8 +2183,8 @@ export class DownloadManager {
|
||||
}
|
||||
if (!this.isBatchMode) {
|
||||
// Single-item download
|
||||
if (this.source === 'huggingface') {
|
||||
return this._downloadHfSingle({
|
||||
if (isExternalModelSource(this.source)) {
|
||||
return this._downloadExternalRepoFiles({
|
||||
modelRoot,
|
||||
targetFolder,
|
||||
useDefaultPaths,
|
||||
@@ -2228,7 +2239,7 @@ export class DownloadManager {
|
||||
if (m.error) return false;
|
||||
if (!m.selectedVersion) return false;
|
||||
// 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;
|
||||
});
|
||||
if (downloadItems.length === 0) {
|
||||
@@ -2255,10 +2266,11 @@ export class DownloadManager {
|
||||
let cancelled = false;
|
||||
const failedItems = [];
|
||||
// Successful CivitAI items are reconciled in place afterwards
|
||||
// (their cards can be matched by model id); HF items keep the
|
||||
// legacy full reload because they have no CivitAI identity (#1078).
|
||||
// (their cards can be matched by model id); externally-sourced items
|
||||
// keep the legacy full reload because they have no CivitAI identity
|
||||
// (#1078).
|
||||
const completedCivitaiItems = [];
|
||||
let hfCompletedCount = 0;
|
||||
let externalCompletedCount = 0;
|
||||
|
||||
loadingManager.showCancelButton(async () => {
|
||||
if (cancelled) return;
|
||||
@@ -2301,15 +2313,15 @@ export class DownloadManager {
|
||||
|
||||
const item = downloadItems[i];
|
||||
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);
|
||||
loadingManager.setStatus(`${i + 1}/${downloadItems.length}: ${name}`);
|
||||
|
||||
try {
|
||||
let response;
|
||||
if (isHf) {
|
||||
const downloadId = Date.now().toString() + '_hf_' + i;
|
||||
if (isExternal) {
|
||||
const downloadId = Date.now().toString() + '_src_' + i;
|
||||
const wsHf = new WebSocket(`${wsProtocol}${window.location.host}/ws/download-progress?id=${downloadId}`);
|
||||
try {
|
||||
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,
|
||||
filename: item.filename,
|
||||
revision: item.revision || 'main',
|
||||
revision: item.revision || this._defaultRevisionFor(item.platform || item.source),
|
||||
modelRoot,
|
||||
relativePath: targetFolder,
|
||||
useDefaultPaths,
|
||||
@@ -2363,8 +2376,8 @@ export class DownloadManager {
|
||||
} else {
|
||||
completedDownloads++;
|
||||
updateProgress(100, completedDownloads, '');
|
||||
if (isHf) {
|
||||
hfCompletedCount++;
|
||||
if (isExternal) {
|
||||
externalCompletedCount++;
|
||||
} else {
|
||||
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, {
|
||||
|
||||
@@ -19,21 +19,35 @@ export const MODEL_SOURCES = [
|
||||
groupPrefix: 'hf',
|
||||
supportsEnrichment: true,
|
||||
supportsDownload: true,
|
||||
defaultRevision: 'main',
|
||||
defaultSubdir: 'huggingface',
|
||||
exampleUrl: 'https://huggingface.co/user/repo',
|
||||
placeholder: 'https://huggingface.co/user/repo',
|
||||
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}`,
|
||||
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',
|
||||
label: 'ModelScope',
|
||||
groupPrefix: 'ms',
|
||||
supportsEnrichment: true,
|
||||
supportsDownload: false,
|
||||
supportsDownload: true,
|
||||
defaultRevision: 'master',
|
||||
defaultSubdir: 'modelscope',
|
||||
exampleUrl: 'https://modelscope.cn/models/user/repo',
|
||||
placeholder: 'https://modelscope.cn/models/user/repo',
|
||||
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}`,
|
||||
filePage: (id, filename) =>
|
||||
`https://modelscope.cn/models/${id}/file/view/master/${filename}`,
|
||||
},
|
||||
{
|
||||
platform: 'tensorart',
|
||||
@@ -41,10 +55,14 @@ export const MODEL_SOURCES = [
|
||||
groupPrefix: 'ta',
|
||||
supportsEnrichment: false,
|
||||
supportsDownload: false,
|
||||
defaultRevision: '',
|
||||
defaultSubdir: '',
|
||||
exampleUrl: 'https://tensor.art/models/827823520299086029',
|
||||
placeholder: 'https://tensor.art/models/827823520299086029',
|
||||
pattern: /^https?:\/\/(?:www\.)?(?:tensor\.art|tusi\.cn)\/models\/(\d+)/i,
|
||||
filePattern: null,
|
||||
canonical: (id) => `https://tensor.art/models/${id}`,
|
||||
filePage: null,
|
||||
},
|
||||
];
|
||||
|
||||
@@ -171,3 +189,103 @@ export function openModelSource(url) {
|
||||
if (!url) return;
|
||||
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(),
|
||||
downloadModelSource: vi.fn(),
|
||||
downloadHfModel: vi.fn(),
|
||||
cancelDownload: vi.fn(),
|
||||
getPageState: vi.fn(() => ({})),
|
||||
@@ -158,7 +159,7 @@ describe('DownloadManager batch download summary flow', () => {
|
||||
// Reset the shared mocks so mockResolvedValueOnce queues and call
|
||||
// history never leak between tests.
|
||||
mockApiClient.downloadModel.mockReset();
|
||||
mockApiClient.downloadHfModel.mockReset();
|
||||
mockApiClient.downloadModelSource.mockReset();
|
||||
mockApiClient.cancelDownload.mockReset();
|
||||
showToastMock.mockClear();
|
||||
showDownloadBatchSummaryMock.mockClear();
|
||||
@@ -406,14 +407,15 @@ describe('DownloadManager batch download summary flow', () => {
|
||||
expect(showToastMock).toHaveBeenCalledWith('toast.loras.downloadCompleted', {}, 'success');
|
||||
});
|
||||
|
||||
it('shows a summary for HF partial failure and retries only the failed files', async () => {
|
||||
manager.hfRepoId = 'user/repo';
|
||||
manager.hfSelectedFiles = ['a.safetensors', 'b.safetensors'];
|
||||
mockApiClient.downloadHfModel
|
||||
it('shows a summary for external repo partial failure and retries only the failed files', async () => {
|
||||
manager.sourcePlatform = 'huggingface';
|
||||
manager.sourceRepoId = 'user/repo';
|
||||
manager.sourceSelectedFiles = ['a.safetensors', 'b.safetensors'];
|
||||
mockApiClient.downloadModelSource
|
||||
.mockResolvedValueOnce({ success: true })
|
||||
.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(showDownloadBatchSummaryMock).toHaveBeenCalledTimes(1);
|
||||
@@ -428,7 +430,9 @@ describe('DownloadManager batch download summary flow', () => {
|
||||
|
||||
await summary.onRetry();
|
||||
|
||||
expect(mockApiClient.downloadHfModel).toHaveBeenCalledTimes(3);
|
||||
expect(mockApiClient.downloadHfModel.mock.calls[2][0].filename).toBe('b.safetensors');
|
||||
expect(mockApiClient.downloadModelSource).toHaveBeenCalledTimes(3);
|
||||
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));
|
||||
});
|
||||
});
|
||||
+61
-8
@@ -1,14 +1,15 @@
|
||||
import { describe, it, expect } from 'vitest';
|
||||
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', () => {
|
||||
const result = DownloadManager.detectUrlType(
|
||||
'https://huggingface.co/dx8152/Flux2-Klein-9B-Consistency/resolve/main/Flux2-Klein-9B-consistency-V2.safetensors'
|
||||
);
|
||||
expect(result).toEqual({
|
||||
type: 'hf-resolve',
|
||||
type: 'model-source-file',
|
||||
platform: 'huggingface',
|
||||
repo: 'dx8152/Flux2-Klein-9B-Consistency',
|
||||
revision: 'main',
|
||||
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'
|
||||
);
|
||||
expect(result).toEqual({
|
||||
type: 'hf-resolve',
|
||||
type: 'model-source-file',
|
||||
platform: 'huggingface',
|
||||
repo: 'user/repo',
|
||||
revision: 'main',
|
||||
filename: 'subdir/model.safetensors',
|
||||
@@ -32,7 +34,8 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
|
||||
'https://huggingface.co/dx8152/Flux2-Klein-9B-Consistency'
|
||||
);
|
||||
expect(result).toEqual({
|
||||
type: 'hf-repo',
|
||||
type: 'model-source-repo',
|
||||
platform: 'huggingface',
|
||||
repo: 'dx8152/Flux2-Klein-9B-Consistency',
|
||||
});
|
||||
});
|
||||
@@ -40,7 +43,8 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
|
||||
it('detects HF repo URL (bare user/repo)', () => {
|
||||
const result = DownloadManager.detectUrlType('dx8152/Flux2-Klein-9B-Consistency');
|
||||
expect(result).toEqual({
|
||||
type: 'hf-repo',
|
||||
type: 'model-source-repo',
|
||||
platform: 'huggingface',
|
||||
repo: 'dx8152/Flux2-Klein-9B-Consistency',
|
||||
});
|
||||
});
|
||||
@@ -50,7 +54,8 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
|
||||
'https://huggingface.co/user/repo/'
|
||||
);
|
||||
expect(result).toEqual({
|
||||
type: 'hf-repo',
|
||||
type: 'model-source-repo',
|
||||
platform: 'huggingface',
|
||||
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'
|
||||
);
|
||||
expect(result).toEqual({
|
||||
type: 'hf-resolve',
|
||||
type: 'model-source-file',
|
||||
platform: 'huggingface',
|
||||
repo: 'Comfy-Org/z_image_turbo',
|
||||
revision: 'main',
|
||||
filename: 'split_files/diffusion_models/z_image_turbo_bf16.safetensors',
|
||||
@@ -115,7 +121,7 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
|
||||
const result = DownloadManager.detectUrlType(
|
||||
'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', () => {
|
||||
@@ -126,4 +132,51 @@ describe('DownloadManager.detectUrlType — HF URL detection', () => {
|
||||
);
|
||||
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();
|
||||
});
|
||||
});
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -14,10 +14,14 @@ from py.services.model_sources import (
|
||||
ModelScopeSource,
|
||||
TensorArtSource,
|
||||
detect_source,
|
||||
downloadable_sources,
|
||||
get_download_source,
|
||||
get_source,
|
||||
get_source_platform,
|
||||
has_external_source,
|
||||
is_valid_source_id,
|
||||
list_sources,
|
||||
ModelSourceError,
|
||||
normalize_metadata_source,
|
||||
resolve_source_ref,
|
||||
source_group_key,
|
||||
@@ -127,11 +131,15 @@ class TestCapabilities:
|
||||
source = get_source("huggingface")
|
||||
assert source.supports_enrichment 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")
|
||||
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):
|
||||
source = get_source("tensorart")
|
||||
@@ -322,3 +330,177 @@ class TestAssetBaseUrl:
|
||||
ModelScopeSource().asset_base_url("u/r")
|
||||
== "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"
|
||||
|
||||
Reference in New Issue
Block a user