only load encoders of AEs

configs/config_cvae.py

@@ -1,7 +1,20 @@
 from configs.config_base import config
 
-config['lr'] = 0.0001
+config['gpu'] = '3'
 config['name'] = 'cvae_19'
+config['notes'] = '\
+offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \
+retrained body encoder with deoriented bodies \
+use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\
+fixed gender issue and retrained body encoder, FINAL\
+remove shape from unposed garment\
+training on 100 garments with offsets over the average of these\
+fixed autoencoders!!!!!!!!!\
+cvae only loads encoders instead AEs\
+'
+
+
+config['lr'] = 0.0001
 config['starting_balancing_coeff'] = 0.0001
 
 config['signature_dynamic'] = { 'id': (),
@@ -19,9 +32,9 @@ config['signature_static'] = { 'uv_static': (512, 256, 3),
                               }
 
 config['avg_garment_path'] = 'custom_data/train_100_no_shape_avg_uv.npy'
-config['static_model'] = 'checkpoints/ae_static_6_copy'
+config['static_model'] = 'checkpoints/ae_static_7_2_x'
-config['body_model'] = 'checkpoints/ae_body_8_copy'
+config['body_model'] = 'checkpoints/ae_body_9_x'
-config['curvature_model'] = 'checkpoints/ae_curvature_7_copy'
+config['curvature_model'] = 'checkpoints/ae_curv_9_x'
 config['freeze_encoders'] = True
 
 config['cvae'] = dict()

cvae.py

@@ -1,5 +1,6 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+from configs.config_cvae import config
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
 
 gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -35,24 +36,16 @@ from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
 from models.cvae_model import CVAE
-from configs.config_cvae import config
 
 
 def init_wandb():
     run = wandb.init(project='vae', name=config['name'],
-            notes='\
+            notes=config['notes'])
-offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \
-retrained body encoder with deoriented bodies \
-use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\
-fixed gender issue and retrained body encoder, FINAL\
-remove shape from unposed garment\
-training on 100 garments with offsets over the average of these\
-fixed autoencoders!!!!!!!!!')
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae.py')
     artifact.add_file('models/cvae_model.py')
     artifact.add_file('losses.py')
-    artifact.add_file('configs/config.py')
+    artifact.add_file('configs/config_base.py')
     artifact.add_file('configs/config_cvae.py')
     artifact.add_file('models/base.py')
     artifact.add_file('models/cond_encoder.py')
@@ -76,7 +69,7 @@ def kl_loss(mean, logvar):
     return kl_loss
 
 
-# @tf.function
+@tf.function
 def train_step(batch):
 
     with tf.GradientTape() as tape:
@@ -321,7 +314,7 @@ def sample(folder):
 
 
 if __name__ == '__main__':
-    # init_wandb()
+    init_wandb()
     allow_memory_growth()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)

models/cvae_model.py

@@ -8,19 +8,7 @@ import scipy as sp
 # import plotly.express as ex
 
 from custom_utils.custom_utils import load_model
-from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
+from models.base import AEEncoder, VariationalEncoder, Decoder
-from models.base import Bottleneck, VariationalEncoder, Decoder
-from models.cond_encoder import AutoEncoder
-# from models.ae_body_model import AutoEncoder as BodyAE
-# from models.ae_curv_model import AutoEncoder as CurvatureAE
-# from models.ae_static_model import AutoEncoder as StaticAE
-# from models.ae_body_model import AutoEncoder32 as BodyAE32
-# from models.ae_curv_model import AutoEncoder32 as CurvatureAE32
-# from models.ae_static_model import AutoEncoder32 as StaticAE32
-# from models.ae_body_model import AutoEncoder128 as BodyAE128
-# from models.ae_curv_model import AutoEncoder128 as CurvatureAE128
-# from models.ae_static_model import AutoEncoder128 as StaticAE128
-
 
 
 class CVAE(tf.keras.Model):
@@ -30,28 +18,23 @@ class CVAE(tf.keras.Model):
         self.latent_dim = config['latent_dim']
         self.size = config['resolution']
 
-        model : AutoEncoder = AutoEncoder(config['cond_ae'])
+        self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
         if config['static_model'] != '':
-            load_model(model, config['static_model'])
+            load_model(self.static_enc, config['static_model'] + '_enc')
-        # self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
-        self.static_enc = model.encoder
 
-        model : AutoEncoder = AutoEncoder(config['cond_ae'])
+        self.body_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
         if config['body_model'] != '':
-            load_model(model, config['body_model'])
+            load_model(self.body_enc, config['body_model'] + '_enc')
-        self.body_enc = model.encoder
 
-        model : AutoEncoder = AutoEncoder(config['cond_ae'])
+        self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
         if config['curvature_model'] != '':
-            load_model(model, config['curvature_model'])
+            load_model(self.curv_enc, config['curvature_model'] + '_enc')
-        self.curv_enc = model.encoder
+
-        
         if config['freeze_encoders']:
             self.static_enc.trainable = False
             self.body_enc.trainable = False
             self.curv_enc.trainable = False
 
-
         self.encoder = VariationalEncoder(config['cvae']['latent_dim'], config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
         self.decoder = Decoder(config['cvae']['decoder_input'], config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'])
 
@@ -65,7 +48,6 @@ class CVAE(tf.keras.Model):
                                                  tf.keras.layers.Dense(128, activation='relu'),
                                                  tf.keras.layers.Dense(config['latent_dim']),])
 
-        # self.shuffle = tf.keras.layers.Lambda(lambda x: tf.stack([tf.gather(x[i], tf.random.shuffle(tf.range(x[i].shape[0]))) for i in range(x.shape[0])]))
 
     # @tf.function
     def sample_single(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False):
@@ -153,7 +135,7 @@ class CVAE(tf.keras.Model):
         return z_cond
 
 
-    def call(self, inputs, training=None, use_mask=False):
+    def call(self, inputs, training=None, use_mask=True):
 
         if use_mask:
             if self.size == 512: