from __future__ import annotations import asyncio import io import logging import re import threading import wave from contextlib import asynccontextmanager from functools import lru_cache from typing import Any, Dict, List, Optional, Union import numpy as np import torch from fastapi import BackgroundTasks, FastAPI, HTTPException from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field, field_validator, model_validator from transformers import AutoTokenizer, VitsModel from langdetect import detect, DetectorFactory # ------------------------- Logging ------------------------- logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) DetectorFactory.seed = 0 # ------------------------- Config ------------------------- DEVICE = "cuda" if torch.cuda.is_available() else "cpu" SAMPLE_RATE = 16000 MAX_CHARS_PER_CHUNK = 900 SILENCE_BETWEEN_SENTENCES_SEC = 0.04 SILENCE_BETWEEN_PARAGRAPHS_SEC = 0.12 # MMS currently exposes one voice per language checkpoint. SUPPORTED_VOICES = {"default"} LANGDETECT_TO_MMS = { "en": "eng", "hi": "hin", "bn": "ben", "as": "asm", "ta": "tam", "te": "tel", "mr": "mar", "gu": "guj", "kn": "kan", "ml": "mal", "pa": "pan", "ur": "urd", "or": "ori", "ne": "nep", "si": "sin", "ru": "rus", "fr": "fra", "de": "deu", "es": "spa", "it": "ita", "pt": "por", "ja": "jpn", "ko": "kor", "zh-cn": "zho", "ar": "ara", "fa": "fas", "tr": "tur", "pl": "pol", "uk": "ukr", "vi": "vie", "th": "tha", "id": "ind", "ms": "msa", "sw": "swa", "am": "amh", "yo": "yor", "ig": "ibo", "ha": "hau", "sn": "sna", "zu": "zul", "so": "som", "km": "khm", "my": "mya", "lo": "lao", "mn": "mon", "kk": "kaz", "uz": "uzb", "tl": "tgl", "ceb": "ceb", "hmn": "hmn", "nv": "nav", "sm": "smo", "ty": "tah", "haw": "haw", "mt": "mlt", "et": "est", "lv": "lav", "lt": "lit", "sl": "slv", "hr": "hrv", "sr": "srp", "bs": "bos", "mk": "mkd", "bg": "bul", "ro": "ron", "hu": "hun", "fi": "fin", "sv": "swe", "no": "nor", "da": "dan", "is": "isl", "ga": "gle", "cy": "cym", "eu": "eus", "ca": "cat", "gl": "glg", "he": "heb", "hy": "hye", "ka": "kat", "sa": "san", } LANGUAGE_MAP: Dict[str, str] = { "asm": "Assamese", "ben": "Bengali", "eng": "English", "hin": "Hindi", "tam": "Tamil", "tel": "Telugu", "mar": "Marathi", "guj": "Gujarati", "kan": "Kannada", "mal": "Malayalam", "pan": "Punjabi", "urd": "Urdu", "ori": "Odia", "nep": "Nepali", "sin": "Sinhala", "rus": "Russian", "fra": "French", "deu": "German", "spa": "Spanish", "ita": "Italian", "por": "Portuguese", "jpn": "Japanese", "kor": "Korean", "zho": "Chinese", "ara": "Arabic", "fas": "Persian", "tur": "Turkish", "pol": "Polish", "ukr": "Ukrainian", "vie": "Vietnamese", "tha": "Thai", "ind": "Indonesian", "msa": "Malay", "swa": "Swahili", "amh": "Amharic", "yor": "Yoruba", "ibo": "Igbo", "hau": "Hausa", "sna": "Shona", "zul": "Zulu", "som": "Somali", "khm": "Khmer", "mya": "Burmese", "lao": "Lao", "mon": "Mongolian", "kaz": "Kazakh", "uzb": "Uzbek", "tgl": "Tagalog", "ceb": "Cebuano", "hmn": "Hmong", "nav": "Navajo", "smo": "Samoan", "tah": "Tahitian", "haw": "Hawaiian", "mlt": "Maltese", "est": "Estonian", "lav": "Latvian", "lit": "Lithuanian", "slv": "Slovenian", "hrv": "Croatian", "srp": "Serbian", "bos": "Bosnian", "mkd": "Macedonian", "bul": "Bulgarian", "ron": "Romanian", "hun": "Hungarian", "fin": "Finnish", "swe": "Swedish", "nor": "Norwegian", "dan": "Danish", "isl": "Icelandic", "gle": "Irish", "cym": "Welsh", "eus": "Basque", "cat": "Catalan", "glg": "Galician", "heb": "Hebrew", "hye": "Armenian", "kat": "Georgian", "san": "Sanskrit", } # ------------------------- Thread-safe caches ------------------------- _model_cache: Dict[str, VitsModel] = {} _tokenizer_cache: Dict[str, Any] = {} _cache_lock = threading.Lock() @lru_cache(maxsize=128) def _build_model_id(language_code: str) -> str: return f"facebook/mms-tts-{language_code}" def get_model_and_tokenizer(language_code: str): with _cache_lock: if language_code in _model_cache: return _model_cache[language_code], _tokenizer_cache[language_code] model_id = _build_model_id(language_code) logger.info("Loading model: %s", model_id) try: model = VitsModel.from_pretrained(model_id) tokenizer = AutoTokenizer.from_pretrained(model_id) model.to(DEVICE) model.eval() except Exception as e: logger.exception("Failed to load %s", model_id) raise HTTPException( status_code=404, detail=f"Language code '{language_code}' is not available in MMS TTS." ) from e _model_cache[language_code] = model _tokenizer_cache[language_code] = tokenizer return model, tokenizer def detect_language(text: str) -> str: try: lang_code = detect(text) return LANGDETECT_TO_MMS.get(lang_code, lang_code) except Exception: return "eng" def normalize_text(text: str) -> str: text = text.replace("\r\n", "\n").replace("\r", "\n") text = re.sub(r"[ \t]+", " ", text) text = re.sub(r"\n{3,}", "\n\n", text) return text.strip() def split_paragraphs(text: str) -> List[str]: text = normalize_text(text) parts = [p.strip() for p in re.split(r"\n\s*\n+", text) if p.strip()] return parts if parts else [text] def split_sentences(paragraph: str) -> List[str]: # English punctuation + Hindi danda. parts = re.split(r"(?<=[.!?。！？।॥])\s+", paragraph.strip()) return [p.strip() for p in parts if p.strip()] def chunk_text(text: str, max_chars: int = MAX_CHARS_PER_CHUNK) -> List[str]: paragraphs = split_paragraphs(text) chunks: List[str] = [] current = "" def flush(): nonlocal current if current.strip(): chunks.append(current.strip()) current = "" for paragraph in paragraphs: sentences = split_sentences(paragraph) or [paragraph] if current: flush() for sent in sentences: if len(sent) <= max_chars: if not current: current = sent elif len(current) + 1 + len(sent) <= max_chars: current += " " + sent else: flush() current = sent else: words = sent.split() temp = "" for word in words: if not temp: temp = word elif len(temp) + 1 + len(word) <= max_chars: temp += " " + word else: if temp.strip(): chunks.append(temp.strip()) temp = word if temp.strip(): if current and len(current) + 1 + len(temp) <= max_chars: current += " " + temp else: if current: flush() current = temp flush() return chunks def sanitize_chunks(chunks: List[str]) -> List[str]: return [normalize_text(c) for c in chunks if normalize_text(c)] def resolve_language( request_language: Optional[str], auto_detect: bool, text: Union[str, List[str]], ) -> tuple[str, Optional[str]]: if not auto_detect: if not request_language: raise HTTPException(status_code=400, detail="Provide 'language' or set 'auto_detect=true'.") return request_language, None sample = "" if isinstance(text, str): sample = next((p for p in split_paragraphs(text) if p.strip()), "") else: sample = next((p for p in text if isinstance(p, str) and p.strip()), "") if not sample: raise HTTPException(status_code=400, detail="Cannot auto-detect language from empty text.") detected = detect_language(sample) return detected, detected def synthesize_single_chunk(text: str, language_code: str) -> np.ndarray: model, tokenizer = get_model_and_tokenizer(language_code) inputs = tokenizer(text, return_tensors="pt") inputs = {k: v.to(DEVICE) for k, v in inputs.items()} with torch.inference_mode(): outputs = model(**inputs) waveform = outputs.waveform[0].detach().cpu().numpy() return waveform.astype(np.float32) def merge_audio_chunks(chunk_texts: List[str], language_code: str) -> np.ndarray: segments: List[np.ndarray] = [] for idx, txt in enumerate(chunk_texts): audio = synthesize_single_chunk(txt, language_code) segments.append(audio) if idx < len(chunk_texts) - 1: silence_sec = SILENCE_BETWEEN_SENTENCES_SEC segments.append(np.zeros(int(SAMPLE_RATE * silence_sec), dtype=np.float32)) return np.concatenate(segments) if segments else np.zeros(0, dtype=np.float32) def wav_bytes_from_float32(audio: np.ndarray) -> bytes: audio_int16 = (audio * 32767.0).clip(-32768, 32767).astype(np.int16) buffer = io.BytesIO() with wave.open(buffer, "wb") as wav_file: wav_file.setnchannels(1) wav_file.setsampwidth(2) wav_file.setframerate(SAMPLE_RATE) wav_file.writeframes(audio_int16.tobytes()) return buffer.getvalue() # ------------------------- Pydantic models ------------------------- class TTSRequest(BaseModel): text: Union[str, List[str]] = Field( ..., description="Text to speak. A single string is auto-chunked; list items are treated as separate blocks." ) language: Optional[str] = Field( None, description="ISO 639-3 language code, for example 'hin', 'eng', 'ben'." ) auto_detect: bool = Field( False, description="Detect language from the first non-empty chunk." ) voice: str = Field( "default", description="Compatibility field. MMS currently exposes only the 'default' voice." ) speed: float = Field(1.0, ge=0.5, le=2.0) @field_validator("text") @classmethod def validate_text(cls, v): if isinstance(v, str) and not v.strip(): raise ValueError("Text cannot be empty.") if isinstance(v, list): cleaned = [x for x in v if isinstance(x, str) and x.strip()] if not cleaned: raise ValueError("Text array cannot be empty.") return v @model_validator(mode="after") def validate_voice(self): if self.voice not in SUPPORTED_VOICES: raise ValueError( f"Unsupported voice '{self.voice}'. Available voices: {sorted(SUPPORTED_VOICES)}" ) return self class TTSResponse(BaseModel): audio_format: str = "audio/wav" language_detected: Optional[str] = None language_used: str num_chunks: int duration_seconds: float chunks: List[str] = Field(default_factory=list) # ------------------------- FastAPI ------------------------- @asynccontextmanager async def lifespan(app: FastAPI): logger.info("Starting MMS-TTS API on %s", DEVICE) yield if DEVICE == "cuda": torch.cuda.empty_cache() app = FastAPI( title="Enhanced MMS-TTS API", description="Production-grade Text-to-Speech for 1,100+ languages using Facebook's MMS models.", version="2.1.0", lifespan=lifespan, ) @app.get("/health", tags=["System"]) async def health_check() -> Dict[str, str]: return {"status": "ok", "device": DEVICE, "sample_rate": str(SAMPLE_RATE)} @app.get("/languages", tags=["Metadata"]) async def get_languages() -> Dict[str, Any]: languages = [{"code": code, "name": name} for code, name in sorted(LANGUAGE_MAP.items())] return { "total_languages": len(languages), "languages": languages, "voices": sorted(SUPPORTED_VOICES), } @app.post("/tts", tags=["Synthesis"]) async def text_to_speech(request: TTSRequest, background_tasks: BackgroundTasks) -> StreamingResponse: try: language_code, detected_lang = resolve_language(request.language, request.auto_detect, request.text) if isinstance(request.text, str): raw_chunks = chunk_text(request.text) else: raw_chunks = [] for part in request.text: raw_chunks.extend(chunk_text(part)) chunks = sanitize_chunks(raw_chunks) if not chunks: raise HTTPException(status_code=400, detail="No valid chunks after processing.") audio_np = await asyncio.to_thread(merge_audio_chunks, chunks, language_code) if request.speed != 1.0 and len(audio_np) > 1: new_len = max(1, int(len(audio_np) / request.speed)) x_old = np.linspace(0, 1, num=len(audio_np), endpoint=True) x_new = np.linspace(0, 1, num=new_len, endpoint=True) audio_np = np.interp(x_new, x_old, audio_np).astype(np.float32) wav_data = wav_bytes_from_float32(audio_np) duration = float(len(audio_np) / SAMPLE_RATE) headers = { "Content-Disposition": "attachment; filename=speech.wav", "X-Language-Used": language_code, "X-Detected-Language": detected_lang or "", "X-Num-Chunks": str(len(chunks)), "X-Duration-Seconds": f"{duration:.3f}", "X-Voice-Used": request.voice, } return StreamingResponse( io.BytesIO(wav_data), media_type="audio/wav", headers=headers, background=background_tasks, ) except HTTPException: raise except Exception as e: logger.exception("Synthesis error") raise HTTPException(status_code=500, detail=f"Speech synthesis failed: {e}") @app.post("/tts/info", tags=["Synthesis"], response_model=TTSResponse) async def tts_info(request: TTSRequest) -> TTSResponse: language_code, detected_lang = resolve_language(request.language, request.auto_detect, request.text) if isinstance(request.text, str): raw_chunks = chunk_text(request.text) else: raw_chunks = [] for part in request.text: raw_chunks.extend(chunk_text(part)) chunks = sanitize_chunks(raw_chunks) return TTSResponse( language_detected=detected_lang, language_used=language_code, num_chunks=len(chunks), duration_seconds=0.0, chunks=chunks, )