from transformers import AutoProcessor, AutoModel
from PIL import Image
import torch
import ImageReward as RM

class ImageRewardScorer(torch.nn.Module):
    def __init__(self, device="cuda", dtype=torch.float32, model_dir=None):
        super().__init__()
        self.model_path = f"{model_dir}/ImageReward/ImageReward.pt"
        self.med_config_path = f"{model_dir}/ImageReward/med_config.json"
        self.device = device
        self.dtype = dtype
        self.model = RM.load(self.model_path, device=device, med_config=self.med_config_path).eval().to(dtype=dtype)
        self.model.requires_grad_(False)
        
    @torch.no_grad()
    def __call__(self, prompt, images, device=None):
        self.model.to(device)
        _, rewards = self.model.inference_rank(prompt, images)
        return rewards