app.py

import os
import tempfile
import time
import asyncio
from typing import List, Dict, Any, Optional
from concurrent.futures import ThreadPoolExecutor

import torch
from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn
from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq
import librosa
import numpy as np
from fastapi.responses import JSONResponse
import gc

# Initialize thread pool for background processing
thread_pool = ThreadPoolExecutor(max_workers=2)

# Environment and model configuration
MODEL_NAME = "nyrahealth/CrisperWhisper"
BATCH_SIZE = 8
FILE_LIMIT_MB = 30
FILE_EXTENSIONS = [".mp3", ".wav", ".m4a", ".ogg", ".flac"]

# Initialize FastAPI app
app = FastAPI(
    title="Speech to Text API",
    description="API for transcribing audio files using the CrisperWhisper model",
    version="1.0.0"
)

# Add CORS support
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Response models
class TranscriptionChunk(BaseModel):
    timestamp: List[float]
    text: str

class TranscriptionResponse(BaseModel):
    text: str
    chunks: List[TranscriptionChunk]

# Setup device and load model
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# Load model and processor at startup
@app.on_event("startup")
async def load_model():
    global processor, model
    print("Loading model and processor...")
    processor = AutoProcessor.from_pretrained(MODEL_NAME)
    model = AutoModelForSpeechSeq2Seq.from_pretrained(MODEL_NAME)
    model.to(device)
    print("Model loaded successfully!")

def load_audio(file_path: str) -> tuple:
    """Load audio file efficiently"""
    try:
        # Use a faster sr=None first to get the original sampling rate,
        # then resample only if needed
        audio_array, orig_sr = librosa.load(file_path, sr=None, mono=True)
        
        # Resample only if needed
        if orig_sr != 16000:
            audio_array = librosa.resample(audio_array, orig_sr=orig_sr, target_sr=16000)
            sampling_rate = 16000
        else:
            sampling_rate = orig_sr
            
        # Convert to float32 if needed
        if audio_array.dtype != np.float32:
            audio_array = audio_array.astype(np.float32)
            
        return audio_array, sampling_rate
        
    except Exception as e:
        print(f"Error loading audio: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error loading audio: {str(e)}")

def process_audio_file(file_path: str) -> Dict:
    """Process audio file and return transcription with timestamps"""
    try:
        # Load audio file efficiently
        audio_array, sampling_rate = load_audio(file_path)
        
        # Process with model
        inputs = processor(audio_array, sampling_rate=sampling_rate, return_tensors="pt")
        inputs = {key: value.to(device) for key, value in inputs.items()}
        
        # Generate transcription with word timestamps
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                return_timestamps=True,
                return_dict_in_generate=True,
                output_scores=True,
                max_new_tokens=256 if len(audio_array) < 160000 else 512,  # Adjust based on audio length
                num_beams=1,  # Use greedy decoding for speed
            )
        
        # Extract timestamps and words
        result = processor.decode(outputs.sequences[0], skip_special_tokens=False, output_word_offsets=True)
        words_with_timestamps = []
        
        for word in result.word_offsets:
            words_with_timestamps.append({
                "text": word["word"].strip(),
                "timestamp": [
                    round(word["start_offset"] / sampling_rate, 2),
                    round(word["end_offset"] / sampling_rate, 2)
                ]
            })
        
        # Create final response format
        response_data = {
            "text": processor.decode(outputs.sequences[0], skip_special_tokens=True),
            "chunks": words_with_timestamps
        }
        
        # Manual garbage collection to free memory
        del inputs, outputs, result
        if device == "cuda":
            torch.cuda.empty_cache()
        gc.collect()
        
        return response_data
    
    except Exception as e:
        print(f"Error processing audio: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error processing audio: {str(e)}")

async def process_in_background(file_path: str):
    """Process audio file in a background thread to prevent blocking"""
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(thread_pool, process_audio_file, file_path)

@app.post("/transcribe", response_model=TranscriptionResponse)
async def transcribe_audio(file: UploadFile = File(...)):
    """
    Transcribe an audio file to text with timestamps for each word.
    
    Accepts .mp3, .wav, .m4a, .ogg or .flac files up to 30MB.
    """
    start_time = time.time()
    
    # Validate file extension
    file_ext = os.path.splitext(file.filename)[1].lower()
    if file_ext not in FILE_EXTENSIONS:
        raise HTTPException(
            status_code=400, 
            detail=f"Unsupported file format. Supported formats: {', '.join(FILE_EXTENSIONS)}"
        )
    
    # Create temp file to store upload
    with tempfile.NamedTemporaryFile(delete=False, suffix=file_ext) as temp_file:
        # Get file content
        content = await file.read()
        
        # Check file size
        if len(content) > FILE_LIMIT_MB * 1024 * 1024:
            raise HTTPException(
                status_code=400,
                detail=f"File too large. Maximum size: {FILE_LIMIT_MB}MB"
            )
        
        # Write to temp file
        temp_file.write(content)
        temp_file_path = temp_file.name
    
    try:
        # Process the audio file in background to prevent blocking
        result = await process_in_background(temp_file_path)
        processing_time = time.time() - start_time
        print(f"Processing completed in {processing_time:.2f} seconds")
        
        return JSONResponse(content=result)
    
    finally:
        # Clean up the temp file
        if os.path.exists(temp_file_path):
            try:
                os.unlink(temp_file_path)
            except Exception as e:
                print(f"Error deleting temp file: {e}")

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {"status": "healthy"}

# Simple root endpoint that shows API is running
@app.get("/")
async def root():
    return {
        "message": "Speech-to-Text API is running",
        "endpoints": {
            "transcribe": "/transcribe (POST)",
            "health": "/health (GET)",
            "docs": "/docs (GET)"
        },
        "model": MODEL_NAME,
        "device": device
    }

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run("app:app", host="0.0.0.0", port=port)