api_server.py

# NVIDIA Parakeet TDT 0.6B v3 - Local ASR Service
# API-compatible with OpenAI Whisper API format

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
import torch
import numpy as np
import tempfile
import os
import io
from pathlib import Path
from pydub import AudioSegment
import logging

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

app = FastAPI(title="Parakeet ASR API", version="1.0.0")

# Global model variable
model = None
device = None

def load_model():
    """Load the Parakeet model"""
    global model, device
    
    try:
        from nemo.collections.asr.models import ASRModel
        
        # Check for CUDA
        device = "cuda" if torch.cuda.is_available() else "cpu"
        logger.info(f"Using device: {device}")
        
        if device == "cpu":
            logger.warning("Running on CPU - applying dynamic quantization to reduce RAM")
        
        logger.info("Loading Parakeet TDT 0.6B v3 model...")
        model = ASRModel.from_pretrained(model_name="nvidia/parakeet-tdt-0.6b-v3")
        
        if device == "cuda":
            model = model.to(device)
            # Use float16 on GPU to save VRAM
            model = model.half()
        else:
            # Apply PyTorch dynamic quantization on CPU to convert fp32 linear layers to int8
            # This drastically reduces RAM usage (often by half) and improves CPU inference speed
            try:
                model.encoder = torch.quantization.quantize_dynamic(
                    model.encoder, {torch.nn.Linear}, dtype=torch.qint8
                )
                model.decoder = torch.quantization.quantize_dynamic(
                    model.decoder, {torch.nn.Linear}, dtype=torch.qint8
                )
                logger.info("Dynamic int8 quantization applied successfully.")
            except Exception as q_err:
                logger.warning(f"Failed to apply dynamic quantization: {q_err}")
                
        model.eval()
        logger.info("Model loaded successfully")
        return True
        
    except Exception as e:
        logger.error(f"Failed to load model: {e}")
        return False

def convert_to_wav(audio_path):
    """Convert uploaded audio to 16kHz mono WAV for the model"""
    try:
        # Load audio with pydub (handles mp3, ogg, m4a, etc.)
        audio = AudioSegment.from_file(audio_path)
        
        # Convert to mono if stereo
        if audio.channels > 1:
            audio = audio.set_channels(1)
        
        # Resample to 16kHz (required by the model)
        if audio.frame_rate != 16000:
            audio = audio.set_frame_rate(16000)
            
        # Overwrite the temp file with proper WAV format
        audio.export(audio_path, format="wav")
        return len(audio) / 1000.0  # Return duration in seconds
        
    except Exception as e:
        logger.error(f"Audio conversion failed: {e}")
        raise

@app.on_event("startup")
async def startup_event():
    """Load model on startup"""
    success = load_model()
    if not success:
        logger.error("Failed to initialize model on startup")

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy" if model is not None else "unhealthy",
        "model": "nvidia/parakeet-tdt-0.6b-v3",
        "device": device or "unknown"
    }

@app.post("/v1/audio/transcriptions")
async def transcribe(
    file: UploadFile = File(...),
    model_name: str = "parakeet-tdt-0.6b-v3",
    language: str = None,
    prompt: str = None,
    response_format: str = "json",
    temperature: float = 0.0,
    timestamp_granularities: list = None
):
    """
    Transcribe audio file - OpenAI-compatible API
    """
    if model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    # Validate file (some clients omit content_type)
    if file.content_type and not file.content_type.startswith('audio/'):
        raise HTTPException(status_code=400, detail="File must be an audio file")
    
    try:
        # Save uploaded file to temp location
        with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as tmp:
            content = await file.read()
            tmp.write(content)
            tmp_path = tmp.name
        
        try:
            # Convert audio to 16kHz WAV
            logger.info(f"Processing audio file: {file.filename}")
            durationSeconds = convert_to_wav(tmp_path)
            
            # Transcribe using NeMo model
            with torch.no_grad():
                result = model.transcribe([tmp_path])
                
                # Handle different return types from NeMo versions
                if isinstance(result, tuple):
                    transcriptions = result[0]
                else:
                    transcriptions = result
                    
                transcription = transcriptions[0]
                
            # Format response (OpenAI-compatible)
            response = {
                "text": str(transcription).strip()
            }
            
            # Add word-level timestamps if requested (simulated based on duration)
            if timestamp_granularities and "word" in timestamp_granularities:
                word_text = str(transcription).strip()
                words = word_text.split() if word_text else []
                avg_word_duration = durationSeconds / len(words) if words else 0
                
                word_timestamps = []
                current_time = 0.0
                for word in words:
                    word_timestamps.append({
                        "word": word,
                        "start": round(current_time, 2),
                        "end": round(current_time + avg_word_duration, 2)
                    })
                    current_time += avg_word_duration
                
                response["words"] = word_timestamps
            
            return response
            
        finally:
            # Cleanup temp file
            if os.path.exists(tmp_path):
                os.unlink(tmp_path)
                
    except Exception as e:
        logger.error(f"Transcription failed: {e}")
        raise HTTPException(status_code=500, detail=f"Transcription failed: {str(e)}")

@app.post("/transcribe")
async def transcribe_simple(file: UploadFile = File(...)):
    """Simple transcription endpoint"""
    result = await transcribe(file)
    return result

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