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