main.py

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: <value>
# 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))