diff --git a/.env.example b/.env.example index 3c08612..e01c52b 100644 --- a/.env.example +++ b/.env.example @@ -1,11 +1,13 @@ # SparkyUI - ComfyUI for DGX Spark (Blackwell GB10) # Copy this to .env and customize paths as needed -# Base path where your existing ComfyUI installation lives (for models) -COMFYUI_HOST_PATH=/path/to/your/ComfyUI +# Base path that holds the models/ directory (defaults to the project root). +# The Model Manager downloads into /models and ComfyUI reads it. +COMFYUI_HOST_PATH=. -# Base path for SparkyUI data (custom_nodes, outputs, inputs, etc.) -SPARKYUI_DATA_PATH=/path/to/SparkyUI +# Base path for SparkyUI data (custom_nodes, outputs, inputs, manager DB, etc.) +# Defaults to the project root. +SPARKYUI_DATA_PATH=. # ComfyUI settings COMFYUI_PORT=8188 @@ -25,3 +27,7 @@ SAGEATTN_REF=main # Access at http://:3000 COMFYUIMINI_PORT=3000 COMFYUIMINI_REF=main + +# Model Manager - StabilityMatrix-style model download/management UI +# Access at http://:8189 +MODEL_MANAGER_PORT=8189 diff --git a/.gitignore b/.gitignore index 480ef5f..12fa27b 100644 --- a/.gitignore +++ b/.gitignore @@ -54,6 +54,13 @@ input/* workflows/* !workflows/.gitkeep +# Downloaded models - ignore contents but keep the directory +models/* +!models/.gitkeep + +# Model Manager persistent data (SQLite DB, stored API keys) +sparkyui-data/ + # Wheels directory - for prebuilt ARM64/sm_121 binaries # Ignore contents except .gitkeep (add wheels explicitly if needed) wheels/* diff --git a/README.md b/README.md index 7637882..b960c1a 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,7 @@ Standard ComfyUI containers and PyTorch wheels don't support sm_121. SparkyUI so - **ComfyUI** (latest master branch) - **ComfyUI-Manager** - auto-installed on first run for easy custom node management - **ComfyUIMini** - mobile-friendly web UI for phones/tablets (separate container) +- **Model Manager** - StabilityMatrix-style UI to download/manage models (separate container) - **SageAttention** - compiled natively for sm_121 (Blackwell tensor cores) - **PyTorch 2.9.1+cu130** - ARM64 wheels with CUDA 13.0 support @@ -77,6 +78,7 @@ docker compose logs -f **Access:** - **ComfyUI (Desktop):** http://localhost:8188 - **ComfyUIMini (Mobile):** http://localhost:3000 +- **Model Manager:** http://localhost:8189 ## Requirements @@ -91,17 +93,28 @@ docker compose logs -f Copy `.env.example` to `.env` and edit: ```bash -# Path to your existing ComfyUI models (mounted read-only) -COMFYUI_HOST_PATH=/path/to/your/ComfyUI +# Base path holding the models/ directory (defaults to the project root). +# The Model Manager downloads into /models; ComfyUI reads it. +COMFYUI_HOST_PATH=. -# Path for SparkyUI data (custom_nodes, outputs, inputs) -SPARKYUI_DATA_PATH=/path/to/SparkyUI +# Path for SparkyUI data (custom_nodes, outputs, inputs, manager DB). +# Defaults to the project root. +SPARKYUI_DATA_PATH=. + +# Ports +COMFYUI_PORT=8188 +COMFYUIMINI_PORT=3000 +MODEL_MANAGER_PORT=8189 # Optional: pin to specific versions COMFYUI_REF=master SAGEATTN_REF=main ``` +Both paths default to the project root, so out of the box models are stored in +`./models` and the Model Manager's database in `./sparkyui-data`. Point +`COMFYUI_HOST_PATH` at an existing ComfyUI install if you'd rather reuse its models. + ## Architecture ``` @@ -214,6 +227,36 @@ docker compose build comfyuimini docker compose up -d comfyuimini ``` +## Model Manager + +SparkyUI includes a **StabilityMatrix-style Model Manager** - a lightweight FastAPI web app +(separate container) for downloading and managing models without touching the command line. + +**Access:** `http://:8189` + +**Features:** +- **Installed Models** - browse what's on disk, grouped by type, with size and delete actions +- **Add / Download** - paste a download URL and pick a type; live progress bars + - **Direct URLs** - any direct download link + - **CivitAI** - paste a model page link (`civitai.com/models/...`) or an + `api/download/models/...` link; the type and filename are auto-detected + - **HuggingFace** - paste a `resolve` URL (works with gated repos via your token) +- **Settings** - store your **CivitAI API key** and **HuggingFace token** persistently + (saved to a SQLite DB under `./sparkyui-data`, never committed to git) + +**How it works:** +- Runs as a FastAPI server in its own container (`python:3.12-slim`) +- Downloads land in the shared `models/` folder, sorted into ComfyUI's standard sub-folders + by type (`checkpoints/`, `loras/`, `vae/`, `controlnet/`, `upscale_models/`, …) - these are + created automatically on first start +- ComfyUI mounts the same `models/` folder read-only, so new downloads appear in its loaders + +**Build only the Model Manager** (if the rest is already built): +```bash +docker compose build model-manager +docker compose up -d model-manager +``` + ## SageAttention Notes SageAttention PR #297 added sm_121 support but was merged then reverted due to stability issues. Our approach: diff --git a/docker-compose.yml b/docker-compose.yml index 6448bc1..0f1c1c5 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -50,20 +50,21 @@ services: CUDA_CACHE_MAXSIZE: "4294967296" volumes: - # Models from existing ComfyUI install (read-only) - - ${COMFYUI_HOST_PATH}/models:/opt/ComfyUI/models:ro + # Models from existing ComfyUI install (read-only). + # Defaults to the project root; the model-manager service writes here. + - ${COMFYUI_HOST_PATH:-.}/models:/opt/ComfyUI/models:ro # Custom nodes - comment out to use container-only (fresh) custom_nodes # If mounted, ComfyUI-Manager installs persist across container restarts - - ${SPARKYUI_DATA_PATH}/custom_nodes:/opt/ComfyUI/custom_nodes + - ${SPARKYUI_DATA_PATH:-.}/custom_nodes:/opt/ComfyUI/custom_nodes # Outputs/inputs/workflows - persistent across restarts - - ${SPARKYUI_DATA_PATH}/output:/opt/ComfyUI/output - - ${SPARKYUI_DATA_PATH}/input:/opt/ComfyUI/input - - ${SPARKYUI_DATA_PATH}/workflows:/opt/ComfyUI/workflows + - ${SPARKYUI_DATA_PATH:-.}/output:/opt/ComfyUI/output + - ${SPARKYUI_DATA_PATH:-.}/input:/opt/ComfyUI/input + - ${SPARKYUI_DATA_PATH:-.}/workflows:/opt/ComfyUI/workflows # Wheel cache (optional - for prebuilt wheels) - - ${SPARKYUI_DATA_PATH}/wheels:/opt/wheels + - ${SPARKYUI_DATA_PATH:-.}/wheels:/opt/wheels # Sparky patches - Grace-Blackwell unified memory optimizations # model_management.py: HIGH_VRAM→NORMAL_VRAM, intermediate_device()→cuda, soft_empty_cache skip, @@ -114,7 +115,7 @@ services: volumes: # Share output directory with ComfyUI for gallery feature (read-only) - - ${SPARKYUI_DATA_PATH}/output:/shared/output:ro + - ${SPARKYUI_DATA_PATH:-.}/output:/shared/output:ro # Persist server-side workflows - comfyuimini_workflows:/app/workflows @@ -127,6 +128,42 @@ services: restart: unless-stopped + # Model Manager - StabilityMatrix-style model download/management UI + # Access at http://:8189 + model-manager: + build: + context: ./model-manager + dockerfile: Dockerfile + + image: sparkyui-model-manager:latest + container_name: model-manager + + ports: + - "${MODEL_MANAGER_PORT:-8189}:8189" + + environment: + MODELS_DIR: /models + DATA_DIR: /data + + volumes: + # Shared models dir - read-WRITE here so downloads land on the host. + # ComfyUI mounts the same host folder read-only and picks up new files. + - ${COMFYUI_HOST_PATH:-.}/models:/models + # Persistent SQLite DB (sources, API keys, download history) + - ${SPARKYUI_DATA_PATH:-.}/sparkyui-data:/data + + networks: + - sparky_net + + healthcheck: + test: ["CMD", "curl", "-f", "http://localhost:8189/api/model-types"] + interval: 30s + timeout: 10s + start_period: 15s + retries: 3 + + restart: unless-stopped + networks: sparky_net: driver: bridge diff --git a/model-manager/.dockerignore b/model-manager/.dockerignore new file mode 100644 index 0000000..709965c --- /dev/null +++ b/model-manager/.dockerignore @@ -0,0 +1,9 @@ +__pycache__/ +*.py[cod] +*.egg-info/ +.venv/ +venv/ +.pytest_cache/ +*.db +*.db-wal +*.db-shm diff --git a/model-manager/Dockerfile b/model-manager/Dockerfile new file mode 100644 index 0000000..0e92ea1 --- /dev/null +++ b/model-manager/Dockerfile @@ -0,0 +1,28 @@ +# SparkyUI Model Manager - StabilityMatrix-style model download/management UI +# Lightweight FastAPI service. Downloads land in the shared host models/ folder +# (mounted read-write here, read-only in the ComfyUI container). +FROM python:3.12-slim + +LABEL maintainer="SparkyUI" +LABEL description="Model Manager for ComfyUI on DGX Spark" + +WORKDIR /app + +# curl is used by the healthcheck. +RUN apt-get update && apt-get install -y --no-install-recommends curl \ + && rm -rf /var/lib/apt/lists/* + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY app ./app + +ENV MODELS_DIR=/models \ + DATA_DIR=/data + +EXPOSE 8189 + +HEALTHCHECK --interval=30s --timeout=10s --start-period=10s --retries=3 \ + CMD curl -f http://localhost:8189/api/model-types || exit 1 + +CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8189"] diff --git a/model-manager/app/__init__.py b/model-manager/app/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/model-manager/app/config.py b/model-manager/app/config.py new file mode 100644 index 0000000..c231682 --- /dev/null +++ b/model-manager/app/config.py @@ -0,0 +1,53 @@ +"""Paths, model-type mapping, and folder bootstrap for the Model Manager.""" +from __future__ import annotations + +import os +from pathlib import Path + +# Root of the host `models/` directory (mounted read-write into this container). +MODELS_DIR = Path(os.environ.get("MODELS_DIR", "/models")).resolve() + +# Persistent data dir for the SQLite database. +DATA_DIR = Path(os.environ.get("DATA_DIR", "/data")).resolve() + +DB_PATH = DATA_DIR / "manager.db" + +# StabilityMatrix-style model types mapped to ComfyUI's standard `models/` subfolders. +# `key` is the stable id used by the API/UI; `folder` is the on-disk subdirectory. +MODEL_TYPES: list[dict[str, str]] = [ + {"key": "checkpoint", "label": "Checkpoint", "folder": "checkpoints"}, + {"key": "lora", "label": "LoRA / LyCORIS", "folder": "loras"}, + {"key": "vae", "label": "VAE", "folder": "vae"}, + {"key": "embedding", "label": "Textual Inversion", "folder": "embeddings"}, + {"key": "controlnet", "label": "ControlNet", "folder": "controlnet"}, + {"key": "upscaler", "label": "Upscaler", "folder": "upscale_models"}, + {"key": "clip", "label": "CLIP", "folder": "clip"}, + {"key": "clip_vision", "label": "CLIP Vision", "folder": "clip_vision"}, + {"key": "text_encoder", "label": "Text Encoder", "folder": "text_encoders"}, + {"key": "diffusion_model", "label": "Diffusion Model (UNET)", "folder": "diffusion_models"}, + {"key": "unet", "label": "UNET", "folder": "unet"}, + {"key": "hypernetwork", "label": "Hypernetwork", "folder": "hypernetworks"}, + {"key": "style_model", "label": "Style Model", "folder": "style_models"}, + {"key": "gligen", "label": "GLIGEN", "folder": "gligen"}, + {"key": "vae_approx", "label": "VAE Approx", "folder": "vae_approx"}, + {"key": "ipadapter", "label": "IP-Adapter", "folder": "ipadapter"}, + {"key": "other", "label": "Other", "folder": "other"}, +] + +# Quick lookups. +TYPE_BY_KEY: dict[str, dict[str, str]] = {t["key"]: t for t in MODEL_TYPES} +FOLDER_BY_KEY: dict[str, str] = {t["key"]: t["folder"] for t in MODEL_TYPES} +KEY_BY_FOLDER: dict[str, str] = {t["folder"]: t["key"] for t in MODEL_TYPES} + + +def folder_for_type(type_key: str) -> str: + """Return the on-disk subfolder for a model type key, defaulting to `other`.""" + return FOLDER_BY_KEY.get(type_key, "other") + + +def ensure_dirs() -> None: + """Create the data dir and all standard model subfolders if missing.""" + DATA_DIR.mkdir(parents=True, exist_ok=True) + MODELS_DIR.mkdir(parents=True, exist_ok=True) + for t in MODEL_TYPES: + (MODELS_DIR / t["folder"]).mkdir(parents=True, exist_ok=True) diff --git a/model-manager/app/db.py b/model-manager/app/db.py new file mode 100644 index 0000000..4f28805 --- /dev/null +++ b/model-manager/app/db.py @@ -0,0 +1,131 @@ +"""SQLite persistence: settings (API keys) and download history/queue.""" +from __future__ import annotations + +import sqlite3 +import threading +import time +from typing import Any, Optional + +from .config import DB_PATH, ensure_dirs + +# A single connection guarded by a lock keeps things simple and safe across the +# async event loop + background download tasks (sqlite writes are serialized). +_lock = threading.Lock() +_conn: Optional[sqlite3.Connection] = None + + +def _connect() -> sqlite3.Connection: + global _conn + if _conn is None: + ensure_dirs() + _conn = sqlite3.connect(str(DB_PATH), check_same_thread=False) + _conn.row_factory = sqlite3.Row + _conn.execute("PRAGMA journal_mode=WAL") + return _conn + + +def init_db() -> None: + with _lock: + conn = _connect() + conn.execute( + """ + CREATE TABLE IF NOT EXISTS settings ( + key TEXT PRIMARY KEY, + value TEXT + ) + """ + ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS downloads ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + url TEXT NOT NULL, + source TEXT, + model_type TEXT, + filename TEXT, + dest_path TEXT, + status TEXT NOT NULL DEFAULT 'queued', + bytes_done INTEGER NOT NULL DEFAULT 0, + bytes_total INTEGER NOT NULL DEFAULT 0, + error TEXT, + created_at REAL NOT NULL, + updated_at REAL NOT NULL + ) + """ + ) + conn.commit() + + +# ---- settings ------------------------------------------------------------- + +def get_setting(key: str) -> Optional[str]: + with _lock: + row = _connect().execute( + "SELECT value FROM settings WHERE key = ?", (key,) + ).fetchone() + return row["value"] if row else None + + +def set_setting(key: str, value: Optional[str]) -> None: + with _lock: + conn = _connect() + if value is None or value == "": + conn.execute("DELETE FROM settings WHERE key = ?", (key,)) + else: + conn.execute( + "INSERT INTO settings(key, value) VALUES(?, ?) " + "ON CONFLICT(key) DO UPDATE SET value = excluded.value", + (key, value), + ) + conn.commit() + + +# ---- downloads ------------------------------------------------------------ + +def create_download(url: str, source: str, model_type: str, filename: str, + dest_path: str) -> int: + now = time.time() + with _lock: + conn = _connect() + cur = conn.execute( + "INSERT INTO downloads(url, source, model_type, filename, dest_path, " + "status, created_at, updated_at) VALUES(?, ?, ?, ?, ?, 'queued', ?, ?)", + (url, source, model_type, filename, dest_path, now, now), + ) + conn.commit() + return int(cur.lastrowid) + + +def update_download(download_id: int, **fields: Any) -> None: + if not fields: + return + fields["updated_at"] = time.time() + cols = ", ".join(f"{k} = ?" for k in fields) + values = list(fields.values()) + [download_id] + with _lock: + conn = _connect() + conn.execute(f"UPDATE downloads SET {cols} WHERE id = ?", values) + conn.commit() + + +def get_download(download_id: int) -> Optional[dict]: + with _lock: + row = _connect().execute( + "SELECT * FROM downloads WHERE id = ?", (download_id,) + ).fetchone() + return dict(row) if row else None + + +def list_downloads() -> list[dict]: + with _lock: + rows = _connect().execute( + "SELECT * FROM downloads ORDER BY id DESC" + ).fetchall() + return [dict(r) for r in rows] + + +def delete_download(download_id: int) -> None: + with _lock: + conn = _connect() + conn.execute("DELETE FROM downloads WHERE id = ?", (download_id,)) + conn.commit() diff --git a/model-manager/app/downloader.py b/model-manager/app/downloader.py new file mode 100644 index 0000000..f721cc8 --- /dev/null +++ b/model-manager/app/downloader.py @@ -0,0 +1,129 @@ +"""Async streaming downloader with progress tracking and cancellation.""" +from __future__ import annotations + +import asyncio +import os +import re +from pathlib import Path +from typing import Optional + +import httpx + +from . import db +from .config import MODELS_DIR, folder_for_type + +# download_id -> asyncio.Event set when a cancel is requested. +_cancel_flags: dict[int, asyncio.Event] = {} + +_FILENAME_STAR_RE = re.compile(r"filename\*=(?:UTF-8'')?([^;]+)", re.IGNORECASE) +_FILENAME_RE = re.compile(r'filename="?([^";]+)"?', re.IGNORECASE) + + +def sanitize_filename(name: str) -> str: + """Strip any directory components and unsafe characters from a filename.""" + name = os.path.basename(name.strip().strip('"')) + name = name.replace("\x00", "") + # Disallow path separators / parent refs that survived basename on other OSes. + name = name.replace("/", "_").replace("\\", "_") + if name in ("", ".", ".."): + name = "model.bin" + return name + + +def _filename_from_disposition(value: str) -> Optional[str]: + from urllib.parse import unquote + m = _FILENAME_STAR_RE.search(value) + if m: + return unquote(m.group(1)) + m = _FILENAME_RE.search(value) + if m: + return m.group(1) + return None + + +def safe_dest(model_type: str, filename: str) -> Path: + """Build a destination path under MODELS_DIR/, guarding traversal.""" + folder = folder_for_type(model_type) + filename = sanitize_filename(filename) + base = (MODELS_DIR / folder).resolve() + dest = (base / filename).resolve() + if not str(dest).startswith(str(base) + os.sep): + raise ValueError("Refusing path outside the model folder") + return dest + + +def request_cancel(download_id: int) -> bool: + """Signal an in-flight download to stop. Returns True if it was active.""" + ev = _cancel_flags.get(download_id) + if ev is not None: + ev.set() + return True + return False + + +async def run_download(download_id: int, url: str, headers: dict[str, str], + model_type: str, filename: Optional[str]) -> None: + """Stream `url` to disk, updating the DB row as it progresses.""" + cancel = asyncio.Event() + _cancel_flags[download_id] = cancel + part_path: Optional[Path] = None + try: + db.update_download(download_id, status="downloading") + async with httpx.AsyncClient(follow_redirects=True, timeout=None, + headers=headers) as client: + async with client.stream("GET", url) as resp: + resp.raise_for_status() + + # Prefer a server-provided filename if we don't have a good one. + disp = resp.headers.get("content-disposition") + if disp: + server_name = _filename_from_disposition(disp) + if server_name: + filename = server_name + if not filename: + filename = os.path.basename(str(resp.url).split("?")[0]) or "model.bin" + + dest = safe_dest(model_type, filename) + part_path = dest.with_suffix(dest.suffix + ".part") + total = int(resp.headers.get("content-length", 0) or 0) + db.update_download(download_id, filename=dest.name, + dest_path=str(dest), bytes_total=total) + + done = 0 + last_report = 0.0 + with open(part_path, "wb") as fh: + async for chunk in resp.aiter_bytes(chunk_size=1024 * 256): + if cancel.is_set(): + raise asyncio.CancelledError() + fh.write(chunk) + done += len(chunk) + # Throttle DB writes to ~5/sec. + now = asyncio.get_event_loop().time() + if now - last_report > 0.2: + db.update_download(download_id, bytes_done=done) + last_report = now + + db.update_download(download_id, bytes_done=done) + os.replace(part_path, dest) + part_path = None + db.update_download(download_id, status="completed") + except asyncio.CancelledError: + db.update_download(download_id, status="canceled", error="Canceled by user") + _cleanup(part_path) + except httpx.HTTPStatusError as exc: + db.update_download(download_id, status="failed", + error=f"HTTP {exc.response.status_code}") + _cleanup(part_path) + except Exception as exc: # noqa: BLE001 - surface any failure to the UI + db.update_download(download_id, status="failed", error=str(exc)) + _cleanup(part_path) + finally: + _cancel_flags.pop(download_id, None) + + +def _cleanup(part_path: Optional[Path]) -> None: + if part_path is not None: + try: + part_path.unlink(missing_ok=True) + except OSError: + pass diff --git a/model-manager/app/main.py b/model-manager/app/main.py new file mode 100644 index 0000000..ae41c46 --- /dev/null +++ b/model-manager/app/main.py @@ -0,0 +1,191 @@ +"""SparkyUI Model Manager - FastAPI app, API routes, and static UI.""" +from __future__ import annotations + +import asyncio +import os +from pathlib import Path +from typing import Optional + +from fastapi import FastAPI, HTTPException +from fastapi.responses import FileResponse +from fastapi.staticfiles import StaticFiles +from pydantic import BaseModel + +from . import db, downloader, registries +from .config import ( + KEY_BY_FOLDER, + MODEL_TYPES, + MODELS_DIR, + TYPE_BY_KEY, + ensure_dirs, +) + +STATIC_DIR = Path(__file__).parent / "static" + +app = FastAPI(title="SparkyUI Model Manager") + + +@app.on_event("startup") +def _startup() -> None: + ensure_dirs() + db.init_db() + # Any download left mid-flight by a restart is no longer running. + for d in db.list_downloads(): + if d["status"] in ("queued", "downloading"): + db.update_download(d["id"], status="failed", + error="Interrupted by restart") + + +# ---- request models ------------------------------------------------------- + +class SettingsIn(BaseModel): + civitai_api_key: Optional[str] = None + huggingface_token: Optional[str] = None + + +class DownloadIn(BaseModel): + url: str + model_type: Optional[str] = None + filename: Optional[str] = None + + +# ---- model types & installed models --------------------------------------- + +@app.get("/api/model-types") +def model_types() -> list[dict]: + return MODEL_TYPES + + +@app.get("/api/models") +def list_models() -> list[dict]: + results: list[dict] = [] + for folder, type_key in KEY_BY_FOLDER.items(): + base = MODELS_DIR / folder + if not base.is_dir(): + continue + for entry in sorted(base.iterdir()): + if entry.name.startswith(".") or entry.name.endswith(".part"): + continue + if not entry.is_file(): + continue + stat = entry.stat() + results.append({ + "name": entry.name, + "type": type_key, + "type_label": TYPE_BY_KEY[type_key]["label"], + "folder": folder, + "size": stat.st_size, + "mtime": stat.st_mtime, + }) + return results + + +class DeleteModelIn(BaseModel): + folder: str + name: str + + +@app.delete("/api/models") +def delete_model(body: DeleteModelIn) -> dict: + if body.folder not in KEY_BY_FOLDER: + raise HTTPException(400, "Unknown model folder") + name = os.path.basename(body.name) + base = (MODELS_DIR / body.folder).resolve() + target = (base / name).resolve() + if not str(target).startswith(str(base) + os.sep): + raise HTTPException(400, "Invalid path") + if not target.is_file(): + raise HTTPException(404, "File not found") + target.unlink() + return {"deleted": name} + + +# ---- settings ------------------------------------------------------------- + +@app.get("/api/settings") +def get_settings() -> dict: + # Never return the secrets themselves, only whether they are configured. + return { + "civitai_api_key_set": bool(db.get_setting("civitai_api_key")), + "huggingface_token_set": bool(db.get_setting("huggingface_token")), + } + + +@app.post("/api/settings") +def save_settings(body: SettingsIn) -> dict: + # `None` leaves a value untouched; empty string clears it. + if body.civitai_api_key is not None: + db.set_setting("civitai_api_key", body.civitai_api_key) + if body.huggingface_token is not None: + db.set_setting("huggingface_token", body.huggingface_token) + return get_settings() + + +# ---- downloads ------------------------------------------------------------ + +@app.get("/api/downloads") +def get_downloads() -> list[dict]: + return db.list_downloads() + + +@app.post("/api/downloads") +async def start_download(body: DownloadIn) -> dict: + url = body.url.strip() + if not url: + raise HTTPException(400, "URL is required") + + try: + resolved = await registries.resolve(url) + except Exception as exc: # noqa: BLE001 - report resolution failure to the user + raise HTTPException(400, f"Could not resolve URL: {exc}") + + # Pick the model type: explicit user choice wins, else registry suggestion. + model_type = body.model_type or resolved.model_type + if not model_type: + raise HTTPException( + 400, "Please choose a model type for this download.") + if model_type not in TYPE_BY_KEY: + raise HTTPException(400, "Unknown model type") + + filename = body.filename or resolved.filename or "" + + download_id = db.create_download( + url=resolved.download_url, + source=resolved.source, + model_type=model_type, + filename=filename, + dest_path="", + ) + + asyncio.create_task(downloader.run_download( + download_id=download_id, + url=resolved.download_url, + headers=resolved.headers, + model_type=model_type, + filename=filename or None, + )) + + row = db.get_download(download_id) + return row or {"id": download_id} + + +@app.delete("/api/downloads/{download_id}") +def delete_download(download_id: int) -> dict: + row = db.get_download(download_id) + if not row: + raise HTTPException(404, "Download not found") + # If it's running, signal cancel; the task will mark it canceled. + if downloader.request_cancel(download_id): + return {"canceled": download_id} + db.delete_download(download_id) + return {"removed": download_id} + + +# ---- static UI ------------------------------------------------------------ + +@app.get("/") +def index() -> FileResponse: + return FileResponse(STATIC_DIR / "index.html") + + +app.mount("/", StaticFiles(directory=str(STATIC_DIR)), name="static") diff --git a/model-manager/app/registries.py b/model-manager/app/registries.py new file mode 100644 index 0000000..a6d37b2 --- /dev/null +++ b/model-manager/app/registries.py @@ -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, + ) diff --git a/model-manager/app/static/app.js b/model-manager/app/static/app.js new file mode 100644 index 0000000..c84d55a --- /dev/null +++ b/model-manager/app/static/app.js @@ -0,0 +1,292 @@ +"use strict"; + +const $ = (sel) => document.querySelector(sel); +const $$ = (sel) => Array.from(document.querySelectorAll(sel)); + +let MODEL_TYPES = []; +let downloadsTimer = null; + +// ---- helpers -------------------------------------------------------------- + +async function api(path, opts) { + const resp = await fetch(path, opts); + const ct = resp.headers.get("content-type") || ""; + const data = ct.includes("application/json") ? await resp.json() : null; + if (!resp.ok) { + const msg = (data && data.detail) ? data.detail : `HTTP ${resp.status}`; + throw new Error(msg); + } + return data; +} + +function fmtBytes(n) { + if (!n) return "0 B"; + const units = ["B", "KB", "MB", "GB", "TB"]; + let i = 0; + while (n >= 1024 && i < units.length - 1) { n /= 1024; i++; } + return `${n.toFixed(n < 10 && i > 0 ? 1 : 0)} ${units[i]}`; +} + +function fmtDate(epoch) { + if (!epoch) return ""; + return new Date(epoch * 1000).toLocaleString(); +} + +// ---- navigation ----------------------------------------------------------- + +function showView(name) { + $$(".nav-item").forEach((b) => b.classList.toggle("active", b.dataset.view === name)); + $$(".view").forEach((v) => v.classList.toggle("active", v.id === `view-${name}`)); + if (name === "installed") loadModels(); + if (name === "download") { loadDownloads(); startDownloadsPolling(); } + else stopDownloadsPolling(); + if (name === "settings") loadSettings(); +} + +$$(".nav-item").forEach((btn) => + btn.addEventListener("click", () => showView(btn.dataset.view))); + +// ---- model types ---------------------------------------------------------- + +async function loadModelTypes() { + MODEL_TYPES = await api("/api/model-types"); + const sel = $("#dlType"); + sel.innerHTML = ""; + for (const t of MODEL_TYPES) { + const opt = document.createElement("option"); + opt.value = t.key; + opt.textContent = t.label; + sel.appendChild(opt); + } +} + +// ---- installed models ----------------------------------------------------- + +async function loadModels() { + const container = $("#modelsContainer"); + try { + const models = await api("/api/models"); + renderModels(models); + } catch (err) { + container.innerHTML = `

Error: ${err.message}

`; + } +} + +function renderModels(models) { + const container = $("#modelsContainer"); + const filter = $("#modelSearch").value.trim().toLowerCase(); + const shown = filter + ? models.filter((m) => m.name.toLowerCase().includes(filter)) + : models; + + if (!shown.length) { + container.innerHTML = `

${ + models.length ? "No models match your filter." : "No models installed yet. Use Add / Download." + }

`; + return; + } + + const groups = {}; + for (const m of shown) (groups[m.type_label] ||= []).push(m); + + container.innerHTML = ""; + for (const label of Object.keys(groups).sort()) { + const group = document.createElement("div"); + group.className = "type-group"; + group.innerHTML = `

${label} · ${groups[label].length}

`; + const grid = document.createElement("div"); + grid.className = "model-grid"; + for (const m of groups[label]) { + const card = document.createElement("div"); + card.className = "model-card"; + card.innerHTML = ` +
+
${fmtBytes(m.size)} · ${fmtDate(m.mtime)}
+
+ ${m.folder}/ + +
`; + card.querySelector(".name").textContent = m.name; + card.querySelector(".danger").addEventListener("click", () => deleteModel(m)); + grid.appendChild(card); + } + group.appendChild(grid); + container.appendChild(group); + } +} + +async function deleteModel(m) { + if (!confirm(`Delete ${m.name}?`)) return; + try { + await api("/api/models", { + method: "DELETE", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ folder: m.folder, name: m.name }), + }); + loadModels(); + } catch (err) { + alert(`Delete failed: ${err.message}`); + } +} + +$("#refreshModels").addEventListener("click", loadModels); +$("#modelSearch").addEventListener("input", () => loadModels()); + +// ---- downloads ------------------------------------------------------------ + +function detectSource(url) { + url = url.toLowerCase(); + if (url.includes("civitai.com")) return "CivitAI (type auto-detected, API key used)"; + if (url.includes("huggingface.co") || url.includes("hf.co")) return "HuggingFace (token used)"; + if (url) return "Direct download"; + return ""; +} + +$("#dlUrl").addEventListener("input", (e) => { + $("#sourceHint").textContent = detectSource(e.target.value.trim()); +}); + +$("#startDownload").addEventListener("click", async () => { + const url = $("#dlUrl").value.trim(); + const msg = $("#dlMsg"); + if (!url) { msg.className = "form-msg err"; msg.textContent = "Enter a URL."; return; } + const btn = $("#startDownload"); + btn.disabled = true; + msg.className = "form-msg"; msg.textContent = "Starting…"; + try { + await api("/api/downloads", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + url, + model_type: $("#dlType").value || null, + filename: $("#dlFilename").value.trim() || null, + }), + }); + msg.className = "form-msg ok"; msg.textContent = "Download started."; + $("#dlUrl").value = ""; $("#dlFilename").value = ""; $("#sourceHint").textContent = ""; + loadDownloads(); + startDownloadsPolling(); + } catch (err) { + msg.className = "form-msg err"; msg.textContent = err.message; + } finally { + btn.disabled = false; + } +}); + +async function loadDownloads() { + let rows; + try { + rows = await api("/api/downloads"); + } catch { + return; + } + renderDownloads(rows); + + const active = rows.some((r) => r.status === "downloading" || r.status === "queued"); + if (!active) stopDownloadsPolling(); +} + +function renderDownloads(rows) { + const container = $("#downloadsContainer"); + if (!rows.length) { + container.innerHTML = `

No downloads yet.

`; + return; + } + container.innerHTML = ""; + for (const r of rows) { + const pct = r.bytes_total ? Math.floor((r.bytes_done / r.bytes_total) * 100) : 0; + const row = document.createElement("div"); + row.className = "dl-row"; + const sizeText = r.bytes_total + ? `${fmtBytes(r.bytes_done)} / ${fmtBytes(r.bytes_total)} (${pct}%)` + : fmtBytes(r.bytes_done); + const sub = r.status === "failed" && r.error ? r.error + : `${r.source} · ${r.model_type} · ${sizeText}`; + row.innerHTML = ` +
+
+
+
+
+
+ ${r.status} + +
+
+ ${r.status === "downloading" + ? `
` : ""}`; + row.querySelector(".dl-name").textContent = r.filename || r.url; + row.querySelector(".dl-sub").textContent = sub; + const actionBtn = row.querySelector(".danger"); + const active = r.status === "downloading" || r.status === "queued"; + actionBtn.textContent = active ? "Cancel" : "Remove"; + actionBtn.addEventListener("click", () => removeDownload(r.id)); + container.appendChild(row); + } +} + +async function removeDownload(id) { + try { + await api(`/api/downloads/${id}`, { method: "DELETE" }); + loadDownloads(); + } catch (err) { + alert(`Failed: ${err.message}`); + } +} + +function startDownloadsPolling() { + if (downloadsTimer) return; + downloadsTimer = setInterval(loadDownloads, 1000); +} +function stopDownloadsPolling() { + if (downloadsTimer) { clearInterval(downloadsTimer); downloadsTimer = null; } +} + +// ---- settings ------------------------------------------------------------- + +async function loadSettings() { + try { + const s = await api("/api/settings"); + setBadge($("#civitaiBadge"), s.civitai_api_key_set); + setBadge($("#hfBadge"), s.huggingface_token_set); + } catch { /* ignore */ } +} + +function setBadge(el, isSet) { + el.textContent = isSet ? "configured" : "not set"; + el.className = `badge ${isSet ? "set" : "unset"}`; +} + +$("#saveSettings").addEventListener("click", async () => { + const msg = $("#settingsMsg"); + const body = {}; + const civ = $("#civitaiKey").value; + const hf = $("#hfToken").value; + // Only send fields the user actually typed into (blank = leave unchanged). + if (civ !== "") body.civitai_api_key = civ; + if (hf !== "") body.huggingface_token = hf; + if (!Object.keys(body).length) { + msg.className = "form-msg"; msg.textContent = "Nothing to save."; + return; + } + try { + await api("/api/settings", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(body), + }); + msg.className = "form-msg ok"; msg.textContent = "Saved."; + $("#civitaiKey").value = ""; $("#hfToken").value = ""; + loadSettings(); + } catch (err) { + msg.className = "form-msg err"; msg.textContent = err.message; + } +}); + +// ---- boot ----------------------------------------------------------------- + +(async function init() { + await loadModelTypes(); + showView("installed"); +})(); diff --git a/model-manager/app/static/index.html b/model-manager/app/static/index.html new file mode 100644 index 0000000..ff2d696 --- /dev/null +++ b/model-manager/app/static/index.html @@ -0,0 +1,115 @@ + + + + + + SparkyUI · Model Manager + + + +
+ + +
+ +
+
+

Installed Models

+
+ + +
+
+
+

Loading…

+
+
+ + +
+

Add / Download Model

+
+ + +
+ +
+
+ + +
CivitAI links auto-detect type; leave as-is to use it.
+
+
+ + +
+
+ +
+ + +
+
+ +

Downloads

+
+

No downloads yet.

+
+
+ + +
+

Settings

+
+

API Keys

+

+ Stored persistently and sent as auth headers when downloading from the + matching site. Leave blank to keep the existing value. +

+ + + + + + + +
+ + +
+
+ +
+

Model Storage

+

+ Models are saved into the project models/ directory, sorted into + ComfyUI's standard sub-folders by type. ComfyUI reads from the same folder. +

+
+
+
+
+ + + diff --git a/model-manager/app/static/style.css b/model-manager/app/static/style.css new file mode 100644 index 0000000..5bbe98f --- /dev/null +++ b/model-manager/app/static/style.css @@ -0,0 +1,152 @@ +:root { + --bg: #14161c; + --bg-2: #1b1e27; + --bg-3: #232735; + --border: #2c3142; + --text: #e6e8ee; + --text-dim: #99a0b3; + --accent: #6c8cff; + --accent-2: #8a6cff; + --green: #36c98b; + --red: #ef5d6b; + --amber: #f0b400; + --radius: 10px; +} + +* { box-sizing: border-box; } + +html, body { + margin: 0; + height: 100%; + background: var(--bg); + color: var(--text); + font-family: "Segoe UI", system-ui, -apple-system, Roboto, sans-serif; + font-size: 14px; +} + +.app { display: flex; height: 100vh; } + +/* ---- sidebar ---- */ +.sidebar { + width: 240px; + flex-shrink: 0; + background: var(--bg-2); + border-right: 1px solid var(--border); + display: flex; + flex-direction: column; + padding: 18px 14px; +} +.brand { display: flex; align-items: center; gap: 12px; padding: 4px 6px 22px; } +.brand .spark { font-size: 26px; } +.brand-title { font-weight: 700; font-size: 16px; } +.brand-sub { color: var(--text-dim); font-size: 12px; } + +nav { display: flex; flex-direction: column; gap: 4px; } +.nav-item { + display: flex; align-items: center; gap: 10px; + background: transparent; color: var(--text-dim); + border: none; border-radius: 8px; + padding: 10px 12px; font-size: 14px; text-align: left; + cursor: pointer; transition: background .12s, color .12s; +} +.nav-item:hover { background: var(--bg-3); color: var(--text); } +.nav-item.active { + background: linear-gradient(90deg, rgba(108,140,255,.18), rgba(138,108,255,.12)); + color: var(--text); +} +.nav-ico { width: 18px; text-align: center; opacity: .85; } +.sidebar-foot { + margin-top: auto; color: var(--text-dim); font-size: 11px; + word-break: break-all; padding: 8px 6px 0; +} + +/* ---- content ---- */ +.content { flex: 1; overflow-y: auto; padding: 26px 32px; } +.view { display: none; } +.view.active { display: block; } +.view-head { display: flex; align-items: center; justify-content: space-between; margin-bottom: 20px; } +.view-head h1 { font-size: 22px; margin: 0; } +.head-actions { display: flex; gap: 10px; } +.subhead { font-size: 15px; color: var(--text-dim); margin: 28px 0 12px; } +.subhead.first { margin-top: 0; } + +/* ---- inputs/buttons ---- */ +input, select { + background: var(--bg-3); color: var(--text); + border: 1px solid var(--border); border-radius: 8px; + padding: 9px 11px; font-size: 14px; width: 100%; +} +input:focus, select:focus { outline: none; border-color: var(--accent); } +input[type=search] { width: 200px; } +label { display: block; margin: 14px 0 6px; color: var(--text-dim); font-size: 13px; } + +.btn { + background: var(--bg-3); color: var(--text); + border: 1px solid var(--border); border-radius: 8px; + padding: 9px 16px; font-size: 14px; cursor: pointer; + transition: background .12s, border-color .12s; +} +.btn:hover { background: #2a3042; } +.btn.primary { background: var(--accent); border-color: var(--accent); color: #0c1020; font-weight: 600; } +.btn.primary:hover { background: #5b7dff; } +.btn.danger { color: var(--red); border-color: transparent; background: transparent; padding: 6px 10px; } +.btn.danger:hover { background: rgba(239,93,107,.12); } +.btn:disabled { opacity: .5; cursor: not-allowed; } + +/* ---- cards ---- */ +.card { + background: var(--bg-2); border: 1px solid var(--border); + border-radius: var(--radius); padding: 20px; margin-bottom: 18px; +} +.form-card { max-width: 720px; } +.row { display: flex; gap: 18px; } +.col { flex: 1; } +.hint { color: var(--text-dim); font-size: 12px; margin-top: 6px; } +.source-hint { font-size: 12px; margin-top: 8px; min-height: 16px; color: var(--accent); } +.form-actions { display: flex; align-items: center; gap: 14px; margin-top: 20px; } +.form-msg { font-size: 13px; } +.form-msg.ok { color: var(--green); } +.form-msg.err { color: var(--red); } +.badge { font-size: 11px; padding: 1px 8px; border-radius: 10px; margin-left: 6px; } +.badge.set { background: rgba(54,201,139,.18); color: var(--green); } +.badge.unset { background: var(--bg-3); color: var(--text-dim); } +code { background: var(--bg-3); padding: 1px 5px; border-radius: 4px; font-size: 12px; } + +/* ---- installed models ---- */ +.models-container { display: flex; flex-direction: column; gap: 22px; } +.type-group h3 { + font-size: 13px; text-transform: uppercase; letter-spacing: .04em; + color: var(--text-dim); margin: 0 0 10px; font-weight: 600; +} +.model-grid { + display: grid; gap: 12px; + grid-template-columns: repeat(auto-fill, minmax(260px, 1fr)); +} +.model-card { + background: var(--bg-2); border: 1px solid var(--border); + border-radius: var(--radius); padding: 14px; + display: flex; flex-direction: column; gap: 6px; +} +.model-card .name { font-weight: 600; word-break: break-all; } +.model-card .meta { color: var(--text-dim); font-size: 12px; } +.model-card .card-foot { display: flex; justify-content: space-between; align-items: center; margin-top: 4px; } +.empty { color: var(--text-dim); } + +/* ---- downloads ---- */ +.downloads-container { display: flex; flex-direction: column; gap: 10px; max-width: 860px; } +.dl-row { + background: var(--bg-2); border: 1px solid var(--border); + border-radius: var(--radius); padding: 12px 14px; +} +.dl-top { display: flex; justify-content: space-between; align-items: center; gap: 12px; } +.dl-name { font-weight: 600; word-break: break-all; } +.dl-sub { color: var(--text-dim); font-size: 12px; margin-top: 2px; } +.dl-status { font-size: 12px; padding: 2px 9px; border-radius: 10px; white-space: nowrap; } +.dl-status.downloading { background: rgba(108,140,255,.18); color: var(--accent); } +.dl-status.completed { background: rgba(54,201,139,.18); color: var(--green); } +.dl-status.failed { background: rgba(239,93,107,.18); color: var(--red); } +.dl-status.canceled { background: var(--bg-3); color: var(--text-dim); } +.dl-status.queued { background: rgba(240,180,0,.16); color: var(--amber); } +.progress { height: 6px; background: var(--bg-3); border-radius: 4px; margin-top: 10px; overflow: hidden; } +.progress > span { display: block; height: 100%; background: linear-gradient(90deg, var(--accent), var(--accent-2)); width: 0; transition: width .3s; } +.dl-actions { display: flex; align-items: center; gap: 8px; } diff --git a/model-manager/requirements.txt b/model-manager/requirements.txt new file mode 100644 index 0000000..534f16c --- /dev/null +++ b/model-manager/requirements.txt @@ -0,0 +1,4 @@ +fastapi==0.115.6 +uvicorn[standard]==0.34.0 +httpx==0.28.1 +pydantic==2.10.4 diff --git a/models/.gitkeep b/models/.gitkeep new file mode 100644 index 0000000..e69de29