added skinning and gnn discriminator

load_models.py

@@ -38,23 +38,20 @@ def load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=T
 
 
 def load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=True):
-    if custom_cuda:
-        from stylegan2.discriminator import Discriminator
-    else:
-        from stylegan2_ref.discriminator import Discriminator
+    # if custom_cuda:
+    #     from stylegan2.discriminator import Discriminator
+    # else:
+    #     from stylegan2_ref.discriminator import Discriminator
+
+    from custom_utils.graph_utils import GNNDiscriminator as Discriminator
 
-    if d_params is None:
-        d_params = {
-            'labels_dim': 0,
-            'resolutions': [4, 8, 16, 32, 64, 128, 256, 512, 1024],
-            'featuremaps': [512, 512, 512, 512, 512, 256, 128, 64, 32],
-        }
 
     res = d_params['resolutions'][-1]
-    test_images = tf.ones((1, 3, res, int(res/2)), dtype=tf.float32)
+    batch_size = d_params['batch_size']
+    test_images = tf.ones((batch_size, 3, res, int(res/2)), dtype=tf.float32)
 
     # build discriminator model
-    discriminator = Discriminator(d_params)
+    discriminator = Discriminator(batch_size, 32)
     _ = discriminator(test_images)
 
     if ckpt_dir is not None: