1from PIL import Image
2from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
3from cursor_ccf import CCFConfig, ccf_predict_bbox, classify_instruction
4import torch
5
6# Load the base model exactly as you would normally
7model_id = "inclusionAI/GUI-G2-3B"
8processor = AutoProcessor.from_pretrained(
9 model_id, min_pixels=3136, max_pixels=12_845_056,
10)
11model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
12 model_id,
13 torch_dtype=torch.bfloat16,
14 attn_implementation="flash_attention_2",
15 device_map="auto",
16)
17model.eval()
18
19
20def predict_gui_g2(image, instruction):
21 """Single forward pass returning ((cx, cy), raw_text). The prompt
22 matches GUI-G2's training format exactly so output coords are in
23 the processor's smart_resize space; we rescale to the original image."""
24 from qwen_vl_utils import process_vision_info
25 import re
26
27 prompt = (
28 "Outline the position corresponding to the instruction: {}. "
29 "The output should be only [x1,y1,x2,y2]."
30 ).format(instruction)
31 messages = [{"role": "user", "content": [
32 {"type": "image", "image": image},
33 {"type": "text", "text": prompt},
34 ]}]
35 text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
36 image_inputs, _ = process_vision_info(messages)
37 inputs = processor(
38 text=[text], images=image_inputs, padding=True, return_tensors="pt",
39 ).to(model.device)
40 with torch.no_grad():
41 output = model.generate(**inputs, max_new_tokens=32, do_sample=False)
42 response = processor.batch_decode(
43 [output[0][inputs.input_ids.shape[1]:]],
44 skip_special_tokens=True,
45 )[0]
46 m = re.search(r"\[(\d+),\s*(\d+),\s*(\d+),\s*(\d+)\]", response)
47 if not m:
48 return (None, None), response
49 x1, y1, x2, y2 = map(int, m.groups())
50 abs_cx, abs_cy = (x1 + x2) / 2, (y1 + y2) / 2
51 # Rescale from processed-pixel space back to original-image pixels
52 proc_w = inputs["image_grid_thw"][0][2].item() * 14
53 proc_h = inputs["image_grid_thw"][0][1].item() * 14
54 orig_w, orig_h = image.size
55 return (abs_cx * orig_w / proc_w, abs_cy * orig_h / proc_h), response
56
57
58# Plug into CCF
59def predict_with_ccf(image, instruction, type_gate=True):
60 cfg = CCFConfig(
61 zoom_factor=2.0,
62 coarse_max_pixels=1_500_000,
63 instruction_classifier_fn=classify_instruction if type_gate else None,
64 )
65 def inner(img, instr):
66 (x, y), raw = predict_gui_g2(img, instr)
67 return (x, y) if x is not None else None, raw
68 result = ccf_predict_bbox(inner, image, instruction, cfg)
69 if result is None:
70 return None
71 return (result.x, result.y), result.stage
72
73
74# Run it
75image = Image.open("screenshot.png").convert("RGB")
76(x, y), stage = predict_with_ccf(image, "click the settings icon")
77print(f"Click at ({x:.0f}, {y:.0f}) [stage={stage}]")
