download_missing_models: fix mid-download deadlock (disable HF_XET_HIGH_PERFORMANCE/hf_transfer, add stall watchdog + resumable curl fallback)

download_missing_models.sh

@@ -30,7 +30,12 @@
 set -euo pipefail
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-export HF_HUB_ENABLE_HF_TRANSFER=1
+# The pod image exports HF_HUB_ENABLE_HF_TRANSFER=1 and HF_XET_HIGH_PERFORMANCE=1.
+# hf_transfer is deprecated in huggingface_hub 1.x, and the Xet high-performance
+# mode has deadlocked mid-download on this pod (~115 parallel connections, all
+# data buffered in RAM, zero bytes written, no error). Stock Xet concurrency is
+# still several hundred MB/s here, which is what the network volume tops out at.
+unset HF_HUB_ENABLE_HF_TRANSFER HF_XET_HIGH_PERFORMANCE
 # never let git sit on an interactive credential/hostkey prompt
 export GIT_TERMINAL_PROMPT=0
 
@@ -47,13 +52,13 @@ else
 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)
+# 2) Ensure deps. huggingface_hub 1.x downloads Xet-backed repos with hf_xet;
+# without it the hub falls back to slower, less robust code paths.
 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
+    echo "Installing huggingface_hub..."
+    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
+python3 -c "import hf_xet" >/dev/null 2>&1 || pip install -q hf_xet >/dev/null 2>&1 || true
 
 # 3) Run the downloader.
 # Load the Python code into a variable (NOT via stdin, which must stay attached
@@ -502,29 +507,102 @@ def cleanup_debris():
     if count:
         print(f"{C_YELLOW}Cleaned up {count} leftover partial-download file(s) ({human(freed)} freed).{C_RESET}")
 
+def bytes_on_disk(to_download):
+    """Bytes actually landed for the planned downloads: finished files,
+    hub .incomplete fragments, and our curl .part files."""
+    done = 0
+    cache = os.path.join(COMFY_DIR, ".cache", "huggingface", "download")
+    for _, _, size, dest in to_download:
+        if os.path.exists(dest):
+            done += min(os.path.getsize(dest), size)
+        if os.path.exists(dest + ".part"):
+            done += os.path.getsize(dest + ".part")
+    for root, _, names in os.walk(cache):
+        for n in names:
+            if n.endswith(".incomplete"):
+                try:
+                    done += os.path.getsize(os.path.join(root, n))
+                except OSError:
+                    pass
+    return done
+
 def download_heartbeat(to_download):
     """Print bytes-on-disk every 20s so a big download never looks hung."""
     total = sum(s for _, _, s, _ in to_download)
     stop = threading.Event()
-    cache = os.path.join(COMFY_DIR, ".cache", "huggingface", "download")
     def run():
         while not stop.wait(20):
-            done = 0
-            for _, _, size, dest in to_download:
-                if os.path.exists(dest):
-                    done += min(os.path.getsize(dest), size)
-            for root, _, names in os.walk(cache):
-                for n in names:
-                    if n.endswith(".incomplete"):
-                        try:
-                            done += os.path.getsize(os.path.join(root, n))
-                        except OSError:
-                            pass
+            done = bytes_on_disk(to_download)
             print(f"  ... still downloading: ~{human(min(done, total))} of {human(total)} on disk", flush=True)
     t = threading.Thread(target=run, daemon=True)
     t.start()
     return stop
 
+STALL_SECS = 180
+
+def hf_download_with_watchdog(patterns, to_download):
+    """Run snapshot_download in a child process and kill it if bytes-on-disk
+    stop growing for STALL_SECS. Both hf_transfer and the Xet backend have
+    wedged mid-file on this pod (data buffered in RAM, file frozen, no error
+    raised), so a blocking in-process call can hang forever. Returns True if
+    the hub download finished."""
+    import multiprocessing as mp
+    def child():
+        from huggingface_hub import snapshot_download
+        snapshot_download(
+            repo_id=REPO_ID,
+            allow_patterns=patterns,
+            local_dir=COMFY_DIR,
+            token=os.environ.get("HF_TOKEN"),
+            max_workers=8,
+        )
+    p = mp.Process(target=child)
+    p.start()
+    last, last_t = -1, time.time()
+    while p.is_alive():
+        p.join(timeout=15)
+        if not p.is_alive():
+            break
+        done = bytes_on_disk(to_download)
+        if done > last:
+            last, last_t = done, time.time()
+        elif time.time() - last_t > STALL_SECS:
+            print(f"\n  {C_YELLOW}⚠ no bytes written for {STALL_SECS}s — hub download is wedged, killing it{C_RESET}", flush=True)
+            p.terminate()
+            p.join(10)
+            if p.is_alive():
+                p.kill()
+                p.join()
+            return False
+    return p.exitcode == 0
+
+def curl_fallback(to_download):
+    """Plain resumable HTTP for whatever is still missing. --speed-limit makes
+    curl abort any transfer that drops below 1 MB/s for 30s (instead of
+    hanging on a dead connection), and -C - resumes from the .part file."""
+    token = os.environ.get("HF_TOKEN", "")
+    for _, match, size, dest in to_download:
+        if os.path.exists(dest) and (size == 0 or os.path.getsize(dest) == size):
+            continue
+        url = f"https://huggingface.co/{REPO_ID}/resolve/main/{match}"
+        part = dest + ".part"
+        os.makedirs(os.path.dirname(dest), exist_ok=True)
+        print(f"  {C_CYAN}fetching {match} with curl (resumable) ...{C_RESET}", flush=True)
+        for attempt in range(1, 11):
+            r = subprocess.run(
+                ["curl", "-L", "--fail", "-C", "-",
+                 "-H", f"Authorization: Bearer {token}",
+                 "--speed-limit", "1000000", "--speed-time", "30",
+                 "--progress-bar", "-o", part, url],
+                stdin=subprocess.DEVNULL)
+            have = os.path.getsize(part) if os.path.exists(part) else 0
+            if r.returncode == 0 and (size == 0 or have == size):
+                os.replace(part, dest)
+                break
+            print(f"  {C_YELLOW}… attempt {attempt} stopped at {human(have)} of {human(size)} — resuming{C_RESET}", flush=True)
+        else:
+            print(f"  {C_RED}✘ gave up on {match} after 10 attempts{C_RESET}")
+
 def main():
     argv = sys.argv[1:]
     dry_run = "--dry-run" in argv
@@ -656,21 +734,16 @@ def main():
 
     # -------- download models --------
     if to_download:
-        # 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
+        print(f"\n{C_CYAN}Downloading with the Hugging Face hub (Xet backend, parallel)...{C_RESET}\n")
         patterns = [m for _, m, _, _ in to_download]
         heartbeat = download_heartbeat(to_download)
         try:
-            snapshot_download(
-                repo_id=REPO_ID,
-                allow_patterns=patterns,
-                local_dir=COMFY_DIR,
-                token=os.environ.get("HF_TOKEN"),
-                max_workers=8,
-            )
+            finished = hf_download_with_watchdog(patterns, to_download)
         finally:
             heartbeat.set()
+        if not finished:
+            print(f"\n{C_YELLOW}Hub download did not finish cleanly — switching to plain resumable HTTP.{C_RESET}")
+            curl_fallback(to_download)
 
     print(f"\n{C_BOLD}Verifying:{C_RESET}")
     ok = True