from transformers import InstructBlipProcessor, InstructBlipForConditionalGeneration
import torch
from PIL import Image



class InstructBlip:
    device = "cuda" if torch.cuda.is_available() else "cpu"

    def __init__(self, model_pretrain:str = "Salesforce/instructblip-vicuna-7b"):                                                  
        self.model = InstructBlipForConditionalGeneration.from_pretrained(model_pretrain
                                                        , device_map={"": 0}, torch_dtype=torch.float16)
        self.processor = InstructBlipProcessor.from_pretrained(model_pretrain)
        
    def image_captioning(self, image: Image.Image) -> str:
        prompt = "What are the features of this picture?"
        inputs = self.processor(images=image, text=prompt, return_tensors="pt").to(self.device)

        outputs = self.model.generate(
                **inputs,
                do_sample=False,
                num_beams=5,
                max_length=256,
                min_length=1,
                top_p=0.9,
                repetition_penalty=1.5,
                length_penalty=1.0,
                temperature=1,
        )
        generated_text = self.processor.batch_decode(outputs, skip_special_tokens=True)[0].strip()

        return generated_text

    def visual_question_answering(self, image: Image.Image, prompt: str) -> str:
        inputs = self.processor(images=image, text=prompt, return_tensors="pt").to(device)

        outputs = self.model.generate(
                **inputs,
                do_sample=False,
                num_beams=5,
                max_length=256,
                min_length=1,
                top_p=0.9,
                repetition_penalty=1.5,
                length_penalty=1.0,
                temperature=1,
        )
        generated_text = self.processor.batch_decode(outputs, skip_special_tokens=True)[0].strip()
        
        return generated_text