Spaces:
Running on Zero
Running on Zero
| """DeskForge models (and their base models), prompt, and action handling. | |
| The model replies with one line of pyautogui code. We never exec() it: the line | |
| is parsed with `ast` into an `Action`, and the action is replayed on the X11 | |
| display with xdotool. | |
| """ | |
| import spaces # MUST come before torch / transformers (ZeroGPU) | |
| import ast | |
| import re | |
| import time | |
| from dataclasses import dataclass | |
| import torch | |
| from PIL import Image, ImageDraw, ImageFont | |
| from transformers import AutoModelForImageTextToText, AutoProcessor | |
| MAX_PIXELS = 2_097_152 # training-time cap on screenshot area | |
| SYSTEM_PROMPT = """You are a computer-use agent operating a desktop graphical interface. At each step you see the user's task, a screenshot of the current screen, and the actions you have already taken. Reply with the next action as pyautogui code and nothing else -- no explanation, no code fence, no commentary. | |
| Coordinates are fractions of the screen, not pixels: x runs from 0.0 at the left edge to 1.0 at the right edge, y from 0.0 at the top to 1.0 at the bottom. Write both with four decimals. | |
| These are the only actions available: | |
| pyautogui.click(x=0.0000, y=0.0000) | |
| pyautogui.doubleClick(x=0.0000, y=0.0000) | |
| pyautogui.rightClick(x=0.0000, y=0.0000) | |
| pyautogui.middleClick(x=0.0000, y=0.0000) | |
| computer.tripleClick(x=0.0000, y=0.0000) | |
| pyautogui.moveTo(x=0.0000, y=0.0000) | |
| pyautogui.dragTo(x=0.0000, y=0.0000, button='left') | |
| pyautogui.scroll(-4) | |
| pyautogui.hscroll(4) | |
| pyautogui.write(message='text to type') | |
| pyautogui.press('enter') | |
| pyautogui.hotkey(['ctrl', 'c']) | |
| computer.wait() | |
| computer.terminate(status='success') | |
| To scroll at a particular place, move there first and then scroll. When the task is finished, or cannot be finished, end with computer.terminate.""" | |
| def fit_qwen(image, max_pixels=MAX_PIXELS): | |
| """Downscale to at most max_pixels (even sides, Lanczos), as in Qwen training.""" | |
| w, h = image.size | |
| if w * h <= max_pixels: | |
| return image | |
| s = (max_pixels / (w * h)) ** 0.5 | |
| return image.resize( | |
| (max(2, int(w * s) // 2 * 2), max(2, int(h * s) // 2 * 2)), Image.LANCZOS | |
| ) | |
| def fit_gemma(image, max_pixels=MAX_PIXELS): | |
| """Downscale to at most max_pixels (bilinear, sides rounded down), as in Gemma training.""" | |
| w, h = image.size | |
| if w * h <= max_pixels: | |
| return image | |
| ratio = w / h | |
| height = (max_pixels / ratio) ** 0.5 | |
| return image.resize((max(1, int(height * ratio)), max(1, int(height))), Image.BILINEAR) | |
| class Family: | |
| deskforge: str # fine-tuned model | |
| base: str # what it was fine-tuned from, for the duel | |
| base_label: str | |
| fit: object | |
| FAMILIES = { | |
| "Qwen3.5-4B": Family("docling-project/DeskForge-Qwen3.5-4B", "Qwen/Qwen3.5-4B", | |
| "Qwen3.5-4B (base)", fit_qwen), | |
| "Gemma4-E4B": Family("docling-project/DeskForge-Gemma4-E4B", "google/gemma-4-E4B-it", | |
| "Gemma4-E4B (base)", fit_gemma), | |
| } | |
| DEFAULT_FAMILY = "Gemma4-E4B" # its base model loses clearly, which makes the clearest duel | |
| def user_text(instruction): | |
| # The model was trained on single-target steps, so every instruction is | |
| # presented as the first step of its own task. | |
| return f"Task: {instruction}\n\nActions already taken:\n(none -- this is the first step)\n\nNext action:" | |
| def chat(instruction): | |
| return [ | |
| {"role": "system", "content": SYSTEM_PROMPT}, | |
| { | |
| "role": "user", | |
| "content": [{"type": "image"}, {"type": "text", "text": user_text(instruction)}], | |
| }, | |
| ] | |
| def _load(name): | |
| return AutoModelForImageTextToText.from_pretrained(name, dtype=torch.bfloat16).eval().to("cuda") | |
| # Each base model uses its fine-tune's processor, so both sides of a duel get | |
| # the screenshot encoded exactly as in DeskForge training (e.g. Gemma's 1,120 | |
| # image tokens), and only the weights differ. | |
| PROCESSORS = {f: AutoProcessor.from_pretrained(fam.deskforge) for f, fam in FAMILIES.items()} | |
| MODELS = {(f, which): _load(getattr(fam, which)) | |
| for f, fam in FAMILIES.items() for which in ("deskforge", "base")} | |
| def _generate(family, which, instruction, screenshot, max_new_tokens): | |
| processor, model = PROCESSORS[family], MODELS[(family, which)] | |
| prompt = processor.apply_chat_template( | |
| chat(instruction.strip()), tokenize=False, add_generation_prompt=True, | |
| enable_thinking=False, | |
| ) | |
| inputs = processor( | |
| text=[prompt], images=[FAMILIES[family].fit(screenshot)], return_tensors="pt" | |
| ).to(model.device) | |
| t0 = time.perf_counter() | |
| with torch.inference_mode(): | |
| output = model.generate(**inputs, max_new_tokens=int(max_new_tokens), do_sample=False) | |
| elapsed = time.perf_counter() - t0 | |
| action = processor.decode( | |
| output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True | |
| ).strip() | |
| return action, elapsed | |
| # four models resident (~50 GB) | |
| def predict(instruction: str, screenshot, max_new_tokens: int = 128, | |
| family: str = DEFAULT_FAMILY, which: str = "deskforge"): | |
| """Run one model once. Returns (pyautogui action code, inference seconds).""" | |
| return _generate(family, which, instruction, screenshot.convert("RGB"), max_new_tokens) | |
| # a pair takes ~10 s; quota is checked against this | |
| def predict_duel(instruction: str, screenshot, max_new_tokens: int = 128, family: str = DEFAULT_FAMILY): | |
| """DeskForge and its base on the same screenshot in one GPU call: {"deskforge": (code, s), "base": (code, s)}.""" | |
| screenshot = screenshot.convert("RGB") | |
| return {which: _generate(family, which, instruction, screenshot, max_new_tokens) | |
| for which in ("deskforge", "base")} | |
| # -------------------------------------------------------------------------- | |
| # Parsing the model's pyautogui line into a structured action | |
| # -------------------------------------------------------------------------- | |
| POINT_KINDS = {"click", "doubleClick", "rightClick", "middleClick", "tripleClick", "moveTo", "dragTo"} | |
| KNOWN_KINDS = POINT_KINDS | {"scroll", "hscroll", "write", "press", "hotkey", "wait", "terminate"} | |
| class Action: | |
| kind: str # one of KNOWN_KINDS, or "unknown" | |
| raw: str | |
| x: float | None = None # fractional screen coordinates | |
| y: float | None = None | |
| amount: int = 0 # scroll clicks | |
| text: str = "" # write / terminate status | |
| keys: tuple = () # press / hotkey | |
| def has_point(self): | |
| return self.x is not None and self.y is not None | |
| def parse_action(code: str) -> Action: | |
| """Parse e.g. `pyautogui.click(x=0.1234, y=0.5678)` without executing it.""" | |
| # Base models often wrap the call in a code fence (```py ... ```); take the call itself. | |
| m = re.search(r"(?:pyautogui|computer)\.\w+\([^\n]*\)", code) | |
| line = m.group(0) if m else next((l.strip() for l in code.splitlines() if l.strip()), "") | |
| try: | |
| call = ast.parse(line, mode="eval").body | |
| if not (isinstance(call, ast.Call) and isinstance(call.func, ast.Attribute)): | |
| raise ValueError | |
| kind = call.func.attr | |
| args = [ast.literal_eval(a) for a in call.args] | |
| kw = {k.arg: ast.literal_eval(k.value) for k in call.keywords} | |
| except (SyntaxError, ValueError): | |
| return Action("unknown", line) | |
| if kind not in KNOWN_KINDS: | |
| return Action("unknown", line) | |
| a = Action(kind, line) | |
| if kind in POINT_KINDS: | |
| x = kw.get("x", args[0] if len(args) > 0 else None) | |
| y = kw.get("y", args[1] if len(args) > 1 else None) | |
| if isinstance(x, (int, float)) and isinstance(y, (int, float)): | |
| a.x, a.y = min(max(float(x), 0.0), 1.0), min(max(float(y), 0.0), 1.0) | |
| elif kind != "moveTo": | |
| return Action("unknown", line) | |
| elif kind in ("scroll", "hscroll"): | |
| a.amount = int(kw.get("clicks", args[0] if args else 0)) | |
| elif kind == "write": | |
| a.text = str(kw.get("message", args[0] if args else "")) | |
| elif kind in ("press", "hotkey"): | |
| keys = kw.get("keys", args[0] if len(args) == 1 else args) | |
| a.keys = tuple(keys) if isinstance(keys, (list, tuple)) else (keys,) | |
| a.keys = tuple(str(k) for k in a.keys) | |
| elif kind == "terminate": | |
| a.text = str(kw.get("status", args[0] if args else "")) | |
| return a | |
| # -------------------------------------------------------------------------- | |
| # Replaying an action with xdotool | |
| # -------------------------------------------------------------------------- | |
| KEYSYMS = { | |
| "enter": "Return", "return": "Return", "esc": "Escape", "escape": "Escape", | |
| "tab": "Tab", "backspace": "BackSpace", "delete": "Delete", "del": "Delete", | |
| "insert": "Insert", "space": "space", "up": "Up", "down": "Down", | |
| "left": "Left", "right": "Right", "home": "Home", "end": "End", | |
| "pageup": "Prior", "pgup": "Prior", "pagedown": "Next", "pgdn": "Next", | |
| "ctrl": "ctrl", "control": "ctrl", "alt": "alt", "shift": "shift", | |
| "win": "super", "super": "super", "cmd": "super", "command": "super", | |
| "+": "plus", "-": "minus", | |
| } | |
| BUTTONS = {"click": ("1", 1), "doubleClick": ("1", 2), "tripleClick": ("1", 3), | |
| "rightClick": ("3", 1), "middleClick": ("2", 1)} | |
| def keysym(key: str) -> str: | |
| k = key.strip().lower() | |
| if re.fullmatch(r"f\d{1,2}", k): | |
| return k.upper() | |
| return KEYSYMS.get(k, key.strip()) | |
| def to_pixels(action: Action, width: int, height: int): | |
| if not action.has_point: | |
| return None | |
| return min(round(action.x * width), width - 1), min(round(action.y * height), height - 1) | |
| def xdotool_commands(action: Action, width: int, height: int) -> list[list[str]]: | |
| """The xdotool invocations (argv lists) that perform `action` on a width x height screen.""" | |
| point = to_pixels(action, width, height) | |
| move = ["mousemove", str(point[0]), str(point[1])] if point else [] | |
| k = action.kind | |
| if k in BUTTONS: | |
| button, repeat = BUTTONS[k] | |
| return [["xdotool", *move, "click", "--repeat", str(repeat), "--delay", "80", button]] | |
| if k == "moveTo": | |
| return [["xdotool", *move]] if move else [] | |
| if k == "dragTo": | |
| return [["xdotool", "mousedown", "1", "sleep", "0.1", *move, "sleep", "0.1", "mouseup", "1"]] | |
| if k in ("scroll", "hscroll") and action.amount: | |
| if k == "scroll": | |
| button = "4" if action.amount > 0 else "5" | |
| else: | |
| button = "7" if action.amount > 0 else "6" | |
| return [["xdotool", "click", "--repeat", str(min(abs(action.amount), 25)), "--delay", "30", button]] | |
| if k == "write" and action.text: | |
| return [["xdotool", "type", "--delay", "25", "--", action.text[:500]]] | |
| if k == "press": | |
| return [["xdotool", "key", "--", keysym(key)] for key in action.keys] | |
| if k == "hotkey" and action.keys: | |
| return [["xdotool", "key", "--", "+".join(keysym(key) for key in action.keys)]] | |
| return [] | |
| # -------------------------------------------------------------------------- | |
| # Drawing | |
| # -------------------------------------------------------------------------- | |
| def draw_marker(image, point): | |
| """Draw a crosshair + circle at the predicted click point. | |
| A white underlay keeps the marker visible on both light and dark UIs. | |
| """ | |
| out = image.copy() | |
| d = ImageDraw.Draw(out) | |
| x, y = point | |
| r = max(10, min(out.width, out.height) // 60) | |
| d.ellipse([x - r, y - r, x + r, y + r], outline=(255, 255, 255), width=6) | |
| d.ellipse([x - r, y - r, x + r, y + r], outline=(255, 64, 0), width=3) | |
| for (x0, y0, x1, y1) in [ | |
| (x - r * 1.8, y, x + r * 1.8, y), # horizontal | |
| (x, y - r * 1.8, x, y + r * 1.8), # vertical | |
| ]: | |
| d.line([x0, y0, x1, y1], fill=(255, 255, 255), width=9) | |
| d.line([x0, y0, x1, y1], fill=(255, 64, 0), width=4) | |
| return out | |
| def annotate(screenshot, action: Action): | |
| point = to_pixels(action, screenshot.width, screenshot.height) | |
| return draw_marker(screenshot, point) if point else screenshot.copy() | |
| DESKFORGE_COLOR = (210, 96, 58) # terracotta, the UI accent | |
| BASE_COLOR = (59, 110, 190) # blue | |
| TARGET_COLOR = (74, 178, 64) # green | |
| def _font(size): | |
| try: | |
| return ImageFont.truetype("DejaVuSans-Bold.ttf", size) | |
| except OSError: | |
| return ImageFont.load_default() | |
| def draw_duel(screenshot, actions: dict, target_box=None): | |
| """Both models' predicted points (orange = DeskForge, blue = base) and the true target box.""" | |
| out = screenshot.copy() | |
| d = ImageDraw.Draw(out) | |
| font = _font(13) | |
| if target_box: | |
| x0, y0, x1, y1 = target_box | |
| d.rectangle([x0 - 2, y0 - 2, x1 + 2, y1 + 2], outline=(255, 255, 255), width=5) | |
| d.rectangle([x0 - 2, y0 - 2, x1 + 2, y1 + 2], outline=TARGET_COLOR, width=3) | |
| r = max(9, min(out.width, out.height) // 70) | |
| for which, color, label in (("base", BASE_COLOR, "Base"), ("deskforge", DESKFORGE_COLOR, "DeskForge")): | |
| action = actions.get(which) | |
| point = to_pixels(action, out.width, out.height) if action else None | |
| if not point: | |
| continue | |
| x, y = point | |
| d.ellipse([x - r, y - r, x + r, y + r], outline=(255, 255, 255), width=6) | |
| d.ellipse([x - r, y - r, x + r, y + r], outline=color, width=3) | |
| d.line([x - r * 1.7, y, x + r * 1.7, y], fill=(255, 255, 255), width=7) | |
| d.line([x, y - r * 1.7, x, y + r * 1.7], fill=(255, 255, 255), width=7) | |
| d.line([x - r * 1.7, y, x + r * 1.7, y], fill=color, width=3) | |
| d.line([x, y - r * 1.7, x, y + r * 1.7], fill=color, width=3) | |
| tw = d.textlength(label, font=font) | |
| lx = min(x + r + 4, out.width - tw - 8) | |
| ly = y - r - 18 if which == "deskforge" else y + r + 2 | |
| d.rounded_rectangle([lx, ly, lx + tw + 8, ly + 17], radius=4, fill=color) | |
| d.text((lx + 4, ly + 1), label, fill=(255, 255, 255), font=font) | |
| return out | |