"""Load and validate YAML config; build ComfyUI paths.""" from pathlib import Path from typing import Any import yaml # Config keys for model lists → ComfyUI subdir under models/ MODEL_TYPE_SUBDIRS = { "diffusion_models": "diffusion_models", "text_encoders": "text_encoders", "vaes": "vae", "upscale_models": "upscale_models", "loras": "loras", } REQUIRED_KEYS = {"comfyui_base_dir"} TOKEN_KEYS = {"huggingface_token", "civitai_token"} def load_config(path: str | Path) -> dict[str, Any]: """Load YAML config from path. Raises FileNotFoundError or yaml error.""" p = Path(path) if not p.exists(): raise FileNotFoundError(f"Config not found: {p}") with open(p, encoding="utf-8") as f: data = yaml.safe_load(f) if not isinstance(data, dict): raise ValueError("Config must be a YAML object") return data def validate_config(data: dict[str, Any]) -> None: """Validate required keys and model list types. Raises ValueError on failure.""" if "comfyui_base_dir" not in data: raise ValueError("Config must contain 'comfyui_base_dir'") base = data["comfyui_base_dir"] if not base or not isinstance(base, str): raise ValueError("comfyui_base_dir must be a non-empty string") for key in MODEL_TYPE_SUBDIRS: val = data.get(key) if val is None: data[key] = [] elif not isinstance(val, list): raise ValueError(f"'{key}' must be a list of URL strings") else: for i, item in enumerate(val): if not isinstance(item, str): raise ValueError(f"'{key}[{i}]' must be a string") def get_model_tasks( data: dict[str, Any], only_types: list[str] | None = None ) -> list[tuple[str, str]]: """Return list of (model_type, url) for all configured model URLs.""" types = only_types if only_types is not None else list(MODEL_TYPE_SUBDIRS) tasks: list[tuple[str, str]] = [] for model_type in types: if model_type not in MODEL_TYPE_SUBDIRS: continue urls = data.get(model_type) if not isinstance(urls, list): continue for url in urls: if isinstance(url, str) and url.strip(): tasks.append((model_type, url.strip())) return tasks def get_download_dir(data: dict[str, Any], model_type: str) -> Path: """Return the absolute directory path for a model type under ComfyUI base.""" base = Path(data["comfyui_base_dir"]).expanduser().resolve() subdir = MODEL_TYPE_SUBDIRS[model_type] return base / "models" / subdir