"""Model upload route functions for ComfyUI.
Provides functionality to add models to ComfyUI by uploading files
or downloading from URLs.
"""
import os
from enum import Enum
from pathlib import Path
import httpx
from src.client.response import ResponseGetData
from src.utils import get_global_logger
from src.utils.logging import log_call
logger = get_global_logger("ComfyUI_MCP.routes.model_upload")
class ModelType(str, Enum):
"""Supported model types and their target directories."""
CHECKPOINT = "checkpoints"
VAE = "vae"
LORA = "loras"
CONTROLNET = "controlnet"
UPSCALER = "upscale_models"
EMBEDDING = "embeddings"
HYPERNETWORK = "hypernetworks"
CLIP = "clip"
CLIP_VISION = "clip_vision"
UNET = "unet"
STYLE_MODELS = "style_models"
def get_default_comfyui_path() -> str | None:
"""Get default ComfyUI path from environment or common locations.
Checks:
1. COMFYUI_MODELS_PATH environment variable (direct path to models dir)
2. COMFYUI_PATH environment variable (base ComfyUI installation)
3. Common installation locations
Returns:
Path to ComfyUI installation, or None if not found
"""
# Check explicit models path first
if models_path := os.getenv("COMFYUI_MODELS_PATH"):
# This should be the models directory itself
if Path(models_path).exists():
logger.info(f"Using COMFYUI_MODELS_PATH: {models_path}")
return str(Path(models_path).parent) # Return parent (ComfyUI root)
# Check base ComfyUI path
if comfyui_path := os.getenv("COMFYUI_PATH"):
if Path(comfyui_path).exists():
logger.info(f"Using COMFYUI_PATH: {comfyui_path}")
return comfyui_path
# Try common locations
common_paths = [
Path.home() / "ComfyUI",
Path.home() / "comfyui",
Path("/opt/ComfyUI"),
Path("/usr/local/ComfyUI"),
Path("C:/ComfyUI") if os.name == "nt" else None,
]
for path in common_paths:
if path and path.exists() and (path / "models").exists():
logger.info(f"Found ComfyUI at common location: {path}")
return str(path)
logger.warning("Could not auto-detect ComfyUI path")
return None
class ModelUploadError(Exception):
"""Raised when model upload operations fail."""
def __init__(self, message: str, response: ResponseGetData | None = None):
super().__init__(message)
self.response = response
def get_model_directory(comfyui_path: str, model_type: ModelType) -> Path:
"""Get the target directory for a specific model type.
Args:
comfyui_path: Base path to ComfyUI installation
model_type: Type of model (checkpoint, lora, etc.)
Returns:
Path to the model directory
Raises:
ModelUploadError: If directory doesn't exist and can't be created
"""
model_dir = Path(comfyui_path) / "models" / model_type.value
# Create directory if it doesn't exist
try:
model_dir.mkdir(parents=True, exist_ok=True)
except Exception as e:
raise ModelUploadError(f"Failed to create model directory {model_dir}: {e}") from e
return model_dir
def validate_model_file(filename: str, model_type: ModelType) -> bool:
"""Validate that the filename has an appropriate extension.
Args:
filename: Name of the model file
model_type: Type of model
Returns:
True if valid, False otherwise
"""
valid_extensions = {
ModelType.CHECKPOINT: [".safetensors", ".ckpt", ".pt", ".pth", ".bin"],
ModelType.VAE: [".safetensors", ".ckpt", ".pt", ".pth"],
ModelType.LORA: [".safetensors", ".ckpt", ".pt", ".pth"],
ModelType.CONTROLNET: [".safetensors", ".ckpt", ".pt", ".pth"],
ModelType.UPSCALER: [".pth", ".pt"],
ModelType.EMBEDDING: [".safetensors", ".pt", ".bin"],
ModelType.HYPERNETWORK: [".safetensors", ".pt"],
ModelType.CLIP: [".safetensors", ".pt"],
ModelType.CLIP_VISION: [".safetensors", ".pt"],
ModelType.UNET: [".safetensors", ".pt"],
ModelType.STYLE_MODELS: [".safetensors", ".ckpt"],
}
file_ext = Path(filename).suffix.lower()
return file_ext in valid_extensions.get(model_type, [])
@log_call(action_name="upload_model_file", level_name="route")
async def upload_model_file(
file_data: bytes,
filename: str,
model_type: ModelType = ModelType.CHECKPOINT,
overwrite: bool = False,
comfyui_path: str | None = None,
) -> dict[str, any]:
"""Upload a model file to ComfyUI.
Saves the provided file data to the appropriate ComfyUI model directory.
Args:
file_data: Binary content of the model file
filename: Name for the saved file
model_type: Type of model (checkpoint, lora, etc.)
overwrite: Whether to overwrite existing file
comfyui_path: Optional base path to ComfyUI installation. If not provided,
will use COMFYUI_PATH or COMFYUI_MODELS_PATH env vars,
or attempt to auto-detect from common locations.
Returns:
Dict with:
- success: True if successful
- path: Full path to saved file
- size_bytes: Size of the file in bytes
- model_type: Type of model uploaded
- comfyui_path: Resolved ComfyUI path used
Raises:
ModelUploadError: If upload fails or ComfyUI path cannot be determined
Example:
>>> with open("sd_xl_base_1.0.safetensors", "rb") as f:
... file_data = f.read()
>>> result = await upload_model_file(
... file_data=file_data,
... filename="sd_xl_base_1.0.safetensors",
... model_type=ModelType.CHECKPOINT,
... )
>>> print(result["path"])
/opt/ComfyUI/models/checkpoints/sd_xl_base_1.0.safetensors
"""
# Resolve ComfyUI path
if comfyui_path is None:
comfyui_path = get_default_comfyui_path()
if comfyui_path is None:
raise ModelUploadError(
"ComfyUI path not specified and could not be auto-detected. "
"Please set COMFYUI_PATH or COMFYUI_MODELS_PATH environment variable, "
"or provide comfyui_path parameter."
)
# Validate filename
if not validate_model_file(filename, model_type):
raise ModelUploadError(
f"Invalid file extension for {model_type.value}. "
f"Expected one of: {', '.join(['.safetensors', '.ckpt', '.pt', '.pth'])}"
)
# Sanitize filename to prevent path traversal
filename = Path(filename).name
# Get target directory
try:
target_dir = get_model_directory(comfyui_path, model_type)
except Exception as e:
raise ModelUploadError(f"Failed to access model directory: {e}") from e
# Check if file exists
target_path = target_dir / filename
if target_path.exists() and not overwrite:
raise ModelUploadError(
f"File {filename} already exists in {model_type.value}. "
f"Set overwrite=True to replace it."
)
# Write file
try:
with open(target_path, "wb") as f:
f.write(file_data)
except Exception as e:
raise ModelUploadError(f"Failed to write file to {target_path}: {e}") from e
return {
"success": True,
"path": str(target_path),
"size_bytes": len(file_data),
"model_type": model_type.value,
"filename": filename,
"comfyui_path": comfyui_path,
}
@log_call(
action_name="download_model_from_url",
level_name="route",
log_params=True,
sensitive_params=["url"], # URL may contain API tokens in query string
)
async def download_model_from_url(
url: str,
filename: str | None = None,
model_type: ModelType = ModelType.CHECKPOINT,
overwrite: bool = False,
comfyui_path: str | None = None,
*,
session: httpx.AsyncClient | None = None,
) -> dict[str, any]:
"""Download a model from a URL and save it to ComfyUI.
Downloads a model file from the provided URL and saves it to the
appropriate ComfyUI model directory.
Args:
url: URL to download the model from
filename: Optional filename (extracted from URL if not provided)
model_type: Type of model (checkpoint, lora, etc.)
overwrite: Whether to overwrite existing file
comfyui_path: Optional base path to ComfyUI installation. If not provided,
will use COMFYUI_PATH or COMFYUI_MODELS_PATH env vars,
or attempt to auto-detect from common locations.
session: Optional HTTPX client for connection pooling
Returns:
Dict with:
- success: True if successful
- path: Full path to saved file
- size_bytes: Size of the downloaded file
- model_type: Type of model uploaded
- download_url: Original download URL
- comfyui_path: Resolved ComfyUI path used
Raises:
ModelUploadError: If download or save fails
Example:
>>> result = await download_model_from_url(
... url="https://huggingface.co/stabilityai/sdxl-base/resolve/main/sd_xl_base_1.0.safetensors",
... model_type=ModelType.CHECKPOINT,
... )
>>> print(result["path"])
/opt/ComfyUI/models/checkpoints/sd_xl_base_1.0.safetensors
"""
# Extract filename from URL if not provided
if not filename:
filename = Path(url).name
if not filename or filename == "":
raise ModelUploadError(
"Could not extract filename from URL. Please provide filename explicitly."
)
# Download file
try:
close_session = False
if session is None:
session = httpx.AsyncClient(timeout=300.0) # 5 minute timeout for large files
close_session = True
try:
response = await session.get(url, follow_redirects=True)
response.raise_for_status()
file_data = response.content
finally:
if close_session:
await session.aclose()
except httpx.HTTPError as e:
raise ModelUploadError(f"Failed to download model from {url}: {e}") from e
# Upload the downloaded file
result = await upload_model_file(
comfyui_path=comfyui_path,
file_data=file_data,
filename=filename,
model_type=model_type,
overwrite=overwrite,
)
# Add download URL to result
result["download_url"] = url
return result
@log_call(action_name="list_installed_models", level_name="route")
async def list_installed_models(
model_type: ModelType | None = None,
comfyui_path: str | None = None,
) -> dict[str, any]:
"""List all installed models in ComfyUI.
Scans the ComfyUI models directory and returns a list of installed models,
optionally filtered by type.
Args:
model_type: Optional model type filter (lists all types if None)
comfyui_path: Optional base path to ComfyUI installation. If not provided,
will use COMFYUI_PATH or COMFYUI_MODELS_PATH env vars,
or attempt to auto-detect from common locations.
Returns:
Dict with:
- models: Dict mapping model type to list of filenames
- total_count: Total number of models
- model_types: List of model types with files
- comfyui_path: Resolved ComfyUI path used
Raises:
ModelUploadError: If directory scan fails or ComfyUI path cannot be determined
Example:
>>> result = await list_installed_models(
... model_type=ModelType.CHECKPOINT,
... )
>>> print(result["models"]["checkpoints"])
[{"filename": "sd_xl_base_1.0.safetensors", "size_bytes": 12345, "modified": 1234567890.0}]
"""
# Resolve ComfyUI path
if comfyui_path is None:
comfyui_path = get_default_comfyui_path()
if comfyui_path is None:
raise ModelUploadError(
"ComfyUI path not specified and could not be auto-detected. "
"Please set COMFYUI_PATH or COMFYUI_MODELS_PATH environment variable, "
"or provide comfyui_path parameter."
)
models: dict[str, list] = {}
total_count = 0
# Determine which model types to scan
types_to_scan = [model_type] if model_type else list(ModelType)
for mtype in types_to_scan:
model_dir = Path(comfyui_path) / "models" / mtype.value
if not model_dir.exists():
continue
try:
files = []
for file_path in model_dir.iterdir():
if file_path.is_file():
files.append(
{
"filename": file_path.name,
"size_bytes": file_path.stat().st_size,
"modified": file_path.stat().st_mtime,
}
)
if files:
models[mtype.value] = files
total_count += len(files)
except Exception as e:
raise ModelUploadError(f"Failed to scan directory {model_dir}: {e}") from e
return {
"models": models,
"total_count": total_count,
"model_types": list(models.keys()),
"comfyui_path": comfyui_path,
}
@log_call(action_name="delete_model", level_name="route")
async def delete_model(
filename: str,
model_type: ModelType,
comfyui_path: str | None = None,
) -> dict[str, any]:
"""Delete a model file from ComfyUI.
Removes the specified model file from the ComfyUI models directory.
Args:
filename: Name of the file to delete
model_type: Type of model
comfyui_path: Optional base path to ComfyUI installation. If not provided,
will use COMFYUI_PATH or COMFYUI_MODELS_PATH env vars,
or attempt to auto-detect from common locations.
Returns:
Dict with:
- success: True if successful
- deleted_file: Path of the deleted file
- model_type: Type of model deleted
- comfyui_path: Resolved ComfyUI path used
Raises:
ModelUploadError: If deletion fails or ComfyUI path cannot be determined
Example:
>>> result = await delete_model(
... filename="old_model.safetensors",
... model_type=ModelType.CHECKPOINT,
... )
>>> print(result["success"])
True
"""
# Resolve ComfyUI path
if comfyui_path is None:
comfyui_path = get_default_comfyui_path()
if comfyui_path is None:
raise ModelUploadError(
"ComfyUI path not specified and could not be auto-detected. "
"Please set COMFYUI_PATH or COMFYUI_MODELS_PATH environment variable, "
"or provide comfyui_path parameter."
)
# Sanitize filename to prevent path traversal
filename = Path(filename).name
# Get target directory
target_dir = get_model_directory(comfyui_path, model_type)
target_path = target_dir / filename
# Check if file exists
if not target_path.exists():
raise ModelUploadError(f"File {filename} not found in {model_type.value}")
# Delete file
try:
target_path.unlink()
except Exception as e:
raise ModelUploadError(f"Failed to delete {target_path}: {e}") from e
return {
"success": True,
"deleted_file": str(target_path),
"model_type": model_type.value,
"filename": filename,
"comfyui_path": comfyui_path,
}
