diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 81f3a126..fb2b578d 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -17,15 +17,19 @@ jobs: - platform: 'macos-latest' args: '--target aarch64-apple-darwin' python-version: '3.12' + backend: 'mlx' - platform: 'macos-15-intel' args: '--target x86_64-apple-darwin' python-version: '3.12' + backend: 'pytorch' # - platform: 'ubuntu-22.04' # args: '' # python-version: '3.12' + # backend: 'pytorch' - platform: 'windows-latest' args: '' python-version: '3.12' + backend: 'pytorch' runs-on: ${{ matrix.platform }} @@ -57,6 +61,11 @@ jobs: pip install pyinstaller 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) if: matrix.platform != 'windows-latest' run: | @@ -133,7 +142,8 @@ jobs: See the assets below to download and install this version. ### 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 - **Linux**: Download the `.AppImage` or `.deb` package diff --git a/backend/backends/__init__.py b/backend/backends/__init__.py new file mode 100644 index 00000000..d7167826 --- /dev/null +++ b/backend/backends/__init__.py @@ -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 diff --git a/backend/backends/mlx_backend.py b/backend/backends/mlx_backend.py new file mode 100644 index 00000000..289901d4 --- /dev/null +++ b/backend/backends/mlx_backend.py @@ -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) diff --git a/backend/backends/pytorch_backend.py b/backend/backends/pytorch_backend.py new file mode 100644 index 00000000..c915761d --- /dev/null +++ b/backend/backends/pytorch_backend.py @@ -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) diff --git a/backend/build_binary.py b/backend/build_binary.py index e9b23733..da63907f 100644 --- a/backend/build_binary.py +++ b/backend/build_binary.py @@ -4,9 +4,15 @@ PyInstaller build script for creating standalone Python server binary. import PyInstaller.__main__ import os +import platform 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(): """Build Python server as standalone binary.""" backend_dir = Path(__file__).parent @@ -24,7 +30,7 @@ def build_server(): args.extend(['--paths', str(qwen_tts_path)]) print(f"Using local qwen_tts source from: {qwen_tts_path}") - # Add hidden imports + # Add common hidden imports args.extend([ '--hidden-import', 'backend', '--hidden-import', 'backend.main', @@ -35,6 +41,9 @@ def build_server(): '--hidden-import', 'backend.history', '--hidden-import', 'backend.tts', '--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.cache', '--hidden-import', 'backend.utils.progress', @@ -59,6 +68,26 @@ def build_server(): # Fix for pkg_resources and jaraco namespace packages '--hidden-import', 'pkg_resources.extern', '--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', '--clean', ]) diff --git a/backend/main.py b/backend/main.py index ca53f348..d34bad46 100644 --- a/backend/main.py +++ b/backend/main.py @@ -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 .utils.progress import get_progress_manager from .utils.tasks import get_task_manager +from .platform import get_backend_type app = FastAPI( title="voicebox API", @@ -73,6 +74,7 @@ async def health(): import os tts_model = tts.get_tts_model() + backend_type = get_backend_type() # Check for GPU availability (CUDA or MPS) has_cuda = torch.cuda.is_available() @@ -84,6 +86,8 @@ async def health(): gpu_type = f"CUDA ({torch.cuda.get_device_name(0)})" elif has_mps: gpu_type = "MPS (Apple Silicon)" + elif backend_type == "mlx": + gpu_type = "Metal (Apple Silicon via MLX)" vram_used = None if has_cuda: @@ -111,7 +115,11 @@ async def health(): model_downloaded = None try: # 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 try: @@ -130,7 +138,8 @@ async def health(): any(repo_cache.rglob("*.bin")) or any(repo_cache.rglob("*.safetensors")) 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 except Exception: @@ -144,6 +153,7 @@ async def health(): gpu_available=gpu_available, gpu_type=gpu_type, vram_used_mb=vram_used, + backend_type=backend_type, ) @@ -1149,6 +1159,8 @@ async def get_model_status(): from pathlib import Path import os + backend_type = get_backend_type() + # Try to import scan_cache_dir (might not be available in older versions) try: 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.""" try: 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: return False @@ -1168,50 +1180,66 @@ async def get_model_status(): """Check if Whisper model is loaded with specific size.""" try: 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: 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_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", "check_loaded": lambda: check_tts_loaded("1.7B"), }, { "model_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", "check_loaded": lambda: check_tts_loaded("0.6B"), }, { "model_name": "whisper-base", "display_name": "Whisper Base", - "hf_repo_id": "openai/whisper-base", + "hf_repo_id": whisper_base_id, "model_size": "base", "check_loaded": lambda: check_whisper_loaded("base"), }, { "model_name": "whisper-small", "display_name": "Whisper Small", - "hf_repo_id": "openai/whisper-small", + "hf_repo_id": whisper_small_id, "model_size": "small", "check_loaded": lambda: check_whisper_loaded("small"), }, { "model_name": "whisper-medium", "display_name": "Whisper Medium", - "hf_repo_id": "openai/whisper-medium", + "hf_repo_id": whisper_medium_id, "model_size": "medium", "check_loaded": lambda: check_whisper_loaded("medium"), }, { "model_name": "whisper-large", "display_name": "Whisper Large", - "hf_repo_id": "openai/whisper-large", + "hf_repo_id": whisper_large_id, "model_size": "large", "check_loaded": lambda: check_whisper_loaded("large"), }, @@ -1256,11 +1284,13 @@ async def get_model_status(): if repo_cache.exists(): # Check for model files (bin, safetensors, or other common model files) + # MLX models may use .npz or .safetensors has_model_files = ( any(repo_cache.rglob("*.bin")) or any(repo_cache.rglob("*.safetensors")) or any(repo_cache.rglob("*.pt")) 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("pytorch_model.bin.index.json")) ) @@ -1533,10 +1563,13 @@ async def get_active_tasks(): def _get_gpu_status() -> str: """Get GPU availability status.""" + backend_type = get_backend_type() if torch.cuda.is_available(): return f"CUDA ({torch.cuda.get_device_name(0)})" elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): return "MPS (Apple Silicon)" + elif backend_type == "mlx": + return "Metal (Apple Silicon via MLX)" return "None (CPU only)" @@ -1546,6 +1579,8 @@ async def startup_event(): print("voicebox API starting up...") database.init_db() 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()}") # Initialize progress manager with main event loop for thread-safe operations diff --git a/backend/models.py b/backend/models.py index 1951b887..7e72cd13 100644 --- a/backend/models.py +++ b/backend/models.py @@ -126,6 +126,7 @@ class HealthResponse(BaseModel): gpu_available: bool gpu_type: Optional[str] = None # GPU type (CUDA, MPS, or None) vram_used_mb: Optional[float] = None + backend_type: Optional[str] = None # Backend type (mlx or pytorch) class ModelStatus(BaseModel): diff --git a/backend/platform.py b/backend/platform.py new file mode 100644 index 00000000..c4db19dc --- /dev/null +++ b/backend/platform.py @@ -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" diff --git a/backend/requirements-mlx.txt b/backend/requirements-mlx.txt new file mode 100644 index 00000000..80ab1730 --- /dev/null +++ b/backend/requirements-mlx.txt @@ -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 diff --git a/backend/transcribe.py b/backend/transcribe.py index 6d4ac060..4d735fce 100644 --- a/backend/transcribe.py +++ b/backend/transcribe.py @@ -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 -import asyncio -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 +from typing import Optional +from .backends import get_stt_backend, STTBackend -class WhisperModel: - """Manages Whisper model loading and transcription.""" +def get_whisper_model() -> STTBackend: + """ + Get STT backend instance (MLX or PyTorch based on platform). - 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 - - 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 + Returns: + STT backend instance + """ + return get_stt_backend() def unload_whisper_model(): """Unload Whisper model to free memory.""" - global _whisper_model - if _whisper_model is not None: - _whisper_model.unload_model() + backend = get_stt_backend() + backend.unload_model() diff --git a/backend/tts.py b/backend/tts.py index f2e9783e..98db3412 100644 --- a/backend/tts.py +++ b/backend/tts.py @@ -1,372 +1,37 @@ """ -TTS inference module using Qwen3-TTS. +TTS inference module - delegates to backend abstraction layer. """ -from typing import Optional, List, Tuple -import asyncio -import torch +from typing import Optional import numpy as np import io import soundfile as sf -from pathlib import Path -from .utils.cache import get_cache_key, get_cached_voice_prompt, cache_voice_prompt -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 +from .backends import get_tts_backend, TTSBackend -class TTSModel: - """Manages Qwen3-TTS model loading and inference.""" +def get_tts_model() -> TTSBackend: + """ + Get TTS backend instance (MLX or PyTorch based on platform). - 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 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 + Returns: + TTS backend instance + """ + return get_tts_backend() def unload_tts_model(): """Unload TTS model to free memory.""" - global _tts_model - if _tts_model is not None: - _tts_model.unload_model() + backend = get_tts_backend() + backend.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: diff --git a/backend/utils/cache.py b/backend/utils/cache.py index cc2f3bd6..a070781f 100644 --- a/backend/utils/cache.py +++ b/backend/utils/cache.py @@ -5,7 +5,7 @@ Voice prompt caching utilities. import hashlib import torch from pathlib import Path -from typing import Optional +from typing import Optional, Union, Dict, Any from .. import config @@ -15,8 +15,8 @@ def _get_cache_dir() -> Path: return config.get_cache_dir() -# In-memory cache -_memory_cache: dict[str, torch.Tensor] = {} +# In-memory cache - can store dict (voice prompt) or tensor (legacy) +_memory_cache: dict[str, Union[torch.Tensor, Dict[str, Any]]] = {} 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( cache_key: str, -) -> Optional[torch.Tensor]: +) -> Optional[Union[torch.Tensor, Dict[str, Any]]]: """ Get cached voice prompt if available. @@ -51,7 +51,7 @@ def get_cached_voice_prompt( cache_key: Cache key Returns: - Cached voice prompt tensor or None + Cached voice prompt (dict or tensor) or None """ # Check in-memory cache if cache_key in _memory_cache: @@ -73,18 +73,18 @@ def get_cached_voice_prompt( def cache_voice_prompt( cache_key: str, - voice_prompt: torch.Tensor, + voice_prompt: Union[torch.Tensor, Dict[str, Any]], ) -> None: """ Cache voice prompt to memory and disk. Args: cache_key: Cache key - voice_prompt: Voice prompt tensor + voice_prompt: Voice prompt (dict or tensor) """ # Store in memory _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" torch.save(voice_prompt, cache_file) diff --git a/tauri/src-tauri/gen/Assets.car b/tauri/src-tauri/gen/Assets.car index 0c06aa52..7e4e205b 100644 Binary files a/tauri/src-tauri/gen/Assets.car and b/tauri/src-tauri/gen/Assets.car differ