Update app.py

app.py

@@ -1,60 +1,46 @@
+import gradio as gr
 import torch
 import torchaudio
 import scipy.io.wavfile
+import numpy as np
 from transformers import AutoProcessor, SeamlessM4Tv2Model
 from pathlib import Path
 from typing import Optional, Union
 
 class SeamlessTranslator:
-    """
-    A wrapper class for Facebook's SeamlessM4T translation model.
-    Handles both text-to-speech and speech-to-speech translation.
-    """
-    
     def __init__(self, model_name: str = "facebook/seamless-m4t-v2-large"):
-        """
-        Initialize the translator with the specified model.
-        
-        Args:
-            model_name (str): Name of the model to use
-        """
         try:
             self.processor = AutoProcessor.from_pretrained(model_name)
             self.model = SeamlessM4Tv2Model.from_pretrained(model_name)
             self.sample_rate = self.model.config.sampling_rate
         except Exception as e:
             raise RuntimeError(f"Failed to initialize model: {str(e)}")
-
-    def translate_text(self, text: str, src_lang: str, tgt_lang: str) -> numpy.ndarray:
-        """
-        Translate text to speech in the target language.
         
-        Args:
-            text (str): Input text to translate
-            src_lang (str): Source language code (e.g., 'eng')
-            tgt_lang (str): Target language code (e.g., 'rus')
-            
-        Returns:
-            numpy.ndarray: Audio waveform array
-        """
+        # Available language pairs
+        self.language_codes = {
+            "English": "eng",
+            "Spanish": "spa",
+            "French": "fra",
+            "German": "deu",
+            "Italian": "ita",
+            "Portuguese": "por",
+            "Russian": "rus",
+            "Chinese": "cmn",
+            "Japanese": "jpn",
+            "Korean": "kor",
+            "Arabic": "ara",
+            "Hindi": "hin",
+        }
+
+    def translate_text(self, text: str, src_lang: str, tgt_lang: str) -> tuple[int, np.ndarray]:
         try:
             inputs = self.processor(text=text, src_lang=src_lang, return_tensors="pt")
             audio_array = self.model.generate(**inputs, tgt_lang=tgt_lang)[0].cpu().numpy().squeeze()
-            return audio_array
+            return self.sample_rate, audio_array
         except Exception as e:
             raise RuntimeError(f"Text translation failed: {str(e)}")
 
-    def translate_audio(self, audio_path: Union[str, Path], tgt_lang: str) -> numpy.ndarray:
-        """
-        Translate audio to speech in the target language.
-        
-        Args:
-            audio_path (str or Path): Path to input audio file
-            tgt_lang (str): Target language code (e.g., 'rus')
-            
-        Returns:
-            numpy.ndarray: Audio waveform array
-        """
+    def translate_audio(self, audio_path: str, tgt_lang: str) -> tuple[int, np.ndarray]:
         try:
             # Load and resample audio
             audio, orig_freq = torchaudio.load(audio_path)
@@ -67,50 +53,103 @@ class SeamlessTranslator:
             # Process and generate translation
             inputs = self.processor(audios=audio, return_tensors="pt")
             audio_array = self.model.generate(**inputs, tgt_lang=tgt_lang)[0].cpu().numpy().squeeze()
-            return audio_array
+            return self.sample_rate, audio_array
         except Exception as e:
             raise RuntimeError(f"Audio translation failed: {str(e)}")
 
-    def save_audio(self, audio_array: numpy.ndarray, output_path: Union[str, Path]) -> None:
-        """
-        Save an audio array to a WAV file.
-        
-        Args:
-            audio_array (numpy.ndarray): Audio data to save
-            output_path (str or Path): Path where to save the WAV file
-        """
-        try:
-            scipy.io.wavfile.write(
-                output_path,
-                rate=self.sample_rate,
-                data=audio_array
-            )
-        except Exception as e:
-            raise RuntimeError(f"Failed to save audio: {str(e)}")
+class GradioInterface:
+    def __init__(self):
+        self.translator = SeamlessTranslator()
+        self.languages = list(self.translator.language_codes.keys())
+
+    def text_to_speech(self, text: str, src_lang: str, tgt_lang: str) -> tuple[int, np.ndarray]:
+        src_code = self.translator.language_codes[src_lang]
+        tgt_code = self.translator.language_codes[tgt_lang]
+        return self.translator.translate_text(text, src_code, tgt_code)
+
+    def speech_to_speech(self, audio_path: str, tgt_lang: str) -> tuple[int, np.ndarray]:
+        tgt_code = self.translator.language_codes[tgt_lang]
+        return self.translator.translate_audio(audio_path, tgt_code)
 
-def main():
-    """Example usage of the SeamlessTranslator class."""
-    try:
-        # Initialize translator
-        translator = SeamlessTranslator()
+    def launch(self):
+        # Create the Gradio interface
+        with gr.Blocks(title="SeamlessM4T Translator") as demo:
+            gr.Markdown("# 🌐 SeamlessM4T Translator")
+            gr.Markdown("Translate text or speech to different languages using Meta's SeamlessM4T model")
 
-        # Example text translation
-        text_audio = translator.translate_text(
-            text="Hello, my dog is cute",
-            src_lang="eng",
-            tgt_lang="rus"
-        )
-        translator.save_audio(text_audio, "output_from_text.wav")
+            with gr.Tabs():
+                # Text-to-Speech tab
+                with gr.TabItem("Text to Speech"):
+                    with gr.Row():
+                        with gr.Column():
+                            text_input = gr.Textbox(
+                                label="Input Text",
+                                placeholder="Enter text to translate...",
+                                lines=3
+                            )
+                            src_lang = gr.Dropdown(
+                                choices=self.languages,
+                                value="English",
+                                label="Source Language"
+                            )
+                            tgt_lang_text = gr.Dropdown(
+                                choices=self.languages,
+                                value="Spanish",
+                                label="Target Language"
+                            )
+                            translate_btn = gr.Button("Translate", variant="primary")
+                        
+                        with gr.Column():
+                            audio_output = gr.Audio(
+                                label="Translated Speech",
+                                type="numpy"
+                            )
 
-        # Example audio translation
-        audio_audio = translator.translate_audio(
-            audio_path="input_audio.wav",
-            tgt_lang="rus"
-        )
-        translator.save_audio(audio_audio, "output_from_audio.wav")
+                    translate_btn.click(
+                        fn=self.text_to_speech,
+                        inputs=[text_input, src_lang, tgt_lang_text],
+                        outputs=audio_output
+                    )
+
+                # Speech-to-Speech tab
+                with gr.TabItem("Speech to Speech"):
+                    with gr.Row():
+                        with gr.Column():
+                            audio_input = gr.Audio(
+                                label="Input Speech",
+                                type="filepath"
+                            )
+                            tgt_lang_speech = gr.Dropdown(
+                                choices=self.languages,
+                                value="Spanish",
+                                label="Target Language"
+                            )
+                            translate_audio_btn = gr.Button("Translate", variant="primary")
+                        
+                        with gr.Column():
+                            audio_output_s2s = gr.Audio(
+                                label="Translated Speech",
+                                type="numpy"
+                            )
+
+                    translate_audio_btn.click(
+                        fn=self.speech_to_speech,
+                        inputs=[audio_input, tgt_lang_speech],
+                        outputs=audio_output_s2s
+                    )
+
+            gr.Markdown(
+                """
+                ### Notes
+                - Text-to-Speech: Enter text and select source/target languages
+                - Speech-to-Speech: Upload an audio file and select target language
+                - Processing may take a few moments depending on input length
+                """
+            )
 
-    except Exception as e:
-        print(f"Translation failed: {str(e)}")
+        # Launch the interface
+        demo.launch(share=True)
 
 if __name__ == "__main__":
-    main()
+    interface = GradioInterface()
+    interface.launch()