File size: 2,163 Bytes
230a1aa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | from transformers import AutoProcessor
import torch
from ..core import ModelConfig
from ..core.device.npu_compatible_device import get_device_type
from ..models.pickscore import PickScoreModel
from .base import Metric
class PickScoreMetric(Metric):
def __init__(self, model: PickScoreModel):
super().__init__()
self.model = model
@classmethod
def from_pretrained(
cls,
model_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="PickScore/model.safetensors"),
processor_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="PickScore/"),
torch_dtype: torch.dtype = None,
device: torch.device = get_device_type(),
max_length: int = 77,
processor_kwargs: dict = None,
vram_limit: float = None,
):
processor_kwargs = processor_kwargs or {}
model_pool = cls.download_and_load_models([model_config], torch_dtype=torch_dtype, device=device, vram_limit=vram_limit)
model = model_pool.fetch_model("image_metrics_clip_hf")
processor_config.download_if_necessary()
processor = AutoProcessor.from_pretrained(processor_config.path, **processor_kwargs)
model = PickScoreModel(model=model, processor=processor, max_length=max_length).eval()
return cls(model)
@torch.no_grad()
def score(
self,
prompt: str | list[str],
images,
):
scores = self.model(prompt, images)
return self.tensor_to_list(scores)
@torch.no_grad()
def probabilities(
self,
prompt: str | list[str],
images,
):
scores = self.model(prompt, images)
probabilities = torch.softmax(scores, dim=-1)
return self.tensor_to_list(probabilities)
def calc_probs(self, prompt: str | list[str], images):
return self.probabilities(prompt, images)
def compute(self, prompt: str | list[str], images):
return self.score(prompt, images)
def forward(self, prompt: str | list[str], images):
return self.score(prompt, images)
|