mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-26 13:45:16 -07:00
ADDED MLX FOR SUPER FAST GENERATIONS ON APPLE SILICON
- Added support for MLX backend on Apple Silicon, enabling optimized performance for TTS and STT tasks. - Updated release workflow to include MLX-specific dependencies and configurations for macOS platforms. - Refactored backend code to dynamically select between MLX and PyTorch based on the runtime environment. - Enhanced model loading and inference logic to accommodate backend-specific requirements, including updated model IDs and hidden imports. - Improved health check and model status reporting to reflect the active backend type. - Streamlined caching mechanisms to support both backend types, ensuring compatibility and performance.
This commit is contained in:
@@ -17,15 +17,19 @@ jobs:
|
|||||||
- platform: 'macos-latest'
|
- platform: 'macos-latest'
|
||||||
args: '--target aarch64-apple-darwin'
|
args: '--target aarch64-apple-darwin'
|
||||||
python-version: '3.12'
|
python-version: '3.12'
|
||||||
|
backend: 'mlx'
|
||||||
- platform: 'macos-15-intel'
|
- platform: 'macos-15-intel'
|
||||||
args: '--target x86_64-apple-darwin'
|
args: '--target x86_64-apple-darwin'
|
||||||
python-version: '3.12'
|
python-version: '3.12'
|
||||||
|
backend: 'pytorch'
|
||||||
# - platform: 'ubuntu-22.04'
|
# - platform: 'ubuntu-22.04'
|
||||||
# args: ''
|
# args: ''
|
||||||
# python-version: '3.12'
|
# python-version: '3.12'
|
||||||
|
# backend: 'pytorch'
|
||||||
- platform: 'windows-latest'
|
- platform: 'windows-latest'
|
||||||
args: ''
|
args: ''
|
||||||
python-version: '3.12'
|
python-version: '3.12'
|
||||||
|
backend: 'pytorch'
|
||||||
|
|
||||||
runs-on: ${{ matrix.platform }}
|
runs-on: ${{ matrix.platform }}
|
||||||
|
|
||||||
@@ -57,6 +61,11 @@ jobs:
|
|||||||
pip install pyinstaller
|
pip install pyinstaller
|
||||||
pip install -r backend/requirements.txt
|
pip install -r backend/requirements.txt
|
||||||
|
|
||||||
|
- name: Install MLX dependencies (Apple Silicon only)
|
||||||
|
if: matrix.backend == 'mlx'
|
||||||
|
run: |
|
||||||
|
pip install -r backend/requirements-mlx.txt
|
||||||
|
|
||||||
- name: Build Python server (Linux/macOS)
|
- name: Build Python server (Linux/macOS)
|
||||||
if: matrix.platform != 'windows-latest'
|
if: matrix.platform != 'windows-latest'
|
||||||
run: |
|
run: |
|
||||||
@@ -133,7 +142,8 @@ jobs:
|
|||||||
See the assets below to download and install this version.
|
See the assets below to download and install this version.
|
||||||
|
|
||||||
### Installation
|
### Installation
|
||||||
- **macOS**: Download the `.dmg` file
|
- **macOS (Apple Silicon)**: Download the `aarch64.dmg` file - uses MLX for fast native inference
|
||||||
|
- **macOS (Intel)**: Download the `x64.dmg` file - uses PyTorch
|
||||||
- **Windows**: Download the `.msi` installer
|
- **Windows**: Download the `.msi` installer
|
||||||
- **Linux**: Download the `.AppImage` or `.deb` package
|
- **Linux**: Download the `.AppImage` or `.deb` package
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,166 @@
|
|||||||
|
"""
|
||||||
|
Backend abstraction layer for TTS and STT.
|
||||||
|
|
||||||
|
Provides a unified interface for MLX and PyTorch backends.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Protocol, Optional, Tuple, List
|
||||||
|
from typing_extensions import runtime_checkable
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from ..platform import get_backend_type
|
||||||
|
|
||||||
|
|
||||||
|
@runtime_checkable
|
||||||
|
class TTSBackend(Protocol):
|
||||||
|
"""Protocol for TTS backend implementations."""
|
||||||
|
|
||||||
|
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,
|
||||||
|
) -> 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."""
|
||||||
|
...
|
||||||
|
|
||||||
|
|
||||||
|
# Global backend instances
|
||||||
|
_tts_backend: Optional[TTSBackend] = None
|
||||||
|
_stt_backend: Optional[STTBackend] = None
|
||||||
|
|
||||||
|
|
||||||
|
def get_tts_backend() -> TTSBackend:
|
||||||
|
"""
|
||||||
|
Get or create TTS backend instance based on platform.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
TTS backend instance (MLX or PyTorch)
|
||||||
|
"""
|
||||||
|
global _tts_backend
|
||||||
|
|
||||||
|
if _tts_backend is None:
|
||||||
|
backend_type = get_backend_type()
|
||||||
|
|
||||||
|
if backend_type == "mlx":
|
||||||
|
from .mlx_backend import MLXTTSBackend
|
||||||
|
_tts_backend = MLXTTSBackend()
|
||||||
|
else:
|
||||||
|
from .pytorch_backend import PyTorchTTSBackend
|
||||||
|
_tts_backend = PyTorchTTSBackend()
|
||||||
|
|
||||||
|
return _tts_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 reset_backends():
|
||||||
|
"""Reset backend instances (useful for testing)."""
|
||||||
|
global _tts_backend, _stt_backend
|
||||||
|
_tts_backend = None
|
||||||
|
_stt_backend = None
|
||||||
@@ -0,0 +1,445 @@
|
|||||||
|
"""
|
||||||
|
MLX backend implementation for TTS and STT using mlx-audio.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Optional, List, Tuple
|
||||||
|
import asyncio
|
||||||
|
import numpy as np
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from . import TTSBackend, STTBackend
|
||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
class MLXTTSBackend:
|
||||||
|
"""MLX-based TTS backend using mlx-audio."""
|
||||||
|
|
||||||
|
def __init__(self, model_size: str = "1.7B"):
|
||||||
|
self.model = None
|
||||||
|
self.model_size = model_size
|
||||||
|
self._current_model_size = None
|
||||||
|
|
||||||
|
def is_loaded(self) -> bool:
|
||||||
|
"""Check if model is loaded."""
|
||||||
|
return self.model is not None
|
||||||
|
|
||||||
|
def _get_model_path(self, model_size: str) -> str:
|
||||||
|
"""
|
||||||
|
Get the MLX model path.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
model_size: Model size (1.7B or 0.6B)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
HuggingFace Hub model ID for MLX
|
||||||
|
"""
|
||||||
|
# MLX model mapping
|
||||||
|
mlx_model_map = {
|
||||||
|
"1.7B": "mlx-community/Qwen3-TTS-12Hz-1.7B-Base-bf16",
|
||||||
|
# 0.6B not yet converted to MLX format
|
||||||
|
"0.6B": "mlx-community/Qwen3-TTS-12Hz-1.7B-Base-bf16", # Fallback to 1.7B
|
||||||
|
}
|
||||||
|
|
||||||
|
if model_size not in mlx_model_map:
|
||||||
|
raise ValueError(f"Unknown model size: {model_size}")
|
||||||
|
|
||||||
|
hf_model_id = mlx_model_map[model_size]
|
||||||
|
print(f"Will download MLX model from HuggingFace Hub: {hf_model_id}")
|
||||||
|
|
||||||
|
return hf_model_id
|
||||||
|
|
||||||
|
async def load_model_async(self, model_size: Optional[str] = None):
|
||||||
|
"""
|
||||||
|
Lazy load the MLX TTS model.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
model_size: Model size to load (1.7B or 0.6B)
|
||||||
|
"""
|
||||||
|
if model_size is None:
|
||||||
|
model_size = self.model_size
|
||||||
|
|
||||||
|
# If already loaded with correct size, return
|
||||||
|
if self.model is not None and self._current_model_size == model_size:
|
||||||
|
return
|
||||||
|
|
||||||
|
# Unload existing model if different size requested
|
||||||
|
if self.model is not None and self._current_model_size != model_size:
|
||||||
|
self.unload_model()
|
||||||
|
|
||||||
|
# Run blocking load in thread pool
|
||||||
|
await asyncio.to_thread(self._load_model_sync, model_size)
|
||||||
|
|
||||||
|
# Alias for compatibility
|
||||||
|
load_model = load_model_async
|
||||||
|
|
||||||
|
def _load_model_sync(self, model_size: str):
|
||||||
|
"""Synchronous model loading."""
|
||||||
|
try:
|
||||||
|
from mlx_audio.tts import load
|
||||||
|
|
||||||
|
# Get model path
|
||||||
|
model_path = self._get_model_path(model_size)
|
||||||
|
|
||||||
|
# Set up progress tracking
|
||||||
|
progress_manager = get_progress_manager()
|
||||||
|
model_name = f"qwen-tts-{model_size}"
|
||||||
|
|
||||||
|
# Start tracking download task
|
||||||
|
task_manager = get_task_manager()
|
||||||
|
task_manager.start_download(model_name)
|
||||||
|
|
||||||
|
print(f"Loading MLX TTS model {model_size}...")
|
||||||
|
|
||||||
|
# Initialize progress state
|
||||||
|
progress_manager.update_progress(
|
||||||
|
model_name=model_name,
|
||||||
|
current=0,
|
||||||
|
total=1,
|
||||||
|
filename="",
|
||||||
|
status="downloading",
|
||||||
|
)
|
||||||
|
|
||||||
|
# Set up progress callback
|
||||||
|
progress_callback = create_hf_progress_callback(model_name, progress_manager)
|
||||||
|
tracker = HFProgressTracker(progress_callback)
|
||||||
|
|
||||||
|
# Use progress tracker during download
|
||||||
|
with tracker.patch_download():
|
||||||
|
# Load MLX model (downloads automatically)
|
||||||
|
self.model = load(model_path)
|
||||||
|
|
||||||
|
self._current_model_size = model_size
|
||||||
|
self.model_size = model_size
|
||||||
|
|
||||||
|
# Mark as complete
|
||||||
|
progress_manager.mark_complete(model_name)
|
||||||
|
task_manager.complete_download(model_name)
|
||||||
|
|
||||||
|
print(f"MLX TTS model {model_size} loaded successfully")
|
||||||
|
|
||||||
|
except ImportError as e:
|
||||||
|
print(f"Error: mlx_audio package not found. Install with: pip install mlx-audio")
|
||||||
|
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 MLX TTS model: {e}")
|
||||||
|
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
|
||||||
|
|
||||||
|
def unload_model(self):
|
||||||
|
"""Unload the model to free memory."""
|
||||||
|
if self.model is not None:
|
||||||
|
del self.model
|
||||||
|
self.model = None
|
||||||
|
self._current_model_size = None
|
||||||
|
print("MLX TTS model unloaded")
|
||||||
|
|
||||||
|
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.
|
||||||
|
|
||||||
|
MLX backend stores voice prompt as a dict with audio path and text.
|
||||||
|
The actual voice prompt processing happens during generation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
audio_path: Path to reference audio file
|
||||||
|
reference_text: Transcript of reference audio
|
||||||
|
use_cache: Whether to use cached prompt if available
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Tuple of (voice_prompt_dict, was_cached)
|
||||||
|
"""
|
||||||
|
await self.load_model_async(None)
|
||||||
|
|
||||||
|
# Check cache if enabled
|
||||||
|
if use_cache:
|
||||||
|
cache_key = get_cache_key(audio_path, reference_text)
|
||||||
|
cached_prompt = get_cached_voice_prompt(cache_key)
|
||||||
|
if cached_prompt is not None:
|
||||||
|
# Return cached prompt (should be dict format)
|
||||||
|
if isinstance(cached_prompt, dict):
|
||||||
|
return cached_prompt, True
|
||||||
|
|
||||||
|
# MLX voice prompt format - store audio path and text
|
||||||
|
# The model will process this during generation
|
||||||
|
voice_prompt_items = {
|
||||||
|
"ref_audio": str(audio_path),
|
||||||
|
"ref_text": reference_text,
|
||||||
|
}
|
||||||
|
|
||||||
|
# Cache if enabled
|
||||||
|
if use_cache:
|
||||||
|
cache_key = get_cache_key(audio_path, reference_text)
|
||||||
|
cache_voice_prompt(cache_key, voice_prompt_items)
|
||||||
|
|
||||||
|
return voice_prompt_items, False
|
||||||
|
|
||||||
|
async def combine_voice_prompts(
|
||||||
|
self,
|
||||||
|
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
|
||||||
|
|
||||||
|
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 using voice prompt.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
text: Text to synthesize
|
||||||
|
voice_prompt: Voice prompt dictionary with ref_audio and ref_text
|
||||||
|
language: Language code (en or zh) - may not be fully supported by MLX
|
||||||
|
seed: Random seed for reproducibility
|
||||||
|
instruct: Natural language instruction (may not be supported by MLX)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Tuple of (audio_array, sample_rate)
|
||||||
|
"""
|
||||||
|
await self.load_model_async(None)
|
||||||
|
|
||||||
|
print(f"Generating audio for text: {text}")
|
||||||
|
|
||||||
|
def _generate_sync():
|
||||||
|
"""Run synchronous generation in thread pool."""
|
||||||
|
# MLX generate() returns a generator yielding GenerationResult objects
|
||||||
|
audio_chunks = []
|
||||||
|
sample_rate = 24000
|
||||||
|
|
||||||
|
# Set seed if provided (MLX uses numpy random)
|
||||||
|
if seed is not None:
|
||||||
|
import mlx.core as mx
|
||||||
|
np.random.seed(seed)
|
||||||
|
mx.random.seed(seed)
|
||||||
|
|
||||||
|
# Extract voice prompt info
|
||||||
|
ref_audio = voice_prompt.get("ref_audio") or voice_prompt.get("ref_audio_path")
|
||||||
|
ref_text = voice_prompt.get("ref_text", "")
|
||||||
|
|
||||||
|
# Check if model supports voice cloning via generate method
|
||||||
|
# MLX API may support ref_audio parameter directly
|
||||||
|
try:
|
||||||
|
# Try with voice cloning parameters if supported
|
||||||
|
if ref_audio:
|
||||||
|
# Check if generate accepts ref_audio parameter
|
||||||
|
import inspect
|
||||||
|
sig = inspect.signature(self.model.generate)
|
||||||
|
if "ref_audio" in sig.parameters:
|
||||||
|
# Generate with voice cloning
|
||||||
|
for result in self.model.generate(text, ref_audio=ref_audio, ref_text=ref_text):
|
||||||
|
audio_chunks.append(np.array(result.audio))
|
||||||
|
sample_rate = result.sample_rate
|
||||||
|
else:
|
||||||
|
# Fallback: generate without voice cloning
|
||||||
|
for result in self.model.generate(text):
|
||||||
|
audio_chunks.append(np.array(result.audio))
|
||||||
|
sample_rate = result.sample_rate
|
||||||
|
else:
|
||||||
|
# No voice prompt, generate normally
|
||||||
|
for result in self.model.generate(text):
|
||||||
|
audio_chunks.append(np.array(result.audio))
|
||||||
|
sample_rate = result.sample_rate
|
||||||
|
except Exception as e:
|
||||||
|
# If voice cloning fails, try without it
|
||||||
|
print(f"Warning: Voice cloning failed, generating without voice prompt: {e}")
|
||||||
|
for result in self.model.generate(text):
|
||||||
|
audio_chunks.append(np.array(result.audio))
|
||||||
|
sample_rate = result.sample_rate
|
||||||
|
|
||||||
|
# Concatenate all chunks
|
||||||
|
if audio_chunks:
|
||||||
|
audio = np.concatenate([np.asarray(chunk, dtype=np.float32) for chunk in audio_chunks])
|
||||||
|
else:
|
||||||
|
# Fallback: empty audio
|
||||||
|
audio = np.array([], dtype=np.float32)
|
||||||
|
|
||||||
|
return audio, sample_rate
|
||||||
|
|
||||||
|
# Run blocking inference in thread pool
|
||||||
|
audio, sample_rate = await asyncio.to_thread(_generate_sync)
|
||||||
|
|
||||||
|
return audio, sample_rate
|
||||||
|
|
||||||
|
|
||||||
|
class MLXSTTBackend:
|
||||||
|
"""MLX-based STT backend using mlx-audio Whisper."""
|
||||||
|
|
||||||
|
def __init__(self, model_size: str = "base"):
|
||||||
|
self.model = None
|
||||||
|
self.model_size = model_size
|
||||||
|
|
||||||
|
def is_loaded(self) -> bool:
|
||||||
|
"""Check if model is loaded."""
|
||||||
|
return self.model is not None
|
||||||
|
|
||||||
|
async def load_model_async(self, model_size: Optional[str] = None):
|
||||||
|
"""
|
||||||
|
Lazy load the MLX Whisper model.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
model_size: Model size (tiny, base, small, medium, large)
|
||||||
|
"""
|
||||||
|
if model_size is None:
|
||||||
|
model_size = self.model_size
|
||||||
|
|
||||||
|
if self.model is not None and self.model_size == model_size:
|
||||||
|
return
|
||||||
|
|
||||||
|
# Run blocking load in thread pool
|
||||||
|
await asyncio.to_thread(self._load_model_sync, model_size)
|
||||||
|
|
||||||
|
# Alias for compatibility
|
||||||
|
load_model = load_model_async
|
||||||
|
|
||||||
|
def _load_model_sync(self, model_size: str):
|
||||||
|
"""Synchronous model loading."""
|
||||||
|
try:
|
||||||
|
from mlx_audio.asr import load
|
||||||
|
|
||||||
|
# MLX Whisper model naming
|
||||||
|
model_name = f"mlx-community/whisper-{model_size}"
|
||||||
|
|
||||||
|
# Set up progress tracking
|
||||||
|
progress_manager = get_progress_manager()
|
||||||
|
progress_model_name = f"whisper-{model_size}"
|
||||||
|
|
||||||
|
# Start tracking download task
|
||||||
|
task_manager = get_task_manager()
|
||||||
|
task_manager.start_download(progress_model_name)
|
||||||
|
|
||||||
|
print(f"Loading MLX Whisper model {model_size}...")
|
||||||
|
|
||||||
|
# Initialize progress state
|
||||||
|
progress_manager.update_progress(
|
||||||
|
model_name=progress_model_name,
|
||||||
|
current=0,
|
||||||
|
total=1,
|
||||||
|
filename="",
|
||||||
|
status="downloading",
|
||||||
|
)
|
||||||
|
|
||||||
|
# Set up progress callback
|
||||||
|
progress_callback = create_hf_progress_callback(progress_model_name, progress_manager)
|
||||||
|
tracker = HFProgressTracker(progress_callback)
|
||||||
|
|
||||||
|
# Use progress tracker during download
|
||||||
|
with tracker.patch_download():
|
||||||
|
self.model = load(model_name)
|
||||||
|
|
||||||
|
self.model_size = model_size
|
||||||
|
|
||||||
|
# Mark as complete
|
||||||
|
progress_manager.mark_complete(progress_model_name)
|
||||||
|
task_manager.complete_download(progress_model_name)
|
||||||
|
|
||||||
|
print(f"MLX Whisper model {model_size} loaded successfully")
|
||||||
|
|
||||||
|
except ImportError as e:
|
||||||
|
print(f"Error: mlx_audio package not found. Install with: pip install mlx-audio")
|
||||||
|
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
|
||||||
|
except Exception as e:
|
||||||
|
print(f"Error loading MLX 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
|
||||||
|
|
||||||
|
def unload_model(self):
|
||||||
|
"""Unload the model to free memory."""
|
||||||
|
if self.model is not None:
|
||||||
|
del self.model
|
||||||
|
self.model = None
|
||||||
|
print("MLX Whisper model unloaded")
|
||||||
|
|
||||||
|
async def transcribe(
|
||||||
|
self,
|
||||||
|
audio_path: str,
|
||||||
|
language: Optional[str] = None,
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Transcribe audio to text.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
audio_path: Path to audio file
|
||||||
|
language: Optional language hint (en or zh)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Transcribed text
|
||||||
|
"""
|
||||||
|
await self.load_model_async(None)
|
||||||
|
|
||||||
|
def _transcribe_sync():
|
||||||
|
"""Run synchronous transcription in thread pool."""
|
||||||
|
# Load audio
|
||||||
|
audio, sr = load_audio(audio_path, sample_rate=16000)
|
||||||
|
|
||||||
|
# MLX Whisper transcription
|
||||||
|
# The API may vary - check mlx-audio documentation
|
||||||
|
# For now, assuming similar API to PyTorch Whisper
|
||||||
|
result = self.model.transcribe(audio, language=language)
|
||||||
|
|
||||||
|
# Extract text from result (format may vary)
|
||||||
|
if isinstance(result, str):
|
||||||
|
return result.strip()
|
||||||
|
elif isinstance(result, dict):
|
||||||
|
return result.get("text", "").strip()
|
||||||
|
else:
|
||||||
|
# Try to get text attribute
|
||||||
|
return str(result).strip()
|
||||||
|
|
||||||
|
# Run blocking transcription in thread pool
|
||||||
|
return await asyncio.to_thread(_transcribe_sync)
|
||||||
@@ -0,0 +1,455 @@
|
|||||||
|
"""
|
||||||
|
PyTorch backend implementation for TTS and STT.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Optional, List, Tuple
|
||||||
|
import asyncio
|
||||||
|
import torch
|
||||||
|
import numpy as np
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from . import TTSBackend, STTBackend
|
||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
class PyTorchTTSBackend:
|
||||||
|
"""PyTorch-based TTS backend using Qwen3-TTS."""
|
||||||
|
|
||||||
|
def __init__(self, model_size: str = "1.7B"):
|
||||||
|
self.model = None
|
||||||
|
self.model_size = model_size
|
||||||
|
self.device = self._get_device()
|
||||||
|
self._current_model_size = None
|
||||||
|
|
||||||
|
def _get_device(self) -> str:
|
||||||
|
"""Get the best available device."""
|
||||||
|
if torch.cuda.is_available():
|
||||||
|
return "cuda"
|
||||||
|
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
||||||
|
# MPS can have issues, use CPU for stability
|
||||||
|
return "cpu"
|
||||||
|
return "cpu"
|
||||||
|
|
||||||
|
def is_loaded(self) -> bool:
|
||||||
|
"""Check if model is loaded."""
|
||||||
|
return self.model is not None
|
||||||
|
|
||||||
|
def _get_model_path(self, model_size: str) -> str:
|
||||||
|
"""
|
||||||
|
Get the HuggingFace Hub model ID.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
model_size: Model size (1.7B or 0.6B)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
HuggingFace Hub model ID
|
||||||
|
"""
|
||||||
|
hf_model_map = {
|
||||||
|
"1.7B": "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
|
||||||
|
"0.6B": "Qwen/Qwen3-TTS-12Hz-0.6B-Base",
|
||||||
|
}
|
||||||
|
|
||||||
|
if model_size not in hf_model_map:
|
||||||
|
raise ValueError(f"Unknown model size: {model_size}")
|
||||||
|
|
||||||
|
return hf_model_map[model_size]
|
||||||
|
|
||||||
|
async def load_model_async(self, model_size: Optional[str] = None):
|
||||||
|
"""
|
||||||
|
Lazy load the TTS model with automatic downloading from HuggingFace Hub.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
model_size: Model size to load (1.7B or 0.6B)
|
||||||
|
"""
|
||||||
|
if model_size is None:
|
||||||
|
model_size = self.model_size
|
||||||
|
|
||||||
|
# If already loaded with correct size, return
|
||||||
|
if self.model is not None and self._current_model_size == model_size:
|
||||||
|
return
|
||||||
|
|
||||||
|
# Unload existing model if different size requested
|
||||||
|
if self.model is not None and self._current_model_size != model_size:
|
||||||
|
self.unload_model()
|
||||||
|
|
||||||
|
# Run blocking load in thread pool
|
||||||
|
await asyncio.to_thread(self._load_model_sync, model_size)
|
||||||
|
|
||||||
|
# Alias for compatibility
|
||||||
|
load_model = load_model_async
|
||||||
|
|
||||||
|
def _load_model_sync(self, model_size: str):
|
||||||
|
"""Synchronous model loading."""
|
||||||
|
try:
|
||||||
|
from qwen_tts import Qwen3TTSModel
|
||||||
|
|
||||||
|
# Get model path (local or HuggingFace Hub ID)
|
||||||
|
model_path = self._get_model_path(model_size)
|
||||||
|
|
||||||
|
# Set up progress tracking
|
||||||
|
progress_manager = get_progress_manager()
|
||||||
|
model_name = f"qwen-tts-{model_size}"
|
||||||
|
|
||||||
|
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)
|
||||||
|
|
||||||
|
# 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",
|
||||||
|
)
|
||||||
|
|
||||||
|
# Set up progress callback
|
||||||
|
progress_callback = create_hf_progress_callback(model_name, progress_manager)
|
||||||
|
tracker = HFProgressTracker(progress_callback)
|
||||||
|
|
||||||
|
# Use progress tracker during download
|
||||||
|
with tracker.patch_download():
|
||||||
|
# Load the model - downloads will happen automatically with progress tracking
|
||||||
|
self.model = Qwen3TTSModel.from_pretrained(
|
||||||
|
model_path,
|
||||||
|
device_map=self.device,
|
||||||
|
torch_dtype=torch.float32 if self.device == "cpu" else torch.bfloat16,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Mark as complete
|
||||||
|
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
|
||||||
|
|
||||||
|
def unload_model(self):
|
||||||
|
"""Unload the model to free memory."""
|
||||||
|
if self.model is not None:
|
||||||
|
del self.model
|
||||||
|
self.model = None
|
||||||
|
self._current_model_size = None
|
||||||
|
|
||||||
|
if torch.cuda.is_available():
|
||||||
|
torch.cuda.empty_cache()
|
||||||
|
|
||||||
|
print("TTS model unloaded")
|
||||||
|
|
||||||
|
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.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
audio_path: Path to reference audio file
|
||||||
|
reference_text: Transcript of reference audio
|
||||||
|
use_cache: Whether to use cached prompt if available
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Tuple of (voice_prompt_dict, was_cached)
|
||||||
|
"""
|
||||||
|
await self.load_model_async(None)
|
||||||
|
|
||||||
|
# Check cache if enabled
|
||||||
|
if use_cache:
|
||||||
|
cache_key = get_cache_key(audio_path, reference_text)
|
||||||
|
cached_prompt = get_cached_voice_prompt(cache_key)
|
||||||
|
if cached_prompt is not None:
|
||||||
|
# Cache stores as torch.Tensor but actual prompt is dict
|
||||||
|
# Convert if needed
|
||||||
|
if isinstance(cached_prompt, dict):
|
||||||
|
return cached_prompt, True
|
||||||
|
elif isinstance(cached_prompt, torch.Tensor):
|
||||||
|
# Legacy cache format - convert to dict
|
||||||
|
# This shouldn't happen in practice, but handle it
|
||||||
|
return {"prompt": cached_prompt}, True
|
||||||
|
|
||||||
|
def _create_prompt_sync():
|
||||||
|
"""Run synchronous voice prompt creation in thread pool."""
|
||||||
|
return self.model.create_voice_clone_prompt(
|
||||||
|
ref_audio=str(audio_path),
|
||||||
|
ref_text=reference_text,
|
||||||
|
x_vector_only_mode=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Run blocking operation in thread pool
|
||||||
|
voice_prompt_items = await asyncio.to_thread(_create_prompt_sync)
|
||||||
|
|
||||||
|
# Cache if enabled
|
||||||
|
if use_cache:
|
||||||
|
cache_key = get_cache_key(audio_path, reference_text)
|
||||||
|
cache_voice_prompt(cache_key, voice_prompt_items)
|
||||||
|
|
||||||
|
return voice_prompt_items, False
|
||||||
|
|
||||||
|
async def combine_voice_prompts(
|
||||||
|
self,
|
||||||
|
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
|
||||||
|
|
||||||
|
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 using voice prompt.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
text: Text to synthesize
|
||||||
|
voice_prompt: Voice prompt dictionary from create_voice_prompt
|
||||||
|
language: Language code (en or zh)
|
||||||
|
seed: Random seed for reproducibility
|
||||||
|
instruct: Natural language instruction for speech delivery control
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Tuple of (audio_array, sample_rate)
|
||||||
|
"""
|
||||||
|
# Load model
|
||||||
|
await self.load_model_async(None)
|
||||||
|
|
||||||
|
def _generate_sync():
|
||||||
|
"""Run synchronous generation in thread pool."""
|
||||||
|
# Set seed if provided
|
||||||
|
if seed is not None:
|
||||||
|
torch.manual_seed(seed)
|
||||||
|
if torch.cuda.is_available():
|
||||||
|
torch.cuda.manual_seed(seed)
|
||||||
|
|
||||||
|
# Generate audio - this is the blocking operation
|
||||||
|
wavs, sample_rate = self.model.generate_voice_clone(
|
||||||
|
text=text,
|
||||||
|
voice_clone_prompt=voice_prompt,
|
||||||
|
instruct=instruct,
|
||||||
|
)
|
||||||
|
return wavs[0], sample_rate
|
||||||
|
|
||||||
|
# Run blocking inference in thread pool to avoid blocking event loop
|
||||||
|
audio, sample_rate = await asyncio.to_thread(_generate_sync)
|
||||||
|
|
||||||
|
return audio, sample_rate
|
||||||
|
|
||||||
|
|
||||||
|
class PyTorchSTTBackend:
|
||||||
|
"""PyTorch-based STT backend using Whisper."""
|
||||||
|
|
||||||
|
def __init__(self, model_size: str = "base"):
|
||||||
|
self.model = None
|
||||||
|
self.processor = None
|
||||||
|
self.model_size = model_size
|
||||||
|
self.device = self._get_device()
|
||||||
|
|
||||||
|
def _get_device(self) -> str:
|
||||||
|
"""Get the best available device."""
|
||||||
|
if torch.cuda.is_available():
|
||||||
|
return "cuda"
|
||||||
|
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
||||||
|
# MPS support for Whisper
|
||||||
|
return "cpu" # Use CPU for stability
|
||||||
|
return "cpu"
|
||||||
|
|
||||||
|
def is_loaded(self) -> bool:
|
||||||
|
"""Check if model is loaded."""
|
||||||
|
return self.model is not None
|
||||||
|
|
||||||
|
async def load_model_async(self, model_size: Optional[str] = None):
|
||||||
|
"""
|
||||||
|
Lazy load the Whisper model.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
model_size: Model size (tiny, base, small, medium, large)
|
||||||
|
"""
|
||||||
|
if model_size is None:
|
||||||
|
model_size = self.model_size
|
||||||
|
|
||||||
|
if self.model is not None and self.model_size == model_size:
|
||||||
|
return
|
||||||
|
|
||||||
|
# Run blocking load in thread pool
|
||||||
|
await asyncio.to_thread(self._load_model_sync, model_size)
|
||||||
|
|
||||||
|
# Alias for compatibility
|
||||||
|
load_model = load_model_async
|
||||||
|
|
||||||
|
def _load_model_sync(self, model_size: str):
|
||||||
|
"""Synchronous model loading."""
|
||||||
|
try:
|
||||||
|
from transformers import WhisperProcessor, WhisperForConditionalGeneration
|
||||||
|
|
||||||
|
model_name = f"openai/whisper-{model_size}"
|
||||||
|
|
||||||
|
# Set up progress tracking
|
||||||
|
progress_manager = get_progress_manager()
|
||||||
|
progress_model_name = f"whisper-{model_size}"
|
||||||
|
|
||||||
|
# Start tracking download task
|
||||||
|
task_manager = get_task_manager()
|
||||||
|
task_manager.start_download(progress_model_name)
|
||||||
|
|
||||||
|
print(f"Loading Whisper model {model_size} on {self.device}...")
|
||||||
|
|
||||||
|
# Initialize progress state to show download has started
|
||||||
|
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",
|
||||||
|
)
|
||||||
|
|
||||||
|
# Set up progress callback
|
||||||
|
progress_callback = create_hf_progress_callback(progress_model_name, progress_manager)
|
||||||
|
tracker = HFProgressTracker(progress_callback)
|
||||||
|
|
||||||
|
# Use progress tracker during download
|
||||||
|
with tracker.patch_download():
|
||||||
|
self.processor = WhisperProcessor.from_pretrained(model_name)
|
||||||
|
self.model = WhisperForConditionalGeneration.from_pretrained(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:
|
||||||
|
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
|
||||||
|
|
||||||
|
def unload_model(self):
|
||||||
|
"""Unload the model to free memory."""
|
||||||
|
if self.model is not None:
|
||||||
|
del self.model
|
||||||
|
del self.processor
|
||||||
|
self.model = None
|
||||||
|
self.processor = None
|
||||||
|
|
||||||
|
if torch.cuda.is_available():
|
||||||
|
torch.cuda.empty_cache()
|
||||||
|
|
||||||
|
print("Whisper model unloaded")
|
||||||
|
|
||||||
|
async def transcribe(
|
||||||
|
self,
|
||||||
|
audio_path: str,
|
||||||
|
language: Optional[str] = None,
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Transcribe audio to text.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
audio_path: Path to audio file
|
||||||
|
language: Optional language hint (en or zh)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Transcribed text
|
||||||
|
"""
|
||||||
|
await self.load_model_async(None)
|
||||||
|
|
||||||
|
def _transcribe_sync():
|
||||||
|
"""Run synchronous transcription in thread pool."""
|
||||||
|
# Load audio
|
||||||
|
audio, sr = load_audio(audio_path, sample_rate=16000)
|
||||||
|
|
||||||
|
# Process audio
|
||||||
|
inputs = self.processor(
|
||||||
|
audio,
|
||||||
|
sampling_rate=16000,
|
||||||
|
return_tensors="pt",
|
||||||
|
)
|
||||||
|
inputs = inputs.to(self.device)
|
||||||
|
|
||||||
|
# Set language if provided
|
||||||
|
forced_decoder_ids = None
|
||||||
|
if language:
|
||||||
|
# Support all languages from frontend: en, zh, ja, ko, de, fr, ru, pt, es, it
|
||||||
|
# Whisper supports these and many more
|
||||||
|
forced_decoder_ids = self.processor.get_decoder_prompt_ids(
|
||||||
|
language=language,
|
||||||
|
task="transcribe",
|
||||||
|
)
|
||||||
|
|
||||||
|
# Generate transcription
|
||||||
|
with torch.no_grad():
|
||||||
|
predicted_ids = self.model.generate(
|
||||||
|
inputs["input_features"],
|
||||||
|
forced_decoder_ids=forced_decoder_ids,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Decode
|
||||||
|
transcription = self.processor.batch_decode(
|
||||||
|
predicted_ids,
|
||||||
|
skip_special_tokens=True,
|
||||||
|
)[0]
|
||||||
|
|
||||||
|
return transcription.strip()
|
||||||
|
|
||||||
|
# Run blocking transcription in thread pool
|
||||||
|
return await asyncio.to_thread(_transcribe_sync)
|
||||||
+30
-1
@@ -4,9 +4,15 @@ PyInstaller build script for creating standalone Python server binary.
|
|||||||
|
|
||||||
import PyInstaller.__main__
|
import PyInstaller.__main__
|
||||||
import os
|
import os
|
||||||
|
import platform
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
def is_apple_silicon():
|
||||||
|
"""Check if running on Apple Silicon."""
|
||||||
|
return platform.system() == "Darwin" and platform.machine() == "arm64"
|
||||||
|
|
||||||
|
|
||||||
def build_server():
|
def build_server():
|
||||||
"""Build Python server as standalone binary."""
|
"""Build Python server as standalone binary."""
|
||||||
backend_dir = Path(__file__).parent
|
backend_dir = Path(__file__).parent
|
||||||
@@ -24,7 +30,7 @@ def build_server():
|
|||||||
args.extend(['--paths', str(qwen_tts_path)])
|
args.extend(['--paths', str(qwen_tts_path)])
|
||||||
print(f"Using local qwen_tts source from: {qwen_tts_path}")
|
print(f"Using local qwen_tts source from: {qwen_tts_path}")
|
||||||
|
|
||||||
# Add hidden imports
|
# Add common hidden imports
|
||||||
args.extend([
|
args.extend([
|
||||||
'--hidden-import', 'backend',
|
'--hidden-import', 'backend',
|
||||||
'--hidden-import', 'backend.main',
|
'--hidden-import', 'backend.main',
|
||||||
@@ -35,6 +41,9 @@ def build_server():
|
|||||||
'--hidden-import', 'backend.history',
|
'--hidden-import', 'backend.history',
|
||||||
'--hidden-import', 'backend.tts',
|
'--hidden-import', 'backend.tts',
|
||||||
'--hidden-import', 'backend.transcribe',
|
'--hidden-import', 'backend.transcribe',
|
||||||
|
'--hidden-import', 'backend.platform',
|
||||||
|
'--hidden-import', 'backend.backends',
|
||||||
|
'--hidden-import', 'backend.backends.pytorch_backend',
|
||||||
'--hidden-import', 'backend.utils.audio',
|
'--hidden-import', 'backend.utils.audio',
|
||||||
'--hidden-import', 'backend.utils.cache',
|
'--hidden-import', 'backend.utils.cache',
|
||||||
'--hidden-import', 'backend.utils.progress',
|
'--hidden-import', 'backend.utils.progress',
|
||||||
@@ -59,6 +68,26 @@ def build_server():
|
|||||||
# Fix for pkg_resources and jaraco namespace packages
|
# Fix for pkg_resources and jaraco namespace packages
|
||||||
'--hidden-import', 'pkg_resources.extern',
|
'--hidden-import', 'pkg_resources.extern',
|
||||||
'--collect-submodules', 'jaraco',
|
'--collect-submodules', 'jaraco',
|
||||||
|
])
|
||||||
|
|
||||||
|
# Add MLX-specific imports if building on Apple Silicon
|
||||||
|
if is_apple_silicon():
|
||||||
|
print("Building for Apple Silicon - including MLX dependencies")
|
||||||
|
args.extend([
|
||||||
|
'--hidden-import', 'backend.backends.mlx_backend',
|
||||||
|
'--hidden-import', 'mlx',
|
||||||
|
'--hidden-import', 'mlx.core',
|
||||||
|
'--hidden-import', 'mlx.nn',
|
||||||
|
'--hidden-import', 'mlx_audio',
|
||||||
|
'--hidden-import', 'mlx_audio.tts',
|
||||||
|
'--hidden-import', 'mlx_audio.asr',
|
||||||
|
'--collect-submodules', 'mlx',
|
||||||
|
'--collect-submodules', 'mlx_audio',
|
||||||
|
])
|
||||||
|
else:
|
||||||
|
print("Building for non-Apple Silicon platform - PyTorch only")
|
||||||
|
|
||||||
|
args.extend([
|
||||||
'--noconfirm',
|
'--noconfirm',
|
||||||
'--clean',
|
'--clean',
|
||||||
])
|
])
|
||||||
|
|||||||
+45
-10
@@ -27,6 +27,7 @@ from . import database, models, profiles, history, tts, transcribe, config, expo
|
|||||||
from .database import get_db, Generation as DBGeneration, VoiceProfile as DBVoiceProfile
|
from .database import get_db, Generation as DBGeneration, VoiceProfile as DBVoiceProfile
|
||||||
from .utils.progress import get_progress_manager
|
from .utils.progress import get_progress_manager
|
||||||
from .utils.tasks import get_task_manager
|
from .utils.tasks import get_task_manager
|
||||||
|
from .platform import get_backend_type
|
||||||
|
|
||||||
app = FastAPI(
|
app = FastAPI(
|
||||||
title="voicebox API",
|
title="voicebox API",
|
||||||
@@ -73,6 +74,7 @@ async def health():
|
|||||||
import os
|
import os
|
||||||
|
|
||||||
tts_model = tts.get_tts_model()
|
tts_model = tts.get_tts_model()
|
||||||
|
backend_type = get_backend_type()
|
||||||
|
|
||||||
# Check for GPU availability (CUDA or MPS)
|
# Check for GPU availability (CUDA or MPS)
|
||||||
has_cuda = torch.cuda.is_available()
|
has_cuda = torch.cuda.is_available()
|
||||||
@@ -84,6 +86,8 @@ async def health():
|
|||||||
gpu_type = f"CUDA ({torch.cuda.get_device_name(0)})"
|
gpu_type = f"CUDA ({torch.cuda.get_device_name(0)})"
|
||||||
elif has_mps:
|
elif has_mps:
|
||||||
gpu_type = "MPS (Apple Silicon)"
|
gpu_type = "MPS (Apple Silicon)"
|
||||||
|
elif backend_type == "mlx":
|
||||||
|
gpu_type = "Metal (Apple Silicon via MLX)"
|
||||||
|
|
||||||
vram_used = None
|
vram_used = None
|
||||||
if has_cuda:
|
if has_cuda:
|
||||||
@@ -111,7 +115,11 @@ async def health():
|
|||||||
model_downloaded = None
|
model_downloaded = None
|
||||||
try:
|
try:
|
||||||
# Check if the default model (1.7B) is cached
|
# Check if the default model (1.7B) is cached
|
||||||
default_model_id = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
|
# Use different model IDs based on backend
|
||||||
|
if backend_type == "mlx":
|
||||||
|
default_model_id = "mlx-community/Qwen3-TTS-12Hz-1.7B-Base-bf16"
|
||||||
|
else:
|
||||||
|
default_model_id = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
|
||||||
|
|
||||||
# Method 1: Try scan_cache_dir if available
|
# Method 1: Try scan_cache_dir if available
|
||||||
try:
|
try:
|
||||||
@@ -130,7 +138,8 @@ async def health():
|
|||||||
any(repo_cache.rglob("*.bin")) or
|
any(repo_cache.rglob("*.bin")) or
|
||||||
any(repo_cache.rglob("*.safetensors")) or
|
any(repo_cache.rglob("*.safetensors")) or
|
||||||
any(repo_cache.rglob("*.pt")) or
|
any(repo_cache.rglob("*.pt")) or
|
||||||
any(repo_cache.rglob("*.pth"))
|
any(repo_cache.rglob("*.pth")) or
|
||||||
|
any(repo_cache.rglob("*.npz")) # MLX models may use npz
|
||||||
)
|
)
|
||||||
model_downloaded = has_model_files
|
model_downloaded = has_model_files
|
||||||
except Exception:
|
except Exception:
|
||||||
@@ -144,6 +153,7 @@ async def health():
|
|||||||
gpu_available=gpu_available,
|
gpu_available=gpu_available,
|
||||||
gpu_type=gpu_type,
|
gpu_type=gpu_type,
|
||||||
vram_used_mb=vram_used,
|
vram_used_mb=vram_used,
|
||||||
|
backend_type=backend_type,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -1149,6 +1159,8 @@ async def get_model_status():
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
import os
|
import os
|
||||||
|
|
||||||
|
backend_type = get_backend_type()
|
||||||
|
|
||||||
# Try to import scan_cache_dir (might not be available in older versions)
|
# Try to import scan_cache_dir (might not be available in older versions)
|
||||||
try:
|
try:
|
||||||
from huggingface_hub import scan_cache_dir
|
from huggingface_hub import scan_cache_dir
|
||||||
@@ -1160,7 +1172,7 @@ async def get_model_status():
|
|||||||
"""Check if TTS model is loaded with specific size."""
|
"""Check if TTS model is loaded with specific size."""
|
||||||
try:
|
try:
|
||||||
tts_model = tts.get_tts_model()
|
tts_model = tts.get_tts_model()
|
||||||
return tts_model.is_loaded() and tts_model.model_size == model_size
|
return tts_model.is_loaded() and getattr(tts_model, 'model_size', None) == model_size
|
||||||
except Exception:
|
except Exception:
|
||||||
return False
|
return False
|
||||||
|
|
||||||
@@ -1168,50 +1180,66 @@ async def get_model_status():
|
|||||||
"""Check if Whisper model is loaded with specific size."""
|
"""Check if Whisper model is loaded with specific size."""
|
||||||
try:
|
try:
|
||||||
whisper_model = transcribe.get_whisper_model()
|
whisper_model = transcribe.get_whisper_model()
|
||||||
return whisper_model.is_loaded() and whisper_model.model_size == model_size
|
return whisper_model.is_loaded() and getattr(whisper_model, 'model_size', None) == model_size
|
||||||
except Exception:
|
except Exception:
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
# Use backend-specific model IDs
|
||||||
|
if backend_type == "mlx":
|
||||||
|
tts_1_7b_id = "mlx-community/Qwen3-TTS-12Hz-1.7B-Base-bf16"
|
||||||
|
tts_0_6b_id = "mlx-community/Qwen3-TTS-12Hz-1.7B-Base-bf16" # Fallback to 1.7B
|
||||||
|
whisper_base_id = "mlx-community/whisper-base"
|
||||||
|
whisper_small_id = "mlx-community/whisper-small"
|
||||||
|
whisper_medium_id = "mlx-community/whisper-medium"
|
||||||
|
whisper_large_id = "mlx-community/whisper-large"
|
||||||
|
else:
|
||||||
|
tts_1_7b_id = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
|
||||||
|
tts_0_6b_id = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
|
||||||
|
whisper_base_id = "openai/whisper-base"
|
||||||
|
whisper_small_id = "openai/whisper-small"
|
||||||
|
whisper_medium_id = "openai/whisper-medium"
|
||||||
|
whisper_large_id = "openai/whisper-large"
|
||||||
|
|
||||||
model_configs = [
|
model_configs = [
|
||||||
{
|
{
|
||||||
"model_name": "qwen-tts-1.7B",
|
"model_name": "qwen-tts-1.7B",
|
||||||
"display_name": "Qwen TTS 1.7B",
|
"display_name": "Qwen TTS 1.7B",
|
||||||
"hf_repo_id": "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
|
"hf_repo_id": tts_1_7b_id,
|
||||||
"model_size": "1.7B",
|
"model_size": "1.7B",
|
||||||
"check_loaded": lambda: check_tts_loaded("1.7B"),
|
"check_loaded": lambda: check_tts_loaded("1.7B"),
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"model_name": "qwen-tts-0.6B",
|
"model_name": "qwen-tts-0.6B",
|
||||||
"display_name": "Qwen TTS 0.6B",
|
"display_name": "Qwen TTS 0.6B",
|
||||||
"hf_repo_id": "Qwen/Qwen3-TTS-12Hz-0.6B-Base",
|
"hf_repo_id": tts_0_6b_id,
|
||||||
"model_size": "0.6B",
|
"model_size": "0.6B",
|
||||||
"check_loaded": lambda: check_tts_loaded("0.6B"),
|
"check_loaded": lambda: check_tts_loaded("0.6B"),
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"model_name": "whisper-base",
|
"model_name": "whisper-base",
|
||||||
"display_name": "Whisper Base",
|
"display_name": "Whisper Base",
|
||||||
"hf_repo_id": "openai/whisper-base",
|
"hf_repo_id": whisper_base_id,
|
||||||
"model_size": "base",
|
"model_size": "base",
|
||||||
"check_loaded": lambda: check_whisper_loaded("base"),
|
"check_loaded": lambda: check_whisper_loaded("base"),
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"model_name": "whisper-small",
|
"model_name": "whisper-small",
|
||||||
"display_name": "Whisper Small",
|
"display_name": "Whisper Small",
|
||||||
"hf_repo_id": "openai/whisper-small",
|
"hf_repo_id": whisper_small_id,
|
||||||
"model_size": "small",
|
"model_size": "small",
|
||||||
"check_loaded": lambda: check_whisper_loaded("small"),
|
"check_loaded": lambda: check_whisper_loaded("small"),
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"model_name": "whisper-medium",
|
"model_name": "whisper-medium",
|
||||||
"display_name": "Whisper Medium",
|
"display_name": "Whisper Medium",
|
||||||
"hf_repo_id": "openai/whisper-medium",
|
"hf_repo_id": whisper_medium_id,
|
||||||
"model_size": "medium",
|
"model_size": "medium",
|
||||||
"check_loaded": lambda: check_whisper_loaded("medium"),
|
"check_loaded": lambda: check_whisper_loaded("medium"),
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"model_name": "whisper-large",
|
"model_name": "whisper-large",
|
||||||
"display_name": "Whisper Large",
|
"display_name": "Whisper Large",
|
||||||
"hf_repo_id": "openai/whisper-large",
|
"hf_repo_id": whisper_large_id,
|
||||||
"model_size": "large",
|
"model_size": "large",
|
||||||
"check_loaded": lambda: check_whisper_loaded("large"),
|
"check_loaded": lambda: check_whisper_loaded("large"),
|
||||||
},
|
},
|
||||||
@@ -1256,11 +1284,13 @@ async def get_model_status():
|
|||||||
|
|
||||||
if repo_cache.exists():
|
if repo_cache.exists():
|
||||||
# Check for model files (bin, safetensors, or other common model files)
|
# Check for model files (bin, safetensors, or other common model files)
|
||||||
|
# MLX models may use .npz or .safetensors
|
||||||
has_model_files = (
|
has_model_files = (
|
||||||
any(repo_cache.rglob("*.bin")) or
|
any(repo_cache.rglob("*.bin")) or
|
||||||
any(repo_cache.rglob("*.safetensors")) or
|
any(repo_cache.rglob("*.safetensors")) or
|
||||||
any(repo_cache.rglob("*.pt")) or
|
any(repo_cache.rglob("*.pt")) or
|
||||||
any(repo_cache.rglob("*.pth")) or
|
any(repo_cache.rglob("*.pth")) or
|
||||||
|
any(repo_cache.rglob("*.npz")) or
|
||||||
any(repo_cache.rglob("model.safetensors.index.json")) or
|
any(repo_cache.rglob("model.safetensors.index.json")) or
|
||||||
any(repo_cache.rglob("pytorch_model.bin.index.json"))
|
any(repo_cache.rglob("pytorch_model.bin.index.json"))
|
||||||
)
|
)
|
||||||
@@ -1533,10 +1563,13 @@ async def get_active_tasks():
|
|||||||
|
|
||||||
def _get_gpu_status() -> str:
|
def _get_gpu_status() -> str:
|
||||||
"""Get GPU availability status."""
|
"""Get GPU availability status."""
|
||||||
|
backend_type = get_backend_type()
|
||||||
if torch.cuda.is_available():
|
if torch.cuda.is_available():
|
||||||
return f"CUDA ({torch.cuda.get_device_name(0)})"
|
return f"CUDA ({torch.cuda.get_device_name(0)})"
|
||||||
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
||||||
return "MPS (Apple Silicon)"
|
return "MPS (Apple Silicon)"
|
||||||
|
elif backend_type == "mlx":
|
||||||
|
return "Metal (Apple Silicon via MLX)"
|
||||||
return "None (CPU only)"
|
return "None (CPU only)"
|
||||||
|
|
||||||
|
|
||||||
@@ -1546,6 +1579,8 @@ async def startup_event():
|
|||||||
print("voicebox API starting up...")
|
print("voicebox API starting up...")
|
||||||
database.init_db()
|
database.init_db()
|
||||||
print(f"Database initialized at {database._db_path}")
|
print(f"Database initialized at {database._db_path}")
|
||||||
|
backend_type = get_backend_type()
|
||||||
|
print(f"Backend: {backend_type.upper()}")
|
||||||
print(f"GPU available: {_get_gpu_status()}")
|
print(f"GPU available: {_get_gpu_status()}")
|
||||||
|
|
||||||
# Initialize progress manager with main event loop for thread-safe operations
|
# Initialize progress manager with main event loop for thread-safe operations
|
||||||
|
|||||||
@@ -126,6 +126,7 @@ class HealthResponse(BaseModel):
|
|||||||
gpu_available: bool
|
gpu_available: bool
|
||||||
gpu_type: Optional[str] = None # GPU type (CUDA, MPS, or None)
|
gpu_type: Optional[str] = None # GPU type (CUDA, MPS, or None)
|
||||||
vram_used_mb: Optional[float] = None
|
vram_used_mb: Optional[float] = None
|
||||||
|
backend_type: Optional[str] = None # Backend type (mlx or pytorch)
|
||||||
|
|
||||||
|
|
||||||
class ModelStatus(BaseModel):
|
class ModelStatus(BaseModel):
|
||||||
|
|||||||
@@ -0,0 +1,33 @@
|
|||||||
|
"""
|
||||||
|
Platform detection for backend selection.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import platform
|
||||||
|
from typing import Literal
|
||||||
|
|
||||||
|
|
||||||
|
def is_apple_silicon() -> bool:
|
||||||
|
"""
|
||||||
|
Check if running on Apple Silicon (arm64 macOS).
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
True if on Apple Silicon, False otherwise
|
||||||
|
"""
|
||||||
|
return platform.system() == "Darwin" and platform.machine() == "arm64"
|
||||||
|
|
||||||
|
|
||||||
|
def get_backend_type() -> Literal["mlx", "pytorch"]:
|
||||||
|
"""
|
||||||
|
Detect the best backend for the current platform.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
"mlx" on Apple Silicon (if MLX is available), "pytorch" otherwise
|
||||||
|
"""
|
||||||
|
if is_apple_silicon():
|
||||||
|
try:
|
||||||
|
import mlx
|
||||||
|
return "mlx"
|
||||||
|
except ImportError:
|
||||||
|
# MLX not installed, fallback to PyTorch
|
||||||
|
return "pytorch"
|
||||||
|
return "pytorch"
|
||||||
@@ -0,0 +1,5 @@
|
|||||||
|
# MLX-specific dependencies (Apple Silicon only)
|
||||||
|
# These should only be installed on aarch64-apple-darwin platforms
|
||||||
|
|
||||||
|
mlx>=0.30.0
|
||||||
|
mlx-audio>=0.3.1
|
||||||
+12
-266
@@ -1,276 +1,22 @@
|
|||||||
"""
|
"""
|
||||||
Whisper ASR module for transcription.
|
STT (Speech-to-Text) module - delegates to backend abstraction layer.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from typing import Optional, List, Dict
|
from typing import Optional
|
||||||
import asyncio
|
from .backends import get_stt_backend, STTBackend
|
||||||
import torch
|
|
||||||
import numpy as np
|
|
||||||
from pathlib import Path
|
|
||||||
from .utils.progress import get_progress_manager
|
|
||||||
from .utils.hf_progress import HFProgressTracker, create_hf_progress_callback
|
|
||||||
from .utils.tasks import get_task_manager
|
|
||||||
|
|
||||||
|
|
||||||
class WhisperModel:
|
def get_whisper_model() -> STTBackend:
|
||||||
"""Manages Whisper model loading and transcription."""
|
"""
|
||||||
|
Get STT backend instance (MLX or PyTorch based on platform).
|
||||||
|
|
||||||
def __init__(self, model_size: str = "base"):
|
Returns:
|
||||||
self.model = None
|
STT backend instance
|
||||||
self.processor = None
|
"""
|
||||||
self.model_size = model_size
|
return get_stt_backend()
|
||||||
self.device = self._get_device()
|
|
||||||
|
|
||||||
def _get_device(self) -> str:
|
|
||||||
"""Get the best available device."""
|
|
||||||
if torch.cuda.is_available():
|
|
||||||
return "cuda"
|
|
||||||
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
|
||||||
# MPS support for Whisper
|
|
||||||
return "cpu" # Use CPU for stability
|
|
||||||
return "cpu"
|
|
||||||
|
|
||||||
def is_loaded(self) -> bool:
|
|
||||||
"""Check if model is loaded."""
|
|
||||||
return self.model is not None
|
|
||||||
|
|
||||||
def load_model(self, model_size: Optional[str] = None):
|
|
||||||
"""
|
|
||||||
Lazy load the Whisper model.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
model_size: Model size (tiny, base, small, medium, large)
|
|
||||||
"""
|
|
||||||
if model_size is None:
|
|
||||||
model_size = self.model_size
|
|
||||||
|
|
||||||
if self.model is not None and self.model_size == model_size:
|
|
||||||
return
|
|
||||||
|
|
||||||
try:
|
|
||||||
from transformers import WhisperProcessor, WhisperForConditionalGeneration
|
|
||||||
|
|
||||||
model_name = f"openai/whisper-{model_size}"
|
|
||||||
|
|
||||||
# Set up progress tracking
|
|
||||||
progress_manager = get_progress_manager()
|
|
||||||
progress_model_name = f"whisper-{model_size}"
|
|
||||||
|
|
||||||
# Start tracking download task
|
|
||||||
task_manager = get_task_manager()
|
|
||||||
task_manager.start_download(progress_model_name)
|
|
||||||
|
|
||||||
print(f"Loading Whisper model {model_size} on {self.device}...")
|
|
||||||
|
|
||||||
# Initialize progress state to show download has started
|
|
||||||
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",
|
|
||||||
)
|
|
||||||
|
|
||||||
# Set up progress callback
|
|
||||||
progress_callback = create_hf_progress_callback(progress_model_name, progress_manager)
|
|
||||||
tracker = HFProgressTracker(progress_callback)
|
|
||||||
|
|
||||||
# Use progress tracker during download
|
|
||||||
with tracker.patch_download():
|
|
||||||
self.processor = WhisperProcessor.from_pretrained(model_name)
|
|
||||||
self.model = WhisperForConditionalGeneration.from_pretrained(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:
|
|
||||||
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
|
|
||||||
|
|
||||||
async def load_model_async(self, model_size: Optional[str] = None):
|
|
||||||
"""
|
|
||||||
Async version of load_model that runs in thread pool.
|
|
||||||
|
|
||||||
This prevents blocking the event loop during model loading.
|
|
||||||
"""
|
|
||||||
if model_size is None:
|
|
||||||
model_size = self.model_size
|
|
||||||
|
|
||||||
# If already loaded with correct size, return immediately
|
|
||||||
if self.model is not None and self.model_size == model_size:
|
|
||||||
return
|
|
||||||
|
|
||||||
# Run the blocking load operation in a thread pool
|
|
||||||
await asyncio.to_thread(self.load_model, model_size)
|
|
||||||
|
|
||||||
def unload_model(self):
|
|
||||||
"""Unload the model to free memory."""
|
|
||||||
if self.model is not None:
|
|
||||||
del self.model
|
|
||||||
del self.processor
|
|
||||||
self.model = None
|
|
||||||
self.processor = None
|
|
||||||
|
|
||||||
if torch.cuda.is_available():
|
|
||||||
torch.cuda.empty_cache()
|
|
||||||
|
|
||||||
print("Whisper model unloaded")
|
|
||||||
|
|
||||||
async def transcribe(
|
|
||||||
self,
|
|
||||||
audio_path: str,
|
|
||||||
language: Optional[str] = None,
|
|
||||||
) -> str:
|
|
||||||
"""
|
|
||||||
Transcribe audio to text.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
audio_path: Path to audio file
|
|
||||||
language: Optional language hint (en or zh)
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
Transcribed text
|
|
||||||
"""
|
|
||||||
await self.load_model_async()
|
|
||||||
|
|
||||||
from .utils.audio import load_audio
|
|
||||||
|
|
||||||
def _transcribe_sync():
|
|
||||||
"""Run synchronous transcription in thread pool."""
|
|
||||||
# Load audio
|
|
||||||
audio, sr = load_audio(audio_path, sample_rate=16000)
|
|
||||||
|
|
||||||
# Process audio
|
|
||||||
inputs = self.processor(
|
|
||||||
audio,
|
|
||||||
sampling_rate=16000,
|
|
||||||
return_tensors="pt",
|
|
||||||
)
|
|
||||||
inputs = inputs.to(self.device)
|
|
||||||
|
|
||||||
# Set language if provided
|
|
||||||
forced_decoder_ids = None
|
|
||||||
if language:
|
|
||||||
# Support all languages from frontend: en, zh, ja, ko, de, fr, ru, pt, es, it
|
|
||||||
# Whisper supports these and many more
|
|
||||||
forced_decoder_ids = self.processor.get_decoder_prompt_ids(
|
|
||||||
language=language,
|
|
||||||
task="transcribe",
|
|
||||||
)
|
|
||||||
|
|
||||||
# Generate transcription
|
|
||||||
with torch.no_grad():
|
|
||||||
predicted_ids = self.model.generate(
|
|
||||||
inputs["input_features"],
|
|
||||||
forced_decoder_ids=forced_decoder_ids,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Decode
|
|
||||||
transcription = self.processor.batch_decode(
|
|
||||||
predicted_ids,
|
|
||||||
skip_special_tokens=True,
|
|
||||||
)[0]
|
|
||||||
|
|
||||||
return transcription.strip()
|
|
||||||
|
|
||||||
# Run blocking transcription in thread pool
|
|
||||||
return await asyncio.to_thread(_transcribe_sync)
|
|
||||||
|
|
||||||
async def transcribe_with_timestamps(
|
|
||||||
self,
|
|
||||||
audio_path: str,
|
|
||||||
language: Optional[str] = None,
|
|
||||||
) -> List[Dict[str, any]]:
|
|
||||||
"""
|
|
||||||
Transcribe audio with word-level timestamps.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
audio_path: Path to audio file
|
|
||||||
language: Optional language hint
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
List of word segments with timestamps
|
|
||||||
"""
|
|
||||||
await self.load_model_async()
|
|
||||||
|
|
||||||
from .utils.audio import load_audio
|
|
||||||
|
|
||||||
def _transcribe_timestamps_sync():
|
|
||||||
"""Run synchronous transcription with timestamps in thread pool."""
|
|
||||||
# Load audio
|
|
||||||
audio, sr = load_audio(audio_path, sample_rate=16000)
|
|
||||||
|
|
||||||
# Process audio
|
|
||||||
inputs = self.processor(
|
|
||||||
audio,
|
|
||||||
sampling_rate=16000,
|
|
||||||
return_tensors="pt",
|
|
||||||
)
|
|
||||||
inputs = inputs.to(self.device)
|
|
||||||
|
|
||||||
# Set language if provided
|
|
||||||
forced_decoder_ids = None
|
|
||||||
if language:
|
|
||||||
# Support all languages from frontend: en, zh, ja, ko, de, fr, ru, pt, es, it
|
|
||||||
# Whisper supports these and many more
|
|
||||||
forced_decoder_ids = self.processor.get_decoder_prompt_ids(
|
|
||||||
language=language,
|
|
||||||
task="transcribe",
|
|
||||||
)
|
|
||||||
|
|
||||||
# Generate with timestamps
|
|
||||||
with torch.no_grad():
|
|
||||||
predicted_ids = self.model.generate(
|
|
||||||
inputs["input_features"],
|
|
||||||
forced_decoder_ids=forced_decoder_ids,
|
|
||||||
return_timestamps=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Parse timestamps (simplified - would need more robust parsing)
|
|
||||||
# For now, return basic transcription
|
|
||||||
# TODO: Implement proper timestamp parsing
|
|
||||||
transcription = self.processor.batch_decode(
|
|
||||||
predicted_ids,
|
|
||||||
skip_special_tokens=True,
|
|
||||||
)[0]
|
|
||||||
|
|
||||||
return [
|
|
||||||
{
|
|
||||||
"text": transcription,
|
|
||||||
"start": 0.0,
|
|
||||||
"end": len(audio) / sr,
|
|
||||||
}
|
|
||||||
]
|
|
||||||
|
|
||||||
# Run blocking transcription in thread pool
|
|
||||||
return await asyncio.to_thread(_transcribe_timestamps_sync)
|
|
||||||
|
|
||||||
|
|
||||||
# Global model instance
|
|
||||||
_whisper_model: Optional[WhisperModel] = None
|
|
||||||
|
|
||||||
|
|
||||||
def get_whisper_model() -> WhisperModel:
|
|
||||||
"""Get or create Whisper model instance."""
|
|
||||||
global _whisper_model
|
|
||||||
if _whisper_model is None:
|
|
||||||
_whisper_model = WhisperModel()
|
|
||||||
return _whisper_model
|
|
||||||
|
|
||||||
|
|
||||||
def unload_whisper_model():
|
def unload_whisper_model():
|
||||||
"""Unload Whisper model to free memory."""
|
"""Unload Whisper model to free memory."""
|
||||||
global _whisper_model
|
backend = get_stt_backend()
|
||||||
if _whisper_model is not None:
|
backend.unload_model()
|
||||||
_whisper_model.unload_model()
|
|
||||||
|
|||||||
+20
-355
@@ -1,372 +1,37 @@
|
|||||||
"""
|
"""
|
||||||
TTS inference module using Qwen3-TTS.
|
TTS inference module - delegates to backend abstraction layer.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from typing import Optional, List, Tuple
|
from typing import Optional
|
||||||
import asyncio
|
|
||||||
import torch
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import io
|
import io
|
||||||
import soundfile as sf
|
import soundfile as sf
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
from .utils.cache import get_cache_key, get_cached_voice_prompt, cache_voice_prompt
|
from .backends import get_tts_backend, TTSBackend
|
||||||
from .utils.audio import normalize_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
|
|
||||||
from . import config
|
|
||||||
|
|
||||||
|
|
||||||
class TTSModel:
|
def get_tts_model() -> TTSBackend:
|
||||||
"""Manages Qwen3-TTS model loading and inference."""
|
"""
|
||||||
|
Get TTS backend instance (MLX or PyTorch based on platform).
|
||||||
|
|
||||||
def __init__(self, model_size: str = "1.7B"):
|
Returns:
|
||||||
self.model = None
|
TTS backend instance
|
||||||
self.model_size = model_size
|
"""
|
||||||
self.device = self._get_device()
|
return get_tts_backend()
|
||||||
self._current_model_size = None
|
|
||||||
|
|
||||||
def _get_device(self) -> str:
|
|
||||||
"""Get the best available device."""
|
|
||||||
if torch.cuda.is_available():
|
|
||||||
return "cuda"
|
|
||||||
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
|
||||||
# MPS can have issues, use CPU for stability
|
|
||||||
return "cpu"
|
|
||||||
return "cpu"
|
|
||||||
|
|
||||||
def is_loaded(self) -> bool:
|
|
||||||
"""Check if model is loaded."""
|
|
||||||
return self.model is not None
|
|
||||||
|
|
||||||
def _get_model_path(self, model_size: str) -> str:
|
|
||||||
"""
|
|
||||||
Get the model path, downloading from HuggingFace Hub if needed.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
model_size: Model size (1.7B or 0.6B)
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
Path to model (either local or HuggingFace Hub ID)
|
|
||||||
"""
|
|
||||||
# HuggingFace Hub model IDs
|
|
||||||
hf_model_map = {
|
|
||||||
"1.7B": "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
|
|
||||||
"0.6B": "Qwen/Qwen3-TTS-12Hz-0.6B-Base",
|
|
||||||
}
|
|
||||||
|
|
||||||
# Local directory names (for backwards compatibility)
|
|
||||||
local_model_map = {
|
|
||||||
"1.7B": "Qwen--Qwen3-TTS-12Hz-1.7B-Base",
|
|
||||||
"0.6B": "Qwen--Qwen3-TTS-12Hz-0.6B-Base",
|
|
||||||
}
|
|
||||||
|
|
||||||
if model_size not in hf_model_map:
|
|
||||||
raise ValueError(f"Unknown model size: {model_size}")
|
|
||||||
|
|
||||||
# Check if model exists locally (backwards compatibility)
|
|
||||||
local_path = config.get_models_dir() / local_model_map[model_size]
|
|
||||||
if local_path.exists():
|
|
||||||
print(f"Found local model at {local_path}")
|
|
||||||
return str(local_path)
|
|
||||||
|
|
||||||
# Use HuggingFace Hub model ID (will auto-download)
|
|
||||||
hf_model_id = hf_model_map[model_size]
|
|
||||||
print(f"Will download model from HuggingFace Hub: {hf_model_id}")
|
|
||||||
|
|
||||||
return hf_model_id
|
|
||||||
|
|
||||||
def load_model(self, model_size: Optional[str] = None):
|
|
||||||
"""
|
|
||||||
Lazy load the TTS model with automatic downloading from HuggingFace Hub.
|
|
||||||
|
|
||||||
The model will be automatically downloaded on first use and cached locally.
|
|
||||||
This works similar to how Whisper models are loaded.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
model_size: Model size to load (1.7B or 0.6B)
|
|
||||||
"""
|
|
||||||
if model_size is None:
|
|
||||||
model_size = self.model_size
|
|
||||||
|
|
||||||
# If already loaded with correct size, return
|
|
||||||
if self.model is not None and self._current_model_size == model_size:
|
|
||||||
return
|
|
||||||
|
|
||||||
# Unload existing model if different size requested
|
|
||||||
if self.model is not None and self._current_model_size != model_size:
|
|
||||||
self.unload_model()
|
|
||||||
|
|
||||||
try:
|
|
||||||
from qwen_tts import Qwen3TTSModel
|
|
||||||
|
|
||||||
# Get model path (local or HuggingFace Hub ID)
|
|
||||||
model_path = self._get_model_path(model_size)
|
|
||||||
|
|
||||||
# Set up progress tracking
|
|
||||||
progress_manager = get_progress_manager()
|
|
||||||
model_name = f"qwen-tts-{model_size}"
|
|
||||||
|
|
||||||
# Check if model is being downloaded from HuggingFace Hub
|
|
||||||
if model_path.startswith("Qwen/"):
|
|
||||||
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)
|
|
||||||
|
|
||||||
# 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",
|
|
||||||
)
|
|
||||||
|
|
||||||
# Set up progress callback
|
|
||||||
progress_callback = create_hf_progress_callback(model_name, progress_manager)
|
|
||||||
tracker = HFProgressTracker(progress_callback)
|
|
||||||
|
|
||||||
# Use progress tracker during download
|
|
||||||
with tracker.patch_download():
|
|
||||||
# Load the model - downloads will happen automatically with progress tracking
|
|
||||||
self.model = Qwen3TTSModel.from_pretrained(
|
|
||||||
model_path,
|
|
||||||
device_map=self.device,
|
|
||||||
torch_dtype=torch.float32 if self.device == "cpu" else torch.bfloat16,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Mark as complete
|
|
||||||
progress_manager.mark_complete(model_name)
|
|
||||||
task_manager.complete_download(model_name)
|
|
||||||
else:
|
|
||||||
# Local model, no download needed
|
|
||||||
print(f"Loading TTS model {model_size} on {self.device}...")
|
|
||||||
self.model = Qwen3TTSModel.from_pretrained(
|
|
||||||
model_path,
|
|
||||||
device_map=self.device,
|
|
||||||
torch_dtype=torch.float32 if self.device == "cpu" else torch.bfloat16,
|
|
||||||
)
|
|
||||||
|
|
||||||
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
|
|
||||||
|
|
||||||
async def load_model_async(self, model_size: Optional[str] = None):
|
|
||||||
"""
|
|
||||||
Async version of load_model that runs in thread pool.
|
|
||||||
|
|
||||||
This prevents blocking the event loop during model loading.
|
|
||||||
"""
|
|
||||||
if model_size is None:
|
|
||||||
model_size = self.model_size
|
|
||||||
|
|
||||||
# If already loaded with correct size, return immediately
|
|
||||||
if self.model is not None and self._current_model_size == model_size:
|
|
||||||
return
|
|
||||||
|
|
||||||
# Run the blocking load operation in a thread pool
|
|
||||||
await asyncio.to_thread(self.load_model, model_size)
|
|
||||||
|
|
||||||
def unload_model(self):
|
|
||||||
"""Unload the model to free memory."""
|
|
||||||
if self.model is not None:
|
|
||||||
del self.model
|
|
||||||
self.model = None
|
|
||||||
self._current_model_size = None
|
|
||||||
|
|
||||||
if torch.cuda.is_available():
|
|
||||||
torch.cuda.empty_cache()
|
|
||||||
|
|
||||||
print("TTS model unloaded")
|
|
||||||
|
|
||||||
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.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
audio_path: Path to reference audio file
|
|
||||||
reference_text: Transcript of reference audio
|
|
||||||
use_cache: Whether to use cached prompt if available
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
Tuple of (voice_prompt_dict, was_cached)
|
|
||||||
"""
|
|
||||||
await self.load_model_async()
|
|
||||||
|
|
||||||
# Check cache if enabled
|
|
||||||
if use_cache:
|
|
||||||
cache_key = get_cache_key(audio_path, reference_text)
|
|
||||||
cached_prompt = get_cached_voice_prompt(cache_key)
|
|
||||||
if cached_prompt is not None:
|
|
||||||
return cached_prompt, True
|
|
||||||
|
|
||||||
def _create_prompt_sync():
|
|
||||||
"""Run synchronous voice prompt creation in thread pool."""
|
|
||||||
return self.model.create_voice_clone_prompt(
|
|
||||||
ref_audio=str(audio_path),
|
|
||||||
ref_text=reference_text,
|
|
||||||
x_vector_only_mode=False,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Run blocking operation in thread pool
|
|
||||||
voice_prompt_items = await asyncio.to_thread(_create_prompt_sync)
|
|
||||||
|
|
||||||
# Cache if enabled
|
|
||||||
if use_cache:
|
|
||||||
cache_voice_prompt(cache_key, voice_prompt_items)
|
|
||||||
|
|
||||||
return voice_prompt_items, False
|
|
||||||
|
|
||||||
async def combine_voice_prompts(
|
|
||||||
self,
|
|
||||||
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)
|
|
||||||
"""
|
|
||||||
from .utils.audio import load_audio
|
|
||||||
|
|
||||||
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
|
|
||||||
|
|
||||||
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 using voice prompt.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
text: Text to synthesize
|
|
||||||
voice_prompt: Voice prompt dictionary from create_voice_prompt
|
|
||||||
language: Language code (en or zh)
|
|
||||||
seed: Random seed for reproducibility
|
|
||||||
instruct: Natural language instruction for speech delivery control
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
Tuple of (audio_array, sample_rate)
|
|
||||||
"""
|
|
||||||
# Load model (already handles async via to_thread if needed)
|
|
||||||
await self.load_model_async()
|
|
||||||
|
|
||||||
def _generate_sync():
|
|
||||||
"""Run synchronous generation in thread pool."""
|
|
||||||
# Set seed if provided
|
|
||||||
if seed is not None:
|
|
||||||
torch.manual_seed(seed)
|
|
||||||
if torch.cuda.is_available():
|
|
||||||
torch.cuda.manual_seed(seed)
|
|
||||||
|
|
||||||
# Generate audio - this is the blocking operation
|
|
||||||
wavs, sample_rate = self.model.generate_voice_clone(
|
|
||||||
text=text,
|
|
||||||
voice_clone_prompt=voice_prompt,
|
|
||||||
instruct=instruct,
|
|
||||||
)
|
|
||||||
return wavs[0], sample_rate
|
|
||||||
|
|
||||||
# Run blocking inference in thread pool to avoid blocking event loop
|
|
||||||
audio, sample_rate = await asyncio.to_thread(_generate_sync)
|
|
||||||
|
|
||||||
return audio, sample_rate
|
|
||||||
|
|
||||||
async def generate_from_reference(
|
|
||||||
self,
|
|
||||||
text: str,
|
|
||||||
audio_path: str,
|
|
||||||
reference_text: str,
|
|
||||||
language: str = "en",
|
|
||||||
seed: Optional[int] = None,
|
|
||||||
) -> Tuple[np.ndarray, int]:
|
|
||||||
"""
|
|
||||||
Generate audio directly from reference (convenience method).
|
|
||||||
|
|
||||||
Args:
|
|
||||||
text: Text to synthesize
|
|
||||||
audio_path: Path to reference audio
|
|
||||||
reference_text: Transcript of reference audio
|
|
||||||
language: Language code
|
|
||||||
seed: Random seed
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
Tuple of (audio_array, sample_rate)
|
|
||||||
"""
|
|
||||||
# Create voice prompt (with caching)
|
|
||||||
voice_prompt, _ = await self.create_voice_prompt(audio_path, reference_text)
|
|
||||||
|
|
||||||
# Generate
|
|
||||||
return await self.generate(text, voice_prompt, language, seed)
|
|
||||||
|
|
||||||
|
|
||||||
# Global model instance
|
|
||||||
_tts_model: Optional[TTSModel] = None
|
|
||||||
|
|
||||||
|
|
||||||
def get_tts_model() -> TTSModel:
|
|
||||||
"""Get or create TTS model instance."""
|
|
||||||
global _tts_model
|
|
||||||
if _tts_model is None:
|
|
||||||
_tts_model = TTSModel()
|
|
||||||
return _tts_model
|
|
||||||
|
|
||||||
|
|
||||||
def unload_tts_model():
|
def unload_tts_model():
|
||||||
"""Unload TTS model to free memory."""
|
"""Unload TTS model to free memory."""
|
||||||
global _tts_model
|
backend = get_tts_backend()
|
||||||
if _tts_model is not None:
|
backend.unload_model()
|
||||||
_tts_model.unload_model()
|
|
||||||
|
|
||||||
|
def audio_to_wav_bytes(audio: np.ndarray, sample_rate: int) -> bytes:
|
||||||
|
"""Convert audio array to WAV bytes."""
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
sf.write(buffer, audio, sample_rate, format="WAV")
|
||||||
|
buffer.seek(0)
|
||||||
|
return buffer.read()
|
||||||
|
|
||||||
|
|
||||||
def audio_to_wav_bytes(audio: np.ndarray, sample_rate: int) -> bytes:
|
def audio_to_wav_bytes(audio: np.ndarray, sample_rate: int) -> bytes:
|
||||||
|
|||||||
@@ -5,7 +5,7 @@ Voice prompt caching utilities.
|
|||||||
import hashlib
|
import hashlib
|
||||||
import torch
|
import torch
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Optional
|
from typing import Optional, Union, Dict, Any
|
||||||
|
|
||||||
from .. import config
|
from .. import config
|
||||||
|
|
||||||
@@ -15,8 +15,8 @@ def _get_cache_dir() -> Path:
|
|||||||
return config.get_cache_dir()
|
return config.get_cache_dir()
|
||||||
|
|
||||||
|
|
||||||
# In-memory cache
|
# In-memory cache - can store dict (voice prompt) or tensor (legacy)
|
||||||
_memory_cache: dict[str, torch.Tensor] = {}
|
_memory_cache: dict[str, Union[torch.Tensor, Dict[str, Any]]] = {}
|
||||||
|
|
||||||
|
|
||||||
def get_cache_key(audio_path: str, reference_text: str) -> str:
|
def get_cache_key(audio_path: str, reference_text: str) -> str:
|
||||||
@@ -43,7 +43,7 @@ def get_cache_key(audio_path: str, reference_text: str) -> str:
|
|||||||
|
|
||||||
def get_cached_voice_prompt(
|
def get_cached_voice_prompt(
|
||||||
cache_key: str,
|
cache_key: str,
|
||||||
) -> Optional[torch.Tensor]:
|
) -> Optional[Union[torch.Tensor, Dict[str, Any]]]:
|
||||||
"""
|
"""
|
||||||
Get cached voice prompt if available.
|
Get cached voice prompt if available.
|
||||||
|
|
||||||
@@ -51,7 +51,7 @@ def get_cached_voice_prompt(
|
|||||||
cache_key: Cache key
|
cache_key: Cache key
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Cached voice prompt tensor or None
|
Cached voice prompt (dict or tensor) or None
|
||||||
"""
|
"""
|
||||||
# Check in-memory cache
|
# Check in-memory cache
|
||||||
if cache_key in _memory_cache:
|
if cache_key in _memory_cache:
|
||||||
@@ -73,18 +73,18 @@ def get_cached_voice_prompt(
|
|||||||
|
|
||||||
def cache_voice_prompt(
|
def cache_voice_prompt(
|
||||||
cache_key: str,
|
cache_key: str,
|
||||||
voice_prompt: torch.Tensor,
|
voice_prompt: Union[torch.Tensor, Dict[str, Any]],
|
||||||
) -> None:
|
) -> None:
|
||||||
"""
|
"""
|
||||||
Cache voice prompt to memory and disk.
|
Cache voice prompt to memory and disk.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
cache_key: Cache key
|
cache_key: Cache key
|
||||||
voice_prompt: Voice prompt tensor
|
voice_prompt: Voice prompt (dict or tensor)
|
||||||
"""
|
"""
|
||||||
# Store in memory
|
# Store in memory
|
||||||
_memory_cache[cache_key] = voice_prompt
|
_memory_cache[cache_key] = voice_prompt
|
||||||
|
|
||||||
# Store on disk
|
# Store on disk (torch.save can handle both dicts and tensors)
|
||||||
cache_file = _get_cache_dir() / f"{cache_key}.prompt"
|
cache_file = _get_cache_dir() / f"{cache_key}.prompt"
|
||||||
torch.save(voice_prompt, cache_file)
|
torch.save(voice_prompt, cache_file)
|
||||||
|
|||||||
Binary file not shown.
Reference in New Issue
Block a user