Update app.py

app.py

@@ -1,20 +1,21 @@
-import streamlit as st
-import base64
 import os
+import base64
+import tempfile
 from flask import Flask, request, jsonify
 from pydub import AudioSegment
 import whisper
 from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
 
-# Define cache directory
-os.environ['HF_HOME'] = '/app/cache'
+# Set the cache directory to a writable location
+os.environ["TRANSFORMERS_CACHE"] = "/app/cache"  # Use a subdirectory in your working directory
+os.environ["HF_HOME"] = "/app/cache"  # Set HF_HOME for Hugging Face libraries
 
-# Load the Whisper model with a specified cache directory
-whisper_model = whisper.load_model("base", download_root="/app/cache")
+# Load the Whisper model
+whisper_model = whisper.load_model("base")
 
 # Load the translation model and tokenizer
-tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M", cache_dir="/app/cache")
-translation_model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M", cache_dir="/app/cache")
+tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")
+translation_model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M")
 
 def preprocess_audio(audio_path):
     """Convert audio to 16kHz mono WAV format."""
@@ -57,3 +58,24 @@ def handle_request(audio_base64, source_lang, target_lang):
     os.remove(audio_file_path)
 
     return {"transcribed_text": spoken_text, "translated_text": translated_text}
+
+# Flask for handling external POST requests
+app = Flask(__name__)
+
+@app.route('/translate', methods=['POST'])
+def translate():
+    """API endpoint for handling audio translation."""
+    data = request.json
+    if 'audio' not in data or 'source_lang' not in data or 'target_lang' not in data:
+        return jsonify({"error": "Invalid request format"}), 400
+    
+    audio_base64 = data['audio']
+    source_lang = data['source_lang']
+    target_lang = data['target_lang']
+
+    # Call the handle_request function to process the request
+    response = handle_request(audio_base64, source_lang, target_lang)
+    return jsonify(response)
+
+if __name__ == "__main__":
+    app.run(host='0.0.0.0', port=7860)