59 lines
2.0 KiB
Python
59 lines
2.0 KiB
Python
import folder_paths
|
|
from pathlib import Path
|
|
from typing import Dict, Any
|
|
from huggingface_hub import snapshot_download
|
|
from ..logger import niko_logger as logger
|
|
|
|
|
|
MODEL_REGISTRY: Dict[str, Dict[str, Any]] = {
|
|
"lineart": {
|
|
"repo_id": "Nikos7766/lineart-models",
|
|
"subfolder": "controlnet/preprocessors/lineart"
|
|
},
|
|
"depthanythingv2": {
|
|
"repo_id": "Nikos7766/DepthAnythingV2",
|
|
"subfolder": "controlnet/preprocessors/depthanythingv2"
|
|
},
|
|
}
|
|
|
|
|
|
def get_model_path(model_key: str, model_name: str) -> str:
|
|
"""Finds or downloads the model using pathlib for better path handling."""
|
|
if model_key not in MODEL_REGISTRY:
|
|
raise ValueError(f"Model key '{model_key}' not found in registry.")
|
|
|
|
config = MODEL_REGISTRY[model_key]
|
|
subfolder = Path(config["subfolder"])
|
|
download_dir = Path(folder_paths.models_dir) / subfolder
|
|
local_model_path = download_dir / model_name
|
|
|
|
# Check in ComfyUI - configured paths
|
|
registered_path = folder_paths.get_full_path(model_key, model_name)
|
|
if registered_path and Path(registered_path).exists():
|
|
logger.info(f"Found - {model_name} in registered path: {registered_path} - Loading ...")
|
|
return str(registered_path)
|
|
|
|
# Check existing paths
|
|
if local_model_path.exists():
|
|
logger.info(f"Found {model_name} at {local_model_path} - Loading ...")
|
|
return str(local_model_path)
|
|
|
|
# Download if missing
|
|
logger.info(f"Model Not Found - Downloading {model_name}...")
|
|
download_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
try:
|
|
downloaded_file = Path(snapshot_download(
|
|
repo_id=config["repo_id"],
|
|
local_dir=str(download_dir), # Convert Path to str for API compatibility
|
|
allow_patterns=[model_name],
|
|
)).resolve()
|
|
|
|
model_path = downloaded_file / model_name
|
|
logger.info(f"Downloaded {model_name} to {downloaded_file}")
|
|
return str(model_path)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Download failed: {e}")
|
|
raise
|