#!/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 = ["", "", "", ""] 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="", 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 == '': parsed_boxes = parsed_answer['']['bboxes'] draw = florence_draw_bboxes(image, parsed_answer) output_text = "None" elif prompt == '': output_text = parsed_answer[''] 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"""

LLM & VLM Demo

Use the different LLMs or VLMs to experience the different models. Note: 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= "") 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()