import gradio as gr
from transformers import pipeline
import torch
from PIL import Image
import requests
from io import BytesIO
import os
# Language options
LANGUAGES = {
"English": "en",
"Spanish": "es",
"German": "de",
"French": "fr",
"Italian": "it",
"Portuguese": "pt",
"Hindi": "hi",
"Chinese": "zh",
"Japanese": "ja",
"Korean": "ko",
"Arabic": "ar",
"Russian": "ru",
"Dutch": "nl",
"Polish": "pl",
"Turkish": "tr",
"Vietnamese": "vi",
"Thai": "th",
"Indonesian": "id",
"Bengali": "bn",
"Tamil": "ta",
}
# Load model
print("Loading TranslateGemma model...")
pipe = pipeline(
"image-text-to-text",
model="google/translategemma-4b-it",
device="cuda" if torch.cuda.is_available() else "cpu",
torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
token=os.environ.get("HF_TOKEN"),
)
print("Model loaded!")
def translate_text(text, source_lang, target_lang):
"""Translate text from source to target language"""
if not text.strip():
return "Please enter text to translate."
src_code = LANGUAGES.get(source_lang, "en")
tgt_code = LANGUAGES.get(target_lang, "es")
messages = [
{
"role": "user",
"content": [
{
"type": "text",
"source_lang_code": src_code,
"target_lang_code": tgt_code,
"text": text,
}
],
}
]
output = pipe(text=messages, max_new_tokens=1000, generate_kwargs={"do_sample": False})
return output[0]["generated_text"][-1]["content"]
def translate_image(image_url, source_lang, target_lang):
"""Translate text in an image"""
if not image_url.strip():
return "Please enter an image URL."
src_code = LANGUAGES.get(source_lang, "en")
tgt_code = LANGUAGES.get(target_lang, "de")
messages = [
{
"role": "user",
"content": [
{
"type": "image",
"source_lang_code": src_code,
"target_lang_code": tgt_code,
"url": image_url,
},
],
}
]
output = pipe(text=messages, max_new_tokens=500, generate_kwargs={"do_sample": False})
return output[0]["generated_text"][-1]["content"]
def load_image_preview(url):
"""Load image from URL for preview"""
if not url.strip():
return None
try:
response = requests.get(url, timeout=10)
img = Image.open(BytesIO(response.content))
return img
except:
return None
def translate_batch(texts, source_lang, target_lang):
"""Translate multiple texts"""
if not texts.strip():
return "Please enter texts to translate (one per line)."
lines = [line.strip() for line in texts.split("\n") if line.strip()]
results = []
for line in lines:
result = translate_text(line, source_lang, target_lang)
results.append(result)
return "\n".join(results)
# Custom CSS
css = """
.gradio-container {
font-family: 'Inter', sans-serif;
}
.header {
text-align: center;
margin-bottom: 20px;
}
.footer {
text-align: center;
margin-top: 20px;
padding: 10px;
color: #666;
}
"""
# Build Gradio interface
with gr.Blocks(css=css, title="TranslateGemma - AI Translator", theme=gr.themes.Soft()) as demo:
# Header
gr.HTML("""
""")
with gr.Tabs():
# Tab 1: Text Translation
with gr.TabItem("📝 Text Translation"):
with gr.Row():
with gr.Column():
text_source = gr.Dropdown(
choices=list(LANGUAGES.keys()),
value="English",
label="From Language"
)
text_input = gr.Textbox(
label="Enter text to translate",
placeholder="Type or paste your text here...",
lines=6
)
with gr.Column():
text_target = gr.Dropdown(
choices=list(LANGUAGES.keys()),
value="Spanish",
label="To Language"
)
text_output = gr.Textbox(
label="Translation",
lines=6,
interactive=False
)
text_btn = gr.Button("🔄 Translate", variant="primary")
text_btn.click(
translate_text,
inputs=[text_input, text_source, text_target],
outputs=text_output
)
# Examples
gr.Examples(
examples=[
["Hello, how are you today?", "English", "Spanish"],
["The weather is beautiful this morning.", "English", "French"],
["Artificial Intelligence is changing the world.", "English", "Hindi"],
],
inputs=[text_input, text_source, text_target],
)
# Tab 2: Image Translation
with gr.TabItem("🖼️ Image Translation"):
with gr.Row():
with gr.Column():
img_source = gr.Dropdown(
choices=list(LANGUAGES.keys()),
value="English",
label="From Language"
)
img_url = gr.Textbox(
label="Image URL",
placeholder="https://example.com/image.png"
)
img_preview = gr.Image(label="Image Preview", type="pil")
img_url.change(load_image_preview, inputs=img_url, outputs=img_preview)
with gr.Column():
img_target = gr.Dropdown(
choices=list(LANGUAGES.keys()),
value="German",
label="To Language"
)
img_output = gr.Textbox(
label="Translation",
lines=6,
interactive=False
)
img_btn = gr.Button("🔄 Translate Image", variant="primary")
img_btn.click(
translate_image,
inputs=[img_url, img_source, img_target],
outputs=img_output
)
gr.Examples(
examples=[
["https://raw.githubusercontent.com/esalesky/vistra-benchmark/refs/heads/main/images/f488c322.png", "English", "German"],
],
inputs=[img_url, img_source, img_target],
)
# Tab 3: Batch Translation
with gr.TabItem("📚 Batch Translation"):
with gr.Row():
with gr.Column():
batch_source = gr.Dropdown(
choices=list(LANGUAGES.keys()),
value="English",
label="From Language"
)
batch_input = gr.Textbox(
label="Enter texts (one per line)",
placeholder="Hello\nHow are you?\nGoodbye",
lines=8
)
with gr.Column():
batch_target = gr.Dropdown(
choices=list(LANGUAGES.keys()),
value="Spanish",
label="To Language"
)
batch_output = gr.Textbox(
label="Translations",
lines=8,
interactive=False
)
batch_btn = gr.Button("🔄 Translate All", variant="primary")
batch_btn.click(
translate_batch,
inputs=[batch_input, batch_source, batch_target],
outputs=batch_output
)
# Footer
gr.HTML("""
""")
# Launch
if __name__ == "__main__":
demo.launch()