main.py

import io
import os
import re
import unicodedata
import logging
import numpy as np
import torch
import soundfile as sf
import torchaudio

from fastapi import FastAPI, HTTPException, UploadFile, File, Form
from fastapi.responses import StreamingResponse
from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
from speechbrain.pretrained import EncoderClassifier
from denoiser import pretrained
from typing import Optional

from phonemizer import phonemize
from phonemizer.backend import EspeakBackend
import re

# ── Logging ───────────────────────────────────────────────────────────────────
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ── Paths — models loaded from filesystem via symlink ./models → /tmp/models ──
MODEL_DIR     = os.environ.get("MODEL_DIR",     "./models/urdu-speecht5-finetuned")
PROCESSOR_DIR = os.environ.get("PROCESSOR_DIR", "./models/urdu-tts-processor")
VOCODER_DIR   = os.environ.get("VOCODER_DIR",   "./models/speecht5-hifigan")
SPKREC_DIR    = os.environ.get("SPKREC_DIR",    "./models/spkrec-xvect-voxceleb")

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
logger.info(f"Using device: {DEVICE}")

# ── Standalone Urdu normalization (no urduhack / tensorflow dependency) ────────
_PUNCT_MAP = str.maketrans({
    "٪": "%", "٫": ".", "٬": ",",
    "\u060c": "،",
    "\u200f": "", "\u200e": "", "\u200b": "", "\ufeff": "",
})
_ARABIC_INDIC = str.maketrans("٠١٢٣٤٥٦٧٨٩", "0123456789")
_EXTENDED_AR  = str.maketrans("۰۱۲۳۴۵۶۷۸۹", "0123456789")

def normalize_urdu(text: str) -> str:
    text = unicodedata.normalize("NFC", text)
    text = text.translate(_PUNCT_MAP)
    text = text.translate(_ARABIC_INDIC)
    text = text.translate(_EXTENDED_AR)
    text = re.sub(r"\s+", " ", text).strip()
    return text


# ── Load all models at startup ────────────────────────────────────────────────
logger.info("Loading processor...")
processor = SpeechT5Processor.from_pretrained(PROCESSOR_DIR)

logger.info("Loading TTS model...")
tts_model = SpeechT5ForTextToSpeech.from_pretrained(MODEL_DIR).to(DEVICE)
tts_model.eval()

logger.info("Loading vocoder...")
vocoder = SpeechT5HifiGan.from_pretrained(VOCODER_DIR).to(DEVICE)
vocoder.eval()

logger.info("Loading speaker encoder...")
speaker_encoder = EncoderClassifier.from_hparams(
    source=SPKREC_DIR,
    run_opts={"device": DEVICE},
)

logger.info("Loading DNS64 denoiser...")
denoiser_model = pretrained.dns64().to(DEVICE)
denoiser_model.eval()

# ── Default speaker embedding — loaded from embeddings/normal_embedding.pt ────
_EMBEDDING_PATH = os.environ.get("SPEAKER_EMBEDDING_PATH", "./embeddings/normal_embedding.pt")
logger.info(f"Loading speaker embedding from {_EMBEDDING_PATH}...")
_emb = torch.load(_EMBEDDING_PATH, map_location="cpu")
if not isinstance(_emb, torch.Tensor):
    _emb = torch.tensor(_emb)
_emb = _emb.float().squeeze()
_emb = torch.nn.functional.normalize(_emb.unsqueeze(0), dim=1)  # (1, 512)
DEFAULT_SPEAKER_EMBEDDING = _emb.to(DEVICE)

# Initialize phonemizer for Urdu
backend = EspeakBackend(
    'ur',
    preserve_punctuation=False,   # punctuation has no sound, skip it
    with_stress=True,             # stress markers improve prosody
)

def urdu_to_phonemes(text):
    """Convert Urdu text to space-separated IPA phoneme string."""
    result = backend.phonemize([text], njobs=1)[0]
    # Clean up extra whitespace
    result = re.sub(r'\s+', ' ', result).strip()
    return result

logger.info("All models loaded ✅")

# ── FastAPI app ───────────────────────────────────────────────────────────────
app = FastAPI(
    title="Urdu TTS Phonemes API",
    description="Convert Urdu text to speech using a fine-tuned SpeechT5 model.",
    version="1.0.0",
)

TARGET_SR = 16000   # SpeechT5 and speaker encoder both expect 16 kHz


def embedding_from_audio(audio_bytes: bytes) -> torch.Tensor:
    """
    Compute a (1, 512) speaker embedding from raw audio file bytes.
    Accepts any format soundfile can read (wav, mp3, flac, ogg, etc.).
    Resamples to 16 kHz mono automatically.
    """
    buf = io.BytesIO(audio_bytes)

    # Read with torchaudio — supports more formats than soundfile
    waveform, sr = torchaudio.load(buf)          # (channels, T)

    # Mix down to mono
    if waveform.shape[0] > 1:
        waveform = waveform.mean(dim=0, keepdim=True)  # (1, T)

    # Resample to 16 kHz if needed
    if sr != TARGET_SR:
        resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=TARGET_SR)
        waveform = resampler(waveform)

    waveform = waveform.to(DEVICE)   # (1, T)

    with torch.no_grad():
        emb = speaker_encoder.encode_batch(waveform)          # (1, 1, 512)
        emb = torch.nn.functional.normalize(emb, dim=2)
        emb = emb.squeeze(0)                                   # (1, 512)

    return emb   # (1, 512) on DEVICE


def run_tts(text: str, speaker_embedding: torch.Tensor) -> np.ndarray:
    """Full pipeline: text → normalize → tokenize → TTS → vocoder → denoise."""
    normalized = normalize_urdu(text)
    phonemized = urdu_to_phonemes(normalized)
    logger.info(f"Phonemes: {phonemized}")

    inputs = processor(text=phonemized, return_tensors="pt")
    input_ids = inputs["input_ids"].to(DEVICE)

    with torch.no_grad():
        speech = tts_model.generate_speech(
            input_ids,
            speaker_embedding,
            vocoder=vocoder,
        )

    # DNS64 expects (batch, channels, time) at 16 kHz
    speech_3d = speech.unsqueeze(0).unsqueeze(0)
    with torch.no_grad():
        denoised = denoiser_model(speech_3d)[0]

    return denoised.squeeze().cpu().numpy()


@app.get("/health")
def health():
    return {"status": "ok"}


@app.post("/synthesize")
async def synthesize(
    text: str = Form(..., description="Urdu text to synthesize"),
    reference_audio: Optional[UploadFile] = File(
        default=None,
        description="Optional reference audio file for voice cloning. "
                    "If omitted, the default speaker voice is used. "
                    "Accepts WAV, MP3, FLAC, OGG etc.",
    ),
):
    """
    Convert Urdu text to speech.

    - Send as **multipart/form-data** with:
      - `text` (required): the Urdu string to synthesize
      - `reference_audio` (optional): an audio file whose voice will be cloned

    - If no `reference_audio` is provided, the pre-loaded default embedding is used.
    - Response is a denoised **audio/wav** file.

    **Example — default voice (curl):**
    ```
    curl -X POST /synthesize -F "text=یہ اردو ہے" --output out.wav
    ```

    **Example — custom voice (curl):**
    ```
    curl -X POST /synthesize -F "text=یہ اردو ہے" -F "reference_audio=@my_voice.wav" --output out.wav
    ```
    """
    if not text or not text.strip():
        raise HTTPException(status_code=400, detail="text must not be empty.")

    # ── Resolve speaker embedding ─────────────────────────────────────────────
    if reference_audio is not None:
        # Validate content type loosely — torchaudio will hard-fail on bad files
        allowed = {"audio/wav", "audio/wave", "audio/mpeg", "audio/mp3",
                   "audio/flac", "audio/ogg", "audio/x-wav", "audio/x-flac"}
        ct = (reference_audio.content_type or "").lower()
        if ct and ct not in allowed:
            raise HTTPException(
                status_code=415,
                detail=f"Unsupported audio type '{ct}'. Use WAV, MP3, FLAC, or OGG.",
            )

        audio_bytes = await reference_audio.read()
        if len(audio_bytes) == 0:
            raise HTTPException(status_code=400, detail="Uploaded audio file is empty.")

        try:
            logger.info(f"Computing embedding from uploaded file: {reference_audio.filename}")
            speaker_embedding = embedding_from_audio(audio_bytes)
        except Exception as e:
            logger.exception("Failed to compute speaker embedding from reference audio")
            raise HTTPException(
                status_code=422,
                detail=f"Could not process reference audio: {str(e)}",
            )
    else:
        logger.info("No reference audio provided — using default speaker embedding")
        speaker_embedding = DEFAULT_SPEAKER_EMBEDDING

    # ── Run TTS ───────────────────────────────────────────────────────────────
    try:
        audio_np = run_tts(text.strip(), speaker_embedding)
    except Exception as e:
        logger.exception("TTS generation failed")
        raise HTTPException(status_code=500, detail=f"TTS generation failed: {str(e)}")

    buffer = io.BytesIO()
    sf.write(buffer, audio_np, samplerate=16000, format="WAV")
    buffer.seek(0)

    return StreamingResponse(
        buffer,
        media_type="audio/wav",
        headers={"Content-Disposition": 'attachment; filename="output.wav"'},
    )