Add download_missing_models.sh — interactive workflow model downloader

download_missing_models.sh

@@ -0,0 +1,277 @@
+#!/usr/bin/env bash
+# download_missing_models.sh
+# Interactively pick workflows from ./workflows, find every model they reference,
+# locate those models in the HF repo aleph65/ComfyUI (which mirrors ComfyUI/models/),
+# and download the missing ones into place, in parallel, with progress.
+#
+# Usage:
+#   ./download_missing_models.sh                 # interactive workflow picker
+#   ./download_missing_models.sh --all           # select every workflow
+#   ./download_missing_models.sh qwen-edit.json  # select specific workflow(s)
+#   ./download_missing_models.sh --dry-run ...   # show the plan, download nothing
+#
+# Auth: uses the HF_TOKEN environment variable if set, otherwise prompts for a token.
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+export HF_HUB_ENABLE_HF_TRANSFER=1
+
+# 1) Hugging Face token: use HF_TOKEN from the environment, otherwise prompt.
+if [[ -n "${HF_TOKEN:-}" ]]; then
+    echo "Using HF token from environment (HF_TOKEN)."
+else
+    read -r -s -p "Enter your Hugging Face token (input hidden): " HF_TOKEN
+    echo
+    [[ -n "$HF_TOKEN" ]] || { echo "No token provided — aborting."; exit 1; }
+fi
+export HF_TOKEN
+
+# 2) Ensure deps (huggingface_hub >= 1.x ships fast Xet downloads built in and no
+# longer provides the hf-transfer extra, so hf_transfer is best-effort only)
+if ! python3 -c "import huggingface_hub" >/dev/null 2>&1; then
+    echo "Installing huggingface_hub[hf_transfer]..."
+    pip install -q -U "huggingface_hub[hf_transfer]" || pip install -q -U huggingface_hub
+fi
+python3 -c "import hf_transfer" >/dev/null 2>&1 || pip install -q hf_transfer >/dev/null 2>&1 || true
+
+# 3) Run the downloader.
+# Load the Python code into a variable (NOT via stdin, which must stay attached
+# to the terminal for the interactive prompts).
+PYCODE=$(cat <<'PYEOF'
+import json
+import os
+import sys
+import glob
+
+WORKSPACE = os.environ.get("COMFY_WORKSPACE", "/workspace")
+WORKFLOWS_DIR = os.path.join(WORKSPACE, "workflows")
+COMFY_DIR = os.path.join(WORKSPACE, "ComfyUI")
+REPO_ID = "aleph65/ComfyUI"
+MODEL_EXTS = (".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".gguf", ".sft", ".onnx")
+SKIP_NODE_TYPES = {"MarkdownNote", "Note", "PrimitiveString", "String"}
+
+C_RESET, C_GREEN, C_YELLOW, C_RED, C_CYAN, C_BOLD = "\033[0m", "\033[32m", "\033[33m", "\033[31m", "\033[36m", "\033[1m"
+
+def human(n):
+    for unit in ("B", "KB", "MB", "GB", "TB"):
+        if n < 1024 or unit == "TB":
+            return f"{n:.1f} {unit}" if unit != "B" else f"{int(n)} B"
+        n /= 1024
+
+# ---------------- workflow selection ----------------
+
+def pick_workflows(files, argv):
+    names = [os.path.basename(f) for f in files]
+    args = [a for a in argv if not a.startswith("--")]
+    if "--all" in argv:
+        return files
+    if args:
+        chosen = []
+        for a in args:
+            matches = [f for f in files if os.path.basename(f) == a or os.path.basename(f) == a + ".json"]
+            if not matches:
+                sys.exit(f"No workflow named '{a}' in {WORKFLOWS_DIR}")
+            chosen += matches
+        return chosen
+    print(f"\n{C_BOLD}Available workflows:{C_RESET}")
+    for i, name in enumerate(names, 1):
+        print(f"  {i}) {name}")
+    print()
+    while True:
+        try:
+            raw = input("Select workflows (e.g. '1 3', '1,3', or 'all'; q to quit): ").strip().lower()
+        except EOFError:
+            sys.exit("\nNo input available. Run interactively, or pass workflow names or --all.")
+        if raw in ("q", "quit", "exit"):
+            sys.exit("Aborted.")
+        if raw in ("all", "a", "*"):
+            return files
+        picks = raw.replace(",", " ").split()
+        if picks and all(p.isdigit() and 1 <= int(p) <= len(names) for p in picks):
+            return [files[int(p) - 1] for p in dict.fromkeys(picks)]
+        print(f"  Invalid selection — enter numbers 1-{len(names)}, or 'all'.")
+
+# ---------------- model extraction ----------------
+
+def looks_like_model(v):
+    if not isinstance(v, str):
+        return False
+    v = v.strip()
+    if not v.lower().endswith(MODEL_EXTS):
+        return False
+    if "\n" in v or "http://" in v or "https://" in v or len(v) > 300:
+        return False
+    return True
+
+def extract_models(path):
+    """Return set of model refs (may include subdirs like 'flux2/klein/x.safetensors')."""
+    with open(path) as fh:
+        data = json.load(fh)
+    refs = set()
+
+    def scan_values(values):
+        stack = [values]
+        while stack:
+            v = stack.pop()
+            if isinstance(v, str) and looks_like_model(v):
+                refs.add(v.strip().replace("\\", "/"))
+            elif isinstance(v, list):
+                stack.extend(v)
+            elif isinstance(v, dict):
+                stack.extend(v.values())
+
+    def scan_nodes(nodes):
+        for n in nodes:
+            if n.get("type") in SKIP_NODE_TYPES or n.get("class_type") in SKIP_NODE_TYPES:
+                continue
+            if "widgets_values" in n:
+                scan_values(n["widgets_values"])
+            if "inputs" in n and isinstance(n["inputs"], dict):
+                scan_values(list(n["inputs"].values()))
+            # subgraph instances store nodes too
+            for sub in (n.get("subgraph") or {}).get("nodes", []) if isinstance(n.get("subgraph"), dict) else []:
+                scan_nodes([sub])
+
+    if isinstance(data, dict) and isinstance(data.get("nodes"), list):
+        scan_nodes(data["nodes"])
+        # subgraph definitions (ComfyUI >= subgraph support)
+        defs = (data.get("definitions") or {})
+        for sg in defs.get("subgraphs", []) or []:
+            if isinstance(sg.get("nodes"), list):
+                scan_nodes(sg["nodes"])
+    elif isinstance(data, dict):
+        # API format: {"1": {"class_type": ..., "inputs": {...}}, ...}
+        for n in data.values():
+            if isinstance(n, dict) and "class_type" in n:
+                if n.get("class_type") in SKIP_NODE_TYPES:
+                    continue
+                scan_values(list((n.get("inputs") or {}).values()))
+    return refs
+
+# ---------------- main ----------------
+
+def cleanup_debris():
+    """Remove partial-download fragments and stale locks left by interrupted runs."""
+    cache = os.path.join(COMFY_DIR, ".cache", "huggingface", "download")
+    freed, count = 0, 0
+    for root, _, names in os.walk(cache):
+        for name in names:
+            if name.endswith((".incomplete", ".lock")):
+                p = os.path.join(root, name)
+                try:
+                    freed += os.path.getsize(p)
+                    os.remove(p)
+                    count += 1
+                except OSError:
+                    pass
+    if count:
+        print(f"{C_YELLOW}Cleaned up {count} leftover partial-download file(s) ({human(freed)} freed).{C_RESET}")
+
+def main():
+    argv = sys.argv[1:]
+    dry_run = "--dry-run" in argv
+
+    cleanup_debris()
+    files = sorted(glob.glob(os.path.join(WORKFLOWS_DIR, "*.json")))
+    if not files:
+        sys.exit(f"No workflow JSONs found in {WORKFLOWS_DIR}")
+
+    chosen = pick_workflows(files, argv)
+    print(f"\n{C_BOLD}Selected workflows:{C_RESET}")
+    wanted = {}  # ref -> [workflow names]
+    for f in chosen:
+        refs = extract_models(f)
+        print(f"  • {os.path.basename(f)}  ({len(refs)} model refs)")
+        for r in refs:
+            wanted.setdefault(r, []).append(os.path.basename(f))
+    if not wanted:
+        sys.exit("No model references found in the selected workflows.")
+
+    print(f"\n{C_CYAN}Listing files in hf.co/{REPO_ID} ...{C_RESET}")
+    from huggingface_hub import HfApi
+    api = HfApi(token=os.environ.get("HF_TOKEN"))
+    repo_files = {}  # repo path -> size
+    for entry in api.list_repo_tree(REPO_ID, recursive=True):
+        if hasattr(entry, "size") and entry.path.startswith("models/"):
+            repo_files[entry.path] = entry.size or 0
+
+    # index by basename and by suffix path
+    by_basename = {}
+    for p in repo_files:
+        by_basename.setdefault(os.path.basename(p), []).append(p)
+
+    to_download, have, not_found = [], [], []
+    for ref, wfs in sorted(wanted.items()):
+        base = os.path.basename(ref)
+        candidates = by_basename.get(base, [])
+        # prefer a repo path that ends with the workflow's relative path (handles lora subdirs)
+        match = next((p for p in candidates if p.endswith("/" + ref) or p == "models/" + ref), None)
+        if match is None and candidates:
+            match = candidates[0]
+        if match is None:
+            not_found.append((ref, wfs))
+            continue
+        dest = os.path.join(COMFY_DIR, match)
+        size = repo_files[match]
+        if os.path.exists(dest) and (size == 0 or os.path.getsize(dest) == size):
+            have.append((ref, match, size))
+        else:
+            to_download.append((ref, match, size, dest))
+
+    print(f"\n{C_BOLD}Plan:{C_RESET}")
+    for ref, match, size in have:
+        print(f"  {C_GREEN}✔ have{C_RESET}     {match}  ({human(size)})")
+    for ref, match, size, dest in to_download:
+        print(f"  {C_YELLOW}↓ fetch{C_RESET}    {match}  ({human(size)})")
+    for ref, wfs in not_found:
+        print(f"  {C_RED}✘ missing{C_RESET}  {ref}  — not in hf.co/{REPO_ID} (used by {', '.join(wfs)})")
+
+    total = sum(s for _, _, s, _ in to_download)
+    print(f"\n{len(have)} already present, {len(to_download)} to download ({human(total)}), {len(not_found)} not in repo.")
+
+    if not to_download:
+        print(f"{C_GREEN}Nothing to download — all set!{C_RESET}")
+        return
+    if dry_run:
+        print("(dry run — nothing downloaded)")
+        return
+
+    if sys.stdin.isatty():
+        resp = input(f"\nDownload {len(to_download)} files ({human(total)}) to {COMFY_DIR}/models? [Y/n] ").strip().lower()
+        if resp not in ("", "y", "yes"):
+            sys.exit("Aborted.")
+
+    # hf_transfer parallelizes chunks within a file; snapshot_download parallelizes across files.
+    print(f"\n{C_CYAN}Downloading with hf_transfer (parallel)...{C_RESET}\n")
+    from huggingface_hub import snapshot_download
+    patterns = [m for _, m, _, _ in to_download]
+    snapshot_download(
+        repo_id=REPO_ID,
+        allow_patterns=patterns,
+        local_dir=COMFY_DIR,
+        token=os.environ.get("HF_TOKEN"),
+        max_workers=8,
+    )
+
+    print(f"\n{C_BOLD}Verifying:{C_RESET}")
+    ok = True
+    for _, match, size, dest in to_download:
+        if os.path.exists(dest) and (size == 0 or os.path.getsize(dest) == size):
+            print(f"  {C_GREEN}✔{C_RESET} {dest}  ({human(os.path.getsize(dest))})")
+        else:
+            ok = False
+            print(f"  {C_RED}✘{C_RESET} {dest}  (incomplete or missing)")
+    if not_found:
+        print(f"\n{C_YELLOW}Note:{C_RESET} {len(not_found)} model(s) were not found in the repo (listed above) — "
+              f"you'll need to source those elsewhere.")
+    print(f"\n{C_GREEN if ok else C_RED}{'Done — all downloads verified.' if ok else 'Done, but some files failed — re-run to retry.'}{C_RESET}")
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        sys.exit("\nInterrupted.")
+PYEOF
+)
+exec python3 -c "$PYCODE" "$@"