mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-10-04 01:25:18 -07:00
The qwen configs already set retries_runaway on MLX, with the reason in a comment: mlx-audio can continue past an EOS miss and emit silence followed by codec noise. The Chatterbox MLX backend added in this PR hits the same failure and did not have the guard wired. Observed on an M4 Max with a cloned pt-BR profile: a 145 character sentence took 18.0s and returned 5.3s of audio for text worth about 8s, and a short reply came back as an endless hiss. With retries_runaway enabled the same sentence takes 3.4s and returns the full 7.4s of speech.
820 lines
26 KiB
Python
820 lines
26 KiB
Python
"""
|
|
Backend abstraction layer for TTS and STT.
|
|
|
|
Provides a unified interface for MLX and PyTorch backends,
|
|
and a model config registry that eliminates per-engine dispatch maps.
|
|
"""
|
|
|
|
# Install HF compatibility patches before any backend imports transformers /
|
|
# huggingface_hub. The module runs ``patch_transformers_mistral_regex`` at
|
|
# import time, which wraps transformers' tokenizer load against the
|
|
# unconditional HuggingFace metadata call that otherwise raises on
|
|
# HF_HUB_OFFLINE=1 and on network failures.
|
|
from ..utils import hf_offline_patch # noqa: F401
|
|
|
|
import threading
|
|
from dataclasses import dataclass, field
|
|
from typing import Protocol, Optional, Tuple, List
|
|
from typing_extensions import runtime_checkable
|
|
import numpy as np
|
|
|
|
DEFAULT_LLM_MAX_TOKENS = 512
|
|
DEFAULT_LLM_TEMPERATURE = 0.7
|
|
|
|
from ..utils.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",
|
|
}
|
|
|
|
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",
|
|
}
|
|
|
|
|
|
@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_size: str = "default"
|
|
size_mb: int = 0
|
|
needs_trim: bool = False
|
|
retries_runaway: bool = False
|
|
supports_instruct: bool = False
|
|
languages: list[str] = field(default_factory=lambda: ["en"])
|
|
|
|
|
|
@runtime_checkable
|
|
class TTSBackend(Protocol):
|
|
"""Protocol for TTS backend implementations."""
|
|
|
|
# Each backend class should define MODEL_CONFIGS as a class variable:
|
|
# MODEL_CONFIGS: list[ModelConfig]
|
|
|
|
async def load_model(self, model_size: str) -> None:
|
|
"""Load TTS model."""
|
|
...
|
|
|
|
async def create_voice_prompt(
|
|
self,
|
|
audio_path: str,
|
|
reference_text: str,
|
|
use_cache: bool = True,
|
|
) -> Tuple[dict, bool]:
|
|
"""
|
|
Create voice prompt from reference audio.
|
|
|
|
Returns:
|
|
Tuple of (voice_prompt_dict, was_cached)
|
|
"""
|
|
...
|
|
|
|
async def combine_voice_prompts(
|
|
self,
|
|
audio_paths: List[str],
|
|
reference_texts: List[str],
|
|
) -> Tuple[np.ndarray, str]:
|
|
"""
|
|
Combine multiple voice prompts.
|
|
|
|
Returns:
|
|
Tuple of (combined_audio_array, combined_text)
|
|
"""
|
|
...
|
|
|
|
async def generate(
|
|
self,
|
|
text: str,
|
|
voice_prompt: dict,
|
|
language: str = "en",
|
|
seed: Optional[int] = None,
|
|
instruct: Optional[str] = None,
|
|
) -> Tuple[np.ndarray, int]:
|
|
"""
|
|
Generate audio from text.
|
|
|
|
Returns:
|
|
Tuple of (audio_array, sample_rate)
|
|
"""
|
|
...
|
|
|
|
def unload_model(self) -> None:
|
|
"""Unload model to free memory."""
|
|
...
|
|
|
|
def is_loaded(self) -> bool:
|
|
"""Check if model is loaded."""
|
|
...
|
|
|
|
def _get_model_path(self, model_size: str) -> str:
|
|
"""
|
|
Get model path for a given size.
|
|
|
|
Returns:
|
|
Model path or HuggingFace Hub ID
|
|
"""
|
|
...
|
|
|
|
|
|
@runtime_checkable
|
|
class STTBackend(Protocol):
|
|
"""Protocol for STT (Speech-to-Text) backend implementations."""
|
|
|
|
async def load_model(self, model_size: str) -> None:
|
|
"""Load STT model."""
|
|
...
|
|
|
|
async def transcribe(
|
|
self,
|
|
audio_path: str,
|
|
language: Optional[str] = None,
|
|
model_size: Optional[str] = None,
|
|
) -> str:
|
|
"""
|
|
Transcribe audio to text.
|
|
|
|
Returns:
|
|
Transcribed text
|
|
"""
|
|
...
|
|
|
|
def unload_model(self) -> None:
|
|
"""Unload model to free memory."""
|
|
...
|
|
|
|
def is_loaded(self) -> bool:
|
|
"""Check if model is loaded."""
|
|
...
|
|
|
|
|
|
@runtime_checkable
|
|
class LLMBackend(Protocol):
|
|
"""Protocol for local LLM (chat/completion) backend implementations."""
|
|
|
|
async def load_model(self, model_size: str) -> None:
|
|
"""Load LLM weights and tokenizer."""
|
|
...
|
|
|
|
async def generate(
|
|
self,
|
|
prompt: str,
|
|
system: Optional[str] = None,
|
|
max_tokens: int = DEFAULT_LLM_MAX_TOKENS,
|
|
temperature: float = DEFAULT_LLM_TEMPERATURE,
|
|
model_size: Optional[str] = None,
|
|
examples: Optional[list[tuple[str, str]]] = None,
|
|
) -> str:
|
|
"""Run a single-turn chat completion and return the assistant reply.
|
|
|
|
``examples`` is an optional list of ``(user, assistant)`` pairs
|
|
prepended to the conversation as proper chat turns — small models
|
|
pattern-match on inline system-prompt examples (echoing them
|
|
verbatim for unrelated inputs), but treat structured turns as
|
|
data and generalize instead. Used by the refinement service.
|
|
"""
|
|
...
|
|
|
|
def unload_model(self) -> None:
|
|
...
|
|
|
|
def is_loaded(self) -> bool:
|
|
...
|
|
|
|
|
|
# Global backend instances
|
|
_tts_backend: Optional[TTSBackend] = None
|
|
_tts_backends: dict[str, TTSBackend] = {}
|
|
_tts_backends_lock = threading.Lock()
|
|
_stt_backend: Optional[STTBackend] = None
|
|
_llm_backends: dict[str, LLMBackend] = {}
|
|
_llm_backends_lock = threading.Lock()
|
|
|
|
# Supported TTS engines — keyed by engine name, value is the backend class import path.
|
|
# The factory function uses this for the if/elif chain; the model configs live on the backend classes.
|
|
TTS_ENGINES = {
|
|
"qwen": "Qwen TTS",
|
|
"qwen_custom_voice": "Qwen CustomVoice",
|
|
"luxtts": "LuxTTS",
|
|
"chatterbox": "Chatterbox TTS",
|
|
"chatterbox_turbo": "Chatterbox Turbo",
|
|
"tada": "TADA",
|
|
"kokoro": "Kokoro",
|
|
}
|
|
|
|
LLM_ENGINES = {
|
|
"qwen_llm": "Qwen3 LLM",
|
|
}
|
|
|
|
|
|
def _get_qwen_model_configs() -> list[ModelConfig]:
|
|
"""Return Qwen model configs with backend-aware HF repo IDs."""
|
|
backend_type = get_backend_type()
|
|
if backend_type == "mlx":
|
|
repo_1_7b = "mlx-community/Qwen3-TTS-12Hz-1.7B-Base-bf16"
|
|
repo_0_6b = "mlx-community/Qwen3-TTS-12Hz-0.6B-Base-bf16"
|
|
else:
|
|
repo_1_7b = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
|
|
repo_0_6b = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
|
|
|
|
# mlx-audio can continue after an EOS miss with silence followed by
|
|
# codec noise. Retry only the affected text as smaller chunks.
|
|
retries_runaway = backend_type == "mlx"
|
|
|
|
return [
|
|
ModelConfig(
|
|
model_name="qwen-tts-1.7B",
|
|
display_name="Qwen TTS 1.7B",
|
|
engine="qwen",
|
|
hf_repo_id=repo_1_7b,
|
|
model_size="1.7B",
|
|
size_mb=3500,
|
|
retries_runaway=retries_runaway,
|
|
supports_instruct=False, # Base model drops instruct silently
|
|
languages=["zh", "en", "ja", "ko", "de", "fr", "ru", "pt", "es", "it"],
|
|
),
|
|
ModelConfig(
|
|
model_name="qwen-tts-0.6B",
|
|
display_name="Qwen TTS 0.6B",
|
|
engine="qwen",
|
|
hf_repo_id=repo_0_6b,
|
|
model_size="0.6B",
|
|
size_mb=1200,
|
|
retries_runaway=retries_runaway,
|
|
supports_instruct=False,
|
|
languages=["zh", "en", "ja", "ko", "de", "fr", "ru", "pt", "es", "it"],
|
|
),
|
|
]
|
|
|
|
|
|
def _get_qwen_custom_voice_configs() -> list[ModelConfig]:
|
|
"""Return Qwen CustomVoice model configs."""
|
|
return [
|
|
ModelConfig(
|
|
model_name="qwen-custom-voice-1.7B",
|
|
display_name="Qwen CustomVoice 1.7B",
|
|
engine="qwen_custom_voice",
|
|
hf_repo_id="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
|
|
model_size="1.7B",
|
|
size_mb=3500,
|
|
supports_instruct=True,
|
|
languages=["zh", "en", "ja", "ko", "de", "fr", "ru", "pt", "es", "it"],
|
|
),
|
|
ModelConfig(
|
|
model_name="qwen-custom-voice-0.6B",
|
|
display_name="Qwen CustomVoice 0.6B",
|
|
engine="qwen_custom_voice",
|
|
hf_repo_id="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice",
|
|
model_size="0.6B",
|
|
size_mb=1200,
|
|
supports_instruct=True,
|
|
languages=["zh", "en", "ja", "ko", "de", "fr", "ru", "pt", "es", "it"],
|
|
),
|
|
]
|
|
|
|
|
|
def _get_non_qwen_tts_configs() -> list[ModelConfig]:
|
|
"""Return model configs for non-Qwen TTS engines."""
|
|
# Chatterbox multilingual follows the same backend-aware split as Qwen: the MLX
|
|
# backend loads pre-converted weights, so the download must match the backend that
|
|
# will consume it.
|
|
on_mlx = get_backend_type() == "mlx"
|
|
if on_mlx:
|
|
chatterbox_repo = "mlx-community/chatterbox-multilingual-v3"
|
|
chatterbox_size_mb = 2600
|
|
else:
|
|
chatterbox_repo = "ResembleAI/chatterbox"
|
|
chatterbox_size_mb = 3200
|
|
|
|
return [
|
|
ModelConfig(
|
|
model_name="luxtts",
|
|
display_name="LuxTTS (Fast, CPU-friendly)",
|
|
engine="luxtts",
|
|
hf_repo_id="YatharthS/LuxTTS",
|
|
size_mb=300,
|
|
languages=["en"],
|
|
),
|
|
ModelConfig(
|
|
model_name="chatterbox-tts",
|
|
display_name="Chatterbox TTS (Multilingual)",
|
|
engine="chatterbox",
|
|
hf_repo_id=chatterbox_repo,
|
|
size_mb=chatterbox_size_mb,
|
|
needs_trim=True,
|
|
# Same EOS miss the qwen configs guard against: on mlx-audio the decoder can run past
|
|
# the end of the sentence and emit silence followed by codec noise, which reaches the
|
|
# listener as an endless hiss. Retrying the affected text as smaller chunks is the
|
|
# existing remedy; it just was not wired for this engine.
|
|
retries_runaway=on_mlx,
|
|
languages=[
|
|
"zh",
|
|
"en",
|
|
"ja",
|
|
"ko",
|
|
"de",
|
|
"fr",
|
|
"ru",
|
|
"pt",
|
|
"es",
|
|
"it",
|
|
"he",
|
|
"ar",
|
|
"da",
|
|
"el",
|
|
"fi",
|
|
"hi",
|
|
"ms",
|
|
"nl",
|
|
"no",
|
|
"pl",
|
|
"sv",
|
|
"sw",
|
|
"tr",
|
|
],
|
|
),
|
|
ModelConfig(
|
|
model_name="chatterbox-turbo",
|
|
display_name="Chatterbox Turbo (English, Tags)",
|
|
engine="chatterbox_turbo",
|
|
hf_repo_id="ResembleAI/chatterbox-turbo",
|
|
size_mb=1500,
|
|
needs_trim=True,
|
|
languages=["en"],
|
|
),
|
|
ModelConfig(
|
|
model_name="tada-1b",
|
|
display_name="TADA 1B (English)",
|
|
engine="tada",
|
|
hf_repo_id="HumeAI/tada-1b",
|
|
model_size="1B",
|
|
size_mb=4000,
|
|
languages=["en"],
|
|
),
|
|
ModelConfig(
|
|
model_name="tada-3b-ml",
|
|
display_name="TADA 3B Multilingual",
|
|
engine="tada",
|
|
hf_repo_id="HumeAI/tada-3b-ml",
|
|
model_size="3B",
|
|
size_mb=8000,
|
|
languages=["en", "ar", "zh", "de", "es", "fr", "it", "ja", "pl", "pt"],
|
|
),
|
|
ModelConfig(
|
|
model_name="kokoro",
|
|
display_name="Kokoro 82M",
|
|
engine="kokoro",
|
|
hf_repo_id="hexgrad/Kokoro-82M",
|
|
size_mb=350,
|
|
languages=["en", "es", "fr", "hi", "it", "pt", "ja", "zh"],
|
|
),
|
|
]
|
|
|
|
|
|
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",
|
|
),
|
|
]
|
|
|
|
|
|
def _get_qwen_llm_configs() -> list[ModelConfig]:
|
|
"""Return Qwen3 LLM configs with backend-aware HF repo IDs.
|
|
|
|
MLX path uses 4-bit community quantizations for Apple Silicon; PyTorch path
|
|
uses the upstream instruct weights.
|
|
"""
|
|
backend_type = get_backend_type()
|
|
if backend_type == "mlx":
|
|
repo_0_6 = "mlx-community/Qwen3-0.6B-4bit"
|
|
repo_1_7 = "mlx-community/Qwen3-1.7B-4bit"
|
|
repo_4 = "mlx-community/Qwen3-4B-4bit"
|
|
else:
|
|
repo_0_6 = "Qwen/Qwen3-0.6B"
|
|
repo_1_7 = "Qwen/Qwen3-1.7B"
|
|
repo_4 = "Qwen/Qwen3-4B"
|
|
|
|
common_languages = [
|
|
"en", "zh", "ja", "ko", "de", "fr", "ru", "pt", "es", "it",
|
|
]
|
|
|
|
return [
|
|
ModelConfig(
|
|
model_name="qwen3-0.6b",
|
|
display_name="Qwen3 0.6B",
|
|
engine="qwen_llm",
|
|
hf_repo_id=repo_0_6,
|
|
model_size="0.6B",
|
|
size_mb=400 if backend_type == "mlx" else 1400,
|
|
languages=common_languages,
|
|
),
|
|
ModelConfig(
|
|
model_name="qwen3-1.7b",
|
|
display_name="Qwen3 1.7B",
|
|
engine="qwen_llm",
|
|
hf_repo_id=repo_1_7,
|
|
model_size="1.7B",
|
|
size_mb=1100 if backend_type == "mlx" else 3500,
|
|
languages=common_languages,
|
|
),
|
|
ModelConfig(
|
|
model_name="qwen3-4b",
|
|
display_name="Qwen3 4B",
|
|
engine="qwen_llm",
|
|
hf_repo_id=repo_4,
|
|
model_size="4B",
|
|
size_mb=2500 if backend_type == "mlx" else 8000,
|
|
languages=common_languages,
|
|
),
|
|
]
|
|
|
|
|
|
def get_all_model_configs() -> list[ModelConfig]:
|
|
"""Return the full list of model configs (TTS + STT + LLM)."""
|
|
return (
|
|
_get_qwen_model_configs()
|
|
+ _get_qwen_custom_voice_configs()
|
|
+ _get_non_qwen_tts_configs()
|
|
+ _get_whisper_configs()
|
|
+ _get_qwen_llm_configs()
|
|
)
|
|
|
|
|
|
def get_tts_model_configs() -> list[ModelConfig]:
|
|
"""Return only TTS model configs."""
|
|
return _get_qwen_model_configs() + _get_qwen_custom_voice_configs() + _get_non_qwen_tts_configs()
|
|
|
|
|
|
def get_llm_model_configs() -> list[ModelConfig]:
|
|
"""Return only LLM model configs."""
|
|
return _get_qwen_llm_configs()
|
|
|
|
|
|
def get_stt_model_configs() -> list[ModelConfig]:
|
|
"""Return only STT (Whisper) model configs."""
|
|
return _get_whisper_configs()
|
|
|
|
|
|
# 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():
|
|
if cfg.model_name == model_name:
|
|
return cfg
|
|
return None
|
|
|
|
|
|
def engine_needs_trim(engine: str) -> bool:
|
|
"""Whether this engine's output should be run through trim_tts_output."""
|
|
for cfg in get_tts_model_configs():
|
|
if cfg.engine == engine:
|
|
return cfg.needs_trim
|
|
return False
|
|
|
|
|
|
def engine_retries_runaway(engine: str) -> bool:
|
|
"""Whether unstable output should be retried in smaller chunks."""
|
|
for cfg in get_tts_model_configs():
|
|
if cfg.engine == engine:
|
|
return cfg.retries_runaway
|
|
return False
|
|
|
|
|
|
def engine_has_model_sizes(engine: str) -> bool:
|
|
"""Whether this engine supports multiple model sizes (only Qwen currently)."""
|
|
configs = [c for c in get_tts_model_configs() if c.engine == engine]
|
|
return len(configs) > 1
|
|
|
|
|
|
async def load_engine_model(engine: str, model_size: str = "default") -> None:
|
|
"""Load a model for the given engine, handling engines with multiple model sizes."""
|
|
backend = get_tts_backend_for_engine(engine)
|
|
if engine in ("qwen", "qwen_custom_voice"):
|
|
await backend.load_model_async(model_size)
|
|
elif engine == "tada":
|
|
await backend.load_model(model_size)
|
|
else:
|
|
await backend.load_model()
|
|
|
|
|
|
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():
|
|
if c.engine == engine and c.model_size == model_size:
|
|
cfg = c
|
|
break
|
|
|
|
if engine in ("qwen", "qwen_custom_voice", "tada"):
|
|
if not backend._is_model_cached(model_size):
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail=f"Model {model_size} is not downloaded yet. Use /generate to trigger a download.",
|
|
)
|
|
else:
|
|
if not backend._is_model_cached():
|
|
display = cfg.display_name if cfg else engine
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail=f"{display} model is not downloaded yet. Use /generate to trigger a download.",
|
|
)
|
|
|
|
|
|
def unload_model_by_config(config: ModelConfig) -> bool:
|
|
"""Unload a model given its config. Returns True if it was loaded, False otherwise."""
|
|
from . import get_tts_backend_for_engine
|
|
from ..services import tts, transcribe, llm as llm_service
|
|
from ..utils.cache import clear_voice_prompt_memory_cache
|
|
|
|
if config.engine == "whisper":
|
|
whisper_model = transcribe.get_whisper_model()
|
|
if whisper_model.is_loaded() and whisper_model.model_size == config.model_size:
|
|
transcribe.unload_whisper_model()
|
|
return True
|
|
return False
|
|
|
|
if config.engine == "qwen_llm":
|
|
backend = llm_service.get_llm_model()
|
|
loaded_size = getattr(backend, "_current_model_size", None) or getattr(backend, "model_size", None)
|
|
if backend.is_loaded() and loaded_size == config.model_size:
|
|
backend.unload_model()
|
|
return True
|
|
return False
|
|
|
|
if config.engine == "qwen":
|
|
tts_model = tts.get_tts_model()
|
|
loaded_size = getattr(tts_model, "_current_model_size", None) or getattr(tts_model, "model_size", None)
|
|
if tts_model.is_loaded() and loaded_size == config.model_size:
|
|
tts.unload_tts_model()
|
|
return True
|
|
return False
|
|
|
|
if config.engine == "qwen_custom_voice":
|
|
backend = get_tts_backend_for_engine(config.engine)
|
|
loaded_size = getattr(backend, "_current_model_size", None) or getattr(backend, "model_size", None)
|
|
if backend.is_loaded() and loaded_size == config.model_size:
|
|
clear_voice_prompt_memory_cache()
|
|
backend.unload_model()
|
|
return True
|
|
return False
|
|
|
|
# All other TTS engines
|
|
backend = get_tts_backend_for_engine(config.engine)
|
|
if backend.is_loaded():
|
|
clear_voice_prompt_memory_cache()
|
|
backend.unload_model()
|
|
return True
|
|
return False
|
|
|
|
|
|
def check_model_loaded(config: ModelConfig) -> bool:
|
|
"""Check if a model is currently loaded."""
|
|
from . import get_tts_backend_for_engine
|
|
from ..services import tts, transcribe, llm as llm_service
|
|
|
|
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
|
|
|
|
if config.engine == "qwen_llm":
|
|
backend = llm_service.get_llm_model()
|
|
loaded_size = getattr(backend, "_current_model_size", None) or getattr(backend, "model_size", None)
|
|
return backend.is_loaded() and loaded_size == config.model_size
|
|
|
|
if config.engine == "qwen":
|
|
tts_model = tts.get_tts_model()
|
|
loaded_size = getattr(tts_model, "_current_model_size", None) or getattr(tts_model, "model_size", None)
|
|
return tts_model.is_loaded() and loaded_size == config.model_size
|
|
|
|
if config.engine == "qwen_custom_voice":
|
|
backend = get_tts_backend_for_engine(config.engine)
|
|
loaded_size = getattr(backend, "_current_model_size", None) or getattr(backend, "model_size", None)
|
|
return backend.is_loaded() and loaded_size == config.model_size
|
|
|
|
backend = get_tts_backend_for_engine(config.engine)
|
|
return backend.is_loaded()
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def get_model_load_func(config: ModelConfig):
|
|
"""Return a callable that loads/downloads the model."""
|
|
from . import get_tts_backend_for_engine
|
|
from ..services import tts, transcribe, llm as llm_service
|
|
|
|
if config.engine == "whisper":
|
|
return lambda: transcribe.get_whisper_model().load_model(config.model_size)
|
|
|
|
if config.engine == "qwen":
|
|
return lambda: tts.get_tts_model().load_model(config.model_size)
|
|
|
|
if config.engine == "qwen_custom_voice":
|
|
return lambda: get_tts_backend_for_engine(config.engine).load_model(config.model_size)
|
|
|
|
if config.engine == "qwen_llm":
|
|
return lambda: llm_service.get_llm_model().load_model(config.model_size)
|
|
|
|
return lambda: get_tts_backend_for_engine(config.engine).load_model()
|
|
|
|
|
|
def get_tts_backend() -> TTSBackend:
|
|
"""
|
|
Get or create the default (Qwen) TTS backend instance based on platform.
|
|
|
|
Returns:
|
|
TTS backend instance (MLX or PyTorch)
|
|
"""
|
|
return get_tts_backend_for_engine("qwen")
|
|
|
|
|
|
def get_tts_backend_for_engine(engine: str) -> TTSBackend:
|
|
"""
|
|
Get or create a TTS backend for the given engine.
|
|
|
|
Args:
|
|
engine: Engine name (e.g. "qwen", "luxtts", "chatterbox", "chatterbox_turbo")
|
|
|
|
Returns:
|
|
TTS backend instance
|
|
"""
|
|
global _tts_backends
|
|
|
|
# Fast path: check without lock
|
|
if engine in _tts_backends:
|
|
return _tts_backends[engine]
|
|
|
|
# Slow path: create with lock to avoid duplicate instantiation
|
|
with _tts_backends_lock:
|
|
# Double-check after acquiring lock
|
|
if engine in _tts_backends:
|
|
return _tts_backends[engine]
|
|
|
|
if engine == "qwen":
|
|
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":
|
|
# Same split the qwen engine already makes: on Apple Silicon the MLX/Metal
|
|
# port renders 7-9x faster than the CPU-pinned PyTorch path.
|
|
if get_backend_type() == "mlx":
|
|
from .chatterbox_mlx_backend import ChatterboxMLXTTSBackend
|
|
|
|
backend = ChatterboxMLXTTSBackend()
|
|
else:
|
|
from .chatterbox_backend import ChatterboxTTSBackend
|
|
|
|
backend = ChatterboxTTSBackend()
|
|
elif engine == "chatterbox_turbo":
|
|
from .chatterbox_turbo_backend import ChatterboxTurboTTSBackend
|
|
|
|
backend = ChatterboxTurboTTSBackend()
|
|
elif engine == "tada":
|
|
from .hume_backend import HumeTadaBackend
|
|
|
|
backend = HumeTadaBackend()
|
|
elif engine == "kokoro":
|
|
from .kokoro_backend import KokoroTTSBackend
|
|
|
|
backend = KokoroTTSBackend()
|
|
elif engine == "qwen_custom_voice":
|
|
from .qwen_custom_voice_backend import QwenCustomVoiceBackend
|
|
|
|
backend = QwenCustomVoiceBackend()
|
|
else:
|
|
raise ValueError(f"Unknown TTS engine: {engine}. Supported: {list(TTS_ENGINES.keys())}")
|
|
|
|
_tts_backends[engine] = backend
|
|
return backend
|
|
|
|
|
|
def get_stt_backend() -> STTBackend:
|
|
"""
|
|
Get or create STT backend instance based on platform.
|
|
|
|
Returns:
|
|
STT backend instance (MLX or PyTorch)
|
|
"""
|
|
global _stt_backend
|
|
|
|
if _stt_backend is None:
|
|
backend_type = get_backend_type()
|
|
|
|
if backend_type == "mlx":
|
|
from .mlx_backend import MLXSTTBackend
|
|
|
|
_stt_backend = MLXSTTBackend()
|
|
else:
|
|
from .pytorch_backend import PyTorchSTTBackend
|
|
|
|
_stt_backend = PyTorchSTTBackend()
|
|
|
|
return _stt_backend
|
|
|
|
|
|
def get_llm_backend() -> LLMBackend:
|
|
"""Get or create the default Qwen3 LLM backend based on platform."""
|
|
return get_llm_backend_for_engine("qwen_llm")
|
|
|
|
|
|
def get_llm_backend_for_engine(engine: str) -> LLMBackend:
|
|
"""Get or create an LLM backend for the given engine."""
|
|
global _llm_backends
|
|
|
|
if engine in _llm_backends:
|
|
return _llm_backends[engine]
|
|
|
|
with _llm_backends_lock:
|
|
if engine in _llm_backends:
|
|
return _llm_backends[engine]
|
|
|
|
if engine == "qwen_llm":
|
|
backend_type = get_backend_type()
|
|
if backend_type == "mlx":
|
|
from .qwen_llm_backend import MLXQwenLLMBackend
|
|
|
|
backend = MLXQwenLLMBackend()
|
|
else:
|
|
from .qwen_llm_backend import PyTorchQwenLLMBackend
|
|
|
|
backend = PyTorchQwenLLMBackend()
|
|
else:
|
|
raise ValueError(f"Unknown LLM engine: {engine}. Supported: {list(LLM_ENGINES.keys())}")
|
|
|
|
_llm_backends[engine] = backend
|
|
return backend
|
|
|
|
|
|
def reset_backends():
|
|
"""Reset backend instances (useful for testing)."""
|
|
global _tts_backend, _tts_backends, _stt_backend, _llm_backends
|
|
_tts_backend = None
|
|
_tts_backends.clear()
|
|
_stt_backend = None
|
|
_llm_backends.clear()
|