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:
+50
-3
@@ -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
|
||||
Reference in New Issue
Block a user