Update app.py

app.py

@@ -7,17 +7,34 @@
 # demo.launch()
 import requests
 from PIL import Image
-from transformers import AutoModelForCausalLM, AutoProcessor
+from transformers import AutoModelForCausalLM, AutoProcessor, MarianMTModel, MarianTokenizer
 import torch
 import gradio as gr
 
 # 设置设备
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# 加载模型和处理器
+# 加载 Florence-2 模型和处理器
 model = AutoModelForCausalLM.from_pretrained("MiaoshouAI/Florence-2-base-PromptGen-v1.5", trust_remote_code=True).to(device)
 processor = AutoProcessor.from_pretrained("MiaoshouAI/Florence-2-base-PromptGen-v1.5", trust_remote_code=True)
 
+# 加载 Helsinki-NLP 的翻译模型（英文到中文）
+translation_model_name = "Helsinki-NLP/opus-mt-en-zh"
+translation_tokenizer = MarianTokenizer.from_pretrained(translation_model_name)
+translation_model = MarianMTModel.from_pretrained(translation_model_name).to(device)
+
+# 翻译函数
+def translate_to_chinese(text):
+    try:
+        # 分词和翻译
+        tokenized_text = translation_tokenizer(text, return_tensors="pt", max_length=512, truncation=True).to(device)
+        translated_tokens = translation_model.generate(**tokenized_text)
+        translated_text = translation_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
+        return translated_text
+    except Exception as e:
+        return f"Translation error: {str(e)}"
+
+# 生成描述并翻译
 def generate_caption(image_url):
     try:
         # 下载并打开图像
@@ -40,7 +57,10 @@ def generate_caption(image_url):
         # 解析生成的文本
         parsed_answer = processor.post_process_generation(generated_text, task=prompt, image_size=(image.width, image.height))
         
-        return parsed_answer
+        # 翻译成中文
+        translated_answer = translate_to_chinese(parsed_answer)
+        
+        return translated_answer
     except Exception as e:
         return f"Error: {str(e)}"
 
@@ -53,9 +73,9 @@ def gradio_interface(image_url):
 iface = gr.Interface(
     fn=gradio_interface,  # 处理函数
     inputs=gr.Textbox(label="Image URL", placeholder="Enter the URL of the image..."),  # 输入组件
-    outputs=gr.Textbox(label="Generated Caption"),  # 输出组件
+    outputs=gr.Textbox(label="Generated Caption (Translated to Chinese)"),  # 输出组件
     title="Florence-2 Prompt Generation",  # 标题
-    description="Generate detailed captions for images using Florence-2 model.",  # 描述
+    description="Generate detailed captions for images using Florence-2 model and translate them to Chinese.",  # 描述
     examples=[
         ["https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true"]
     ]  # 示例