import torch
from aesthetic_predictor_v2_5 import convert_v2_5_from_siglip
from PIL import Image


class AestheticScorer:
    def __init__(self):
        self.model, self.preprocessor = convert_v2_5_from_siglip(
            low_cpu_mem_usage=True,
            trust_remote_code=True,
        )
        self.model = self.model.to(torch.bfloat16).cuda()

    def __call__(self, image: Image.Image | str, device) -> float:
        if isinstance(image, str):
            image = Image.open(image).convert("RGB")
        elif not isinstance(image, Image.Image):
            raise ValueError("Input must be a file path or a PIL Image.")
        pixel_values = (
            self.preprocessor(images=image, return_tensors="pt")
            .pixel_values.to(torch.bfloat16)
            .cuda()
        )
        with torch.inference_mode():
            self.model.to(device)
            score = self.model(pixel_values).logits.squeeze().float().cpu().numpy()
        return float(score)