|
"""
|
|
This file defines the core research contribution
|
|
"""
|
|
import matplotlib
|
|
matplotlib.use('Agg')
|
|
import math
|
|
|
|
import torch
|
|
from torch import nn
|
|
from model.encoder.encoders import psp_encoders
|
|
from model.stylegan.model import Generator
|
|
|
|
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)
|
|
|
|
self.opts.n_styles = int(math.log(self.opts.output_size, 2)) * 2 - 2
|
|
|
|
self.encoder = self.set_encoder()
|
|
self.decoder = Generator(self.opts.output_size, 512, 8)
|
|
self.face_pool = torch.nn.AdaptiveAvgPool2d((256, 256))
|
|
|
|
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:
|
|
pass
|
|
'''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, z_plus_latent=False, return_z_plus_latent=True):
|
|
if input_code:
|
|
codes = x
|
|
else:
|
|
codes = self.encoder(x)
|
|
|
|
|
|
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
|
|
if z_plus_latent:
|
|
input_is_latent = False
|
|
images, result_latent = self.decoder([codes],
|
|
input_is_latent=input_is_latent,
|
|
randomize_noise=randomize_noise,
|
|
return_latents=return_latents,
|
|
z_plus_latent=z_plus_latent)
|
|
|
|
if resize:
|
|
images = self.face_pool(images)
|
|
|
|
if return_latents:
|
|
if z_plus_latent and return_z_plus_latent:
|
|
return images, codes
|
|
if z_plus_latent and not return_z_plus_latent:
|
|
return images, result_latent
|
|
else:
|
|
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
|
|
|
