slotread/local.py

"""Reading slots in-process with transformers: no server. One forward pass computes every slot of a layout
(only the logits at the slot positions, logits_to_keep), exactly as training reads them. Used for the merged
models and the Hugging Face Space (ZeroGPU runs PyTorch, not vLLM).

  from slotread.local import LocalSlotModel
  m = LocalSlotModel.from_pretrained("infinitylogesh/GUI-Decisions-31B")          # merged model
  m = LocalSlotModel.from_pretrained("google/gemma-4-31B-it", adapter="infinitylogesh/GUI-Decisions-31B-LoRA")
  pred = m.decide(Image.open("shot.png"), "open the settings")    # {'action': 'click', 'x': ..., 'y': ...}

The slot template and label ids come from slotread_config.json in the model (or adapter) repo: the
config.json that slotread.train writes, renamed so it does not collide with the model's own config.json.
"""

import json, time

import torch

from . import labels as L
from .client import agent_user, load_config
from .images import cap_image
from .labels import Arm
from .tasks import AgentTask, BoxTask, SegTask

class LocalSlotModel:
    def __init__(self, model, processor, cfg):
        self.model, self.proc, self.cfg = model, processor, cfg
        self.arm = Arm(cfg["arm"], cfg["coarse"], cfg["bins"], cfg["tokens"], cfg["fine"])
        self.tasks = {"agent": AgentTask(), "box": BoxTask(), "seg": SegTask(cfg.get("rays", 24))}
        L.PH = cfg["placeholder"]  # the placeholder id of this model's tokenizer

    @classmethod
    def from_pretrained(cls, model_id, adapter=None, processor=None, config=None, dtype=torch.bfloat16, device_map=None):
        """config: a slotread config file; default: slotread_config.json in the adapter or model repo."""
        from transformers import AutoModelForImageTextToText, AutoProcessor
        model = AutoModelForImageTextToText.from_pretrained(model_id, dtype=dtype, device_map=device_map or {"": 0},
                                                            attn_implementation="sdpa")
        if adapter:
            from peft import PeftModel
            # on the model's device: under ZeroGPU CUDA looks available before a GPU is attached
            kw = {"torch_device": "cpu"} if device_map == "cpu" else {}
            model = PeftModel.from_pretrained(model, adapter, **kw)
        model.eval()
        cfg = json.load(open(config)) if config else load_config(adapter or model_id)
        return cls(model, AutoProcessor.from_pretrained(processor or model_id), cfg)

    @torch.no_grad()
    def read(self, task, image, user):
        """{slot: log-probabilities over its labels} for the task's whole layout, one forward pass;
        also the forward's milliseconds."""
        t = self.tasks[task]
        c = self.cfg["tasks"][task]
        image = cap_image(image.convert("RGB"), self.cfg.get("image_cap") or 0)  # as trained (GUI-Owl: 1 MP)
        ex = {"agent": dict(image=image, user=user), "box": dict(image=image, instruction=user),
              "seg": dict(image=image, sentence=user)}[task]
        msgs = [{"role": "system", "content": [{"type": "text", "text": c["system"]}]},
                {"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": t.user(ex)}]},
                {"role": "assistant", "content": [{"type": "text", "text": " ".join(f"{s}: ___" for s in c["slots"])}]}]
        enc = self.proc.apply_chat_template(msgs, tokenize=True, return_dict=True, return_tensors="pt",
                                           add_generation_prompt=False, continue_final_message=True, enable_thinking=False)
        ids = enc["input_ids"][0]
        keep = (ids == self.cfg["placeholder"]).nonzero().flatten()[-len(c["slots"]):] - 1
        dev = next(self.model.parameters()).device
        if dev.type == "cuda":
            torch.cuda.synchronize()
        t0 = time.perf_counter()
        out = self.model(**{k: v.to(dev) for k, v in enc.items()}, logits_to_keep=keep.to(dev), use_cache=False)
        lv = torch.log_softmax(out.logits[0].float(), -1)
        if dev.type == "cuda":
            torch.cuda.synchronize()
        ms = (time.perf_counter() - t0) * 1000
        return {s: lv[i, c["labels"][s]].cpu() for i, s in enumerate(c["slots"])}, ms

    def decide(self, image, instruction, previous="None"):
        """A computer-use step from one forward pass: action and point, plus the swipe end, key and scroll
        slots (meaningful only for their actions); `ms` is the forward pass."""
        from .data.aguvis import KEYS, SDIR
        lp, ms = self.read("agent", image, agent_user(instruction, previous))
        d = self.tasks["agent"].decode(lp, self.arm)
        d.update(key=KEYS[int(lp["key"].argmax())], sdir=SDIR[int(lp["sdir"].argmax())], samt=int(lp["samt"].argmax()))
        return dict(d, ms=round(ms, 1))

    def box(self, image, description):
        lp, ms = self.read("box", image, description)
        return dict(box=self.tasks["box"].decode(lp, self.arm), ms=round(ms, 1))

    @torch.no_grad()
    def write_text(self, image, instruction, action, previous="None", max_new_tokens=24):
        """The text a type / answer / open step needs, written by the base model: the adapter is switched
        off, so this needs base + adapter (a merged model cannot switch it off)."""
        from .agent import TEXT_ASK
        if not hasattr(self.model, "disable_adapter"):
            raise RuntimeError("text steps need the base model: load with adapter=...")
        image = cap_image(image.convert("RGB"), self.cfg.get("image_cap") or 0)
        msgs = [{"role": "user", "content": [{"type": "image", "image": image},
                                             {"type": "text", "text": agent_user(instruction, previous) + "\n\n" + TEXT_ASK[action]}]}]
        enc = self.proc.apply_chat_template(msgs, tokenize=True, return_dict=True, return_tensors="pt",
                                           add_generation_prompt=True, enable_thinking=False)
        dev = next(self.model.parameters()).device
        enc = {k: v.to(dev) for k, v in enc.items()}
        with self.model.disable_adapter():
            out = self.model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False)
        t = self.proc.decode(out[0, enc["input_ids"].shape[1]:], skip_special_tokens=True)
        return t.strip().strip('"').split("\n")[0][:120]

    @torch.no_grad()
    def vanilla_decide(self, image, instruction, previous="None", max_new_tokens=96):
        """The base model (adapter off) writing the step as JSON in Gemma 4's native [y, x] format, as in
        the post's comparisons: (action, point, args, text, raw output, ms for the whole generation)."""
        from .agent import json_system, parse_step
        if not hasattr(self.model, "disable_adapter"):
            raise RuntimeError("the vanilla comparison needs the base model: load with adapter=...")
        image = cap_image(image.convert("RGB"), self.cfg.get("image_cap") or 0)
        msgs = [{"role": "system", "content": [{"type": "text", "text": json_system(True)}]},
                {"role": "user", "content": [{"type": "image", "image": image},
                                             {"type": "text", "text": agent_user(instruction, previous)}]}]
        enc = self.proc.apply_chat_template(msgs, tokenize=True, return_dict=True, return_tensors="pt",
                                           add_generation_prompt=True, enable_thinking=False)
        dev = next(self.model.parameters()).device
        enc = {k: v.to(dev) for k, v in enc.items()}
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        with self.model.disable_adapter():
            out = self.model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False)
        torch.cuda.synchronize()
        ms = (time.perf_counter() - t0) * 1000
        txt = self.proc.decode(out[0, enc["input_ids"].shape[1]:], skip_special_tokens=True)
        return parse_step(txt, True) + (txt, round(ms, 1))