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 <iamdanhenry@gmail.com>

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/`) |
+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/<type>` directory must already exist under your ComfyUI base; the tool does not create these directories and will fail if a directory is missing.
-|--------------------|--------------------------------------------------|
-| diffusion_models   | diffusion_models                                 |
-| text_encoders      | text_encoders                                    |
-| vaes               | vae                                              |
-| upscale_models     | upscale_models                                   |
-| loras              | loras                                            |
 
 Tokens stay in your local `config.yaml` only; keep them out of git.
 

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: []
+# models:
-text_encoders: []
+#   - url: "https://..."
-vaes: []
+#     type: loras
-upscale_models: []
+models: []
-loras: []

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:

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:
+    if "models" not in data:
-        val = data.get(key)
+        data["models"] = []
-        if val is None:
+    models = data["models"]
-            data[key] = []
+    if not isinstance(models, list):
-        elif not isinstance(val, list):
+        raise ValueError("'models' must be a list")
-            raise ValueError(f"'{key}' must be a list of URL strings")
+    for i, item in enumerate(models):
-        else:
+        if not isinstance(item, dict):
-            for i, item in enumerate(val):
+            raise ValueError(f"models[{i}] must be an object with url and type")
-                if not isinstance(item, str):
+        url = item.get("url")
-                    raise ValueError(f"'{key}[{i}]' must be a string")
+        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:
+    for item in models:
-        if model_type not in MODEL_TYPE_SUBDIRS:
+        if not isinstance(item, dict):
             continue
-        urls = data.get(model_type)
+        url = item.get("url")
-        if not isinstance(urls, list):
+        t = item.get("type")
+        if not isinstance(url, str) or not url.strip():
             continue
-        for url in urls:
+        if not isinstance(t, str) or not t.strip() or not _is_safe_type(t):
-            if isinstance(url, str) and url.strip():
+            continue
-                tasks.append((model_type, url.strip()))
+        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" / model_type
-    return base / "models" / subdir

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/<type> under your ComfyUI base.")
+        return False
 
     # Single GET with stream (Civitai redirects to S3 and often doesn't support HEAD)
     try: