import os import logging import tempfile import traceback from typing import Optional, List, Union from contextlib import asynccontextmanager from fastapi import FastAPI, UploadFile, File, Form, HTTPException, status from fastapi.responses import JSONResponse from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, Field from faster_whisper import WhisperModel # --- Configuration & Logging --- logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger("ASR_API") # Model Settings # Default to 'base' for best balance of speed and accuracy on CPU. # 'tiny' is faster but less accurate. 'small' is slower but more accurate. MODEL_SIZE = os.getenv("WHISPER_MODEL_SIZE", "base") # int8 is recommended for CPU to improve speed COMPUTE_TYPE = os.getenv("WHISPER_COMPUTE_TYPE", "int8") DEVICE = "cpu" # Enforcing CPU usage as requested # Global model instance model_instance = None # --- Lifespan Management --- @asynccontextmanager async def lifespan(app: FastAPI): # Startup: Load the model into memory global model_instance logger.info(f"Loading Whisper Model (Size: {MODEL_SIZE}, Device: {DEVICE}, Compute: {COMPUTE_TYPE})...") try: # faster-whisper handles model downloading automatically if not present model_instance = WhisperModel( MODEL_SIZE, device=DEVICE, compute_type=COMPUTE_TYPE, download_root=".cache/models" # Cache in space directory ) logger.info("Model loaded successfully.") except Exception as e: logger.error(f"Failed to load model: {e}") raise RuntimeError("Failed to initialize ASR model") yield # Shutdown: Cleanup (if necessary) logger.info("Shutting down application...") del model_instance app = FastAPI( title="Enterprise ASR API", description="High-performance Speech-to-Text API using Faster Whisper on CPU", version="1.0.0", lifespan=lifespan ) # Enable CORS for all origins app.add_middleware( CORSMiddleware, allow_origins=["*"], # Allow all origins allow_credentials=True, allow_methods=["*"], # Allow all methods (GET, POST, etc.) allow_headers=["*"], # Allow all headers ) # --- Pydantic Schemas --- class TranscriptionSegment(BaseModel): id: int seek: int start: float end: float text: str tokens: List[int] temperature: float avg_logprob: float compression_ratio: float no_speech_prob: float class Config: # Ensure fields like tokens are serialized correctly json_schema_extra = { "example": { "id": 0, "seek": 0, "start": 0.0, "end": 2.0, "text": "Hello world", "tokens": [50364, 1917, 1002, 50664], "temperature": 0.0, "avg_logprob": -0.5, "compression_ratio": 1.0, "no_speech_prob": 0.1 } } class TranscriptionResponse(BaseModel): text: str = Field(..., description="The full transcribed text") language: str = Field(..., description="Detected language code") segments: List[TranscriptionSegment] = Field(default_factory=list, description="Detailed segments of transcription") duration: float = Field(..., description="Duration of audio in seconds") # --- Helper Functions --- def clean_up_temp_file(filepath: str): """Safely remove a temporary file.""" try: if os.path.exists(filepath): os.remove(filepath) logger.debug(f"Deleted temporary file: {filepath}") except Exception as e: logger.warning(f"Failed to delete temporary file {filepath}: {e}") # --- Endpoints --- @app.get("/health", tags=["System"]) async def health_check(): """Check if the API is running and the model is loaded.""" if model_instance is None: raise HTTPException(status_code=503, detail="Model not loaded yet") return {"status": "healthy", "model": MODEL_SIZE, "device": DEVICE} @app.post( "/v1/transcribe", response_model=TranscriptionResponse, tags=["ASR Operations"], summary="Transcribe or Translate Audio" ) async def transcribe_audio( file: UploadFile = File(..., description="Audio file (mp3, wav, m4a, etc.)"), task: str = Form("transcribe", description="Task: 'transcribe' or 'translate'"), language: Optional[str] = Form(None, description="Language code (e.g., 'en', 'es'). Auto-detect if None."), initial_prompt: Optional[str] = Form(None, description="Optional initial text to guide the model."), vad_filter: bool = Form(True, description="Enable Voice Activity Detection (VAD) to silence gaps."), word_timestamps: bool = Form(False, description="Extract word-level timestamps."), temperature: float = Form(0.0, description="Sampling temperature (0.0 for greedy).") ): """ Endpoint to process audio files for transcription or translation. - **file**: The audio file to process. - **task**: Use 'transcribe' for same-language text, 'translate' for English. - **language**: Specifying a language improves speed and accuracy. - **vad_filter**: Strongly recommended for long audio to cut out silence. """ if model_instance is None: raise HTTPException(status_code=503, detail="Model is initializing or failed to load.") # 1. Validate input file extension (basic check) filename = file.filename or "audio.blob" logger.info(f"Received request: File='{filename}', Task='{task}', Lang='{language}'") # Create a temporary file to store the upload temp_filepath = None try: # We write the uploaded file to a temp disk because faster-whisper expects a file path # and handles loading optimized chunks from disk better than loading full byte-strings in memory with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(filename)[1]) as tmp: content = await file.read() tmp.write(content) temp_filepath = tmp.name # 2. Perform Inference # We run the blocking CPU task in a thread pool to avoid blocking the async event loop segments_generator, info = await run_inference( temp_filepath, task, language, initial_prompt, vad_filter, word_timestamps, temperature ) # 3. Format Results full_text = "" segment_objects = [] for segment in segments_generator: full_text += segment.text segment_objects.append(TranscriptionSegment( id=segment.id, seek=segment.seek, start=segment.start, end=segment.end, text=segment.text.strip(), tokens=segment.tokens, temperature=segment.temperature, avg_logprob=segment.avg_logprob, compression_ratio=segment.compression_ratio, no_speech_prob=segment.no_speech_prob )) return TranscriptionResponse( text=full_text.strip(), language=info.language, duration=info.duration, segments=segment_objects ) except HTTPException: raise except Exception as e: logger.error(f"Error during processing: {str(e)}") logger.error(traceback.format_exc()) raise HTTPException( status_code=500, detail=f"Internal Server Error during processing: {str(e)}" ) finally: # 4. Cleanup if temp_filepath: clean_up_temp_file(temp_filepath) async def run_inference( filepath: str, task: str, language: Optional[str], initial_prompt: Optional[str], vad_filter: bool, word_timestamps: bool, temperature: float ): """ Wrapper to run the synchronous faster-whisper model in a thread pool executor. This prevents the API from freezing while the CPU crunches numbers. """ def _compute(): return model_instance.transcribe( audio=filepath, task=task, language=language, initial_prompt=initial_prompt, vad_filter=vad_filter, word_timestamps=word_timestamps, temperature=temperature ) # Run in a separate thread from fastapi.concurrency import run_in_threadpool return await run_in_threadpool(_compute)