209 lines
8.7 KiB
Python
209 lines
8.7 KiB
Python
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 |