diff --git a/backend/backends/mlx_backend.py b/backend/backends/mlx_backend.py index 2584df0b..c4ecc090 100644 --- a/backend/backends/mlx_backend.py +++ b/backend/backends/mlx_backend.py @@ -54,20 +54,43 @@ class MLXTTSBackend: def _is_model_cached(self, model_size: str) -> bool: """ - Check if the model is already cached locally. + Check if the model is already cached locally AND fully downloaded. Args: model_size: Model size to check Returns: - True if model is cached, False otherwise + True if model is fully cached, False if missing or incomplete """ try: from huggingface_hub import constants as hf_constants model_path = self._get_model_path(model_size) repo_cache = Path(hf_constants.HF_HUB_CACHE) / ("models--" + model_path.replace("/", "--")) - return repo_cache.exists() - except Exception: + + if not repo_cache.exists(): + return False + + # Check for .incomplete files - if any exist, download is still in progress + blobs_dir = repo_cache / "blobs" + if blobs_dir.exists() and any(blobs_dir.glob("*.incomplete")): + print(f"[_is_model_cached] Found .incomplete files for {model_size}, treating as not cached") + return False + + # Check that actual model weight files exist in snapshots + snapshots_dir = repo_cache / "snapshots" + if snapshots_dir.exists(): + has_weights = ( + any(snapshots_dir.rglob("*.safetensors")) or + any(snapshots_dir.rglob("*.bin")) or + any(snapshots_dir.rglob("*.npz")) + ) + if not has_weights: + print(f"[_is_model_cached] No model weights found for {model_size}, treating as not cached") + return False + + return True + except Exception as e: + print(f"[_is_model_cached] Error checking cache for {model_size}: {e}") return False async def load_model_async(self, model_size: Optional[str] = None): @@ -121,9 +144,15 @@ class MLXTTSBackend: # Start tracking download task task_manager.start_download(model_name) - # 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 + # Initialize progress state so SSE endpoint has initial data to send + # This provides immediate feedback while HuggingFace fetches metadata + progress_manager.update_progress( + model_name=model_name, + current=0, + total=0, # Will be updated once actual total is known + filename="Connecting to HuggingFace...", + status="downloading", + ) # IMPORTANT: Patch tqdm BEFORE importing mlx_audio # Otherwise mlx_audio caches reference to original tqdm @@ -363,20 +392,43 @@ class MLXSTTBackend: def _is_model_cached(self, model_size: str) -> bool: """ - Check if the Whisper model is already cached locally. + Check if the Whisper model is already cached locally AND fully downloaded. Args: model_size: Model size to check Returns: - True if model is cached, False otherwise + True if model is fully cached, False if missing or incomplete """ try: from huggingface_hub import constants as hf_constants model_name = f"openai/whisper-{model_size}" repo_cache = Path(hf_constants.HF_HUB_CACHE) / ("models--" + model_name.replace("/", "--")) - return repo_cache.exists() - except Exception: + + if not repo_cache.exists(): + return False + + # Check for .incomplete files - if any exist, download is still in progress + blobs_dir = repo_cache / "blobs" + if blobs_dir.exists() and any(blobs_dir.glob("*.incomplete")): + print(f"[_is_model_cached] Found .incomplete files for whisper-{model_size}, treating as not cached") + return False + + # Check that actual model weight files exist in snapshots + snapshots_dir = repo_cache / "snapshots" + if snapshots_dir.exists(): + has_weights = ( + any(snapshots_dir.rglob("*.safetensors")) or + any(snapshots_dir.rglob("*.bin")) or + any(snapshots_dir.rglob("*.npz")) + ) + if not has_weights: + print(f"[_is_model_cached] No model weights found for whisper-{model_size}, treating as not cached") + return False + + return True + except Exception as e: + print(f"[_is_model_cached] Error checking cache for whisper-{model_size}: {e}") return False async def load_model_async(self, model_size: Optional[str] = None): @@ -431,8 +483,14 @@ class MLXSTTBackend: # Start tracking download task task_manager.start_download(progress_model_name) - # Note: Don't initialize progress here with fake total=1 - # Let the actual tqdm callbacks set the real values + # Initialize progress state so SSE endpoint has initial data to send + progress_manager.update_progress( + model_name=progress_model_name, + current=0, + total=0, + filename="Connecting to HuggingFace...", + status="downloading", + ) # Load the model (tqdm is patched, but filters out non-download progress) try: diff --git a/backend/backends/pytorch_backend.py b/backend/backends/pytorch_backend.py index 7f8c2379..26f38726 100644 --- a/backend/backends/pytorch_backend.py +++ b/backend/backends/pytorch_backend.py @@ -60,20 +60,42 @@ class PyTorchTTSBackend: def _is_model_cached(self, model_size: str) -> bool: """ - Check if the model is already cached locally. + Check if the model is already cached locally AND fully downloaded. Args: model_size: Model size to check Returns: - True if model is cached, False otherwise + True if model is fully cached, False if missing or incomplete """ try: from huggingface_hub import constants as hf_constants model_path = self._get_model_path(model_size) repo_cache = Path(hf_constants.HF_HUB_CACHE) / ("models--" + model_path.replace("/", "--")) - return repo_cache.exists() - except Exception: + + if not repo_cache.exists(): + return False + + # Check for .incomplete files - if any exist, download is still in progress + blobs_dir = repo_cache / "blobs" + if blobs_dir.exists() and any(blobs_dir.glob("*.incomplete")): + print(f"[_is_model_cached] Found .incomplete files for {model_size}, treating as not cached") + return False + + # Check that actual model weight files exist in snapshots + snapshots_dir = repo_cache / "snapshots" + if snapshots_dir.exists(): + has_weights = ( + any(snapshots_dir.rglob("*.safetensors")) or + any(snapshots_dir.rglob("*.bin")) + ) + if not has_weights: + print(f"[_is_model_cached] No model weights found for {model_size}, treating as not cached") + return False + + return True + except Exception as e: + print(f"[_is_model_cached] Error checking cache for {model_size}: {e}") return False async def load_model_async(self, model_size: Optional[str] = None): @@ -133,12 +155,12 @@ class PyTorchTTSBackend: # Start tracking download task task_manager.start_download(model_name) - # Initialize progress state to show download has started + # Initialize progress state so SSE endpoint has initial data to send progress_manager.update_progress( model_name=model_name, current=0, - total=1, # Set to 1 initially, will be updated by callback - filename="", + total=0, # Will be updated once actual total is known + filename="Connecting to HuggingFace...", status="downloading", ) @@ -347,20 +369,42 @@ class PyTorchSTTBackend: def _is_model_cached(self, model_size: str) -> bool: """ - Check if the Whisper model is already cached locally. + Check if the Whisper model is already cached locally AND fully downloaded. Args: model_size: Model size to check Returns: - True if model is cached, False otherwise + True if model is fully cached, False if missing or incomplete """ try: from huggingface_hub import constants as hf_constants model_name = f"openai/whisper-{model_size}" repo_cache = Path(hf_constants.HF_HUB_CACHE) / ("models--" + model_name.replace("/", "--")) - return repo_cache.exists() - except Exception: + + if not repo_cache.exists(): + return False + + # Check for .incomplete files - if any exist, download is still in progress + blobs_dir = repo_cache / "blobs" + if blobs_dir.exists() and any(blobs_dir.glob("*.incomplete")): + print(f"[_is_model_cached] Found .incomplete files for whisper-{model_size}, treating as not cached") + return False + + # Check that actual model weight files exist in snapshots + snapshots_dir = repo_cache / "snapshots" + if snapshots_dir.exists(): + has_weights = ( + any(snapshots_dir.rglob("*.safetensors")) or + any(snapshots_dir.rglob("*.bin")) + ) + if not has_weights: + print(f"[_is_model_cached] No model weights found for whisper-{model_size}, treating as not cached") + return False + + return True + except Exception as e: + print(f"[_is_model_cached] Error checking cache for whisper-{model_size}: {e}") return False async def load_model_async(self, model_size: Optional[str] = None): @@ -422,18 +466,15 @@ class PyTorchSTTBackend: if not is_cached: # Start tracking download task task_manager.start_download(progress_model_name) - print(f"[DEBUG] Task manager started download") - # Initialize progress state to show download has started - print(f"[DEBUG] Calling update_progress...") + # Initialize progress state so SSE endpoint has initial data to send progress_manager.update_progress( model_name=progress_model_name, current=0, - total=1, # Set to 1 initially, will be updated by callback - filename="", + total=0, # Will be updated once actual total is known + filename="Connecting to HuggingFace...", status="downloading", ) - print(f"[DEBUG] update_progress called, listeners: {len(progress_manager._listeners.get(progress_model_name, []))}") # Load models (tqdm is patched, but filters out non-download progress) try: diff --git a/backend/main.py b/backend/main.py index 40c39a33..59fb9e18 100644 --- a/backend/main.py +++ b/backend/main.py @@ -1156,15 +1156,14 @@ async def get_model_progress(model_name: str): @app.get("/models/status", response_model=models.ModelStatusListResponse) async def get_model_status(): """Get status of all available models.""" - from huggingface_hub import hf_hub_download, constants as hf_constants + from huggingface_hub import constants as hf_constants from pathlib import Path - import os backend_type = get_backend_type() task_manager = get_task_manager() - # Get set of currently downloading models - active_downloads = {task.model_name for task in task_manager.get_active_downloads()} + # Get set of currently downloading model names + active_download_names = {task.model_name for task in task_manager.get_active_downloads()} # Try to import scan_cache_dir (might not be available in older versions) try: @@ -1251,6 +1250,13 @@ async def get_model_status(): }, ] + # Build a mapping of model_name -> hf_repo_id so we can check if shared repos are downloading + model_to_repo = {cfg["model_name"]: cfg["hf_repo_id"] for cfg in model_configs} + + # Get the set of hf_repo_ids that are currently being downloaded + # This handles the case where multiple models share the same repo (e.g., 0.6B and 1.7B on MLX) + active_download_repos = {model_to_repo.get(name) for name in active_download_names if name in model_to_repo} + # Get HuggingFace cache info (if available) cache_info = None if use_scan_cache: @@ -1273,13 +1279,37 @@ async def get_model_status(): repo_id = config["hf_repo_id"] for repo in cache_info.repos: if repo.repo_id == repo_id: - downloaded = True - # Calculate size from cache info + # Check if actual model weight files exist (not just config files) + # scan_cache_dir only shows completed files, so check if any are model weights + has_model_weights = False + for rev in repo.revisions: + for f in rev.files: + fname = f.file_name.lower() + if fname.endswith(('.safetensors', '.bin', '.pt', '.pth', '.npz')): + has_model_weights = True + break + if has_model_weights: + break + + # Also check for .incomplete files in blobs directory (downloads in progress) + has_incomplete = False try: - total_size = sum(revision.size_on_disk for revision in repo.revisions) - size_mb = total_size / (1024 * 1024) + cache_dir = hf_constants.HF_HUB_CACHE + blobs_dir = Path(cache_dir) / ("models--" + repo_id.replace("/", "--")) / "blobs" + if blobs_dir.exists(): + has_incomplete = any(blobs_dir.glob("*.incomplete")) except Exception: pass + + # Only mark as downloaded if we have model weights AND no incomplete files + if has_model_weights and not has_incomplete: + downloaded = True + # Calculate size from cache info + try: + total_size = sum(revision.size_on_disk for revision in repo.revisions) + size_mb = total_size / (1024 * 1024) + except Exception: + pass break # Method 2: Fallback to checking cache directory directly (using HuggingFace's OS-specific cache location) @@ -1289,42 +1319,40 @@ async def get_model_status(): repo_cache = Path(cache_dir) / ("models--" + config["hf_repo_id"].replace("/", "--")) if repo_cache.exists(): - # Check for model files (bin, safetensors, or other common model files) - # MLX models may use .npz or .safetensors - has_model_files = ( - any(repo_cache.rglob("*.bin")) or - any(repo_cache.rglob("*.safetensors")) or - any(repo_cache.rglob("*.pt")) or - any(repo_cache.rglob("*.pth")) or - any(repo_cache.rglob("*.npz")) or - any(repo_cache.rglob("model.safetensors.index.json")) or - any(repo_cache.rglob("pytorch_model.bin.index.json")) - ) + # Check for .incomplete files - if any exist, download is still in progress + blobs_dir = repo_cache / "blobs" + has_incomplete = blobs_dir.exists() and any(blobs_dir.glob("*.incomplete")) - if has_model_files: - downloaded = True - # Calculate size - try: - total_size = sum(f.stat().st_size for f in repo_cache.rglob("*") if f.is_file()) - size_mb = total_size / (1024 * 1024) - except Exception: - pass + if not has_incomplete: + # Check for actual model weight files (not just index files) + # in the snapshots directory (symlinks to completed blobs) + snapshots_dir = repo_cache / "snapshots" + has_model_files = False + if snapshots_dir.exists(): + has_model_files = ( + any(snapshots_dir.rglob("*.bin")) or + any(snapshots_dir.rglob("*.safetensors")) or + any(snapshots_dir.rglob("*.pt")) or + any(snapshots_dir.rglob("*.pth")) or + any(snapshots_dir.rglob("*.npz")) + ) + + if has_model_files: + downloaded = True + # Calculate size (exclude .incomplete files) + try: + total_size = sum( + f.stat().st_size for f in repo_cache.rglob("*") + if f.is_file() and not f.name.endswith('.incomplete') + ) + size_mb = total_size / (1024 * 1024) + except Exception: + pass except Exception: pass - # Method 3: Try to check if model can be loaded locally (last resort) - if not downloaded: - try: - # Try to download with local_files_only=True to check if cached - hf_hub_download( - repo_id=config["hf_repo_id"], - filename="config.json", # Try a common file - local_files_only=True, - ) - downloaded = True - except Exception: - # File not found locally, model not downloaded - pass + # Method 3 removed - checking for config.json is too lenient + # Methods 1 and 2 properly verify that model weight files exist # Check if loaded in memory try: @@ -1332,12 +1360,13 @@ async def get_model_status(): except Exception: loaded = False - # Check if this model is currently being downloaded - is_downloading = config["model_name"] in active_downloads + # Check if this model (or its shared repo) is currently being downloaded + is_downloading = config["hf_repo_id"] in active_download_repos # If downloading, don't report as downloaded (partial files exist) if is_downloading: downloaded = False + size_mb = None # Don't show partial size during download statuses.append(models.ModelStatus( model_name=config["model_name"], @@ -1354,8 +1383,8 @@ async def get_model_status(): except Exception: loaded = False - # Check if this model is currently being downloaded - is_downloading = config["model_name"] in active_downloads + # Check if this model (or its shared repo) is currently being downloaded + is_downloading = config["hf_repo_id"] in active_download_repos statuses.append(models.ModelStatus( model_name=config["model_name"], @@ -1375,6 +1404,7 @@ async def trigger_model_download(request: models.ModelDownloadRequest): import asyncio task_manager = get_task_manager() + progress_manager = get_progress_manager() model_configs = { "qwen-tts-1.7B": { @@ -1422,6 +1452,18 @@ async def trigger_model_download(request: models.ModelDownloadRequest): # Start tracking download task_manager.start_download(request.model_name) + + # Initialize progress state so SSE endpoint has initial data to send. + # This fixes a race condition where the frontend connects to SSE before + # any progress callbacks have fired (especially for large models like Qwen + # where huggingface_hub takes time to fetch metadata for all files). + progress_manager.update_progress( + model_name=request.model_name, + current=0, + total=0, # Will be updated once actual total is known + filename="Connecting to HuggingFace...", + status="downloading", + ) # Start download in background task (don't await) asyncio.create_task(download_in_background()) diff --git a/backend/utils/hf_progress.py b/backend/utils/hf_progress.py index 2c88abe7..7fc88edd 100644 --- a/backend/utils/hf_progress.py +++ b/backend/utils/hf_progress.py @@ -22,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.""" @@ -90,35 +91,66 @@ class HFProgressTracker: total = getattr(self, "total", 0) if total and total > 0: - # Determine if we should report this progress - should_report = True + # 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: - # Filter out non-download progress bars (e.g., "Segment 1/1" during generation) - # Only report progress for actual file downloads from HuggingFace - should_report = self._is_download_progress(filename) + if not self._is_download_progress(filename): + return result - if should_report: - # Update per-file tracking - tracker._file_sizes[filename] = total - tracker._file_downloaded[filename] = current - - # Calculate totals across all files - tracker._total_size = sum(tracker._file_sizes.values()) - tracker._total_downloaded = sum(tracker._file_downloaded.values()) - - # Call progress callback - if tracker.progress_callback: - tracker.progress_callback( - tracker._total_downloaded, - tracker._total_size, - filename - ) + # Update per-file tracking + tracker._file_sizes[filename] = total + tracker._file_downloaded[filename] = current + + # Calculate totals across all files + 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( + tracker._total_downloaded, + tracker._total_size, + filename + ) 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 download progress bar vs internal processing.""" + """Check if this is a real file download progress bar vs internal processing.""" if not filename or filename == "unknown": return False @@ -132,7 +164,7 @@ class HFProgressTracker: filename_lower = filename.lower() has_extension = any(filename_lower.endswith(ext) for ext in download_extensions) - # Skip internal progress indicators + # Skip generation-related progress indicators skip_patterns = ['segment', 'processing', 'generating', 'loading'] has_skip_pattern = any(pattern in filename_lower for pattern in skip_patterns) @@ -177,22 +209,79 @@ class HFProgressTracker: tqdm_module.auto.tqdm = tracked_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) + 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 + # 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 except ImportError: @@ -209,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 @@ -225,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 diff --git a/backend/utils/progress.py b/backend/utils/progress.py index 7edf9eb9..418a88c7 100644 --- a/backend/utils/progress.py +++ b/backend/utils/progress.py @@ -89,7 +89,14 @@ class ProgressManager: 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,