1git clone https://github.com/deepglint/UniME-v2.git
2cd UniME-v2
1conda create -n uniMEv2 python=3.10 -y
2conda activate uniMEv2
3pip install -r requirements.txt
4
5# Optional: Install Flash Attention for acceleration
6# wget https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
7# pip install flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
1import torch
2from torch.nn import functional as F
3from utils.utils import init_model_and_processor, prepare_stage_data, parse_answer_index
4
5device="cuda"
6embedding=False # adjust embedding model or rerank model
7if embedding:
8 model_name="models/UniME-V2_qwen2VL_2B"
9 # model_name="models/UniME-V2_qwen2VL_7B"
10 # model_name="models/UniME-V2_LLaVA_onevision_8B"
11 text = "A man is crossing the street with a red car parked nearby."
12 image_path = "Figures/demo.png"
13else:
14 model_name="models/UniME-v2-rerank_qwen25VL_7B"
15 text = ["A man is crossing the street with a red car parked nearby.", #! Target text
16 "A woman is walking her dog with a blue bicycle leaning nearby.",
17 "A child is riding a scooter past a green truck stopped nearby.",
18 "A couple is waiting for the bus beside a yellow taxi parked nearby.",
19 "A jogger is running along the path with a black motorcycle parked nearby."]
20 image_path = "Figures/demo.png"
21
22model, processor = init_model_and_processor(model_name, device, embedding=embedding)
23
24if embedding:
25 inputs_image, inputs_txt = prepare_stage_data(model_name, processor, text, image_path, embedding=embedding)
26 inputs_image = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_image.items()}
27 inputs_txt = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_txt.items()}
28 with torch.no_grad():
29 emb_text = model(**inputs_txt, output_hidden_states=True, return_dict=True).hidden_states[-1][:, -1, :]
30 emb_image = model(**inputs_image, output_hidden_states=True, return_dict=True).hidden_states[-1][:, -1, :]
31 emb_text = F.normalize(emb_text, dim=-1)
32 emb_image = F.normalize(emb_image, dim=-1)
33 Score = emb_image @ emb_text.T
34 print("Score: ", Score.item()) # qwen2VL 2B : Score: 0.62109375
35else:
36 inputs = prepare_stage_data(model_name, processor, text, image_path, embedding=embedding)
37 inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
38 with torch.no_grad():
39 generated_ids = model.generate(**inputs, max_new_tokens=128, output_scores=True, return_dict_in_generate=True, do_sample=False).sequences
40 generated_ids_trimmed = [
41 out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs['input_ids'], generated_ids)
42 ]
43 output_text = processor.batch_decode(
44 generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
45 )
46 print("Rerank Answer: ", parse_answer_index(output_text[0])) # qwen25VL 7B: Rerank Answer: 0
If you find this repository useful, please use the following BibTeX entry for citation.
1@misc{gu2025unimev2mllmasajudgeuniversalmultimodal,
2 title={UniME-V2: MLLM-as-a-Judge for Universal Multimodal Embedding Learning},
3 author={Tiancheng Gu and Kaicheng Yang and Kaichen Zhang and Xiang An and Ziyong Feng and Yueyi Zhang and Weidong Cai and Jiankang Deng and Lidong Bing},
4 year={2025},
5 eprint={2510.13515},
6 archivePrefix={arXiv},
7 primaryClass={cs.CV},
8 url={https://arxiv.org/abs/2510.13515},
9}
10
11@inproceedings{unime,
12 title={Breaking the Modality Barrier: Universal Embedding Learning with Multimodal LLMs},
13 author={Gu, Tiancheng and Yang, Kaicheng and Feng, Ziyong and Wang, Xingjun and Zhang, Yanzhao and Long, Dingkun and Chen, Yingda and Cai, Weidong and Deng, Jiankang},
14 booktitle={ACM MM},
15 year={2025}
16}
17
