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

Refactor model download progress tracking and enhance SSE handling

Browse Source 
- Rearranged imports for consistency in useModelDownloadToast hook.
- Improved logging in useModelDownloadToast for better debugging during download events.
- Updated progress calculation to handle cases where progress exceeds 100%.
- Enhanced toast notifications to reflect download completion and error states.
- Introduced throttling in ProgressManager to optimize SSE updates and prevent overwhelming clients.
- Added new test scripts for monitoring SSE events during model downloads, ensuring accurate progress reporting.
This commit is contained in:
Jamie Pine
2026-01-30 20:18:53 -08:00
parent 07c0aba883
commit d3393fb940
9 changed files with 414 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/lib/hooks/useModelDownloadToast.tsx
+42 -27
View File
@@ -1,9 +1,9 @@
import { CheckCircle2, Loader2, XCircle } from 'lucide-react';
import { useCallback, useEffect, useRef } from 'react';
import { useToast } from '@/components/ui/use-toast';
import { useServerStore } from '@/stores/serverStore';
import { Progress } from '@/components/ui/progress';
import { Loader2, CheckCircle2, XCircle } from 'lucide-react';
import { useToast } from '@/components/ui/use-toast';
import type { ModelProgress } from '@/lib/api/types';
import { useServerStore } from '@/stores/serverStore';
interface UseModelDownloadToastOptions {
modelName: string;
@@ -36,19 +36,24 @@ export function useModelDownloadToast({
const k = 1024;
const sizes = ['B', 'KB', 'MB', 'GB'];
const i = Math.floor(Math.log(bytes) / Math.log(k));
return `${(bytes / Math.pow(k, i)).toFixed(1)} ${sizes[i]}`;
return `${(bytes / k ** i).toFixed(1)} ${sizes[i]}`;
}, []);
useEffect(() => {
console.log('[useModelDownloadToast] useEffect triggered', { enabled, serverUrl, modelName, displayName });
console.log('[useModelDownloadToast] useEffect triggered', {
enabled,
serverUrl,
modelName,
displayName,
});
if (!enabled || !serverUrl || !modelName) {
console.log('[useModelDownloadToast] Not enabled, skipping');
return;
}
console.log('[useModelDownloadToast] Creating toast and EventSource for:', modelName);
// Create initial toast
const toastResult = toast({
title: displayName,
@@ -131,30 +136,40 @@ export function useModelDownloadToast({
});
// Close connection and dismiss toast on completion or error
if (progress.status === 'complete' || progress.status === 'error') {
// Also treat progress >= 100% as complete
const isComplete = progress.status === 'complete' || progress.progress >= 100;
const isError = progress.status === 'error';
if (isComplete || isError) {
console.log('[useModelDownloadToast] Download finished:', {
isComplete,
isError,
progress: progress.progress,
});
eventSource.close();
eventSourceRef.current = null;
// Call callbacks
if (progress.status === 'complete' && onComplete) {
console.log('[useModelDownloadToast] Download complete, calling onComplete callback');
onComplete();
} else if (progress.status === 'error' && onError) {
console.log('[useModelDownloadToast] Download error, calling onError callback');
onError();
// Update toast to show completion state before callbacks
if (isComplete && toastUpdateRef.current) {
toastUpdateRef.current({
title: (
<div className="flex items-center gap-2">
<CheckCircle2 className="h-4 w-4 text-green-500" />
<span>{displayName}</span>
</div>
),
description: 'Download complete',
duration: 3000,
});
}
// Auto-dismiss on completion after delay
if (progress.status === 'complete') {
setTimeout(() => {
if (toastIdRef.current && toastUpdateRef.current) {
toastUpdateRef.current({
open: false,
});
toastIdRef.current = null;
toastUpdateRef.current = null;
}
}, 5000);
// Call callbacks
if (isComplete && onComplete) {
console.log('[useModelDownloadToast] Download complete, calling onComplete callback');
onComplete();
} else if (isError && onError) {
console.log('[useModelDownloadToast] Download error, calling onError callback');
onError();
}
}
}
@@ -198,4 +213,4 @@ export function useModelDownloadToast({
return {
isTracking: enabled && eventSourceRef.current !== null,
};
}
}
backend/backends/mlx_backend.py
+20 -25
View File
@@ -97,9 +97,7 @@ class MLXTTSBackend:
def _load_model_sync(self, model_size: str):
"""Synchronous model loading."""
try:
from mlx_audio.tts import load
# Get model path
# Get model path BEFORE importing mlx_audio
model_path = self._get_model_path(model_size)
# Set up progress tracking
@@ -123,19 +121,24 @@ class MLXTTSBackend:
# Start tracking download task
task_manager.start_download(model_name)
# Initialize progress state
progress_manager.update_progress(
model_name=model_name,
current=0,
total=1,
filename="",
status="downloading",
)
# Note: Don't initialize progress here with fake total=1
# Let the actual tqdm callbacks set the real values
# This avoids crazy percentages when current > 1 but total is still 1
# Use progress tracker (tqdm is patched, but filters out non-download progress)
with tracker.patch_download():
# Load MLX model (downloads automatically)
# IMPORTANT: Patch tqdm BEFORE importing mlx_audio
# Otherwise mlx_audio caches reference to original tqdm
tracker_context = tracker.patch_download()
tracker_context.__enter__()
# Import mlx_audio AFTER patching tqdm
from mlx_audio.tts import load
# Load MLX model (downloads automatically)
try:
self.model = load(model_path)
finally:
# Exit the patch context
tracker_context.__exit__(None, None, None)
# Only mark download as complete if we were tracking it
if not is_cached:
@@ -412,10 +415,8 @@ class MLXSTTBackend:
tracker = HFProgressTracker(progress_callback, filter_non_downloads=is_cached)
# Patch tqdm BEFORE importing mlx_audio
print("[DEBUG] Starting tqdm patch BEFORE mlx_audio import")
tracker_context = tracker.patch_download()
tracker_context.__enter__()
print("[DEBUG] tqdm patched")
# Import mlx_audio
from mlx_audio.stt import load
@@ -429,15 +430,9 @@ class MLXSTTBackend:
if not is_cached:
# Start tracking download task
task_manager.start_download(progress_model_name)
# Initialize progress state
progress_manager.update_progress(
model_name=progress_model_name,
current=0,
total=1,
filename="",
status="downloading",
)
# Note: Don't initialize progress here with fake total=1
# Let the actual tqdm callbacks set the real values
# Load the model (tqdm is patched, but filters out non-download progress)
try:
backend/tests/test_check_progress_state.py
-48
View File
@@ -1,48 +0,0 @@
"""
Check the internal state of ProgressManager.
"""
import sys
import time
# Add backend to path
sys.path.insert(0, '.')
from utils.progress import get_progress_manager
from utils.tasks import get_task_manager
def main():
pm = get_progress_manager()
tm = get_task_manager()
print("=" * 60)
print("ProgressManager State Inspector")
print("=" * 60)
# Check all progress
print("\nAll progress entries:")
print(f" _progress dict: {pm._progress}")
# Check listeners
print("\nActive listeners:")
print(f" _listeners dict: {pm._listeners}")
# Check main loop
print(f"\nMain loop set: {pm._main_loop is not None}")
if pm._main_loop:
print(f" Loop running: {pm._main_loop.is_running()}")
# Check task manager
print("\nTaskManager state:")
print(f" Active downloads: {tm.get_active_downloads()}")
print(f" Active generations: {tm.get_active_generations()}")
# Check for specific model
print("\nChecking whisper-base specifically:")
progress = pm.get_progress("whisper-base")
if progress:
print(f" Progress: {progress}")
else:
print(" No progress data found")
if __name__ == "__main__":
main()
backend/tests/test_generation_progress.py → backend/tests/test_generation_download.py
View File
backend/tests/test_qwen_download.py
+317
View File
@@ -0,0 +1,317 @@
"""
Test Qwen TTS model download with SSE progress monitoring.
This specifically tests the MLX TTS backend download progress tracking,
which requires tqdm to be patched BEFORE mlx_audio is imported.
Usage:
cd backend && python -m tests.test_qwen_download
Prerequisites:
- Server must be running: cd backend && python main.py
- Delete model first for fresh download test:
curl -X DELETE http://localhost:8000/models/qwen-tts-0.6B
"""
import asyncio
import json
import httpx
import time
from typing import List, Dict, Optional
async def monitor_sse_stream(model_name: str, timeout: int = 600) -> List[Dict]:
"""
Monitor SSE stream for a model download.
Args:
model_name: Name of the model to monitor
timeout: Maximum time to wait for download (seconds)
Returns:
List of SSE events received
"""
events: List[Dict] = []
url = f"http://localhost:8000/models/progress/{model_name}"
last_progress = -1
print(f"\n📡 Connecting to SSE endpoint: {url}")
try:
async with httpx.AsyncClient(timeout=timeout) as client:
async with client.stream("GET", url) as response:
print(f" SSE connected, status: {response.status_code}")
if response.status_code != 200:
print(f" ❌ Error: SSE endpoint returned {response.status_code}")
return events
async for line in response.aiter_lines():
if not line:
continue
if line.startswith("data: "):
try:
data = json.loads(line[6:])
events.append(data)
# Print progress (only when it changes significantly)
progress = data.get('progress', 0)
status = data.get('status', 'unknown')
filename = data.get('filename', '')
current = data.get('current', 0)
total = data.get('total', 0)
# Print every 5% change or status change
if abs(progress - last_progress) >= 5 or status in ('complete', 'error'):
current_mb = current / (1024 * 1024)
total_mb = total / (1024 * 1024)
print(f" 📊 {status:12} {progress:6.1f}% ({current_mb:.1f}MB / {total_mb:.1f}MB) {filename[:50]}")
last_progress = progress
# Stop if complete or error
if status in ("complete", "error"):
if status == "complete":
print(f" ✅ Download complete!")
else:
print(f" ❌ Download error: {data.get('error', 'unknown')}")
break
except json.JSONDecodeError as e:
print(f" ⚠️ Error parsing JSON: {e}")
elif line.startswith(": heartbeat"):
# Heartbeat every 1 second, don't spam
pass
except asyncio.CancelledError:
print(" ⏹️ SSE monitor cancelled")
except Exception as e:
print(f" ❌ SSE error: {e}")
return events
async def trigger_download(model_name: str) -> bool:
"""Trigger a model download via the API."""
url = "http://localhost:8000/models/download"
print(f"\n🚀 Triggering download for: {model_name}")
try:
async with httpx.AsyncClient(timeout=30) as client:
response = await client.post(url, json={"model_name": model_name})
result = response.json()
print(f" Response: {response.status_code} - {result}")
return response.status_code == 200
except Exception as e:
print(f" ❌ Error triggering download: {e}")
return False
async def delete_model(model_name: str) -> bool:
"""Delete a model from cache."""
url = f"http://localhost:8000/models/{model_name}"
print(f"\n🗑️ Deleting model: {model_name}")
try:
async with httpx.AsyncClient(timeout=30) as client:
response = await client.delete(url)
if response.status_code == 200:
print(f" ✅ Model deleted")
return True
elif response.status_code == 404:
print(f" ℹ️ Model not found (already deleted)")
return True
else:
print(f" ⚠️ Delete response: {response.status_code} - {response.text}")
return False
except Exception as e:
print(f" ❌ Error deleting model: {e}")
return False
async def check_model_status(model_name: str) -> Optional[Dict]:
"""Check the status of a model."""
try:
async with httpx.AsyncClient(timeout=10) as client:
response = await client.get("http://localhost:8000/models/status")
if response.status_code == 200:
data = response.json()
for model in data.get("models", []):
if model["model_name"] == model_name:
return model
except Exception as e:
print(f" ⚠️ Error checking model status: {e}")
return None
async def check_server() -> bool:
"""Check if the server is running."""
try:
async with httpx.AsyncClient(timeout=5) as client:
response = await client.get("http://localhost:8000/health")
return response.status_code == 200
except Exception:
return False
async def main():
print("=" * 70)
print("🧪 Qwen TTS Model Download Progress Test")
print("=" * 70)
print("\nThis test verifies that MLX TTS download progress tracking works.")
print("It specifically tests the tqdm patching for mlx_audio.tts imports.")
# Check if server is running
print("\n📡 Checking if server is running...")
if not await check_server():
print(" ❌ Server is not running on http://localhost:8000")
print("\n Please start the server first:")
print(" cd backend && python main.py")
return False
print(" ✅ Server is running")
# Test model
model_name = "qwen-tts-0.6B" # Note: 0.6B currently maps to 1.7B on MLX
# Check current status
print(f"\n📊 Checking status of {model_name}...")
status = await check_model_status(model_name)
if status:
print(f" Downloaded: {status.get('downloaded', False)}")
print(f" Downloading: {status.get('downloading', False)}")
print(f" Loaded: {status.get('loaded', False)}")
if status.get('size_mb'):
print(f" Size: {status['size_mb']:.1f} MB")
else:
print(" ⚠️ Could not get model status")
# Ask if user wants to delete first
print("\n" + "-" * 70)
if status and status.get('downloaded'):
print("⚠️ Model is already downloaded. Delete it for a fresh download test?")
print(" [y] Yes, delete and download fresh")
print(" [n] No, just test SSE connection")
print(" [q] Quit")
choice = input("\nChoice [y/n/q]: ").strip().lower()
if choice == 'q':
print("Exiting...")
return True
if choice == 'y':
if not await delete_model(model_name):
print("Failed to delete model. Continue anyway? [y/n]")
if input().strip().lower() != 'y':
return False
else:
print("Model not downloaded. Will perform fresh download test.")
input("Press Enter to continue...")
# Run the test
print("\n" + "=" * 70)
print("🏃 Starting Download Test")
print("=" * 70)
async def run_test():
# Start SSE monitor in background FIRST
monitor_task = asyncio.create_task(monitor_sse_stream(model_name, timeout=600))
# Wait for SSE to connect
await asyncio.sleep(1)
# Trigger download
success = await trigger_download(model_name)
if not success:
print(" ❌ Failed to trigger download")
monitor_task.cancel()
try:
await monitor_task
except asyncio.CancelledError:
pass
return []
# Wait for SSE monitor to complete
print("\n⏳ Waiting for download to complete (this may take several minutes)...")
events = await monitor_task
return events
start_time = time.time()
events = await run_test()
elapsed = time.time() - start_time
# Results
print("\n" + "=" * 70)
print("📋 Test Results")
print("=" * 70)
print(f"\n⏱️ Elapsed time: {elapsed:.1f} seconds")
print(f"📨 Total SSE events received: {len(events)}")
if not events:
print("\n❌ FAILED - No SSE events received!")
print("\nPossible causes:")
print(" 1. SSE endpoint not working")
print(" 2. tqdm not patched before mlx_audio import")
print(" 3. Progress callbacks not firing")
print(" 4. Model already fully downloaded")
print("\nDebug steps:")
print(" 1. Check server logs for [DEBUG] messages")
print(" 2. Look for 'tqdm patched' before 'mlx_audio.tts import'")
print(f" 3. Delete model: curl -X DELETE http://localhost:8000/models/{model_name}")
return False
# Analyze events
first_event = events[0]
last_event = events[-1]
print(f"\n📊 First event:")
print(f" Status: {first_event.get('status')}")
print(f" Progress: {first_event.get('progress', 0):.1f}%")
print(f"\n📊 Last event:")
print(f" Status: {last_event.get('status')}")
print(f" Progress: {last_event.get('progress', 0):.1f}%")
# Check for expected behaviors
has_progress_updates = len(events) > 2
has_increasing_progress = False
has_complete = any(e.get('status') == 'complete' for e in events)
has_100_percent = any(e.get('progress', 0) >= 100 for e in events)
# Check if progress increased over time
if len(events) >= 2:
progress_values = [e.get('progress', 0) for e in events]
has_increasing_progress = progress_values[-1] > progress_values[0]
print("\n📋 Checks:")
print(f" {'✅' if has_progress_updates else '❌'} Multiple progress updates received ({len(events)} events)")
print(f" {'✅' if has_increasing_progress else '❌'} Progress increased over time")
print(f" {'✅' if has_100_percent else '❌'} Reached 100% progress")
print(f" {'✅' if has_complete else '❌'} Received 'complete' status")
# Overall result
success = has_progress_updates and has_complete
if success:
print("\n" + "=" * 70)
print("✅ TEST PASSED - Qwen TTS download progress tracking works!")
print("=" * 70)
else:
print("\n" + "=" * 70)
print("❌ TEST FAILED - Progress tracking has issues")
print("=" * 70)
print("\nCheck the server logs for debug output.")
return success
if __name__ == "__main__":
result = asyncio.run(main())
exit(0 if result else 1)
backend/tests/test_real_download.py → backend/tests/test_whisper_download.py
View File
backend/utils/hf_progress.py
+1 -17
View File
@@ -32,7 +32,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:
@@ -81,7 +80,6 @@ class HFProgressTracker:
}
def update(self, n=1):
print(f"[DEBUG TrackedTqdm] update called with n={n}, filter_non_downloads={tracker.filter_non_downloads}")
result = super().update(n)
# Report progress
@@ -121,10 +119,7 @@ class HFProgressTracker:
def _is_download_progress(self, filename: str) -> bool:
"""Check if this is a real download progress bar vs internal processing."""
print(f"[DEBUG _is_download_progress] Checking filename: '{filename}'")
if not filename or filename == "unknown":
print(f"[DEBUG _is_download_progress] Rejected: empty/unknown filename")
return False
# Real downloads have file extensions
@@ -141,9 +136,7 @@ class HFProgressTracker:
skip_patterns = ['segment', 'processing', 'generating', 'loading']
has_skip_pattern = any(pattern in filename_lower for pattern in skip_patterns)
result = has_extension and not has_skip_pattern
print(f"[DEBUG _is_download_progress] has_extension={has_extension}, has_skip_pattern={has_skip_pattern}, result={result}")
return result
return has_extension and not has_skip_pattern
def close(self):
with tracker._lock:
@@ -156,13 +149,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:
@@ -175,22 +166,18 @@ 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
self._patched_modules = {}
patched_count = 0
for module_name in list(sys.modules.keys()):
if "huggingface" in module_name or module_name.startswith("tqdm"):
try:
@@ -203,11 +190,8 @@ class HFProgressTracker:
):
self._patched_modules[module_name] = attr
setattr(module, "tqdm", tracked_tqdm)
patched_count += 1
print(f"[DEBUG HFProgressTracker] Patched {module_name}.tqdm")
except (AttributeError, TypeError):
pass
print(f"[DEBUG HFProgressTracker] Patched {patched_count} modules in sys.modules")
yield
backend/utils/progress.py
+34 -10
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,6 +86,7 @@ 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
@@ -90,25 +101,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}%")
tauri/src-tauri/gen/Assets.car
BIN
View File
Binary file not shown.