Spaces:
Running
Running
| import os | |
| os.system("cd open_flamingo && pip install .") | |
| os.system("cd transformers && pip install .") | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from open_flamingo.train.distributed import init_distributed_device, world_info_from_env | |
| import string | |
| import cv2 | |
| import gradio as gr | |
| import torch | |
| from PIL import Image | |
| from huggingface_hub import hf_hub_download, login | |
| from open_flamingo.src.factory import create_model_and_transforms | |
| flamingo, image_processor, tokenizer, vis_embed_size = create_model_and_transforms( | |
| "ViT-L-14", | |
| "datacomp_xl_s13b_b90k", | |
| "facebook/opt-350m", | |
| "facebook/opt-350m", | |
| add_visual_grounding=True, | |
| location_token_num=1000, | |
| add_visual_token = True, | |
| use_format_v2 = True, | |
| ) | |
| checkpoint_path = hf_hub_download("chendl/mm", "checkpoint_opt350m_v2.pt") | |
| checkpoint = torch.load(checkpoint_path, map_location="cpu") | |
| model_state_dict = {} | |
| for key in checkpoint.keys(): | |
| model_state_dict[key.replace("module.", "")] = checkpoint[key] | |
| if "vision_encoder.logit_scale"in model_state_dict: | |
| # previous checkpoint has some unnecessary weights | |
| del model_state_dict["vision_encoder.logit_scale"] | |
| del model_state_dict["vision_encoder.visual.proj"] | |
| del model_state_dict["vision_encoder.visual.ln_post.weight"] | |
| del model_state_dict["vision_encoder.visual.ln_post.bias"] | |
| flamingo.load_state_dict(model_state_dict, strict=True) | |
| def get_outputs( | |
| model, | |
| batch_images, | |
| attention_mask, | |
| max_generation_length, | |
| min_generation_length, | |
| num_beams, | |
| length_penalty, | |
| input_ids, | |
| image_start_index_list=None, | |
| image_nums=None, | |
| bad_words_ids=None, | |
| ): | |
| # and torch.cuda.amp.autocast(dtype=torch.float16) | |
| with torch.inference_mode(): | |
| outputs = model.generate( | |
| batch_images, | |
| input_ids, | |
| attention_mask=attention_mask, | |
| max_new_tokens=max_generation_length, | |
| min_length=min_generation_length, | |
| num_beams=num_beams, | |
| length_penalty=length_penalty, | |
| image_start_index_list=image_start_index_list, | |
| image_nums=image_nums, | |
| bad_words_ids=bad_words_ids, | |
| ) | |
| outputs = outputs[:, len(input_ids[0]) :] | |
| return outputs | |
| def generate( | |
| idx, | |
| image, | |
| text, | |
| vis_embed_size=256, | |
| rank=0, | |
| world_size=1, | |
| ): | |
| if image is None: | |
| raise gr.Error("Please upload an image.") | |
| flamingo.eval() | |
| loc_token_ids = [] | |
| for i in range(1000): | |
| loc_token_ids.append(int(tokenizer(f"<loc_{i}>", add_special_tokens=False)["input_ids"][-1])) | |
| media_token_id = tokenizer("<|#image#|>", add_special_tokens=False)["input_ids"][-1] | |
| endofchunk_token_id = tokenizer("<|endofchunk|>", add_special_tokens=False)["input_ids"][-1] | |
| endofmedia_token_id = tokenizer("<|#endofimage#|>", add_special_tokens=False)["input_ids"][-1] | |
| pad_token_id = tokenizer(tokenizer.pad_token, add_special_tokens=False)["input_ids"][-1] | |
| bos_token_id = tokenizer(tokenizer.bos_token, add_special_tokens=False)["input_ids"][-1] | |
| all_ids = set(range(flamingo.lang_encoder.lm_head.out_features)) | |
| bad_words_ids = list(all_ids - set(loc_token_ids)) | |
| bad_words_ids = [[b] for b in bad_words_ids] | |
| loc_word_ids = list(set(loc_token_ids)) | |
| loc_word_ids = [[b] for b in loc_word_ids] | |
| min_loc_token_id = min(loc_token_ids) | |
| max_loc_token_id = max(loc_token_ids) | |
| image_ori = image | |
| image = image.convert("RGB") | |
| width = image.width | |
| height = image.height | |
| image = image.resize((224, 224)) | |
| batch_images = image_processor(image).unsqueeze(0).unsqueeze(1).unsqueeze(0) | |
| if idx == 1: | |
| prompt = [f"<|#image#|>{tokenizer.pad_token*vis_embed_size}<|#endofimage#|><|#obj#|>{text.rstrip('.')}<|#loc#|>"] | |
| bad_words_ids = None | |
| max_generation_length = 5 | |
| else: | |
| prompt = [f"<|#image#|>{tokenizer.pad_token * vis_embed_size}<|#endofimage#|>{text.rstrip('.')}"] | |
| bad_words_ids = loc_word_ids | |
| max_generation_length = 100 | |
| encodings = tokenizer( | |
| prompt, | |
| padding="longest", | |
| truncation=True, | |
| return_tensors="pt", | |
| max_length=2000, | |
| ) | |
| input_ids = encodings["input_ids"] | |
| attention_mask = encodings["attention_mask"] | |
| image_start_index_list = ((input_ids == media_token_id).nonzero(as_tuple=True)[-1] + 1).tolist() | |
| image_start_index_list = [[x] for x in image_start_index_list] | |
| image_nums = [1] * len(input_ids) | |
| outputs = get_outputs( | |
| model=flamingo, | |
| batch_images=batch_images, | |
| attention_mask=attention_mask, | |
| max_generation_length=max_generation_length, | |
| min_generation_length=4, | |
| num_beams=1, | |
| length_penalty=1.0, | |
| input_ids=input_ids, | |
| bad_words_ids=bad_words_ids, | |
| image_start_index_list=image_start_index_list, | |
| image_nums=image_nums, | |
| ) | |
| box = [] | |
| out_image = image_ori | |
| for o in outputs[0]: | |
| if o >= min_loc_token_id and o <= max_loc_token_id: | |
| box.append(o.item() - min_loc_token_id) | |
| if len(box) == 4: | |
| break | |
| # else: | |
| # tqdm.write(f"output: {tokenizer.batch_decode(outputs)}") | |
| # tqdm.write(f"prompt: {prompt}") | |
| if len(box) == 4: | |
| img = cv2.cvtColor(np.array(image_ori), cv2.COLOR_RGB2BGR) | |
| out = cv2.rectangle(img, (int(box[0] * width / 1000), int(box[1] * height / 1000)), | |
| (int(box[2] * width / 1000), int(box[3] * height / 1000)), color=(255, 0, 255), thickness=2) | |
| out = cv2.cvtColor(out, cv2.COLOR_BGR2RGB) | |
| out_image = Image.fromarray(out) | |
| # else: | |
| # tqdm.write(f"output: {tokenizer.batch_decode(outputs)}") | |
| # tqdm.write(f"prompt: {prompt}") | |
| gen_text = tokenizer.batch_decode(outputs) | |
| if idx == 1: | |
| return f"Output:{gen_text}", out_image | |
| elif idx == 2: | |
| return (f"Question: {text.strip()} Answer: {gen_text}") | |
| else: | |
| return (f"Output:{gen_text}") | |
| with gr.Blocks() as demo: | |
| gr.Markdown( | |
| """ | |
| 🍜 Object Centric Pretraining Demo | |
| 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. | |
| 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. | |
| """ | |
| ) | |
| with gr.Accordion("See terms and conditions"): | |
| gr.Markdown( | |
| """**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.""") | |
| with gr.Tab("📷 Image Captioning"): | |
| with gr.Row(): | |
| query_image = gr.Image(type="pil") | |
| with gr.Row(): | |
| chat_input = gr.Textbox(lines=1, label="Chat Input") | |
| text_output = gr.Textbox(value="Output:", label="Model output") | |
| run_btn = gr.Button("Run model") | |
| def on_click_fn(img,text): return generate(0, img, text) | |
| run_btn.click(on_click_fn, inputs=[query_image,chat_input], outputs=[text_output]) | |
| with gr.Tab("🦓 Grounding"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| query_image = gr.Image(type="pil") | |
| with gr.Column(scale=1): | |
| out_image = gr.Image(type="pil") | |
| with gr.Row(): | |
| chat_input = gr.Textbox(lines=1, label="Chat Input") | |
| text_output = gr.Textbox(value="Output:", label="Model output") | |
| run_btn = gr.Button("Run model") | |
| def on_click_fn(img, text): return generate(1, img, text) | |
| run_btn.click(on_click_fn, inputs=[query_image, chat_input], outputs=[text_output, out_image]) | |
| with gr.Tab("🔢 Counting objects"): | |
| with gr.Row(): | |
| query_image = gr.Image(type="pil") | |
| with gr.Row(): | |
| chat_input = gr.Textbox(lines=1, label="Chat Input") | |
| text_output = gr.Textbox(value="Output:", label="Model output") | |
| run_btn = gr.Button("Run model") | |
| def on_click_fn(img,text): return generate(0, img, text) | |
| run_btn.click(on_click_fn, inputs=[query_image, chat_input], outputs=[text_output]) | |
| with gr.Tab("🕵️ Visual Question Answering"): | |
| with gr.Row(): | |
| query_image = gr.Image(type="pil") | |
| with gr.Row(): | |
| question = gr.Textbox(lines=1, label="Question") | |
| text_output = gr.Textbox(value="Output:", label="Model output") | |
| run_btn = gr.Button("Run model") | |
| def on_click_fn(img, txt): return generate(2, img, txt) | |
| run_btn.click( | |
| on_click_fn, inputs=[query_image, question], outputs=[text_output] | |
| ) | |
| with gr.Tab("🌎 Custom"): | |
| gr.Markdown( | |
| """### Customize the demonstration by uploading your own images and text samples. | |
| ### **Note: Any text prompt you use will be prepended with an 'Output:', so you don't need to include it in your prompt.**""" | |
| ) | |
| with gr.Row(): | |
| query_image = gr.Image(type="pil") | |
| with gr.Row(): | |
| question = gr.Textbox(lines=1, label="Question") | |
| text_output = gr.Textbox(value="Output:", label="Model output") | |
| run_btn = gr.Button("Run model") | |
| def on_click_fn(img, txt): return generate(2, img, txt) | |
| run_btn.click( | |
| on_click_fn, inputs=[query_image, question], outputs=[text_output] | |
| ) | |
| demo.queue(concurrency_count=1) | |
| demo.launch() |