app.py

# import gradio as gr

# def greet(name):
#     return "Hello " + name + "!!"

# demo = gr.Interface(fn=greet, inputs="text", outputs="text")
# demo.launch()
import requests
from PIL import Image
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:
        # 下载并打开图像
        image = Image.open(requests.get(image_url, stream=True).raw)
        
        # 准备输入
        prompt = "<MORE_DETAILED_CAPTION>"
        inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
        
        # 生成文本
        generated_ids = model.generate(
            input_ids=inputs["input_ids"],
            pixel_values=inputs["pixel_values"],
            max_new_tokens=1024,
            do_sample=False,
            num_beams=3
        )
        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
        
        # 解析生成的文本
        parsed_answer = processor.post_process_generation(generated_text, task=prompt, image_size=(image.width, image.height))
        
        # 翻译成中文
        translated_answer = translate_to_chinese(parsed_answer)
        
        return translated_answer
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio 界面
def gradio_interface(image_url):
    result = generate_caption(image_url)
    return result

# 创建 Gradio 应用
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 (Translated to Chinese)"),  # 输出组件
    title="Florence-2 Prompt Generation",  # 标题
    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"]
    ]  # 示例
)

# 启动 Gradio 应用
iface.launch()