mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-20 18:51:26 -03:00
Will Miao
d572292142
A ModelScope or Hugging Face download landed as a bare filename, hash and
source link; the model card stayed empty until the user ran "Enrich
Metadata with AI" by hand. But everything that makes a CivitAI download
useful — the display name, the description, the tags, the trigger words,
the example images, the preview — is already published by those sites'
public APIs, so asking for it at download time is deterministic work, not
model work.
Add `py/services/model_sources/hydration.py`, called by
`_save_source_metadata()` once the sidecar exists and the file is in the
scanner cache. It fetches the model card plus the site's card extras and
hands them to the same `PostProcessor` the AI skill uses, with an empty
`llm_output`, so the two paths cannot drift apart. What lands:
* `model_name` from the site's own display name (ModelScope's `Name`), so
the card stops showing the local filename — written only while the value
still equals the file stem, since once a user renames a model that
choice is theirs to keep
* `civitai.name` from the matched version's label (`showName`), which the
card renders as the version chip
* `civitai.description` / `modelDescription` from the author summary plus
the README as HTML
* `civitai.images` / `preview_url` from the per-file example images
* `civitai.trainedWords` from the per-file trigger words
* `base_model`, `tags` and `usage_tips` as before
Provenance stays honest: the pass records
`metadata_source = "source:<platform>"` rather than the skill's
`agent:enrich_hf_metadata`, and — because no provider ran — it no longer
stamps `llm_enriched_at`; that stamp is now conditional on the LLM
actually answering, which is what the field means. The five hand-rolled
`civitai` dict merges in the post-processor collapse into one
`_merge_civitai()` helper.
Two guards keep it safe. Only a model whose stored
`source_platform`/`source_url` match the repository being downloaded is
updated, so a local file that merely shares a name never receives another
model's card; and a file already on disk is topped up too, which
back-fills models downloaded before this existed. READMEs and detail
payloads describe the repository rather than the file, so a short-lived
process-wide `ModelSourceCache` (300 s, 32 entries) keeps a batch over one
repository to two HTTP requests. Every failure is logged and swallowed:
hydration can never fail a download.
Fix the hash policy while here. `_save_source_metadata()` went straight to
`MetadataManager.create_default_metadata()`, bypassing the per-type
factory on the owning scanner, so a checkpoint paid a full SHA256 inside
the download request — `CheckpointScanner`/`OtherScanner` deliberately
record `hash_status="pending"` with an empty `sha256` for their multi-GB
files. Metadata is now created through `scanner._create_default_metadata()`.
Hydration copes with the empty hash: `_matching_versions()` falls back to
the repository basename, which is exactly what the download just wrote.
Report both post-transfer stages, which advance no byte counter and so
read as a stall: the bar sat at 100% showing `0 B/s` for the seconds spent
hashing and fetching. `_report_phase()` broadcasts
`{"status": "metadata", "stage": "indexing" | "source", "platform": ...}`,
and `LoadingManager` names the stage in the status line (keeping the batch
position), retitles the item line, replaces the dead speed figure and runs
a sheen over the bar. `stage`/`platform` are machine-readable; the wording
is localised in the frontend.
Finally, `modelscope.ai` is its own catalogue rather than an alias of
`modelscope.cn` — `referall13/EM1` exists only on `.ai` and
`jj3550945163/Krea-2-LORA` only on `.cn` — so its URLs were rejected with
"Invalid model URL format". Register it as `ModelScopeIntlSource`
(`platform="modelscope-ai"`, `msai:` group prefix, its own default
download directory) and derive every URL either deployment builds from a
per-class `base_url`. `modelscope.com` stays an alias of `.cn`, which is
what it redirects to. The frontend source table, the link dialog hints and
the docs mirror the split.
Verified against the live APIs: both reported `.ai` repositories list
their files, read their READMEs and yield name / version / base model /
trigger words / example images. Backend 3092 passed; frontend 1259 JS +
91 Vue passed. The nine locales carry the new progress copy in the next
commit.
640 lines
26 KiB
Python
640 lines
26 KiB
Python
"""Handlers for external model sources: linking, file listing and downloads.
|
|
|
|
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
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
from typing import Any
|
|
|
|
from aiohttp import web
|
|
|
|
from ...config import config
|
|
from ...services.downloader import (
|
|
DownloadProgress,
|
|
get_downloader,
|
|
)
|
|
from ...services.aria2_downloader import Aria2Downloader
|
|
from ...services.model_sources import (
|
|
ModelSourceError,
|
|
SourceRef,
|
|
detect_source,
|
|
get_download_source,
|
|
hydrate_from_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.metadata_manager import MetadataManager
|
|
from ...utils.models import LoraMetadata, CheckpointMetadata, EmbeddingMetadata
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_DEFAULT_MODEL_CLASS = LoraMetadata
|
|
_DEFAULT_SCANNER_GETTER = "get_lora_scanner"
|
|
|
|
|
|
def _infer_model_type(model_root: str) -> tuple[Any, str]:
|
|
"""Determine model class and scanner by matching ``model_root`` against the
|
|
configured root paths for each model type (from ``Config``).
|
|
|
|
The ``model_root`` value comes from the frontend's model-root dropdown,
|
|
which is populated from the current page's scanner roots. By checking
|
|
which scanner's root list it belongs to, we avoid fragile heuristics
|
|
like substring-matching path names.
|
|
"""
|
|
norm = os.path.normpath(model_root).replace(os.sep, "/")
|
|
|
|
# LoRA roots
|
|
for p in (config.loras_roots or []) + (config.extra_loras_roots or []):
|
|
if os.path.normpath(p).replace(os.sep, "/") == norm:
|
|
return LoraMetadata, "get_lora_scanner"
|
|
|
|
# Checkpoint / UNet roots
|
|
for p in (
|
|
(config.checkpoints_roots or [])
|
|
+ (config.extra_checkpoints_roots or [])
|
|
+ (config.unet_roots or [])
|
|
+ (config.extra_unet_roots or [])
|
|
):
|
|
if os.path.normpath(p).replace(os.sep, "/") == norm:
|
|
return CheckpointMetadata, "get_checkpoint_scanner"
|
|
|
|
# Embedding roots
|
|
for p in (config.embeddings_roots or []) + (config.extra_embeddings_roots or []):
|
|
if os.path.normpath(p).replace(os.sep, "/") == norm:
|
|
return EmbeddingMetadata, "get_embedding_scanner"
|
|
|
|
# Fallback — should not happen in normal use
|
|
logger.warning(
|
|
"Could not determine model type for root '%s'; defaulting to LoRA",
|
|
model_root,
|
|
)
|
|
return _DEFAULT_MODEL_CLASS, _DEFAULT_SCANNER_GETTER
|
|
|
|
|
|
async def _report_phase(
|
|
download_id: str | None, stage: str, platform: str = ""
|
|
) -> None:
|
|
"""Tell the progress UI which post-transfer stage is running.
|
|
|
|
A download's byte counter stops the moment the last byte lands, but the
|
|
backend still has to index the file and read the model site's API. Without
|
|
this the bar sits at 100% reporting "0 B/s" and the download looks stuck for
|
|
several seconds. *stage* is machine-readable — the UI localises it — and
|
|
*platform* lets it name the site the metadata comes from.
|
|
"""
|
|
|
|
if not download_id:
|
|
return
|
|
try:
|
|
await ws_manager.broadcast_download_progress(
|
|
download_id,
|
|
{
|
|
"status": "metadata",
|
|
"stage": stage,
|
|
"platform": platform,
|
|
"progress": 100,
|
|
},
|
|
)
|
|
except Exception as exc: # pragma: no cover - progress must never be fatal
|
|
logger.debug("Failed to report the '%s' phase: %s", stage, exc)
|
|
|
|
|
|
async def _save_source_metadata(
|
|
dest_path: str, ref: SourceRef, model_root: str, *, download_id: str | None = None
|
|
) -> None:
|
|
"""Create a proper .metadata.json and add the model to the scanner cache.
|
|
|
|
The metadata is created through the owning scanner rather than
|
|
``MetadataManager.create_default_metadata()``, because that is the only
|
|
factory that knows when hashing must be deferred: ``CheckpointScanner`` and
|
|
``OtherScanner`` deliberately record ``hash_status="pending"`` with an empty
|
|
``sha256`` for their multi-GB files, and the generic helper would read a
|
|
10 GB checkpoint end to end *inside the download request*. Scanners for the
|
|
small types delegate straight back to it, so nothing changes for them.
|
|
|
|
The external-source fields are then overlaid and the model is registered in
|
|
the in-memory scanner cache so it appears immediately without a full
|
|
filesystem walk.
|
|
|
|
Finally the site's own published metadata is applied (see
|
|
:func:`~py.services.model_sources.hydration.hydrate_from_source`), so a
|
|
ModelScope or Hugging Face download lands with the same populated model
|
|
card a CivitAI download produces instead of a bare filename and hash.
|
|
|
|
Both post-transfer stages are reported through *download_id* when the UI is
|
|
watching one, because neither advances the byte counter.
|
|
"""
|
|
try:
|
|
model_class, scanner_getter_name = _infer_model_type(model_root)
|
|
|
|
scanner = None
|
|
scanner_getter = getattr(ServiceRegistry, scanner_getter_name, None)
|
|
if scanner_getter is not None:
|
|
scanner = await scanner_getter()
|
|
|
|
# 1. Create proper metadata (reads safetensors headers; hashes only for
|
|
# the model types whose scanner does not defer it)
|
|
await _report_phase(download_id, "indexing", ref.platform)
|
|
create_metadata = getattr(scanner, "_create_default_metadata", None)
|
|
if create_metadata is not None:
|
|
metadata = await create_metadata(dest_path)
|
|
else:
|
|
metadata = await MetadataManager.create_default_metadata(
|
|
dest_path, model_class=model_class
|
|
)
|
|
if metadata is None:
|
|
logger.warning("create_default_metadata returned None for %s", dest_path)
|
|
return
|
|
|
|
# 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 %s metadata (source=%s, hash_status=%s) for %s",
|
|
ref.platform, ref.url, getattr(metadata, "hash_status", "?"), dest_path,
|
|
)
|
|
|
|
# 4. Determine relative folder path for cache
|
|
# model_root is an absolute path; dest_path is under it
|
|
folder = ""
|
|
if os.path.isabs(model_root) and dest_path.startswith(model_root):
|
|
rel = os.path.relpath(os.path.dirname(dest_path), model_root)
|
|
folder = rel.replace(os.sep, "/") if rel != "." else ""
|
|
|
|
# 5. Add to scanner cache (same as CivitAI's _execute_download does)
|
|
if scanner is not None:
|
|
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)
|
|
|
|
# 6. Top up from the site's public API. Runs last so the scanner-cache
|
|
# refresh it performs lands on the entry created above. It never
|
|
# raises and never fails the download.
|
|
await _report_phase(download_id, "source", ref.platform)
|
|
await hydrate_from_source(dest_path, ref=ref)
|
|
|
|
except Exception as exc:
|
|
logger.warning("Failed to save source metadata for %s: %s", dest_path, exc)
|
|
|
|
|
|
def _find_matching_root(dest_dir: str) -> str | None:
|
|
"""Walk up *dest_dir* to find which configured scanner root it belongs to."""
|
|
norm = os.path.normpath(dest_dir).replace(os.sep, "/")
|
|
all_roots = []
|
|
for root_list in (
|
|
config.loras_roots or [],
|
|
config.extra_loras_roots or [],
|
|
config.checkpoints_roots or [],
|
|
config.extra_checkpoints_roots or [],
|
|
config.unet_roots or [],
|
|
config.extra_unet_roots or [],
|
|
config.embeddings_roots or [],
|
|
config.extra_embeddings_roots or [],
|
|
):
|
|
all_roots.extend([os.path.normpath(p).replace(os.sep, "/") for p in root_list])
|
|
# Find the longest matching prefix
|
|
match: str | None = None
|
|
for root in all_roots:
|
|
if norm.startswith(root):
|
|
if match is None or len(root) > len(match):
|
|
match = root
|
|
return match
|
|
|
|
|
|
async def _add_to_scanner_cache(dest_path: str, metadata: dict[str, Any]) -> None:
|
|
model_dir = os.path.dirname(dest_path)
|
|
model_root = _find_matching_root(model_dir)
|
|
if not model_root:
|
|
raise ValueError(f"File path {dest_path} is not within any configured scanner root")
|
|
scanner_getter_name = _infer_model_type(model_root)[1]
|
|
scanner_getter = getattr(ServiceRegistry, scanner_getter_name, None)
|
|
if scanner_getter is None:
|
|
raise RuntimeError(f"Scanner getter '{scanner_getter_name}' not found in ServiceRegistry")
|
|
scanner = await scanner_getter()
|
|
if scanner is None:
|
|
raise RuntimeError(f"Scanner '{scanner_getter_name}' returned None")
|
|
await scanner.update_single_model_cache(dest_path, dest_path, metadata)
|
|
|
|
|
|
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, to
|
|
explain which sites support AI metadata enrichment, and to pick the
|
|
right download endpoint/revision.
|
|
"""
|
|
|
|
return web.json_response([
|
|
{
|
|
"platform": source.platform,
|
|
"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"
|
|
),
|
|
}
|
|
for source in list_sources()
|
|
])
|
|
|
|
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. 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:
|
|
payload: dict[str, Any] = await request.json()
|
|
except json.JSONDecodeError:
|
|
return web.json_response({"success": False, "error": "Invalid JSON"}, status=400)
|
|
|
|
file_path = (payload.get("file_path") or "").strip()
|
|
raw_url = (
|
|
payload.get("source_url")
|
|
or payload.get("hf_url")
|
|
or payload.get("url")
|
|
or ""
|
|
)
|
|
source_url = raw_url.strip() if isinstance(raw_url, str) else ""
|
|
|
|
if not file_path or not source_url:
|
|
return web.json_response(
|
|
{
|
|
"success": False,
|
|
"error": "Missing required fields: 'file_path' and 'source_url'",
|
|
},
|
|
status=400,
|
|
)
|
|
|
|
ref = detect_source(source_url, strict=True)
|
|
if ref is None:
|
|
return web.json_response(
|
|
{
|
|
"success": False,
|
|
"error": (
|
|
"Unsupported model URL. Supported formats: "
|
|
+ ", ".join(
|
|
f"{s.label} ({s.canonical_url('user/repo')})"
|
|
if s.platform != "tensorart"
|
|
else f"{s.label} (https://tensor.art/models/<id>)"
|
|
for s in list_sources()
|
|
)
|
|
),
|
|
},
|
|
status=400,
|
|
)
|
|
|
|
if not os.path.isfile(file_path):
|
|
return web.json_response(
|
|
{"success": False, "error": f"File not found: {file_path}"},
|
|
status=404,
|
|
)
|
|
|
|
model_root = _find_matching_root(os.path.dirname(file_path))
|
|
if not model_root:
|
|
return web.json_response(
|
|
{
|
|
"success": False,
|
|
"error": "File is not within any configured model directory. Cannot link to a model source.",
|
|
},
|
|
status=400,
|
|
)
|
|
|
|
try:
|
|
existing = await MetadataManager.load_metadata_payload(file_path)
|
|
|
|
already_linked = (
|
|
(existing.get("source_url") or "").strip() == ref.url
|
|
and (existing.get("source_platform") or "").strip().lower()
|
|
== ref.platform
|
|
) or (
|
|
not existing.get("source_url")
|
|
and ref.platform == "huggingface"
|
|
and (existing.get("hf_url") or "").strip() == ref.url
|
|
)
|
|
if already_linked:
|
|
return web.json_response({
|
|
"success": True,
|
|
"message": "source_url already set",
|
|
"source_url": ref.url,
|
|
"source_platform": ref.platform,
|
|
"hf_url": ref.url if ref.platform == "huggingface" else "",
|
|
})
|
|
|
|
existing["source_url"] = ref.url
|
|
existing["source_platform"] = ref.platform
|
|
if ref.platform == "huggingface":
|
|
existing["hf_url"] = ref.url
|
|
else:
|
|
existing.pop("hf_url", None)
|
|
normalize_metadata_source(existing)
|
|
|
|
# NOTE: deliberately do NOT touch `from_civitai` here. It records
|
|
# where the metadata came from, and the UI must show the CivitAI
|
|
# link whenever CivitAI data is present — linking an external
|
|
# source must not hide it (#1094). Source provenance is tracked
|
|
# via `source_platform` / `source_url`.
|
|
await MetadataManager.save_metadata(file_path, existing)
|
|
|
|
await _add_to_scanner_cache(file_path, existing)
|
|
|
|
logger.info(
|
|
"Linked %s to %s source (%s)", file_path, ref.platform, ref.url
|
|
)
|
|
return web.json_response({
|
|
"success": True,
|
|
"message": f"Linked to {ref.url}",
|
|
"source_url": ref.url,
|
|
"source_platform": ref.platform,
|
|
"hf_url": existing.get("hf_url", ""),
|
|
})
|
|
except Exception as exc:
|
|
logger.error("Failed to link %s to a model source: %s", file_path, exc)
|
|
return web.json_response(
|
|
{"success": False, "error": str(exc)},
|
|
status=500,
|
|
)
|
|
|
|
async def list_model_source_files(self, request: web.Request) -> web.Response:
|
|
"""List the downloadable weight files of an external repository.
|
|
|
|
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.
|
|
"""
|
|
|
|
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 owner/name)"},
|
|
status=400,
|
|
)
|
|
|
|
try:
|
|
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 list %s files in %s: %s", platform, repo, exc)
|
|
return web.json_response({"error": str(exc)}, status=502)
|
|
|
|
return web.json_response(files)
|
|
|
|
async def download_model_source(self, request: web.Request) -> web.Response:
|
|
"""Download a single file from an external repository.
|
|
|
|
POST JSON body::
|
|
|
|
{
|
|
"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.
|
|
|
|
Respects the ``download_backend`` setting (``aria2`` or ``default``).
|
|
"""
|
|
try:
|
|
payload: dict[str, Any] = await request.json()
|
|
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 "").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_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
|
|
)
|
|
|
|
# `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:
|
|
return web.json_response({"error": "Invalid filename"}, status=400)
|
|
|
|
# Validate relative_path — must not be absolute or escape base directory
|
|
if relative_path:
|
|
if os.path.isabs(relative_path):
|
|
return web.json_response({"error": "relative_path must not be absolute"}, status=400)
|
|
if ".." in relative_path.split("/") or "\\" in relative_path:
|
|
return web.json_response({"error": "Invalid relative_path"}, status=400)
|
|
|
|
# Use model_root directly as the base directory — same approach as
|
|
# CivitAI's download path (download_manager.py). No realpath, no
|
|
# allowed-roots validation, no path-traversal check; those are
|
|
# 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_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, source.default_subdir, owner, repo_name)
|
|
elif relative_path:
|
|
target_dir = os.path.join(base_dir, relative_path)
|
|
else:
|
|
target_dir = base_dir
|
|
|
|
# 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)
|
|
dest_path = os.path.join(target_dir, file_base)
|
|
|
|
# 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)
|
|
)
|
|
|
|
# Check if already exists (simple skip)
|
|
if os.path.exists(dest_path) and os.path.getsize(dest_path) > 0:
|
|
logger.info("download_model_source: file already exists, skipping — %s", dest_path)
|
|
# The sidecar may predate the source metadata being fetched, or may
|
|
# have been deleted, so top it up instead of skipping past it.
|
|
# Hydration no-ops when there is no sidecar to update.
|
|
await _report_phase(download_id, "source", source.platform)
|
|
await hydrate_from_source(dest_path, ref=ref)
|
|
return web.json_response({
|
|
"success": True,
|
|
"message": f"File already exists: {dest_path}",
|
|
"path": dest_path,
|
|
})
|
|
|
|
# Set up progress callback if download_id is provided
|
|
progress_callback = None
|
|
if download_id:
|
|
|
|
async def _progress_callback(
|
|
progress: float | DownloadProgress,
|
|
snapshot: DownloadProgress | None = None,
|
|
) -> None:
|
|
percent = 0.0
|
|
metrics = snapshot if isinstance(snapshot, DownloadProgress) else None
|
|
|
|
if isinstance(progress, DownloadProgress):
|
|
percent = progress.percent_complete
|
|
metrics = progress
|
|
elif isinstance(snapshot, DownloadProgress):
|
|
percent = snapshot.percent_complete
|
|
else:
|
|
percent = float(progress)
|
|
|
|
broadcast: dict[str, Any] = {
|
|
"status": "progress",
|
|
"progress": round(percent),
|
|
}
|
|
if metrics:
|
|
broadcast["bytes_downloaded"] = metrics.bytes_downloaded
|
|
broadcast["total_bytes"] = metrics.total_bytes
|
|
broadcast["bytes_per_second"] = metrics.bytes_per_second
|
|
|
|
await ws_manager.broadcast_download_progress(download_id, broadcast)
|
|
|
|
progress_callback = _progress_callback
|
|
|
|
# Respect download backend setting (aria2 vs default)
|
|
download_backend = (
|
|
get_settings_manager().get("download_backend", "default")
|
|
)
|
|
|
|
if download_backend == "aria2":
|
|
aria2 = await Aria2Downloader.get_instance()
|
|
aid = download_id or f"{source.platform}_{repo}_{filename}"
|
|
try:
|
|
ok, result = await aria2.download_file(
|
|
url=resolve_url,
|
|
save_path=dest_path,
|
|
download_id=aid,
|
|
progress_callback=progress_callback,
|
|
)
|
|
if ok:
|
|
await _save_source_metadata(
|
|
dest_path, ref, model_root, download_id=download_id
|
|
)
|
|
return web.json_response({
|
|
"success": True,
|
|
"message": f"Downloaded to {dest_path}",
|
|
"path": dest_path,
|
|
})
|
|
return web.json_response(
|
|
{"success": False, "error": result or "aria2 download failed"},
|
|
status=500,
|
|
)
|
|
except Exception as exc:
|
|
logger.error("%s download (aria2) failed: %s", platform, exc)
|
|
return web.json_response(
|
|
{"success": False, "error": str(exc)}, status=500
|
|
)
|
|
|
|
# Default: use built-in aiohttp Downloader
|
|
downloader = await get_downloader()
|
|
try:
|
|
success, result = await downloader.download_file(
|
|
url=resolve_url,
|
|
save_path=dest_path,
|
|
use_auth=False,
|
|
allow_resume=True,
|
|
progress_callback=progress_callback,
|
|
)
|
|
if success:
|
|
await _save_source_metadata(
|
|
dest_path, ref, model_root, download_id=download_id
|
|
)
|
|
return web.json_response({
|
|
"success": True,
|
|
"message": f"Downloaded to {result}",
|
|
"path": result,
|
|
})
|
|
return web.json_response(
|
|
{"success": False, "error": result or "Download failed"},
|
|
status=500,
|
|
)
|
|
except Exception as exc:
|
|
logger.error("%s download failed: %s", platform, exc)
|
|
return web.json_response(
|
|
{"success": False, "error": str(exc)}, status=500
|
|
)
|