Upload app.py

app.py

@@ -14,6 +14,7 @@
 # # 加载模型和处理器
 # # 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):
@@ -50,46 +51,53 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 _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)
+
+processor = BlipProcessor.from_pretrained("IDEA-CCNL/Taiyi-BLIP-750M-Chinese", use_auth_token=HF_TOKEN)
+model = BlipForConditionalGeneration.from_pretrained("IDEA-CCNL/Taiyi-BLIP-750M-Chinese", use_auth_token=HF_TOKEN).eval().to(device)
+
+# 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)
+        inputs = processor(raw_image ,return_tensors= "pt").to(device)
         with torch.no_grad():
             if strategy == "Beam search":
+                # Beam search，即集束搜索，每次生成多个词，然后选择概率最大的前 k 个词，然后继续生成，直到生成结束
                 config = GenerationConfig(
                     do_sample=False,
                     num_beams=3,
                     max_length=50,
                     min_length=5,
                 )
-                captions = model.generate(**input, generation_config=config)
+                captions = model.generate(**inputs ,generation_config=config)
             else:
+                # Nucleus sampling，即 top-p sampling，只保留累积概率大于 p 的词，然后重新归一化，得到一个新的概率分布，再从中采样，这样可以保证采样的结果更多样
                 config = GenerationConfig(
                     do_sample=True,
                     top_p=0.9,
                     max_length=50,
                     min_length=5,
                 )
-                captions = model.generate(**input, generation_config=config)
+                captions = model.generate(**inputs ,generation_config=config)
             caption = processor.decode(captions[0], skip_special_tokens=True)
             caption = caption.replace(' ', '')
             print(caption)
-            return 'caption: ' + caption
+            return caption
 
 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")
+    gr.Radio(choices=['Image Captioning'], value="Image Captioning", label="Task"),# 任务选择，目前只有图片描述生成
+    gr.Radio(choices=['Beam search', 'Nucleus sampling'], value="Nucleus sampling", label="Caption Decoding Strategy")# 两种生成策略，Beam search 和 Nucleus sampling，前者生成的结果更准确，后者更多样
 ]
 outputs = gr.Textbox(label="Output")
 
-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>"
+title = "图片描述生成器"
 
-gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, examples=[
+gradio_app=gr.Interface(inference, inputs, outputs, title=title, examples=[
     ['demo.jpg', "Image Captioning", "Nucleus sampling"]
-]).launch()
+])
+
+if __name__ == "__main__":
+    gradio_app.launch()