Spaces:
Paused
Paused
| import os | |
| import io | |
| import wave | |
| import soundfile as sf | |
| import numpy as np | |
| import asyncio | |
| import sys | |
| import traceback | |
| from typing import AsyncIterable | |
| from livekit.agents import tts | |
| from livekit.agents.types import DEFAULT_API_CONNECT_OPTIONS | |
| import uuid | |
| from gradio_client import Client | |
| # Fallback language mapping: API voice -> MMS model for fallback | |
| FALLBACK_MMS = { | |
| "stano03/habari-xtts": "facebook/mms-tts-swh", | |
| "sta-03/habari-f5tts": "facebook/mms-tts-swh", | |
| "Stanley03/habari-yarngpt": "facebook/mms-tts-yor", | |
| "Stanley-07/habari-nursetoto": "facebook/mms-tts-swh", | |
| "Stanley-07/habari-kisii": "facebook/mms-tts-swh", | |
| } | |
| class ApiTTS(tts.TTS): | |
| def __init__(self, space_id: str, sample_rate: int = 24000, **kwargs): | |
| super().__init__( | |
| capabilities=tts.TTSCapabilities(streaming=False), | |
| sample_rate=sample_rate, | |
| num_channels=1, | |
| ) | |
| self.space_id = space_id | |
| self.kwargs = kwargs | |
| self.client = None | |
| self._fallback_tts = None # lazy MMS fallback | |
| self._api_failed = False # track if API is down | |
| print(f"[ApiTTS] Initializing for: {space_id}", flush=True) | |
| def _get_fallback(self): | |
| """Lazy-load an MMS-TTS fallback engine.""" | |
| if self._fallback_tts is None: | |
| from tts_mms import MmsTTS | |
| fallback_model = FALLBACK_MMS.get(self.space_id, "facebook/mms-tts-swh") | |
| print(f"[ApiTTS] Loading MMS fallback: {fallback_model}", flush=True) | |
| self._fallback_tts = MmsTTS(model_id=fallback_model) | |
| return self._fallback_tts | |
| def synthesize(self, text: str, *, conn_options=DEFAULT_API_CONNECT_OPTIONS) -> tts.ChunkedStream: | |
| # If we know the API is down, go straight to fallback | |
| if self._api_failed: | |
| print(f"[ApiTTS] API known down, using MMS fallback", flush=True) | |
| return self._get_fallback().synthesize(text, conn_options=conn_options) | |
| return _ApiChunkedStream(tts_engine=self, text=text, conn_options=conn_options) | |
| class _ApiChunkedStream(tts.ChunkedStream): | |
| def __init__(self, *, tts_engine: ApiTTS, text: str, conn_options=DEFAULT_API_CONNECT_OPTIONS): | |
| super().__init__(tts=tts_engine, input_text=text, conn_options=conn_options) | |
| self.text = text | |
| self.tts_engine = tts_engine | |
| async def _run(self, output_emitter: tts.AudioEmitter) -> None: | |
| try: | |
| # 1. Initialize client if not already done (with timeout) | |
| if self.tts_engine.client is None: | |
| print(f"[ApiTTS] Connecting to {self.tts_engine.space_id}...", flush=True) | |
| try: | |
| self.tts_engine.client = await asyncio.wait_for( | |
| asyncio.to_thread(Client, self.tts_engine.space_id), | |
| timeout=120.0 | |
| ) | |
| print(f"[ApiTTS] Connected OK", flush=True) | |
| except (asyncio.TimeoutError, Exception) as conn_err: | |
| print(f"[ApiTTS] Connection failed: {conn_err}", flush=True) | |
| self.tts_engine._api_failed = True | |
| # Use fallback | |
| await self._run_fallback(output_emitter) | |
| return | |
| # 2. Call the API with timeout and retry for "models still loading" | |
| print(f"[ApiTTS] Generating: '{self.text[:60]}'", flush=True) | |
| kwargs = dict(self.tts_engine.kwargs) | |
| kwargs["text"] = self.text | |
| result = None | |
| max_retries = 3 | |
| for attempt in range(max_retries): | |
| try: | |
| result = await asyncio.wait_for( | |
| asyncio.to_thread( | |
| self.tts_engine.client.predict, | |
| api_name="/generate", | |
| **kwargs | |
| ), | |
| timeout=180.0 | |
| ) | |
| break # Success | |
| except Exception as api_err: | |
| err_msg = str(api_err) | |
| if "still loading" in err_msg and attempt < max_retries - 1: | |
| wait_time = 15 * (attempt + 1) | |
| print(f"[ApiTTS] Models loading, retry {attempt+1}/{max_retries} in {wait_time}s...", flush=True) | |
| await asyncio.sleep(wait_time) | |
| continue | |
| print(f"[ApiTTS] API call failed: {api_err}", flush=True) | |
| self.tts_engine._api_failed = True | |
| self.tts_engine.client = None | |
| await self._run_fallback(output_emitter) | |
| return | |
| audio_path = result | |
| print(f"[ApiTTS] Got audio: {audio_path}", flush=True) | |
| # 3. Read and convert audio | |
| data, sample_rate = sf.read(audio_path, dtype='float32') | |
| if len(data.shape) > 1: | |
| data = data.mean(axis=1) | |
| pcm_int16 = (data * 32767).clip(-32768, 32767).astype(np.int16) | |
| pcm_bytes = pcm_int16.tobytes() | |
| print(f"[ApiTTS] PCM: {len(pcm_int16)} samples @ {sample_rate}Hz", flush=True) | |
| # 4. Push audio | |
| output_emitter.initialize( | |
| request_id=str(uuid.uuid4()), | |
| sample_rate=sample_rate, | |
| num_channels=1, | |
| mime_type="audio/pcm" | |
| ) | |
| output_emitter.push(pcm_bytes) | |
| output_emitter.flush() | |
| print(f"[ApiTTS] Audio pushed OK", flush=True) | |
| except Exception as e: | |
| print(f"[ApiTTS] Unexpected error: {e}", flush=True) | |
| traceback.print_exc() | |
| sys.stdout.flush() | |
| # Try fallback as last resort | |
| try: | |
| await self._run_fallback(output_emitter) | |
| except Exception as fb_err: | |
| print(f"[ApiTTS] Fallback also failed: {fb_err}", flush=True) | |
| raise | |
| async def _run_fallback(self, output_emitter: tts.AudioEmitter) -> None: | |
| """Use MMS-TTS fallback when API is unavailable.""" | |
| print(f"[ApiTTS] Using MMS fallback for: '{self.text[:60]}'", flush=True) | |
| fallback = self.tts_engine._get_fallback() | |
| # Create MMS stream and directly call its _run with our emitter | |
| from tts_mms import MmsStream | |
| fallback_stream = MmsStream(tts=fallback, input_text=self.text) | |
| await fallback_stream._run(output_emitter) | |