Update app.py

app.py

@@ -1,42 +1,47 @@
-# from transformers import AutoModel, AutoTokenizer
-
-# tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
-# model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
-# model = model.eval().cuda()
-
-
-# input your test image
-# image_file = 'car.jpg'
-
-# plain texts OCR
-# res = model.chat(tokenizer, image_file, ocr_type='ocr')
-
-# format texts OCR:
-# res = model.chat(tokenizer, image_file, ocr_type='format')
-
-# fine-grained OCR:
-# res = model.chat(tokenizer, image_file, ocr_type='ocr', ocr_box='')
-# res = model.chat(tokenizer, image_file, ocr_type='format', ocr_box='')
-# res = model.chat(tokenizer, image_file, ocr_type='ocr', ocr_color='')
-# res = model.chat(tokenizer, image_file, ocr_type='format', ocr_color='')
-
-# multi-crop OCR:
-# res = model.chat_crop(tokenizer, image_file, ocr_type='ocr')
-# res = model.chat_crop(tokenizer, image_file, ocr_type='format')
+import torch
+from PIL import Image
+from transformers import AutoModel, AutoTokenizer
+
+model = AutoModel.from_pretrained('openbmb/MiniCPM-V-2_6', trust_remote_code=True,
+    attn_implementation='sdpa', torch_dtype=torch.bfloat16) # sdpa or flash_attention_2, no eager
+model = model.eval().cuda()
+tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-V-2_6', trust_remote_code=True)
+
+image = Image.open('car.jpg').convert('RGB')
+question = 'What is in the image?'
+msgs = [{'role': 'user', 'content': [image, question]}]
+
+res = model.chat(
+    image=None,
+    msgs=msgs,
+    tokenizer=tokenizer
+)
+print(res)
+
+## if you want to use streaming, please make sure sampling=True and stream=True
+## the model.chat will return a generator
+res = model.chat(
+    image=None,
+    msgs=msgs,
+    tokenizer=tokenizer,
+    sampling=True,
+    stream=True
+)
 
-# render the formatted OCR results:
-# res = model.chat(tokenizer, image_file, ocr_type='format', render=True, save_render_file = './demo.html')
+generated_text = ""
+for new_text in res:
+    generated_text += new_text
+    print(new_text, flush=True, end='')
 
-# print(res)
 
-import google.generativeai as genai
-import os
+#import google.generativeai as genai
+#import os
 
-genai.configure(api_key=os.environ["AIzaSyB5WiEJf_yLMD1dMQf305EAbaPTzF_QD-I"])
+#genai.configure(api_key=os.environ["AIzaSyB5WiEJf_yLMD1dMQf305EAbaPTzF_QD-I"])
 
-model = genai.GenerativeModel('gemini-1.5-flash')
-response = model.generate_content(
-    text_input="the color of the car is ?", 
-    image_input="car.jpg"
-)
-print(response)
+#model = genai.GenerativeModel('gemini-1.5-flash')
+#response = model.generate_content(
+#    text_input="the color of the car is ?", 
+#    image_input="car.jpg"
+#)
+#print(response)