comment cleanup

This commit is contained in:
James Pine
2026-03-16 01:46:19 -07:00
parent fe19a9ca47
commit b7781951df
10 changed files with 738 additions and 694 deletions
+83 -24
View File
@@ -14,9 +14,16 @@ import numpy as np
from ..platform_detect import get_backend_type
LANGUAGE_CODE_TO_NAME = {
"zh": "chinese", "en": "english", "ja": "japanese", "ko": "korean",
"de": "german", "fr": "french", "ru": "russian", "pt": "portuguese",
"es": "spanish", "it": "italian",
"zh": "chinese",
"en": "english",
"ja": "japanese",
"ko": "korean",
"de": "german",
"fr": "french",
"ru": "russian",
"pt": "portuguese",
"es": "spanish",
"it": "italian",
}
WHISPER_HF_REPOS = {
@@ -31,10 +38,11 @@ WHISPER_HF_REPOS = {
@dataclass
class ModelConfig:
"""Declarative config for a downloadable model variant."""
model_name: str # e.g. "luxtts", "chatterbox-tts"
display_name: str # e.g. "LuxTTS (Fast, CPU-friendly)"
engine: str # e.g. "luxtts", "chatterbox"
hf_repo_id: str # e.g. "YatharthS/LuxTTS"
model_name: str # e.g. "luxtts", "chatterbox-tts"
display_name: str # e.g. "LuxTTS (Fast, CPU-friendly)"
engine: str # e.g. "luxtts", "chatterbox"
hf_repo_id: str # e.g. "YatharthS/LuxTTS"
model_size: str = "default"
size_mb: int = 0
needs_trim: bool = False
@@ -160,10 +168,6 @@ TTS_ENGINES = {
}
# ---------------------------------------------------------------------------
# Model config registry
# ---------------------------------------------------------------------------
def _get_qwen_model_configs() -> list[ModelConfig]:
"""Return Qwen model configs with backend-aware HF repo IDs."""
backend_type = get_backend_type()
@@ -220,9 +224,29 @@ def _get_non_qwen_tts_configs() -> list[ModelConfig]:
size_mb=3200,
needs_trim=True,
languages=[
"zh", "en", "ja", "ko", "de", "fr", "ru", "pt", "es", "it",
"he", "ar", "da", "el", "fi", "hi", "ms", "nl", "no", "pl",
"sv", "sw", "tr",
"zh",
"en",
"ja",
"ko",
"de",
"fr",
"ru",
"pt",
"es",
"it",
"he",
"ar",
"da",
"el",
"fi",
"hi",
"ms",
"nl",
"no",
"pl",
"sv",
"sw",
"tr",
],
),
ModelConfig(
@@ -240,11 +264,41 @@ def _get_non_qwen_tts_configs() -> list[ModelConfig]:
def _get_whisper_configs() -> list[ModelConfig]:
"""Return Whisper STT model configs."""
return [
ModelConfig(model_name="whisper-base", display_name="Whisper Base", engine="whisper", hf_repo_id="openai/whisper-base", model_size="base"),
ModelConfig(model_name="whisper-small", display_name="Whisper Small", engine="whisper", hf_repo_id="openai/whisper-small", model_size="small"),
ModelConfig(model_name="whisper-medium", display_name="Whisper Medium", engine="whisper", hf_repo_id="openai/whisper-medium", model_size="medium"),
ModelConfig(model_name="whisper-large", display_name="Whisper Large", engine="whisper", hf_repo_id="openai/whisper-large-v3", model_size="large"),
ModelConfig(model_name="whisper-turbo", display_name="Whisper Turbo", engine="whisper", hf_repo_id="openai/whisper-large-v3-turbo", model_size="turbo"),
ModelConfig(
model_name="whisper-base",
display_name="Whisper Base",
engine="whisper",
hf_repo_id="openai/whisper-base",
model_size="base",
),
ModelConfig(
model_name="whisper-small",
display_name="Whisper Small",
engine="whisper",
hf_repo_id="openai/whisper-small",
model_size="small",
),
ModelConfig(
model_name="whisper-medium",
display_name="Whisper Medium",
engine="whisper",
hf_repo_id="openai/whisper-medium",
model_size="medium",
),
ModelConfig(
model_name="whisper-large",
display_name="Whisper Large",
engine="whisper",
hf_repo_id="openai/whisper-large-v3",
model_size="large",
),
ModelConfig(
model_name="whisper-turbo",
display_name="Whisper Turbo",
engine="whisper",
hf_repo_id="openai/whisper-large-v3-turbo",
model_size="turbo",
),
]
@@ -260,6 +314,7 @@ def get_tts_model_configs() -> list[ModelConfig]:
# Lookup helpers — these replace the if/elif chains in main.py
def get_model_config(model_name: str) -> Optional[ModelConfig]:
"""Look up a model config by model_name."""
for cfg in get_all_model_configs():
@@ -294,6 +349,7 @@ async def load_engine_model(engine: str, model_size: str = "default") -> None:
async def ensure_model_cached_or_raise(engine: str, model_size: str = "default") -> None:
"""Check if a model is cached, raise HTTPException if not. Used by streaming endpoint."""
from fastapi import HTTPException
backend = get_tts_backend_for_engine(engine)
cfg = None
for c in get_tts_model_configs():
@@ -352,7 +408,7 @@ def check_model_loaded(config: ModelConfig) -> bool:
try:
if config.engine == "whisper":
whisper_model = transcribe.get_whisper_model()
return whisper_model.is_loaded() and getattr(whisper_model, 'model_size', None) == config.model_size
return whisper_model.is_loaded() and getattr(whisper_model, "model_size", None) == config.model_size
if config.engine == "qwen":
tts_model = tts.get_tts_model()
@@ -379,10 +435,6 @@ def get_model_load_func(config: ModelConfig):
return lambda: get_tts_backend_for_engine(config.engine).load_model()
# ---------------------------------------------------------------------------
# Backend factory
# ---------------------------------------------------------------------------
def get_tts_backend() -> TTSBackend:
"""
Get or create the default (Qwen) TTS backend instance based on platform.
@@ -419,18 +471,23 @@ def get_tts_backend_for_engine(engine: str) -> TTSBackend:
backend_type = get_backend_type()
if backend_type == "mlx":
from .mlx_backend import MLXTTSBackend
backend = MLXTTSBackend()
else:
from .pytorch_backend import PyTorchTTSBackend
backend = PyTorchTTSBackend()
elif engine == "luxtts":
from .luxtts_backend import LuxTTSBackend
backend = LuxTTSBackend()
elif engine == "chatterbox":
from .chatterbox_backend import ChatterboxTTSBackend
backend = ChatterboxTTSBackend()
elif engine == "chatterbox_turbo":
from .chatterbox_turbo_backend import ChatterboxTurboTTSBackend
backend = ChatterboxTurboTTSBackend()
else:
raise ValueError(f"Unknown TTS engine: {engine}. Supported: {list(TTS_ENGINES.keys())}")
@@ -453,9 +510,11 @@ def get_stt_backend() -> STTBackend:
if backend_type == "mlx":
from .mlx_backend import MLXSTTBackend
_stt_backend = MLXSTTBackend()
else:
from .pytorch_backend import PyTorchSTTBackend
_stt_backend = PyTorchSTTBackend()
return _stt_backend
+4 -23
View File
@@ -21,10 +21,6 @@ from ..utils.tasks import get_task_manager
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# HuggingFace cache checking
# ---------------------------------------------------------------------------
def is_model_cached(
hf_repo: str,
*,
@@ -46,9 +42,7 @@ def is_model_cached(
try:
from huggingface_hub import constants as hf_constants
repo_cache = Path(hf_constants.HF_HUB_CACHE) / (
"models--" + hf_repo.replace("/", "--")
)
repo_cache = Path(hf_constants.HF_HUB_CACHE) / ("models--" + hf_repo.replace("/", "--"))
if not repo_cache.exists():
return False
@@ -83,10 +77,6 @@ def is_model_cached(
return False
# ---------------------------------------------------------------------------
# Device detection
# ---------------------------------------------------------------------------
def get_torch_device(
*,
allow_xpu: bool = False,
@@ -114,6 +104,7 @@ def get_torch_device(
if allow_xpu:
try:
import intel_extension_for_pytorch # noqa: F401
if hasattr(torch, "xpu") and torch.xpu.is_available():
return "xpu"
except ImportError:
@@ -122,6 +113,7 @@ def get_torch_device(
if allow_directml:
try:
import torch_directml
if torch_directml.device_count() > 0:
return torch_directml.device(0)
except ImportError:
@@ -134,10 +126,6 @@ def get_torch_device(
return "cpu"
# ---------------------------------------------------------------------------
# Voice prompt combination
# ---------------------------------------------------------------------------
async def combine_voice_prompts(
audio_paths: List[str],
reference_texts: List[str],
@@ -169,10 +157,6 @@ async def combine_voice_prompts(
return mixed, combined_text
# ---------------------------------------------------------------------------
# Model loading progress tracking
# ---------------------------------------------------------------------------
@contextmanager
def model_load_progress(
model_name: str,
@@ -237,10 +221,6 @@ def model_load_progress(
tracker_context.__exit__(None, None, None)
# ---------------------------------------------------------------------------
# Chatterbox f32 dtype patches
# ---------------------------------------------------------------------------
def patch_chatterbox_f32(model) -> None:
"""
Patch float64 -> float32 dtype mismatches in upstream chatterbox.
@@ -261,6 +241,7 @@ def patch_chatterbox_f32(model) -> None:
def _f32_log_mel(self_tokzr, audio, padding=0):
import torch as _torch
if _torch.is_tensor(audio):
audio = audio.float()
return _orig_log_mel(self_tokzr, audio, padding)