Spaces:
Sleeping
Sleeping
| import subprocess | |
| subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True) | |
| import gradio as gr | |
| import json | |
| import os | |
| from typing import Any, List | |
| import spaces | |
| from PIL import Image, ImageDraw | |
| import requests | |
| from transformers import AutoModelForImageTextToText, AutoProcessor | |
| from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize | |
| import torch | |
| import re | |
| # --- Configuration --- | |
| MODEL_ID = "Hcompany/Holo1-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, | |
| attn_implementation="flash_attention_2", | |
| 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) | |
| # --- Helper functions from the model card (or adapted) --- | |
| def get_localization_prompt(pil_image: Image.Image, instruction: str) -> List[dict[str, Any]]: | |
| """ | |
| Prepares the prompt structure for the Holo1 model for localization tasks. | |
| The `pil_image` argument here is primarily for semantic completeness in the prompt structure, | |
| as the actual image tensor is handled by the processor later. | |
| """ | |
| guidelines: str = "Localize an element on the GUI image according to my instructions and output a click position as Click(x, y) with x num pixels from the left edge and y num pixels from the top edge." | |
| return [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| { | |
| "type": "image", | |
| "image": pil_image, | |
| }, | |
| {"type": "text", "text": f"{guidelines}\n{instruction}"}, | |
| ], | |
| } | |
| ] | |
| 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 predict_click_location(input_pil_image: Image.Image, instruction: 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 instruction or instruction.strip() == "": | |
| return "No instruction provided. Please type an instruction.", 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. | |
| 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 instruction | |
| messages = get_localization_prompt(resized_image, instruction) | |
| # 3. Run inference | |
| # Pass `messages` (which includes the image object for template processing) | |
| # and `resized_image` (for actual tensor conversion). | |
| try: | |
| coordinates_str = run_inference_localization(messages, 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 on the image | |
| output_image_with_click = resized_image.copy().convert("RGB") # Ensure it's RGB for drawing | |
| parsed_coords = None | |
| # Expected format from the model: "Click(x, y)" | |
| match = re.search(r"Click\((\d+),\s*(\d+)\)", coordinates_str) | |
| if match: | |
| try: | |
| x = int(match.group(1)) | |
| y = int(match.group(2)) | |
| parsed_coords = (x, y) | |
| draw = ImageDraw.Draw(output_image_with_click) | |
| # Make the marker somewhat responsive to image size, but not too small/large | |
| radius = max(5, min(resized_width // 100, resized_height // 100, 15)) | |
| # Define the bounding box for the ellipse (circle) | |
| bbox = (x - radius, y - radius, x + radius, y + radius) | |
| draw.ellipse(bbox, outline="red", width=max(2, radius // 4)) # Draw a red circle | |
| print(f"Predicted and drawn click at: ({x}, {y}) on resized image ({resized_width}x{resized_height})") | |
| except ValueError: | |
| print(f"Could not parse integers from coordinates: {coordinates_str}") | |
| # Keep original coordinates_str, output_image_with_click will be the resized image without a mark | |
| except Exception as e: | |
| print(f"Error drawing on image: {e}") | |
| else: | |
| print(f"Could not parse 'Click(x, y)' from model output: {coordinates_str}") | |
| return coordinates_str, output_image_with_click | |
| # --- Load Example Data --- | |
| example_image = None | |
| example_instruction = "Select July 14th as the check-out date" | |
| try: | |
| example_image_url = "https://huggingface.co/Hcompany/Holo1-7B/resolve/main/calendar_example.jpg" | |
| example_image = Image.open(requests.get(example_image_url, stream=True).raw) | |
| except Exception as e: | |
| print(f"Could not load example image from URL: {e}") | |
| # Create a placeholder image if loading fails, so Gradio example still works | |
| try: | |
| example_image = Image.new("RGB", (200, 150), color="lightgray") | |
| draw = ImageDraw.Draw(example_image) | |
| draw.text((10, 10), "Example image\nfailed to load", fill="black") | |
| except: # If PIL itself is an issue (unlikely here but good for robustness) | |
| pass | |
| # --- Gradio Interface Definition --- | |
| title = "Holo1-7B: Action VLM Localization Demo" | |
| description = """ | |
| This demo showcases **Holo1-7B**, an Action Vision-Language Model developed by HCompany, fine-tuned from Qwen/Qwen2.5-VL-7B-Instruct. | |
| It's designed to interact with web interfaces like a human user. Here, we demonstrate its UI localization capability. | |
| **How to use:** | |
| 1. Upload an image (e.g., a screenshot of a UI, like the calendar example). | |
| 2. Provide a textual instruction (e.g., "Select July 14th as the check-out date"). | |
| 3. The model will predict the click coordinates in the format `Click(x, y)`. | |
| 4. The predicted click point will be marked with a red circle on the (resized) image. | |
| The model processes a resized version of your input image. Coordinates are relative to this resized image. | |
| """ | |
| article = f""" | |
| <p style='text-align: center'> | |
| Model: <a href='https://huggingface.co/{MODEL_ID}' target='_blank'>{MODEL_ID}</a> by HCompany | | |
| Paper: <a href='https://cdn.prod.website-files.com/67e2dbd9acff0c50d4c8a80c/683ec8095b353e8b38317f80_h_tech_report_v1.pdf' target='_blank'>HCompany Tech Report</a> | | |
| Blog: <a href='https://www.hcompany.ai/surfer-h' target='_blank'>Surfer-H Blog Post</a> | |
| </p> | |
| """ | |
| if not model_loaded: | |
| with gr.Blocks() as demo: | |
| gr.Markdown(f"# <center>⚠️ Error: Model Failed to Load ⚠️</center>") | |
| gr.Markdown(f"<center>{load_error_message}</center>") | |
| gr.Markdown("<center>Please check the console output for more details. Reloading the space might help if it's a temporary issue.</center>") | |
| else: | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| gr.Markdown(f"<h1 style='text-align: center;'>{title}</h1>") | |
| # gr.Markdown(description) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| input_image_component = gr.Image(type="pil", label="Input UI Image", height=400) | |
| instruction_component = gr.Textbox( | |
| label="Instruction", | |
| placeholder="e.g., Click the 'Login' button", | |
| info="Type the action you want the model to localize on the image." | |
| ) | |
| submit_button = gr.Button("Localize Click", variant="primary") | |
| with gr.Column(scale=1): | |
| output_coords_component = gr.Textbox(label="Predicted Coordinates (Format: Click(x,y))", interactive=False) | |
| output_image_component = gr.Image(type="pil", label="Image with Predicted Click Point", height=400, interactive=False) | |
| if example_image: | |
| gr.Examples( | |
| examples=[[example_image, example_instruction]], | |
| inputs=[input_image_component, instruction_component], | |
| outputs=[output_coords_component, output_image_component], | |
| fn=predict_click_location, | |
| cache_examples="lazy", | |
| ) | |
| gr.Markdown(article) | |
| submit_button.click( | |
| fn=predict_click_location, | |
| inputs=[input_image_component, instruction_component], | |
| outputs=[output_coords_component, output_image_component] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(debug=True) |