deploy

app.py

@@ -1,3 +1,6 @@
+import asyncio
+import logging
+import time
 from fastapi import FastAPI, File, UploadFile
 from transformers import WhisperProcessor, WhisperForConditionalGeneration
 import torch
@@ -5,9 +8,7 @@ import io
 import soundfile as sf
 import numpy as np
 import torchaudio
-import logging
-import timeout_decorator
-import time
+import psutil
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -30,80 +31,71 @@ except Exception as e:
 model.config.forced_decoder_ids = processor.get_decoder_prompt_ids(language="ur", task="transcribe")
 logger.info("Set forced_decoder_ids for Urdu transcription")
 
-# Move model to CPU (free Spaces don’t have GPU)
+# Move model to CPU
 device = "cpu"
 model.to(device)
 logger.info(f"Model loaded and moved to {device}")
 
+# Log memory usage
+def log_memory_usage():
+    process = psutil.Process()
+    mem_info = process.memory_info()
+    logger.info(f"Memory usage: {mem_info.rss / 1024**2:.2f} MB")
+
 @app.post("/transcribe")
 async def transcribe_audio(file: UploadFile = File(...)):
     try:
         start_time = time.time()
-        # Read audio file (supports WAV, MP3, etc.)
+        log_memory_usage()
+
+        # Read audio file
         logger.info("Reading audio file")
-        try:
-            audio_data, sample_rate = sf.read(io.BytesIO(await file.read()))
-        except Exception as e:
-            logger.error(f"Failed to read audio file: {str(e)}")
-            return {"error": f"Invalid or unsupported audio file: {str(e)}. Supported formats: WAV, MP3, FLAC."}
+        audio_data, sample_rate = sf.read(io.BytesIO(await file.read()))
         logger.info(f"Audio read in {time.time() - start_time:.2f} seconds")
 
         # Ensure audio is mono
         if len(audio_data.shape) > 1:
-            audio_data = np.mean(audio_data, axis=1)  # Convert to mono
+            audio_data = np.mean(audio_data, axis=1)
 
         # Resample to 16kHz if necessary
         target_sample_rate = 16000
         if sample_rate != target_sample_rate:
             logger.info(f"Resampling audio from {sample_rate} Hz to {target_sample_rate} Hz")
-            step_time = time.time()
             audio_tensor = torch.from_numpy(audio_data).float()
             resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_sample_rate)
-            audio_tensor = resampler(audio_tensor)
-            audio_data = audio_tensor.numpy()
+            audio_data = resampler(audio_tensor).numpy()
             sample_rate = target_sample_rate
-            logger.info(f"Resampling completed in {time.time() - step_time:.2f} seconds")
-
-        # Trim silence (simplified for torchaudio 2.0.2)
-        logger.info("Trimming silence")
-        step_time = time.time()
-        audio_tensor = torch.from_numpy(audio_data).float()
-        vad = torchaudio.transforms.Vad(sample_rate=sample_rate)  # No threshold
-        audio_tensor = vad(audio_tensor)
-        audio_data = audio_tensor.numpy()
-        logger.info(f"Silence trimming completed in {time.time() - step_time:.2f} seconds")
 
         # Process audio input
         logger.info("Processing audio input")
-        step_time = time.time()
         inputs = processor(audio_data, sampling_rate=sample_rate, return_tensors="pt")
         input_features = inputs.input_features.to(device)
-        logger.info(f"Input processing completed in {time.time() - step_time:.2f} seconds")
 
-        # Generate transcription with timeout
+        # Generate transcription with async timeout
         logger.info("Generating transcription")
-        step_time = time.time()
-        @timeout_decorator.timeout(15, timeout_exception=TimeoutError)  # 15-second timeout
-        def generate_transcription():
+        async def generate_transcription():
             with torch.no_grad():
                 generated_ids = model.generate(
                     input_features,
                     max_new_tokens=225,
-                    num_beams=1,  # Disable beam search
-                    length_penalty=0.0  # Faster decoding
+                    num_beams=1,
+                    length_penalty=0.0
                 )
             return generated_ids
 
-        generated_ids = generate_transcription()
-        logger.info(f"Transcription generated in {time.time() - step_time:.2f} seconds")
-        transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        try:
+            async with asyncio.timeout(60):  # 60-second timeout
+                generated_ids = await generate_transcription()
+        except asyncio.TimeoutError:
+            logger.error("Transcription timed out after 60 seconds")
+            return {"error": "Transcription took too long. Try a smaller model or upgrade your Space."}
 
+        transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
         total_time = time.time() - start_time
         logger.info(f"Total transcription time: {total_time:.2f} seconds")
+        log_memory_usage()
         return {"transcription": transcription}
-    except TimeoutError:
-        logger.error("Transcription timed out after 15 seconds")
-        return {"error": "Transcription took too long. Try a faster model or check Space performance."}
+
     except Exception as e:
         logger.error(f"Error during transcription: {str(e)}")
         return {"error": str(e)}

requirements.txt

@@ -1,11 +1,11 @@
-transformers==4.38.2
-torch==2.0.1 
-fastapi==0.103.0 
-uvicorn==0.23.2 
-pydantic==2.3.0 
-soundfile==0.12.1 
-python-multipart==0.0.9 
-numpy==1.26.4 
-timeout-decorator==0.5.0 
-torchaudio==2.0.2 
+transformers==4.44.2
+torch==2.4.1
+torchaudio==2.4.1
+fastapi==0.103.0
+uvicorn==0.23.2
+pydantic==2.3.0
+soundfile==0.12.1
+python-multipart==0.0.9
+numpy==1.26.4
+psutil==6.0.0
 accelerate==0.30.1