Download progress bar + fixed local_dir

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)
This commit is contained in:
Adrien Toupet
2025-07-17 21:48:05 -04:00
committed by GitHub
parent 3f26640f1d
commit afe33ddf0d
+50 -3
View File
@@ -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