From afe33ddf0daf5eca55ae0fca0eb42da0e22e8d4d Mon Sep 17 00:00:00 2001
From: Adrien Toupet <cayus.net@gmail.com>
Date: Thu, 17 Jul 2025 21:48:05 -0400
Subject: [PATCH] Download progress bar + fixed local_dir
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Download progress is now displaying in the logs, e.g.
📥 Downloading model: seedvr2_ema_3b_fp16.safetensors
seedvr2_ema_3b_fp16.safetensors:   6% @ 45.9MB/s

local_dir is using base_cache_dir if model_dir is None else model_dir (was hardcoded to base_cache_dir before)
---
 src/utils/downloads.py | 53 +++++++++++++++++++++++++++++++++++++++---
 1 file changed, 50 insertions(+), 3 deletions(-)

diff --git a/src/utils/downloads.py b/src/utils/downloads.py
index a252027..6210f98 100644
--- a/src/utils/downloads.py
+++ b/src/utils/downloads.py
@@ -6,6 +6,7 @@ Extracted from: seedvr2.py (line 968-1015)
 """
 
 import os
+import sys
 from huggingface_hub import hf_hub_download
 try:
     import folder_paths
@@ -17,6 +18,38 @@ try:
 except:
     base_cache_dir = "./seedvr2_models"
 
+
+class ProgressCapture:
+    """Capture download progress from tqdm"""
+    def __init__(self, desc):
+        self.desc = desc
+        self.shown_percentages = set()
+        
+    def write(self, text):
+        if not text or text.isspace():
+            return
+        # Skip HF duplicate messages
+        if any(skip in text for skip in ["Downloading '", "Download complete.", "to '/data/"]):
+            return
+        # Extract and show progress
+        import re
+        match = re.search(r'(\d+)%', text)
+        if match:
+            percent = int(match.group(1))
+            if percent not in self.shown_percentages:
+                self.shown_percentages.add(percent)
+                speed_match = re.search(r'(\d+\.?\d*[KMGT]?B/s)', text)
+                speed = f" @ {speed_match.group(1)}" if speed_match else ""
+                sys.stdout.write(f"\r{self.desc}: {percent:3d}%{speed}")
+                sys.stdout.flush()
+                if percent == 100:
+                    sys.stdout.write("\n")
+                    sys.stdout.flush()
+    
+    def flush(self):
+        pass
+
+
 def download_weight(model, model_dir=None):
     """
     Télécharge un modèle SeedVR2 et son VAE associé depuis HuggingFace Hub
@@ -45,21 +78,35 @@ def download_weight(model, model_dir=None):
     # 🚀 Téléchargement du modèle principal
     if not os.path.exists(model_path):
         print(f"📥 Downloading model: {model}")
-        hf_hub_download(repo_id=repo_id, filename=model, local_dir=base_cache_dir)
+        
+        # Add progress capture
+        old_stderr = sys.stderr
+        sys.stderr = ProgressCapture(model)
+        try:
+            hf_hub_download(repo_id=repo_id, filename=model, local_dir=base_cache_dir if model_dir is None else model_dir)
+        finally:
+            sys.stderr = old_stderr
+            
         print(f"✅ Downloaded: {model}")
     
     # 🚀 Téléchargement du VAE avec stratégie de fallback
     if not os.path.exists(vae_fp16_path):
         print("📥 Downloading FP16 VAE SafeTensors...")
+        
+        # Add progress capture
+        old_stderr = sys.stderr
+        sys.stderr = ProgressCapture("ema_vae_fp16.safetensors")
         try:
             hf_hub_download(
                 repo_id=repo_id, 
                 filename="ema_vae_fp16.safetensors", 
-                local_dir=base_cache_dir
+                local_dir=base_cache_dir if model_dir is None else model_dir
             )
             print("✅ Downloaded: ema_vae_fp16.safetensors (FP16 SafeTensors)")
         except Exception as e:
             print(f"⚠️ FP16 SafeTensors VAE not available: {e}")
+        finally:
+            sys.stderr = old_stderr
     
     return
 
@@ -71,4 +118,4 @@ def get_base_cache_dir():
     Returns:
         str: Chemin du répertoire de cache
     """
-    return base_cache_dir
+    return base_cache_dir
\ No newline at end of file