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

Enhance model caching checks and progress tracking for downloads

Browse Source 
- Updated caching methods in MLX, PyTorch, and backend to ensure models are fully downloaded before being marked as cached.
- Improved progress tracking to filter out non-download progress and provide accurate feedback during model downloads.
- Enhanced HFProgressTracker to skip non-byte progress bars and ensure meaningful progress reporting.
- Refactored progress initialization to provide immediate feedback while fetching metadata from HuggingFace.
- Added error handling and logging for better debugging during cache checks and download processes.
This commit is contained in:
Jamie Pine
2026-01-30 21:17:01 -08:00
parent d3393fb940
commit 60a03c56a9
5 changed files with 368 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/backends/mlx_backend.py
+71 -13
View File
@@ -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:
backend/backends/pytorch_backend.py
+58 -17
View File
@@ -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: