code/docs/inference.md

Inference

Inference defaults to the complete PRIMA checkpoint, PLAN-Lab/PRIMA. The package includes merged model weights, both tokenizers, and the vision and grounding components. Set merged_ckpt_path to a downloaded model directory to load a local copy, or use hub_revision to pin a Hub commit.

Q-Former weights, its configuration, and its tokenizer are already included. No separate Q-Former checkpoint URL or BERT download is needed when loading the merged model.

image_paths is a list of image groups; prompts contains one prompt for each group. The model processes images in their listed order and numbers them from one. Each group starts a separate conversation. The default interface accepts one to five images per group; supported performance beyond the training distribution must be evaluated separately.

merged_ckpt_path: PLAN-Lab/PRIMA
image_paths:
  - - /path/to/first_image.jpg
    - /path/to/second_image.jpg
prompts:
  - Identify and segment the objects shared by these images.

Run inference with python run.py --config run_config.yaml. Command-line options override configuration values.

You can also load PRIMA directly with Transformers, using the inference code included with the checkpoint:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BertTokenizer

model_id = "PLAN-Lab/PRIMA"
tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
qformer_tokenizer = BertTokenizer.from_pretrained(
    model_id, subfolder="qformer_tokenizer", truncation_side="left"
)
model = AutoModelForCausalLM.from_pretrained(
    model_id, trust_remote_code=True,
    torch_dtype=torch.bfloat16, low_cpu_mem_usage=True,
).to(device="cuda", dtype=torch.bfloat16).eval()
result = model.predict(
    images=["/path/to/first_image.jpg", "/path/to/second_image.jpg"],
    question="Identify and segment the objects shared by these images.",
    tokenizer=tokenizer, qformer_tokenizer=qformer_tokenizer,
)
print(result["response"])

See the model card for the checkpoint contents and output format.

For initialization or older trainable-state exports, set merged_ckpt_path: "" and configure version, ckpt_path, eva_weights, q_former_model, and qformer_bert_path. A trainable-state export from further fine-tuning of the merged release can be supplied as ckpt_path alongside merged_ckpt_path; its LoRA branches are reconstructed before loading the export.

Each result directory contains a JSON record with the ordered images, input prompt, generated response, and mask paths. Binary masks are NumPy arrays in the original image resolution. Within each image, mask order follows the generated grounding routes. A response may have no grounding for a particular image; its mask list is empty. Ungrounded or malformed generations must not be interpreted as successful localization.