import gradio as gr
from transformers import AutoModelForVision2Seq, AutoProcessor, BitsAndBytesConfig
import torch


model_id = "HuggingFaceM4/idefics2-8b"

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.float16
)
processor = AutoProcessor.from_pretrained(model_id)
model = AutoModelForVision2Seq.from_pretrained(model_id, torch_dtype=torch.float16, quantization_config=quantization_config)


def respond(multimodal_input):
    images = multimodal_input["files"]
    content = [{"type": "image"} for _ in images]
    content.append({"type": "text", "text": multimodal_input["text"]})
    messages = [{"role": "user", "content": content}]
    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
    inputs = processor(text=prompt, images=[images], return_tensors="pt")
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    num_tokens = len(inputs["input_ids"][0])
    with torch.inference_mode():
        generated_ids = model.generate(**inputs, max_new_tokens=500)

    new_tokens = generated_ids[:, num_tokens:]
    generated_text = processor.batch_decode(new_tokens, skip_special_tokens=True)[0]
    return generated_text


gr.Interface(
    respond,
    inputs=[gr.MultimodalTextbox(file_types=["image"], show_label=False)],
    outputs="text",
    title="IDEFICS2-8B DPO",
    description="Try IDEFICS2-8B fine-tuned using direct preference optimization (DPO) in this demo. Learn more about vision language model DPO integration of TRL [here](https://huggingface.co/blog/dpo_vlm).",
    examples=[
        {"text": "What is the type of flower in the image and what insect is on it?", "files": ["./bee.jpg"]},
        {"text": "Describe the image", "files": ["./howl.jpg"]},
    ],
).launch()