from typing import Dict, List, Any
from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
from PIL import Image
import requests
import torch
import gc

class EndpointHandler:
    def __init__(self, path=""):
        self.processor = AutoProcessor.from_pretrained(
            path,
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
            device_map='auto'
        )
        self.model = AutoModelForCausalLM.from_pretrained(
            path,
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
            device_map='auto',
            low_cpu_mem_usage=True
        )

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        # Clear CUDA cache
        torch.cuda.empty_cache()
        gc.collect()

        # Extract inputs from the request data
        inputs = data.get("inputs", {})
        image_url = inputs.get("image_url")
        text_prompt = inputs.get("text_prompt", "Describe this image.")

        if not image_url:
            return [{"error": "No image_url provided in inputs"}]

        # Download and process the image
        try:
            image = Image.open(requests.get(image_url, stream=True).raw)
            if image.mode != "RGB":
                image = image.convert("RGB")
        except Exception as e:
            return [{"error": f"Failed to load image: {str(e)}"}]

        # Process the image and text
        try:
            with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16):
                inputs = self.processor.process(
                    images=[image],
                    text=text_prompt
                )
                
                # Move inputs to the correct device and make a batch of size 1
                inputs = {k: v.to(self.model.device).unsqueeze(0) for k, v in inputs.items()}

                # Generate output
                output = self.model.generate_from_batch(
                    inputs,
                    GenerationConfig(max_new_tokens=200, stop_strings="<|endoftext|>"),
                    tokenizer=self.processor.tokenizer
                )

            # Decode the generated tokens
            generated_tokens = output[0, inputs['input_ids'].size(1):]
            generated_text = self.processor.tokenizer.decode(generated_tokens, skip_special_tokens=True)

            # Clear CUDA cache again
            torch.cuda.empty_cache()
            gc.collect()

            return [{"generated_text": generated_text}]
        except Exception as e:
            return [{"error": f"Error during generation: {str(e)}"}]