lot of things

stylegan2/generator.py

@@ -1,10 +1,19 @@
 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)
@@ -115,3 +124,86 @@ class Generator(tf.keras.Model):
 
     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