lot of things

load_models.py

@@ -2,10 +2,13 @@ import tensorflow as tf
 
 
 def load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=True):
-    if custom_cuda:
-        from stylegan2.generator import Generator
-    else:
-        from stylegan2_ref.generator import Generator
+    # if custom_cuda:
+    #     from stylegan2.generator import Generator
+    # else:
+    #     from stylegan2_ref.generator import Generator
+    from stylegan2.generator import ConvNeXtReverseConditional as Generator
+
+
 
     if g_params is None:
         g_params = {
@@ -23,7 +26,7 @@ def load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=T
     test_conditions = { 'pose_shape': tf.ones((1, 83), dtype=tf.float32), 'uv_static': tf.ones((1, 3, 512, 256), dtype=tf.float32), 'uv_mask': tf.ones((1, 512, 256), dtype=tf.float32), 'uv_shaped_body': tf.ones((1, 3, 512, 256), dtype=tf.float32) }
 
     # build generator model
-    generator = Generator(g_params)
+    generator = Generator()
     _ = generator([test_latent, test_conditions])
 
     if ckpt_dir is not None:
@@ -44,7 +47,8 @@ def load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=True):
     # else:
     #     from stylegan2_ref.discriminator import Discriminator
 
-    from custom_utils.graph_utils import GNNDiscriminator as Discriminator
+    # from custom_utils.graph_utils import GNNDiscriminator as Discriminator
+    from convnext.models.convnext_tf import ConvNeXtDiscriminator as Discriminator
 
 
     res = d_params['resolutions'][-1]
@@ -52,7 +56,8 @@ def load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=True):
     test_images = tf.ones((batch_size, 3, res, int(res/2)), dtype=tf.float32)
 
     # build discriminator model
-    discriminator = Discriminator(batch_size, 32)
+    # discriminator = Discriminator(batch_size, 32)
+    discriminator = Discriminator(depths=[3, 3, 27, 3], dims=[96, 192, 256, 8], include_top=False)
     _ = discriminator(test_images)
 
     if ckpt_dir is not None: