import torch
from PIL import Image
from transformers import AutoModel, AutoTokenizer

torch.manual_seed(100)

model = AutoModel.from_pretrained('./MiniCPM-V-4_5', trust_remote_code=True, # or openbmb/MiniCPM-o-2_6
    attn_implementation='sdpa', torch_dtype=torch.bfloat16) # sdpa or flash_attention_2, no eager
model = model.eval().cuda()
tokenizer = AutoTokenizer.from_pretrained('./MiniCPM-V-4_5', trust_remote_code=True) # or openbmb/MiniCPM-o-2_6

ckpt_path = "/user/qinyihua/checkpoints/model_tunnel/llava1.5_downsample_mlp_sft/sft/job_27496_ckpt_5000"
state_dict = torch.load(ckpt_path, map_location="cpu")

image = Image.open('./test_xx.jpg').convert('RGB')

enable_thinking=False # If `enable_thinking=True`, the thinking mode is enabled.

# First round chat 
question = "What is the landform in the picture?"
msgs = [{'role': 'user', 'content': [image, question]}]

answer = model.chat(
    msgs=msgs,
    tokenizer=tokenizer,
    enable_thinking=enable_thinking
)
print(answer)

# # Second round chat, pass history context of multi-turn conversation
# msgs.append({"role": "assistant", "content": [answer]})
# msgs.append({"role": "user", "content": ["What should I pay attention to when traveling here?"]})

# answer = model.chat(
#     msgs=msgs,
#     tokenizer=tokenizer
# )
# print(answer)