ai-art

  1"""Stripped version of https://github.com/richzhang/PerceptualSimilarity/tree/master/models"""
  2
  3import os
  4
  5import torch
  6import torch.nn as nn
  7from collections import namedtuple
  8
  9from core.utils.loader import download
 10from core.taming.utils import normalize_tensor, spatial_average, load_vgg
 11
 12
 13class LPIPS(nn.Module):
 14    # Learned perceptual metric
 15    def __init__(self, model_dir="/models", use_dropout=True):
 16        super().__init__()
 17        self.scaling_layer = ScalingLayer()
 18        self.chns = [64, 128, 256, 512, 512]  # vg16 features
 19        self.net = VGG16(model_dir=model_dir, pretrained=True, requires_grad=False)
 20        self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
 21        self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
 22        self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
 23        self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
 24        self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
 25        self.load_from_pretrained(model_dir)
 26        for param in self.parameters():
 27            param.requires_grad = False
 28
 29    def load_from_pretrained(self, model_dir="/models"):
 30        ckpt = f"{model_dir}/vgg.pth"
 31        if not os.path.exists(ckpt):
 32            download("https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1", ckpt)
 33        self.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
 34        print(f"Loaded pretrained LPIPS loss from '{ckpt}'")
 35
 36    def forward(self, input, target):
 37        in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target))
 38        outs0, outs1 = self.net(in0_input), self.net(in1_input)
 39        feats0, feats1, diffs = {}, {}, {}
 40        lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
 41        for kk in range(len(self.chns)):
 42            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
 43            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2
 44
 45        res = [spatial_average(lins[kk].model(diffs[kk]), keepdim=True) for kk in range(len(self.chns))]
 46        val = res[0]
 47        for l in range(1, len(self.chns)):
 48            val += res[l]
 49        return val
 50
 51
 52class ScalingLayer(nn.Module):
 53    def __init__(self):
 54        super(ScalingLayer, self).__init__()
 55        self.register_buffer('shift', torch.Tensor([-.030, -.088, -.188])[None, :, None, None])
 56        self.register_buffer('scale', torch.Tensor([.458, .448, .450])[None, :, None, None])
 57
 58    def forward(self, inp):
 59        return (inp - self.shift) / self.scale
 60
 61
 62class NetLinLayer(nn.Module):
 63    """ A single linear layer which does a 1x1 conv """
 64    def __init__(self, chn_in, chn_out=1, use_dropout=False):
 65        super(NetLinLayer, self).__init__()
 66        layers = [nn.Dropout(), ] if (use_dropout) else []
 67        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False), ]
 68        self.model = nn.Sequential(*layers)
 69
 70
 71class VGG16(torch.nn.Module):
 72    def __init__(self, model_dir="/models", requires_grad=False, pretrained=True):
 73        super(VGG16, self).__init__()
 74        vgg_pretrained_features = load_vgg(model_dir=model_dir, pretrained=pretrained).features
 75        self.slice1 = torch.nn.Sequential()
 76        self.slice2 = torch.nn.Sequential()
 77        self.slice3 = torch.nn.Sequential()
 78        self.slice4 = torch.nn.Sequential()
 79        self.slice5 = torch.nn.Sequential()
 80        self.N_slices = 5
 81        for x in range(4):
 82            self.slice1.add_module(str(x), vgg_pretrained_features[x])
 83        for x in range(4, 9):
 84            self.slice2.add_module(str(x), vgg_pretrained_features[x])
 85        for x in range(9, 16):
 86            self.slice3.add_module(str(x), vgg_pretrained_features[x])
 87        for x in range(16, 23):
 88            self.slice4.add_module(str(x), vgg_pretrained_features[x])
 89        for x in range(23, 30):
 90            self.slice5.add_module(str(x), vgg_pretrained_features[x])
 91        if not requires_grad:
 92            for param in self.parameters():
 93                param.requires_grad = False
 94
 95    def forward(self, X):
 96        h = self.slice1(X)
 97        h_relu1_2 = h
 98        h = self.slice2(h)
 99        h_relu2_2 = h
100        h = self.slice3(h)
101        h_relu3_3 = h
102        h = self.slice4(h)
103        h_relu4_3 = h
104        h = self.slice5(h)
105        h_relu5_3 = h
106        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3'])
107        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
108        return out