| import numpy as np |
| import torch |
| from PIL import Image |
| from io import BytesIO |
| from tqdm.auto import tqdm |
| from transformers import CLIPFeatureExtractor, CLIPImageProcessor |
| from transformers import CLIPConfig |
| from dataclasses import dataclass |
| from transformers import CLIPModel as HFCLIPModel |
| from safetensors.torch import load_file |
| from torch import nn, einsum |
|
|
| from .trainer.models.base_model import BaseModelConfig |
|
|
| from transformers import CLIPConfig |
| from transformers import AutoProcessor, AutoModel, AutoTokenizer |
| from typing import Any, Optional, Tuple, Union, List |
| import torch |
|
|
| from .trainer.models.cross_modeling import Cross_model |
| from .trainer.models import clip_model |
| import torch.nn.functional as F |
| import gc |
| import json |
| from .config import MODEL_PATHS |
|
|
| class MPScore(torch.nn.Module): |
| def __init__(self, device: Union[str, torch.device], path: str = MODEL_PATHS, condition: str = 'overall'): |
| super().__init__() |
| """Initialize the MPSModel with a processor, tokenizer, and model. |
| |
| Args: |
| device (Union[str, torch.device]): The device to load the model on. |
| """ |
| self.device = device |
| processor_name_or_path = path.get("clip") |
| self.image_processor = CLIPImageProcessor.from_pretrained(processor_name_or_path) |
| self.tokenizer = AutoTokenizer.from_pretrained(processor_name_or_path, trust_remote_code=True) |
| self.model = clip_model.CLIPModel(processor_name_or_path, config_file=True) |
| state_dict = load_file(path.get("mps")) |
| self.model.load_state_dict(state_dict, strict=False) |
| self.model.to(device) |
| self.condition = condition |
|
|
| def _calculate_score(self, image: torch.Tensor, prompt: str) -> float: |
| """Calculate the reward score for a single image and prompt. |
| |
| Args: |
| image (torch.Tensor): The processed image tensor. |
| prompt (str): The prompt text. |
| |
| Returns: |
| float: The reward score. |
| """ |
| def _tokenize(caption): |
| input_ids = self.tokenizer( |
| caption, |
| max_length=self.tokenizer.model_max_length, |
| padding="max_length", |
| truncation=True, |
| return_tensors="pt" |
| ).input_ids |
| return input_ids |
|
|
| text_input = _tokenize(prompt).to(self.device) |
| if self.condition == 'overall': |
| condition_prompt = 'light, color, clarity, tone, style, ambiance, artistry, shape, face, hair, hands, limbs, structure, instance, texture, quantity, attributes, position, number, location, word, things' |
| elif self.condition == 'aesthetics': |
| condition_prompt = 'light, color, clarity, tone, style, ambiance, artistry' |
| elif self.condition == 'quality': |
| condition_prompt = 'shape, face, hair, hands, limbs, structure, instance, texture' |
| elif self.condition == 'semantic': |
| condition_prompt = 'quantity, attributes, position, number, location' |
| else: |
| raise ValueError( |
| f"Unsupported condition: {self.condition}. Choose 'overall', 'aesthetics', 'quality', or 'semantic'.") |
| condition_batch = _tokenize(condition_prompt).repeat(text_input.shape[0], 1).to(self.device) |
|
|
| with torch.no_grad(): |
| text_f, text_features = self.model.model.get_text_features(text_input) |
|
|
| image_f = self.model.model.get_image_features(image.half()) |
| condition_f, _ = self.model.model.get_text_features(condition_batch) |
|
|
| sim_text_condition = 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].cpu().numpy().item() |
|
|
| @torch.no_grad() |
| def score(self, images: Union[str, List[str], Image.Image, List[Image.Image]], prompt: str) -> List[float]: |
| """Score the images based on the prompt. |
| |
| Args: |
| images (Union[str, List[str], Image.Image, List[Image.Image]]): Path(s) to the image(s) or PIL image(s). |
| prompt (str): The prompt text. |
| |
| Returns: |
| List[float]: List of reward scores for the images. |
| """ |
| if isinstance(images, (str, Image.Image)): |
| |
| if isinstance(images, str): |
| image = self.image_processor(Image.open(images), return_tensors="pt")["pixel_values"].to(self.device) |
| else: |
| image = self.image_processor(images, return_tensors="pt")["pixel_values"].to(self.device) |
| return [self._calculate_score(image, prompt)] |
| elif isinstance(images, list): |
| |
| scores = [] |
| for one_images in images: |
| if isinstance(one_images, str): |
| image = self.image_processor(Image.open(one_images), return_tensors="pt")["pixel_values"].to(self.device) |
| elif isinstance(one_images, Image.Image): |
| image = self.image_processor(one_images, return_tensors="pt")["pixel_values"].to(self.device) |
| else: |
| raise TypeError("The type of parameter images is illegal.") |
| scores.append(self._calculate_score(image, prompt)) |
| return scores |
| else: |
| raise TypeError("The type of parameter images is illegal.") |
|
|
