update main script

configs/config_cvae.py

@@ -1,7 +1,7 @@
 from configs.config_base import config
 
-config['gpu'] = '4'
+config['gpu'] = '0'
-config['name'] = 'v2_4'
+config['name'] = 'v2_7_4_3'
 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 | \
@@ -12,6 +12,13 @@ training on 100 garments with offsets over the average of these | \
 fixed autoencoders!!!!!!!!! | \
 cvae only loads encoders instead AEs | \
 increase balancing coeff sooner | \
+use 3d mesh loss | \
+add final gc layer | \
+starting from v2_5_2, add tanh to cond_enc last layer | \
+NOT using one-hot skinning | \
+using new average garment with one-hot versions | \
+fix collision for plotting | \
+use 3d mesh loss | \
 '
 config['dataset_in_memory'] = False
 
@@ -21,18 +28,23 @@ config['dataset_in_memory'] = False
 # DON'T FORGET TO CHECK IF STORE IN MEMORY
 # DON'T FORGET TO CHECK REG LOSS COEFF
 # DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
 
+config['pretrained_model'] = 'checkpoints/v2_5_2'
 
 config['resolution'] = 512
 config['is_skirt'] = False
 config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['l2_coeff'] = 1
 
 if config['resolution'] != 512:
   config['notes'] += '\
 using custom enhancer for upscaling, using not filled versions\
 '
     
-config['lr'] = 0.0001
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
 config['starting_balancing_coeff'] = 0.0001
 
 config['signature_dynamic'] = { 'id': (),
@@ -58,7 +70,7 @@ if config['resolution'] != 512:
   config['signature_static']['uv_mask' + res_modifier] = (config['resolution'], config['resolution']//2)
 
 if config['resolution'] == 512:
-  config['avg_garment_path'] = 'custom_data/train_no_shape_avg_uv.npy'
+  config['avg_garment_path'] = 'custom_data/train_no_shape_avg_uv.npy' # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
 else:
   config['avg_garment_path'] = ''
 
@@ -67,10 +79,13 @@ config['body_model'] = 'checkpoints/body_1'
 config['curvature_model'] = 'checkpoints/curv_1'
 config['upscale_enhancer'] = ''
 config['freeze_encoders'] = True
+config['use_one_hot'] = False
 
 config['cvae'] = dict()
 config['cvae']['is_label'] = False
 config['cvae']['latent_dim'] = config['latent_dim']
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.025
 
 if config['resolution'] == 512:
   config['cvae']['kernel_size'] = 7

configs/config_cvae_test.py

@@ -0,0 +1,101 @@
+config = dict()
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
+
+config['batch_size'] = 4
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['use_filled'] = False
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+
+config['gpu'] = '1'
+config['name'] = 'ab_1'
+config['notes'] = ''
+config['pretrained_model'] = 'checkpoints/v2_5_2'
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['l2_coeff'] = 1
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+
+config['signature_dynamic'] = { 'id': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                          'mesh_unposed': (14475, 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (14475,),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['avg_garment_path'] = 'custom_data/train_no_shape_avg_uv.npy' # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+
+
+config['static_model'] = 'checkpoints/static_1'
+config['body_model'] = 'checkpoints/body_1'
+config['curvature_model'] = 'checkpoints/curv_1'
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.025
+
+config['cvae']['kernel_size'] = 7
+config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = [256, 128, 64, 16]
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

cvae.py

@@ -1,5 +1,7 @@
 import os
-from configs.config_cvae import config
+import sys
+from configs.config_cvae_test import config
+
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
 
@@ -18,23 +20,23 @@ if gpus:
 import time
 import argparse
 import wandb
+import tqdm
 import random
 import numpy as np
-import scipy as sp
-from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces
 import plotly.graph_objects as go
 import plotly.express as ex
 from PIL import Image
+from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
+from postprocessing import fix_collision
 
 if config['is_skirt']:
     from losses import edge_loss_skirt as edge_loss
 else:
     from losses import edge_loss
 
-from utils import str_to_bool
+
-from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
+
-from load_models import load_generator, load_discriminator
+from losses import collision_ratio, kl_loss #, concept_loss
-from losses import uv_mask_loss, bend_loss, kl_loss
 # from losses import collision_penalty as collision_penalty
 from losses import collision_loss as collision_penalty
 from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d
@@ -44,7 +46,7 @@ from models.cvae_model import CVAE
 
 
 def init_wandb():
-    run = wandb.init(project='cvae', name=config['name'], notes=config['notes'])
+    run = wandb.init(project='cvae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae.py')
     artifact.add_file('models/cvae_model.py')
@@ -55,15 +57,12 @@ def init_wandb():
     artifact.add_file('models/cond_encoder.py')
     artifact.add_file('dataset.py')
     artifact.add_file('convnext/models/convnext_tf.py')
+    config_file_name = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    assert len(config_file_name) == 1
+    artifact.add_file('/'.join(config_file_name[0].split('.')) + '.py')
     run.log_artifact(artifact)
 
-uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'])
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False)
-
-
-def log_normal_pdf(sample, mean, logvar, raxis=1):
-    log2pi = tf.math.log(2. * np.pi)
-    res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis)
-    return res
 
 
 def kl_loss(mean, logvar):
@@ -94,23 +93,27 @@ def train_step(batch):
         e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
         c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
         l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
+        l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
+        # conc_loss = concept_loss(pred, batch['uv_unposed'])
         kl = kl_loss(mean, logvar)
 
-        reconstr_loss = l2_loss + (e_loss + c_loss) * config['reg_loss_coeff']
+        reconstr_loss = (l2_loss + l2_mesh_loss) * config['l2_coeff'] + (e_loss + c_loss) * config['reg_loss_coeff']
         loss = reconstr_loss + g_balancing_coeff * kl
 
         mean = tf.reduce_mean(mean)
         logvar = tf.reduce_mean(logvar)
 
-    ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
+    # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
     metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    # c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
     gradients = tape.gradient(loss, model.trainable_variables)
     optimizer.apply_gradients(zip(gradients, model.trainable_variables))
-    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
+    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric,
-             'train/reconstr_loss': reconstr_loss,
+             'train/reconstr_loss': reconstr_loss, 
              'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
-             'train/l2_loss': l2_loss }
+             'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss 
-        
+             # 'train/ms_ssim': ms_ssim, 'train/collision_ratio': c_ratio,
+            }
 
 
 @tf.function
@@ -132,9 +135,54 @@ def valid_step(batch):
     e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
     c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
     l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
+    l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
+    # conc_loss = concept_loss(pred, batch['uv_unposed'])
     kl = kl_loss(mean, logvar)
 
-    reconstr_loss = l2_loss + (e_loss + c_loss) * config['reg_loss_coeff']
+    reconstr_loss = (l2_loss + l2_mesh_loss) * config['l2_coeff'] + (e_loss + c_loss) * config['reg_loss_coeff']
+    loss = reconstr_loss + g_balancing_coeff * kl
+
+    mean = tf.reduce_mean(mean)
+    logvar = tf.reduce_mean(logvar)
+
+    # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
+    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    # c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
+
+    return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/euclid_error': metric, 
+             'valid/reconstr_loss': reconstr_loss, 'valid/l2_mesh_loss': l2_mesh_loss,
+             'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar,
+             # 'valid/ms_ssim': ms_ssim, 'valid/collision_ratio': c_ratio,
+              }
+
+
+@tf.function
+def test_step(batch):
+    pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
+
+    mesh_pred = uv2mesh(pred)
+    mesh_orig = batch['mesh_unposed']
+
+    shape = batch['pose_shape'][:, 72:]
+    gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+    shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+
+    mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
+    mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
+    mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+    mesh_orig_shaped = mesh_orig + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
+    mesh_orig_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_orig_shaped)
+
+    # is the output already masked?
+
+    e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+    c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
+    l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
+    l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
+    # conc_loss = concept_loss(pred, batch['uv_unposed'])
+    kl = kl_loss(mean, logvar)
+
+    reconstr_loss = (l2_loss + l2_mesh_loss) * config['l2_coeff'] + (e_loss + c_loss) * config['reg_loss_coeff']
     loss = reconstr_loss + g_balancing_coeff * kl
 
     mean = tf.reduce_mean(mean)
@@ -142,10 +190,14 @@ def valid_step(batch):
 
     ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
     metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    euclid_posed = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_posed - mesh_pred_posed) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    euclid_posed_fixed = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_posed - mesh_pred_posed_fixed) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
 
-    return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim,
+    return { 'test/loss': loss, 'test/edge_loss': e_loss, 'test/collision_loss': c_loss, 'test/euclid_error': metric, 'test/ms_ssim': ms_ssim,
-             'train/reconstr_loss': reconstr_loss,
+             'test/reconstr_loss': reconstr_loss, 'test/collision_ratio': c_ratio, 'test/l2_mesh_loss': l2_mesh_loss, 
-             'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar }
+             'test/mean': mean, 'test/logvar': logvar, 'test/kl_loss': kl, 'test/var': np.e ** logvar,
+             'test/euclid_posed': euclid_posed, 'test/euclid_posed_fixed': euclid_posed_fixed }
 
 
 def plot_everything(batch: dict, type: str):
@@ -175,18 +227,23 @@ def plot_everything(batch: dict, type: str):
     for i in range(body.shape[0]):
         mesh_pred[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
     mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
+    mesh_posed_fixed = fix_collision(mesh_posed, body, mask).numpy()
 
     garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
     collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
     collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))]
     collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
 
+    garment_posed_fixed = [mesh_posed_fixed[i, mask[i], :] for i in range(body.shape[0])]
+    collision_fixed = [np.concatenate((body[i], garment_posed_fixed[i]), axis=0) for i in range(len(garment_posed))]
+
     figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
-    
+    figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision[i], collision_fixed[i], collision_faces[i], collision_color[i]) for i in range(4) }
+
     mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0), num_bins=model.latent_dim)
     logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0), num_bins=model.latent_dim)
 
-    return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
+    return { **figs, **figs_coll, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
 
 
 def sample(folder):
@@ -219,8 +276,6 @@ def sample(folder):
         mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
 
 
-
-
         meshes = []
         meshes.append(go.Mesh3d(
                     # 8 vertices of a cube
@@ -314,7 +369,10 @@ def sample(folder):
 
 
 def test():
-    load_model(model, 'checkpoints/' + config['name'].split('_')[0])
+    str_split = config['name'].split('_')
+    test_idx = str_split.index('test')
+    checkpoint_name = '_'.join(str_split[:test_idx])
+    load_model(model, 'checkpoints/' + checkpoint_name)
 
     dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False, do_repeat=False)
     static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False)
@@ -322,23 +380,20 @@ def test():
     step = 0
     losses = []
 
-    for batch in dataset_test:
+    for batch in tqdm.tqdm(dataset_test):
         batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
 
-        test_loss = valid_step(batch)
+        test_loss = test_step(batch)
         losses.append(test_loss)
 
         to_log = { **{ 'current_step': step }, **test_loss }
-        if step % 100 == 0:
+        if step % 1000 == 0:
             figs = plot_everything(batch, 'test')
             to_log = { **to_log, **figs }
 
         wandb.log(to_log, step = step)
-
-        if step > 30:
-            break
         
-        print(step)
+        # print(step)
         step += 1
    
     # averages = dict()
@@ -362,20 +417,19 @@ def test():
 
 
 if __name__ == '__main__':
-    init_wandb()
+    # init_wandb()
-    allow_memory_growth()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())
 
-    smpl = SMPL()
+    smpl = SMPL(use_one_hot=config['use_one_hot'])
     skirt_modifier = '_skirt' if config['is_skirt'] else ''
 
     factory = DatasetFactory()
 
-    dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'])
+    dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], use_one_hot=config['use_one_hot'])
     static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True)
-    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'])
+    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], use_one_hot=config['use_one_hot'])
     static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True)
 
     model = CVAE(config)
@@ -385,43 +439,46 @@ if __name__ == '__main__':
     g_balancing_coeff = config['starting_balancing_coeff']
 
     # sample('checkpoints/cvae_13')
-    test()
+    # test()
 
-    # ckpt = tf.train.Checkpoint(model=model)
+    ckpt = tf.train.Checkpoint(model=model)
-    # manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20)
+    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20)
 
-    # valid_iterator = iter(dataset_valid)
+    if config['pretrained_model'] != '':
-    # step = 0
+        load_model(model, config['pretrained_model'])
-    # for batch in dataset_train:
-    #     batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
 
-    #     train_loss = train_step(batch)
+    valid_iterator = iter(dataset_valid)
+    step = 0
+    for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
 
-    #     data_valid = valid_iterator.get_next()
+        train_loss = train_step(batch)
-    #     data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
-    #     valid_loss = valid_step(data_valid)
 
-    #     to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
+        data_valid = valid_iterator.get_next()
-    #     if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
+        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
-    #         train_figs = plot_everything(batch, 'train')
+        valid_loss = valid_step(data_valid)
-    #         valid_figs = plot_everything(data_valid, 'valid')
-    #         to_log = { **to_log, **train_figs, **valid_figs }
 
-    #     if step % 10000 == 0:
+        to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
-    #         manager.save(checkpoint_number=step)
+        if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
-    #     if step == 0:
+            train_figs = plot_everything(batch, 'train')
-    #         to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+            valid_figs = plot_everything(data_valid, 'valid')
+            to_log = { **to_log, **train_figs, **valid_figs }
 
-    #     if train_loss['train/euclid_error'] < 0.01:
+        if step % 10000 == 0:
-    #         if balancing_activated_step == 0:
+            manager.save(checkpoint_number=step)
-    #             balancing_activated_step = step
+        if step == 0:
-    #         g_balancing_coeff = g_balancing_coeff + 1 / 15000      
+            to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
-    #     to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } }
+        if train_loss['train/euclid_error'] < 0.019:
+            if balancing_activated_step == 0:
+                balancing_activated_step = step
+            g_balancing_coeff = g_balancing_coeff + 1 / 10000      
+
+        to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } }
         
-    #     wandb.log(to_log, step = step)
+        wandb.log(to_log, step = step)
 
-    #     print(step)
+        print(step)
 
-    #     step += 1
+        step += 1