mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-18 22:30:40 -07:00
fix: download progress tracking for all engines and inline progress UI
- Add HFProgressTracker to LuxTTS and Chatterbox backends so tqdm-based file-level download progress reaches the frontend (previously only Qwen had this, LuxTTS/Chatterbox showed a static spinner) - Add progress/current/total/filename fields to ActiveDownloadTask so the /tasks/active polling endpoint carries progress data - Show inline progress bar + bytes in the model list and detail modal, poll at 1s during active downloads (5s otherwise) - Fix GpuAcceleration crash: cudaStatusLoading was referenced before initialization in its own useQuery declaration
This commit is contained in:
@@ -103,19 +103,27 @@ class ChatterboxTTSBackend:
|
||||
|
||||
def _load_model_sync(self):
|
||||
"""Synchronous model loading."""
|
||||
from ..utils.hf_progress import HFProgressTracker, create_hf_progress_callback
|
||||
|
||||
progress_manager = get_progress_manager()
|
||||
task_manager = get_task_manager()
|
||||
model_name = "chatterbox-tts"
|
||||
|
||||
is_cached = self._is_model_cached()
|
||||
|
||||
# Set up HF progress tracking (intercepts tqdm for file-level progress)
|
||||
progress_callback = create_hf_progress_callback(model_name, progress_manager)
|
||||
tracker = HFProgressTracker(progress_callback, filter_non_downloads=is_cached)
|
||||
tracker_context = tracker.patch_download()
|
||||
tracker_context.__enter__()
|
||||
|
||||
if not is_cached:
|
||||
task_manager.start_download(model_name)
|
||||
progress_manager.update_progress(
|
||||
model_name=model_name,
|
||||
current=0,
|
||||
total=0,
|
||||
filename="Downloading Chatterbox model...",
|
||||
filename="Connecting to HuggingFace...",
|
||||
status="downloading",
|
||||
)
|
||||
|
||||
@@ -131,25 +139,28 @@ class ChatterboxTTSBackend:
|
||||
# Monkey-patch torch.load for CPU loading. The model's .pt files
|
||||
# were saved on CUDA; from_pretrained() doesn't pass map_location
|
||||
# so loading on CPU fails without this.
|
||||
if device == "cpu":
|
||||
_orig_torch_load = torch.load
|
||||
try:
|
||||
if device == "cpu":
|
||||
_orig_torch_load = torch.load
|
||||
|
||||
def _patched_load(*args, **kwargs):
|
||||
kwargs.setdefault("map_location", "cpu")
|
||||
return _orig_torch_load(*args, **kwargs)
|
||||
def _patched_load(*args, **kwargs):
|
||||
kwargs.setdefault("map_location", "cpu")
|
||||
return _orig_torch_load(*args, **kwargs)
|
||||
|
||||
with ChatterboxTTSBackend._load_lock:
|
||||
torch.load = _patched_load
|
||||
try:
|
||||
self.model = ChatterboxMultilingualTTS.from_pretrained(
|
||||
device=device,
|
||||
)
|
||||
finally:
|
||||
torch.load = _orig_torch_load
|
||||
else:
|
||||
self.model = ChatterboxMultilingualTTS.from_pretrained(
|
||||
device=device,
|
||||
)
|
||||
with ChatterboxTTSBackend._load_lock:
|
||||
torch.load = _patched_load
|
||||
try:
|
||||
self.model = ChatterboxMultilingualTTS.from_pretrained(
|
||||
device=device,
|
||||
)
|
||||
finally:
|
||||
torch.load = _orig_torch_load
|
||||
else:
|
||||
self.model = ChatterboxMultilingualTTS.from_pretrained(
|
||||
device=device,
|
||||
)
|
||||
finally:
|
||||
tracker_context.__exit__(None, None, None)
|
||||
|
||||
# Fix: transformers >= 4.36 defaults LlamaModel to sdpa attention
|
||||
# which doesn't support output_attentions=True (needed by
|
||||
|
||||
@@ -94,19 +94,27 @@ class LuxTTSBackend:
|
||||
|
||||
def _load_model_sync(self):
|
||||
"""Synchronous model loading."""
|
||||
from ..utils.hf_progress import HFProgressTracker, create_hf_progress_callback
|
||||
|
||||
progress_manager = get_progress_manager()
|
||||
task_manager = get_task_manager()
|
||||
model_name = "luxtts"
|
||||
|
||||
is_cached = self._is_model_cached()
|
||||
|
||||
# Set up HF progress tracking (intercepts tqdm for file-level progress)
|
||||
progress_callback = create_hf_progress_callback(model_name, progress_manager)
|
||||
tracker = HFProgressTracker(progress_callback, filter_non_downloads=is_cached)
|
||||
tracker_context = tracker.patch_download()
|
||||
tracker_context.__enter__()
|
||||
|
||||
if not is_cached:
|
||||
task_manager.start_download(model_name)
|
||||
progress_manager.update_progress(
|
||||
model_name=model_name,
|
||||
current=0,
|
||||
total=0,
|
||||
filename="Downloading LuxTTS model...",
|
||||
filename="Connecting to HuggingFace...",
|
||||
status="downloading",
|
||||
)
|
||||
|
||||
@@ -117,19 +125,22 @@ class LuxTTSBackend:
|
||||
logger.info(f"Loading LuxTTS on {device}...")
|
||||
|
||||
# LuxTTS constructor downloads model and loads everything
|
||||
if device == "cpu":
|
||||
import os
|
||||
threads = os.cpu_count() or 4
|
||||
self.model = LuxTTS(
|
||||
model_path=LUXTTS_HF_REPO,
|
||||
device="cpu",
|
||||
threads=min(threads, 8),
|
||||
)
|
||||
else:
|
||||
self.model = LuxTTS(
|
||||
model_path=LUXTTS_HF_REPO,
|
||||
device=device,
|
||||
)
|
||||
try:
|
||||
if device == "cpu":
|
||||
import os
|
||||
threads = os.cpu_count() or 4
|
||||
self.model = LuxTTS(
|
||||
model_path=LUXTTS_HF_REPO,
|
||||
device="cpu",
|
||||
threads=min(threads, 8),
|
||||
)
|
||||
else:
|
||||
self.model = LuxTTS(
|
||||
model_path=LUXTTS_HF_REPO,
|
||||
device=device,
|
||||
)
|
||||
finally:
|
||||
tracker_context.__exit__(None, None, None)
|
||||
|
||||
if not is_cached:
|
||||
progress_manager.mark_complete(model_name)
|
||||
|
||||
@@ -1912,11 +1912,22 @@ async def get_active_tasks():
|
||||
pm_data = progress_manager._progress.get(model_name)
|
||||
if pm_data:
|
||||
error = pm_data.get("error")
|
||||
# Include progress data if available
|
||||
prog = progress or {}
|
||||
if not prog:
|
||||
with progress_manager._lock:
|
||||
pm_data = progress_manager._progress.get(model_name)
|
||||
if pm_data:
|
||||
prog = pm_data
|
||||
active_downloads.append(models.ActiveDownloadTask(
|
||||
model_name=model_name,
|
||||
status=task.status,
|
||||
started_at=task.started_at,
|
||||
error=error,
|
||||
progress=prog.get("progress"),
|
||||
current=prog.get("current"),
|
||||
total=prog.get("total"),
|
||||
filename=prog.get("filename"),
|
||||
))
|
||||
elif progress:
|
||||
# Progress exists but no task - create from progress data
|
||||
@@ -1934,6 +1945,10 @@ async def get_active_tasks():
|
||||
status=progress.get("status", "downloading"),
|
||||
started_at=started_at,
|
||||
error=progress.get("error"),
|
||||
progress=progress.get("progress"),
|
||||
current=progress.get("current"),
|
||||
total=progress.get("total"),
|
||||
filename=progress.get("filename"),
|
||||
))
|
||||
|
||||
# Get active generations
|
||||
|
||||
@@ -158,6 +158,10 @@ class ActiveDownloadTask(BaseModel):
|
||||
status: str
|
||||
started_at: datetime
|
||||
error: Optional[str] = None
|
||||
progress: Optional[float] = None # 0-100 percentage
|
||||
current: Optional[int] = None # bytes downloaded
|
||||
total: Optional[int] = None # total bytes
|
||||
filename: Optional[str] = None # current file being downloaded
|
||||
|
||||
|
||||
class ActiveGenerationTask(BaseModel):
|
||||
|
||||
Reference in New Issue
Block a user