adapters/molmoact2.py

"""MolmoAct2 adapter — a faithful port of the proven single-model server
(cloud_inference/space/molmoact2_server.py): same predict_action call, same
RTC per-session guidance cache, same OOM->507 behavior. Semantics UNCHANGED;
only the HTTP layer moved into serve/server.py."""

from __future__ import annotations

import os
import re
import time
from typing import Optional

import numpy as np
from fastapi import HTTPException

from adapters.base import resolve_source

REPO_ID = os.environ.get("MOLMOACT_REPO", "allenai/MolmoAct2-SO100_101")
NORM_TAG = os.environ.get("MOLMOACT_NORM_TAG", "so100_so101_molmoact2")
ACTION_HORIZON = 30          # this checkpoint's chunk length (30 moves @ 30 Hz)
CHUNK_HZ = 30.0
MAX_NUM_STEPS = 50
MAX_IMAGES = 6
_RTC_SESSION_CAP = 8


class MolmoAct2Adapter:
    def __init__(self, model_path: str):
        self._probe_path = model_path
        self._source: Optional[str] = None
        self._model = None
        self._processor = None
        self._rtc_state = None
        self._rtc_prev: dict = {}    # session id -> previous normalized chunk
        self._dtype = None

    def load(self) -> None:
        import torch
        from transformers import AutoModelForImageTextToText, AutoProcessor

        import rtc as rtcmod

        self._dtype = (torch.bfloat16
                       if os.environ.get("MOLMOACT_BF16", "1") == "1" else torch.float32)
        self._source = resolve_source(self._probe_path, REPO_ID)
        print(f"[molmoact2] loading from {self._source}", flush=True)
        proc = AutoProcessor.from_pretrained(self._source, trust_remote_code=True)
        model = (AutoModelForImageTextToText
                 .from_pretrained(self._source, trust_remote_code=True, dtype=self._dtype)
                 .to("cuda").eval())
        self._rtc_state = rtcmod.RTCState()
        try:
            if rtcmod.install_rtc(model, self._rtc_state) is None:
                print("[molmoact2] RTC: flow loop not found — serving un-guided", flush=True)
        except Exception as exc:
            print(f"[molmoact2] RTC install failed ({exc}) — serving un-guided", flush=True)
        self._processor, self._model = proc, model

    def meta(self) -> dict:
        return {"kind": "molmoact2", "chunk_hz": CHUNK_HZ, "horizon": ACTION_HORIZON,
                "dof": 6, "cameras": ["wrist", "third_person"],
                "max_images": MAX_IMAGES, "source": self._source,
                "norm_tag": NORM_TAG, "supports_point": True, "supports_rtc": True}

    def act(self, *, images, state, instruction, num_steps, extras):
        import torch
        import rtc as rtcmod

        num_steps = max(1, min(int(num_steps), MAX_NUM_STEPS))
        rtc_req = extras.get("rtc")
        rtc_horizon = None
        rtc_note = None
        if rtc_req is not None:
            rtc_horizon = rtcmod.pick_execution_horizon(
                int(rtc_req.get("delay", 0)), ACTION_HORIZON)
            if rtc_horizon is None:
                rtc_note = (f"skipped: delay {rtc_req.get('delay')} too large for "
                            f"horizon {ACTION_HORIZON} (needs d <= H/2)")
        use_rtc = rtc_req is not None and rtc_horizon is not None
        cuda_graph = False if use_rtc else bool(extras.get("enable_cuda_graph", True))
        grad_ctx = (torch.enable_grad()
                    if (use_rtc or extras.get("enable_grad")) else torch.no_grad())
        st = self._rtc_state
        try:
            with grad_ctx:
                if use_rtc:
                    session = str(rtc_req.get("session"))
                    st.prev = self._rtc_prev.get(session)
                    st.enabled = st.prev is not None
                    st.consumed = max(0, int(rtc_req.get("consumed", 0)))
                    st.delay = max(0, int(rtc_req.get("delay", 0)))
                    st.execution_horizon = rtc_horizon
                    st.applied = 0
                out = self._model.predict_action(
                    processor=self._processor, images=list(images), task=instruction,
                    state=state, norm_tag=NORM_TAG,
                    inference_action_mode="continuous", num_steps=num_steps,
                    normalize_language=True, enable_cuda_graph=cuda_graph)
                if use_rtc:
                    if (len(self._rtc_prev) >= _RTC_SESSION_CAP
                            and session not in self._rtc_prev):
                        self._rtc_prev.pop(next(iter(self._rtc_prev)))
                    self._rtc_prev[session] = st.prev
                    rtc_note = {"guided_steps": st.applied, "execution_horizon": rtc_horizon,
                                "delay": st.delay, "consumed": st.consumed,
                                "had_target": bool(st.enabled)}
                    st.prev = None
                    st.enabled = False
        except torch.cuda.OutOfMemoryError as e:
            torch.cuda.empty_cache()
            raise HTTPException(status_code=507, detail=f"CUDA OOM: {e}") from e
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        acts = out.actions
        if torch.is_tensor(acts):
            acts = acts.detach().float().cpu().numpy()
        return np.asarray(acts, dtype=np.float32), {
            "rtc": (rtc_note if isinstance(rtc_note, dict)
                    else ({"skipped": rtc_note} if rtc_note else None))}

    def point(self, *, image, query, max_new_tokens) -> dict:
        import torch
        from PIL import Image

        img = Image.fromarray(image)
        prompt = f"Point to {query}."
        t0 = time.time()
        try:
            with torch.inference_mode():
                try:
                    inputs = self._processor.apply_chat_template(
                        [{"role": "user",
                          "content": [{"type": "image", "image": img},
                                      {"type": "text", "text": prompt}]}],
                        add_generation_prompt=True, tokenize=True,
                        return_dict=True, return_tensors="pt")
                except Exception:
                    inputs = self._processor.process(images=[img], text=prompt)
                    inputs = {k: (v.unsqueeze(0) if hasattr(v, "dim") and v.dim() in (1, 3) else v)
                              for k, v in inputs.items()}
                inputs = {k: (v.to(self._model.device) if hasattr(v, "to") else v)
                          for k, v in inputs.items()}
                out = self._model.generate(**inputs, max_new_tokens=int(max_new_tokens))
                n_in = inputs["input_ids"].shape[1] if "input_ids" in inputs else 0
                text = self._processor.tokenizer.decode(out[0][n_in:], skip_special_tokens=False)
        except Exception as exc:
            raise HTTPException(status_code=500,
                                detail=f"pointing failed: {type(exc).__name__}: {exc}")
        pts = [[float(x), float(y)] for x, y in
               re.findall(r'x\d*="([0-9.]+)"\s+y\d*="([0-9.]+)"', text)]
        return {"raw": text, "points": pts,
                "compute_ms": round((time.time() - t0) * 1000.0, 1)}