Views
No views yet
1git clone https://github.com/BwLiu01/IDMR.git
2cd IDMR
3pip install -r requirements.txt1import torch
2import numpy as np
3from PIL import Image
4from src.model import IDMRModel
5from src.vlm_backbone.intern_vl import InternVLProcessor
6from src.arguments import ModelArguments
7from transformers import AutoTokenizer, AutoImageProcessor
8
9device = "cuda"
10IMAGE_TOKEN = "<image>"
11
12# Load model and processor
13model_args = ModelArguments(model_name="lbw18601752667/IDMR-8B", model_backbone="internvl_2_5")
14
15# Initialize processor
16tokenizer = AutoTokenizer.from_pretrained(model_args.model_name, trust_remote_code=True)
17image_processor = AutoImageProcessor.from_pretrained(model_args.model_name, trust_remote_code=True, use_fast=False)
18processor = InternVLProcessor(image_processor=image_processor, tokenizer=tokenizer)
19
20# Load model
21model = IDMRModel.load(model_args).to(device, dtype=torch.bfloat16).eval()
22
23def get_embedding(text, image=None, type="qry"):
24 """Get embedding for text and/or image input"""
25 inputs = processor(
26 text=f"{IMAGE_TOKEN}\n {text}" if text else f"{IMAGE_TOKEN}\n Represent the given image.",
27 images=[image] if image else None,
28 return_tensors="pt",
29 max_length=1024,
30 truncation=True
31 )
32 inputs = {key: value.to(device) for key, value in inputs.items()}
33 inputs["image_flags"] = torch.tensor([1 if image else 0], dtype=torch.long).to(device)
34
35 with torch.no_grad(), torch.autocast(device_type=device, dtype=torch.bfloat16):
36 if type == "qry":
37 output = model(qry=inputs)["qry_reps"]
38 else:
39 output = model(tgt=inputs)["tgt_reps"]
40 return output.float()
41
42
43# Query
44query_text = "your query text"
45query_image = Image.open("your query image path")
46query_embedding = get_embedding(query_text, query_image, type="qry")
47
48# Target
49target_image = Image.open("your target image path")
50target_embedding = get_embedding(None, target_image, type="tgt")
51
52print(model.compute_similarity(query_embedding, target_embedding))
53