Labyricorn/voicebox
1
0
Fork 0
mirror of https://github.com/jamiepine/voicebox.git synced 2026-09-19 14:50:38 -07:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #25 from jamiepine/fix-dl-notification-when-generating-from-already-cached-model

Browse Source 
Fix dl notification when generating from already cached model
This commit is contained in:
Jamie Pine
2026-01-30 21:20:25 -08:00
committed by GitHub
parent 20851ccc2b 7d0557a099
commit 97eb570b28
27 changed files with 2232 additions and 295 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/utils/hf_progress.py
+153 -31
View File
@@ -11,8 +11,9 @@ import sys
class HFProgressTracker:
"""Tracks HuggingFace Hub download progress by intercepting tqdm."""
def __init__(self, progress_callback: Optional[Callable] = None):
def __init__(self, progress_callback: Optional[Callable] = None, filter_non_downloads: bool = False):
self.progress_callback = progress_callback
self.filter_non_downloads = filter_non_downloads # Only filter if True
self._original_tqdm_class = None
self._lock = threading.Lock()
self._total_downloaded = 0
@@ -21,6 +22,7 @@ class HFProgressTracker:
self._file_downloaded = {} # Track downloaded bytes per file
self._current_filename = ""
self._active_tqdms = {} # Track active tqdm instances
self._hf_tqdm_original_update = None # For monkey-patching hf's tqdm
def _create_tracked_tqdm_class(self):
"""Create a tqdm subclass that tracks progress."""
@@ -31,7 +33,6 @@ class HFProgressTracker:
"""A tqdm subclass that reports progress to our tracker."""
def __init__(self, *args, **kwargs):
print(f"[DEBUG TrackedTqdm] __init__ called with desc: {kwargs.get('desc', '')}")
# Extract filename from desc before passing to parent
desc = kwargs.get("desc", "")
if not desc and args:
@@ -80,7 +81,6 @@ class HFProgressTracker:
}
def update(self, n=1):
print(f"[DEBUG TrackedTqdm] update called with n={n}")
result = super().update(n)
# Report progress
@@ -91,6 +91,16 @@ class HFProgressTracker:
total = getattr(self, "total", 0)
if total and total > 0:
# Always filter out non-byte progress bars (e.g., "Fetching 12 files")
# These cause crazy percentages because they're counting files, not bytes
if self._is_non_byte_progress(filename):
return result
# When model is cached, also filter out generation-related progress
if tracker.filter_non_downloads:
if not self._is_download_progress(filename):
return result
# Update per-file tracking
tracker._file_sizes[filename] = total
tracker._file_downloaded[filename] = current
@@ -99,6 +109,13 @@ class HFProgressTracker:
tracker._total_size = sum(tracker._file_sizes.values())
tracker._total_downloaded = sum(tracker._file_downloaded.values())
# Only report progress once we have a meaningful total (at least 1MB)
# This avoids the "100% at 0MB" issue when small config
# files are counted before the real model files
MIN_TOTAL_BYTES = 1_000_000 # 1MB
if tracker._total_size < MIN_TOTAL_BYTES:
return result
# Call progress callback
if tracker.progress_callback:
tracker.progress_callback(
@@ -109,6 +126,50 @@ class HFProgressTracker:
return result
def _is_non_byte_progress(self, filename: str) -> bool:
"""Check if this progress bar should be SKIPPED (returns True to skip).
We want to track byte-based progress bars. This method identifies
progress bars that count files/items instead of bytes, which would
cause crazy percentages if mixed with our byte counting.
Returns:
True = SKIP this bar (it's not byte-based)
False = TRACK this bar (it counts bytes)
"""
if not filename:
return False
filename_lower = filename.lower()
# Skip "Fetching X files" - it counts files (total=12), not bytes
# Don't skip "Downloading (incomplete total...)" - that IS byte-based
skip_patterns = [
'fetching', # "Fetching 12 files" has total=12 files, not bytes
]
return any(pattern in filename_lower for pattern in skip_patterns)
def _is_download_progress(self, filename: str) -> bool:
"""Check if this is a real file download progress bar vs internal processing."""
if not filename or filename == "unknown":
return False
# Real downloads have file extensions
download_extensions = [
'.safetensors', '.bin', '.pt', '.pth', # Model weights
'.json', '.txt', '.py', # Config files
'.msgpack', '.h5', # Other formats
]
filename_lower = filename.lower()
has_extension = any(filename_lower.endswith(ext) for ext in download_extensions)
# Skip generation-related progress indicators
skip_patterns = ['segment', 'processing', 'generating', 'loading']
has_skip_pattern = any(pattern in filename_lower for pattern in skip_patterns)
return has_extension and not has_skip_pattern
def close(self):
with tracker._lock:
if id(self) in tracker._active_tqdms:
@@ -120,13 +181,11 @@ class HFProgressTracker:
@contextmanager
def patch_download(self):
"""Context manager to patch tqdm for progress tracking."""
print("[DEBUG HFProgressTracker] patch_download called")
try:
import tqdm as tqdm_module
# Store original tqdm class
self._original_tqdm_class = tqdm_module.tqdm
print(f"[DEBUG HFProgressTracker] Original tqdm class: {self._original_tqdm_class}")
# Reset totals
with self._lock:
@@ -139,39 +198,89 @@ class HFProgressTracker:
# Create our tracked tqdm class
tracked_tqdm = self._create_tracked_tqdm_class()
print(f"[DEBUG HFProgressTracker] Created TrackedTqdm class: {tracked_tqdm}")
# Patch tqdm.tqdm
tqdm_module.tqdm = tracked_tqdm
print(f"[DEBUG HFProgressTracker] Patched tqdm.tqdm")
# Also patch tqdm.auto.tqdm if it exists (used by huggingface_hub)
self._original_tqdm_auto = None
if hasattr(tqdm_module, "auto") and hasattr(tqdm_module.auto, "tqdm"):
self._original_tqdm_auto = tqdm_module.auto.tqdm
tqdm_module.auto.tqdm = tracked_tqdm
print(f"[DEBUG HFProgressTracker] Patched tqdm.auto.tqdm")
# Patch in sys.modules to catch already-imported references
# huggingface_hub uses: from tqdm.auto import tqdm as base_tqdm
# So we need to patch both 'tqdm' and 'base_tqdm' attributes
self._patched_modules = {}
tqdm_attr_names = ['tqdm', 'base_tqdm', 'old_tqdm'] # Various names used
patched_count = 0
for module_name in list(sys.modules.keys()):
if "huggingface" in module_name or module_name.startswith("tqdm"):
try:
module = sys.modules[module_name]
if hasattr(module, "tqdm"):
attr = getattr(module, "tqdm")
# Only patch if it's the original tqdm class (not already patched)
if attr is self._original_tqdm_class or (
hasattr(attr, "__name__") and attr.__name__ == "tqdm"
):
self._patched_modules[module_name] = attr
setattr(module, "tqdm", tracked_tqdm)
patched_count += 1
print(f"[DEBUG HFProgressTracker] Patched {module_name}.tqdm")
for attr_name in tqdm_attr_names:
if hasattr(module, attr_name):
attr = getattr(module, attr_name)
# Only patch if it's a tqdm class (not already patched)
is_tqdm_class = (
attr is self._original_tqdm_class or
(self._original_tqdm_auto and attr is self._original_tqdm_auto) or
(hasattr(attr, "__name__") and attr.__name__ == "tqdm" and
hasattr(attr, "update")) # tqdm classes have update method
)
if is_tqdm_class:
key = f"{module_name}.{attr_name}"
self._patched_modules[key] = (module, attr_name, attr)
setattr(module, attr_name, tracked_tqdm)
patched_count += 1
except (AttributeError, TypeError):
pass
print(f"[DEBUG HFProgressTracker] Patched {patched_count} modules in sys.modules")
# ALSO monkey-patch the update method on huggingface_hub's tqdm class
# This is needed because the class was already defined at import time
self._hf_tqdm_original_update = None
try:
from huggingface_hub.utils import tqdm as hf_tqdm_module
if hasattr(hf_tqdm_module, 'tqdm'):
hf_tqdm_class = hf_tqdm_module.tqdm
self._hf_tqdm_original_update = hf_tqdm_class.update
# Create a wrapper that calls our tracking
tracker = self # Reference to HFProgressTracker instance
def patched_update(tqdm_self, n=1):
result = tracker._hf_tqdm_original_update(tqdm_self, n)
# Track this progress
with tracker._lock:
desc = getattr(tqdm_self, 'desc', '') or ''
current = getattr(tqdm_self, 'n', 0)
total = getattr(tqdm_self, 'total', 0) or 0
# Skip non-byte progress bars
if 'fetching' in desc.lower():
return result
# Skip until we have a meaningful total (at least 1MB)
# This avoids the "100% at 0MB" issue when small config
# files are counted before the real model files
MIN_TOTAL_BYTES = 1_000_000 # 1MB
if total >= MIN_TOTAL_BYTES:
tracker._total_downloaded = current
tracker._total_size = total
if tracker.progress_callback:
tracker.progress_callback(current, total, desc)
return result
hf_tqdm_class.update = patched_update
patched_count += 1
print(f"[HFProgressTracker] Monkey-patched huggingface_hub.utils.tqdm.tqdm.update")
except (ImportError, AttributeError) as e:
print(f"[HFProgressTracker] Could not monkey-patch hf_tqdm: {e}")
print(f"[HFProgressTracker] Patched {patched_count} tqdm references")
yield
@@ -189,15 +298,24 @@ class HFProgressTracker:
tqdm_module.auto.tqdm = self._original_tqdm_auto
# Restore patched modules
for module_name, original in self._patched_modules.items():
for key, (module, attr_name, original) in self._patched_modules.items():
try:
module = sys.modules.get(module_name)
if module and original:
setattr(module, "tqdm", original)
setattr(module, attr_name, original)
except (AttributeError, TypeError):
pass
self._patched_modules = {}
# Restore hf_tqdm's original update method
if self._hf_tqdm_original_update:
try:
from huggingface_hub.utils import tqdm as hf_tqdm_module
if hasattr(hf_tqdm_module, 'tqdm'):
hf_tqdm_module.tqdm.update = self._hf_tqdm_original_update
except (ImportError, AttributeError):
pass
self._hf_tqdm_original_update = None
except (ImportError, AttributeError):
pass
@@ -205,13 +323,17 @@ class HFProgressTracker:
def create_hf_progress_callback(model_name: str, progress_manager):
"""Create a progress callback for HuggingFace downloads."""
def callback(downloaded: int, total: int, filename: str = ""):
"""Progress callback."""
if total > 0:
progress_manager.update_progress(
model_name=model_name,
current=downloaded,
total=total,
filename=filename or "",
status="downloading",
)
"""Progress callback.
Note: We send updates even when total=0 (unknown) to provide feedback
during the "incomplete total" phase of huggingface_hub downloads.
The frontend handles total=0 gracefully.
"""
progress_manager.update_progress(
model_name=model_name,
current=downloaded,
total=total,
filename=filename or "",
status="downloading",
)
return callback
backend/utils/progress.py
+42 -11
View File
@@ -16,11 +16,17 @@ class ProgressManager:
Thread-safe: can be called from background threads (e.g., via asyncio.to_thread).
"""
# Throttle settings to prevent overwhelming SSE clients
THROTTLE_INTERVAL_SECONDS = 0.5 # Minimum time between updates
THROTTLE_PROGRESS_DELTA = 1.0 # Minimum progress change (%) to force update
def __init__(self):
self._progress: Dict[str, Dict] = {}
self._listeners: Dict[str, list] = {}
self._lock = threading.Lock() # Thread-safe lock for progress dict
self._main_loop: Optional[asyncio.AbstractEventLoop] = None
self._last_notify_time: Dict[str, float] = {} # Last notification time per model
self._last_notify_progress: Dict[str, float] = {} # Last notified progress per model
def _set_main_loop(self, loop: asyncio.AbstractEventLoop):
"""Set the main event loop for thread-safe operations."""
@@ -67,6 +73,10 @@ class ProgressManager:
Update progress for a model download.
Thread-safe: can be called from background threads.
Progress updates are throttled to prevent overwhelming SSE clients.
Updates are sent at most every THROTTLE_INTERVAL_SECONDS, or when
progress changes by at least THROTTLE_PROGRESS_DELTA percent.
Args:
model_name: Name of the model (e.g., "qwen-tts-1.7B", "whisper-base")
@@ -76,9 +86,17 @@ class ProgressManager:
status: Status string (downloading, extracting, complete, error)
"""
import logging
import time
logger = logging.getLogger(__name__)
progress_pct = (current / total * 100) if total > 0 else 0
# Calculate progress percentage, clamped to 0-100 range
# This prevents crazy percentages from edge cases like:
# - current > total temporarily during aggregation
# - mixing file-count progress with byte-count progress
if total > 0:
progress_pct = min(100.0, max(0.0, (current / total * 100)))
else:
progress_pct = 0
progress_data = {
"model_name": model_name,
@@ -90,25 +108,38 @@ class ProgressManager:
"timestamp": datetime.now().isoformat(),
}
print(f"[DEBUG] update_progress called: {model_name}, {progress_pct:.1f}%")
# Thread-safe update of progress dict
# Thread-safe update of progress dict (always update internal state)
with self._lock:
self._progress[model_name] = progress_data
# Check if we should notify listeners (throttling)
current_time = time.time()
last_time = self._last_notify_time.get(model_name, 0)
last_progress = self._last_notify_progress.get(model_name, -100)
time_delta = current_time - last_time
progress_delta = abs(progress_pct - last_progress)
# Always notify for complete/error status, or if throttle conditions are met
should_notify = (
status in ("complete", "error") or
time_delta >= self.THROTTLE_INTERVAL_SECONDS or
progress_delta >= self.THROTTLE_PROGRESS_DELTA
)
if not should_notify:
return # Skip this update (throttled)
# Update throttle tracking
self._last_notify_time[model_name] = current_time
self._last_notify_progress[model_name] = progress_pct
# Notify all listeners (thread-safe)
listener_count = len(self._listeners.get(model_name, []))
print(f"[DEBUG] Listener count for {model_name}: {listener_count}")
print(f"[DEBUG] All listeners: {list(self._listeners.keys())}")
print(f"[DEBUG] Main loop set: {self._main_loop is not None}")
if self._main_loop:
print(f"[DEBUG] Main loop running: {self._main_loop.is_running()}")
if listener_count > 0:
logger.debug(f"Notifying {listener_count} listeners for {model_name}: {progress_pct:.1f}% ({filename})")
print(f"[DEBUG] About to notify listeners...")
self._notify_listeners_threadsafe(model_name, progress_data)
print(f"[DEBUG] Notified listeners")
else:
logger.debug(f"No listeners for {model_name}, progress update stored: {progress_pct:.1f}%")