import torch
from PIL import Image
from transformers import AutoProcessor, LlavaForConditionalGeneration

# Load the model in half-precision
model = LlavaForConditionalGeneration.from_pretrained(
    "llava-hf/llava-1.5-7b-hf",
    cache_dir="/root/llm-project/utils/models/hub",
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    attn_implementation="flash_attention_2",
    device_map="auto")
processor = AutoProcessor.from_pretrained(
    "llava-hf/llava-1.5-7b-hf",
    cache_dir="/root/llm-project/utils/models/hub",
)
image1 = Image.open("/root/llm-project/LVLM/eval/Extended_CHAIR/images/chair-500/000000006763.jpg")
image2 = Image.open("./demo/material/1.jpg")

# Prepare a batch of two prompts
conversation_1 = [
    {
        "role": "user",
        "content": [
            {
                "type": "image"
            },
            {
                "type": "text",
                "text": "What is shown in this image? Please tell me.",
            },
        ],
    },
]

conversation_2 = [
    {
        "role": "user",
        "content": [
            {
                "type": "image"
            },
            {
                "type": "text",
                "text": "Describe this image."
            },
        ],
    },
]

prompt_1 = processor.apply_chat_template(conversation_1, add_generation_prompt=True)
prompt_2 = processor.apply_chat_template(conversation_2, add_generation_prompt=True)
prompts = [prompt_1, prompt_2]

# We can simply feed images in the order they have to be used in the text prompt
inputs = processor(
    images=[image1, image2], text=prompts, padding=True, return_tensors="pt").to(model.device, torch.float16)

# Generate
generate_ids = model.generate(**inputs, max_new_tokens=200)
out = processor.batch_decode(generate_ids, skip_special_tokens=True)
print(out)

['USER:  \nWhat is shown in this image? Please tell me. ASSISTANT: The image shows a man and a woman standing close to each other, posing for a picture. The man is wearing a tie, and they are both smiling for the camera. The scene takes place in a restaurant, as there are dining tables and chairs visible in the background.', 
'USER:  \nDescribe this image. ASSISTANT: The, the image, the image, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the, the,']

['USER:  \nWhat is shown in this image? Please tell me. ASSISTANT: The image shows a man and a woman standing close to each other, posing for a picture. The man is wearing a tie, and they are both smiling for the camera. The scene takes place in a restaurant, as there are dining tables and chairs visible in the background.', 
'USER:  \nDescribe this image. ASSISTANT: The image features a group of birds perched on a tree branch. There are five birds in total, with some sitting closer to the front of the branch and others further back. The birds are of various sizes and colors, creating a diverse and lively scene. The birds appear to be engaged in conversation or simply enjoying their time together on the branch.']