app.py

"""
KES Transcription API — HuggingFace Space (Free CPU tier).

Exposes a POST /transcribe endpoint that accepts audio via URL or base64,
runs Whisper + optional pyannote diarization, and returns the rich JSON
shape expected by the transcription microservice.
"""

import base64
import os
import tempfile
import warnings
from typing import Any, Dict, List, Optional

import requests
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

warnings.filterwarnings("ignore")

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
WHISPER_MODEL = os.environ.get("WHISPER_MODEL", "small")
HF_TOKEN = os.environ.get("HF_TOKEN", "")
ENABLE_DIARIZATION = os.environ.get("ENABLE_DIARIZATION", "false").lower() == "true"

# ---------------------------------------------------------------------------
# Global model references (loaded once at startup)
# ---------------------------------------------------------------------------
whisper_model = None
diarization_pipeline = None


def load_whisper():
    global whisper_model
    from faster_whisper import WhisperModel

    import torch

    device = "cuda" if torch.cuda.is_available() else "cpu"
    compute_type = "float16" if device == "cuda" else "int8"
    print(f"Loading Whisper '{WHISPER_MODEL}' on {device} ({compute_type})...")
    whisper_model = WhisperModel(WHISPER_MODEL, device=device, compute_type=compute_type)
    print("Whisper loaded.")


def patch_torchaudio():
    """Patch torchaudio for compatibility with pyannote 3.x on newer PyTorch."""
    import torchaudio

    if not hasattr(torchaudio, "AudioMetaData"):
        from dataclasses import dataclass

        @dataclass
        class AudioMetaData:
            sample_rate: int = 0
            num_frames: int = 0
            num_channels: int = 0
            bits_per_sample: int = 0
            encoding: str = ""

        torchaudio.AudioMetaData = AudioMetaData

    if not hasattr(torchaudio, "info"):
        import soundfile as sf

        def _info(filepath):
            info = sf.info(filepath)
            return torchaudio.AudioMetaData(
                sample_rate=info.samplerate,
                num_frames=info.frames,
                num_channels=info.channels,
                bits_per_sample=16,
                encoding=info.subtype or "PCM_16",
            )

        torchaudio.info = _info

    if not hasattr(torchaudio, "list_audio_backends"):
        torchaudio.list_audio_backends = lambda: ["soundfile"]


def load_diarization():
    global diarization_pipeline
    if not ENABLE_DIARIZATION:
        print("Diarization disabled (set ENABLE_DIARIZATION=true to enable).")
        return
    if not HF_TOKEN:
        print("Warning: HF_TOKEN not set, diarization disabled.")
        return

    try:
        patch_torchaudio()

        import torch
        from pyannote.audio import Pipeline

        os.environ["HF_TOKEN"] = HF_TOKEN
        print("Loading pyannote diarization pipeline...")
        diarization_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1")
        print("Diarization loaded.")
    except Exception as exc:
        print(f"Warning: diarization disabled: {exc}")
        diarization_pipeline = None



# ---------------------------------------------------------------------------
# FastAPI app
# ---------------------------------------------------------------------------
app = FastAPI(title="KES Transcription API")


class TranscribeRequest(BaseModel):
    audio_url: Optional[str] = None
    audio_base64: Optional[str] = None
    filename: Optional[str] = None
    task: Optional[str] = None
    language: Optional[str] = None
    num_speakers: int = 2


class WordItem(BaseModel):
    text: str
    start: float
    end: float
    type: str = "word"
    speaker_id: str = "SPEAKER_00"
    confidence: float = 0.0


class TranscribeResponse(BaseModel):
    text: str
    language: str
    language_probability: float
    duration: float
    avg_confidence: float
    words: List[WordItem]
    speakers: List[str]


@app.on_event("startup")
def startup():
    load_whisper()
    load_diarization()


@app.post("/")
async def root_transcribe(body: dict):
    """Accept Inference Endpoint format for compatibility with the microservice.

    Expected body:
    {
      "inputs": {"audio_url": "..."} or {"audio_base64": "...", "filename": "..."},
      "parameters": {"task": "translate", "language": "en", "num_speakers": 2}
    }
    """
    inputs = body.get("inputs", body)
    parameters = body.get("parameters") or {}

    req = TranscribeRequest(
        audio_url=inputs.get("audio_url") if isinstance(inputs, dict) else None,
        audio_base64=inputs.get("audio_base64") if isinstance(inputs, dict) else None,
        filename=inputs.get("filename") if isinstance(inputs, dict) else None,
        task=parameters.get("task"),
        language=parameters.get("language"),
        num_speakers=int(parameters.get("num_speakers") or 2),
    )
    return transcribe(req)


@app.get("/health")
def health():
    return {
        "status": "ok",
        "whisper_model": WHISPER_MODEL,
        "diarization": diarization_pipeline is not None,
    }


@app.post("/transcribe", response_model=TranscribeResponse)
def transcribe(req: TranscribeRequest):
    if not req.audio_url and not req.audio_base64:
        raise HTTPException(status_code=400, detail="Provide audio_url or audio_base64")

    # Materialise audio to a temp file
    audio_path = materialise_audio(req)
    try:
        # Transcribe
        transcription = run_transcription(audio_path, req.language, req.task)

        # Diarize (if enabled)
        speaker_segments = run_diarization(audio_path, req.num_speakers)

        # Merge
        words = merge_words_with_speakers(transcription["words"], speaker_segments)
        speakers = sorted({w["speaker_id"] for w in words}) if words else ["SPEAKER_00"]

        avg_conf = 0.0
        confs = [w["confidence"] for w in words]
        if confs:
            avg_conf = round(sum(confs) / len(confs), 4)

        return TranscribeResponse(
            text=transcription["text"],
            language=transcription["language"],
            language_probability=transcription["language_probability"],
            duration=transcription["duration"],
            avg_confidence=avg_conf,
            words=[WordItem(**w) for w in words],
            speakers=speakers,
        )
    finally:
        try:
            os.unlink(audio_path)
        except OSError:
            pass


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def materialise_audio(req: TranscribeRequest) -> str:
    suffix = ".wav"
    if req.audio_url:
        resp = requests.get(req.audio_url, timeout=300)
        resp.raise_for_status()
        audio_bytes = resp.content
        clean = req.audio_url.split("?", 1)[0]
        ext = os.path.splitext(clean)[1]
        if ext:
            suffix = ext
    else:
        audio_bytes = base64.b64decode(req.audio_base64)
        if req.filename:
            ext = os.path.splitext(req.filename)[1]
            if ext:
                suffix = ext

    handle = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
    try:
        handle.write(audio_bytes)
        return handle.name
    finally:
        handle.close()


def run_transcription(
    audio_path: str, language: Optional[str], task: Optional[str]
) -> Dict[str, Any]:
    kwargs: Dict[str, Any] = {"word_timestamps": True, "vad_filter": True}
    if language:
        kwargs["language"] = language
    if task in {"transcribe", "translate"}:
        kwargs["task"] = task

    segments_gen, info = whisper_model.transcribe(audio_path, **kwargs)

    words: List[Dict[str, Any]] = []
    text_parts: List[str] = []
    for segment in segments_gen:
        text_parts.append(segment.text.strip())
        for word in segment.words or []:
            words.append(
                {
                    "text": word.word,
                    "start": round(word.start, 3),
                    "end": round(word.end, 3),
                    "type": "word",
                    "speaker_id": "SPEAKER_00",
                    "confidence": round(word.probability, 4),
                }
            )

    return {
        "text": " ".join(text_parts).strip(),
        "language": info.language,
        "language_probability": round(info.language_probability, 4),
        "duration": round(info.duration, 2),
        "words": words,
    }


def run_diarization(audio_path: str, num_speakers: int) -> Optional[List[Dict[str, Any]]]:
    if diarization_pipeline is None:
        return None
    try:
        import soundfile as sf
        import torch

        data, sample_rate = sf.read(audio_path)
        if data.ndim == 1:
            waveform = torch.from_numpy(data).float().unsqueeze(0)
        else:
            waveform = torch.from_numpy(data.T).float()

        diarization = diarization_pipeline(
            {"waveform": waveform, "sample_rate": sample_rate},
            num_speakers=num_speakers,
        )

        segments: List[Dict[str, Any]] = []
        for turn, _, speaker in diarization.itertracks(yield_label=True):
            segments.append(
                {
                    "start": round(turn.start, 3),
                    "end": round(turn.end, 3),
                    "speaker": speaker,
                }
            )
        return segments
    except Exception as exc:
        print(f"Warning: diarization failed: {exc}")
        return None


def merge_words_with_speakers(
    words: List[Dict[str, Any]],
    speaker_segments: Optional[List[Dict[str, Any]]],
) -> List[Dict[str, Any]]:
    if not speaker_segments:
        return words

    for word in words:
        midpoint = (word["start"] + word["end"]) / 2
        best_speaker = "SPEAKER_00"
        best_overlap = -1.0

        for seg in speaker_segments:
            if seg["start"] <= midpoint <= seg["end"]:
                overlap = min(word["end"], seg["end"]) - max(word["start"], seg["start"])
                if overlap > best_overlap:
                    best_overlap = overlap
                    best_speaker = seg["speaker"]

        word["speaker_id"] = best_speaker

    return words