import torch
from transformers import AutoModelForImageTextToText, AutoProcessor


class TransformersModel:
    def __init__(self, model_id: str, to_device: str = "cuda"):
        self.model_id = model_id
        self.processor = AutoProcessor.from_pretrained(model_id)
        self.processor.image_processor.size = {"longest_edge": 3 * 384}
        self.model = AutoModelForImageTextToText.from_pretrained(model_id, torch_dtype=torch.bfloat16).to(to_device)

    def generate(self, messages: list[dict], **kwargs):
        inputs = self.processor.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
        ).to(self.model.device, dtype=torch.bfloat16)
        generated_ids = self.model.generate(**inputs, **kwargs)
        return self.processor.batch_decode(
                generated_ids[:, len(inputs["input_ids"][0]) :], skip_special_tokens=True
            )[0]