mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-19 23:00:45 -07:00
refactor: remove dead code, deduplicate backends
Phase 1 - delete dead code: - studio.py, migrate_add_instruct.py, utils/validation.py - duplicate _profile_to_response in main.py, duplicate asyncio import - pointless _get_profiles_dir/_get_generations_dir wrappers - duplicate LANGUAGE_CODE_TO_NAME and WHISPER_HF_REPOS constants Phase 2 - extract backends/base.py with shared utilities: - is_model_cached() replaces 7 copy-pasted HF cache checks - get_torch_device() replaces 5 device detection methods - combine_voice_prompts() replaces 5 identical implementations - model_load_progress() ctx manager replaces progress boilerplate in all backends - patch_chatterbox_f32() replaces identical monkey-patches in both chatterbox backends net -1078 lines across the backend
This commit is contained in:
Jamie Pine
@@ -6,20 +6,11 @@ from typing import Optional, List, Tuple
|
||||
import asyncio
|
||||
import torch
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
|
||||
from . import TTSBackend, STTBackend
|
||||
from . import TTSBackend, STTBackend, LANGUAGE_CODE_TO_NAME, WHISPER_HF_REPOS
|
||||
from .base import is_model_cached, get_torch_device, combine_voice_prompts as _combine_voice_prompts, model_load_progress
|
||||
from ..utils.cache import get_cache_key, get_cached_voice_prompt, cache_voice_prompt
|
||||
from ..utils.audio import normalize_audio, load_audio
|
||||
from ..utils.progress import get_progress_manager
|
||||
from ..utils.hf_progress import HFProgressTracker, create_hf_progress_callback
|
||||
from ..utils.tasks import get_task_manager
|
||||
|
||||
LANGUAGE_CODE_TO_NAME = {
|
||||
"zh": "chinese", "en": "english", "ja": "japanese", "ko": "korean",
|
||||
"de": "german", "fr": "french", "ru": "russian", "pt": "portuguese",
|
||||
"es": "spanish", "it": "italian",
|
||||
}
|
||||
from ..utils.audio import load_audio
|
||||
|
||||
|
||||
class PyTorchTTSBackend:
|
||||
@@ -33,26 +24,7 @@ class PyTorchTTSBackend:
|
||||
|
||||
def _get_device(self) -> str:
|
||||
"""Get the best available device."""
|
||||
if torch.cuda.is_available():
|
||||
return "cuda"
|
||||
# Intel Arc / Intel Xe GPU via intel-extension-for-pytorch (IPEX)
|
||||
try:
|
||||
import intel_extension_for_pytorch # noqa: F401
|
||||
if hasattr(torch, 'xpu') and torch.xpu.is_available():
|
||||
return "xpu"
|
||||
except ImportError:
|
||||
pass
|
||||
# Any GPU on Windows via DirectML (torch-directml)
|
||||
try:
|
||||
import torch_directml
|
||||
if torch_directml.device_count() > 0:
|
||||
return torch_directml.device(0)
|
||||
except ImportError:
|
||||
pass
|
||||
# MPS (Apple Silicon) — kept for completeness but MLX backend is preferred
|
||||
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
||||
return "cpu" # MPS disabled for stability; MLX backend handles Apple Silicon
|
||||
return "cpu"
|
||||
return get_torch_device(allow_xpu=True, allow_directml=True)
|
||||
|
||||
def is_loaded(self) -> bool:
|
||||
"""Check if model is loaded."""
|
||||
@@ -79,44 +51,7 @@ 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 AND fully downloaded.
|
||||
|
||||
Args:
|
||||
model_size: Model size to check
|
||||
|
||||
Returns:
|
||||
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("/", "--"))
|
||||
|
||||
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
|
||||
return is_model_cached(self._get_model_path(model_size))
|
||||
|
||||
async def load_model_async(self, model_size: Optional[str] = None):
|
||||
"""
|
||||
@@ -144,94 +79,30 @@ class PyTorchTTSBackend:
|
||||
|
||||
def _load_model_sync(self, model_size: str):
|
||||
"""Synchronous model loading."""
|
||||
try:
|
||||
progress_manager = get_progress_manager()
|
||||
task_manager = get_task_manager()
|
||||
model_name = f"qwen-tts-{model_size}"
|
||||
model_name = f"qwen-tts-{model_size}"
|
||||
is_cached = self._is_model_cached(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, filter_non_downloads=is_cached)
|
||||
|
||||
# Patch tqdm BEFORE importing qwen_tts
|
||||
tracker_context = tracker.patch_download()
|
||||
tracker_context.__enter__()
|
||||
|
||||
# Import qwen_tts
|
||||
with model_load_progress(model_name, is_cached):
|
||||
from qwen_tts import Qwen3TTSModel
|
||||
|
||||
# Get model path (local or HuggingFace Hub ID)
|
||||
model_path = self._get_model_path(model_size)
|
||||
|
||||
print(f"Loading TTS model {model_size} on {self.device}...")
|
||||
|
||||
# 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 so SSE endpoint has initial data to send
|
||||
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",
|
||||
if self.device == "cpu":
|
||||
self.model = Qwen3TTSModel.from_pretrained(
|
||||
model_path,
|
||||
torch_dtype=torch.float32,
|
||||
low_cpu_mem_usage=False,
|
||||
)
|
||||
else:
|
||||
self.model = Qwen3TTSModel.from_pretrained(
|
||||
model_path,
|
||||
device_map=self.device,
|
||||
torch_dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
# Load the model (tqdm is patched, but filters out non-download progress)
|
||||
try:
|
||||
# Don't pass device_map on CPU: accelerate's meta-tensor mechanism
|
||||
# causes "Cannot copy out of meta tensor" when moving to CPU.
|
||||
# Instead load directly then call .to(device) if needed.
|
||||
if self.device == "cpu":
|
||||
self.model = Qwen3TTSModel.from_pretrained(
|
||||
model_path,
|
||||
torch_dtype=torch.float32,
|
||||
low_cpu_mem_usage=False,
|
||||
)
|
||||
else:
|
||||
self.model = Qwen3TTSModel.from_pretrained(
|
||||
model_path,
|
||||
device_map=self.device,
|
||||
torch_dtype=torch.bfloat16,
|
||||
)
|
||||
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:
|
||||
progress_manager.mark_complete(model_name)
|
||||
task_manager.complete_download(model_name)
|
||||
|
||||
self._current_model_size = model_size
|
||||
self.model_size = model_size
|
||||
|
||||
print(f"TTS model {model_size} loaded successfully")
|
||||
|
||||
except ImportError as e:
|
||||
print(f"Error: qwen_tts package not found. Install with: pip install git+https://github.com/QwenLM/Qwen3-TTS.git")
|
||||
progress_manager = get_progress_manager()
|
||||
task_manager = get_task_manager()
|
||||
model_name = f"qwen-tts-{model_size}"
|
||||
progress_manager.mark_error(model_name, str(e))
|
||||
task_manager.error_download(model_name, str(e))
|
||||
raise
|
||||
except Exception as e:
|
||||
print(f"Error loading TTS model: {e}")
|
||||
print(f"Tip: The model will be automatically downloaded from HuggingFace Hub on first use.")
|
||||
progress_manager = get_progress_manager()
|
||||
task_manager = get_task_manager()
|
||||
model_name = f"qwen-tts-{model_size}"
|
||||
progress_manager.mark_error(model_name, str(e))
|
||||
task_manager.error_download(model_name, str(e))
|
||||
raise
|
||||
self._current_model_size = model_size
|
||||
self.model_size = model_size
|
||||
print(f"TTS model {model_size} loaded successfully")
|
||||
|
||||
def unload_model(self):
|
||||
"""Unload the model to free memory."""
|
||||
@@ -303,31 +174,7 @@ class PyTorchTTSBackend:
|
||||
audio_paths: List[str],
|
||||
reference_texts: List[str],
|
||||
) -> Tuple[np.ndarray, str]:
|
||||
"""
|
||||
Combine multiple reference samples for better quality.
|
||||
|
||||
Args:
|
||||
audio_paths: List of audio file paths
|
||||
reference_texts: List of reference texts
|
||||
|
||||
Returns:
|
||||
Tuple of (combined_audio, combined_text)
|
||||
"""
|
||||
combined_audio = []
|
||||
|
||||
for audio_path in audio_paths:
|
||||
audio, sr = load_audio(audio_path)
|
||||
audio = normalize_audio(audio)
|
||||
combined_audio.append(audio)
|
||||
|
||||
# Concatenate audio
|
||||
mixed = np.concatenate(combined_audio)
|
||||
mixed = normalize_audio(mixed)
|
||||
|
||||
# Combine texts
|
||||
combined_text = " ".join(reference_texts)
|
||||
|
||||
return mixed, combined_text
|
||||
return await _combine_voice_prompts(audio_paths, reference_texts)
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
@@ -376,15 +223,6 @@ class PyTorchTTSBackend:
|
||||
return audio, sample_rate
|
||||
|
||||
|
||||
WHISPER_HF_REPOS = {
|
||||
"base": "openai/whisper-base",
|
||||
"small": "openai/whisper-small",
|
||||
"medium": "openai/whisper-medium",
|
||||
"large": "openai/whisper-large-v3",
|
||||
"turbo": "openai/whisper-large-v3-turbo",
|
||||
}
|
||||
|
||||
|
||||
class PyTorchSTTBackend:
|
||||
"""PyTorch-based STT backend using Whisper."""
|
||||
|
||||
@@ -396,69 +234,15 @@ class PyTorchSTTBackend:
|
||||
|
||||
def _get_device(self) -> str:
|
||||
"""Get the best available device."""
|
||||
if torch.cuda.is_available():
|
||||
return "cuda"
|
||||
# Intel Arc / Intel Xe GPU via intel-extension-for-pytorch (IPEX)
|
||||
try:
|
||||
import intel_extension_for_pytorch # noqa: F401
|
||||
if hasattr(torch, 'xpu') and torch.xpu.is_available():
|
||||
return "xpu"
|
||||
except ImportError:
|
||||
pass
|
||||
# Any GPU on Windows via DirectML (torch-directml)
|
||||
try:
|
||||
import torch_directml
|
||||
if torch_directml.device_count() > 0:
|
||||
return torch_directml.device(0)
|
||||
except ImportError:
|
||||
pass
|
||||
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
||||
return "cpu" # MPS disabled for stability
|
||||
return "cpu"
|
||||
return get_torch_device(allow_xpu=True, allow_directml=True)
|
||||
|
||||
def is_loaded(self) -> bool:
|
||||
"""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 AND fully downloaded.
|
||||
|
||||
Args:
|
||||
model_size: Model size to check
|
||||
|
||||
Returns:
|
||||
True if model is fully cached, False if missing or incomplete
|
||||
"""
|
||||
try:
|
||||
from huggingface_hub import constants as hf_constants
|
||||
hf_repo = WHISPER_HF_REPOS.get(model_size, f"openai/whisper-{model_size}")
|
||||
repo_cache = Path(hf_constants.HF_HUB_CACHE) / ("models--" + hf_repo.replace("/", "--"))
|
||||
|
||||
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
|
||||
hf_repo = WHISPER_HF_REPOS.get(model_size, f"openai/whisper-{model_size}")
|
||||
return is_model_cached(hf_repo)
|
||||
|
||||
async def load_model_async(self, model_size: Optional[str] = None):
|
||||
"""
|
||||
@@ -467,94 +251,33 @@ class PyTorchSTTBackend:
|
||||
Args:
|
||||
model_size: Model size (tiny, base, small, medium, large)
|
||||
"""
|
||||
print(f"[DEBUG] load_model_async called with size: {model_size}")
|
||||
if model_size is None:
|
||||
model_size = self.model_size
|
||||
|
||||
print(f"[DEBUG] Model already loaded? {self.model is not None}, current size: {self.model_size}, requested: {model_size}")
|
||||
if self.model is not None and self.model_size == model_size:
|
||||
print(f"[DEBUG] Early return - model already loaded")
|
||||
return
|
||||
|
||||
print(f"[DEBUG] Calling asyncio.to_thread for _load_model_sync")
|
||||
# Run blocking load in thread pool
|
||||
await asyncio.to_thread(self._load_model_sync, model_size)
|
||||
print(f"[DEBUG] asyncio.to_thread completed")
|
||||
|
||||
# Alias for compatibility
|
||||
load_model = load_model_async
|
||||
|
||||
def _load_model_sync(self, model_size: str):
|
||||
"""Synchronous model loading."""
|
||||
print(f"[DEBUG] _load_model_sync called for Whisper {model_size}")
|
||||
try:
|
||||
progress_manager = get_progress_manager()
|
||||
task_manager = get_task_manager()
|
||||
progress_model_name = f"whisper-{model_size}"
|
||||
progress_model_name = f"whisper-{model_size}"
|
||||
is_cached = self._is_model_cached(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, filter_non_downloads=is_cached)
|
||||
|
||||
# Patch tqdm BEFORE importing transformers
|
||||
print("[DEBUG] Starting tqdm patch BEFORE transformers import")
|
||||
tracker_context = tracker.patch_download()
|
||||
tracker_context.__enter__()
|
||||
print("[DEBUG] tqdm patched, now importing transformers")
|
||||
|
||||
# Import transformers
|
||||
with model_load_progress(progress_model_name, is_cached):
|
||||
from transformers import WhisperProcessor, WhisperForConditionalGeneration
|
||||
|
||||
model_name = WHISPER_HF_REPOS.get(model_size, f"openai/whisper-{model_size}")
|
||||
print(f"[DEBUG] Model name: {model_name}")
|
||||
|
||||
print(f"Loading Whisper model {model_size} on {self.device}...")
|
||||
|
||||
# Only track download progress if model is NOT cached
|
||||
if not is_cached:
|
||||
# Start tracking download task
|
||||
task_manager.start_download(progress_model_name)
|
||||
self.processor = WhisperProcessor.from_pretrained(model_name)
|
||||
self.model = WhisperForConditionalGeneration.from_pretrained(model_name)
|
||||
|
||||
# Initialize progress state so SSE endpoint has initial data to send
|
||||
progress_manager.update_progress(
|
||||
model_name=progress_model_name,
|
||||
current=0,
|
||||
total=0, # Will be updated once actual total is known
|
||||
filename="Connecting to HuggingFace...",
|
||||
status="downloading",
|
||||
)
|
||||
|
||||
# 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)
|
||||
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:
|
||||
progress_manager.mark_complete(progress_model_name)
|
||||
task_manager.complete_download(progress_model_name)
|
||||
|
||||
self.model.to(self.device)
|
||||
self.model_size = model_size
|
||||
|
||||
print(f"Whisper model {model_size} loaded successfully")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error loading Whisper model: {e}")
|
||||
progress_manager = get_progress_manager()
|
||||
task_manager = get_task_manager()
|
||||
progress_model_name = f"whisper-{model_size}"
|
||||
progress_manager.mark_error(progress_model_name, str(e))
|
||||
task_manager.error_download(progress_model_name, str(e))
|
||||
raise
|
||||
self.model.to(self.device)
|
||||
self.model_size = model_size
|
||||
print(f"Whisper model {model_size} loaded successfully")
|
||||
|
||||
def unload_model(self):
|
||||
"""Unload the model to free memory."""
|
||||
|
||||
Reference in New Issue
Block a user