import gradio as gr import re import torch from PIL import Image from transformers import AutoTokenizer, FuyuForCausalLM, FuyuImageProcessor, FuyuProcessor model_id = "adept/fuyu-8b" dtype = torch.bfloat16 device = "cuda" tokenizer = AutoTokenizer.from_pretrained(model_id) model = FuyuForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=dtype) processor = FuyuProcessor(image_processor=FuyuImageProcessor(), tokenizer=tokenizer) CAPTION_PROMPT = "Generate a coco-style caption.\n" DETAILED_CAPTION_PROMPT = "What is happening in this image?" def resize_to_max(image, max_width=1920, max_height=1080): width, height = image.size if width <= max_width and height <= max_height: return image scale = min(max_width/width, max_height/height) width = int(width*scale) height = int(height*scale) return image.resize((width, height), Image.LANCZOS) def pad_to_size(image, canvas_width=1920, canvas_height=1080): width, height = image.size if width >= canvas_width and height >= canvas_height: return image # Paste at (0, 0) canvas = Image.new("RGB", (canvas_width, canvas_height)) canvas.paste(image) return canvas def predict(image, prompt): # image = image.convert('RGB') model_inputs = processor(text=prompt, images=[image]) model_inputs = {k: v.to(dtype=dtype if torch.is_floating_point(v) else v.dtype, device=device) for k,v in model_inputs.items()} generation_output = model.generate(**model_inputs, max_new_tokens=50) prompt_len = model_inputs["input_ids"].shape[-1] return tokenizer.decode(generation_output[0][prompt_len:], skip_special_tokens=True) def caption(image, detailed_captioning): if detailed_captioning: caption_prompt = DETAILED_CAPTION_PROMPT else: caption_prompt = CAPTION_PROMPT return predict(image, caption_prompt).lstrip() def set_example_image(example: list) -> dict: return gr.Image.update(value=example[0]) def scale_factor_to_fit(original_size, target_size=(1920, 1080)): width, height = original_size max_width, max_height = target_size if width <= max_width and height <= max_height: return 1.0 return min(max_width/width, max_height/height) def tokens_to_box(tokens, original_size): bbox_start = tokenizer.convert_tokens_to_ids("<0x00>") bbox_end = tokenizer.convert_tokens_to_ids("<0x01>") try: # Assumes a single box bbox_start_pos = (tokens == bbox_start).nonzero(as_tuple=True)[0].item() bbox_end_pos = (tokens == bbox_end).nonzero(as_tuple=True)[0].item() if bbox_end_pos != bbox_start_pos + 5: return tokens # Retrieve transformed coordinates from tokens coords = tokenizer.convert_ids_to_tokens(tokens[bbox_start_pos+1:bbox_end_pos]) # Scale back to original image size and multiply by 2 scale = scale_factor_to_fit(original_size) top, left, bottom, right = [2 * int(float(c)/scale) for c in coords] # Replace the IDs so they get detokenized right replacement = f" {top}, {left}, {bottom}, {right}" replacement = tokenizer.tokenize(replacement)[1:] replacement = tokenizer.convert_tokens_to_ids(replacement) replacement = torch.tensor(replacement).to(tokens) tokens = torch.cat([tokens[:bbox_start_pos], replacement, tokens[bbox_end_pos+1:]], 0) return tokens except: gr.Error("Can't convert tokens.") return tokens def coords_from_response(response): # y1, x1, y2, x2 pattern = r"(\d+),\s*(\d+),\s*(\d+),\s*(\d+)" match = re.search(pattern, response) if match: # Unpack and change order y1, x1, y2, x2 = [int(coord) for coord in match.groups()] return (x1, y1, x2, y2) else: gr.Error("The string is malformed or does not match the expected pattern.") def localize(image, query): prompt = f"When presented with a box, perform OCR to extract text contained within it. If provided with text, generate the corresponding bounding box.\n{query}" # Downscale and/or pad to 1920x1080 padded = resize_to_max(image) padded = pad_to_size(padded) model_inputs = processor(text=prompt, images=[padded]) model_inputs = {k: v.to(dtype=dtype if torch.is_floating_point(v) else v.dtype, device=device) for k,v in model_inputs.items()} generation_output = model.generate(**model_inputs, max_new_tokens=40) prompt_len = model_inputs["input_ids"].shape[-1] tokens = generation_output[0][prompt_len:] tokens = tokens_to_box(tokens, image.size) decoded = tokenizer.decode(tokens, skip_special_tokens=True) coords = coords_from_response(decoded) return image, [(coords, f"Location of \"{query}\"")] css = """ #mkd { height: 500px; overflow: auto; border: 1px solid #ccc; } """ with gr.Blocks(css=css) as demo: gr.HTML( """

Fuyu Multimodal Demo

Fuyu-8B is a multimodal model that supports a variety of tasks combining text and image prompts.

For example, you can use it for captioning by asking it to describe an image. You can also ask it questions about an image, a task known as Visual Question Answering, or VQA. This demo lets you explore captioning and VQA, with more tasks coming soon :) Learn more about the model in our blog post.

Note: This is a raw model release. We have not added further instruction-tuning, postprocessing or sampling strategies to control for undesirable outputs. The model may hallucinate, and you should expect to have to fine-tune the model for your use-case!

Play with Fuyu-8B in this demo! 💬

""" ) with gr.Tab("Visual Question Answering"): with gr.Row(): with gr.Column(): image_input = gr.Image(label="Upload your Image", type="pil") text_input = gr.Textbox(label="Ask a Question") vqa_output = gr.Textbox(label="Output") vqa_btn = gr.Button("Answer Visual Question") gr.Examples( [["assets/vqa_example_1.png", "How is this made?"], ["assets/vqa_example_2.png", "What is this flower and where is it's origin?"], ["assets/docvqa_example.png", "How many items are sold?"], ["assets/screen2words_ui_example.png", "What is this app about?"]], inputs = [image_input, text_input], outputs = [vqa_output], fn=predict, cache_examples=True, label='Click on any Examples below to get VQA results quickly 👇' ) with gr.Tab("Image Captioning"): with gr.Row(): with gr.Column(): captioning_input = gr.Image(label="Upload your Image", type="pil") detailed_captioning_checkbox = gr.Checkbox(label="Enable detailed captioning") captioning_output = gr.Textbox(label="Output") captioning_btn = gr.Button("Generate Caption") gr.Examples( [["assets/captioning_example_1.png", False], ["assets/captioning_example_2.png", True]], inputs = [captioning_input, detailed_captioning_checkbox], outputs = [captioning_output], fn=caption, cache_examples=True, label='Click on any Examples below to get captioning results quickly 👇' ) captioning_btn.click(fn=caption, inputs=[captioning_input, detailed_captioning_checkbox], outputs=captioning_output) vqa_btn.click(fn=predict, inputs=[image_input, text_input], outputs=vqa_output) with gr.Tab("Find Text in Screenshots"): with gr.Row(): with gr.Column(): localization_input = gr.Image(label="Upload your Image", type="pil") query_input = gr.Textbox(label="Text to find") localization_btn = gr.Button("Locate Text") with gr.Column(): with gr.Row(height=800): localization_output = gr.AnnotatedImage(label="Text Position") gr.Examples( [["assets/localization_example_1.jpeg", "Share your repair"], ["assets/screen2words_ui_example.png", "statistics"]], inputs = [localization_input, query_input], outputs = [localization_output], fn=localize, cache_examples=True, label='Click on any Examples below to get localization results quickly 👇' ) localization_btn.click(fn=localize, inputs=[localization_input, query_input], outputs=localization_output) demo.launch(share = True)