Add command line option to run omniparser gradio app.py without VM (--no_vm)

.gitignore

@@ -8,5 +8,8 @@ __pycache__/
 debug.ipynb
 util/__pycache__/
 index.html?linkid=2289031
+tmp/
 wget-log
-weights/icon_caption_florence_v2/
+weights/icon_caption_florence_v2/
+venv/
+*.bat

README.md

@@ -29,6 +29,8 @@ conda create -n "omni" python==3.12
 conda activate omni
 pip install -r requirements.txt
 ```
+If you have an nvidia gpu with cuda toolkit 12.4 installed, use pip install -r requirements_gpu.txt instead to enable gpu support.  You can modify the requirements_gpu for your cuda version if you have a version other than 12.4.
+
 
 Ensure you have the V2 weights downloaded in weights folder (ensure caption weights folder is called icon_caption_florence). If not download them with:
 ```
@@ -48,6 +50,7 @@ python weights/convert_safetensor_to_pt.py
 For v1.5: 
 download 'model_v1_5.pt' from https://huggingface.co/microsoft/OmniParser/tree/main/icon_detect_v1_5, make a new dir: weights/icon_detect_v1_5, and put it inside the folder. No weight conversion is needed. 
 ``` -->
+If you have a windows machine you can run weights_download.ps1 from a powershell terminal inside of the root of the project instead to download the model weights.
 
 ## Examples:
 We put together a few simple examples in the demo.ipynb. 
@@ -58,6 +61,20 @@ To run gradio demo, simply run:
 python gradio_demo.py
 ```
 
+## Webui App
+To run the web ui app demo, run:
+```python
+python omnitool/gradio/app.py
+```
+You can use the command line argument --no_vm to run the webui actions on your local computer rather than in a windows virtual machine:
+```python
+python omnitool/gradio/app.py --no_vm
+```
+Remember for either option you have to start the omniparserserver first before running the web app:
+```python
+python omnitool/omniparserserver/omniparserserver.py
+```
+
 ## Model Weights License
 For the model checkpoints on huggingface model hub, please note that icon_detect model is under AGPL license since it is a license inherited from the original yolo model. And icon_caption_blip2 & icon_caption_florence is under MIT license. Please refer to the LICENSE file in the folder of each model: https://huggingface.co/microsoft/OmniParser.
 

omnitool/gradio/agent/llm_utils/omniparserclient.py

@@ -1,3 +1,4 @@
+import os
 import requests
 import base64
 from pathlib import Path
@@ -10,9 +11,10 @@ class OmniParserClient:
     def __init__(self, 
                  url: str) -> None:
         self.url = url
+        self.use_vm = (os.getenv("OMNIPARSER_NO_VM", "False") == "False")
 
     def __call__(self,):
-        screenshot, screenshot_path = get_screenshot()
+        screenshot, screenshot_path = get_screenshot(using_vm = self.use_vm)
         screenshot_path = str(screenshot_path)
         image_base64 = encode_image(screenshot_path)
         response = requests.post(self.url, json={"base64_image": image_base64})

omnitool/gradio/app.py

@@ -31,15 +31,17 @@
 
 Type a message and press submit to start OmniTool. Press stop to pause, and press the trash icon in the chat to clear the message history.
 '''
+args = None
 
 def parse_arguments():
 
     parser = argparse.ArgumentParser(description="Gradio App")
     parser.add_argument("--windows_host_url", type=str, default='localhost:8006')
     parser.add_argument("--omniparser_server_url", type=str, default="localhost:8000")
+    parser.add_argument("--no_vm", action='store_true', help='Run omniparser on local machine without a virtual windows machine')
     return parser.parse_args()
 args = parse_arguments()
-
+os.environ["OMNIPARSER_NO_VM"] = str(args.no_vm) 
 
 class Sender(StrEnum):
     USER = "user"
@@ -195,8 +197,13 @@ def valid_params(user_input, state):
             url = f'http://{url}/probe'
             response = requests.get(url, timeout=3)
             if response.status_code != 200:
-                errors.append(f"{server_name} is not responding")
+                # Throw error if either the OmniParser server does not respond or the Windows VM does not respond and the user has not selected the no_vm command line argument
+                if server_name == 'OmniParser Server' or (server_name == 'Windows Host' and not args.no_vm):
+                    errors.append(f"{server_name} is not responding")
         except RequestException as e:
+            # Skip error handling if server is Windows VM but user has selected no_vm command line argument
+            if server_name == 'Windows Host' and args.no_vm:
+                continue
             errors.append(f"{server_name} is not responding")
 
     if not state["api_key"].strip():

omnitool/gradio/tools/computer.py

@@ -1,3 +1,4 @@
+import os
 import base64
 import time
 from enum import StrEnum
@@ -11,12 +12,15 @@
 from .screen_capture import get_screenshot
 import requests
 import re
+import pyautogui
 
 OUTPUT_DIR = "./tmp/outputs"
 
 TYPING_DELAY_MS = 12
 TYPING_GROUP_SIZE = 50
 
+use_vm: bool = True
+
 Action = Literal[
     "key",
     "type",
@@ -90,13 +94,15 @@ def to_params(self) -> BetaToolComputerUse20241022Param:
 
     def __init__(self, is_scaling: bool = False):
         super().__init__()
-
+        global use_vm
+        use_vm = (os.getenv("OMNIPARSER_NO_VM", "False") == "False")
         # Get screen width and height using Windows command
         self.display_num = None
         self.offset_x = 0
         self.offset_y = 0
         self.is_scaling = is_scaling
         self.width, self.height = self.get_screen_size()
+
         print(f"screen size: {self.width}, {self.height}")
 
         self.key_conversion = {"Page_Down": "pagedown",
@@ -141,11 +147,18 @@ async def __call__(
             print(f"mouse move to {x}, {y}")
 
             if action == "mouse_move":
-                self.send_to_vm(f"pyautogui.moveTo({x}, {y})")
+                if use_vm:
+                    self.send_to_vm(f"pyautogui.moveTo({x}, {y})")
+                else:
+                    pyautogui.moveTo(x, y)
                 return ToolResult(output=f"Moved mouse to ({x}, {y})")
             elif action == "left_click_drag":
-                current_x, current_y = self.send_to_vm("pyautogui.position()")
-                self.send_to_vm(f"pyautogui.dragTo({x}, {y}, duration=0.5)")
+                if use_vm:
+                    current_x, current_y = self.send_to_vm("pyautogui.position()")
+                    self.send_to_vm(f"pyautogui.dragTo({x}, {y}, duration=0.5)")
+                else:
+                    current_x, current_y = pyautogui.position()
+                    pyautogui.dragTo(x, y, duration=0.5) 
                 return ToolResult(output=f"Dragged mouse from ({current_x}, {current_y}) to ({x}, {y})")
 
         if action in ("key", "type"):
@@ -162,19 +175,31 @@ async def __call__(
                 for key in keys:
                     key = self.key_conversion.get(key.strip(), key.strip())
                     key = key.lower()
-                    self.send_to_vm(f"pyautogui.keyDown('{key}')")  # Press down each key
+                    if use_vm:
+                        self.send_to_vm(f"pyautogui.keyDown('{key}')")  # Press down each key
+                    else:
+                        pyautogui.keyDown(key)
                 for key in reversed(keys):
                     key = self.key_conversion.get(key.strip(), key.strip())
                     key = key.lower()
-                    self.send_to_vm(f"pyautogui.keyUp('{key}')")    # Release each key in reverse order
+                    if use_vm:
+                        self.send_to_vm(f"pyautogui.keyUp('{key}')")    # Release each key in reverse order
+                    else:
+                        pyautogui.keyUp(key)
                 return ToolResult(output=f"Pressed keys: {text}")
 
             elif action == "type":
                 # default click before type TODO: check if this is needed
-                self.send_to_vm("pyautogui.click()")
-                self.send_to_vm(f"pyautogui.typewrite('{text}', interval={TYPING_DELAY_MS / 1000})")
-                self.send_to_vm("pyautogui.press('enter')")
-                screenshot_base64 = (await self.screenshot()).base64_image
+                if use_vm:
+                    self.send_to_vm("pyautogui.click()")
+                    self.send_to_vm(f"pyautogui.typewrite('{text}', interval={TYPING_DELAY_MS / 1000})")
+                    self.send_to_vm("pyautogui.press('enter')")
+                else:
+                    pyautogui.click()
+                    pyautogui.typewrite(text, interval=(TYPING_DELAY_MS/1000))
+                    pyautogui.press('enter')
+                screenshot = await self.screenshot()
+                screenshot_base64 = screenshot.base64_image
                 return ToolResult(output=text, base64_image=screenshot_base64)
 
         if action in (
@@ -194,28 +219,54 @@ async def __call__(
             if action == "screenshot":
                 return await self.screenshot()
             elif action == "cursor_position":
-                x, y = self.send_to_vm("pyautogui.position()")
-                x, y = self.scale_coordinates(ScalingSource.COMPUTER, x, y)
+                if use_vm:
+                    x, y = self.send_to_vm("pyautogui.position()")
+                    x, y = self.scale_coordinates(ScalingSource.COMPUTER, x, y)
+                else:
+                    x, y = pyautogui.position()
                 return ToolResult(output=f"X={x},Y={y}")
             else:
                 if action == "left_click":
-                    self.send_to_vm("pyautogui.click()")
+                    if use_vm:
+                        self.send_to_vm("pyautogui.click()")
+                    else:
+                        pyautogui.click()
                 elif action == "right_click":
-                    self.send_to_vm("pyautogui.rightClick()")
+                    if use_vm:
+                        self.send_to_vm("pyautogui.rightClick()")
+                    else:
+                        pyautogui.rightClick()
                 elif action == "middle_click":
-                    self.send_to_vm("pyautogui.middleClick()")
+                    if use_vm:
+                        self.send_to_vm("pyautogui.middleClick()")
+                    else:
+                        pyautogui.middleClick()
                 elif action == "double_click":
-                    self.send_to_vm("pyautogui.doubleClick()")
+                    if use_vm:
+                        self.send_to_vm("pyautogui.doubleClick()")
+                    else:
+                        pyautogui.doubleClick()
                 elif action == "left_press":
-                    self.send_to_vm("pyautogui.mouseDown()")
-                    time.sleep(1)
-                    self.send_to_vm("pyautogui.mouseUp()")
+                    if use_vm:
+                        self.send_to_vm("pyautogui.mouseDown()")
+                        time.sleep(1)
+                        self.send_to_vm("pyautogui.mouseUp()")
+                    else:
+                        pyautogui.mouseDown()
+                        time.sleep(1)
+                        pyautogui.mouseUp()
                 return ToolResult(output=f"Performed {action}")
         if action in ("scroll_up", "scroll_down"):
             if action == "scroll_up":
-                self.send_to_vm("pyautogui.scroll(100)")
+                if use_vm:
+                    self.send_to_vm("pyautogui.scroll(100)")
+                else:
+                    pyautogui.scroll(100)
             elif action == "scroll_down":
-                self.send_to_vm("pyautogui.scroll(-100)")
+                if use_vm:
+                    self.send_to_vm("pyautogui.scroll(-100)")
+                else:
+                    pyautogui.scroll(-100)
             return ToolResult(output=f"Performed {action}")
         if action == "hover":
             return ToolResult(output=f"Performed {action}")
@@ -261,7 +312,7 @@ async def screenshot(self):
             screenshot = self.padding_image(screenshot)
             self.target_dimension = MAX_SCALING_TARGETS["WXGA"]
         width, height = self.target_dimension["width"], self.target_dimension["height"]
-        screenshot, path = get_screenshot(resize=True, target_width=width, target_height=height)
+        screenshot, path = get_screenshot(resize=True, target_width=width, target_height=height, using_vm=use_vm)
         time.sleep(0.7) # avoid async error as actions take time to complete
         return ToolResult(base64_image=base64.b64encode(path.read_bytes()).decode())
 
@@ -310,20 +361,25 @@ def scale_coordinates(self, source: ScalingSource, x: int, y: int):
     def get_screen_size(self):
         """Return width and height of the screen"""
         try:
-            response = requests.post(
-                f"http://localhost:5000/execute",
-                headers={'Content-Type': 'application/json'},
-                json={"command": ["python", "-c", "import pyautogui; print(pyautogui.size())"]},
-                timeout=90
-            )
-            if response.status_code != 200:
-                raise ToolError(f"Failed to get screen size. Status code: {response.status_code}")
-
-            output = response.json()['output'].strip()
-            match = re.search(r'Size\(width=(\d+),\s*height=(\d+)\)', output)
-            if not match:
-                raise ToolError(f"Could not parse screen size from output: {output}")
-            width, height = map(int, match.groups())
-            return width, height
+            # Use VM to get screensize if using VM, otherwise use direct code
+            if use_vm:
+                response = requests.post(
+                    f"http://localhost:5000/execute",
+                    headers={'Content-Type': 'application/json'},
+                    json={"command": ["python", "-c", "import pyautogui; print(pyautogui.size())"]},
+                    timeout=90
+                )
+                if response.status_code != 200:
+                    raise ToolError(f"Failed to get screen size. Status code: {response.status_code}")
+
+                output = response.json()['output'].strip()
+                match = re.search(r'Size\(width=(\d+),\s*height=(\d+)\)', output)
+                if not match:
+                    raise ToolError(f"Could not parse screen size from output: {output}")
+                width, height = map(int, match.groups())
+                return width, height
+            else:
+                width, height = pyautogui.size()
+                return width, height
         except requests.exceptions.RequestException as e:
             raise ToolError(f"An error occurred while trying to get screen size: {str(e)}")

omnitool/gradio/tools/screen_capture.py

@@ -4,23 +4,24 @@
 from PIL import Image
 from .base import BaseAnthropicTool, ToolError
 from io import BytesIO
-
+import pyautogui
 OUTPUT_DIR = "./tmp/outputs"
 
-def get_screenshot(resize: bool = False, target_width: int = 1920, target_height: int = 1080):
+def get_screenshot(resize: bool = False, target_width: int = 1920, target_height: int = 1080, using_vm: bool = True):
     """Capture screenshot by requesting from HTTP endpoint - returns native resolution unless resized"""
     output_dir = Path(OUTPUT_DIR)
     output_dir.mkdir(parents=True, exist_ok=True)
     path = output_dir / f"screenshot_{uuid4().hex}.png"
-
     try:
-        response = requests.get('http://localhost:5000/screenshot')
-        if response.status_code != 200:
-            raise ToolError(f"Failed to capture screenshot: HTTP {response.status_code}")
-
-        # (1280, 800)
-        screenshot = Image.open(BytesIO(response.content))
-
+        if using_vm:
+            response = requests.get('http://localhost:5000/screenshot')
+            if response.status_code != 200:
+                raise ToolError(f"Failed to capture screenshot: HTTP {response.status_code}")
+
+            # (1280, 800)
+            screenshot = Image.open(BytesIO(response.content))
+        else:
+            screenshot = pyautogui.screenshot()
         if resize and screenshot.size != (target_width, target_height):
             screenshot = screenshot.resize((target_width, target_height))
         screenshot.save(path)

requirements_cuda.txt

@@ -0,0 +1,34 @@
+--extra-index-url https://download.pytorch.org/whl/cu124
+torch
+torchvision 
+torchaudio
+easyocr
+supervision==0.18.0
+openai==1.3.5
+transformers
+ultralytics==8.3.70
+azure-identity
+numpy==1.26.4
+opencv-python
+opencv-python-headless
+gradio
+dill
+accelerate
+timm
+einops==0.8.0
+paddlepaddle
+paddleocr
+ruff==0.6.7
+pre-commit==3.8.0
+pytest==8.3.3
+pytest-asyncio==0.23.6
+pyautogui==0.9.54
+streamlit>=1.38.0
+anthropic[bedrock,vertex]>=0.37.1
+jsonschema==4.22.0
+boto3>=1.28.57
+google-auth<3,>=2
+screeninfo
+uiautomation
+dashscope
+groq