| | 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.") |
| |
|
