TBNilles
359043ad67
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>
54 lines
2.5 KiB
Python
54 lines
2.5 KiB
Python
"""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)
|