app.py

import subprocess

subprocess.run(
    "pip install flash-attn --no-build-isolation", env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"}, shell=True
)

from typing import Any, List, Literal

import gradio as gr
import requests
import spaces
import torch
from PIL import Image, ImageDraw
from pydantic import BaseModel, Field
from transformers import AutoModelForImageTextToText, AutoProcessor
from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize

# --- 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) ---


@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 ""


SYSTEM_PROMPT: str = """Imagine you are a robot browsing the web, just like humans. Now you need to complete a task.
In each iteration, you will receive an Observation that includes the last  screenshots of a web browser and the current memory of the agent.
You have also information about the step that the agent is trying to achieve to solve the task.
Carefully analyze the visual information to identify what to do, then follow the guidelines to choose the following action.
You should detail your thought (i.e. reasoning steps) before taking the action.
Also detail in the notes field of the action the extracted information relevant to solve the task.
Once you have enough information in the notes to answer the task, return an answer action with the detailed answer in the notes field.
This will be evaluated by an evaluator and should match all the criteria or requirements of the task.

Guidelines:
- store in the notes all the relevant information to solve the task that fulfill the task criteria. Be precise
- Use both the task and the step information to decide what to do
- if you want to write in a text field and the text field already has text, designate the text field by the text it contains and its type
- If there is a cookies notice, always accept all the cookies first
- The observation is the screenshot of the current page and the memory of the agent.
- If you see relevant information on the screenshot to answer the task, add it to the notes field of the action.
- If there is no relevant information on the screenshot to answer the task, add an empty string to the notes field of the action.
- If you see buttons that allow to navigate directly to relevant information, like jump to ... or go to ... , use them to navigate faster.
- In the answer action, give as many details a possible relevant to answering the task.
- if you want to write, don't click before. Directly use the write action
- to write, identify the web element which is type and the text it already contains
- If you want to use a search bar, directly write text in the search bar
- Don't scroll too much. Don't scroll if the number of scrolls is greater than 3
- Don't scroll if you are at the end of the webpage
- Only refresh if you identify a rate limit problem
- If you are looking for a single flights, click on round-trip to select 'one way'
- Never try to login, enter email or password. If there is a need to login, then go back.
- If you are facing a captcha on a website, try to solve it.

- if you have enough information in the screenshot and in the notes to answer the task, return an answer action with the detailed answer in the notes field
- The current date is {timestamp}.

# <output_json_format>
# ```json
# {output_format}
# ```
# </output_json_format>

"""


class ClickElementAction(BaseModel):
    """Click at absolute coordinates of a web element with its description"""

    action: Literal["click_element"] = Field(description="Click at absolute coordinates of a web element")
    element: str = Field(description="text description of the element")
    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 log(self):
        return f"I have clicked on the element '{self.element}' at absolute coordinates {self.x}, {self.y}"


class WriteElementAction(BaseModel):
    """Write content at absolute coordinates of a web element identified by its description, then press Enter."""

    action: Literal["write_element_abs"] = Field(description="Write content at absolute coordinates of a web page")
    content: str = Field(description="Content to write")
    element: str = Field(description="Text description of the element")
    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 log(self):
        return f"I have written '{self.content}' in the element '{self.element}' at absolute coordinates {self.x}, {self.y}"


class ScrollAction(BaseModel):
    """Scroll action with no required element"""

    action: Literal["scroll"] = Field(description="Scroll the page or a specific element")
    direction: Literal["down", "up", "left", "right"] = Field(description="The direction to scroll in")

    def log(self):
        return f"I have scrolled {self.direction}"


class GoBackAction(BaseModel):
    """Action to navigate back in browser history"""

    action: Literal["go_back"] = Field(description="Navigate to the previous page")

    def log(self):
        return "I have gone back to the previous page"


class RefreshAction(BaseModel):
    """Action to refresh the current page"""

    action: Literal["refresh"] = Field(description="Refresh the current page")

    def log(self):
        return "I have refreshed the page"


class GotoAction(BaseModel):
    """Action to go to a particular URL"""

    action: Literal["goto"] = Field(description="Goto a particular URL")
    url: str = Field(description="A url starting with http:// or https://")

    def log(self):
        return f"I have navigated to the URL {self.url}"


class WaitAction(BaseModel):
    """Action to wait for a particular amount of time"""

    action: Literal["wait"] = Field(description="Wait for a particular amount of time")
    seconds: int = Field(default=2, ge=0, le=10, description="The number of seconds to wait")

    def log(self):
        return f"I have waited for {self.seconds} seconds"


class RestartAction(BaseModel):
    """Restart the task from the beginning."""

    action: Literal["restart"] = "restart"

    def log(self):
        return "I have restarted the task from the beginning"


class AnswerAction(BaseModel):
    """Return a final answer to the task. This is the last action to call in an episode."""

    action: Literal["answer"] = "answer"
    content: str = Field(description="The answer content")

    def log(self):
        return f"I have answered the task with '{self.content}'"


ActionSpace = (
    ClickElementAction
    | WriteElementAction
    | ScrollAction
    | GoBackAction
    | RefreshAction
    | WaitAction
    | RestartAction
    | AnswerAction
    | GotoAction
)


class NavigationStep(BaseModel):
    note: str = Field(
        default="",
        description="Task-relevant information extracted from the previous observation. Keep empty if no new info.",
    )
    thought: str = Field(description="Reasoning about next steps (<4 lines)")
    action: ActionSpace = Field(description="Next action to take")


def get_navigation_prompt(task, image, step=1):
    """
    Get the prompt for the navigation task.
    - task: The task to complete
    - image: The current screenshot of the web page
    - step: The current step of the task
    """
    system_prompt = SYSTEM_PROMPT.format(
        output_format=NavigationStep.model_json_schema(),
        timestamp="2025-06-04 14:16:03",
    )
    return [
        {
            "role": "system",
            "content": [
                {"type": "text", "text": system_prompt},
            ],
        },
        {
            "role": "user",
            "content": [
                {"type": "text", "text": f"<task>\n{task}\n</task>\n"},
                {"type": "text", "text": f"<observation step={step}>\n"},
                {"type": "text", "text": "<screenshot>\n"},
                {
                    "type": "image",
                    "image": image,
                },
                {"type": "text", "text": "\n</screenshot>\n"},
                {"type": "text", "text": "\n</observation>\n"},
            ],
        },
    ]


# --- Gradio processing function ---
def navigate(input_pil_image: Image.Image, 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.
    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_navigation_prompt(task, resized_image, step=1)

    # 3. Run inference
    #    Pass `messages` (which includes the image object for template processing)
    #    and `resized_image` (for actual tensor conversion).
    try:
        navigation_str = 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")

    return navigation_str
    # return NavigationStep(**json.loads(navigation_str))


# --- Load Example Data ---
example_image = None
example_task = "Book a hotel in Paris on August 3rd for 3 nights"
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 Navigation 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 task (e.g., "Book a hotel in Paris on August 3rd for 3 nights").
3. The model will predict the  navigation step.

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)
                task_component = gr.Textbox(
                    label="task",
                    placeholder="e.g., Book a hotel in Paris on August 3rd for 3 nights",
                    info="Type the task you want the model to complete.",
                )
                submit_button = gr.Button("Navigate", variant="primary")

            with gr.Column(scale=1):
                output_coords_component = gr.Textbox(label="Navigation Step", interactive=False)

        if example_image:
            gr.Examples(
                examples=[[example_image, example_task]],
                inputs=[input_image_component, task_component],
                outputs=[output_coords_component],
                fn=navigate,
                cache_examples="lazy",
            )

        gr.Markdown(article)

        submit_button.click(
            fn=navigate,
            inputs=[input_image_component, task_component],
            outputs=[output_coords_component],
        )

if __name__ == "__main__":
    demo.launch(debug=True)