Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, AutoProcessor, Qwen2VLForConditionalGeneration | |
| from utils import image_to_base64, rescale_bounding_boxes, draw_bounding_boxes, florence_draw_bboxes | |
| from qwen_vl_utils import process_vision_info | |
| import re | |
| import base64 | |
| import os | |
| llms = { | |
| "Qwen2-1.5B": {"model": "Qwen/Qwen2-1.5B-Instruct", "prefix": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."}, | |
| "Qwen2-3B": {"model": "Qwen/Qwen2-3B-Instruct", "prefix": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."}, | |
| "Qwen2-7B": {"model": "Qwen/Qwen2-7B-Instruct", "prefix": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."}, | |
| "Qwen2.5-1.5B": {"model": "Qwen/Qwen2.5-1.5B-Instruct", "prefix": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."}, | |
| "Qwen2.5-3B": {"model": "Qwen/Qwen2.5-3B-Instruct", "prefix": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."}, | |
| "DeepSeek-Coder-1.3B": {"model": "deepseek-ai/deepseek-coder-1.3b-instruct", "prefix": "You are a helpful assistant."}, | |
| "DeepSeek-r1-Qwen-1.5B": {"model": "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B", "prefix": "You are a helpful assistant."}, | |
| } | |
| vlms = { | |
| "Florence-2-base": {"model": "microsoft/Florence-2-base", "prefix": "help me"}, | |
| "Florence-2-large": {"model": "microsoft/Florence-2-large", "prefix": "help me"}, | |
| "Qwen2-vl-2B": {"model": "Qwen/Qwen2-VL-2B-Instruct", "prefix": "You are a helpfull assistant to detect objects in images. When asked to detect elements based on a description you return bounding boxes for all elements in the form of [xmin, ymin, xmax, ymax] whith the values beeing scaled to 1000 by 1000 pixels. When there are more than one result, answer with a list of bounding boxes in the form of [[xmin, ymin, xmax, ymax], [xmin, ymin, xmax, ymax], ...]."}, | |
| "Qwen2-vl-7B": {"model": "Qwen/Qwen2-VL-7B-Instruct", "prefix": "You are a helpfull assistant to detect objects in images. When asked to detect elements based on a description you return bounding boxes for all elements in the form of [xmin, ymin, xmax, ymax] whith the values beeing scaled to 1000 by 1000 pixels. When there are more than one result, answer with a list of bounding boxes in the form of [[xmin, ymin, xmax, ymax], [xmin, ymin, xmax, ymax], ...]."}, | |
| "Qwen2.5-vl-3B": {"model": "Qwen/Qwen2.5-VL-3B-Instruct", "prefix": "You are a helpfull assistant to detect objects in images. When asked to detect elements based on a description you return bounding boxes for all elements in the form of [xmin, ymin, xmax, ymax] whith the values beeing scaled to 1000 by 1000 pixels. When there are more than one result, answer with a list of bounding boxes in the form of [[xmin, ymin, xmax, ymax], [xmin, ymin, xmax, ymax], ...]."} | |
| } | |
| tasks = ["<OD>", "<OCR>", "<CAPTION>", "<OCR_WITH_REGION>"] | |
| def get_image_base64(image_path): | |
| with open(image_path, "rb") as image_file: | |
| encoded_string = base64.b64encode(image_file.read()).decode() | |
| return encoded_string | |
| # At the top of your file, after imports | |
| current_dir = os.path.dirname(os.path.abspath(__file__)) | |
| image_path = os.path.join(current_dir, "assets", "hailo_logo.gif") | |
| image_base64 = get_image_base64(image_path) | |
| def run_llm(text_input, model_id="Qwen2-1.5B", prefix=None): | |
| global messages | |
| tokenizer = AutoTokenizer.from_pretrained(llms[model_id]["model"], trust_remote_code=True) | |
| model = AutoModelForCausalLM.from_pretrained(llms[model_id]["model"], trust_remote_code=True) | |
| # Use the provided prefix if available, otherwise fall back to the default | |
| system_prefix = prefix if prefix is not None else llms[model_id]["prefix"] | |
| if messages is None: | |
| messages = [ | |
| {"role": "system", "content": system_prefix}, | |
| {"role": "user", "content": text_input}, | |
| ] | |
| else: | |
| messages.append({"role": "user", "content": text_input}) | |
| text = tokenizer.apply_chat_template ( | |
| messages, | |
| tokenize=False, | |
| add_generation_prompt=True, | |
| ) | |
| model_inputs = tokenizer([text], return_tensors="pt").to(model.device) | |
| generated_ids = model.generate( | |
| **model_inputs, | |
| max_new_tokens=512, | |
| ) | |
| generated_ids = [ | |
| output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids) | |
| ] | |
| response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] | |
| return response | |
| def run_vlm(image, text_input, model_id="Qwen2-vl-2B", prompt="<OD>", custom_prefix=None): | |
| if "Qwen" in model_id: | |
| model = Qwen2VLForConditionalGeneration.from_pretrained(vlms[model_id]["model"], torch_dtype="auto", device_map="auto") | |
| else: | |
| model = AutoModelForCausalLM.from_pretrained(vlms[model_id]["model"], trust_remote_code=True) | |
| processor = AutoProcessor.from_pretrained(vlms[model_id]["model"], trust_remote_code=True) | |
| if "Qwen" in model_id: | |
| # Use custom prefix if provided, otherwise use default from vlms dictionary | |
| prefix_to_use = custom_prefix if custom_prefix is not None else vlms[model_id]["prefix"] | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "image", "image": f"data:image;base64,{image_to_base64(image)}"}, | |
| {"type": "text", "text": prefix_to_use}, | |
| {"type": "text", "text": text_input}, | |
| ], | |
| } | |
| ] | |
| text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) | |
| image_inputs, video_inputs = process_vision_info(messages) | |
| inputs = processor( | |
| text=[text], | |
| images=image_inputs, | |
| videos=video_inputs, | |
| padding=True, | |
| return_tensors="pt", | |
| ).to(model.device) | |
| generated_ids = model.generate(**inputs, max_new_tokens=256) | |
| generated_ids_trimmed = [ | |
| out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids) | |
| ] | |
| output_text = processor.batch_decode( | |
| generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False | |
| ) | |
| print(output_text) | |
| pattern = r'\[\s*([.\d]+)\s*,\s*([.\d]+)\s*,\s*([.\d]+)\s*,\s*([.\d]+)\s*\]' | |
| matches = re.findall(pattern, str(output_text)) | |
| parsed_boxes = [[float(num) for num in match] for match in matches] | |
| scaled_boxes = rescale_bounding_boxes(parsed_boxes, image.width, image.height) | |
| print(scaled_boxes) | |
| draw = draw_bounding_boxes(image, scaled_boxes) | |
| else: | |
| messages = prompt + text_input | |
| inputs = processor(text=messages, images=image, return_tensors="pt").to(model.device) | |
| generated_ids = model.generate( | |
| input_ids=inputs["input_ids"], | |
| pixel_values=inputs["pixel_values"], | |
| max_new_tokens=1024, | |
| early_stopping=False, | |
| do_sample=False, | |
| num_beams=3, | |
| ) | |
| generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0] | |
| parsed_answer = processor.post_process_generation( | |
| generated_text, | |
| task=prompt, | |
| image_size=(image.width, image.height) | |
| ) | |
| print(parsed_answer) | |
| if prompt == '<OD>': | |
| parsed_boxes = parsed_answer['<OD>']['bboxes'] | |
| draw = florence_draw_bboxes(image, parsed_answer) | |
| output_text = "None" | |
| elif prompt == '<OCR>': | |
| output_text = parsed_answer['<OCR>'] | |
| draw = image | |
| parsed_boxes = None | |
| return output_text, parsed_boxes, draw | |
| messages = list() | |
| def reset_conversation(): | |
| global messages | |
| messages = list() | |
| def update_task_dropdown(model): | |
| if "Florence" in model: | |
| return [gr.Dropdown(visible=True), gr.Textbox(value=vlms[model]["prefix"])] | |
| elif model in vlms: | |
| return [gr.Dropdown(visible=False), gr.Textbox(value=vlms[model]["prefix"])] | |
| return [gr.Dropdown(visible=False), gr.Textbox(value="")] | |
| def update_prefix_llm(model): | |
| if model in llms: | |
| return gr.Textbox(value=llms[model]["prefix"], visible=True) | |
| return gr.Textbox(visible=True) | |
| with gr.Blocks() as demo: | |
| gr.Markdown( | |
| f""" | |
| <div style="display: flex; align-items: center; gap: 10px;"> | |
| <img src="data:image/gif;base64,{image_base64}" height="40px" style="margin-right: 10px;"> | |
| <h1 style="margin: 0;">LLM & VLM Demo</h1> | |
| </div> | |
| Use the different LLMs or VLMs to experience the different models. | |
| <u>Note</u>: first use of any model will take more time, for the downloading of the weights. | |
| """) | |
| with gr.Tab(label="LLM"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| model_selector = gr.Dropdown(choices=list(llms.keys()), label="Model", value="Qwen2-1.5B") | |
| text_input = gr.Textbox(label="User Prompt") | |
| prefix_input = gr.Textbox(label="Prefix", value=llms["Qwen2.5-1.5B"]["prefix"]) | |
| submit_btn = gr.Button(value="Submit", variant='primary') | |
| reset_btn = gr.Button(value="Reset conversation", variant='stop') | |
| with gr.Column(): | |
| model_output_text = gr.Textbox(label="Model Output Text") | |
| model_selector.change(update_prefix_llm, inputs=model_selector, outputs=prefix_input) | |
| submit_btn.click(run_llm, | |
| [text_input, model_selector, prefix_input], | |
| [model_output_text]) | |
| reset_btn.click(reset_conversation) | |
| with gr.Tab(label="VLM (WIP)"): | |
| # taken from https://huggingface.co./spaces/maxiw/Qwen2-VL-Detection/blob/main/app.py | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_img = gr.Image(label="Input Image", type="pil", scale=2, height=400) | |
| model_selector = gr.Dropdown(choices=list(vlms.keys()), label="Model", value="Qwen2-vl-2B") | |
| task_select = gr.Dropdown(choices=tasks, label="task", value= "<OD>") | |
| text_input = gr.Textbox(label="User Prompt") | |
| prefix_input = gr.Textbox(label="Prefix") | |
| submit_btn = gr.Button(value="Submit", variant='primary') | |
| with gr.Column(): | |
| model_output_text = gr.Textbox(label="Model Output Text") | |
| parsed_boxes = gr.Textbox(label="Parsed Boxes") | |
| annotated_image = gr.Image(label="Annotated Image", scale=2, height=400) | |
| model_selector.change(update_task_dropdown, | |
| inputs=model_selector, | |
| outputs=[task_select, prefix_input]) | |
| submit_btn.click(run_vlm, | |
| [input_img, text_input, model_selector, task_select, prefix_input], | |
| [model_output_text, parsed_boxes, annotated_image]) | |
| if __name__ == "__main__": | |
| demo.launch() | |