chendl/compositional_test
1
1import os
2import sys
3from pathlib import Path
4os.system("python -m pip install --upgrade pip")
5os.system("cd multimodal && pip install .")
6os.system("cd multimodal/YOLOX && pip install .")
7import numpy as np
8import torch
9from PIL import Image
10import tempfile
11
12import string
13import cv2
14
15import gradio as gr
16import torch
17from PIL import Image
18from huggingface_hub import hf_hub_download, login
19
20from open_flamingo.src.factory import create_model_and_transforms
21from open_flamingo.chat.conversation import ChatBOT, CONV_VISION
22
23sys.path.append(str(Path(__file__).parent.parent.parent))
24TEMP_FILE_DIR = Path(__file__).parent / 'temp'
25TEMP_FILE_DIR.mkdir(parents=True, exist_ok=True)
26
27SHARED_UI_WARNING = f'''### [NOTE] It is possible that you are waiting in a lengthy queue.
28
29You can duplicate and use it with a paid private GPU.
30
31<a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/Vision-CAIR/minigpt4?duplicate=true"><img style="margin-top:0;margin-bottom:0" src="https://huggingface.co/datasets/huggingface/badges/raw/main/duplicate-this-space-xl-dark.svg" alt="Duplicate Space"></a>
32
33Alternatively, you can also use the demo on our [project page](https://compositionalvlm.github.io/).
34'''
35
36flamingo, image_processor, tokenizer, vis_embed_size = create_model_and_transforms(
37 "ViT-L-14",
38 "datacomp_xl_s13b_b90k",
39 "EleutherAI/pythia-1.4b",
40 "EleutherAI/pythia-1.4b",
41 location_token_num=1000,
42 lora=False,
43 lora_r=16,
44 use_sam=None,
45 add_visual_token=True,
46 use_format_v2=True,
47 add_box=True,
48 add_pe=False,
49 add_relation=False,
50 enhance_data=False,
51)
52
53model_name = "pythiaS"
54checkpoint_path = hf_hub_download("chendl/compositional_test", "pythiaS.pt")
55checkpoint = torch.load(checkpoint_path, map_location="cpu")["model_state_dict"]
56model_state_dict = {}
57for key in checkpoint.keys():
58 model_state_dict[key.replace("module.", "")] = checkpoint[key]
59if "vision_encoder.logit_scale" in model_state_dict:
60 # previous checkpoint has some unnecessary weights
61 del model_state_dict["vision_encoder.logit_scale"]
62 del model_state_dict["vision_encoder.visual.proj"]
63 del model_state_dict["vision_encoder.visual.ln_post.weight"]
64 del model_state_dict["vision_encoder.visual.ln_post.bias"]
65flamingo.load_state_dict(model_state_dict, strict=True)
66chat = ChatBOT(flamingo, image_processor, tokenizer, vis_embed_size,model_name)
67
68
69def get_outputs(
70 model,
71 batch_images,
72 attention_mask,
73 max_generation_length,
74 min_generation_length,
75 num_beams,
76 length_penalty,
77 input_ids,
78 image_start_index_list=None,
79 image_nums=None,
80 bad_words_ids=None,
81):
82 # and torch.cuda.amp.autocast(dtype=torch.float16)
83 with torch.inference_mode():
84 outputs = model(
85 vision_x=batch_images,
86 lang_x=input_ids,
87 attention_mask=attention_mask,
88 labels=None,
89 image_nums=image_nums,
90 image_start_index_list=image_start_index_list,
91 added_bbox_list=None,
92 add_box=False,
93 )
94 # outputs = model.generate(
95 # batch_images,
96 # input_ids,
97 # attention_mask=attention_mask,
98 # max_new_tokens=max_generation_length,
99 # min_length=min_generation_length,
100 # num_beams=num_beams,
101 # length_penalty=length_penalty,
102 # image_start_index_list=image_start_index_list,
103 # image_nums=image_nums,
104 # bad_words_ids=bad_words_ids,
105 # )
106
107 return outputs
108
109
110def generate(
111 idx,
112 image,
113 text,
114 vis_embed_size=256,
115 rank=0,
116 world_size=1,
117):
118 if image is None:
119 raise gr.Error("Please upload an image.")
120 flamingo.eval()
121 loc_token_ids = []
122 for i in range(1000):
123 loc_token_ids.append(int(tokenizer(f"<loc_{i}>", add_special_tokens=False)["input_ids"][-1]))
124 media_token_id = tokenizer("<|#image#|>", add_special_tokens=False)["input_ids"][-1]
125 endofmedia_token_id = tokenizer("<|#endofimage#|>", add_special_tokens=False)["input_ids"][-1]
126 pad_token_id = tokenizer(tokenizer.pad_token, add_special_tokens=False)["input_ids"][-1]
127 bos_token_id = tokenizer(tokenizer.bos_token, add_special_tokens=False)["input_ids"][-1]
128 prebox_token_id = tokenizer("<|#prebox#|>", add_special_tokens=False)["input_ids"][-1]
129
130 image_ori = image
131 image = image.convert("RGB")
132 width = image.width
133 height = image.height
134 image = image.resize((224, 224))
135 batch_images = image_processor(image).unsqueeze(0).unsqueeze(1).unsqueeze(0)
136 if idx == 1:
137 prompt = [
138 f"{tokenizer.bos_token}<|#image#|>{tokenizer.pad_token * vis_embed_size}<|#endofimage#|><|#object#|> {text.rstrip('.').strip()}<|#endofobject#|><|#visual#|>"]
139 bad_words_ids = None
140 max_generation_length = 5
141 else:
142 prompt = [f"<|#image#|>{tokenizer.pad_token * vis_embed_size}<|#endofimage#|>{text.rstrip('.')}"]
143 bad_words_ids = loc_word_ids
144 max_generation_length = 30
145 encodings = tokenizer(
146 prompt,
147 padding="longest",
148 truncation=True,
149 return_tensors="pt",
150 max_length=2000,
151 )
152 input_ids = encodings["input_ids"]
153 attention_mask = encodings["attention_mask"]
154 image_start_index_list = ((input_ids == media_token_id).nonzero(as_tuple=True)[-1] + 1).tolist()
155 image_start_index_list = [[x] for x in image_start_index_list]
156 image_nums = [1] * len(input_ids)
157 outputs = get_outputs(
158 model=flamingo,
159 batch_images=batch_images,
160 attention_mask=attention_mask,
161 max_generation_length=max_generation_length,
162 min_generation_length=4,
163 num_beams=1,
164 length_penalty=1.0,
165 input_ids=input_ids,
166 bad_words_ids=bad_words_ids,
167 image_start_index_list=image_start_index_list,
168 image_nums=image_nums,
169 )
170
171 boxes = outputs["boxes"]
172 scores = outputs["scores"]
173 if len(scores) > 0:
174 box = boxes[scores.argmax()] / 224
175 print(f"{box}")
176
177 if idx == 1:
178 open_cv_image = np.array(image_ori)
179 # Convert RGB to BGR
180 open_cv_image = open_cv_image[:, :, ::-1].copy()
181 box = box * [width, height, width, height]
182 # for box in boxes:
183 open_cv_image = cv2.rectangle(open_cv_image, box[:2].astype(int), box[2:].astype(int), (255, 0, 0), 2)
184 out_image = Image.fromarray(cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2RGB))
185 return f"Output:{box}", out_image
186 elif idx == 2:
187 gen_text = tokenizer.batch_decode(outputs)
188 return (f"Question: {text.strip()} Answer: {gen_text}")
189 else:
190 gen_text = tokenizer.batch_decode(outputs)
191 return (f"Output:{gen_text}")
192
193
194title = """<h1 align="center">Demo of Compositional-VLM</h1>"""
195description = """<h3>This is the demo of Compositional-VLM. Upload your images and start chatting!</h3>"""
196article = """<div style='display:flex; gap: 0.25rem; '><a href='https://vis-www.cs.umass.edu/CoVLM/'><img src='https://img.shields.io/badge/Project-Page-Green'></a><a href='https://github.com/UMass-Foundation-Model/CoVLM'><img src='https://img.shields.io/badge/Github-Code-blue'></a><a href='https://arxiv.org/abs/2311.03354'><img src='https://img.shields.io/badge/Paper-PDF-red'></a></div>
197"""
198
199
200# TODO show examples below
201
202# ========================================
203# Gradio Setting
204# ========================================
205
206def gradio_reset(chat_state, img_list):
207 if chat_state is not None:
208 chat_state = []
209 if img_list is not None:
210 img_list = []
211 return None, gr.update(value=None, interactive=True), gr.update(placeholder='Please upload your image first',
212 interactive=False), gr.update(
213 value="Upload & Start Chat", interactive=True), chat_state, img_list
214
215
216def build_image(image):
217 if image is None:
218 return None
219 # res = draw_bounding_boxes(image=image, boxes=boxes_to_draw, colors=color_to_draw, width=8)
220 from torchvision.transforms import ToPILImage
221 # res = ToPILImage()(res)
222 _, path = tempfile.mkstemp(suffix='.jpg', dir=TEMP_FILE_DIR)
223 image.save(path)
224
225 return path
226
227
228def upload_img(gr_img, text_input, chat_state, chatbot):
229 if gr_img is None:
230 return None, None, gr.update(interactive=True), chat_state, None
231 chat_state = []
232 img_list = []
233 path = build_image(gr_img)
234 chatbot = chatbot + [[(path,), None]]
235 llm_message = chat.upload_img(gr_img, chat_state, img_list)
236 return gr.update(interactive=False), gr.Textbox(placeholder='Type and press Enter', interactive=True), gr.update(
237 value="Start Chatting", interactive=False), chat_state, img_list, chatbot
238
239
240def gradio_ask(user_message, chatbot, chat_state, radio):
241 # if len(user_message) == 0:
242 # return gr.update(interactive=True, placeholder='Input should not be empty!'), chatbot, chat_state
243
244 chat.ask(user_message, chat_state, radio)
245 chatbot = chatbot + [[user_message, None]]
246 return chatbot, chat_state
247
248
249def generate_ans(user_message, chatbot, chat_state, img_list, radio, text, num_beams, temperature):
250 # if len(user_message) == 0:
251 # return gr.update(interactive=True, placeholder='Input should not be empty!'), chatbot, chat_state
252
253 chat.ask(user_message, chat_state, radio)
254 chatbot = chatbot + [[user_message, None]]
255 # return chatbot, chat_state
256 image = None
257 llm_message, image = \
258 chat.answer(conv=chat_state, img_list=img_list, max_new_tokens=300, num_beams=1, temperature=temperature,
259 max_length=2000, radio=radio, text_input=text)
260
261 chatbot[-1][1] = llm_message
262 if chat_state[-1]["from"] == "gpt":
263 chat_state[-1]["value"] = llm_message
264 if image == None:
265 return "", chatbot, chat_state, img_list
266 else:
267 path = build_image(image)
268 chatbot = chatbot + [[None, (path,)]]
269 return "", chatbot, chat_state, img_list
270
271
272def gradio_answer(chatbot, chat_state, img_list, radio, text, num_beams, temperature):
273 image = None
274 llm_message, image = \
275 chat.answer(conv=chat_state, img_list=img_list, max_new_tokens=300, num_beams=1, temperature=temperature,
276 max_length=2000, radio=radio, text_input=text)
277
278 chatbot[-1][1] = llm_message
279 if chat_state[-1]["from"] == "gpt":
280 chat_state[-1]["value"] = llm_message
281 if image == None:
282 return "", chatbot, chat_state, img_list
283 else:
284 path = build_image(image)
285 chatbot = chatbot + [[None, (path,)]]
286 return "", chatbot, chat_state, img_list
287
288
289task_template = {
290 "Cap": "Summarize the content of the photo <image>.",
291 "VQA": "For this image <image>, I want a simple and direct answer to my question: <question>",
292 "REC": "Can you point out <expr> in the image <image> and provide the coordinates of its location?",
293 "GC": "Can you give me a description of the region <boxes> in image <image>?",
294 "Advanced": "<question>",
295}
296
297with gr.Blocks() as demo:
298 gr.Markdown(title)
299 gr.Markdown(SHARED_UI_WARNING)
300 gr.Markdown(description)
301 gr.Markdown(article)
302
303 with gr.Row():
304 with gr.Column(scale=0.5):
305 image = gr.Image(type="pil")
306 upload_button = gr.Button(value="Upload & Start Chat", interactive=True, variant="primary")
307 clear = gr.Button("Restart")
308 radio = gr.Radio(
309 ["Cap", "VQA", "REC", "Advanced"], label="Task Template", value='Cap',
310 )
311
312 num_beams = gr.Slider(
313 minimum=1,
314 maximum=5,
315 value=1,
316 step=1,
317 interactive=True,
318 label="beam search numbers)",
319 )
320
321 temperature = gr.Slider(
322 minimum=0.1,
323 maximum=2.0,
324 value=1.0,
325 step=0.1,
326 interactive=True,
327 label="Temperature",
328 )
329
330 with gr.Column():
331 chat_state = gr.State()
332 img_list = gr.State()
333 chatbot = gr.Chatbot(label='Compositional-VLM')
334
335 # template = gr.Textbox(label='Template', show_label=True, lines=1, interactive=False,
336 # value='Provide a comprehensive description of the image <image> and specify the positions of any mentioned objects in square brackets.')
337 # text_input = gr.Textbox(label='<question>', show_label=True, placeholder="Please upload your image firstοΌ then input...", lines=3,
338 # value=None, visible=False, interactive=False)
339 # with gr.Row():
340 text_input = gr.Textbox(label='User', placeholder='Please upload your image first, then input...',
341 interactive=False)
342 # submit_button = gr.Button(value="Submit", interactive=True, variant="primary")
343
344 upload_button.click(upload_img, [image, text_input, chat_state, chatbot],
345 [image, text_input, upload_button, chat_state, img_list, chatbot])
346 # submit_button.click(gradio_ask, [text_input, chatbot, chat_state,radio], [chatbot, chat_state]).then(
347 # gradio_answer, [chatbot, chat_state, img_list, radio, text_input,num_beams, temperature], [text_input,chatbot, chat_state, img_list]
348 # )
349
350 text_input.submit(generate_ans,
351 [text_input, chatbot, chat_state, img_list, radio, text_input, num_beams, temperature],
352 [text_input, chatbot, chat_state, img_list])
353
354 # text_input.submit(gradio_ask, [text_input, chatbot, chat_state, radio], [chatbot, chat_state]).then(
355 # gradio_answer, [chatbot, chat_state, img_list, radio, text_input, num_beams, temperature],
356 # [text_input, chatbot, chat_state, img_list]
357 # )
358 clear.click(gradio_reset, [chat_state, img_list], [chatbot, image, text_input, upload_button, chat_state, img_list],
359 queue=False)
360
361demo.launch(share=True)
362#
363# with gr.Blocks() as demo:
364# gr.Markdown(
365# """
366# π Object Centric Pretraining Demo
367# In this demo we showcase the in-context learning and grounding capabilities of the Object-Centric Pretrained model, a large multimodal model. Note that we add two additional demonstrations to the ones presented to improve the demo experience.
368# The model is trained on an interleaved mixture of text, images and bounding box and is able to generate text conditioned on sequences of images/text.
369# """
370# )
371#
372# with gr.Accordion("See terms and conditions"):
373# gr.Markdown(
374# """**Please read the following information carefully before proceeding.**This demo does NOT store any personal information on its users, and it does NOT store user queries.""")
375#
376# with gr.Tab("π· Image Captioning"):
377# with gr.Row():
378#
379#
380# query_image = gr.Image(type="pil")
381# with gr.Row():
382# chat_input = gr.Textbox(lines=1, label="Chat Input")
383# text_output = gr.Textbox(value="Output:", label="Model output")
384#
385# run_btn = gr.Button("Run model")
386#
387#
388#
389# def on_click_fn(img,text): return generate(0, img, text)
390#
391# run_btn.click(on_click_fn, inputs=[query_image,chat_input], outputs=[text_output])
392#
393# with gr.Tab("π¦ Grounding"):
394# with gr.Row():
395# with gr.Column(scale=1):
396# query_image = gr.Image(type="pil")
397# with gr.Column(scale=1):
398# out_image = gr.Image(type="pil")
399# with gr.Row():
400# chat_input = gr.Textbox(lines=1, label="Chat Input")
401# text_output = gr.Textbox(value="Output:", label="Model output")
402#
403# run_btn = gr.Button("Run model")
404#
405#
406# def on_click_fn(img, text): return generate(1, img, text)
407#
408#
409# run_btn.click(on_click_fn, inputs=[query_image, chat_input], outputs=[text_output, out_image])
410#
411# with gr.Tab("π’ Counting objects"):
412# with gr.Row():
413# query_image = gr.Image(type="pil")
414# with gr.Row():
415# chat_input = gr.Textbox(lines=1, label="Chat Input")
416# text_output = gr.Textbox(value="Output:", label="Model output")
417#
418# run_btn = gr.Button("Run model")
419#
420#
421# def on_click_fn(img,text): return generate(0, img, text)
422#
423#
424# run_btn.click(on_click_fn, inputs=[query_image, chat_input], outputs=[text_output])
425#
426# with gr.Tab("π΅οΈ Visual Question Answering"):
427# with gr.Row():
428# query_image = gr.Image(type="pil")
429# with gr.Row():
430# question = gr.Textbox(lines=1, label="Question")
431# text_output = gr.Textbox(value="Output:", label="Model output")
432#
433# run_btn = gr.Button("Run model")
434#
435#
436# def on_click_fn(img, txt): return generate(2, img, txt)
437#
438#
439# run_btn.click(
440# on_click_fn, inputs=[query_image, question], outputs=[text_output]
441# )
442#
443# with gr.Tab("π Custom"):
444# gr.Markdown(
445# """### Customize the demonstration by uploading your own images and text samples.
446# ### **Note: Any text prompt you use will be prepended with an 'Output:', so you don't need to include it in your prompt.**"""
447# )
448# with gr.Row():
449# query_image = gr.Image(type="pil")
450# with gr.Row():
451# question = gr.Textbox(lines=1, label="Question")
452# text_output = gr.Textbox(value="Output:", label="Model output")
453#
454# run_btn = gr.Button("Run model")
455#
456#
457# def on_click_fn(img, txt): return generate(2, img, txt)
458#
459#
460# run_btn.click(
461# on_click_fn, inputs=[query_image, question], outputs=[text_output]
462# )
463#
464# demo.queue(concurrency_count=1)
465# demo.launch()
466