app.py

"""
IndicConformer STT API for Hugging Face Spaces
"""

from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.responses import JSONResponse
from transformers import AutoModel
import torch
import librosa
import io
import time
import numpy as np
import asyncio
from concurrent.futures import ThreadPoolExecutor
import os
from huggingface_hub import login

# Authenticate with Hugging Face
hf_token = os.getenv("HF_TOKEN")
if hf_token:
    login(token=hf_token)
    print("✓ Authenticated with Hugging Face")
else:
    print("⚠ Warning: HF_TOKEN not found. Model loading may fail for gated repos.")

# Initialize FastAPI app
app = FastAPI(
    title="IndicConformer STT API",
    description="Speech-to-Text API for 22 Indian languages",
    version="1.0"
)

# Global variables
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MODEL = None
INFERENCE_EXECUTOR = ThreadPoolExecutor(max_workers=10)

# Audio chunking configuration
CHUNK_DURATION = 30
OVERLAP_DURATION = 2

# Supported languages
SUPPORTED_LANGUAGES = [
    "as", "bn", "brx", "doi", "gu", "hi", "kn", "kok",
    "ks", "mai", "ml", "mni", "mr", "ne", "or", "pa",
    "sa", "sat", "sd", "ta", "te", "ur"
]


@app.on_event("startup")
async def load_model():
    """Load model on startup"""
    global MODEL
    
    print("Loading IndicConformer model...")
    MODEL = AutoModel.from_pretrained(
        "ai4bharat/indic-conformer-600m-multilingual",
        trust_remote_code=True
    )
    MODEL = MODEL.to(DEVICE)
    
    # Warm-up the model
    print("Warming up model...")
    dummy_audio = torch.randn(1, 16000).to(DEVICE)
    _ = MODEL(dummy_audio, "hi", "rnnt")
    print(f"Model loaded successfully on {DEVICE}")


def split_audio_into_chunks(wav_np, sample_rate=16000, chunk_duration=30, overlap_duration=2):
    """Split audio into overlapping chunks"""
    chunk_samples = int(chunk_duration * sample_rate)
    overlap_samples = int(overlap_duration * sample_rate)
    step_samples = chunk_samples - overlap_samples

    chunks = []
    total_samples = len(wav_np)

    start = 0
    while start < total_samples:
        end = min(start + chunk_samples, total_samples)
        chunk = wav_np[start:end]

        chunks.append({
            'audio': chunk,
            'start_time': start / sample_rate,
            'end_time': end / sample_rate
        })

        if end >= total_samples:
            break

        start += step_samples

    return chunks


def merge_transcriptions_smart(transcriptions, max_overlap_words=10):
    """Merge chunk transcriptions with smart overlap removal"""
    if not transcriptions:
        return ""

    if len(transcriptions) == 1:
        return transcriptions[0].strip()

    result = transcriptions[0].strip()

    for i in range(1, len(transcriptions)):
        current = transcriptions[i].strip()

        if not current:
            continue

        result_words = result.split()
        current_words = current.split()

        max_check = min(len(result_words), len(current_words), max_overlap_words)

        best_overlap = 0
        for overlap_size in range(max_check, 0, -1):
            if result_words[-overlap_size:] == current_words[:overlap_size]:
                best_overlap = overlap_size
                break

        if best_overlap > 0:
            result += " " + " ".join(current_words[best_overlap:])
        else:
            result += " " + current

    return result


def run_inference(wav, language):
    """Run model inference"""
    if DEVICE == 'cuda':
        torch.cuda.synchronize()

    transcription = MODEL(wav, language, "rnnt")

    if DEVICE == 'cuda':
        torch.cuda.synchronize()

    return transcription


async def process_chunk(chunk_data, language, loop):
    """Process a single audio chunk"""
    wav_chunk = torch.tensor(chunk_data['audio']).unsqueeze(0)

    if DEVICE == 'cuda':
        wav_chunk = wav_chunk.to(DEVICE)

    transcription = await loop.run_in_executor(
        INFERENCE_EXECUTOR,
        run_inference,
        wav_chunk,
        language
    )

    return transcription


@app.get("/")
async def root():
    """Root endpoint with API information"""
    return {
        "message": "IndicConformer STT API",
        "version": "1.0",
        "model": "ai4bharat/indic-conformer-600m-multilingual",
        "decoder": "RNNT",
        "parallel_workers": 10,
        "chunk_processing": True,
        "chunk_duration": CHUNK_DURATION,
        "overlap_duration": OVERLAP_DURATION,
        "max_audio_duration": "30 minutes",
        "supported_languages": SUPPORTED_LANGUAGES,
        "device": DEVICE,
        "endpoints": {
            "transcribe": "/transcribe",
            "health": "/health",
            "docs": "/docs"
        }
    }


@app.get("/health")
async def health():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "model_loaded": MODEL is not None,
        "device": DEVICE,
        "parallel_enabled": True,
        "max_workers": 10
    }


@app.post("/transcribe")
async def transcribe_audio(
    file: UploadFile = File(...),
    language: str = Form(default="hi")
):
    """
    Transcribe audio file (supports up to 30 minutes)

    Parameters:
    - file: Audio file (WAV, MP3, FLAC, M4A)
    - language: Language code (hi=Hindi, te=Telugu, bn=Bengali, etc.)
    
    Returns:
    - transcription: Transcribed text
    - metadata: Processing information
    """
    
    try:
        # Validate file format
        if not file.filename.endswith(('.wav', '.mp3', '.flac', '.m4a')):
            raise HTTPException(
                status_code=400,
                detail="Invalid file format. Supported: WAV, MP3, FLAC, M4A"
            )

        # Validate language
        if language not in SUPPORTED_LANGUAGES:
            raise HTTPException(
                status_code=400,
                detail=f"Unsupported language: {language}. Supported: {', '.join(SUPPORTED_LANGUAGES)}"
            )

        # Read and process audio
        audio_bytes = await file.read()
        wav_np, sr = librosa.load(io.BytesIO(audio_bytes), sr=16000, mono=True)
        audio_duration = len(wav_np) / 16000

        print(f"Processing audio: {audio_duration:.2f}s ({audio_duration/60:.1f} minutes)")

        # Check duration limit
        if audio_duration > 1800:  # 30 minutes
            raise HTTPException(
                status_code=400,
                detail=f"Audio too long: {audio_duration/60:.1f} minutes. Maximum: 30 minutes"
            )

        # Split audio into chunks
        chunks = split_audio_into_chunks(
            wav_np,
            sample_rate=16000,
            chunk_duration=CHUNK_DURATION,
            overlap_duration=OVERLAP_DURATION
        )

        print(f"Split into {len(chunks)} chunks")

        # Process chunks in parallel
        start_time = time.time()

        loop = asyncio.get_event_loop()
        tasks = [process_chunk(chunk, language, loop) for chunk in chunks]
        chunk_transcriptions = await asyncio.gather(*tasks)

        inference_time = time.time() - start_time
        rtf = inference_time / audio_duration

        # Merge transcriptions
        full_transcription = merge_transcriptions_smart(chunk_transcriptions)

        print(f"Completed in {inference_time:.2f}s (RTF: {rtf:.4f})")

        return JSONResponse({
            "success": True,
            "transcription": full_transcription,
            "metadata": {
                "audio_duration": round(audio_duration, 2),
                "audio_duration_minutes": round(audio_duration / 60, 2),
                "inference_time": round(inference_time, 4),
                "rtf": round(rtf, 4),
                "language": language,
                "decoder": "rnnt",
                "num_chunks": len(chunks)
            }
        })

    except HTTPException:
        raise
    except Exception as e:
        print(f"Error: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Transcription failed: {str(e)}"
        )


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)