import os import re import uuid import shutil import tempfile import logging import numpy as np import torch import librosa import soundfile as sf from fastapi import ( FastAPI, UploadFile, File, HTTPException, Header, BackgroundTasks, ) from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import FileResponse from pydantic import BaseModel, Field from huggingface_hub import hf_hub_download from transformers import Wav2Vec2BertProcessor, Wav2Vec2BertForCTC from TTS.utils.synthesizer import Synthesizer from TTS.tts.utils.text import cleaners logging.basicConfig(level=logging.INFO) logger = logging.getLogger("KhmerSpeechAPI") # Optional API key. If the API_KEY env var is set, every request to the # ASR/TTS endpoints must include the header: X-API-Key: # If it is not set, the API is open (fine for testing). API_KEY = os.environ.get("API_KEY", "").strip() OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "outputs") os.makedirs(OUTPUT_DIR, exist_ok=True) app = FastAPI( title="Khmer ASR and TTS API", version="1.1.0", ) app.add_middleware( CORSMiddleware, allow_origins=["*"], # tighten to your frontend origin in production allow_credentials=False, allow_methods=["*"], allow_headers=["*"], ) speech_models = {} def check_api_key(x_api_key: str = ""): if API_KEY and x_api_key != API_KEY: raise HTTPException(status_code=401, detail="Invalid or missing API key.") def khmer_cleaners(text): text = " ".join(text.split()) text = re.sub( r"[^\u1780-\u17FF\u19E0-\u19FF0-9\s.,!?;:\-()'\"]", "", text ) return text.strip() setattr(cleaners, "khmer_cleaners", khmer_cleaners) @app.on_event("startup") async def load_models(): device = "cuda" if torch.cuda.is_available() else "cpu" use_cuda = torch.cuda.is_available() ASR_REPO = "Prakmlis/w2v-bert-2.0-khmer-customize-data-final" TTS_REPO = "Prakmlis/khmer-vits-mms" TTS_CHECKPOINT = ( "vits_khmer/" "vits_khmer-July-14-2026_02+08AM-0000000/" "best_model.pth" ) TTS_CONFIG = ( "vits_khmer/" "vits_khmer-July-14-2026_02+08AM-0000000/" "config.json" ) logger.info(f"Using device: {device}") logger.info("Loading ASR model...") speech_models["asr_processor"] = Wav2Vec2BertProcessor.from_pretrained(ASR_REPO) speech_models["asr_model"] = Wav2Vec2BertForCTC.from_pretrained(ASR_REPO).to(device) speech_models["asr_model"].eval() speech_models["device"] = device logger.info("Downloading TTS model...") checkpoint_path = hf_hub_download(repo_id=TTS_REPO, filename=TTS_CHECKPOINT) config_path = hf_hub_download(repo_id=TTS_REPO, filename=TTS_CONFIG) logger.info("Loading TTS synthesizer...") synth = Synthesizer( tts_checkpoint=checkpoint_path, tts_config_path=config_path, use_cuda=use_cuda ) synth.tts_model.length_scale = 0.90 synth.tts_model.noise_scale = 0.3 synth.tts_model.inference_noise_scale = 0.3 synth.tts_model.noise_scale_dp = 0.55 synth.tts_model.inference_noise_scale_dp = 0.55 speech_models["tts_synth"] = synth logger.info("Models loaded successfully.") @app.get("/") def root(): return { "message": "Khmer ASR and TTS API is running", "asr_api": "/api/v1/speech/asr", "tts_api": "/api/v1/speech/tts", "docs": "/docs", "auth": "X-API-Key header required" if API_KEY else "open", } @app.get("/health") def health(): ready = "asr_model" in speech_models and "tts_synth" in speech_models return {"status": "ok" if ready else "loading", "models_loaded": ready} class ASRResponse(BaseModel): transcription: str # NOTE: plain `def` (not `async def`) so FastAPI runs these in a # threadpool and heavy inference does not block the event loop. @app.post("/api/v1/speech/asr", response_model=ASRResponse) def asr_endpoint(file: UploadFile = File(...), x_api_key: str = Header(default="")): check_api_key(x_api_key) if "asr_model" not in speech_models: raise HTTPException(status_code=503, detail="ASR model is not loaded yet. Try again shortly.") suffix = os.path.splitext(file.filename or "")[1] or ".wav" with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_audio: shutil.copyfileobj(file.file, temp_audio) temp_audio_path = temp_audio.name try: speech_array, _ = librosa.load(temp_audio_path, sr=16000) # Guard against very long uploads tying up the CPU (max 5 minutes) if len(speech_array) > 16000 * 300: raise HTTPException(status_code=413, detail="Audio too long (max 5 minutes).") processor = speech_models["asr_processor"] model = speech_models["asr_model"] device = speech_models["device"] inputs = processor( speech_array, sampling_rate=16000, return_tensors="pt", padding=True ) inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): logits = model(**inputs).logits predicted_ids = torch.argmax(logits, dim=-1) decoded_text = processor.batch_decode(predicted_ids)[0] return ASRResponse(transcription=decoded_text.strip()) except HTTPException: raise except Exception as e: logger.error(f"ASR error: {e}") raise HTTPException(status_code=500, detail=str(e)) finally: if os.path.exists(temp_audio_path): os.remove(temp_audio_path) class TTSRequest(BaseModel): text: str = Field(..., min_length=1, max_length=1000, example="សួស្តី") @app.post("/api/v1/speech/tts") def tts_endpoint( payload: TTSRequest, background_tasks: BackgroundTasks, x_api_key: str = Header(default=""), ): check_api_key(x_api_key) if "tts_synth" not in speech_models: raise HTTPException(status_code=503, detail="TTS model is not loaded yet. Try again shortly.") if not payload.text.strip(): raise HTTPException(status_code=400, detail="Text is empty.") try: synth = speech_models["tts_synth"] clean_text = payload.text.replace(" ", ", ") wav = synth.tts(clean_text) wav = np.array(wav, dtype=np.float32) if wav.size == 0: raise HTTPException(status_code=422, detail="Generated waveform is empty.") if np.isnan(wav).any(): raise HTTPException(status_code=422, detail="Generated waveform contains NaN.") peak = np.max(np.abs(wav)) if peak > 0: wav = wav / peak wav = wav * 0.9 wav = np.clip(wav, -1.0, 1.0) output_path = os.path.join(OUTPUT_DIR, f"tts_{uuid.uuid4()}.wav") sf.write( output_path, wav, synth.output_sample_rate, format="WAV" ) logger.info("TTS generated successfully") logger.info(f"Output: {output_path}") logger.info(f"Duration: {len(wav) / synth.output_sample_rate:.2f}s") # Delete the file after the response has been sent, # so the outputs directory doesn't grow forever. background_tasks.add_task(os.remove, output_path) return FileResponse( output_path, media_type="audio/wav", filename="tts_output.wav", ) except HTTPException: raise except Exception as e: logger.error(f"TTS error: {e}") raise HTTPException(status_code=500, detail=str(e))