import os import re import traceback from datetime import datetime from typing import Any, Literal import gradio as gr import numpy as np import requests import spaces import torch from PIL import Image, ImageDraw from pydantic import BaseModel, Field from transformers import AutoProcessor from transformers.models.auto.modeling_auto import AutoModelForImageTextToText from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize # --- Configuration --- MODEL_ID = "Hcompany/Holo1.5-7B" # --- Model and Processor Loading (Load once) --- print(f"Loading model and processor for {MODEL_ID}...") model = None processor = None model_loaded = False load_error_message = "" try: model = AutoModelForImageTextToText.from_pretrained( MODEL_ID, torch_dtype=torch.bfloat16, trust_remote_code=True ).to("cuda") processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True) model_loaded = True print("Model and processor loaded successfully.") except Exception as e: load_error_message = ( f"Error loading model/processor: {e}\n" "This might be due to network issues, an incorrect model ID, or missing dependencies (like flash_attention_2 if enabled by default in some config).\n" "Ensure you have a stable internet connection and the necessary libraries installed." ) print(load_error_message) title = "Holo1.5-7B: Localization VLM Demo" description = """ This demo showcases [**Holo1.5-7B**](https://huggingface.co/Hcompany/Holo1.5-7B), a new version of the Action Vision-Language Model developed by HCompany, fine-tuned from Qwen/Qwen2.5-VL-7B-Instruct. It's designed to perform complex navigation tasks in Web, Android, and Desktop interfaces. **How to use:** 1. Upload an image (e.g., a screenshot of a UI, see example below). 2. Provide a target UI element (e.g., "Docs tab"). 3. The model will predict the coordinates of the element on the screenshot. The model processor resizes your input image. Coordinates are relative to this resized image. """ def array_to_image_path(image_array): if image_array is None: raise ValueError("No image provided. Please upload an image before submitting.") # Convert numpy array to PIL Image img = Image.fromarray(np.uint8(image_array)) # Generate a unique filename using timestamp timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") filename = f"image_{timestamp}.png" # Save the image img.save(filename) # Get the full path of the saved image full_path = os.path.abspath(filename) return full_path LOCALIZATION_PROMPT: str = """Localize an element on the GUI image according to the provided target and output a click position. * Only output the click position, do not output any other text. * The click position should be in the format 'Click(x, y)' with x: num pixels from the left edge and y: num pixels from the top edge Your target is:""" class ClickAbsoluteAction(BaseModel): """Click at absolute coordinates.""" action: Literal["click_absolute"] = "click_absolute" x: int = Field(description="The x coordinate, number of pixels from the left edge.") y: int = Field(description="The y coordinate, number of pixels from the top edge.") def get_localization_prompt(component, image, step=1): """ Get the prompt for the localization task. - component: The component to localize - image: The current screenshot of the web page - step: The current step of the task """ return [ { "role": "user", "content": [ { "type": "image", "image": image, }, {"type": "text", "text": LOCALIZATION_PROMPT + "\n" + component}, ], }, ] def array_to_image(image_array: np.ndarray) -> Image.Image: if image_array is None: raise ValueError("No image provided. Please upload an image before submitting.") # Convert numpy array to PIL Image img = Image.fromarray(np.uint8(image_array)) return img @spaces.GPU(duration=20) def run_inference_localization( messages_for_template: list[dict[str, Any]], pil_image_for_processing: Image.Image ) -> str: model.to("cuda") torch.cuda.set_device(0) """ Runs inference using the Holo1 model. - messages_for_template: The prompt structure, potentially including the PIL image object (which apply_chat_template converts to an image tag). - pil_image_for_processing: The actual PIL image to be processed into tensors. """ # 1. Apply chat template to messages. This will create the text part of the prompt, # including image tags if the image was part of `messages_for_template`. text_prompt = processor.apply_chat_template(messages_for_template, tokenize=False, add_generation_prompt=True) # 2. Process text and image together to get model inputs inputs = processor( text=[text_prompt], images=[pil_image_for_processing], # Provide the actual image data here padding=True, return_tensors="pt", ) inputs = inputs.to(model.device) # 3. Generate response # Using do_sample=False for more deterministic output, as in the model card's structured output example generated_ids = model.generate(**inputs, max_new_tokens=128, do_sample=False) # 4. Trim input_ids from generated_ids to get only the generated part generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)] # 5. Decode the generated tokens decoded_output = processor.batch_decode( generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False ) return decoded_output[0] if decoded_output else "" # --- Gradio processing function --- def localize(input_numpy_image: np.ndarray, task: str) -> str: # if not model_loaded or not processor or not model: # return f"Model not loaded. Error: {load_error_message}", None # if not input_pil_image: # return "No image provided. Please upload an image.", None # if not task or task.strip() == "": # return "No task provided. Please type an task.", input_pil_image.copy().convert("RGB") # 1. Prepare image: Resize according to model's image processor's expected properties # This ensures predicted coordinates match the (resized) image dimensions. input_pil_image = array_to_image(input_numpy_image) assert isinstance(input_pil_image, Image.Image) image_proc_config = processor.image_processor try: resized_height, resized_width = smart_resize( input_pil_image.height, input_pil_image.width, factor=image_proc_config.patch_size * image_proc_config.merge_size, min_pixels=image_proc_config.min_pixels, max_pixels=image_proc_config.max_pixels, ) # Using LANCZOS for resampling as it's generally good for downscaling. # The model card used `resample=None`, which might imply nearest or default. # For visual quality in the demo, LANCZOS is reasonable. resized_image = input_pil_image.resize( size=(resized_width, resized_height), resample=Image.Resampling.LANCZOS, # type: ignore ) except Exception as e: print(f"Error resizing image: {e}") return f"Error resizing image: {e}", input_pil_image.copy().convert("RGB") # 2. Create the prompt using the resized image (for correct image tagging context) and task prompt = get_localization_prompt(task, resized_image, step=1) print("Prompt:") print(prompt) # 3. Run inference # Pass `messages` (which includes the image object for template processing) # and `resized_image` (for actual tensor conversion). try: localization = run_inference_localization(prompt, resized_image) except Exception as e: print(f"Error during model inference: {e}") return f"Error during model inference: {e}", resized_image.copy().convert("RGB") # 4) Parse coordinates and draw marker output_image_with_click = resized_image.copy().convert("RGB") match = re.search(r"Click\((\d+),\s*(\d+)\)", localization) if match: try: x = int(match.group(1)) y = int(match.group(2)) draw = ImageDraw.Draw(output_image_with_click) radius = max(5, min(resized_width // 100, resized_height // 100, 15)) bbox = (x - radius, y - radius, x + radius, y + radius) draw.ellipse(bbox, outline="red", width=max(2, radius // 4)) print(f"Predicted and drawn click at: ({x}, {y}) on resized image ({resized_width}x{resized_height})") except Exception as e: print(f"Error drawing on image: {e}") traceback.print_exc() else: print(f"Could not parse 'Click(x, y)' from model output: {localization}") return localization, output_image_with_click # --- Load Example Data --- example_image_url = "https://huggingface.co/spaces/Hcompany/Holo1.5-Localization/resolve/main/desktop_3.png" example_image = Image.open(requests.get(example_image_url, stream=True).raw) example_task = "Email quote for Hyundai Kona" with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown(f"

{title}

") gr.Markdown(description) with gr.Row(): with gr.Column(): input_image_component = gr.Image(label="Input UI Image", height=400) task_component = gr.Textbox( label="component", placeholder="Email quote for Hyundai Kona", info="Describe the UI component to find.", ) submit_button = gr.Button("Localize", variant="primary") with gr.Column(): output_coords_component = gr.Textbox(label="Localization Step") output_image_component = gr.Image( type="pil", label="Image with coordinates of the component", height=400, interactive=False ) submit_button.click( localize, [input_image_component, task_component], [output_coords_component, output_image_component] ) gr.Examples( examples=[[example_image, example_task]], inputs=[input_image_component, task_component], outputs=[output_coords_component, output_image_component], fn=localize, cache_examples="lazy", ) demo.queue(api_open=False) demo.launch(debug=True)