1import os
2import sys
3import torch
4import numpy as np
5import json
6import hydra
7from hydra.core.global_hydra import GlobalHydra
8from PIL import Image
9
10# Add parent directory to path for src imports
11sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
12
13from src.arguments import ModelArguments, DataArguments, TrainingArguments
14from src.model.model import MMEBModel
15from src.model.processor import load_processor, VLM_IMAGE_TOKENS, get_backbone_name, process_vlm_inputs_fns
16from transformers import AutoConfig
17
18
19# Initialize Hydra for SAM2 loading
20if not GlobalHydra().is_initialized():
21 hydra.initialize(config_path="./configs", version_base=None)
22
23# Determinism
24torch.manual_seed(42)
25torch.cuda.manual_seed_all(42)
26torch.backends.cudnn.deterministic = True
27torch.backends.cudnn.benchmark = False
28np.random.seed(42)
29
30model_dir = 'Sony/VIRTUE-2B-SCaR'
31device = 'cuda' if torch.cuda.is_available() else 'cpu'
32
33config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True, token=True)
34
35# Build arguments directly (no YAML required)
36model_args = ModelArguments(
37 model_name=model_dir,
38 checkpoint_path=None,
39 pooling="last",
40 normalize=True,
41 lora=False,
42 model_backbone='qwen2_vl',
43)
44persisted_sam = config.virtue_sam
45
46model_args.sam = True
47model_args.sam_config = {
48 "config_path": persisted_sam.get('config_path') if persisted_sam else None,
49 "checkpoint": persisted_sam.get('checkpoint') if persisted_sam else None,
50 "points_per_side": (persisted_sam.get('points_per_side') if persisted_sam else 16),
51 "feature_levels": (persisted_sam.get('feature_levels') if persisted_sam else 3),
52}
53
54data_args = DataArguments()
55training_args = TrainingArguments()
56
57processor = load_processor(model_args, data_args)
58model = MMEBModel.load(model_args, is_trainable=False, processor=processor)
59model.eval()
60model = model.to(device, dtype=torch.bfloat16)
61
62# Get model backbone and image token
63model_backbone = get_backbone_name(hf_config=config)
64image_token = VLM_IMAGE_TOKENS[model_backbone]
65
66# Image + Text -> Text
67image_path = '../assets/example.jpg'
68image = Image.open(image_path).convert('RGB')
69
70model_inputs = {
71 'text': [f"{image_token}\nRepresent the given image with the following question: What is in the image"],
72 'images': [image]
73}
74
75process_fn = process_vlm_inputs_fns[model_backbone]
76inputs = process_fn(model_inputs, processor=processor, max_length=512)
77device = next(model.parameters()).device
78inputs = {k: v.to(device) if torch.is_tensor(v) else v for k, v in inputs.items()}
79
80with torch.no_grad():
81 with torch.autocast(enabled=True, dtype=torch.bfloat16, device_type="cuda"):
82 qry_output = model(qry=inputs)["qry_reps"]
83
84# Candidates for all scenarios
85test_strings = ['A cat', 'A dog', 'A tiger']
86
87# Scenario 1: No visual prompts (image only)
88print("\n--- Similarities (no visual prompts) ---")
89for string in test_strings:
90 cand_inputs = process_fn({'text': [string], 'images': [None]}, processor=processor)
91 cand_inputs = {k: v.to(device) if torch.is_tensor(v) else v for k, v in cand_inputs.items()}
92 with torch.no_grad():
93 with torch.autocast(enabled=True, dtype=torch.bfloat16, device_type="cuda"):
94 tgt_output = model(tgt=cand_inputs)["tgt_reps"]
95 sim = model.compute_similarity(qry_output, tgt_output)
96 print(f"no-prompt | {string} = {sim}")
97
98'''
99--- Similarities (no visual prompts) ---
100no-prompt | A cat = tensor([[0.3030]], device='cuda:0')
101no-prompt | A dog = tensor([[0.2453]], device='cuda:0')
102no-prompt | A tiger = tensor([[0.1714]], device='cuda:0')
103'''
104
105# Scenario 2: Point prompts — two examples (left/right)
106print("\n--- Similarities (point prompts) ---")
107sam_size = 1024 # SAM2Transforms output size
108point_examples = [(0.25, 0.5), (0.75, 0.5)]
109for (px, py) in point_examples:
110 point_text = f"{image_token}\nFind the caption that best describes the segmented object, considering both local details and global context in the given image.\nReferring object point: ({int(px*image.size[0])}, {int(py*image.size[1])})"
111 q_inputs = process_fn({'text': [point_text], 'images': [image]}, processor=processor)
112 q_inputs['point'] = [px * sam_size, py * sam_size]
113 q_inputs = {k: v.to(device) if torch.is_tensor(v) else v for k, v in q_inputs.items()}
114 with torch.no_grad():
115 with torch.autocast(enabled=True, dtype=torch.bfloat16, device_type="cuda"):
116 point_qry = model(qry=q_inputs)["qry_reps"]
117 for string in test_strings:
118 cand_inputs = process_fn({'text': [string], 'images': [None]}, processor=processor)
119 cand_inputs = {k: v.to(device) if torch.is_tensor(v) else v for k, v in cand_inputs.items()}
120 with torch.no_grad():
121 with torch.autocast(enabled=True, dtype=torch.bfloat16, device_type="cuda"):
122 tgt_output = model(tgt=cand_inputs)["tgt_reps"]
123 sim = model.compute_similarity(point_qry, tgt_output)
124 print(f"point ({px:.2f},{py:.2f}) | {string} = {sim}")
125
126'''
127--- Similarities (point prompts) ---
128point (0.25,0.50) | A cat = tensor([[0.1793]], device='cuda:0')
129point (0.25,0.50) | A dog = tensor([[0.1339]], device='cuda:0')
130point (0.25,0.50) | A tiger = tensor([[0.1314]], device='cuda:0')
131point (0.75,0.50) | A cat = tensor([[0.2232]], device='cuda:0')
132point (0.75,0.50) | A dog = tensor([[0.1742]], device='cuda:0')
133point (0.75,0.50) | A tiger = tensor([[0.1692]], device='cuda:0')
134'''
135
136# Scenario 3: BBox prompts — two examples (left/right)
137print("\n--- Similarities (bbox prompts) ---")
138bbox_examples = [
139 (0.05, 0.20, 0.45, 0.80), # left
140 (0.55, 0.20, 0.95, 0.80), # right
141]
142for (x1, y1, x2, y2) in bbox_examples:
143 bbox_text = f"{image_token}\nFind the caption that best describes the object in the bounding box, considering both local details and global context in the given image.\nReferring object bbox: ({int(x1*image.size[0])}, {int(y1*image.size[1])}, {int(x2*image.size[0])}, {int(y2*image.size[1])})"
144 q_inputs = process_fn({'text': [bbox_text], 'images': [image]}, processor=processor)
145 q_inputs['bbox'] = [x1 * sam_size, y1 * sam_size, x2 * sam_size, y2 * sam_size]
146 q_inputs = {k: v.to(device) if torch.is_tensor(v) else v for k, v in q_inputs.items()}
147 with torch.no_grad():
148 with torch.autocast(enabled=True, dtype=torch.bfloat16, device_type="cuda"):
149 bbox_qry = model(qry=q_inputs)["qry_reps"]
150 for string in test_strings:
151 cand_inputs = process_fn({'text': [string], 'images': [None]}, processor=processor)
152 cand_inputs = {k: v.to(device) if torch.is_tensor(v) else v for k, v in cand_inputs.items()}
153 with torch.no_grad():
154 with torch.autocast(enabled=True, dtype=torch.bfloat16, device_type="cuda"):
155 tgt_output = model(tgt=cand_inputs)["tgt_reps"]
156 sim = model.compute_similarity(bbox_qry, tgt_output)
157 print(f"bbox ({x1:.2f},{y1:.2f},{x2:.2f},{y2:.2f}) | {string} = {sim}")
158
159'''
160--- Similarities (bbox prompts) ---
161bbox (0.05,0.20,0.45,0.80) | A cat = tensor([[0.2100]], device='cuda:0')
162bbox (0.05,0.20,0.45,0.80) | A dog = tensor([[0.1512]], device='cuda:0')
163bbox (0.05,0.20,0.45,0.80) | A tiger = tensor([[0.1719]], device='cuda:0')
164bbox (0.55,0.20,0.95,0.80) | A cat = tensor([[0.1583]], device='cuda:0')
165bbox (0.55,0.20,0.95,0.80) | A dog = tensor([[0.1953]], device='cuda:0')
166bbox (0.55,0.20,0.95,0.80) | A tiger = tensor([[0.1225]], device='cuda:0')
167'''
