app.py

import gradio as gr
import torch
import torchaudio
import numpy as np
from transformers import AutoProcessor, SeamlessM4Tv2Model

class SeamlessTranslator:
    def __init__(self):
        self.model_name = "facebook/seamless-m4t-v2-large"
        print("Loading model...")
        self.processor = AutoProcessor.from_pretrained(self.model_name)
        self.model = SeamlessM4Tv2Model.from_pretrained(self.model_name)
        self.sample_rate = self.model.config.sampling_rate

        self.languages = {
            "English": "eng",
            "Spanish": "spa",
            "French": "fra",
            "German": "deu",
            "Italian": "ita",
            "Portuguese": "por",
            "Russian": "rus",
            "Chinese": "cmn",
            "Japanese": "jpn",
            "Korean": "kor"
        }

    def translate_text(self, text, src_lang, tgt_lang, progress=gr.Progress()):
        progress(0.3, desc="Processing input...")
        try:
            inputs = self.processor(text=text, src_lang=self.languages[src_lang], return_tensors="pt")
            progress(0.6, desc="Generating audio...")
            audio_array = self.model.generate(**inputs, tgt_lang=self.languages[tgt_lang])[0].cpu().numpy().squeeze()
            progress(1.0, desc="Done!")
            return (self.sample_rate, audio_array)
        except Exception as e:
            raise gr.Error(str(e))

    def translate_audio(self, audio_path, tgt_lang, progress=gr.Progress()):
        progress(0.3, desc="Loading audio...")
        try:
            audio, orig_freq = torchaudio.load(audio_path)
            audio = torchaudio.functional.resample(audio, orig_freq=orig_freq, new_freq=16000)
            
            progress(0.6, desc="Translating...")
            inputs = self.processor(audios=audio, return_tensors="pt")
            audio_array = self.model.generate(**inputs, tgt_lang=self.languages[tgt_lang])[0].cpu().numpy().squeeze()
            progress(1.0, desc="Done!")
            return (self.sample_rate, audio_array)
        except Exception as e:
            raise gr.Error(str(e))

css = """
#component-0 {
    max-width: 1200px;
    margin: auto;
    padding: 20px;
}

.container {
    border-radius: 12px;
    padding: 20px;
}

.gr-form {
    border-color: #e5e7eb !important;
}

.gr-button {
    border-radius: 8px !important;
    background: linear-gradient(to right, #2563eb, #4f46e5) !important;
    color: white !important;
    font-weight: 600 !important;
}

.gr-button:hover {
    box-shadow: 0 4px 6px -1px rgb(0 0 0 / 0.1) !important;
    transform: translateY(-1px);
}

.gr-input, .gr-select {
    border-radius: 8px !important;
}

.gr-panel {
    border-radius: 12px !important;
}

.title {
    text-align: center;
    font-size: 2.5rem;
    font-weight: bold;
    margin: 1rem 0;
    background: linear-gradient(to right, #2563eb, #4f46e5);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}

.subtitle {
    text-align: center;
    color: #6b7280;
    margin-bottom: 2rem;
}

.tab-nav {
    border-bottom: 2px solid #e5e7eb;
    margin-bottom: 1rem;
}

.output-label {
    font-weight: 600;
    color: #374151;
    margin-bottom: 0.5rem;
}

.footer {
    text-align: center;
    margin-top: 2rem;
    padding-top: 1rem;
    border-top: 1px solid #e5e7eb;
    color: #6b7280;
    font-size: 0.875rem;
}
"""

def create_ui():
    translator = SeamlessTranslator()

    with gr.Blocks(css=css, title="A.R.I.S. Translator") as demo:
        gr.HTML(
            """
            <div class="title">A.R.I.S. Translator</div>
            <div class="subtitle">Advanced Real-time Interpretation System</div>
            """
        )

        with gr.Tabs() as tabs:
            # Text to Speech Tab
            with gr.Tab("Text Translation", id=1):
                with gr.Row():
                    with gr.Column():
                        text_input = gr.Textbox(
                            label="Text to Translate",
                            placeholder="Enter your text here...",
                            lines=5
                        )
                        with gr.Row():
                            src_lang = gr.Dropdown(
                                choices=list(translator.languages.keys()),
                                value="English",
                                label="Source Language"
                            )
                            tgt_lang = gr.Dropdown(
                                choices=list(translator.languages.keys()),
                                value="Spanish",
                                label="Target Language"
                            )
                        translate_btn = gr.Button("Translate", variant="primary")

                    with gr.Column():
                        gr.HTML('<div class="output-label">Translation Output</div>')
                        audio_output = gr.Audio(
                            label="Translated Audio",
                            type="numpy"
                        )

            # Audio to Speech Tab
            with gr.Tab("Audio Translation", id=2):
                with gr.Row():
                    with gr.Column():
                        audio_input = gr.Audio(
                            label="Upload Audio",
                            type="filepath"
                        )
                        tgt_lang_audio = gr.Dropdown(
                            choices=list(translator.languages.keys()),
                            value="English",
                            label="Target Language"
                        )
                        translate_audio_btn = gr.Button("Translate Audio", variant="primary")

                    with gr.Column():
                        gr.HTML('<div class="output-label">Translation Output</div>')
                        audio_output_from_audio = gr.Audio(
                            label="Translated Audio",
                            type="numpy"
                        )

        gr.HTML(
            """
            <div class="footer">
                Powered by Meta's SeamlessM4T model | Built with Gradio
            </div>
            """
        )

        # Event handlers
        translate_btn.click(
            fn=translator.translate_text,
            inputs=[text_input, src_lang, tgt_lang],
            outputs=audio_output
        )

        translate_audio_btn.click(
            fn=translator.translate_audio,
            inputs=[audio_input, tgt_lang_audio],
            outputs=audio_output_from_audio
        )

    return demo

if __name__ == "__main__":
    demo = create_ui()
    demo.queue()
    demo.launch()