Create vtoonify/model/encoder/psp.py

vtoonify/model/encoder/psp.py

@@ -0,0 +1,127 @@
+
+"""
+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)
+		# 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:
+			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)
+			#print(codes.shape)
+			# 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
+		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
+