Update audio handling: Save synthesized audio as WAV, return URL, set 10-min expiration with 5-min cleanup interval

app.py

@@ -10,9 +10,13 @@ import torch
 import numpy as np
 import soundfile as sf
 import librosa
-from fastapi import FastAPI, HTTPException, UploadFile, File, Form
+import wave
+import time
+from fastapi import FastAPI, HTTPException, UploadFile, File, Form, BackgroundTasks
 from fastapi.responses import JSONResponse
+from fastapi.staticfiles import StaticFiles
 from typing import Dict, Any, Optional, Tuple
+from datetime import datetime, timedelta
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -20,6 +24,11 @@ logger = logging.getLogger("talklas-api")
 
 app = FastAPI(title="Talklas API")
 
+# Mount a directory to serve audio files
+AUDIO_DIR = "audio_output"
+os.makedirs(AUDIO_DIR, exist_ok=True)
+app.mount("/audio_output", StaticFiles(directory=AUDIO_DIR), name="audio_output")
+
 # Global variables to track application state
 models_loaded = False
 loading_in_progress = False
@@ -58,6 +67,37 @@ NLLB_LANGUAGE_CODES = {
     "pag": "pag_Latn"
 }
 
+# Function to save PCM data as a WAV file
+def save_pcm_to_wav(pcm_data: list, sample_rate: int, output_path: str):
+    with wave.open(output_path, 'wb') as wav_file:
+        # Set WAV parameters: 1 channel (mono), 2 bytes per sample (16-bit), sample rate
+        wav_file.setnchannels(1)
+        wav_file.setsampwidth(2)  # 16-bit audio
+        wav_file.setframerate(sample_rate)
+        # Convert PCM data (list of integers) to bytes
+        wav_file.writeframes(bytes(pcm_data))
+
+# Function to clean up old audio files
+def cleanup_old_audio_files():
+    logger.info("Starting cleanup of old audio files...")
+    expiration_time = datetime.now() - timedelta(minutes=10)  # Files older than 10 minutes
+    for filename in os.listdir(AUDIO_DIR):
+        file_path = os.path.join(AUDIO_DIR, filename)
+        if os.path.isfile(file_path):
+            file_mtime = datetime.fromtimestamp(os.path.getmtime(file_path))
+            if file_mtime < expiration_time:
+                try:
+                    os.unlink(file_path)
+                    logger.info(f"Deleted old audio file: {file_path}")
+                except Exception as e:
+                    logger.error(f"Error deleting file {file_path}: {str(e)}")
+
+# Background task to periodically clean up audio files
+def schedule_cleanup():
+    while True:
+        cleanup_old_audio_files()
+        time.sleep(300)  # Run every 5 minutes (300 seconds)
+
 # Function to load models in background
 def load_models_task():
     global models_loaded, loading_in_progress, model_status, error_message
@@ -158,11 +198,18 @@ def start_model_loading():
         loading_thread.daemon = True
         loading_thread.start()
 
-# Start the background process when the app starts
+# Start the background cleanup task
+def start_cleanup_task():
+    cleanup_thread = threading.Thread(target=schedule_cleanup)
+    cleanup_thread.daemon = True
+    cleanup_thread.start()
+
+# Start the background processes when the app starts
 @app.on_event("startup")
 async def startup_event():
     logger.info("Application starting up...")
     start_model_loading()
+    start_cleanup_task()
 
 @app.get("/")
 async def root():
@@ -300,7 +347,7 @@ async def translate_text(text: str = Form(...), source_lang: str = Form(...), ta
         logger.warning("MT model not loaded, skipping translation")
 
     # Convert translated text to speech
-    output_audio = None
+    output_audio_url = None
     if model_status["tts"].startswith("loaded") and tts_model is not None and tts_tokenizer is not None:
         try:
             inputs = tts_tokenizer(translated_text, return_tensors="pt").to(device)
@@ -308,11 +355,20 @@ async def translate_text(text: str = Form(...), source_lang: str = Form(...), ta
                 output = tts_model(**inputs)
             speech = output.waveform.cpu().numpy().squeeze()
             speech = (speech * 32767).astype(np.int16)
-            output_audio = (tts_model.config.sampling_rate, speech.tolist())
+            sample_rate = tts_model.config.sampling_rate
+
+            # Save the audio as a WAV file
+            output_filename = f"{request_id}.wav"
+            output_path = os.path.join(AUDIO_DIR, output_filename)
+            save_pcm_to_wav(speech.tolist(), sample_rate, output_path)
+            logger.info(f"Saved synthesized audio to {output_path}")
+
+            # Generate a URL to the WAV file
+            output_audio_url = f"https://jerich-talklasapp.hf.space/audio_output/{output_filename}"
             logger.info("TTS conversion completed")
         except Exception as e:
             logger.error(f"Error during TTS conversion: {str(e)}")
-            output_audio = None
+            output_audio_url = None
     
     return {
         "request_id": request_id,
@@ -320,7 +376,7 @@ async def translate_text(text: str = Form(...), source_lang: str = Form(...), ta
         "message": "Translation and TTS completed (or partially completed).",
         "source_text": text,
         "translated_text": translated_text,
-        "output_audio": output_audio
+        "output_audio": output_audio_url
     }
 
 @app.post("/translate-audio")
@@ -355,7 +411,7 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
     
     transcription = "Transcription not available"
     translated_text = "Translation not available"
-    output_audio = None
+    output_audio_url = None
     
     try:
         # Step 1: Transcribe the audio (STT)
@@ -415,11 +471,20 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
                     output = tts_model(**inputs)
                 speech = output.waveform.cpu().numpy().squeeze()
                 speech = (speech * 32767).astype(np.int16)
-                output_audio = (tts_model.config.sampling_rate, speech.tolist())
+                sample_rate = tts_model.config.sampling_rate
+
+                # Save the audio as a WAV file
+                output_filename = f"{request_id}.wav"
+                output_path = os.path.join(AUDIO_DIR, output_filename)
+                save_pcm_to_wav(speech.tolist(), sample_rate, output_path)
+                logger.info(f"Saved synthesized audio to {output_path}")
+
+                # Generate a URL to the WAV file
+                output_audio_url = f"https://jerich-talklasapp.hf.space/audio_output/{output_filename}"
                 logger.info("TTS conversion completed")
             except Exception as e:
                 logger.error(f"Error during TTS conversion: {str(e)}")
-                output_audio = None
+                output_audio_url = None
         
         return {
             "request_id": request_id,
@@ -427,7 +492,7 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
             "message": "Transcription, translation, and TTS completed (or partially completed).",
             "source_text": transcription,
             "translated_text": translated_text,
-            "output_audio": output_audio
+            "output_audio": output_audio_url
         }
     except Exception as e:
         logger.error(f"Error during processing: {str(e)}")
@@ -437,7 +502,7 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
             "message": f"Processing failed: {str(e)}",
             "source_text": transcription,
             "translated_text": translated_text,
-            "output_audio": output_audio
+            "output_audio": output_audio_url
         }
     finally:
         logger.info(f"Cleaning up temporary file: {temp_path}")