mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-20 07:10:40 -07:00
Add progress tracking and caching checks for model downloads
- Introduced methods to check if models are cached locally in MLX and PyTorch backends. - Enhanced progress tracking during model loading to filter out non-download progress when models are cached. - Updated HFProgressTracker to conditionally report progress based on download status. - Added test scripts for monitoring SSE events during model downloads and verifying progress tracking functionality. - Improved overall error handling and logging for better debugging during model download processes.
This commit is contained in:
Jamie Pine
@@ -58,6 +58,24 @@ class PyTorchTTSBackend:
|
||||
|
||||
return hf_model_map[model_size]
|
||||
|
||||
def _is_model_cached(self, model_size: str) -> bool:
|
||||
"""
|
||||
Check if the model is already cached locally.
|
||||
|
||||
Args:
|
||||
model_size: Model size to check
|
||||
|
||||
Returns:
|
||||
True if model is cached, False otherwise
|
||||
"""
|
||||
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:
|
||||
return False
|
||||
|
||||
async def load_model_async(self, model_size: Optional[str] = None):
|
||||
"""
|
||||
Lazy load the TTS model with automatic downloading from HuggingFace Hub.
|
||||
@@ -85,20 +103,24 @@ class PyTorchTTSBackend:
|
||||
def _load_model_sync(self, model_size: str):
|
||||
"""Synchronous model loading."""
|
||||
try:
|
||||
# IMPORTANT: Set up progress tracking BEFORE importing qwen_tts
|
||||
# This ensures tqdm is patched before any HuggingFace Hub imports
|
||||
progress_manager = get_progress_manager()
|
||||
task_manager = get_task_manager()
|
||||
model_name = f"qwen-tts-{model_size}"
|
||||
|
||||
# Check if model is already cached
|
||||
is_cached = self._is_model_cached(model_size)
|
||||
|
||||
# Set up progress callback and tracker
|
||||
# If cached: filter out non-download progress (like "Segment 1/1" during generation)
|
||||
# If not cached: report all progress (we're actually downloading)
|
||||
progress_callback = create_hf_progress_callback(model_name, progress_manager)
|
||||
tracker = HFProgressTracker(progress_callback)
|
||||
tracker = HFProgressTracker(progress_callback, filter_non_downloads=is_cached)
|
||||
|
||||
# Patch tqdm BEFORE importing qwen_tts
|
||||
tracker_context = tracker.patch_download()
|
||||
tracker_context.__enter__()
|
||||
|
||||
# NOW import qwen_tts - it will use our patched tqdm
|
||||
# Import qwen_tts
|
||||
from qwen_tts import Qwen3TTSModel
|
||||
|
||||
# Get model path (local or HuggingFace Hub ID)
|
||||
@@ -106,20 +128,21 @@ class PyTorchTTSBackend:
|
||||
|
||||
print(f"Loading TTS model {model_size} on {self.device}...")
|
||||
|
||||
# Start tracking download task
|
||||
task_manager = get_task_manager()
|
||||
task_manager.start_download(model_name)
|
||||
# Only track download progress if model is NOT cached
|
||||
if not is_cached:
|
||||
# Start tracking download task
|
||||
task_manager.start_download(model_name)
|
||||
|
||||
# Initialize progress state to show download has started
|
||||
progress_manager.update_progress(
|
||||
model_name=model_name,
|
||||
current=0,
|
||||
total=1, # Set to 1 initially, will be updated by callback
|
||||
filename="",
|
||||
status="downloading",
|
||||
)
|
||||
# Initialize progress state to show download has started
|
||||
progress_manager.update_progress(
|
||||
model_name=model_name,
|
||||
current=0,
|
||||
total=1, # Set to 1 initially, will be updated by callback
|
||||
filename="",
|
||||
status="downloading",
|
||||
)
|
||||
|
||||
# Load the model (tqdm is already patched from above)
|
||||
# Load the model (tqdm is patched, but filters out non-download progress)
|
||||
try:
|
||||
self.model = Qwen3TTSModel.from_pretrained(
|
||||
model_path,
|
||||
@@ -130,9 +153,10 @@ class PyTorchTTSBackend:
|
||||
# Exit the patch context
|
||||
tracker_context.__exit__(None, None, None)
|
||||
|
||||
# Mark as complete
|
||||
progress_manager.mark_complete(model_name)
|
||||
task_manager.complete_download(model_name)
|
||||
# Only mark download as complete if we were tracking it
|
||||
if not is_cached:
|
||||
progress_manager.mark_complete(model_name)
|
||||
task_manager.complete_download(model_name)
|
||||
|
||||
self._current_model_size = model_size
|
||||
self.model_size = model_size
|
||||
@@ -321,6 +345,24 @@ class PyTorchSTTBackend:
|
||||
"""Check if model is loaded."""
|
||||
return self.model is not None
|
||||
|
||||
def _is_model_cached(self, model_size: str) -> bool:
|
||||
"""
|
||||
Check if the Whisper model is already cached locally.
|
||||
|
||||
Args:
|
||||
model_size: Model size to check
|
||||
|
||||
Returns:
|
||||
True if model is cached, False otherwise
|
||||
"""
|
||||
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:
|
||||
return False
|
||||
|
||||
async def load_model_async(self, model_size: Optional[str] = None):
|
||||
"""
|
||||
Lazy load the Whisper model.
|
||||
@@ -349,14 +391,18 @@ class PyTorchSTTBackend:
|
||||
"""Synchronous model loading."""
|
||||
print(f"[DEBUG] _load_model_sync called for Whisper {model_size}")
|
||||
try:
|
||||
# IMPORTANT: Set up progress tracking BEFORE importing transformers
|
||||
# This ensures tqdm is patched before any HuggingFace Hub imports
|
||||
progress_manager = get_progress_manager()
|
||||
task_manager = get_task_manager()
|
||||
progress_model_name = f"whisper-{model_size}"
|
||||
|
||||
# Check if model is already cached
|
||||
is_cached = self._is_model_cached(model_size)
|
||||
|
||||
# Set up progress callback and tracker
|
||||
# If cached: filter out non-download progress
|
||||
# If not cached: report all progress (we're actually downloading)
|
||||
progress_callback = create_hf_progress_callback(progress_model_name, progress_manager)
|
||||
tracker = HFProgressTracker(progress_callback)
|
||||
tracker = HFProgressTracker(progress_callback, filter_non_downloads=is_cached)
|
||||
|
||||
# Patch tqdm BEFORE importing transformers
|
||||
print("[DEBUG] Starting tqdm patch BEFORE transformers import")
|
||||
@@ -364,31 +410,32 @@ class PyTorchSTTBackend:
|
||||
tracker_context.__enter__()
|
||||
print("[DEBUG] tqdm patched, now importing transformers")
|
||||
|
||||
# NOW import transformers - it will use our patched tqdm
|
||||
# Import transformers
|
||||
from transformers import WhisperProcessor, WhisperForConditionalGeneration
|
||||
|
||||
model_name = f"openai/whisper-{model_size}"
|
||||
print(f"[DEBUG] Model name: {model_name}")
|
||||
|
||||
# Start tracking download task
|
||||
task_manager = get_task_manager()
|
||||
task_manager.start_download(progress_model_name)
|
||||
print(f"[DEBUG] Task manager started download")
|
||||
|
||||
print(f"Loading Whisper model {model_size} on {self.device}...")
|
||||
|
||||
# Initialize progress state to show download has started
|
||||
print(f"[DEBUG] Calling update_progress...")
|
||||
progress_manager.update_progress(
|
||||
model_name=progress_model_name,
|
||||
current=0,
|
||||
total=1, # Set to 1 initially, will be updated by callback
|
||||
filename="",
|
||||
status="downloading",
|
||||
)
|
||||
print(f"[DEBUG] update_progress called, listeners: {len(progress_manager._listeners.get(progress_model_name, []))}")
|
||||
# Only track download progress if model is NOT cached
|
||||
if not is_cached:
|
||||
# Start tracking download task
|
||||
task_manager.start_download(progress_model_name)
|
||||
print(f"[DEBUG] Task manager started download")
|
||||
|
||||
# Load models (tqdm is already patched from above)
|
||||
# Initialize progress state to show download has started
|
||||
print(f"[DEBUG] Calling update_progress...")
|
||||
progress_manager.update_progress(
|
||||
model_name=progress_model_name,
|
||||
current=0,
|
||||
total=1, # Set to 1 initially, will be updated by callback
|
||||
filename="",
|
||||
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:
|
||||
self.processor = WhisperProcessor.from_pretrained(model_name)
|
||||
self.model = WhisperForConditionalGeneration.from_pretrained(model_name)
|
||||
@@ -396,13 +443,14 @@ class PyTorchSTTBackend:
|
||||
# Exit the patch context
|
||||
tracker_context.__exit__(None, None, None)
|
||||
|
||||
# Only mark download as complete if we were tracking it
|
||||
if not is_cached:
|
||||
progress_manager.mark_complete(progress_model_name)
|
||||
task_manager.complete_download(progress_model_name)
|
||||
|
||||
self.model.to(self.device)
|
||||
self.model_size = model_size
|
||||
|
||||
# Mark as complete
|
||||
progress_manager.mark_complete(progress_model_name)
|
||||
task_manager.complete_download(progress_model_name)
|
||||
|
||||
print(f"Whisper model {model_size} loaded successfully")
|
||||
|
||||
except Exception as e:
|
||||
|
||||
Reference in New Issue
Block a user