Expose the Hugging Face Code as an API

app.py

@@ -1,8 +1,10 @@
 import os
 import torch
-import gradio as gr
 import numpy as np
 import soundfile as sf
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
 from transformers import (
     AutoModelForSeq2SeqLM,
     AutoTokenizer,
@@ -13,13 +15,11 @@ from transformers import (
     WhisperForConditionalGeneration
 )
 from typing import Optional, Tuple, Dict, List
+import base64
+import io
 
+# Your existing TalklasTranslator class (unchanged)
 class TalklasTranslator:
-    """
-    Speech-to-Speech translation pipeline for Philippine languages.
-    Uses MMS/Whisper for STT, NLLB for MT, and MMS for TTS.
-    """
-
     LANGUAGE_MAPPING = {
         "English": "eng",
         "Tagalog": "tgl",
@@ -50,45 +50,34 @@ class TalklasTranslator:
         self.sample_rate = 16000
 
         print(f"Initializing Talklas Translator on {self.device}")
-
-        # Initialize models
         self._initialize_stt_model()
         self._initialize_mt_model()
         self._initialize_tts_model()
 
     def _initialize_stt_model(self):
-        """Initialize speech-to-text model with fallback to Whisper"""
         try:
             print("Loading STT model...")
             try:
-                # Try loading MMS model first
                 self.stt_processor = AutoProcessor.from_pretrained("facebook/mms-1b-all")
                 self.stt_model = AutoModelForCTC.from_pretrained("facebook/mms-1b-all")
-
-                # Set language if available
                 if self.source_lang in self.stt_processor.tokenizer.vocab.keys():
                     self.stt_processor.tokenizer.set_target_lang(self.source_lang)
                     self.stt_model.load_adapter(self.source_lang)
                     print(f"Loaded MMS STT model for {self.source_lang}")
                 else:
                     print(f"Language {self.source_lang} not in MMS, using default")
-
             except Exception as mms_error:
                 print(f"MMS loading failed: {mms_error}")
-                # Fallback to Whisper
                 print("Loading Whisper as fallback...")
                 self.stt_processor = WhisperProcessor.from_pretrained("openai/whisper-small")
                 self.stt_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
                 print("Loaded Whisper STT model")
-
             self.stt_model.to(self.device)
-
         except Exception as e:
             print(f"STT model initialization failed: {e}")
             raise RuntimeError("Could not initialize STT model")
 
     def _initialize_mt_model(self):
-        """Initialize machine translation model"""
         try:
             print("Loading NLLB Translation model...")
             self.mt_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
@@ -100,7 +89,6 @@ class TalklasTranslator:
             raise
 
     def _initialize_tts_model(self):
-        """Initialize text-to-speech model"""
         try:
             print("Loading TTS model...")
             try:
@@ -112,102 +100,78 @@ class TalklasTranslator:
                 print("Falling back to English TTS")
                 self.tts_model = VitsModel.from_pretrained("facebook/mms-tts-eng")
                 self.tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng")
-
             self.tts_model.to(self.device)
         except Exception as e:
             print(f"TTS model initialization failed: {e}")
             raise
 
     def update_languages(self, source_lang: str, target_lang: str) -> str:
-        """Update languages and reinitialize models if needed"""
         if source_lang == self.source_lang and target_lang == self.target_lang:
             return "Languages already set"
-
         self.source_lang = source_lang
         self.target_lang = target_lang
-
-        # Only reinitialize models that depend on language
         self._initialize_stt_model()
         self._initialize_tts_model()
-
         return f"Languages updated to {source_lang} → {target_lang}"
 
     def speech_to_text(self, audio_path: str) -> str:
-        """Convert speech to text using loaded STT model"""
         try:
             waveform, sample_rate = sf.read(audio_path)
-
             if sample_rate != 16000:
                 import librosa
                 waveform = librosa.resample(waveform, orig_sr=sample_rate, target_sr=16000)
-
             inputs = self.stt_processor(
                 waveform,
                 sampling_rate=16000,
                 return_tensors="pt"
             ).to(self.device)
-
             with torch.no_grad():
-                if isinstance(self.stt_model, WhisperForConditionalGeneration):  # Whisper model
+                if isinstance(self.stt_model, WhisperForConditionalGeneration):
                     generated_ids = self.stt_model.generate(**inputs)
                     transcription = self.stt_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-                else:  # MMS model (Wav2Vec2ForCTC)
+                else:
                     logits = self.stt_model(**inputs).logits
                     predicted_ids = torch.argmax(logits, dim=-1)
                     transcription = self.stt_processor.batch_decode(predicted_ids)[0]
-
             return transcription
-
         except Exception as e:
             print(f"Speech recognition failed: {e}")
             raise RuntimeError("Speech recognition failed")
 
     def translate_text(self, text: str) -> str:
-        """Translate text using NLLB model"""
         try:
             source_code = self.NLLB_LANGUAGE_CODES[self.source_lang]
             target_code = self.NLLB_LANGUAGE_CODES[self.target_lang]
-
             self.mt_tokenizer.src_lang = source_code
             inputs = self.mt_tokenizer(text, return_tensors="pt").to(self.device)
-
             with torch.no_grad():
                 generated_tokens = self.mt_model.generate(
                     **inputs,
                     forced_bos_token_id=self.mt_tokenizer.convert_tokens_to_ids(target_code),
                     max_length=448
                 )
-
             return self.mt_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
-
         except Exception as e:
             print(f"Translation failed: {e}")
             raise RuntimeError("Text translation failed")
 
     def text_to_speech(self, text: str) -> Tuple[int, np.ndarray]:
-        """Convert text to speech"""
         try:
             inputs = self.tts_tokenizer(text, return_tensors="pt").to(self.device)
-
             with torch.no_grad():
                 output = self.tts_model(**inputs)
-
             speech = output.waveform.cpu().numpy().squeeze()
             speech = (speech * 32767).astype(np.int16)
-
             return self.tts_model.config.sampling_rate, speech
-
         except Exception as e:
             print(f"Speech synthesis failed: {e}")
             raise RuntimeError("Speech synthesis failed")
 
     def translate_speech(self, audio_path: str) -> Dict:
-        """Full speech-to-speech translation"""
         try:
             source_text = self.speech_to_text(audio_path)
             translated_text = self.translate_text(source_text)
             sample_rate, audio = self.text_to_speech(translated_text)
-
             return {
                 "source_text": source_text,
                 "translated_text": translated_text,
@@ -223,11 +187,9 @@ class TalklasTranslator:
             }
 
     def translate_text_only(self, text: str) -> Dict:
-        """Text-to-speech translation"""
         try:
             translated_text = self.translate_text(text)
             sample_rate, audio = self.text_to_speech(translated_text)
-
             return {
                 "source_text": text,
                 "translated_text": translated_text,
@@ -251,168 +213,88 @@ class TranslatorSingleton:
             cls._instance = TalklasTranslator()
         return cls._instance
 
-def process_audio(audio_path, source_lang, target_lang):
-    """Process audio through the full translation pipeline"""
-    # Validate input
-    if not audio_path:
-        return None, "No audio provided", "No translation available", "Please provide audio input"
-
-    # Update languages
-    source_code = TalklasTranslator.LANGUAGE_MAPPING[source_lang]
-    target_code = TalklasTranslator.LANGUAGE_MAPPING[target_lang]
-
-    translator = TranslatorSingleton.get_instance()
-    status = translator.update_languages(source_code, target_code)
-
-    # Process the audio
-    results = translator.translate_speech(audio_path)
-
-    return results["output_audio"], results["source_text"], results["translated_text"], results["performance"]
-
-def process_text(text, source_lang, target_lang):
-    """Process text through the translation pipeline"""
-    # Validate input
-    if not text:
-        return None, "No text provided", "No translation available", "Please provide text input"
-
-    # Update languages
-    source_code = TalklasTranslator.LANGUAGE_MAPPING[source_lang]
-    target_code = TalklasTranslator.LANGUAGE_MAPPING[target_lang]
-
-    translator = TranslatorSingleton.get_instance()
-    status = translator.update_languages(source_code, target_code)
-
-    # Process the text
-    results = translator.translate_text_only(text)
-
-    return results["output_audio"], results["source_text"], results["translated_text"], results["performance"]
-
-def create_gradio_interface():
-    """Create and launch Gradio interface"""
-    # Define language options
-    languages = list(TalklasTranslator.LANGUAGE_MAPPING.keys())
-
-    # Define the interface
-    demo = gr.Blocks(title="Talklas - Speech & Text Translation")
-
-    with demo:
-        gr.Markdown("# Talklas: Speech-to-Speech Translation System")
-        gr.Markdown("### Translate between Philippine Languages and English")
-
-        with gr.Row():
-            with gr.Column():
-                source_lang = gr.Dropdown(
-                    choices=languages,
-                    value="English",
-                    label="Source Language"
-                )
-
-                target_lang = gr.Dropdown(
-                    choices=languages,
-                    value="Tagalog",
-                    label="Target Language"
-                )
-
-                language_status = gr.Textbox(label="Language Status")
-                update_btn = gr.Button("Update Languages")
-
-        with gr.Tabs():
-            with gr.TabItem("Audio Input"):
-                with gr.Row():
-                    with gr.Column():
-                        gr.Markdown("### Audio Input")
-                        audio_input = gr.Audio(
-                            type="filepath",
-                            label="Upload Audio File"
-                        )
-                        audio_translate_btn = gr.Button("Translate Audio", variant="primary")
-
-                    with gr.Column():
-                        gr.Markdown("### Output")
-                        audio_output = gr.Audio(
-                            label="Translated Speech",
-                            type="numpy",
-                            autoplay=True
-                        )
-
-            with gr.TabItem("Text Input"):
-                with gr.Row():
-                    with gr.Column():
-                        gr.Markdown("### Text Input")
-                        text_input = gr.Textbox(
-                            label="Enter text to translate",
-                            lines=3
-                        )
-                        text_translate_btn = gr.Button("Translate Text", variant="primary")
-
-                    with gr.Column():
-                        gr.Markdown("### Output")
-                        text_output = gr.Audio(
-                            label="Translated Speech",
-                            type="numpy",
-                            autoplay=True
-                        )
-
-        with gr.Row():
-            with gr.Column():
-                source_text = gr.Textbox(label="Source Text")
-                translated_text = gr.Textbox(label="Translated Text")
-                performance_info = gr.Textbox(label="Performance Metrics")
-
-        # Set up events
-        update_btn.click(
-            lambda source_lang, target_lang: TranslatorSingleton.get_instance().update_languages(
-                TalklasTranslator.LANGUAGE_MAPPING[source_lang],
-                TalklasTranslator.LANGUAGE_MAPPING[target_lang]
-            ),
-            inputs=[source_lang, target_lang],
-            outputs=[language_status]
-        )
-
-        # Audio translate button click
-        audio_translate_btn.click(
-            process_audio,
-            inputs=[audio_input, source_lang, target_lang],
-            outputs=[audio_output, source_text, translated_text, performance_info]
-        ).then(
-            None,
-            None,
-            None,
-            js="""() => {
-                const audioElements = document.querySelectorAll('audio');
-                if (audioElements.length > 0) {
-                    const lastAudio = audioElements[audioElements.length - 1];
-                    lastAudio.play().catch(error => {
-                        console.warn('Autoplay failed:', error);
-                        alert('Audio may require user interaction to play');
-                    });
-                }
-            }"""
-        )
-
-        # Text translate button click
-        text_translate_btn.click(
-            process_text,
-            inputs=[text_input, source_lang, target_lang],
-            outputs=[text_output, source_text, translated_text, performance_info]
-        ).then(
-            None,
-            None,
-            None,
-            js="""() => {
-                const audioElements = document.querySelectorAll('audio');
-                if (audioElements.length > 0) {
-                    const lastAudio = audioElements[audioElements.length - 1];
-                    lastAudio.play().catch(error => {
-                        console.warn('Autoplay failed:', error);
-                        alert('Audio may require user interaction to play');
-                    });
-                }
-            }"""
-        )
-
-    return demo
+# FastAPI application
+app = FastAPI(title="Talklas API", description="Speech-to-Speech Translation API")
+
+class TranslationRequest(BaseModel):
+    source_lang: str
+    target_lang: str
+    text: Optional[str] = None
+
+@app.post("/translate/audio")
+async def translate_audio(file: UploadFile = File(...), source_lang: str = "English", target_lang: str = "Tagalog"):
+    try:
+        # Validate languages
+        if source_lang not in TalklasTranslator.LANGUAGE_MAPPING or target_lang not in TalklasTranslator.LANGUAGE_MAPPING:
+            raise HTTPException(status_code=400, detail="Invalid language selection")
+
+        # Save uploaded audio file temporarily
+        audio_path = f"temp_{file.filename}"
+        with open(audio_path, "wb") as f:
+            f.write(await file.read())
+
+        # Update languages
+        source_code = TalklasTranslator.LANGUAGE_MAPPING[source_lang]
+        target_code = TalklasTranslator.LANGUAGE_MAPPING[target_lang]
+        translator = TranslatorSingleton.get_instance()
+        translator.update_languages(source_code, target_code)
+
+        # Process the audio
+        results = translator.translate_speech(audio_path)
+
+        # Clean up temporary file
+        os.remove(audio_path)
+
+        # Convert audio to base64 for response
+        sample_rate, audio = results["output_audio"]
+        buffer = io.BytesIO()
+        sf.write(buffer, audio, sample_rate, format="wav")
+        audio_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+        return JSONResponse(content={
+            "source_text": results["source_text"],
+            "translated_text": results["translated_text"],
+            "audio_base64": audio_base64,
+            "sample_rate": sample_rate,
+            "performance": results["performance"]
+        })
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Translation failed: {str(e)}")
+
+@app.post("/translate/text")
+async def translate_text(request: TranslationRequest):
+    try:
+        # Validate input
+        if not request.text:
+            raise HTTPException(status_code=400, detail="Text input is required")
+        if request.source_lang not in TalklasTranslator.LANGUAGE_MAPPING or request.target_lang not in TalklasTranslator.LANGUAGE_MAPPING:
+            raise HTTPException(status_code=400, detail="Invalid language selection")
+
+        # Update languages
+        source_code = TalklasTranslator.LANGUAGE_MAPPING[request.source_lang]
+        target_code = TalklasTranslator.LANGUAGE_MAPPING[request.target_lang]
+        translator = TranslatorSingleton.get_instance()
+        translator.update_languages(source_code, target_code)
+
+        # Process the text
+        results = translator.translate_text_only(request.text)
+
+        # Convert audio to base64 for response
+        sample_rate, audio = results["output_audio"]
+        buffer = io.BytesIO()
+        sf.write(buffer, audio, sample_rate, format="wav")
+        audio_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+        return JSONResponse(content={
+            "source_text": results["source_text"],
+            "translated_text": results["translated_text"],
+            "audio_base64": audio_base64,
+            "sample_rate": sample_rate,
+            "performance": results["performance"]
+        })
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Translation failed: {str(e)}")
 
 if __name__ == "__main__":
-    demo = create_gradio_interface()
-    demo.launch(share=True, debug=True)
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)