Update app.py

app.py

@@ -4,100 +4,150 @@ import torchaudio
 import soundfile as sf
 from flask import Flask, request, jsonify, send_file
 from flask_cors import CORS
-from transformers import VitsModel, AutoTokenizer
+from transformers import Wav2Vec2ForCTC, AutoProcessor, VitsModel, AutoTokenizer
 
-# Set ALL cache directories to /tmp (writable in Hugging Face Spaces)
+# Set cache directories
 os.environ["HF_HOME"] = "/tmp/hf_home"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/transformers_cache"
 os.environ["HUGGINGFACE_HUB_CACHE"] = "/tmp/huggingface_cache"
 os.environ["TORCH_HOME"] = "/tmp/torch_home"
 
 app = Flask(__name__)
-CORS(app)  # Allow external requests
+CORS(app)
+
+# ASR Model (facebook/mms-1b-all)
+ASR_MODEL_ID = "Coco-18/mms-asr-tgl-en-safetensor"
+asr_processor = AutoProcessor.from_pretrained(ASR_MODEL_ID)
+asr_model = Wav2Vec2ForCTC.from_pretrained(ASR_MODEL_ID)
+
+# Language-specific configurations
+LANGUAGE_CODES = {
+    "kapampangan": "pam",
+    "tagalog": "tgl",
+    "english": "eng"
+}
 
-# Model paths for different languages (Hugging Face Hub)
-MODELS = {
+# TTS Models (Kapampangan, Tagalog, English)
+TTS_MODELS = {
     "kapampangan": "facebook/mms-tts-pam",
     "tagalog": "facebook/mms-tts-tgl",
     "english": "facebook/mms-tts-eng"
 }
 
-loaded_models = {}
-loaded_processors = {}
-
-for lang, path in MODELS.items():
+tts_models = {}
+tts_processors = {}
+for lang, model_id in TTS_MODELS.items():
     try:
-        print(f"🔄 Loading {lang} model: {path}...")
-        # Force models to save in /tmp
-        loaded_models[lang] = VitsModel.from_pretrained(path, cache_dir="/tmp/huggingface_cache")
-        loaded_processors[lang] = AutoTokenizer.from_pretrained(path, cache_dir="/tmp/huggingface_cache")
-        print(f"✅ {lang.capitalize()} model loaded successfully!")
+        tts_models[lang] = VitsModel.from_pretrained(model_id, cache_dir="/tmp/huggingface_cache")
+        tts_processors[lang] = AutoTokenizer.from_pretrained(model_id, cache_dir="/tmp/huggingface_cache")
+        print(f"✅ TTS Model loaded: {lang}")
     except Exception as e:
-        print(f"❌ Error loading {lang} model: {str(e)}")
-        loaded_models[lang] = None  # Mark as unavailable
-        loaded_processors[lang] = None
+        print(f"❌ Error loading {lang} TTS model: {e}")
+        tts_models[lang] = None
 
 # Constants
+SAMPLE_RATE = 16000
 OUTPUT_DIR = "/tmp/"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
+
 @app.route("/", methods=["GET"])
 def home():
-    """Root route to check if the API is running"""
-    return jsonify({"message": "TTS API is running. Use /tts to generate speech."})
+    return jsonify({"message": "Speech API is running."})
+
+
+@app.route("/asr", methods=["POST"])
+def transcribe_audio():
+    try:
+        if "audio" not in request.files:
+            return jsonify({"error": "No audio file uploaded"}), 400
+
+        audio_file = request.files["audio"]
+        language = request.form.get("language", "english").lower()
+
+        # Validate language
+        if language not in LANGUAGE_CODES:
+            return jsonify({"error": f"Unsupported language: {language}"}), 400
+
+        # Get the language code for the ASR model
+        lang_code = LANGUAGE_CODES[language]
+
+        # Save audio file temporarily
+        audio_path = os.path.join(OUTPUT_DIR, "input_audio.wav")
+        audio_file.save(audio_path)
+
+        # Load and process audio
+        try:
+            waveform, sr = torchaudio.load(audio_path)
+            if sr != SAMPLE_RATE:
+                waveform = torchaudio.transforms.Resample(sr, SAMPLE_RATE)(waveform)
+
+            # Normalize audio (recommended for Wav2Vec2)
+            waveform = waveform / torch.max(torch.abs(waveform))
+
+            # Process audio for ASR
+            inputs = asr_processor(
+                waveform.squeeze().numpy(),
+                sampling_rate=SAMPLE_RATE,
+                return_tensors="pt",
+                language=lang_code  # Set the language code
+            )
+        except Exception as e:
+            return jsonify({"error": f"Error processing audio: {str(e)}"}), 400
+
+        # Transcribe
+        with torch.no_grad():
+            logits = asr_model(**inputs).logits
+        ids = torch.argmax(logits, dim=-1)[0]
+        transcription = asr_processor.decode(ids)
+
+        # Log the transcription
+        print(f"Transcription ({language}): {transcription}")
+
+        return jsonify({"transcription": transcription})
+    except Exception as e:
+        print(f"ASR error: {str(e)}")
+        return jsonify({"error": f"ASR failed: {str(e)}"}), 500
+
 
 @app.route("/tts", methods=["POST"])
 def generate_tts():
-    """API endpoint to generate TTS audio"""
     try:
-        # Get request data
         data = request.get_json()
         text_input = data.get("text", "").strip()
         language = data.get("language", "kapampangan").lower()
 
-        # Validate inputs
-        if language not in MODELS:
-            return jsonify({"error": "Invalid language. Choose 'kapampangan', 'tagalog', or 'english'."}), 400
+        if language not in TTS_MODELS:
+            return jsonify({"error": "Invalid language"}), 400
         if not text_input:
             return jsonify({"error": "No text provided"}), 400
-        if loaded_models[language] is None:
-            return jsonify({"error": f"Model for {language} failed to load"}), 500
-
-        print(f"🔄 Generating speech for '{text_input}' in {language}...")
+        if tts_models[language] is None:
+            return jsonify({"error": "TTS model not available"}), 500
 
-        # Process text input
-        processor = loaded_processors[language]
-        model = loaded_models[language]
+        processor = tts_processors[language]
+        model = tts_models[language]
         inputs = processor(text_input, return_tensors="pt")
 
-        # Generate speech - using model(**inputs) instead of model.generate()
         with torch.no_grad():
-            output = model(**inputs).waveform
-            waveform = output.squeeze().cpu().numpy()
-
-        # Save to file
-        output_filename = os.path.join(OUTPUT_DIR, f"{language}_output.wav")
-        # Use the model's sampling rate
-        sampling_rate = model.config.sampling_rate
-        sf.write(output_filename, waveform, sampling_rate)
-        print(f"✅ Speech generated! File saved: {output_filename}")
-
-        return jsonify({
-            "message": "TTS audio generated",
-            "file_url": f"/download/{language}_output.wav"
-        })
+            output = model.generate(**inputs)
+
+        waveform = output.cpu().numpy().flatten()
+        output_filename = os.path.join(OUTPUT_DIR, f"{language}_tts.wav")
+        sf.write(output_filename, waveform, SAMPLE_RATE)
+
+        return jsonify({"file_url": f"/download/{language}_tts.wav"})
     except Exception as e:
-        print(f"❌ Error generating TTS: {e}")
-        return jsonify({"error": f"Internal server error: {str(e)}"}), 500
+        return jsonify({"error": f"TTS failed: {e}"}), 500
+
 
 @app.route("/download/<filename>", methods=["GET"])
 def download_audio(filename):
-    """Serve generated audio files"""
     file_path = os.path.join(OUTPUT_DIR, filename)
     if os.path.exists(file_path):
         return send_file(file_path, mimetype="audio/wav", as_attachment=True)
     return jsonify({"error": "File not found"}), 404
 
+
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860, debug=True)