import os
import torch
from transformers import AutoProcessor, SeamlessM4TForTextToText, SeamlessM4TProcessor
import gradio as gr

MODEL_NAME = "facebook/hf-seamless-m4t-medium"
device = "cuda" if torch.cuda.is_available() else "cpu"

processor = SeamlessM4TProcessor.from_pretrained(MODEL_NAME)
model     = SeamlessM4TForTextToText.from_pretrained(MODEL_NAME).to(device).eval()

def translate(text, sourceLang, targetLang, auto_detect):
    src = None if auto_detect else sourceLang
    inputs = processor(text=text, src_lang=src, return_tensors="pt").to(device)
    tokens = model.generate(**inputs, tgt_lang=targetLang)
    return processor.decode(tokens[0].tolist(), skip_special_tokens=True)

iface = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(label="Text to translate"),
        gr.Textbox(label="Source Language (e.g. eng)"),
        gr.Textbox(label="Target Language (e.g. fra)"),
        gr.Checkbox(label="Auto-detect source")
    ],
    outputs=gr.Textbox(label="Translated Text"),
    title="iVoice Translate"
).queue()  # <— this turns on the /run/predict REST API

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", share=True, server_port=int(os.environ.get("PORT", 7860)))