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>

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:
-        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

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: