import tensorflow as tf import numpy as np from stylegan2.utils import lerp from stylegan2.layers.mapping_block import Mapping from stylegan2.layers.synthesis_block import Synthesis from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D, GlobalAveragePooling2D, Layer, LayerNormalization, Flatten, Reshape) # from stylegan2.layers.mapping_block import Mapping from stylegan2.layers.embedding import PoseShapeEmbedding, GarmentEmbedding from convnext.models.convnext_tf import Block, DropPath class Generator(tf.keras.Model): def __init__(self, g_params, **kwargs): super(Generator, self).__init__(**kwargs) self.z_dim = g_params['z_dim'] self.w_dim = g_params['w_dim'] self.labels_dim = g_params['labels_dim'] self.n_mapping = g_params['n_mapping'] self.resolutions = g_params['resolutions'] self.featuremaps = g_params['featuremaps'] self.w_ema_decay = 0.995 self.style_mixing_prob = 0.9 self.zero_body = tf.ones((4,3,512,256)) * tf.transpose(g_params['zero_body'], [2,0,1])[tf.newaxis,:,:,:] self.n_broadcast = len(self.resolutions) * 2 self.mixing_layer_indices = tf.range(self.n_broadcast, dtype=tf.int32)[tf.newaxis, :, tf.newaxis] self.g_mapping = Mapping(self.w_dim, self.labels_dim, self.n_mapping, name='g_mapping') self.broadcast = tf.keras.layers.Lambda(lambda x: tf.tile(x[:, tf.newaxis], [1, self.n_broadcast, 1])) self.synthesis = Synthesis(self.resolutions, self.featuremaps, name='g_synthesis') def build(self, input_shape): # w_avg self.w_avg = tf.Variable(tf.zeros(shape=[self.w_dim], dtype=tf.dtypes.float32), name='w_avg', trainable=False, synchronization=tf.VariableSynchronization.ON_READ, aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA) @tf.function def set_as_moving_average_of(self, src_net): beta, beta_nontrainable = 0.99, 0.0 for cw, sw in zip(self.weights, src_net.weights): assert sw.shape == cw.shape # print('{} <=> {}'.format(cw.name, sw.name)) if 'w_avg' in cw.name: cw.assign(lerp(sw, cw, beta_nontrainable)) else: cw.assign(lerp(sw, cw, beta)) return def update_moving_average_of_w(self, w_broadcasted): # compute average of current w batch_avg = tf.reduce_mean(w_broadcasted[:, 0], axis=0) # compute moving average of w and update(assign) w_avg self.w_avg.assign(lerp(batch_avg, self.w_avg, self.w_ema_decay)) return def style_mixing_regularization(self, latents1, conditions, w_broadcasted1): # get another w and broadcast it latents2 = tf.random.normal(shape=tf.shape(latents1), dtype=tf.dtypes.float32) dlatents2 = self.g_mapping([latents2, conditions]) w_broadcasted2 = self.broadcast(dlatents2) # find mixing limit index # mixing_cutoff_index = tf.cond( # pred=tf.less(tf.random.uniform([], 0.0, 1.0), self.style_mixing_prob), # true_fn=lambda: tf.random.uniform([], 1, self.n_broadcast, dtype=tf.dtypes.int32), # false_fn=lambda: tf.constant(self.n_broadcast, dtype=tf.dtypes.int32)) if tf.random.uniform([], 0.0, 1.0) < self.style_mixing_prob: mixing_cutoff_index = tf.random.uniform([], 1, self.n_broadcast, dtype=tf.dtypes.int32) else: mixing_cutoff_index = tf.constant(self.n_broadcast, dtype=tf.dtypes.int32) # mix it mixed_w_broadcasted = tf.where( condition=tf.broadcast_to(self.mixing_layer_indices < mixing_cutoff_index, tf.shape(w_broadcasted1)), x=w_broadcasted1, y=w_broadcasted2) return mixed_w_broadcasted def truncation_trick(self, w_broadcasted, truncation_psi, truncation_cutoff=None): ones = tf.ones_like(self.mixing_layer_indices, dtype=tf.float32) tpsi = ones * truncation_psi if truncation_cutoff is None: truncation_coefs = tpsi else: indices = tf.range(self.n_broadcast) truncation_coefs = tf.where(condition=tf.less(indices, truncation_cutoff), x=tpsi, y=ones) truncated_w_broadcasted = lerp(self.w_avg, w_broadcasted, truncation_coefs) return truncated_w_broadcasted def call(self, inputs, ret_w_broadcasted=False, truncation_psi=1.0, truncation_cutoff=None, training=None, mask=None): latents, conditions = inputs dlatents = self.g_mapping([latents, conditions]) w_broadcasted = self.broadcast(dlatents) if training: self.update_moving_average_of_w(w_broadcasted) w_broadcasted = self.style_mixing_regularization(latents, conditions, w_broadcasted) if not training: w_broadcasted = self.truncation_trick(w_broadcasted, truncation_psi, truncation_cutoff) image_out = self.synthesis(w_broadcasted) # image_out = image_out + conditions[2] image_out = image_out + conditions['uv_shaped_body'] image_out = image_out * conditions['uv_mask'][:,tf.newaxis,:,:] if ret_w_broadcasted: return image_out, w_broadcasted else: return image_out def compute_output_shape(self, input_shape): return input_shape[0][0], 3, self.resolutions[-1], self.resolutions[-1] class ConvNeXtReverseConditional(tf.keras.Model): r""" ConvNeXt A Tensorflow keras impl of : `A ConvNet for the 2020s` - https://arxiv.org/pdf/2201.03545.pdf Args: num_classes (int): Number of classes for classification head. Default: 1000 depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3] dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768] include_top (bool): whether to add head or just use it as feature extractor. Default: True drop_path_rate (float): Stochastic depth rate. Default: 0. layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1. """ def __init__(self, depths=[3, 3, 9, 3], dims=[64, 64, 32, 4], drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., ): super().__init__() self.pose_shape_embedding = PoseShapeEmbedding() self.garment_embedding = GarmentEmbedding((128 * 3) - (216 + 11)) self.reshape = Reshape((16, 8, 4)) self.upsample_layers = [] # stem and 3 intermediate downsampling conv layers stem = tf.keras.Sequential([ Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'), LayerNormalization(epsilon=1e-6)] ) self.upsample_layers.append(stem) for i in range(3): downsample_layer = tf.keras.Sequential([ LayerNormalization(epsilon=1e-6), Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')] ) self.upsample_layers.append(downsample_layer) self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))] cur = 0 for i in range(4): stage = tf.keras.Sequential( [Block(dim=dims[i], drop_path=dp_rates[cur + j], layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])] ) self.stages.append(stage) cur += depths[i] self.last = Dense(3, activation='tanh') def call(self, inputs, ret_w_broadcasted=False, truncation_psi=1.0, truncation_cutoff=None,): latents, conditions = inputs pose_shape_latent = self.pose_shape_embedding([conditions['pose_shape'][:, :72], conditions['pose_shape'][:, 72:]]) garment_latent = self.garment_embedding(conditions['uv_static']) x = tf.concat([latents, pose_shape_latent, garment_latent], axis=1) x = self.reshape(x) for i in reversed(range(4)): x = self.upsample_layers[i](x) x = self.stages[i](x) x = self.last(x) x = tf.transpose(x, [0,3,1,2]) x = x + conditions['uv_shaped_body'] x = x * conditions['uv_mask'][:,tf.newaxis,:,:] return x def set_as_moving_average_of(self, src_net): return def update_moving_average_of_w(self, w_broadcasted): return def style_mixing_regularization(self, latents1, conditions, w_broadcasted1): return None