# from modeling_llavanext_for_embedding import LLaVANextForEmbedding
# from transformers import LlavaNextProcessor

# model = LLaVANextForEmbedding.from_pretrained("/share/junjie/code/VISTA2/240920mllmemb/llm_dense_retriever/MMRet-release/test").cuda()
# processor = LlavaNextProcessor.from_pretrained("/share/junjie/code/VISTA2/240920mllmemb/llm_dense_retriever/MMRet-release/test")

# texts = "find a image of a dog"

# inputs = processor(texts, return_tensors="pt").to("cuda")
# outputs = model(**inputs)
# print(outputs)


from transformers import LlavaNextProcessor, AutoModel

model = AutoModel.from_pretrained("/share/junjie/code/VISTA2/240920mllmemb/llm_dense_retriever/MMRet-release/MMRet-MLLM", trust_remote_code=True).cuda()
processor = LlavaNextProcessor.from_pretrained("/share/junjie/code/VISTA2/240920mllmemb/llm_dense_retriever/MMRet-release/MMRet-MLLM")

texts = "find a image of a dog"

inputs = processor(texts, return_tensors="pt").to("cuda")
outputs = model(**inputs)
print(outputs)