Upload app.py

app.py

@@ -1,35 +1,95 @@
-import os
+# import os
+# import gradio as gr
+# from transformers import BlipProcessor ,BlipForConditionalGeneration
+# from PIL import Image
+# from transformers import CLIPProcessor, ChineseCLIPVisionModel ,AutoProcessor
+#
+# # 设置环境变量 HF_HOME 和 HF_ENDPOINT
+# # os.environ['HF_HOME'] = 'D:/AI/OCR/img2text/models'
+# # os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
+#
+#
+# # model = ChineseCLIPVisionModel.from_pretrained("OFA-Sys/chinese-clip-vit-base-patch16")
+# # processor = AutoProcessor.from_pretrained("OFA-Sys/chinese-clip-vit-base-patch16")
+# # 加载模型和处理器
+# # processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+# # model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+# processor = BlipProcessor.from_pretrained("IDEA-CCNL/Taiyi-BLIP-750M-Chinese")
+# model = BlipForConditionalGeneration.from_pretrained("IDEA-CCNL/Taiyi-BLIP-750M-Chinese")
+# def generate_caption(image):
+#     # 确保 image 是 PIL.Image 类型
+#     if not isinstance(image, Image.Image):
+#         raise ValueError("Input must be a PIL.Image")
+#
+#     inputs = processor(image, return_tensors="pt")
+#     input_ids = inputs.get("input_ids")
+#     if input_ids is None:
+#         raise ValueError("Processor did not return input_ids")
+#
+#     outputs = model.generate(input_ids=input_ids, max_length=50)
+#     description = processor.decode(outputs[0], skip_special_tokens=True)
+#     return description
+#
+# # 创建Gradio接口
+# gradio_app = gr.Interface(
+#     fn=generate_caption,
+#     inputs=gr.Image(type="pil"),
+#     outputs="text",
+#     title="图片描述生成器",
+#     description="上传一张图片，生成相应的描述。"
+# )
+#
+# if __name__ == "__main__":
+#     gradio_app.launch()
 import gradio as gr
-from transformers import BlipProcessor ,BlipForConditionalGeneration
-from PIL import Image
-from transformers import CLIPProcessor, ChineseCLIPVisionModel ,AutoProcessor
+import torch
+import os
+from transformers import BlipForConditionalGeneration, BlipProcessor, GenerationConfig
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# 设置环境变量 HF_HOME 和 HF_ENDPOINT
-# os.environ['HF_HOME'] = 'D:/AI/OCR/img2text/models'
-# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
+_MODEL_PATH = 'IDEA-CCNL/Taiyi-BLIP-750M-Chinese'
+HF_TOKEN = os.getenv('HF_TOKEN')
+processor = BlipProcessor.from_pretrained(_MODEL_PATH, use_auth_token=HF_TOKEN)
+model = BlipForConditionalGeneration.from_pretrained(
+    _MODEL_PATH, use_auth_token=HF_TOKEN).eval().to(device)
 
+def inference(raw_image, model_n, strategy):
+    if model_n == 'Image Captioning':
+        input = processor(raw_image, return_tensors="pt").to(device)
+        with torch.no_grad():
+            if strategy == "Beam search":
+                config = GenerationConfig(
+                    do_sample=False,
+                    num_beams=3,
+                    max_length=50,
+                    min_length=5,
+                )
+                captions = model.generate(**input, generation_config=config)
+            else:
+                config = GenerationConfig(
+                    do_sample=True,
+                    top_p=0.9,
+                    max_length=50,
+                    min_length=5,
+                )
+                captions = model.generate(**input, generation_config=config)
+            caption = processor.decode(captions[0], skip_special_tokens=True)
+            caption = caption.replace(' ', '')
+            print(caption)
+            return 'caption: ' + caption
 
-# model = ChineseCLIPVisionModel.from_pretrained("OFA-Sys/chinese-clip-vit-base-patch16")
-# processor = AutoProcessor.from_pretrained("OFA-Sys/chinese-clip-vit-base-patch16")
-# 加载模型和处理器
-# processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-# model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
-processor = BlipProcessor.from_pretrained("IDEA-CCNL/Taiyi-BLIP-750M-Chinese")
-model = BlipForConditionalGeneration.from_pretrained("IDEA-CCNL/Taiyi-BLIP-750M-Chinese")
-def generate_caption(image):
-    inputs = processor(image, return_tensors="pt")
-    outputs = model(**inputs)
-    description = processor.decode(outputs[0], skip_special_tokens=True)
-    return description
+inputs = [
+    gr.Image(type='pil', label="Upload Image"),
+    gr.Radio(choices=['Image Captioning'], value="Image Captioning", label="Task"),
+    gr.Radio(choices=['Beam search', 'Nucleus sampling'], value="Nucleus sampling", label="Caption Decoding Strategy")
+]
+outputs = gr.Textbox(label="Output")
 
-# 创建Gradio接口
-gradio_app = gr.Interface(
-    fn=generate_caption,
-    inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="图片描述生成器",
-    description="上传一张图片，生成相应的描述。"
-)
+title = "BLIP"
+description = "Gradio demo for BLIP: Bootstrapping Language-Image Pre-training for Unified Vision-Language Understanding and Generation (Salesforce Research). To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
+article = "<p style='text-align: center'><a href='https://github.com/IDEA-CCNL/Fengshenbang-LM' target='_blank'>Github Repo</a></p>"
 
-if __name__ == "__main__":
-    gradio_app.launch()
+gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, examples=[
+    ['demo.jpg', "Image Captioning", "Nucleus sampling"]
+]).launch()