File size: 2,090 Bytes
4e2a1b3 | 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 61 62 63 64 | import torch
from ..core import ModelConfig
from ..core.device.npu_compatible_device import get_device_type
from ..models.lpips import LPIPSModel, LPIPS_NET_CHOICES, LPIPSCompute
from .base import Metric
_LPIPS_DEFAULT_FILES = {
"alex": "LPIPS/alexnet.safetensors",
"vgg": "LPIPS/vgg.safetensors",
"squeeze": "LPIPS/squeezenet.safetensors",
}
_LPIPS_MODEL_NAMES = {
"alex": "image_metrics_lpips_alex",
"vgg": "image_metrics_lpips_vgg",
"squeeze": "image_metrics_lpips_squeeze",
}
class LPIPSMetric(Metric):
def __init__(self, model: LPIPSCompute):
super().__init__()
self.model = model
@classmethod
def from_pretrained(
cls,
net: str = "alex",
model_config: ModelConfig = None,
device: torch.device = get_device_type(),
batch_size: int = 16,
target_size: int = 512,
vram_limit: float = None,
):
if net not in LPIPS_NET_CHOICES:
raise ValueError(f"net must be one of {LPIPS_NET_CHOICES}, got {net!r}")
if model_config is None:
model_config = ModelConfig(
model_id="DiffSynth-Studio/ImageMetrics",
origin_file_pattern=_LPIPS_DEFAULT_FILES[net],
)
model_pool = cls.download_and_load_models([model_config], torch_dtype=torch.float32, device=device, vram_limit=vram_limit)
backbone = model_pool.fetch_model(_LPIPS_MODEL_NAMES[net])
if backbone is None:
raise RuntimeError(
f"Failed to load LPIPS model for net={net!r}. The provided weights do not match the registered hash for {_LPIPS_MODEL_NAMES[net]}."
)
compute_model = LPIPSCompute(
model=backbone,
device=device,
batch_size=batch_size,
target_size=target_size,
)
return cls(compute_model)
@torch.no_grad()
def compute(self, image_a, image_b) -> float:
return self.model.compute(image_a, image_b)
def forward(self, image_a, image_b):
return self.compute(image_a, image_b)
|