scripts/_vlm_material_worker.py

#!/usr/bin/env python
"""Classify object-crop material with Qwen3-VL-2B-Instruct, batched, calibrated.

Runs inside the isolated ``vla457`` conda env (transformers 4.57.0) -- NEVER
the ``fpgm`` env, which must not have ``transformers`` installed at all (its
numpy<2 / torch stack is load-bearing for SAM3/TAPNext/pyrender; see
``fpgm.physics.types``'s module docstring for the same isolation argument
applied to MuJoCo). Exactly the same subprocess-and-JSON pattern already
established by ``scripts/_mujoco_settle_worker.py``: a small, dependency-light
worker invoked as a subprocess from the ``fpgm``-side orchestrator
(``fpgm.physics.priors.VlmPriorProposer``) with the isolated env's own
interpreter, talking JSON files rather than in-process imports. This script
deliberately imports NOTHING from ``fpgm`` -- the ``vla457`` env does not have
it installed, and mixing envs defeats the whole point of the isolation.

--- Why one process classifies several objects, not one process per object ---
Model load (weights off disk, onto the GPU, ~2B params) dominates wall clock
for anything less than dozens of images -- a few seconds of inference work
behind many more seconds of process/import/load overhead if paid per object.
So the CLI takes a *list* of ``--images``/``--labels`` and classifies all of
them in one process, one model load. The caller
(``fpgm.physics.priors.VlmPriorProposer``) is expected to batch every object
crop it has cached-miss on into a single worker invocation rather than
shelling out per object.

--- Why logits over candidate tokens, not "please output a probability" ------
Instruction-tuned VLMs asked to "output calibrated probabilities" produce
confident-sounding fabrications -- there is no mechanism that makes the
numbers they type out correspond to anything the model's own uncertainty
actually reflects. What the model's forward pass *does* genuinely encode is
its next-token distribution. So this worker uses the standard multiple-choice
technique (the same one MMLU-style evals use for calibration): present the
material classes as a lettered, closed-set list, ask for a single-letter
answer, and read the softmax over just the candidate letter tokens' logits at
the answer position -- never generation, never a model-typed number. This is
a genuine model-internal quantity (the model's own belief about which
completion is most likely) rather than a post-hoc verbalisation of one, and
its whole justification lives in how it is computed: a probability read off
logits cannot lie about itself in the way a typed-out "87%" can, because
there is no separate "assess and describe your confidence" step for the model
to bluff at -- it's the literal probability of the token.

**Honest limitation**: this is a calibration technique, not a calibration
guarantee. Next-token probabilities from an instruction-tuned model are known
to be *overconfident* on many benchmarks (RLHF sharpens the distribution).
Nothing here corrects that; it is the reason ``fpgm.physics.materials``
imposes its own wide floors (moment-matched mixture variance from the
material table's own gsd) rather than trusting the VLM's spread directly to
set prior width on its own -- the mixture only ever *widens* what the letter
logits report, never narrows it. A raw multi-token-word-likelihood approach
was considered (score each full material *word* via teacher-forcing instead
of a letter) and rejected: it conflates "the model thinks this word is
likely" with "this word is a common/short token sequence" (e.g. "wood" is
one BPE token, "cardboard" may be several), which is exactly the kind of
confound the letter-option framing sidesteps by construction (every option
is a single character, A-J).

Usage:
    /home/quang/miniconda3/envs/vla457/bin/python scripts/_vlm_material_worker.py \\
        --images crop_brick.png crop_books.png \\
        --labels brick books \\
        --out verdicts.json
"""

from __future__ import annotations

import argparse
import json
import string
import sys
import time
from pathlib import Path
from typing import Any

import torch
from PIL import Image
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration

MODEL_ID = "Qwen/Qwen3-VL-2B-Instruct"

#: MUST mirror ``fpgm.physics.materials.CLASS_NAMES`` minus ``"unknown"``
#: exactly (kept out of the VLM's own answer set on purpose -- see below).
#: Cannot be imported from that module: this env has no ``fpgm`` install.
CLASS_NAMES: tuple[str, ...] = (
    "wood",
    "plastic",
    "cardboard",
    "metal",
    "glass",
    "ceramic",
    "rubber",
    "foam",
    "fabric",
    "stone",
)

#: "unknown" is deliberately never offered as a VLM answer choice. It is a
#: pure fallback bucket for "no VLM read was possible at all" (see
#: ``fpgm.physics.priors.VlmPriorProposer.fallback_verdict``); offering it
#: to a model looking at a real photograph would only dilute probability
#: mass away from the ten classes that actually carry physical information,
#: with no corresponding gain -- a model uncertain between two materials
#: should spread mass across THOSE, not opt out.
_LETTERS = string.ascii_uppercase[: len(CLASS_NAMES)]

_PROMPT_TEMPLATE = """You are labelling the surface material of a single rigid or \
semi-rigid object, shown cropped and isolated from its background (transparent = not-object).

Object label (for context only, not a hint about material): "{label}"

Choose the single option below that best describes what the object is primarily made \
of. Options:
{options}

Respond with only the letter of your choice, nothing else."""


def _build_prompt(label: str) -> str:
    options = "\n".join(
        f"{letter}) {name}" for letter, name in zip(_LETTERS, CLASS_NAMES, strict=True)
    )
    return _PROMPT_TEMPLATE.format(label=label, options=options)


def _letter_token_ids(tokenizer: Any, letter: str) -> list[int]:
    """Every single-token encoding of ``letter`` worth checking (bare and space-prefixed).

    Verified empirically for this tokenizer that both ``"A"`` and ``" A"``
    encode to exactly one token each (and differ), so both are read and their
    probabilities summed -- whichever the model's own next-token position
    naturally prefers (immediately after ``"...\\n\\n"``, bare is typical, but
    reading only one variant risks silently discarding real mass the model
    put on the other).
    """
    ids: list[int] = []
    for candidate in (letter, f" {letter}"):
        enc = tokenizer.encode(candidate, add_special_tokens=False)
        if len(enc) == 1 and enc[0] not in ids:
            ids.append(enc[0])
    if not ids:
        raise RuntimeError(f"tokenizer has no single-token encoding for letter {letter!r}")
    return ids


def _select_device() -> str:
    if not torch.cuda.is_available():
        return "cpu"
    best_idx, best_free = None, -1
    for i in range(torch.cuda.device_count()):
        free, _total = torch.cuda.mem_get_info(i)
        if free > best_free:
            best_idx, best_free = i, free
    return f"cuda:{best_idx}" if best_idx is not None else "cpu"


def classify_batch(
    image_paths: list[Path], labels: list[str], device: str | None = None
) -> tuple[list[dict[str, Any]], dict[str, float]]:
    """Classify each ``(image, label)`` pair; returns (verdicts, timing)."""
    timing: dict[str, Any] = {}
    t_load0 = time.perf_counter()

    device = device or _select_device()
    dtype = torch.bfloat16 if device != "cpu" else torch.float32
    # use_fast=False: the fast (torch-based) image processor in this
    # transformers release calls torch.compiler.is_compiling(), added only in
    # torch>=2.3; the vla457 env pins torch 2.2.0 (transformers 4.57.0 is the
    # env's whole reason to exist, and its own torch pin is untouched here).
    # The slow PIL/numpy-based processor has no such dependency and produces
    # the same preprocessing, just without the torch-native fast path.
    processor = AutoProcessor.from_pretrained(MODEL_ID, use_fast=False)
    model = Qwen3VLForConditionalGeneration.from_pretrained(MODEL_ID, dtype=dtype)
    model = model.to(device)
    model.eval()

    letter_token_ids = {
        letter: _letter_token_ids(processor.tokenizer, letter) for letter in _LETTERS
    }
    timing["model_load_seconds"] = time.perf_counter() - t_load0
    timing["device"] = device  # type: ignore[assignment]
    timing["dtype"] = str(dtype)  # type: ignore[assignment]

    verdicts: list[dict[str, Any]] = []
    per_image_seconds: list[float] = []

    for image_path, label in zip(image_paths, labels, strict=True):
        t0 = time.perf_counter()
        image = Image.open(image_path).convert("RGBA")
        # Composite onto a neutral grey so a fully-transparent background
        # (alpha=0, the crop convention -- see fpgm.objects.crop) renders as
        # flat grey rather than the arbitrary colour PIL leaves in the RGB
        # channels under transparency; a flat neutral background keeps the
        # material read about the object, not a background-colour artefact.
        bg = Image.new("RGB", image.size, (128, 128, 128))
        bg.paste(image, mask=image.split()[3])
        image_rgb = bg

        prompt = _build_prompt(label)
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image_rgb},
                    {"type": "text", "text": prompt},
                ],
            }
        ]
        chat_text = processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        inputs = processor(text=[chat_text], images=[image_rgb], return_tensors="pt").to(device)

        with torch.no_grad():
            out = model(**inputs)
        next_token_logits = out.logits[0, -1, :].float()

        raw_letter_logits: dict[str, float] = {}
        for letter in _LETTERS:
            ids = letter_token_ids[letter]
            raw_letter_logits[letter] = float(torch.logsumexp(next_token_logits[ids], dim=0))

        logit_tensor = torch.tensor([raw_letter_logits[letter] for letter in _LETTERS])
        probs = torch.softmax(logit_tensor, dim=0).tolist()
        classes = [[CLASS_NAMES[i], float(probs[i])] for i in range(len(CLASS_NAMES))]

        # Short free-generation purely as a human-auditable trail -- NOT used
        # to derive any probability (see module docstring on why generation
        # itself is not trusted for calibration).
        with torch.no_grad():
            gen_ids = model.generate(
                **inputs,
                max_new_tokens=24,
                do_sample=False,
                pad_token_id=processor.tokenizer.eos_token_id,
            )
        gen_text = processor.tokenizer.decode(
            gen_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True
        ).strip()

        elapsed = time.perf_counter() - t0
        per_image_seconds.append(elapsed)

        best_letter = max(_LETTERS, key=lambda ltr: raw_letter_logits[ltr])
        verdicts.append(
            {
                "label": label,
                "classes": classes,
                "source": f"vlm:{MODEL_ID}",
                "raw": {
                    "prompt": prompt,
                    "letter_logits": raw_letter_logits,
                    "letter_probs": {letter: probs[i] for i, letter in enumerate(_LETTERS)},
                    "argmax_letter": best_letter,
                    "argmax_class": CLASS_NAMES[_LETTERS.index(best_letter)],
                    "generated_text": gen_text,
                    "image_path": str(image_path),
                    "inference_seconds": elapsed,
                },
            }
        )

    timing["n_images"] = float(len(image_paths))
    timing["total_inference_seconds"] = sum(per_image_seconds)
    timing["mean_inference_seconds"] = (
        sum(per_image_seconds) / len(per_image_seconds) if per_image_seconds else 0.0
    )
    return verdicts, timing


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument("--images", type=Path, nargs="+", required=True)
    p.add_argument("--labels", type=str, nargs="+", required=True)
    p.add_argument("--out", type=Path, required=True)
    p.add_argument(
        "--device", type=str, default=None, help="e.g. 'cuda:0' or 'cpu'; default: auto-select"
    )
    return p.parse_args()


def main() -> int:
    args = parse_args()
    if len(args.images) != len(args.labels):
        print(
            f"error: --images has {len(args.images)} entries but --labels has "
            f"{len(args.labels)}",
            file=sys.stderr,
        )
        return 2

    t_total0 = time.perf_counter()
    verdicts, timing = classify_batch(args.images, args.labels, device=args.device)
    timing["total_wall_seconds"] = time.perf_counter() - t_total0

    args.out.write_text(json.dumps(verdicts, indent=2))
    print(
        f"vlm material worker: {len(verdicts)} object(s) classified on "
        f"{timing['device']} ({timing['dtype']}), model load "
        f"{timing['model_load_seconds']:.2f}s, mean inference "
        f"{timing['mean_inference_seconds']:.2f}s/image, total "
        f"{timing['total_wall_seconds']:.2f}s -> {args.out}"
    )
    print(json.dumps({k: v for k, v in timing.items() if k not in ("device", "dtype")}, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())