| from transformers import AutoProcessor |
| import torch |
| from ..core import ModelConfig |
| from ..core.device.npu_compatible_device import get_device_type |
| from ..models.hpsv3 import HPSv3Model |
| from .base import Metric |
|
|
|
|
| class HPSv3Metric(Metric): |
| def __init__(self, model: HPSv3Model): |
| super().__init__() |
| self.model = model |
|
|
| @classmethod |
| def from_pretrained( |
| cls, |
| model_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="HPSv3/model.safetensors"), |
| processor_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="HPSv3/"), |
| torch_dtype: torch.dtype = torch.bfloat16, |
| device: torch.device = get_device_type(), |
| score_index: int = 0, |
| use_special_tokens: bool = True, |
| max_pixels: int = 256 * 28 * 28, |
| min_pixels: int = 256 * 28 * 28, |
| 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_hpsv3") |
| processor_config.download_if_necessary() |
| processor = AutoProcessor.from_pretrained(processor_config.path, padding_side="right", **processor_kwargs) |
| if use_special_tokens: |
| special_tokens = ["<|Reward|>"] |
| processor.tokenizer.add_special_tokens({"additional_special_tokens": special_tokens}) |
| model.special_token_ids = processor.tokenizer.convert_tokens_to_ids(special_tokens) |
| model.reward_token = "special" |
| model.config.tokenizer_padding_side = processor.tokenizer.padding_side |
| model.config.pad_token_id = processor.tokenizer.pad_token_id |
| if hasattr(model.config, "text_config"): |
| model.config.text_config.pad_token_id = processor.tokenizer.pad_token_id |
| model.rm_head.to(torch.float32) |
| model = HPSv3Model( |
| model=model, |
| processor=processor, |
| use_special_tokens=use_special_tokens, |
| max_pixels=max_pixels, |
| min_pixels=min_pixels, |
| score_index=score_index, |
| ).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) |
|
|
| 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) |
|
|
