inference.py

"""
Inference with fine-tuned Qwen3-VL-2B affordance model.
Tests the LoRA adapter on sample images.
"""

import os
import json
import torch
from PIL import Image
from unsloth import FastVisionModel

# ============================================================
# CONFIG
# ============================================================
ADAPTER_PATH = "c:/Fidelity_dynamics/qwen_affordance_exp/output/lora_adapter"
DATASET_DIR = "c:/Fidelity_dynamics/qwen_affordance_exp/dataset"
VAL_FILE = os.path.join(DATASET_DIR, "val_msswift.jsonl")

# ============================================================
# LOAD MODEL + ADAPTER
# ============================================================
print("Loading model with LoRA adapter...")

model, tokenizer = FastVisionModel.from_pretrained(
    ADAPTER_PATH,
    load_in_4bit=True,
    max_seq_length=2048,
)

FastVisionModel.for_inference(model)
print("Model loaded!\n")

# ============================================================
# RUN INFERENCE ON VALIDATION SAMPLES
# ============================================================
# Load a few val samples
val_data = []
with open(VAL_FILE, "r", encoding="utf-8") as f:
    for i, line in enumerate(f):
        if i >= 5:  # Test on first 5
            break
        val_data.append(json.loads(line.strip()))

print("=" * 60)
print("INFERENCE RESULTS")
print("=" * 60)

for i, sample in enumerate(val_data):
    img_path = sample["images"][0]
    ground_truth = sample["messages"][2]["content"]  # assistant response

    # Load image
    image = Image.open(img_path).convert("RGB")
    image = image.resize((384, 216), Image.LANCZOS)

    # Build prompt
    messages = [
        {
            "role": "system",
            "content": [{"type": "text", "text": "You are a visual affordance detection assistant. Given a first-person view image, identify what actions can be performed on the visible objects."}],
        },
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": "What affordances are present in this scene? Describe what actions can be performed on the visible objects."},
            ],
        },
    ]

    # Tokenize
    input_text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(
        input_text,
        return_tensors="pt",
        add_special_tokens=False,
    ).to(model.device)

    # Generate
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=256,
            temperature=0.7,
            do_sample=True,
        )

    # Decode only new tokens
    new_tokens = output_ids[0][inputs["input_ids"].shape[1]:]
    response = tokenizer.decode(new_tokens, skip_special_tokens=True)

    print(f"\n--- Sample {i+1} ---")
    print(f"Image: {os.path.basename(img_path)}")
    print(f"Ground Truth: {ground_truth[:150]}...")
    print(f"Prediction:   {response[:150]}...")
    print()

print("=" * 60)
print("Inference complete!")