"""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) @dataclass(frozen=True) 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 @spaces.GPU(duration=30, size="xlarge") # 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) @spaces.GPU(duration=30, size="xlarge") # 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"} @dataclass 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 @property 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