feat: add StabilityMatrix-style Model Manager service
New FastAPI container (port 8189) to download and manage models: - Installed Models, Add/Download (CivitAI/HuggingFace/direct URL), Settings views - Persistent SQLite storage for API keys and download history (./sparkyui-data) - Downloads land in ./models, auto-sorted into ComfyUI's standard subfolders - Default COMFYUI_HOST_PATH and SPARKYUI_DATA_PATH to the project root - Wire docker-compose service, env defaults, gitignore, README docs Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,142 @@
|
||||
"""Resolve a user-supplied URL into a concrete download (direct URL + auth + metadata).
|
||||
|
||||
Supports direct URLs, CivitAI (model/version pages or api download links), and
|
||||
HuggingFace `resolve` URLs. API keys are read from the settings table.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
from urllib.parse import unquote, urlparse, parse_qs
|
||||
|
||||
import httpx
|
||||
|
||||
from . import db
|
||||
|
||||
# CivitAI model `type` -> our model-type key.
|
||||
CIVITAI_TYPE_MAP = {
|
||||
"Checkpoint": "checkpoint",
|
||||
"LORA": "lora",
|
||||
"LoCon": "lora",
|
||||
"DoRA": "lora",
|
||||
"TextualInversion": "embedding",
|
||||
"Hypernetwork": "hypernetwork",
|
||||
"VAE": "vae",
|
||||
"Controlnet": "controlnet",
|
||||
"Upscaler": "upscaler",
|
||||
"MotionModule": "other",
|
||||
"Poses": "other",
|
||||
"Wildcards": "other",
|
||||
"Other": "other",
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class Resolved:
|
||||
download_url: str
|
||||
source: str
|
||||
headers: dict[str, str] = field(default_factory=dict)
|
||||
filename: Optional[str] = None
|
||||
model_type: Optional[str] = None # suggested type if the registry tells us
|
||||
|
||||
|
||||
def detect_source(url: str) -> str:
|
||||
host = (urlparse(url).hostname or "").lower()
|
||||
if "civitai.com" in host:
|
||||
return "civitai"
|
||||
if "huggingface.co" in host or "hf.co" in host:
|
||||
return "huggingface"
|
||||
return "direct"
|
||||
|
||||
|
||||
def _filename_from_url(url: str) -> Optional[str]:
|
||||
path = urlparse(url).path
|
||||
name = unquote(path.rsplit("/", 1)[-1]) if path else ""
|
||||
return name or None
|
||||
|
||||
|
||||
async def resolve(url: str) -> Resolved:
|
||||
source = detect_source(url)
|
||||
if source == "civitai":
|
||||
return await _resolve_civitai(url)
|
||||
if source == "huggingface":
|
||||
return _resolve_huggingface(url)
|
||||
return _resolve_direct(url)
|
||||
|
||||
|
||||
def _resolve_direct(url: str) -> Resolved:
|
||||
return Resolved(download_url=url, source="direct",
|
||||
filename=_filename_from_url(url))
|
||||
|
||||
|
||||
def _resolve_huggingface(url: str) -> Resolved:
|
||||
headers: dict[str, str] = {}
|
||||
token = db.get_setting("huggingface_token")
|
||||
if token:
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
return Resolved(download_url=url, source="huggingface", headers=headers,
|
||||
filename=_filename_from_url(url))
|
||||
|
||||
|
||||
_CIVITAI_VERSION_RE = re.compile(r"/api/download/models/(\d+)")
|
||||
_CIVITAI_MODEL_RE = re.compile(r"/models/(\d+)")
|
||||
|
||||
|
||||
async def _resolve_civitai(url: str) -> Resolved:
|
||||
headers: dict[str, str] = {}
|
||||
api_key = db.get_setting("civitai_api_key")
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
|
||||
# Case 1: already a direct api/download URL -> use as-is.
|
||||
if _CIVITAI_VERSION_RE.search(urlparse(url).path):
|
||||
return Resolved(download_url=url, source="civitai", headers=headers,
|
||||
filename=_filename_from_url(url))
|
||||
|
||||
# Case 2: a model page URL. Find the version id (explicit query or first version).
|
||||
parsed = urlparse(url)
|
||||
qs = parse_qs(parsed.query)
|
||||
version_id: Optional[str] = None
|
||||
if "modelVersionId" in qs:
|
||||
version_id = qs["modelVersionId"][0]
|
||||
|
||||
if version_id is None:
|
||||
m = _CIVITAI_MODEL_RE.search(parsed.path)
|
||||
if not m:
|
||||
# Can't understand it; fall back to treating it as a direct link.
|
||||
return Resolved(download_url=url, source="civitai", headers=headers,
|
||||
filename=_filename_from_url(url))
|
||||
model_id = m.group(1)
|
||||
async with httpx.AsyncClient(timeout=30, headers=headers) as client:
|
||||
resp = await client.get(f"https://civitai.com/api/v1/models/{model_id}")
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
versions = data.get("modelVersions") or []
|
||||
if not versions:
|
||||
raise ValueError("CivitAI model has no downloadable versions")
|
||||
version_id = str(versions[0]["id"])
|
||||
|
||||
# Resolve the version to a concrete file + type.
|
||||
async with httpx.AsyncClient(timeout=30, headers=headers) as client:
|
||||
resp = await client.get(
|
||||
f"https://civitai.com/api/v1/model-versions/{version_id}")
|
||||
resp.raise_for_status()
|
||||
version = resp.json()
|
||||
|
||||
files = version.get("files") or []
|
||||
# Prefer the primary file, else the first.
|
||||
chosen = next((f for f in files if f.get("primary")), files[0] if files else None)
|
||||
if not chosen:
|
||||
raise ValueError("CivitAI version has no files")
|
||||
|
||||
civ_type = (version.get("model") or {}).get("type") or "Other"
|
||||
model_type = CIVITAI_TYPE_MAP.get(civ_type, "other")
|
||||
|
||||
return Resolved(
|
||||
download_url=chosen["downloadUrl"],
|
||||
source="civitai",
|
||||
headers=headers,
|
||||
filename=chosen.get("name") or _filename_from_url(chosen["downloadUrl"]),
|
||||
model_type=model_type,
|
||||
)
|
||||
Reference in New Issue
Block a user