From 93e53ad838430c96f7f6f4c25766c3d232dad27d Mon Sep 17 00:00:00 2001 From: Daniel Henry Date: Sat, 31 Jan 2026 15:09:12 -0600 Subject: [PATCH] Refactor model configuration structure and update README - Changed the configuration format to use a single 'models' list with entries containing 'url' and 'type'. - Updated validation logic to ensure 'models' entries are correctly structured. - Modified download logic to check for existing directories before downloading. - Revised README to reflect new configuration format and usage instructions. Signed-off-by: Daniel Henry --- README.md | 14 +++----- config.example.yaml | 12 +++---- model_downloader/__main__.py | 4 +-- model_downloader/config.py | 70 ++++++++++++++++++++---------------- model_downloader/download.py | 4 ++- 5 files changed, 55 insertions(+), 49 deletions(-) diff --git a/README.md b/README.md index 91ad38e..ff9d790 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ Copy the example config and edit it: cp config.example.yaml config.yaml ``` -Edit `config.yaml`: set `huggingface_token`, `civitai_token`, `comfyui_base_dir`, and add model URLs under the right keys. **Do not commit `config.yaml`** (it is gitignored). +Edit `config.yaml`: set `huggingface_token`, `civitai_token`, `comfyui_base_dir`, and add entries under `models` with `url` and `type` for each model. **Do not commit `config.yaml`** (it is gitignored). ## Run @@ -30,17 +30,11 @@ python -m model_downloader - `--config PATH` – config file (default: `config.yaml` in current directory). - `--dry-run` – print what would be downloaded and where; no writes. -- `--only TYPE ...` – only process these model types (e.g. `--only diffusion_models loras`). +- `--only TYPE ...` – only process these model types (e.g. `--only loras diffusion_models`). -## Config keys and ComfyUI folders +## Config and ComfyUI folders -| Config key | ComfyUI subdir (under `comfyui_base_dir/models/`) | -|--------------------|--------------------------------------------------| -| diffusion_models | diffusion_models | -| text_encoders | text_encoders | -| vaes | vae | -| upscale_models | upscale_models | -| loras | loras | +Each model entry has `url` and `type`. The `type` value is the subdirectory name under `comfyui_base_dir/models/` (e.g. `type: loras` → `models/loras/`). Each `models/` directory must already exist under your ComfyUI base; the tool does not create these directories and will fail if a directory is missing. Tokens stay in your local `config.yaml` only; keep them out of git. diff --git a/config.example.yaml b/config.example.yaml index dcf5097..b5a3667 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -4,10 +4,10 @@ huggingface_token: "" # from https://huggingface.co/settings/tokens civitai_token: "" # from https://civitai.com/user/account comfyui_base_dir: "/path/to/ComfyUI" -# Lists of model URLs. Civitai: use API URL or add ?token=... is optional (token from config is used). +# Single list of models. Each entry: url and type (type = ComfyUI subdir under models/). +# Civitai: use API URL or add ?token=... is optional (token from config is used). # Hugging Face: use resolve/main/... URLs. -diffusion_models: [] -text_encoders: [] -vaes: [] -upscale_models: [] -loras: [] +# models: +# - url: "https://..." +# type: loras +models: [] diff --git a/model_downloader/__main__.py b/model_downloader/__main__.py index 1bb15a3..e143175 100644 --- a/model_downloader/__main__.py +++ b/model_downloader/__main__.py @@ -27,7 +27,7 @@ def main() -> int: "--only", nargs="+", metavar="TYPE", - help="Only process these model types (e.g. diffusion_models loras)", + help="Only process these model types (e.g. loras diffusion_models)", ) args = parser.parse_args() @@ -48,7 +48,7 @@ def main() -> int: tasks = get_model_tasks(config, only_types=args.only) if not tasks: - print("No model URLs to download. Add URLs under diffusion_models, text_encoders, vaes, upscale_models, or loras in your config.") + print("No model URLs to download. Add entries under models with url and type in your config.") return 0 if args.dry_run: diff --git a/model_downloader/config.py b/model_downloader/config.py index 913ff72..fe46756 100644 --- a/model_downloader/config.py +++ b/model_downloader/config.py @@ -5,19 +5,19 @@ 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 _is_safe_type(t: str) -> bool: + """Reject type values that could escape models/ (path traversal).""" + if not t or not isinstance(t, str): + return False + if ".." in t or "/" in t or "\\" in t: + return False + return True + + def load_config(path: str | Path) -> dict[str, Any]: """Load YAML config from path. Raises FileNotFoundError or yaml error.""" p = Path(path) @@ -31,44 +31,54 @@ def load_config(path: str | Path) -> dict[str, Any]: def validate_config(data: dict[str, Any]) -> None: - """Validate required keys and model list types. Raises ValueError on failure.""" + """Validate required keys and models list. 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") + if "models" not in data: + data["models"] = [] + models = data["models"] + if not isinstance(models, list): + raise ValueError("'models' must be a list") + for i, item in enumerate(models): + if not isinstance(item, dict): + raise ValueError(f"models[{i}] must be an object with url and type") + url = item.get("url") + if not isinstance(url, str) or not url.strip(): + raise ValueError(f"models[{i}] must have a non-empty 'url' string") + t = item.get("type") + if not isinstance(t, str) or not t.strip(): + raise ValueError(f"models[{i}] must have a non-empty 'type' string") + if not _is_safe_type(t): + raise ValueError(f"models[{i}] 'type' must not contain /, \\, or ..") 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) + models = data.get("models") + if not isinstance(models, list): + return [] tasks: list[tuple[str, str]] = [] - for model_type in types: - if model_type not in MODEL_TYPE_SUBDIRS: + for item in models: + if not isinstance(item, dict): continue - urls = data.get(model_type) - if not isinstance(urls, list): + url = item.get("url") + t = item.get("type") + if not isinstance(url, str) or not url.strip(): continue - for url in urls: - if isinstance(url, str) and url.strip(): - tasks.append((model_type, url.strip())) + if not isinstance(t, str) or not t.strip() or not _is_safe_type(t): + continue + if only_types is not None and t not in only_types: + continue + tasks.append((t.strip(), 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 + return base / "models" / model_type diff --git a/model_downloader/download.py b/model_downloader/download.py index 0087a40..5df683d 100644 --- a/model_downloader/download.py +++ b/model_downloader/download.py @@ -129,7 +129,9 @@ def download_one( print(f" Would download to: {dest_dir / filename_guess}") return True - dest_dir.mkdir(parents=True, exist_ok=True) + if not dest_dir.is_dir(): + print(f" Skip: Directory does not exist: {dest_dir}. Create models/ under your ComfyUI base.") + return False # Single GET with stream (Civitai redirects to S3 and often doesn't support HEAD) try: