diff --git a/app/src/lib/hooks/useModelDownloadToast.tsx b/app/src/lib/hooks/useModelDownloadToast.tsx index b78fdc9c..2df221e1 100644 --- a/app/src/lib/hooks/useModelDownloadToast.tsx +++ b/app/src/lib/hooks/useModelDownloadToast.tsx @@ -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: ( +
+ + {displayName} +
+ ), + 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, }; -} \ No newline at end of file +} diff --git a/backend/backends/mlx_backend.py b/backend/backends/mlx_backend.py index 8d84776c..2584df0b 100644 --- a/backend/backends/mlx_backend.py +++ b/backend/backends/mlx_backend.py @@ -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: diff --git a/backend/tests/test_check_progress_state.py b/backend/tests/test_check_progress_state.py deleted file mode 100644 index 6737547f..00000000 --- a/backend/tests/test_check_progress_state.py +++ /dev/null @@ -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() diff --git a/backend/tests/test_generation_progress.py b/backend/tests/test_generation_download.py similarity index 100% rename from backend/tests/test_generation_progress.py rename to backend/tests/test_generation_download.py diff --git a/backend/tests/test_qwen_download.py b/backend/tests/test_qwen_download.py new file mode 100644 index 00000000..b6c42e17 --- /dev/null +++ b/backend/tests/test_qwen_download.py @@ -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) diff --git a/backend/tests/test_real_download.py b/backend/tests/test_whisper_download.py similarity index 100% rename from backend/tests/test_real_download.py rename to backend/tests/test_whisper_download.py diff --git a/backend/utils/hf_progress.py b/backend/utils/hf_progress.py index 014645d5..2c88abe7 100644 --- a/backend/utils/hf_progress.py +++ b/backend/utils/hf_progress.py @@ -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 diff --git a/backend/utils/progress.py b/backend/utils/progress.py index 6535c958..7edf9eb9 100644 --- a/backend/utils/progress.py +++ b/backend/utils/progress.py @@ -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}%") diff --git a/tauri/src-tauri/gen/Assets.car b/tauri/src-tauri/gen/Assets.car index add24cf4..de0e9a08 100644 Binary files a/tauri/src-tauri/gen/Assets.car and b/tauri/src-tauri/gen/Assets.car differ