uodate app py

app.py

@@ -1,23 +1,41 @@
 import os
-# 强制安装依赖
-os.system("pip install gradio==4.0.0 transformers==4.30.2 torch==1.13.1 pillow==9.4.0")
-
+os.system("pip install gradio==4.0.0 transformers==4.36.2 torch==2.0.1 pillow==9.4.0 accelerate==0.30.0 bitsandbytes==0.43.0")
 import gradio as gr
-from transformers import pipeline
+import torch
 from PIL import Image
+from transformers import AutoProcessor, AutoModelForVisionAndLanguageGeneration
 
-# 加载模型
-image_to_text = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
+# 加载DAM-3B模型和处理器
+model_name = "nvidia/DAM-3B"
+processor = AutoProcessor.from_pretrained(model_name)
+# 4-bit量化适配免费空间，降低内存占用
+model = AutoModelForVisionAndLanguageGeneration.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16,
+    load_in_4bit=True,
+    device_map="auto"
+)
 
-def generate_caption(image):
-    result = image_to_text(image)[0]
-    return f"图像描述：{result['generated_text']}"
+def generate_detailed_caption(image):
+    # 模拟全图描述（DAM-3B支持区域指定，此处简化为全图细节描述）
+    inputs = processor(images=image, return_tensors="pt").to(model.device, torch.float16)
+    # 生成详细描述，设置长文本参数
+    outputs = model.generate(
+        **inputs,
+        max_length=200,  # 延长描述长度，保留更多细节
+        num_beams=4,     # 束搜索提升描述连贯性
+        no_repeat_ngram_size=3,  # 避免重复内容
+        early_stopping=True
+    )
+    caption = processor.decode(outputs[0], skip_special_tokens=True)
+    return f"图像细节描述：{caption}"
 
-with gr.Blocks(title="图像文本描述工具") as demo:
-    gr.Markdown("# 图像文本描述工具")
+# 构建Gradio界面
+with gr.Blocks(title="图像细节描述工具") as demo:
+    gr.Markdown("# 图像细节描述工具（DAM-3B优化版）")
     image_input = gr.Image(type="pil", label="上传图片")
-    text_output = gr.Textbox(label="生成描述")
-    gr.Button("生成").click(fn=generate_caption, inputs=image_input, outputs=text_output)
+    text_output = gr.Textbox(label="生成细节描述", lines=5)
+    gr.Button("生成详细描述").click(fn=generate_detailed_caption, inputs=image_input, outputs=text_output)
 
 if __name__ == "__main__":
     demo.launch()