diff --git a/.gitignore b/.gitignore index fe4caf45..2d180855 100644 --- a/.gitignore +++ b/.gitignore @@ -8,5 +8,8 @@ __pycache__/ debug.ipynb util/__pycache__/ index.html?linkid=2289031 +tmp/ wget-log -weights/icon_caption_florence_v2/ \ No newline at end of file +weights/icon_caption_florence_v2/ +venv/ +*.bat \ No newline at end of file diff --git a/README.md b/README.md index 3174c61f..3161edc0 100644 --- a/README.md +++ b/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. diff --git a/omnitool/gradio/agent/llm_utils/omniparserclient.py b/omnitool/gradio/agent/llm_utils/omniparserclient.py index e90ddef8..d01e7829 100644 --- a/omnitool/gradio/agent/llm_utils/omniparserclient.py +++ b/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}) diff --git a/omnitool/gradio/app.py b/omnitool/gradio/app.py index 4711a128..9408b80b 100644 --- a/omnitool/gradio/app.py +++ b/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(): diff --git a/omnitool/gradio/tools/computer.py b/omnitool/gradio/tools/computer.py index 6b91bad2..de8acfc8 100644 --- a/omnitool/gradio/tools/computer.py +++ b/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)}") \ No newline at end of file diff --git a/omnitool/gradio/tools/screen_capture.py b/omnitool/gradio/tools/screen_capture.py index 1c1ad04a..80678ddc 100644 --- a/omnitool/gradio/tools/screen_capture.py +++ b/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) diff --git a/requirements_cuda.txt b/requirements_cuda.txt new file mode 100644 index 00000000..8fe23c35 --- /dev/null +++ b/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 diff --git a/weights_download.ps1 b/weights_download.ps1 new file mode 100644 index 00000000..b1068dfd --- /dev/null +++ b/weights_download.ps1 @@ -0,0 +1,14 @@ +$files = @( + "icon_detect/train_args.yaml", + "icon_detect/model.pt", + "icon_detect/model.yaml", + "icon_caption/config.json", + "icon_caption/generation_config.json", + "icon_caption/model.safetensors" +) + +foreach ($f in $files) { + huggingface-cli download microsoft/OmniParser-v2.0 $f --local-dir weights +} + +Move-Item -Path "weights/icon_caption" -Destination "weights/icon_caption_florence"