evalscope 0.14.0__py3-none-any.whl → 0.15.1__py3-none-any.whl

evalscope/metrics/t2v_metrics/models/clipscore_models/mps_model.py

@@ -0,0 +1,85 @@
+import torch
+from modelscope import AutoProcessor
+from transformers import CLIPConfig
+from typing import List
+
+from ...constants import CACHE_DIR
+from ..model import ScoreModel
+
+MPS_MODELS = ['mps']
+
+
+class MPSModel(ScoreModel):
+    'A wrapper for MPS Score models'
+
+    def __init__(self, model_name='mps', device='cuda', cache_dir=CACHE_DIR):
+        assert model_name in MPS_MODELS
+        super().__init__(model_name=model_name, device=device, cache_dir=cache_dir)
+
+    def load_model(self):
+        """Load the model, tokenizer, image transform
+        """
+        from ..utils import download_file
+        from .build_mps_model.clip_model import CLIPModel
+
+        assert self.model_name == 'mps'
+
+        processor_name_or_path = 'laion/CLIP-ViT-H-14-laion2B-s32B-b79K'
+        self.processor = AutoProcessor.from_pretrained(processor_name_or_path)
+
+        config = download_file('AI-ModelScope/MPS', file_name='config.json', cache_dir=self.cache_dir)
+        model_pretrained_path = download_file(
+            'AI-ModelScope/MPS', file_name='MPS_overall_state_dict.pt', cache_dir=self.cache_dir)  # modelscope model
+        model_weight = torch.load(model_pretrained_path, weights_only=True, map_location='cpu')
+
+        self.model = CLIPModel(config=CLIPConfig.from_json_file(config))
+        self.model.load_state_dict(model_weight, strict=False)
+        self.model.eval().to(self.device)
+
+    def load_images(self, image: List[str]) -> torch.Tensor:
+        """Load the image(s), and return a tensor (no preprocessing!!) put on self.device
+        """
+        image = [self.image_loader(x) for x in image]
+        image = self.processor(images=image, return_tensors='pt')['pixel_values']
+        return image
+
+    def process_text(self, text: List[str]) -> dict:
+        """Process the text(s), and return a tensor (after preprocessing) put on self.device
+        """
+        text_inputs = self.processor(
+            text=text,
+            padding='max_length',
+            truncation=True,
+            return_tensors='pt',
+        ).input_ids
+        return text_inputs
+
+    @torch.no_grad()
+    def forward(self, images: List[str], texts: List[str], condition=None) -> torch.Tensor:
+        """Forward pass of the model to return n scores for n (image, text) pairs (in PyTorch Tensor)
+        """
+        assert len(images) == len(texts)
+        image_input = self.load_images(images).to(self.device)
+        text_input = self.process_text(texts).to(self.device)
+        if condition is None:
+            condition = 'light, color, clarity, tone, style, ambiance, artistry, shape, face, hair, hands, limbs, structure, instance, texture, quantity, attributes, position, number, location, word, things.'
+        condition_batch = self.process_text(condition).repeat(text_input.shape[0], 1).to(self.device)
+
+        # embed
+        text_f, text_features = self.model.model.get_text_features(text_input)
+
+        image_f = self.model.model.get_image_features(image_input.half())
+        condition_f, _ = self.model.model.get_text_features(condition_batch)
+
+        sim_text_condition = torch.einsum('b i d, b j d -> b j i', text_f, condition_f)
+        sim_text_condition = torch.max(sim_text_condition, dim=1, keepdim=True)[0]
+        sim_text_condition = sim_text_condition / sim_text_condition.max()
+        mask = torch.where(sim_text_condition > 0.3, 0, float('-inf'))
+        mask = mask.repeat(1, image_f.shape[1], 1)
+        image_features = self.model.cross_model(image_f, text_f, mask.half())[:, 0, :]
+
+        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
+        image_score = self.model.logit_scale.exp() * text_features @ image_features.T
+
+        return image_score[0]

evalscope/metrics/t2v_metrics/models/clipscore_models/pickscore_model.py

@@ -0,0 +1,62 @@
+import torch
+from abc import abstractmethod
+from modelscope import AutoModel, AutoProcessor
+from typing import List
+
+from ...constants import CACHE_DIR
+from ..model import ScoreModel
+
+PICKSCORE_MODELS = ['pickscore-v1']
+
+
+class PickScoreModel(ScoreModel):
+    'A wrapper for PickScore models'
+
+    def __init__(self, model_name='pickscore-v1', device='cuda', cache_dir=CACHE_DIR):
+        assert model_name in PICKSCORE_MODELS
+        super().__init__(model_name=model_name, device=device, cache_dir=cache_dir)
+
+    def load_model(self):
+        """Load the model, tokenizer, image transform
+        """
+        assert self.model_name == 'pickscore-v1'
+        processor_name_or_path = 'laion/CLIP-ViT-H-14-laion2B-s32B-b79K'
+        # model_pretrained_name_or_path = "yuvalkirstain/PickScore_v1"
+        model_pretrained_name_or_path = 'AI-ModelScope/PickScore_v1'  # modelscope model
+
+        self.processor = AutoProcessor.from_pretrained(processor_name_or_path)
+        self.model = AutoModel.from_pretrained(model_pretrained_name_or_path).eval().to(self.device)
+
+    def load_images(self, image: List[str]) -> torch.Tensor:
+        """Load the image(s), and return a tensor (no preprocessing!!) put on self.device
+        """
+        image = [self.image_loader(x) for x in image]
+        image = self.processor(
+            images=image, padding=True, truncation=True, max_length=77, return_tensors='pt').to(self.device)
+        # image = torch.stack(image, dim=0).to(self.device)
+        return image
+
+    @torch.no_grad()
+    def forward(self, images: List[str], texts: List[str]) -> torch.Tensor:
+        """Forward pass of the model to return n scores for n (image, text) pairs (in PyTorch Tensor)
+        """
+        assert len(images) == len(texts)
+        image = self.load_images(images)
+        text_inputs = self.processor(
+            text=texts,
+            padding=True,
+            truncation=True,
+            max_length=77,
+            return_tensors='pt',
+        ).to(self.device)
+
+        # embed
+        image_embs = self.model.get_image_features(**image)
+        image_embs = image_embs / torch.norm(image_embs, dim=-1, keepdim=True)
+
+        text_embs = self.model.get_text_features(**text_inputs)
+        text_embs = text_embs / torch.norm(text_embs, dim=-1, keepdim=True)
+
+        # score
+        scores = (image_embs * text_embs).sum(dim=-1)
+        return scores

evalscope/metrics/t2v_metrics/models/itmscore_models/__init__.py

@@ -0,0 +1,26 @@
+from ...constants import CACHE_DIR
+from .blip2_itm_model import BLIP2_ITM_MODELS, BLIP2ITMScoreModel
+from .fga_blip2_model import FGA_BLIP2_MODELS, FGA_BLIP2ScoreModel
+from .image_reward_model import IMAGE_REWARD_MODELS, ImageRewardScoreModel
+
+ALL_ITM_MODELS = [
+    BLIP2_ITM_MODELS,
+    IMAGE_REWARD_MODELS,
+    FGA_BLIP2_MODELS,
+]
+
+
+def list_all_itmscore_models():
+    return [model for models in ALL_ITM_MODELS for model in models]
+
+
+def get_itmscore_model(model_name, device='cuda', cache_dir=CACHE_DIR):
+    assert model_name in list_all_itmscore_models()
+    if model_name in BLIP2_ITM_MODELS:
+        return BLIP2ITMScoreModel(model_name, device=device, cache_dir=cache_dir)
+    elif model_name in IMAGE_REWARD_MODELS:
+        return ImageRewardScoreModel(model_name, device=device, cache_dir=cache_dir)
+    elif model_name in FGA_BLIP2_MODELS:
+        return FGA_BLIP2ScoreModel(model_name, device=device, cache_dir=cache_dir)
+    else:
+        raise NotImplementedError()

evalscope/metrics/t2v_metrics/models/itmscore_models/blip2_itm_model.py

@@ -0,0 +1,84 @@
+import os
+import torch
+from torchvision import transforms
+from typing import List
+
+from ...constants import CACHE_DIR
+from ..model import ScoreModel
+from ..vqascore_models.lavis.models import load_model
+
+BLIP2_ITM_MODELS = {
+    'blip2-itm': {
+        'variant': 'pretrain'
+    },
+    'blip2-itm-vitL': {
+        'variant': 'pretrain_vitL'
+    },
+    'blip2-itm-coco': {
+        'variant': 'coco'
+    },
+}
+
+
+class BLIP2ITMScoreModel(ScoreModel):
+    'A wrapper for BLIP-2 ITMScore models'
+
+    def __init__(self, model_name='blip2-itm', device='cuda', cache_dir=CACHE_DIR):
+        assert model_name in BLIP2_ITM_MODELS, f'Model name must be one of {BLIP2_ITM_MODELS.keys()}'
+        os.environ['TORCH_HOME'] = cache_dir
+        super().__init__(model_name=model_name, device=device, cache_dir=cache_dir)
+
+    def load_model(self):
+        """Load the model, tokenizer, image transform
+        """
+        self.variant = BLIP2_ITM_MODELS[self.model_name]['variant']
+        self.model = load_model('blip2', self.variant, is_eval=True, device=self.device)
+        if self.variant == 'coco':
+            size = 364
+        else:
+            size = 224
+        self.image_preprocess = transforms.Compose([
+            transforms.Resize((size, size), interpolation=transforms.functional.InterpolationMode.BICUBIC),
+            transforms.ToTensor(),
+            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
+        ])
+
+    def load_images(self, image: List[str]) -> torch.Tensor:
+        """Load the image(s), and return a tensor (after preprocessing) put on self.device
+        """
+        image = [self.image_loader(x) for x in image]
+        image = [self.image_preprocess(image) for image in image]
+        assert all(x.shape == image[0].shape for x in image)
+        image = torch.stack(image, dim=0).to(self.device)
+        return image
+
+    @torch.no_grad()
+    @torch.autocast(device_type='cuda', dtype=torch.float16)
+    def forward(self, images: List[str], texts: List[str]) -> torch.Tensor:
+        """Forward pass of the model to return n scores for n (image, text) pairs (in PyTorch Tensor)
+        """
+        assert len(images) == len(texts), 'Number of images and texts must match'
+
+        images = self.load_images(images)
+        image_feat = self.model.ln_vision(self.model.visual_encoder(images))
+        image_att = torch.ones(image_feat.size()[:-1], dtype=torch.long).to(self.device)
+        query_token = self.model.query_tokens.expand(image_feat.shape[0], -1, -1)
+        query_att = torch.ones(query_token.size()[:-1], dtype=torch.long).to(query_token.device)
+
+        text_input = self.model.tokenizer(
+            texts, padding='max_length', truncation=True, max_length=35, return_tensors='pt').to(self.device)
+
+        attention_mask_all = torch.cat([query_att, text_input.attention_mask], dim=1)
+        output_itm = self.model.Qformer.bert(
+            text_input.input_ids,
+            query_embeds=query_token,
+            attention_mask=attention_mask_all,
+            encoder_hidden_states=image_feat,
+            encoder_attention_mask=image_att,
+            return_dict=True,
+        )
+        vl_embeddings = output_itm.last_hidden_state[:, :query_token.size(1), :]
+        vl_output = self.model.itm_head(vl_embeddings)
+        itm_logits = vl_output.mean(dim=1)
+        itm_prob = torch.nn.functional.softmax(itm_logits, dim=-1)[:, 1]
+        return itm_prob

evalscope/metrics/t2v_metrics/models/itmscore_models/fga_blip2_model.py

@@ -0,0 +1,97 @@
+import os
+import torch
+from modelscope import AutoTokenizer
+from typing import List, Union
+
+from ...constants import CACHE_DIR
+from ..model import ScoreModel
+from ..vqascore_models.lavis.models import load_model_and_preprocess
+
+FGA_BLIP2_MODELS = {
+    'fga_blip2': {
+        'variant': 'coco'
+    },
+}
+
+
+def get_index(list1, list2):
+    len_list1 = len(list1)
+    len_list2 = len(list2)
+    for i in range(len_list2 - len_list1 + 1):
+        if list2[i:i + len_list1] == list1:
+            return i
+    return 0
+
+
+class FGA_BLIP2ScoreModel(ScoreModel):
+    'A wrapper for FGA BLIP-2 ITMScore models'
+
+    def __init__(self, model_name='fga_blip2', device='cuda', cache_dir=CACHE_DIR):
+        assert model_name in FGA_BLIP2_MODELS, f'Model name must be one of {FGA_BLIP2_MODELS.keys()}'
+        os.environ['TORCH_HOME'] = cache_dir
+        super().__init__(model_name=model_name, device=device, cache_dir=cache_dir)
+
+    def load_model(self):
+        """Load the model, tokenizer, image transform
+        """
+        from ..utils import download_file
+
+        # load tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained('AI-ModelScope/bert-base-uncased', truncation_side='right')
+        self.tokenizer.add_special_tokens({'bos_token': '[DEC]'})
+        # load model
+        self.variant = FGA_BLIP2_MODELS[self.model_name]['variant']
+        self.model, self.vis_processors, self.text_processors = load_model_and_preprocess(
+            'fga_blip2', self.variant, is_eval=True, device=self.device)
+        # load pretrained weights
+        model_weight_path = download_file(
+            'AI-ModelScope/FGA-BLIP2', file_name='fga_blip2.pth', cache_dir=self.cache_dir)
+        self.model.load_checkpoint(model_weight_path)
+        self.model.eval()
+
+    def load_images(self, image):
+        pass
+
+    @torch.no_grad()
+    @torch.autocast(device_type='cuda', dtype=torch.float16)
+    def forward(self, images: List[str], texts: List[Union[str, dict]]) -> torch.Tensor:
+        """Forward pass of the model to return n scores for n (image, text) pairs (in PyTorch Tensor)
+        """
+        assert len(images) == len(texts), 'Number of images and texts must match'
+
+        result_list = []
+        for image_path, text in zip(images, texts):
+            if isinstance(text, str):
+                elements = []  # elements scores
+                prompt = text
+            else:
+                elements = text['tags']
+                prompt = text['prompt']
+
+            image = self.image_loader(image_path)
+            image = self.vis_processors['eval'](image).to(self.device)
+            prompt = self.text_processors['eval'](prompt)
+            prompt_ids = self.tokenizer(prompt).input_ids
+
+            alignment_score, scores = self.model.element_score(image.unsqueeze(0), [prompt])
+
+            elements_score = dict()
+            for element in elements:
+                element_ = element.rpartition('(')[0]
+                element_ids = self.tokenizer(element_).input_ids[1:-1]
+
+                idx = get_index(element_ids, prompt_ids)
+                if idx:
+                    mask = [0] * len(prompt_ids)
+                    mask[idx:idx + len(element_ids)] = [1] * len(element_ids)
+
+                    mask = torch.tensor(mask).to(self.device)
+                    elements_score[element] = (scores * mask).sum() / mask.sum()
+                else:
+                    elements_score[element] = torch.tensor(0.0).to(self.device)
+            if elements_score:
+                result_list.append({'overall_score': alignment_score, **elements_score})
+            else:
+                result_list.append(alignment_score)
+
+        return result_list

evalscope/metrics/t2v_metrics/models/itmscore_models/image_reward/ImageReward.py

@@ -0,0 +1,171 @@
+# adapted from https://github.com/THUDM/ImageReward/tree/main/ImageReward
+
+import os
+import torch
+import torch.nn as nn
+from PIL import Image
+from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor
+
+from .blip_pretrain import BLIP_Pretrain
+
+try:
+    from torchvision.transforms import InterpolationMode
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+
+
+def _convert_image_to_rgb(image):
+    return image.convert('RGB')
+
+
+def _transform(n_px):
+    return Compose([
+        Resize(n_px, interpolation=BICUBIC),
+        CenterCrop(n_px),
+        _convert_image_to_rgb,
+        ToTensor(),
+        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+
+class MLP(nn.Module):
+
+    def __init__(self, input_size):
+        super().__init__()
+        self.input_size = input_size
+
+        self.layers = nn.Sequential(
+            nn.Linear(self.input_size, 1024),
+            #nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(1024, 128),
+            #nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(128, 64),
+            #nn.ReLU(),
+            nn.Dropout(0.1),
+            nn.Linear(64, 16),
+            #nn.ReLU(),
+            nn.Linear(16, 1))
+
+        # initial MLP param
+        for name, param in self.layers.named_parameters():
+            if 'weight' in name:
+                nn.init.normal_(param, mean=0.0, std=1.0 / (self.input_size + 1))
+            if 'bias' in name:
+                nn.init.constant_(param, val=0)
+
+    def forward(self, input):
+        return self.layers(input)
+
+
+class ImageReward(nn.Module):
+
+    def __init__(self, med_config, device='cpu'):
+        super().__init__()
+        self.device = device
+
+        self.blip = BLIP_Pretrain(image_size=224, vit='large', med_config=med_config)
+        self.preprocess = _transform(224)
+        self.mlp = MLP(768)
+
+        self.mean = 0.16717362830052426
+        self.std = 1.0333394966054072
+
+    def score_gard(self, prompt_ids, prompt_attention_mask, image):
+
+        image_embeds = self.blip.visual_encoder(image)
+        # text encode cross attention with image
+        image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(self.device)
+        text_output = self.blip.text_encoder(
+            prompt_ids,
+            attention_mask=prompt_attention_mask,
+            encoder_hidden_states=image_embeds,
+            encoder_attention_mask=image_atts,
+            return_dict=True,
+        )
+
+        txt_features = text_output.last_hidden_state[:, 0, :]  # (feature_dim)
+        rewards = self.mlp(txt_features)
+        rewards = (rewards - self.mean) / self.std
+
+        return rewards
+
+    def score(self, prompt, image):
+
+        if (type(image).__name__ == 'list'):
+            _, rewards = self.inference_rank(prompt, image)
+            return rewards
+
+        # text encode
+        text_input = self.blip.tokenizer(
+            prompt, padding='max_length', truncation=True, max_length=35, return_tensors='pt').to(self.device)
+
+        # image encode
+        if isinstance(image, Image.Image):
+            pil_image = image
+        elif isinstance(image, str) and os.path.isfile(image):
+            pil_image = Image.open(image)
+        else:
+            raise TypeError(
+                r'This image parameter type has not been supportted yet. Please pass PIL.Image or file path str.')
+
+        image = self.preprocess(pil_image).unsqueeze(0).to(self.device)
+        image_embeds = self.blip.visual_encoder(image)
+
+        # text encode cross attention with image
+        image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(self.device)
+        text_output = self.blip.text_encoder(
+            text_input.input_ids,
+            attention_mask=text_input.attention_mask,
+            encoder_hidden_states=image_embeds,
+            encoder_attention_mask=image_atts,
+            return_dict=True,
+        )
+
+        txt_features = text_output.last_hidden_state[:, 0, :].float()  # (feature_dim)
+        rewards = self.mlp(txt_features)
+        rewards = (rewards - self.mean) / self.std
+
+        return rewards.detach().cpu().numpy().item()
+
+    def inference_rank(self, prompt, generations_list):
+
+        text_input = self.blip.tokenizer(
+            prompt, padding='max_length', truncation=True, max_length=35, return_tensors='pt').to(self.device)
+
+        txt_set = []
+        for generation in generations_list:
+            # image encode
+            if isinstance(generation, Image.Image):
+                pil_image = generation
+            elif isinstance(generation, str):
+                if os.path.isfile(generation):
+                    pil_image = Image.open(generation)
+            else:
+                raise TypeError(
+                    r'This image parameter type has not been supportted yet. Please pass PIL.Image or file path str.')
+            image = self.preprocess(pil_image).unsqueeze(0).to(self.device)
+            image_embeds = self.blip.visual_encoder(image)
+
+            # text encode cross attention with image
+            image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(self.device)
+            text_output = self.blip.text_encoder(
+                text_input.input_ids,
+                attention_mask=text_input.attention_mask,
+                encoder_hidden_states=image_embeds,
+                encoder_attention_mask=image_atts,
+                return_dict=True,
+            )
+            txt_set.append(text_output.last_hidden_state[:, 0, :])
+
+        txt_features = torch.cat(txt_set, 0).float()  # [image_num, feature_dim]
+        rewards = self.mlp(txt_features)  # [image_num, 1]
+        rewards = (rewards - self.mean) / self.std
+        rewards = torch.squeeze(rewards)
+        _, rank = torch.sort(rewards, dim=0, descending=True)
+        _, indices = torch.sort(rank, dim=0)
+        indices = indices + 1
+
+        return indices.detach().cpu().numpy().tolist(), rewards.detach().cpu().numpy().tolist()

evalscope/metrics/t2v_metrics/models/itmscore_models/image_reward/blip_pretrain.py

@@ -0,0 +1,80 @@
+'''
+ * Adapted from BLIP (https://github.com/salesforce/BLIP)
+'''
+
+from modelscope import AutoTokenizer
+from torch import nn
+
+from ...vqascore_models.lavis.models.med import BertConfig, BertModel
+from ...vqascore_models.lavis.models.vit import VisionTransformer
+
+
+def init_tokenizer():
+    tokenizer = AutoTokenizer.from_pretrained('AI-ModelScope/bert-base-uncased')
+    tokenizer.add_special_tokens({'bos_token': '[DEC]'})
+    tokenizer.add_special_tokens({'additional_special_tokens': ['[ENC]']})
+    tokenizer.enc_token_id = tokenizer.additional_special_tokens_ids[0]
+    return tokenizer
+
+
+def create_vit(vit, image_size, use_grad_checkpointing=False, ckpt_layer=0, drop_path_rate=0):
+
+    assert vit in ['base', 'large'], 'vit parameter must be base or large'
+    if vit == 'base':
+        vision_width = 768
+        visual_encoder = VisionTransformer(
+            img_size=image_size,
+            patch_size=16,
+            embed_dim=vision_width,
+            depth=12,
+            num_heads=12,
+            use_grad_checkpointing=use_grad_checkpointing,
+            ckpt_layer=ckpt_layer,
+            drop_path_rate=0 or drop_path_rate)
+    elif vit == 'large':
+        vision_width = 1024
+        visual_encoder = VisionTransformer(
+            img_size=image_size,
+            patch_size=16,
+            embed_dim=vision_width,
+            depth=24,
+            num_heads=16,
+            use_grad_checkpointing=use_grad_checkpointing,
+            ckpt_layer=ckpt_layer,
+            drop_path_rate=0.1 or drop_path_rate)
+    return visual_encoder, vision_width
+
+
+class BLIP_Pretrain(nn.Module):
+
+    def __init__(
+        self,
+        med_config='med_config.json',
+        image_size=224,
+        vit='base',
+        vit_grad_ckpt=False,
+        vit_ckpt_layer=0,
+        embed_dim=256,
+        queue_size=57600,
+        momentum=0.995,
+    ):
+        """
+        Args:
+            med_config (str): path for the mixture of encoder-decoder model's configuration file
+            image_size (int): input image size
+            vit (str): model size of vision transformer
+        """
+        super().__init__()
+
+        self.visual_encoder, vision_width = create_vit(vit, image_size, vit_grad_ckpt, vit_ckpt_layer, 0)
+
+        self.tokenizer = init_tokenizer()
+        encoder_config = BertConfig.from_json_file(med_config)
+        encoder_config.encoder_width = vision_width
+        encoder_config.add_type_embeddings = False
+        self.text_encoder = BertModel(config=encoder_config, add_pooling_layer=False)
+
+        text_width = self.text_encoder.config.hidden_size
+
+        self.vision_proj = nn.Linear(vision_width, embed_dim)
+        self.text_proj = nn.Linear(text_width, embed_dim)

evalscope/metrics/t2v_metrics/models/itmscore_models/image_reward_model.py

@@ -0,0 +1,73 @@
+import torch
+from typing import List
+
+from ...constants import CACHE_DIR
+from ..model import ScoreModel
+
+IMAGE_REWARD_MODELS = {
+    'image-reward-v1': {
+        'variant': 'ImageReward-v1.0'
+    },
+}
+
+
+class ImageRewardScoreModel(ScoreModel):
+    'A wrapper for ImageReward ITMScore (finetuned on human preference) models'
+
+    def __init__(self, model_name='image-reward-v1', device='cuda', cache_dir=CACHE_DIR):
+        assert model_name in IMAGE_REWARD_MODELS, f'Model name must be one of {IMAGE_REWARD_MODELS.keys()}'
+
+        super().__init__(model_name=model_name, device=device, cache_dir=cache_dir)
+
+    def load_model(self):
+        """Load the model, tokenizer, image transform
+        """
+        from ..utils import download_file
+        from .image_reward.ImageReward import ImageReward
+
+        self.variant = IMAGE_REWARD_MODELS[self.model_name]['variant']
+
+        self.model_path = download_file('ZhipuAI/ImageReward', file_name='ImageReward.pt', cache_dir=self.cache_dir)
+        self.med_config = download_file('ZhipuAI/ImageReward', file_name='med_config.json', cache_dir=self.cache_dir)
+
+        state_dict = torch.load(self.model_path, map_location='cpu')
+        self.model = ImageReward(device=self.device, med_config=self.med_config).to(self.device)
+        msg = self.model.load_state_dict(state_dict, strict=False)
+        self.model.eval()
+
+    def load_images(self, image: List[str]) -> torch.Tensor:
+        """Load the image(s), and return a tensor (after preprocessing) put on self.device
+        """
+        image = [self.image_loader(x) for x in image]
+        image = [self.model.preprocess(image) for image in image]
+        assert all(x.shape == image[0].shape for x in image)
+        image = torch.stack(image, dim=0).to(self.device)
+        return image
+
+    @torch.no_grad()
+    def forward(self, images: List[str], texts: List[str]) -> torch.Tensor:
+        """Forward pass of the model to return n scores for n (image, text) pairs (in PyTorch Tensor)
+        """
+        assert len(images) == len(texts), 'Number of images and texts must match'
+        rewards = torch.zeros(len(texts), dtype=torch.float32).to(self.device)
+        images = self.load_images(images)
+        for index in range(len(texts)):
+            text_input = self.model.blip.tokenizer(
+                texts[index], padding='max_length', truncation=True, max_length=35, return_tensors='pt').to(self.device)
+            image_embeds = self.model.blip.visual_encoder(images[index].unsqueeze(0))
+            image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(self.device)
+            text_output = self.model.blip.text_encoder(
+                text_input.input_ids,
+                attention_mask=text_input.attention_mask,
+                encoder_hidden_states=image_embeds,
+                encoder_attention_mask=image_atts,
+                return_dict=True,
+            )
+
+            txt_features = text_output.last_hidden_state[:, 0, :].float()  # (feature_dim)
+            reward_score = self.model.mlp(txt_features)
+            reward_score = (reward_score - self.model.mean) / self.model.std
+
+            rewards[index] = reward_score
+
+        return rewards