TBNilles
399acabd58
ComfyUI caches the last model when RAM is plentiful (unified memory), so memory doesn't drop after switching models even though models are being swapped, not accumulated. Add a sidebar "Free GPU memory" button that proxies ComfyUI's POST /free (unload_models + free_memory) via a new /api/comfyui/free endpoint (COMFYUI_URL env). Verified it releases ~7GB. README documents this plus the --disable-smart-memory auto-unload option. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
81 lines
3.7 KiB
Python
81 lines
3.7 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()
|
|
|
|
# Generated images directory (ComfyUI output), mounted read-write for the gallery.
|
|
OUTPUT_DIR = Path(os.environ.get("OUTPUT_DIR", "/output")).resolve()
|
|
|
|
DB_PATH = DATA_DIR / "manager.db"
|
|
|
|
# Ports of the sibling services, used by the device-routing landing page.
|
|
COMFYUI_PORT = os.environ.get("COMFYUI_PORT", "8188")
|
|
COMFYUIMINI_PORT = os.environ.get("COMFYUIMINI_PORT", "3000")
|
|
|
|
# Internal URL of the ComfyUI container (reachable over the docker network),
|
|
# used to proxy the "free memory / unload models" action.
|
|
COMFYUI_URL = os.environ.get("COMFYUI_URL", "http://comfyui:8188")
|
|
|
|
# Image file types shown in the gallery.
|
|
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp"}
|
|
|
|
# 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)
|
|
|
|
|
|
def safe_output_path(rel_path: str) -> Path:
|
|
"""Resolve a gallery-relative path to an absolute one under OUTPUT_DIR.
|
|
|
|
Raises ValueError on any attempt to escape the output directory.
|
|
"""
|
|
# Normalize and strip leading separators so it's treated as relative.
|
|
rel = rel_path.replace("\\", "/").lstrip("/")
|
|
target = (OUTPUT_DIR / rel).resolve()
|
|
if target != OUTPUT_DIR and not str(target).startswith(str(OUTPUT_DIR) + os.sep):
|
|
raise ValueError("Path is outside the output directory")
|
|
return target
|