app.py

import gradio as gr
from transformers import pipeline, AutoProcessor
import torch

# 加载模型（第一次会自动下载，需接受许可）
model_id = "google/translategemma-4b-it"
pipe = pipeline(
    "image-text-to-text",
    model=model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto"  # 自动用 GPU 如果有
)

processor = AutoProcessor.from_pretrained(model_id)

def translate(text, image=None, source_lang="English", target_lang="Chinese"):
    # 构建提示（TranslateGemma 用特定格式提示翻译）
    prompt = f"Translate the following text from {source_lang} to {target_lang}:\n\n{text}"
    
    if image is not None:
        # 如果有图像，结合图像+文本
        inputs = processor(text=prompt, images=image, return_tensors="pt").to(pipe.device)
    else:
        inputs = processor(text=prompt, return_tensors="pt").to(pipe.device)
    
    outputs = pipe.generate(**inputs, max_new_tokens=512)
    translated = processor.decode(outputs[0], skip_special_tokens=True)
    
    return translated

demo = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(label="输入文本 (Input Text)"),
        gr.Image(label="可选：上传图像 (Optional Image for Visual Translation)"),
        gr.Dropdown(["English", "Chinese", "French", "Spanish", "Japanese"], label="源语言", value="English"),
        gr.Dropdown(["Chinese", "English", "French", "Spanish", "Japanese"], label="目标语言", value="Chinese")
    ],
    outputs=gr.Textbox(label="翻译结果 (Translated Text)"),
    title="TranslateGemma 4B - 55语言离线翻译 Demo",
    description="基于 Google TranslateGemma-4b-it，支持文本+图像翻译。模型完全本地运行在 HF Space。"
)

demo.launch()