translator.py

# translator.py - Handles ASR, TTS, and translation tasks

import os
import sys
import logging
import traceback
import torch
import torchaudio
import tempfile
import soundfile as sf
from pydub import AudioSegment
from flask import jsonify
from transformers import Wav2Vec2ForCTC, AutoProcessor, VitsModel, AutoTokenizer
from transformers import MarianMTModel, MarianTokenizer

# Configure logging
logger = logging.getLogger("speech_api")

# Global variables to store models and processors
asr_model = None
asr_processor = None
tts_models = {}
tts_processors = {}
translation_models = {}
translation_tokenizers = {}

# Language-specific configurations
LANGUAGE_CODES = {
    "kapampangan": "pam",
    "filipino": "fil",
    "english": "eng",
    "tagalog": "tgl",
}

# TTS Models (Kapampangan, Tagalog, English)
TTS_MODELS = {
    "kapampangan": "facebook/mms-tts-pam",
    "tagalog": "facebook/mms-tts-tgl",
    "english": "facebook/mms-tts-eng"
}

# Translation Models
TRANSLATION_MODELS = {
    "pam-eng": "Coco-18/opus-mt-pam-en",
    "eng-pam": "Coco-18/opus-mt-en-pam",
    "tgl-eng": "Helsinki-NLP/opus-mt-tl-en",
    "eng-tgl": "Helsinki-NLP/opus-mt-en-tl",
    "phi": "Coco-18/opus-mt-phi"
}

def init_models(device):
    """Initialize all models required for the API"""
    global asr_model, asr_processor, tts_models, tts_processors, translation_models, translation_tokenizers

    # Initialize ASR model
    ASR_MODEL_ID = "Coco-18/mms-asr-tgl-en-safetensor"
    logger.info(f"🔄 Loading ASR model: {ASR_MODEL_ID}")

    try:
        asr_processor = AutoProcessor.from_pretrained(
            ASR_MODEL_ID,
            cache_dir=os.environ.get("TRANSFORMERS_CACHE")
        )
        logger.info("✅ ASR processor loaded successfully")

        asr_model = Wav2Vec2ForCTC.from_pretrained(
            ASR_MODEL_ID,
            cache_dir=os.environ.get("TRANSFORMERS_CACHE")
        )
        asr_model.to(device)
        logger.info(f"✅ ASR model loaded successfully on {device}")
    except Exception as e:
        logger.error(f"❌ Error loading ASR model: {str(e)}")
        logger.debug(f"Stack trace: {traceback.format_exc()}")

    # Initialize TTS models
    for lang, model_id in TTS_MODELS.items():
        logger.info(f"🔄 Loading TTS model for {lang}: {model_id}")
        try:
            tts_processors[lang] = AutoTokenizer.from_pretrained(
                model_id,
                cache_dir=os.environ.get("TRANSFORMERS_CACHE")
            )
            logger.info(f"✅ {lang} TTS processor loaded")

            tts_models[lang] = VitsModel.from_pretrained(
                model_id,
                cache_dir=os.environ.get("TRANSFORMERS_CACHE")
            )
            tts_models[lang].to(device)
            logger.info(f"✅ {lang} TTS model loaded on {device}")
        except Exception as e:
            logger.error(f"❌ Failed to load {lang} TTS model: {str(e)}")
            logger.debug(f"Stack trace: {traceback.format_exc()}")
            tts_models[lang] = None

    # Initialize translation models
    for model_key, model_id in TRANSLATION_MODELS.items():
        logger.info(f"🔄 Loading Translation model: {model_id}")

        try:
            translation_tokenizers[model_key] = MarianTokenizer.from_pretrained(
                model_id,
                cache_dir=os.environ.get("TRANSFORMERS_CACHE")
            )
            logger.info(f"✅ Translation tokenizer loaded successfully for {model_key}")

            translation_models[model_key] = MarianMTModel.from_pretrained(
                model_id,
                cache_dir=os.environ.get("TRANSFORMERS_CACHE")
            )
            translation_models[model_key].to(device)
            logger.info(f"✅ Translation model loaded successfully on {device} for {model_key}")
        except Exception as e:
            logger.error(f"❌ Error loading Translation model for {model_key}: {str(e)}")
            logger.debug(f"Stack trace: {traceback.format_exc()}")
            translation_models[model_key] = None
            translation_tokenizers[model_key] = None


def check_model_status():
    """Check and return the status of all models"""
    # Initialize direct language pair statuses based on loaded models
    translation_status = {}

    # Add status for direct model pairs
    for lang_pair in ["pam-eng", "eng-pam", "tgl-eng", "eng-tgl"]:
        translation_status[lang_pair] = "loaded" if lang_pair in translation_models and translation_models[
            lang_pair] is not None else "failed"

    # Add special phi model status
    phi_status = "loaded" if "phi" in translation_models and translation_models["phi"] is not None else "failed"
    translation_status["pam-fil"] = phi_status
    translation_status["fil-pam"] = phi_status
    translation_status["pam-tgl"] = phi_status  # Using phi model but replacing tgl with fil
    translation_status["tgl-pam"] = phi_status  # Using phi model but replacing tgl with fil

    return {
        "asr_model": "loaded" if asr_model is not None else "failed",
        "tts_models": {lang: "loaded" if model is not None else "failed"
                       for lang, model in tts_models.items()},
        "translation_models": translation_status
    }


def handle_asr_request(request, output_dir, sample_rate):
    """Handle ASR (Automatic Speech Recognition) requests"""
    if asr_model is None or asr_processor is None:
        logger.error("❌ ASR endpoint called but models aren't loaded")
        return jsonify({"error": "ASR model not available"}), 503

    try:
        if "audio" not in request.files:
            logger.warning("⚠️ ASR request missing audio file")
            return jsonify({"error": "No audio file uploaded"}), 400

        audio_file = request.files["audio"]
        language = request.form.get("language", "english").lower()

        if language not in LANGUAGE_CODES:
            logger.warning(f"⚠️ Unsupported language requested: {language}")
            return jsonify(
                {"error": f"Unsupported language: {language}. Available: {list(LANGUAGE_CODES.keys())}"}), 400

        lang_code = LANGUAGE_CODES[language]
        logger.info(f"🔄 Processing {language} audio for ASR")

        # Save the uploaded file temporarily
        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(audio_file.filename)[-1]) as temp_audio:
            temp_audio.write(audio_file.read())
            temp_audio_path = temp_audio.name
            logger.debug(f"📁 Temporary audio saved to {temp_audio_path}")

        # Convert to WAV if necessary
        wav_path = temp_audio_path
        if not audio_file.filename.lower().endswith(".wav"):
            wav_path = os.path.join(output_dir, "converted_audio.wav")
            logger.info(f"🔄 Converting audio to WAV format: {wav_path}")
            try:
                audio = AudioSegment.from_file(temp_audio_path)
                audio = audio.set_frame_rate(sample_rate).set_channels(1)
                audio.export(wav_path, format="wav")
            except Exception as e:
                logger.error(f"❌ Audio conversion failed: {str(e)}")
                return jsonify({"error": f"Audio conversion failed: {str(e)}"}), 500

# Load and process the WAV file
        try:
            waveform, sr = torchaudio.load(wav_path)
            logger.debug(f"✅ Audio loaded: {wav_path} (Sample rate: {sr}Hz)")

            # Resample if needed
            if sr != sample_rate:
                logger.info(f"🔄 Resampling audio from {sr}Hz to {sample_rate}Hz")
                waveform = torchaudio.transforms.Resample(sr, sample_rate)(waveform)

            waveform = waveform / torch.max(torch.abs(waveform))
        except Exception as e:
            logger.error(f"❌ Failed to load or process audio: {str(e)}")
            return jsonify({"error": f"Audio processing failed: {str(e)}"}), 500

        # Process audio for ASR
        try:
            inputs = asr_processor(
                waveform.squeeze().numpy(),
                sampling_rate=sample_rate,
                return_tensors="pt",
                language=lang_code
            )
            inputs = {k: v.to(asr_model.device) for k, v in inputs.items()}
        except Exception as e:
            logger.error(f"❌ ASR preprocessing failed: {str(e)}")
            return jsonify({"error": f"ASR preprocessing failed: {str(e)}"}), 500

        # Perform ASR
        try:
            with torch.no_grad():
                logits = asr_model(**inputs).logits
            ids = torch.argmax(logits, dim=-1)[0]
            transcription = asr_processor.decode(ids)

            logger.info(f"✅ Transcription ({language}): {transcription}")

            # Clean up temp files
            try:
                os.unlink(temp_audio_path)
                if wav_path != temp_audio_path:
                    os.unlink(wav_path)
            except Exception as e:
                logger.warning(f"⚠️ Failed to clean up temp files: {str(e)}")

            return jsonify({
                "transcription": transcription,
                "language": language,
                "language_code": lang_code
            })
        except Exception as e:
            logger.error(f"❌ ASR inference failed: {str(e)}")
            logger.debug(f"Stack trace: {traceback.format_exc()}")
            return jsonify({"error": f"ASR inference failed: {str(e)}"}), 500

    except Exception as e:
        logger.error(f"❌ Unhandled exception in ASR endpoint: {str(e)}")
        logger.debug(f"Stack trace: {traceback.format_exc()}")
        return jsonify({"error": f"Internal server error: {str(e)}"}), 500

def handle_tts_request(request, output_dir):
    """Handle TTS (Text-to-Speech) requests"""
    try:
        data = request.get_json()
        if not data:
            logger.warning("⚠️ TTS endpoint called with no JSON data")
            return jsonify({"error": "No JSON data provided"}), 400

        text_input = data.get("text", "").strip()
        language = data.get("language", "kapampangan").lower()

        if not text_input:
            logger.warning("⚠️ TTS request with empty text")
            return jsonify({"error": "No text provided"}), 400

        if language not in TTS_MODELS:
            logger.warning(f"⚠️ TTS requested for unsupported language: {language}")
            return jsonify({"error": f"Invalid language. Available options: {list(TTS_MODELS.keys())}"}), 400

        if tts_models[language] is None:
            logger.error(f"❌ TTS model for {language} not loaded")
            return jsonify({"error": f"TTS model for {language} not available"}), 503

        logger.info(f"🔄 Generating TTS for language: {language}, text: '{text_input}'")

        try:
            processor = tts_processors[language]
            model = tts_models[language]
            inputs = processor(text_input, return_tensors="pt")
            inputs = {k: v.to(model.device) for k, v in inputs.items()}
        except Exception as e:
            logger.error(f"❌ TTS preprocessing failed: {str(e)}")
            return jsonify({"error": f"TTS preprocessing failed: {str(e)}"}), 500

        # Generate speech
        try:
            with torch.no_grad():
                output = model(**inputs).waveform
                waveform = output.squeeze().cpu().numpy()
        except Exception as e:
            logger.error(f"❌ TTS inference failed: {str(e)}")
            logger.debug(f"Stack trace: {traceback.format_exc()}")
            return jsonify({"error": f"TTS inference failed: {str(e)}"}), 500

        # Save to file
        try:
            output_filename = os.path.join(output_dir, f"{language}_output.wav")
            sampling_rate = model.config.sampling_rate
            sf.write(output_filename, waveform, sampling_rate)
            logger.info(f"✅ Speech generated! File saved: {output_filename}")
        except Exception as e:
            logger.error(f"❌ Failed to save audio file: {str(e)}")
            return jsonify({"error": f"Failed to save audio file: {str(e)}"}), 500

        return jsonify({
            "message": "TTS audio generated",
            "file_url": f"/download/{os.path.basename(output_filename)}",
            "language": language,
            "text_length": len(text_input)
        })
    except Exception as e:
        logger.error(f"❌ Unhandled exception in TTS endpoint: {str(e)}")
        logger.debug(f"Stack trace: {traceback.format_exc()}")
        return jsonify({"error": f"Internal server error: {str(e)}"}), 500

def handle_tts_request(request, output_dir):
    """Handle TTS (Text-to-Speech) requests"""
    try:
        data = request.get_json()
        if not data:
            logger.warning("⚠️ TTS endpoint called with no JSON data")
            return jsonify({"error": "No JSON data provided"}), 400

        text_input = data.get("text", "").strip()
        language = data.get("language", "kapampangan").lower()

        if not text_input:
            logger.warning("⚠️ TTS request with empty text")
            return jsonify({"error": "No text provided"}), 400

        if language not in TTS_MODELS:
            logger.warning(f"⚠️ TTS requested for unsupported language: {language}")
            return jsonify({"error": f"Invalid language. Available options: {list(TTS_MODELS.keys())}"}), 400

        if tts_models[language] is None:
            logger.error(f"❌ TTS model for {language} not loaded")
            return jsonify({"error": f"TTS model for {language} not available"}), 503

        logger.info(f"🔄 Generating TTS for language: {language}, text: '{text_input}'")

        try:
            processor = tts_processors[language]
            model = tts_models[language]
            inputs = processor(text_input, return_tensors="pt")
            inputs = {k: v.to(model.device) for k, v in inputs.items()}
        except Exception as e:
            logger.error(f"❌ TTS preprocessing failed: {str(e)}")
            return jsonify({"error": f"TTS preprocessing failed: {str(e)}"}), 500

        # Generate speech
        try:
            with torch.no_grad():
                output = model(**inputs).waveform
                waveform = output.squeeze().cpu().numpy()
        except Exception as e:
            logger.error(f"❌ TTS inference failed: {str(e)}")
            logger.debug(f"Stack trace: {traceback.format_exc()}")
            return jsonify({"error": f"TTS inference failed: {str(e)}"}), 500

        # Save to file with a unique name to prevent overwriting
        try:
            # Create a unique filename using timestamp and text hash
            import hashlib
            import time
            text_hash = hashlib.md5(text_input.encode()).hexdigest()[:8]
            timestamp = int(time.time())
            
            output_filename = os.path.join(output_dir, f"{language}_{text_hash}_{timestamp}.wav")
            sampling_rate = model.config.sampling_rate
            sf.write(output_filename, waveform, sampling_rate)
            logger.info(f"✅ Speech generated! File saved: {output_filename}")
        except Exception as e:
            logger.error(f"❌ Failed to save audio file: {str(e)}")
            return jsonify({"error": f"Failed to save audio file: {str(e)}"}), 500

        # Add cache-busting parameter to URL
        return jsonify({
            "message": "TTS audio generated",
            "file_url": f"/download/{os.path.basename(output_filename)}?t={timestamp}",
            "language": language,
            "text_length": len(text_input)
        })
    except Exception as e:
        logger.error(f"❌ Unhandled exception in TTS endpoint: {str(e)}")
        logger.debug(f"Stack trace: {traceback.format_exc()}")
        return jsonify({"error": f"Internal server error: {str(e)}"}), 500

def get_asr_model():
    return asr_model

def get_asr_processor():
    return asr_processor