Spaces:
Runtime error
Runtime error
File size: 5,071 Bytes
d9778ff |
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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 |
"""
This file defines the core research contribution
"""
import matplotlib
matplotlib.use("Agg")
import math
import torch
from torch import nn
from pixel2style2pixel.models.encoders import psp_encoders
from pixel2style2pixel.models.stylegan2.model import Generator
from pixel2style2pixel.configs.paths_config import model_paths
def get_keys(d, name):
if "state_dict" in d:
d = d["state_dict"]
d_filt = {k[len(name) + 1 :]: v for k, v in d.items() if k[: len(name)] == name}
return d_filt
class pSp(nn.Module):
def __init__(self, opts):
super(pSp, self).__init__()
self.set_opts(opts)
# compute number of style inputs based on the output resolution
self.opts.n_styles = int(math.log(self.opts.output_size, 2)) * 2 - 2
# Define architecture
self.encoder = self.set_encoder()
self.decoder = Generator(self.opts.output_size, 512, 8)
self.face_pool = torch.nn.AdaptiveAvgPool2d((256, 256))
# Load weights if needed
self.load_weights()
def set_encoder(self):
if self.opts.encoder_type == "GradualStyleEncoder":
encoder = psp_encoders.GradualStyleEncoder(50, "ir_se", self.opts)
elif self.opts.encoder_type == "BackboneEncoderUsingLastLayerIntoW":
encoder = psp_encoders.BackboneEncoderUsingLastLayerIntoW(
50, "ir_se", self.opts
)
elif self.opts.encoder_type == "BackboneEncoderUsingLastLayerIntoWPlus":
encoder = psp_encoders.BackboneEncoderUsingLastLayerIntoWPlus(
50, "ir_se", self.opts
)
else:
raise Exception("{} is not a valid encoders".format(self.opts.encoder_type))
return encoder
def load_weights(self):
if self.opts.checkpoint_path is not None:
print("Loading pSp from checkpoint: {}".format(self.opts.checkpoint_path))
ckpt = torch.load(self.opts.checkpoint_path, map_location="cpu")
self.encoder.load_state_dict(get_keys(ckpt, "encoder"), strict=True)
self.decoder.load_state_dict(get_keys(ckpt, "decoder"), strict=True)
self.__load_latent_avg(ckpt)
else:
print("Loading encoders weights from irse50!")
encoder_ckpt = torch.load(model_paths["ir_se50"])
# if input to encoder is not an RGB image, do not load the input layer weights
if self.opts.label_nc != 0:
encoder_ckpt = {
k: v for k, v in encoder_ckpt.items() if "input_layer" not in k
}
self.encoder.load_state_dict(encoder_ckpt, strict=False)
print("Loading decoder weights from pretrained!")
ckpt = torch.load(self.opts.stylegan_weights)
self.decoder.load_state_dict(ckpt["g_ema"], strict=False)
if self.opts.learn_in_w:
self.__load_latent_avg(ckpt, repeat=1)
else:
self.__load_latent_avg(ckpt, repeat=self.opts.n_styles)
def forward(
self,
x,
resize=True,
latent_mask=None,
input_code=False,
randomize_noise=True,
inject_latent=None,
return_latents=False,
alpha=None,
):
if input_code:
codes = x
else:
codes = self.encoder(x)
# normalize with respect to the center of an average face
if self.opts.start_from_latent_avg:
if self.opts.learn_in_w:
codes = codes + self.latent_avg.repeat(codes.shape[0], 1)
else:
codes = codes + self.latent_avg.repeat(codes.shape[0], 1, 1)
if latent_mask is not None:
for i in latent_mask:
if inject_latent is not None:
if alpha is not None:
codes[:, i] = (
alpha * inject_latent[:, i] + (1 - alpha) * codes[:, i]
)
else:
codes[:, i] = inject_latent[:, i]
else:
codes[:, i] = 0
input_is_latent = not input_code
images, result_latent = self.decoder(
[codes],
input_is_latent=input_is_latent,
randomize_noise=randomize_noise,
return_latents=return_latents,
)
if resize:
images = self.face_pool(images)
if return_latents:
return images, result_latent
else:
return images
def set_opts(self, opts):
self.opts = opts
def __load_latent_avg(self, ckpt, repeat=None):
if "latent_avg" in ckpt:
self.latent_avg = ckpt["latent_avg"].to(self.opts.device)
if repeat is not None:
self.latent_avg = self.latent_avg.repeat(repeat, 1)
else:
self.latent_avg = None
|