diff --git a/auto_encoder.py b/auto_encoder.py
index 70d76df..d5fd861 100644
--- a/auto_encoder.py
+++ b/auto_encoder.py
@@ -21,7 +21,7 @@ from stylegan2_ae.generator import Generator
 from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
 from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
 from dataset import DatasetFactory, StaticDataset
-from custom_utils.custom_utils import batch_to_dict
+from custom_utils.custom_utils import batch_to_dict, plot_3d
 from custom_utils.smpl import SMPL
 
 name='ae_body_1'
@@ -113,53 +113,6 @@ def filter_resolutions_featuremaps(resolutions, featuremaps, res):
 #     return fig
 
 
-def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
-    mesh_pred_unposed[:,1] += 1
-    collision[:,1] += 2
-    fig = go.Figure(data=[
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=mesh_gt[:,0],
-            y=mesh_gt[:,1],
-            z=mesh_gt[:,2],
-            # i, j and k give the vertices of triangles
-            i = faces[:,0],
-            j = faces[:,1],
-            k = faces[:,2]
-        ),
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=mesh_pred_unposed[:,0],
-            y=mesh_pred_unposed[:,1],
-            z=mesh_pred_unposed[:,2],
-            # i, j and k give the vertices of triangles
-            i = faces[:,0],
-            j = faces[:,1],
-            k = faces[:,2],
-        ),
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=collision[:,0],
-            y=collision[:,1],
-            z=collision[:,2],
-            # i, j and k give the vertices of triangles
-            i = collision_faces[:,0],
-            j = collision_faces[:,1],
-            k = collision_faces[:,2],
-            vertexcolor=collision_color
-        )
-    ])
-
-    fig.update_layout(scene_aspectmode='manual',
-                    scene_aspectratio=dict(x=1, y=2, z=1))
-    fig.update_layout(
-        scene = dict(
-            xaxis = dict(nticks=4, range=[-1,1],),
-                        yaxis = dict(nticks=4, range=[-1,3],),
-                        zaxis = dict(nticks=4, range=[-1,1],),)
-                        )
-
-    return fig
 
 class Bottleneck(tf.keras.Model):
     
diff --git a/auto_encoder_body.py b/auto_encoder_body.py
index 775acc4..dadd84f 100644
--- a/auto_encoder_body.py
+++ b/auto_encoder_body.py
@@ -1,5 +1,6 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+from configs.config_ae_body import config
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 
 import time
 import argparse
@@ -8,26 +9,20 @@ import random
 import cv2
 import numpy as np
 import tensorflow as tf
-from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
+from custom_utils.custom_utils import UV2Mesh, load_model, maskFacesBatch
 import plotly.graph_objects as go
 import plotly.express as ex
 
 from tf_utils import check_tf_version, allow_memory_growth
 from losses import edge_loss, uv_mask_loss
 from dataset import DatasetFactory, StaticDataset
-from custom_utils.custom_utils import batch_to_dict
-from custom_utils.smpl import SMPL
-from models.ae_body_model import AutoEncoder, AutoEncoder128, AutoEncoder32
+from custom_utils.custom_utils import batch_to_dict, plot_3d
+from models.cond_encoder import AutoEncoder
 
 
-name='ae_body_8'
 def init_wandb():
-    run = wandb.init(project='vae', name=name,
-            notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \
-    fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \
-    . lowered learning rate because it goes to nan. shallow third stage. slower lr decay. with validation. higher lr rate.\
-    two more downscaling than up, so decoder output is resized wto final resolution\
-    larger batch_size.\
+    run = wandb.init(project='vae', name=config['name'],
+            notes='convnext + convnext, add dense with tanh at the end of reverse convnext,\
     flattened bottleneck with 128 latent_dim\ remove global orientation\
     fixed gender issue with deoriented bodies, again, aaand again\
     trained on 100 garments\
@@ -36,7 +31,10 @@ def init_wandb():
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('auto_encoder_body.py')
     artifact.add_file('losses.py')
-    artifact.add_file('models/ae_body_model.py')
+    artifact.add_file('configs/config.py')
+    artifact.add_file('configs/config_ae_body.py')
+    artifact.add_file('models/base.py')
+    artifact.add_file('models/cond_encoder.py')
     artifact.add_file('dataset.py')
     artifact.add_file('convnext/models/convnext_tf.py')
     run.log_artifact(artifact)
@@ -44,90 +42,24 @@ def init_wandb():
 uv2mesh = UV2Mesh(res=256)
 
 
-
-def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
-    mesh_pred_unposed[:,1] += 1
-    collision[:,1] += 2
-    fig = go.Figure(data=[
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=mesh_gt[:,0],
-            y=mesh_gt[:,1],
-            z=mesh_gt[:,2],
-            # i, j and k give the vertices of triangles
-            i = faces[:,0],
-            j = faces[:,1],
-            k = faces[:,2]
-        ),
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=mesh_pred_unposed[:,0],
-            y=mesh_pred_unposed[:,1],
-            z=mesh_pred_unposed[:,2],
-            # i, j and k give the vertices of triangles
-            i = faces[:,0],
-            j = faces[:,1],
-            k = faces[:,2],
-        ),
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=collision[:,0],
-            y=collision[:,1],
-            z=collision[:,2],
-            # i, j and k give the vertices of triangles
-            i = collision_faces[:,0],
-            j = collision_faces[:,1],
-            k = collision_faces[:,2],
-            vertexcolor=collision_color
-        )
-    ])
-
-    fig.update_layout(scene_aspectmode='manual',
-                    scene_aspectratio=dict(x=1, y=2, z=1))
-    fig.update_layout(
-        scene = dict(
-            xaxis = dict(nticks=4, range=[-1,1],),
-                        yaxis = dict(nticks=4, range=[-1,3],),
-                        zaxis = dict(nticks=4, range=[-1,1],),)
-                        )
-
-    return fig
-
-
 @tf.function
 def train_step(batch):
-
     with tf.GradientTape() as tape:
-
-        pred = model(batch['uv_body_posed' + resolution_modifier])
-        # pred = (pred + zero_body) * batch['mesh_mask'][:,:,tf.newaxis]
+        pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body])
 
         mesh_pred = uv2mesh(pred)
-        # mesh_pred = pred
         mesh_orig = batch['mesh_body_posed']
 
+        e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones((pred.shape[0], 14475))))
+        l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed' + resolution_modifier] - pred))), axis=1))
+        ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_body_posed' + resolution_modifier] + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
 
-        e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones_like(batch['mesh_mask'])))
-        c_loss = 0 # 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
-        # l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['mesh_unposed'] - pred)), axis=-1) * batch['mesh_mask'])
-        l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed' + resolution_modifier] - pred) * uv_mask_body[tf.newaxis,tf.newaxis,:,:])), axis=1))
-        # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed' + resolution_modifier] * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, max_val=2))
-        ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed' + resolution_modifier] * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
-        fft_loss = 0 # tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_body_posed'], tf.zeros_like(batch['uv_body_posed'])))), axis=[1,2,3]))
-        # loss = fft_loss_pc # + e_loss + c_loss
         loss = 1e3 * (1 - ms_ssim)
 
-    mesh_orig_2 = uv2mesh(batch['uv_body_posed' + resolution_modifier])
-    # mesh_skinned = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
     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'])
-    metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
-    euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
     gradients = tape.gradient(loss, model.trainable_variables)
-    # gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
-    # gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
     optimizer.apply_gradients(zip(gradients, model.trainable_variables))
-    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
-             'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/euclid_error_posed': euclid_posed }
+    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, }
         
 
 
@@ -155,35 +87,27 @@ def valid_step(batch):
 
 def plot_everything(batch: dict, type: str):
 
-    pred = model(batch['uv_body_posed' + resolution_modifier])
+    pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body])
     
     mesh_pred = uv2mesh(pred)
     mesh_orig = batch['mesh_body_posed']
     
-    orig_img = tf.transpose(batch['uv_body_posed' + resolution_modifier], [0,2,3,1])
-    pred_img = tf.transpose(pred, [0,2,3,1])
+    orig_img = batch['uv_body_posed' + resolution_modifier]
+    pred_img = pred
     img = tf.concat((orig_img, pred_img), axis=2)
     imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
 
     F = np.load("custom_utils/smpl_faces.npy")
-    mask = batch['mesh_mask'].numpy().astype(bool)
-    # garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
     garment_unposed, faces = mesh_pred.numpy(), F
-    # gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
     gt = [mesh_orig[i].numpy()[:, :] for i in range(pred.shape[0])]
     
     body = uv2mesh(batch['uv_body_posed' + resolution_modifier]).numpy()
-    # body = batch['mesh_body_posed'].numpy()
-    # mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
-    # garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
     garment_posed = garment_unposed
     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_faces = [np.concatenate((F, faces + 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))]
 
     figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces, collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
-    # 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) }
     
     return { **figs, **{ type + '/images': imgs } }
 
@@ -194,11 +118,11 @@ def train():
 
     step = 0
     for batch in dataset_train:
-        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
         train_loss = train_step(batch)
 
         # data_valid = valid_iterator.get_next()
-        # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
+        # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
         valid_loss = {}#valid_step(data_valid)
 
         if step % 2500 == 0:
@@ -211,52 +135,37 @@ def train():
         if step % 10000 == 0:
             manager.save(checkpoint_number=step)
 
-
-
         print(step)
 
         step += 1
 
+
 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()
+
     uv_mask_body = np.load('custom_utils/mask_256x512.npy')
     uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
 
-    # uv_mask_body = tf.ones((128, 64))
+    # TODO: different resolutions not working yet, model should be ready though
     resolution_modifier = ''
-    resolutions = [512]
+    if config['resolution'] != 512:
+        resolution_modifier = '_' + str(config['resolution'])
 
     factory = DatasetFactory()
-    signature_dynamic = { 'id': (),
-                          'pose_shape': (72 + 11,),
-                          'uv_body_posed': (3, 512, 256,),
-                          'mesh_body_posed': (14475, 3,),
-                        }
-    signature_static = { 
-                         'mesh_mask': (14475,),
-                        #  'uv_mask': (512, 256,),
-                        }
 
-    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=32, signature=signature_dynamic, use_filled=False, resolutions=resolutions)
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=resolutions)
-    dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=32, signature=signature_dynamic)
-    static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
+    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']])
+    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']])
+    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'])
+    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'])
 
-    model = AutoEncoder()
-    # model = AutoEncoder32()
-    # model = AutoEncoder128()
-
-    optimizer = tf.optimizers.Adam(learning_rate=0.0001)
-    folder = 'models/ae_logs/'
+    model = AutoEncoder(config['cond_ae'])
+    optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
 
     ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
-    ckpt.restore(manager.latest_checkpoint)
+    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3)
 
     train()
\ No newline at end of file
diff --git a/auto_encoder_curv.py b/auto_encoder_curv.py
index 5def6ea..b7d9703 100644
--- a/auto_encoder_curv.py
+++ b/auto_encoder_curv.py
@@ -5,7 +5,7 @@ import wandb
 import random
 import numpy as np
 import tensorflow as tf
-os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
 
 from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
@@ -15,20 +15,22 @@ import plotly.express as ex
 from tf_utils import check_tf_version, allow_memory_growth
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.custom_utils import batch_to_dict
-from custom_utils.smpl import SMPL
-from models.ae_curv_model import AutoEncoder
+from models.cond_encoder import AutoEncoder
+from configs.config_ae_curv import config
 
 
-name='ae_curvature_x_2'
 testing = False
 def init_wandb():
     if not testing:
-        run = wandb.init(project='vae', name=name, notes='convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments. same as 7')
+        run = wandb.init(project='vae', name=config['name'], notes='convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments. same as 7')
         artifact = wandb.Artifact('code', type='code')
         artifact.add_file('auto_encoder_curv.py')
-        artifact.add_file('dataset.py')
         artifact.add_file('losses.py')
-        artifact.add_file('stylegan2_ae/layers/mapping_block.py')
+        artifact.add_file('configs/config.py')
+        artifact.add_file('configs/config_ae_curv.py')
+        artifact.add_file('models/base.py')
+        artifact.add_file('models/cond_encoder.py')
+        artifact.add_file('dataset.py')
         artifact.add_file('convnext/models/convnext_tf.py')
         run.log_artifact(artifact)
 
@@ -110,7 +112,7 @@ def dice_dense(prediction, ground_truth, weight_map=None):
     return tf.reduce_mean(dice_score)
 
 
-@tf.function
+# @tf.function
 def train_step(batch):
 
     with tf.GradientTape() as tape:
@@ -121,22 +123,14 @@ def train_step(batch):
 
         ce_uv_loss = scce(batch['uv_curvature'], pred)
         ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1))
-        # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)[:,:,:,tf.newaxis], batch['uv_curvature'][:,:,:,tf.newaxis], max_val=7*2))
-
-        # loss =  1e3 * (1 - ms_ssim)
         loss = ce_uv_loss
 
-    # mesh_pred = uv2mesh(tf.tile(tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)[:,:,:,tf.newaxis], [1,1,1,3]))[:,:,0]
-    # mesh_orig = tf.one_hot(batch['mesh_curvature'] + tf.cast(batch['mesh_mask'], tf.int32), 7)
-    # ce_mesh_loss = scce(mesh_orig, mesh_pred)
-
     mean_iou = get_mean_iou(batch['uv_curvature'], tf.math.argmax(pred, axis=-1))
     dice = dice_dense(pred, uv_curvature)
 
     gradients = tape.gradient(loss, model.trainable_variables)
     optimizer.apply_gradients(zip(gradients, model.trainable_variables))
-    return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/ssim': ms_ssim, 'train/mean_iou': mean_iou, 'train/dice': dice  #'ce_mesh_loss': ce_mesh_loss
-             }
+    return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/ssim': ms_ssim, 'train/mean_iou': mean_iou, 'train/dice': dice }
         
 
 
@@ -157,11 +151,10 @@ def valid_step(batch):
              }
 
 
-
 def train():
     step = 0
     for batch in dataset_train:
-        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
 
         train_loss = train_step(batch)
 
@@ -194,7 +187,7 @@ def train():
 def test():
     ckpt.restore(manager.latest_checkpoint)
     for batch in dataset_train:
-        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
         uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
 
         pred = model([uv_curvature, batch['uv_mask']])
@@ -208,57 +201,31 @@ 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()
+
     scce = tf.keras.losses.SparseCategoricalCrossentropy()
 
+    resolution_modifier = ''
+    if config['resolution'] != 512:
+        resolution_modifier = '_' + str(config['resolution'])
+
     factory = DatasetFactory()
-    signature_dynamic = { 'id': (),
-                        #   'pose_shape': (72 + 11,),
-                        #   'uv_posed': (3, 512, 256,),
-                        #   'uv_body_posed': (3, 512, 256,),
-                        #   'mesh_body_posed': (14475, 3,),
-                        #   'mesh_posed': (14475, 3,),
-                        #   'mesh_unposed': (14475, 3,),
-                        #   'uv_unposed': (3, 512, 256,)
-                          'uv_curvature': (512, 256),
-                          'mesh_curvature':  (14475)
-                        }
-    signature_static = { #'uv_static': (3, 512, 256,),
-                        #  'mesh_static': (14475, 3,),
-                         'mesh_mask': (14475,),
-                         'uv_mask': (512, 256,),
-                        #  'uv_shaped_body': (3, 512, 256,),
-                        }
 
-    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100',
-                                  batch_size=8,
-                                  signature=signature_dynamic,
-                                  use_filled=False)
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100',
-                                   signature=signature_static,
-                                   use_filled=False)
-    dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
-                                  batch_size=8,
-                                  signature=signature_dynamic)
-    static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
-                                   signature=signature_static)
+    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']])
+    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']])
+    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'])
+    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'])
 
-    valid_iterator = iter(dataset_valid)
-
-    model = AutoEncoder()
-    # optimizer = tf.optimizers.Adam(learning_rate=0.002, beta_1=0.0, beta_2=0.99, epsilon=1e-8)
-    lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.001, decay_steps=10000, decay_rate=0.1, staircase=False)
-    optimizer = tf.optimizers.Adam(learning_rate=lr_schedule)
-    folder = 'models/ae_logs/'
+    model = AutoEncoder(config['cond_ae'])
+    optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
 
     ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=5)
+    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3)
+
 
     if testing:
         test()
diff --git a/auto_encoder_static.py b/auto_encoder_static.py
index 3027962..cc636cd 100644
--- a/auto_encoder_static.py
+++ b/auto_encoder_static.py
@@ -1,5 +1,7 @@
-import imp
 import os
+from configs.config_ae_static import config
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
+
 import time
 import argparse
 import wandb
@@ -13,24 +15,22 @@ import plotly.express as ex
 from tf_utils import check_tf_version, allow_memory_growth
 from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss
 from dataset import DatasetFactory, StaticDataset
-from custom_utils.custom_utils import batch_to_dict
-from custom_utils.smpl import SMPL
-from models.ae_static_model import AutoEncoder
+from custom_utils.custom_utils import batch_to_dict, plot_3d
+from models.cond_encoder import AutoEncoder
 
 
-name='ae_static_5'
 def init_wandb():
-    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-    run = wandb.init(project='vae', name=name,
-            notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \
-    fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \
-    . lowered learning rate because it goes to nan. shallow third stage. slower lr decay. with validation. higher lr rate.\
-    convnext first stage down/upscale by factor fo 2 instead of 4, and 5 stages instead of 4 so bottlneck same size\
+    run = wandb.init(project='vae', name=config['name'],
+            notes='convnext + convnext, add dense with tanh at the end of reverse convnext,\
     flattened latent dim to 128\
-    trai on 100 garments')
+    train on 100 garments')
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('auto_encoder_static.py')
     artifact.add_file('losses.py')
+    artifact.add_file('configs/config_base.py')
+    artifact.add_file('configs/config_ae_static.py')
+    artifact.add_file('models/base.py')
+    artifact.add_file('models/cond_encoder.py')
     artifact.add_file('dataset.py')
     artifact.add_file('convnext/models/convnext_tf.py')
     run.log_artifact(artifact)
@@ -39,86 +39,24 @@ uv2mesh = UV2Mesh(res=256)
 
 
 
-def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
-    mesh_pred_unposed[:,1] += 1
-    collision[:,1] += 2
-    fig = go.Figure(data=[
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=mesh_gt[:,0],
-            y=mesh_gt[:,1],
-            z=mesh_gt[:,2],
-            # i, j and k give the vertices of triangles
-            i = faces[:,0],
-            j = faces[:,1],
-            k = faces[:,2]
-        ),
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=mesh_pred_unposed[:,0],
-            y=mesh_pred_unposed[:,1],
-            z=mesh_pred_unposed[:,2],
-            # i, j and k give the vertices of triangles
-            i = faces[:,0],
-            j = faces[:,1],
-            k = faces[:,2],
-        ),
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=collision[:,0],
-            y=collision[:,1],
-            z=collision[:,2],
-            # i, j and k give the vertices of triangles
-            i = collision_faces[:,0],
-            j = collision_faces[:,1],
-            k = collision_faces[:,2],
-            vertexcolor=collision_color
-        )
-    ])
-
-    fig.update_layout(scene_aspectmode='manual',
-                    scene_aspectratio=dict(x=1, y=2, z=1))
-    fig.update_layout(
-        scene = dict(
-            xaxis = dict(nticks=4, range=[-1,1],),
-                        yaxis = dict(nticks=4, range=[-1,3],),
-                        zaxis = dict(nticks=4, range=[-1,1],),)
-                        )
-
-    return fig
-
-
-
 @tf.function
 def train_step(batch):
 
     with tf.GradientTape() as tape:
-        pred = model(batch['uv_static'])
-
-        uv_mask = batch['uv_mask']
+        pred = model([batch['uv_static'], batch['uv_mask']])
 
         mesh_pred = uv2mesh(pred)
         mesh_orig = batch['mesh_static']
 
         e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-        c_loss = 0 #1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
-        # l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['mesh_unposed'] - pred)), axis=-1) * batch['mesh_mask'])
-        l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1) * uv_mask)
-        ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) * uv_mask[:,:,:,tf.newaxis] + 1, tf.transpose(batch['uv_static'], [0,2,3,1]) * uv_mask[:,:,:,tf.newaxis] + 1, max_val=2))
-        fft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_static'], tf.zeros_like(batch['uv_static'])))), axis=[1,2,3]))
+        l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1))
+        ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_static'] + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
         loss = 1e3 * (1 - ms_ssim)
 
-    mesh_orig_2 = uv2mesh(batch['uv_static'])
-    # mesh_skinned = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
     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'])
-    metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
-    euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
     gradients = tape.gradient(loss, model.trainable_variables)
-    # gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
-    # gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
     optimizer.apply_gradients(zip(gradients, model.trainable_variables))
-    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
-             'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/euclid_error_posed': euclid_posed }
+    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, }
         
 
 
@@ -145,14 +83,13 @@ def valid_step(batch):
 
 
 def plot_everything(batch: dict, type: str):
-
-    pred = model(batch['uv_static'])
+    pred = model([batch['uv_static'], batch['uv_mask']])
     
     mesh_pred = uv2mesh(pred)
     mesh_orig = batch['mesh_static']
     
-    orig_img = tf.transpose(batch['uv_static'], [0,2,3,1])
-    pred_img = tf.transpose(pred, [0,2,3,1])
+    orig_img = batch['uv_static']
+    pred_img = pred
     img = tf.concat((orig_img, pred_img), axis=2)
     imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
 
@@ -175,20 +112,23 @@ def plot_everything(batch: dict, type: str):
 def train():
     step = 0
     for batch in dataset_train:
-        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
         train_loss = train_step(batch)
 
         # data_valid = valid_iterator.get_next()
         # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
         valid_loss = {}#valid_step(data_valid)
 
-        if step % 1000 == 0:
+        if step % 2500 == 0:
             train_figs = plot_everything(batch, 'train')
             valid_figs = {}#plot_everything(data_valid, 'valid')
             wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step)
         else:
             wandb.log({ **train_loss, **valid_loss }, step=step)
             
+        if step % 10000 == 0:
+            manager_enc.save(checkpoint_number=step)
+            manager_dec.save(checkpoint_number=step)
 
         print(step)
 
@@ -201,41 +141,25 @@ if __name__ == '__main__':
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())
     
-    smpl = SMPL()
+    # TODO: different resolutions not working yet, model should be ready though
 
-    resolution_modifier = '_32'
-    resolutions = [32]
+    resolution_modifier = ''
+    if config['resolution'] != 512:
+        resolution_modifier = '_' + str(config['resolution'])
 
     factory = DatasetFactory()
-    signature_dynamic = { 'id': (),
-                        #   'pose_shape': (72 + 11,),
-                        #   'uv_posed': (3, 512, 256,),
-                        #   'uv_body_posed': (3, 512, 256,),
-                        #   'mesh_body_posed': (14475, 3,),
-                        #   'mesh_posed': (14475, 3,),
-                        #   'mesh_unposed': (14475, 3,),
-                        #   'uv_unposed': (3, 512, 256,)
-                        }
-    signature_static = { 'uv_static': (3, 512, 256,),
-                         'mesh_static': (14475, 3,),
-                         'mesh_mask': (14475,),
-                         'uv_mask': (512, 256,),
-                        #  'uv_shaped_body': (3, 512, 256,),
-                        }
 
-    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=8, signature=signature_dynamic)
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static)
-    dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=8, signature=signature_dynamic)
-    static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
+    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']])
+    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']])
+    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'])
+    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'])
 
+    model : AutoEncoder = AutoEncoder(config['cond_ae'])
+    optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
 
-    model = AutoEncoder()
-    lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0005, decay_steps=30000, decay_rate=0.1, staircase=False)
-    optimizer = tf.optimizers.Adam(learning_rate=lr_schedule)
-    folder = 'models/ae_logs/'
-    valid_iterator = iter(dataset_valid)
-
-    ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=5)
+    ckpt_enc = tf.train.Checkpoint(model=model.encoder)
+    manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=3)
+    ckpt_dec = tf.train.Checkpoint(model=model.decoder)
+    manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=3)
 
     train()
\ No newline at end of file
diff --git a/configs/config_ae_body.py b/configs/config_ae_body.py
new file mode 100644
index 0000000..89b2aa5
--- /dev/null
+++ b/configs/config_ae_body.py
@@ -0,0 +1,14 @@
+from configs.config_base import config
+
+config['gpu'] = '1'
+config['lr'] = 0.0001
+config['name'] = 'ae_body_9'
+
+config['signature_dynamic'] = { 'id': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (3, 512, 256,),
+                          'mesh_body_posed': (14475, 3,),
+                        }
+config['signature_static'] = { 
+                         'mesh_mask': (14475,),
+                        }
\ No newline at end of file
diff --git a/configs/config_ae_curv.py b/configs/config_ae_curv.py
new file mode 100644
index 0000000..d00fe0c
--- /dev/null
+++ b/configs/config_ae_curv.py
@@ -0,0 +1,15 @@
+from configs.config_base import config
+
+config['lr'] = 0.001
+config['name'] = 'ae_body_9'
+
+config['cond_ae']['is_label'] = True
+
+config['signature_dynamic'] = { 'id': (),
+                          'uv_curvature': (512, 256),
+                          'mesh_curvature':  (14475)
+                        }
+config['signature_static'] = { 'mesh_mask': (14475,),
+                         'uv_mask': (512, 256,),
+                        }
+
diff --git a/configs/config_ae_static.py b/configs/config_ae_static.py
new file mode 100644
index 0000000..2a1e15c
--- /dev/null
+++ b/configs/config_ae_static.py
@@ -0,0 +1,12 @@
+from configs.config_base import config
+
+config['gpu'] = '0'
+config['lr'] = 0.0005
+config['name'] = 'ae_static_7'
+
+config['signature_dynamic'] = { 'id': (), }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                         'mesh_static': (14475, 3,),
+                         'mesh_mask': (14475,),
+                         'uv_mask': (512, 256,),
+                        }
\ No newline at end of file
diff --git a/configs/config_base.py b/configs/config_base.py
new file mode 100644
index 0000000..60d6e13
--- /dev/null
+++ b/configs/config_base.py
@@ -0,0 +1,21 @@
+config = dict()
+
+
+config['batch_size'] = 4
+config['train_dir'] = '/home/workspace/local-data/train_100'
+config['valid_dir'] = '/home/workspace/local-data/train_100'
+config['use_filled'] = False
+config['resolution'] = 512
+
+config['latent_dim'] = 128
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = 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]
\ No newline at end of file
diff --git a/configs/config_cvae.py b/configs/config_cvae.py
new file mode 100644
index 0000000..56d84d1
--- /dev/null
+++ b/configs/config_cvae.py
@@ -0,0 +1,53 @@
+from configs.config_base import config
+
+config['lr'] = 0.0001
+config['name'] = 'cvae_19'
+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),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (14475,),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['avg_garment_path'] = 'custom_data/train_100_no_shape_avg_uv.npy'
+config['static_model'] = 'checkpoints/ae_static_6_copy'
+config['body_model'] = 'checkpoints/ae_body_8_copy'
+config['curvature_model'] = 'checkpoints/ae_curvature_7_copy'
+config['freeze_encoders'] = True
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
+
+if config['resolution'] == 512:
+  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]
+elif config['resolution'] == 128:
+  config['cvae']['kernel_size'] = 5
+  config['cvae']['encoder_depths'] = [16, 16, 9, 9]
+  config['cvae']['encoder_dims'] = [64, 128, 256, 256]
+  config['cvae']['decoder_input'] = (4, 2, 128)
+  config['cvae']['decoder_depths'] = [16, 16, 9, 9]
+  config['cvae']['decoder_dims'] = [256, 256, 128, 64]
+elif config['resolution'] == 32:
+  config['cvae']['kernel_size'] = 3
+  config['cvae']['encoder_depths'] = [27, 27]
+  config['cvae']['encoder_dims'] = [256, 256]
+  config['cvae']['decoder_input'] = (4, 2, 256)
+  config['cvae']['decoder_depths'] = [27, 27]
+  config['cvae']['decoder_dims'] = [256, 256]
+else:
+  raise NotImplementedError
diff --git a/convnext/models/convnext_tf.py b/convnext/models/convnext_tf.py
index e2781c9..f413ea7 100644
--- a/convnext/models/convnext_tf.py
+++ b/convnext/models/convnext_tf.py
@@ -104,16 +104,18 @@ class ConvNeXt(tf.keras.Model):
         super().__init__()
         self.include_top = include_top
         self.downsample_layers = []  # stem and 3 intermediate downsampling conv layers
-        stem = tf.keras.Sequential([
-            Conv2D(dims[0], kernel_size=4, strides=4, padding='same'),
-            LayerNormalization(epsilon=1e-6)]
-        )
+        # stem = tf.keras.Sequential([
+        #     Conv2D(dims[0], kernel_size=4, strides=4, padding='same'),
+        #     LayerNormalization(epsilon=1e-6)]
+        # )
+        stem = [Conv2D(dims[0], kernel_size=4, strides=4, padding='same'), LayerNormalization(epsilon=1e-6)]
         self.downsample_layers.append(stem)
         for i in range(len(depths) - 1):
-            downsample_layer = tf.keras.Sequential([
-                LayerNormalization(epsilon=1e-6),
-                Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
-            )
+            # downsample_layer = tf.keras.Sequential([
+            #     LayerNormalization(epsilon=1e-6),
+            #     Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
+            # )
+            downsample_layer = [LayerNormalization(epsilon=1e-6), Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
             self.downsample_layers.append(downsample_layer)
 
         self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
@@ -138,7 +140,8 @@ class ConvNeXt(tf.keras.Model):
 
     def forward_features(self, x):
         for i in range(len(self.downsample_layers)):
-            x = self.downsample_layers[i](x)
+            x = self.downsample_layers[i][0](x)
+            x = self.downsample_layers[i][1](x)
             x = self.stages[i](x)
         return x
 
diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py
index 887385c..30482b8 100644
--- a/custom_utils/custom_utils.py
+++ b/custom_utils/custom_utils.py
@@ -1,3 +1,4 @@
+from faulthandler import disable
 import numpy as np
 import tensorflow as tf
 import plotly.graph_objects as go
@@ -135,14 +136,14 @@ def maskFacesBatch(V, F, mask):
 
 
 def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
-    mesh_pred_unposed[:,1] += 1
-    collision[:,1] += 2
+    mesh_pred_unposed[:,2] += 1
+    collision[:,2] += 2
     fig = go.Figure(data=[
         go.Mesh3d(
             # 8 vertices of a cube
             x=mesh_gt[:,0],
-            y=mesh_gt[:,1],
-            z=mesh_gt[:,2],
+            y=mesh_gt[:,2],
+            z=mesh_gt[:,1],
             # i, j and k give the vertices of triangles
             i = faces[:,0],
             j = faces[:,1],
@@ -151,8 +152,8 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
         go.Mesh3d(
             # 8 vertices of a cube
             x=mesh_pred_unposed[:,0],
-            y=mesh_pred_unposed[:,1],
-            z=mesh_pred_unposed[:,2],
+            y=mesh_pred_unposed[:,2],
+            z=mesh_pred_unposed[:,1],
             # i, j and k give the vertices of triangles
             i = faces[:,0],
             j = faces[:,1],
@@ -161,8 +162,8 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
         go.Mesh3d(
             # 8 vertices of a cube
             x=collision[:,0],
-            y=collision[:,1],
-            z=collision[:,2],
+            y=collision[:,2],
+            z=collision[:,1],
             # i, j and k give the vertices of triangles
             i = collision_faces[:,0],
             j = collision_faces[:,1],
@@ -202,7 +203,10 @@ def sparse_to_tensor(L):
     return tf.SparseTensor(idx, data, shape)
 
 
-def load_model(model, checkpoint):
+def load_model(model, checkpoint, disable_assert=False):
     ckpt = tf.train.Checkpoint(model=model)
     manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
-    ckpt.restore(manager.latest_checkpoint)
+    if not disable_assert:
+        assert manager.latest_checkpoint is not None
+        assert manager.latest_checkpoint.split('/')[-1] != 'ckpt-0'
+    ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()
diff --git a/custom_utils/fill_uv.py b/custom_utils/fill_uv.py
index fee3e73..79ec71e 100644
--- a/custom_utils/fill_uv.py
+++ b/custom_utils/fill_uv.py
@@ -5,12 +5,13 @@ import numpy as np
 import plotly.graph_objects as go
 import plotly.express as ex
 
-
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+import tensorflow as tf
+import tensorflow_addons as tfa
 def fill(img, mask):
     coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
     img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0))
     img_data = img[img_idx[:,0], img_idx[:,1], :]
-
     interpolator =  spi.NearestNDInterpolator(x=img_idx, y=img_data)
     img_2 = interpolator(coords)
     res = np.zeros((512,256,3), dtype=np.float32)
@@ -19,6 +20,19 @@ def fill(img, mask):
 
     return res
 
+# @tf.function
+def fill_tf(img, mask):
+    img = img[tf.newaxis,:,:,:]
+    coords_int = tf.transpose(tf.stack(tf.experimental.numpy.nonzero(1-mask)), (1,0))
+    coords = tf.cast(coords_int, dtype=tf.float32)
+    img_2 = tfa.image.interpolate_bilinear(img, coords[tf.newaxis,:,:])
+    # coords_2 = list(zip(tf.ones(coords_int.shape[0]), coords_int[:,0], coords_int[:,1]))
+    # coords = tf.transpose(tf.stack([coords_int[:,0], coords_int[:,1]]))
+    out = tf.tensor_scatter_nd_update(img[0], coords_int, img_2[0])
+
+    return out
+
+
 
 def fill_label(img, mask):
     coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
@@ -44,7 +58,13 @@ def fill_uv_maps_posed_body(folder, type):
         id, _ = frame.split('_')
         # mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
         img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient']
-        filled = fill(img, mask)
+        # filled = fill(img, mask)
+        img = np.zeros((100, 100, 3), dtype=np.float32)
+        mask = np.zeros((100, 100), dtype=np.int32)
+        img[25:75, 25:75, :] = 1
+        mask[25:75, 25:75] = 1
+
+        filled = fill_tf(img, mask)
         np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_posed_body_no_orient_filled = filled)
 
 
diff --git a/cvae.py b/cvae.py
index 5c47d7e..30c82ba 100644
--- a/cvae.py
+++ b/cvae.py
@@ -1,5 +1,5 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "2"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 import tensorflow as tf
 
 gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -31,83 +31,37 @@ from load_models import load_generator, load_discriminator
 from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_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
+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
 
 
-name='cvae_17'
 def init_wandb():
-    run = wandb.init(project='vae', name=name,
-            notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline. \
-new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1 from 0.1. \
+    run = wandb.init(project='vae', name=config['name'],
+            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')
+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('dataset.py')
     artifact.add_file('losses.py')
+    artifact.add_file('configs/config.py')
+    artifact.add_file('configs/config_cvae.py')
+    artifact.add_file('models/base.py')
+    artifact.add_file('models/cond_encoder.py')
+    artifact.add_file('dataset.py')
     artifact.add_file('convnext/models/convnext_tf.py')
     run.log_artifact(artifact)
 
 uv2mesh = UV2Mesh(res=256)
 
-def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
-    mesh_pred_unposed[:,2] += 1
-    collision[:,2] += 2
-    fig = go.Figure(data=[
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=mesh_gt[:,0],
-            y=mesh_gt[:,2],
-            z=mesh_gt[:,1],
-            # i, j and k give the vertices of triangles
-            i = faces[:,0],
-            j = faces[:,1],
-            k = faces[:,2]
-        ),
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=mesh_pred_unposed[:,0],
-            y=mesh_pred_unposed[:,2],
-            z=mesh_pred_unposed[:,1],
-            # i, j and k give the vertices of triangles
-            i = faces[:,0],
-            j = faces[:,1],
-            k = faces[:,2],
-        ),
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=collision[:,0],
-            y=collision[:,2],
-            z=collision[:,1],
-            # i, j and k give the vertices of triangles
-            i = collision_faces[:,0],
-            j = collision_faces[:,1],
-            k = collision_faces[:,2],
-            vertexcolor=collision_color
-        )
-    ])
-
-    fig.update_layout(scene_aspectmode='manual',
-                    scene_aspectratio=dict(x=1, y=2, z=1))
-    fig.update_layout(
-        scene = dict(
-            xaxis = dict(nticks=4, range=[-1,1],),
-                        yaxis = dict(nticks=4, range=[-1,3],),
-                        zaxis = dict(nticks=4, range=[-1,1],),)
-                        )
-
-    return fig
-
-
-
 
 def log_normal_pdf(sample, mean, logvar, raxis=1):
     log2pi = tf.math.log(2. * np.pi)
@@ -122,13 +76,14 @@ def kl_loss(mean, logvar):
     return kl_loss
 
 
-@tf.function
+# @tf.function
 def train_step(batch):
 
     with tf.GradientTape() as tape:
         uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
 
-        pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], training=True, use_mask=True)
+        pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']],
+                                      training=True, use_mask=True)
 
         mesh_pred = uv2mesh(pred)
         mesh_orig = batch['mesh_unposed']
@@ -141,38 +96,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']))
-        l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
-        l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred) ** 2), axis=-1) * batch['uv_mask'], axis=[1, 2])
+        l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
+        l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1) * batch['uv_mask'], axis=[1, 2])
         # ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
-        ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))
+        ms_ssim = tf.image.ssim_multiscale(pred + 1, batch['uv_unposed'] + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))
         # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale((tf.transpose(pred, [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], (tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], max_val=2))
-        fft_loss = 0 #tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3])
-        tv_loss = tf.image.total_variation(pred)
 
         reconstr_loss = l2_loss
         reconstr_loss = tf.reduce_mean(reconstr_loss)
         kl = kl_loss(mean, logvar)
         loss = reconstr_loss + balancing_coeff * kl
-        # loss = reconstr_loss + balancing_coeff * kl
         ms_ssim = tf.reduce_mean(ms_ssim)
         l1_loss = tf.reduce_mean(l1_loss)
         l2_loss = tf.reduce_mean(l2_loss)
-        fft_loss = tf.reduce_mean(fft_loss)
-        tv_loss = tf.reduce_mean(tv_loss)
         mean = tf.reduce_mean(mean)
         logvar = tf.reduce_mean(logvar)
 
-    mesh_orig_2 = uv2mesh(batch['uv_unposed'])
     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'])
-    metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
-    euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
     gradients = tape.gradient(loss, model.trainable_variables)
-    # gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
-    # gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
     optimizer.apply_gradients(zip(gradients, model.trainable_variables))
     return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
-             'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2,
-             'train/euclid_error_posed': euclid_posed, 'train/reconstr_loss': reconstr_loss, 'train/tv_loss': tv_loss,
+             'train/reconstr_loss': reconstr_loss,
              'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/reconst_loss_norm': reconstr_loss / (2 * (gamma ** 2)), 'train/var': np.e ** logvar,
              'train/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss }
         
@@ -183,7 +127,7 @@ def valid_step(batch):
     uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
     pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
 
-    mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+    mesh_pred = uv2mesh(pred)
     mesh_orig = batch['mesh_unposed']
     e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
     c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
@@ -197,10 +141,6 @@ def valid_step(batch):
     metric_2 = 0 #tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
   
     reconstr_loss = l1_loss
-    # logpx_z = -reconstr_loss
-    # logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.)
-    # logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar)
-    # loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x))
     kl = kl_loss(mean, logvar)
     # loss = reconstr_loss / (2 * (gamma ** 2)) + kl
     loss = reconstr_loss + balancing_coeff * kl
@@ -224,9 +164,9 @@ def plot_everything(batch: dict, type: str):
     mesh_pred = uv2mesh(pred).numpy()
     mesh_orig = batch['mesh_unposed']
     
-    orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
-    uv_static_img = tf.transpose(batch['uv_static'], [0,2,3,1])
-    uv_body_posed_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1])
+    orig_img = batch['uv_unposed']
+    uv_static_img = batch['uv_static']
+    uv_body_posed_img = batch['uv_body_posed']
     pred_img = pred
     img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
     imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
@@ -259,8 +199,6 @@ def plot_everything(batch: dict, type: str):
     return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
 
 
-
-
 def sample(folder):
     # run = wandb.init(project='vae_plots', name=folder.split('/')[-1])
     load_model(model,folder)
@@ -274,7 +212,7 @@ def sample(folder):
         pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
         pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
 
-        mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+        mesh_pred = uv2mesh(pred)
         mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy()
         mesh_orig = batch['mesh_posed'][0].numpy()
         mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0]
@@ -286,9 +224,8 @@ def sample(folder):
         mask = batch['mesh_mask'].numpy().astype(bool)[0]
         _, faces = _maskFaces(np.ones((14475, 3)), F, mask)
 
-
         pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
-        mesh_pred_zero = uv2mesh(tf.transpose(pred_zero, [0,3,1,2]))
+        mesh_pred_zero = uv2mesh(pred_zero)
         mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
 
 
@@ -382,67 +319,42 @@ def sample(folder):
 
         print()
 
+
 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())
 
-    batch_size = 4
+    batch_size = config['batch_size']
     
     smpl = SMPL()
 
     factory = DatasetFactory()
-    signature_dynamic = { 'id': (),
-                          'pose_shape': (72 + 11,),
-                        #   'uv_posed': (3, 512, 256,),
-                          'uv_body_posed': (3, 512, 256,),
-                          'mesh_body_posed': (14475, 3,),
-                          'mesh_posed': (14475, 3,),
-                          'mesh_unposed': (14475, 3,),
-                          'uv_unposed': (3, 512, 256,),
-                          'uv_curvature': (512, 256),
-                        }
-    signature_static = { 'uv_static': (3, 512, 256,),
-                        #  'mesh_static': (14475, 3,),
-                         'mesh_mask': (14475,),
-                         'uv_mask': (512, 256,),
-                        #  'uv_shaped_body': (3, 512, 256,),
-                        }
 
-    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False)
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False)
-    dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=batch_size, signature=signature_dynamic)
-    static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
+    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']])
+    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']])
+    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'])
+    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'])
 
+    model = CVAE(config)
 
-    model = CVAE(avg_garment_path='custom_data/train_100_no_shape_avg_uv.npy',
-                 latent_dim=128, static_model='checkpoints/ae_static_4',
-                 body_model='checpoints/ae_body_8',
-                 curvature_model='checpoints/ae_curvature_7')
-
-
-    # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
-    optimizer = tf.optimizers.Adam(learning_rate=0.0001)
-    balancing_coeff = 0.0001
+    optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
+    balancing_coeff = config['starting_balancing_coeff']
     reg_coeff = 1
-    folder = 'models/ae_logs/'
     valid_iterator = iter(dataset_valid)
     gamma = 1
     balancing_activated_step = 0
-    # avg_jumpsuit = np.load('custom_data/jumpsuit_20_avg_uv.npy')
-
 
     # sample('checkpoints/cvae_13')
 
-
     ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
+    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3)
 
     step = 0
     for batch in dataset_train:
-        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
 
         train_loss = train_step(batch)
 
@@ -464,19 +376,7 @@ if __name__ == '__main__':
         if train_loss['train/euclid_error'] < 0.01:
             if balancing_activated_step == 0:
                 balancing_activated_step = step
-            balancing_coeff = balancing_coeff + 1 / 15000
-        # if train_loss['train/euclid_error'] < 0.015:
-        #     reg_coeff = reg_coeff + 100 / 15000
-        # if step > 15000 and balancing_coeff <= 1e-3:
-        #     balancing_coeff = 1e-2
-        # if step > 20000 and balancing_coeff <= 1e-3:
-        #     balancing_coeff = 1e-2
-        # if step > 25000 and balancing_coeff <= 1e-2:
-        #     balancing_coeff = 1e-1
-        # if step > 30000 and balancing_coeff <= 1e-1:
-        #     balancing_coeff = 1
-
-        # gamma = 0.99 * gamma + 0.01 * train_loss['train/reconstr_loss']        
+            balancing_coeff = balancing_coeff + 1 / 15000      
 
         to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } }
         
diff --git a/cvae_curv.py b/cvae_curv.py
index 3e23054..7a1b613 100644
--- a/cvae_curv.py
+++ b/cvae_curv.py
@@ -370,7 +370,7 @@ if __name__ == '__main__':
     static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
 
 
-    model = CVAE(latent_dim=128, static_model='checkpoints/ae_static_4',
+    model = CVAE(latent_dim=128, static_model='checkpoints/ae_static_5',
                  body_model='checpoints/ae_body_8',)
 
 
diff --git a/cvae_low.py b/cvae_low.py
index b27067e..d8ef11f 100644
--- a/cvae_low.py
+++ b/cvae_low.py
@@ -21,24 +21,22 @@ import random
 import numpy as np
 import scipy as sp
 from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
-from tensorboardX import SummaryWriter
 import plotly.graph_objects as go
 import plotly.express as ex
 from PIL import Image
 
-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 d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_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
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
-from models.cvae_model import CVAE, CVAE128, CVAE32
+from models.cvae_model import CVAE#, CVAE128, CVAE32
+from models.upscaler import UpscalerBasic as Upscaler
 
 
-name='cvae_low_1_d128'
+name='cvae_low_4_2_d32'
 def init_wandb():
     run = wandb.init(project='vae', name=name,
             notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline. \
@@ -49,7 +47,8 @@ use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\
 fixed gender issue and retrained body encoder\
 remove shape from unposed garment\
 training on 100 garments , absolute, not offsets\
-using filled versions')
+using custom enhancer for upscaling, using not filled versions\
+calculating loss at full res')
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae.py')
     artifact.add_file('models/cvae_model.py')
@@ -130,62 +129,57 @@ def train_step(batch):
     with tf.GradientTape() as tape:
         uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
         uv_curvature_low = tf.one_hot(batch['uv_curvature' + res_modifier], 7)
-        pred, mean, logvar, z = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], uv_curvature_low, batch['uv_static'], batch['uv_body_posed'], uv_curvature])
+        pred, mean, logvar, z = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], uv_curvature_low,
+                                       batch['uv_mask' + res_modifier], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], use_mask=True)
         # pred = pred * batch['uv_mask' + res_modifier][:,:,:,tf.newaxis]
-    
-        mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
-        mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
+
+        pred_enhanced = enhance(upscale_enhancer, pred, batch['uv_mask' + res_modifier], training=False)
+        pred_full_res = tf.image.resize(pred_enhanced, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
+        mesh_pred = uv2mesh(pred_enhanced)
         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_posed = mesh_pred_posed + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
+        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)
 
         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']))
-        # l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask' + res_modifier], axis=[1, 2])
-        l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) - pred)), axis=-1), axis=[1, 2])
-        # l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) - pred) ** 2), axis=-1) * batch['uv_mask' + res_modifier], axis=[1, 2])
-        l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) - pred) ** 2), axis=-1), axis=[1, 2])
-        # ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
-        # ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) + 1, max_val=3, power_factors=(0.6, 0.4)) #(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)
-        ms_ssim = tf.image.ssim(pred + 1, tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) + 1, max_val=3) #(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)
-        # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale((tf.transpose(pred, [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], (tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], max_val=2))
-        rfft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.rfft3d(pred) - tf.signal.rfft3d(batch['uv_unposed'])), axis=[1,2,3]))
-        fft_loss = tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]))))), axis=[1,2,3])
-        fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed'])
-        tv_loss = tf.image.total_variation(pred)
+        l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed' + res_modifier] - pred)), axis=-1), axis=[1, 2])
+        l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed' + res_modifier] - pred) ** 2), axis=-1), axis=[1, 2])
+        l2_loss_full = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred_full_res) ** 2), axis=-1), axis=[1, 2])
+        if resolution == 32:
+            ms_ssim = tf.image.ssim(pred + 1, batch['uv_unposed' + res_modifier] + 1, max_val=3) 
+            ms_ssim_full = tf.image.ssim(pred_full_res + 1, batch['uv_unposed'] + 1, max_val=3) 
+        else:
+            ms_ssim = tf.image.ssim_multiscale(pred + 1, batch['uv_unposed' + res_modifier] + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)) 
 
-        reconstr_loss = l2_loss
+        reconstr_loss = l2_loss_full
         reconstr_loss = tf.reduce_mean(reconstr_loss)
-        # logpx_z = -reconstr_loss
-        # logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.)
-        # logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar)
-        # loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x))
         kl = kl_loss(mean, logvar)
         loss = reconstr_loss + balancing_coeff * kl
-        # loss = reconstr_loss + balancing_coeff * kl
         ms_ssim = tf.reduce_mean(ms_ssim)
+        ms_ssim_full = tf.reduce_mean(ms_ssim_full)
         l1_loss = tf.reduce_mean(l1_loss)
         l2_loss = tf.reduce_mean(l2_loss)
-        fft_loss = tf.reduce_mean(fft_loss)
-        tv_loss = tf.reduce_mean(tv_loss)
+        l2_loss_full = tf.reduce_mean(l2_loss_full)
         mean = tf.reduce_mean(mean)
         logvar = tf.reduce_mean(logvar)
 
-    mesh_orig_2 = uv2mesh(batch['uv_unposed' + res_modifier])
+    # mesh_orig_2 = uv2mesh(batch['uv_unposed' + res_modifier])
     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'])
-    metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    # metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
     euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
     gradients = tape.gradient(loss, model.trainable_variables)
     # gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
     # gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
     optimizer.apply_gradients(zip(gradients, model.trainable_variables))
     return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
-             'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/rfft_loss': rfft_loss, 'train/fft_oc_low': fft_loss_pc_low, 'train/fft_pc_high': fft_loss_pc_high,
-             'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed, 'train/reconstr_loss': reconstr_loss, 'train/tv_loss': tv_loss,
-             'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/reconst_loss_norm': reconstr_loss / (2 * (gamma ** 2)), 'train/var': np.e ** logvar,
+             #'train/euclid_error_reconst': metric_2,
+             'train/l2_loss_full_res': l2_loss_full, 'train/ms_ssim_full_res': ms_ssim_full,
+             'train/euclid_error_posed': euclid_posed, 'train/reconstr_loss': reconstr_loss,
+             'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
              'train/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss }
         
 
@@ -235,29 +229,36 @@ def plot_everything(batch: dict, type: str):
     pred, mean, logvar, z = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], uv_curvature_low, batch['uv_static'], batch['uv_body_posed'], uv_curvature])
     # pred = pred * batch['uv_mask' + res_modifier][:,:,:,tf.newaxis]
     
-    mesh_pred = uv2mesh(pred)
+    pred = enhance(upscale_enhancer, pred, batch['uv_mask' + res_modifier], training=False)
+    mesh_pred = uv2mesh(pred).numpy()
     mesh_orig = batch['mesh_unposed']
     
-    orig_img = tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1])
-    uv_static_img = tf.transpose(batch['uv_static' + res_modifier], [0,2,3,1])
-    uv_body_posed_img = tf.transpose(batch['uv_body_posed' + res_modifier], [0,2,3,1])
-    pred_img = pred
+    orig_img = batch['uv_unposed' + res_modifier]
+    uv_static_img = batch['uv_static' + res_modifier]
+    uv_body_posed_img = batch['uv_body_posed' + res_modifier]
+    pred_img = pred.numpy()
+    pred_img[pred_img > 1] = 0
+    pred_img[pred_img < -1] = 0
     img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
     imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
-
-    F = np.load("custom_utils/smpl_faces.npy")
-    mask = batch['mesh_mask'].numpy().astype(bool)
-    garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
-    gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
     
-    body = batch['mesh_body_posed'].numpy()
-    mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
-
     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])])
 
-    garment_posed = [mesh_posed[i, mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(body.shape[0])]
+    F = np.load("custom_utils/smpl_faces.npy")
+    mask = batch['mesh_mask'].numpy().astype(bool)
+    garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
+    garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
+    gt = [mesh_orig[i].numpy()[mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(pred.shape[0])]
+    
+    for i in range(pred.shape[0]):
+        mesh_pred[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
+    body = batch['mesh_body_posed'].numpy()
+    mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).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))]
@@ -393,6 +394,28 @@ def sample(folder):
 
         print()
 
+
+def get_upscale_enhancer(checkpoint, res):
+    upscale_enhancer = Upscaler(512, res)
+    load_model(upscale_enhancer, checkpoint)
+    upscale_enhancer = tf.keras.Sequential([
+        upscale_enhancer.flatten,
+        upscale_enhancer.seq,
+        upscale_enhancer.reshape
+    ])
+    upscale_enhancer.trainable = False
+
+    return upscale_enhancer
+
+
+def enhance(upscale_enhancer, input, mask, training=None):
+    x = input * mask[:,:,:,tf.newaxis]
+    x = upscale_enhancer(x, training=training)
+    x = x * (1 - mask[:,:,:,tf.newaxis])
+    x = x + input
+    return x
+
+
 if __name__ == '__main__':
     init_wandb()
     allow_memory_growth()
@@ -404,7 +427,7 @@ if __name__ == '__main__':
     
     smpl = SMPL()
 
-    resolution = 128
+    resolution = 32
     res_modifier = ''
     if resolution != 512:
         res_modifier = '_' + str(resolution)
@@ -412,37 +435,41 @@ if __name__ == '__main__':
     factory = DatasetFactory()
     signature_dynamic = { 'id': (),
                           'pose_shape': (72 + 11,),
-                        #   'uv_posed': (3, 512, 256,),
-                          'uv_body_posed': (3, 512, 256,),
-                          'uv_body_posed' + res_modifier: (3, resolution, resolution//2,),
+                        #   'uv_posed': (512, 256, 3),
+                          'uv_body_posed': (512, 256, 3),
+                          'uv_body_posed' + res_modifier: (resolution, resolution//2, 3),
                           'mesh_body_posed': (14475, 3,),
                           'mesh_posed': (14475, 3,),
                           'mesh_unposed': (14475, 3,),
-                          'uv_unposed': (3, 512, 256,),
-                          'uv_unposed' + res_modifier: (3, resolution, resolution//2,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_unposed' + res_modifier: (resolution, resolution//2, 3),
                           'uv_curvature': (512, 256,),
                           'uv_curvature' + res_modifier: (resolution, resolution//2),
                         }
-    signature_static = { 'uv_static': (3, 512, 256,),
-                         'uv_static' + res_modifier: (3, resolution, resolution//2,),
+    signature_static = { 'uv_static': (512, 256, 3),
+                         'uv_static' + res_modifier: (resolution, resolution//2, 3),
                         #  'mesh_static': (14475, 3,),
                          'mesh_mask': (14475,),
                          'uv_mask' + res_modifier: (resolution, resolution//2,),
-                        #  'uv_shaped_body': (3, 512, 256,),
+                         'uv_mask': (512, 256,),
                         }
 
-    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=True, resolutions=[512, resolution])
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=True, resolutions=[512, resolution])
+    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False, resolutions=[512, resolution])
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[512, resolution])
     dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=batch_size, signature=signature_dynamic)
     static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
 
 
-    model = CVAE128(#avg_garment_path='custom_data/train_100_no_shape_avg_uv.npy',
-                 latent_dim=128, static_model='checkpoints/ae_static_4',
-                 body_model='checpoints/ae_body_8',
-                 curvature_model='checpoints/ae_curvature_7',
+    model = CVAE(#avg_garment_path='custom_data/train_100_no_shape_avg_uv.npy',
+                 size=resolution,
+                 latent_dim=128, static_model='checkpoints/ae_static_6_copy',
+                 body_model='checkpoints/ae_body_8_copy',
+                 curvature_model='checkpoints/ae_curvature_7_copy',
                  freeze_encoders=True,
-                 use_full_res_encoders=True)
+                 encoder_res=512)
+
+    upscale_enhancer = get_upscale_enhancer('checkpoints/up_2_3', res=resolution)
+
 
     # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
     optimizer = tf.optimizers.Adam(learning_rate=0.0001)
diff --git a/cvae_pyramid.py b/cvae_pyramid.py
index 0833510..bbdbc10 100644
--- a/cvae_pyramid.py
+++ b/cvae_pyramid.py
@@ -1,5 +1,5 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "2"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 import tensorflow as tf
 
 gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -37,7 +37,7 @@ from custom_utils.smpl import SMPL
 from models.cvae_pyramid_model import Pyramid
 
 
-name='cvae_pyr_3_d32'
+name='cvae_pyr_4_d32'
 def init_wandb():
     run = wandb.init(project='vae', name=name,
             notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline. \
@@ -130,7 +130,7 @@ def kl_loss(mean, logvar):
     return kl_loss
 
 
-@tf.function
+# @tf.function
 def train_step(batch):
 
     with tf.GradientTape() as tape:
@@ -154,17 +154,17 @@ def train_step(batch):
 
         e_loss = tf.reduce_mean(edge_loss(mesh_pred_full, mesh_orig, batch['mesh_mask']))
         c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed_full, batch['mesh_body_posed'], batch['mesh_mask']))
-        l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred_full)), axis=-1), axis=[1, 2])
-        l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred_full) ** 2), axis=-1), axis=[1, 2])
-        ms_ssim = tf.image.ssim_multiscale(pred_full + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3)
+        l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred_full)), axis=-1), axis=[1, 2])
+        l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred_full) ** 2), axis=-1), axis=[1, 2])
+        ms_ssim = tf.image.ssim_multiscale(pred_full + 1, batch['uv_unposed'] + 1, max_val=3)
         kl = kl_loss(mean_full, logvar_full)
 
         pred_low = pred_low * batch['uv_mask' + resolution_modifier][:,:,:,tf.newaxis]
         e_loss_low = tf.reduce_mean(edge_loss(mesh_pred_low, mesh_orig, batch['mesh_mask']))
         c_loss_low = tf.reduce_mean(collision_penalty(mesh_pred_posed_low, batch['mesh_body_posed'], batch['mesh_mask']))
-        l1_loss_low = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed' + resolution_modifier], [0,2,3,1]) - pred_low)), axis=-1), axis=[1, 2])
-        l2_loss_low = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed' + resolution_modifier], [0,2,3,1]) - pred_low) ** 2), axis=-1), axis=[1, 2])
-        ms_ssim_low = tf.image.ssim(pred_low + 1, tf.transpose(batch['uv_unposed' + resolution_modifier], [0,2,3,1]) + 1, max_val=3)
+        l1_loss_low = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed' + resolution_modifier] - pred_low)), axis=-1), axis=[1, 2])
+        l2_loss_low = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed' + resolution_modifier] - pred_low) ** 2), axis=-1), axis=[1, 2])
+        ms_ssim_low = tf.image.ssim(pred_low + 1, batch['uv_unposed' + resolution_modifier] + 1, max_val=3)
         kl_low = kl_loss(mean_full, logvar_full)
 
         reconstr_loss = l2_loss + l2_loss_low
@@ -246,9 +246,9 @@ def plot_everything(batch: dict, type: str):
     mesh_pred_low = uv2mesh(pred_low)
     mesh_orig = batch['mesh_unposed']
     
-    orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
-    uv_static_img = tf.transpose(batch['uv_static'], [0,2,3,1])
-    uv_body_posed_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1])
+    orig_img = batch['uv_unposed']
+    uv_static_img = batch['uv_static']
+    uv_body_posed_img = batch['uv_body_posed']
     pred_img = pred_full
     pred_low_up = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
     img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img, pred_low_up), axis=2)
@@ -264,12 +264,14 @@ def plot_everything(batch: dict, type: str):
     garment_unposed_low, faces = maskFacesBatch(mesh_pred_low.numpy(), F, mask)
     garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
     garment_unposed_low = [garment_unposed_low[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
-    gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred_full.shape[0])]
+    gt = [mesh_orig[i].numpy()[mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
     
+    for i in range(mask.shape[0]):
+        mesh_pred_full[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
     body = batch['mesh_body_posed'].numpy()
     mesh_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full).numpy()
 
-    garment_posed = [mesh_posed_full[i, mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(body.shape[0])]
+    garment_posed = [mesh_posed_full[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))]
@@ -405,7 +407,7 @@ def sample(folder):
 
 
 if __name__ == '__main__':
-    init_wandb()
+    # init_wandb()
     allow_memory_growth()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
@@ -420,24 +422,24 @@ if __name__ == '__main__':
     factory = DatasetFactory()
     signature_dynamic = { 'id': (),
                           'pose_shape': (72 + 11,),
-                        #   'uv_posed': (3, 512, 256,),
-                          'uv_body_posed': (3, 512, 256,),
-                          'uv_body_posed' + resolution_modifier: (3, low_resolution, low_resolution//2,),
+                        #   'uv_posed': (512, 256, 3),
+                          'uv_body_posed': (512, 256, 3),
+                          'uv_body_posed' + resolution_modifier: (low_resolution, low_resolution//2, 3),
                           'mesh_body_posed': (14475, 3,),
                           'mesh_posed': (14475, 3,),
                           'mesh_unposed': (14475, 3,),
-                          'uv_unposed': (3, 512, 256,),
-                          'uv_unposed' + resolution_modifier: (3, low_resolution, low_resolution//2,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_unposed' + resolution_modifier: (low_resolution, low_resolution//2, 3),
                           'uv_curvature': (512, 256),
                           'uv_curvature' + resolution_modifier: (low_resolution, low_resolution//2),
                         }
-    signature_static = { 'uv_static': (3, 512, 256,),
-                         'uv_static' + resolution_modifier: (3, low_resolution, low_resolution//2,),
+    signature_static = { 'uv_static': (512, 256, 3),
+                         'uv_static' + resolution_modifier: (low_resolution, low_resolution//2, 3),
                         #  'mesh_static': (14475, 3,),
                          'mesh_mask': (14475,),
                          'uv_mask': (512, 256,),
                          'uv_mask' + resolution_modifier: (low_resolution, low_resolution//2,),
-                        #  'uv_shaped_body': (3, 512, 256,),
+                        #  'uv_shaped_body': (512, 256, 3),
                         }
 
     dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False, resolutions=[512, low_resolution])
@@ -449,9 +451,9 @@ if __name__ == '__main__':
     model = Pyramid(latent_dim=128,
                     low_res=low_resolution, 
                     freeze_encoders=True, 
-                    static_model='checkpoints/ae_static_4',
-                    body_model='checpoints/ae_body_7',
-                    curvature_model='checpoints/ae_curvature_7',
+                    static_model='checkpoints/ae_static_6',
+                    body_model='checkpoints/ae_body_8',
+                    curvature_model='checkpoints/ae_curvature_7',
                     upscaler_model='checkpoints/up_2_3')
 
     # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
diff --git a/cvae_vae.py b/cvae_vae.py
index 4b97230..44f7d7f 100644
--- a/cvae_vae.py
+++ b/cvae_vae.py
@@ -36,7 +36,7 @@ from stylegan2_ae.generator import Generator
 from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder, sparse_to_tensor
 from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
 from dataset import DatasetFactory, StaticDataset
-from custom_utils.custom_utils import batch_to_dict
+from custom_utils.custom_utils import batch_to_dict, plot_3d
 from custom_utils.smpl import SMPL
 from auto_encoder_body import AutoEncoder as BodyAE
 from auto_encoder_curv import AutoEncoder as CurvatureAE
@@ -61,53 +61,6 @@ rebalanced l1 loss. retrained body encoder with deoriented bodies \
 
 uv2mesh = UV2Mesh(res=256)
 
-def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
-    mesh_pred_unposed[:,1] += 1
-    collision[:,1] += 2
-    fig = go.Figure(data=[
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=mesh_gt[:,0],
-            y=mesh_gt[:,1],
-            z=mesh_gt[:,2],
-            # i, j and k give the vertices of triangles
-            i = faces[:,0],
-            j = faces[:,1],
-            k = faces[:,2]
-        ),
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=mesh_pred_unposed[:,0],
-            y=mesh_pred_unposed[:,1],
-            z=mesh_pred_unposed[:,2],
-            # i, j and k give the vertices of triangles
-            i = faces[:,0],
-            j = faces[:,1],
-            k = faces[:,2],
-        ),
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=collision[:,0],
-            y=collision[:,1],
-            z=collision[:,2],
-            # i, j and k give the vertices of triangles
-            i = collision_faces[:,0],
-            j = collision_faces[:,1],
-            k = collision_faces[:,2],
-            vertexcolor=collision_color
-        )
-    ])
-
-    fig.update_layout(scene_aspectmode='manual',
-                    scene_aspectratio=dict(x=1, y=2, z=1))
-    fig.update_layout(
-        scene = dict(
-            xaxis = dict(nticks=4, range=[-1,1],),
-                        yaxis = dict(nticks=4, range=[-1,3],),
-                        zaxis = dict(nticks=4, range=[-1,1],),)
-                        )
-
-    return fig
 
 
 class Bottleneck(tf.keras.Model):
diff --git a/dataset.py b/dataset.py
index c5a952b..e36f2eb 100644
--- a/dataset.py
+++ b/dataset.py
@@ -10,46 +10,46 @@ from custom_utils.custom_utils import nest_conditions, batch_to_dict
 
 
 
-class DatasetSimple(tf.data.Dataset):
+# class DatasetSimple(tf.data.Dataset):
 
-    def _generator(folder):
-        # Opening the file
-        folder = folder.decode("utf-8")
-        uv_folder = os.path.join(folder, 'uv_posed') 
-        pose_hape_folder = os.path.join(folder, 'poses_shapes') 
-        frames = os.listdir(uv_folder)
-        random.shuffle(frames)
+#     def _generator(folder):
+#         # Opening the file
+#         folder = folder.decode("utf-8")
+#         uv_folder = os.path.join(folder, 'uv_posed') 
+#         pose_hape_folder = os.path.join(folder, 'poses_shapes') 
+#         frames = os.listdir(uv_folder)
+#         random.shuffle(frames)
 
-        for frame in frames:
-            id, frame_nr = frame.split('.')[0].split('_')
-            uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
-            uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1])
-            pose_shape = np.load(os.path.join(folder, 'pose_shape',  id + '_' + frame_nr + '.npy'))
-            uv_static = np.load(os.path.join(folder, 'uv_static', id +  '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1])
-            mesh_static = np.load(os.path.join(folder, 'mesh_static', id +  '.npz')); mesh_static = mesh_static['mesh_static']
-            mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
-            mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
-            uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
-            mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed']
-            mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
-            uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
+#         for frame in frames:
+#             id, frame_nr = frame.split('.')[0].split('_')
+#             uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
+#             uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1])
+#             pose_shape = np.load(os.path.join(folder, 'pose_shape',  id + '_' + frame_nr + '.npy'))
+#             uv_static = np.load(os.path.join(folder, 'uv_static', id +  '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1])
+#             mesh_static = np.load(os.path.join(folder, 'mesh_static', id +  '.npz')); mesh_static = mesh_static['mesh_static']
+#             mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
+#             mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
+#             uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
+#             mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed']
+#             mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
+#             uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
 
-            # uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask
-            yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask, uv_body_posed), mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed
+#             # uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask
+#             yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask, uv_body_posed), mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed
 
-    def __new__(cls, folder, res=256):
-        return tf.data.Dataset.from_generator(
-            cls._generator,
-            output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), 
-                                (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ),
-                                tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), 
-                                tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), 
-                                tf.TensorSpec(shape = (14475,), dtype = tf.float32), 
-                                tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
-                                tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
-                                tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), ),
-            args=(folder, )
-        )
+#     def __new__(cls, folder, res=256):
+#         return tf.data.Dataset.from_generator(
+#             cls._generator,
+#             output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), 
+#                                 (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ),
+#                                 tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), 
+#                                 tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), 
+#                                 tf.TensorSpec(shape = (14475,), dtype = tf.float32), 
+#                                 tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
+#                                 tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
+#                                 tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), ),
+#             args=(folder, )
+#         )
 
 
 
@@ -59,7 +59,7 @@ class StaticDataset():
         self.signature = signature
         self.folder = folder
         assert set(resolutions).issubset(set([32, 128, 512]))
-        self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_128' 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_mask_32', 'uv_mask_128', 'uv_shaped_body']
+        self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_128', 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_mask_32', 'uv_mask_128', 'uv_shaped_body']
         self.keys = [value for value in signature.keys() if value in self.possible_keys]
         if 'uv_static_128' in signature.keys():
             self.keys.append('uv_static_128')
@@ -78,7 +78,7 @@ class StaticDataset():
                 res_string = '_' + str(res)
 
             if 'uv_static' + res_string in self.signature.keys():
-                uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id +  '.npz')); uv_static = uv_static['uv_static' + self.filled]; uv_static = np.transpose(uv_static, [2, 0, 1]); data['uv_static' + res_string] = uv_static
+                uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id +  '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static
             if 'uv_mask' + res_string in self.signature.keys():
                 uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask
 
@@ -88,14 +88,11 @@ class StaticDataset():
         if 'mesh_mask' in self.signature.keys():
             mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask
         if 'uv_shaped_body' in self.signature.keys():
-            uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; uv_shaped_body = np.transpose(uv_shaped_body, [2, 0, 1]); data['uv_shaped_body'] = uv_shaped_body
+            uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body
 
         return data
 
     def __getitem__(self, id) -> dict:
-        # id = tf.strings.unicode_decode(id, 'UTF-8')
-        # id = tf.strings.as_string(id)
-        # id = id.numpy()
         if id not in self.store:
             self.store[id] = self.read(id)
 
@@ -170,7 +167,7 @@ class DatasetFactory():
 
                 item['pose_shape'] = pose_shape
             if 'uv_posed' in signature.keys():
-                uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
+                uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
                 item['uv_posed'] = uv_posed
             # if 'uv_body_posed' in signature.keys():
             #     uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1])
@@ -218,11 +215,11 @@ class DatasetFactory():
                     item['uv_curvature' + res_string] = uv_curvature.astype(np.int32)
 
                 if 'uv_unposed' + res_string in signature.keys():
-                    uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled]; uv_unposed_no_shape = np.transpose(uv_unposed_no_shape, [2, 0, 1])
+                    uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled]
                     item['uv_unposed' + res_string] = uv_unposed_no_shape
 
                 if 'uv_body_posed' + res_string in signature.keys():
-                    uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient' + filled]; uv_posed_body_no_orient = np.transpose(uv_posed_body_no_orient, [2, 0, 1])
+                    uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient' + filled]
                     item['uv_body_posed' + res_string] = uv_posed_body_no_orient
 
 
diff --git a/downscale_unpose_test.py b/downscale_unpose_test.py
index 882166b..379b301 100644
--- a/downscale_unpose_test.py
+++ b/downscale_unpose_test.py
@@ -49,7 +49,7 @@ if __name__ == '__main__':
         mask_idx = np.array(range(14475))[mesh_mask]
         
         uv_posed_filled = fill(uv_posed, uv_mask)
-        down = tf.image.resize(uv_posed_filled, size=(32, 16), method=tf.image.ResizeMethod.LANCZOS5)
+        down = tf.image.resize(uv_posed_filled, size=(32, 16), method=tf.image.ResizeMethod.BILINEAR)
         mesh_posed_low = uv2mesh(down[tf.newaxis,:,:,:])
 
         unposed = unpose(mesh_static[mesh_mask,:], mesh_posed[np.newaxis,mesh_mask,:], pose[:,np.newaxis], shape, gender, mesh_mask, mask_idx, smpl_2, id)
diff --git a/downscaling_test.py b/downscaling_test.py
index b795755..9ca036d 100644
--- a/downscaling_test.py
+++ b/downscaling_test.py
@@ -44,11 +44,11 @@ def downscale_1(img):
 def downscale_2(img):
     results = [img]
     for i in range(5):
-        # print(np.array(img.shape[:2]) / (2 ** (i + 1)))
-        # results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 ** (i + 1)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR, antialias=True, preserve_aspect_ratio=True))
-        size = (np.array(img.shape[:2]) / (2 ** (i + 1))).astype(np.int32)
-        size = (size[1], size[0])
-        results.append(cv2.resize(img, size, interpolation=cv2.INTER_NEAREST))
+        print(np.array(img.shape[:2]) / (2 ** (i + 1)))
+        results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 ** (i + 1)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR, antialias=True, preserve_aspect_ratio=True))
+        # size = (np.array(img.shape[:2]) / (2 ** (i + 1))).astype(np.int32)
+        # size = (size[1], size[0])
+        # results.append(cv2.resize(img, size, interpolation=cv2.INTER_NEAREST))
     return results
 
 def upscale_1(imgs):
@@ -127,27 +127,30 @@ if __name__ == '__main__':
     factory = DatasetFactory()
     signature_dynamic = { 'id': (),
                           'pose_shape': (72 + 11,),
-                          'uv_posed': (3, 512, 256,),
+                        #   'uv_posed': (3, 512, 256,),
                         #   'uv_body_posed': (3, 512, 256,),
                           'mesh_body_posed': (14475, 3,),
                           'mesh_posed': (14475, 3,),
                         #   'mesh_unposed': (14475, 3,),
                         #   'uv_unposed': (3, 512, 256,)
                         }
-    signature_static = { #'uv_static': (3, 512, 256,),
-                        #  'mesh_static': (14475, 3,),
+    signature_static = {  'uv_static': (3, 512, 256,),
+                         # 'mesh_static': (14475, 3,),
                          'mesh_mask': (14475,),
                          'uv_mask': (512, 256,),
-                        #  'uv_shaped_body': (3, 512, 256,),
+                         'uv_mask_32': (32, 16,),
+                         # 'uv_shaped_body': (3, 512, 256,),
                         }
 
     dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100',
                                   batch_size=1,
                                   signature=signature_dynamic,
-                                  use_filled=True)
+                                  use_filled=True,
+                                  resolutions=[512, 32])
     static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100',
                                    signature=signature_static,
-                                   use_filled=True)
+                                   use_filled=True,
+                                   resolutions=[512, 32])
 
 
     step = 0
@@ -155,7 +158,7 @@ if __name__ == '__main__':
         # batch = batch_to_dict(batch, signature_dynamic)
         batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
 
-        img = tf.transpose(tf.squeeze(batch['uv_posed']), [1,2,0]).numpy()
+        img = tf.transpose(tf.squeeze(batch['uv_static']), [1,2,0]).numpy()
         coords = np.transpose(np.stack(np.nonzero((1 - batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
         img_idx = np.transpose(np.stack(np.nonzero((batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
         img_data = img[img_idx[:,0], img_idx[:,1], :]
@@ -169,7 +172,8 @@ if __name__ == '__main__':
         img_2 = res
 
         uv_mask = batch['uv_mask'].numpy()[0]
-        uv_mask = downscale_2(uv_mask)
+        uv_mask_32 = batch['uv_mask_32'].numpy()[0]
+        # uv_mask = downscale_2(uv_mask)
         # for i in range(0,3):
         #     uv_mask[i] = cv2.dilate(uv_mask[i].astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)#.astype(bool)[:, :, np.newaxis]
         # for i in range(3,6):
@@ -178,9 +182,9 @@ if __name__ == '__main__':
 
         down_1 = downscale_1(img)
         down_2 = downscale_1(img_2)
-        up_1 = upscale_1(down_1)
+        up_1 = upscale_1([down_1[-2]])
         # down_2 = [down_2[i] * uv_mask[i][:,:,np.newaxis] for i in range(len(down_2))]
-        up_2 = upscale_1(down_2)
+        up_2 = upscale_1([down_1[-2] * uv_mask_32[:,:,tf.newaxis]])
 
         mask = batch['mesh_mask'].numpy().astype(bool)[0]
         mesh_2 = [uv2mesh(tf.transpose(up_2[i], [2,0,1])[tf.newaxis,:,:,:])[0] for i in range(len(up_2))]
@@ -230,7 +234,7 @@ if __name__ == '__main__':
 
         fig = make_subplots(rows=2, cols=6, start_cell="top-left")
 
-        for i, d2 in enumerate(uv_mask):
+        for i, d2 in enumerate(up_1):
             fig.add_trace(ex.imshow((d2 + 1) * 128).data[0], row=1, col=i+1)
 
         for i, u2 in enumerate(up_2):
@@ -240,7 +244,7 @@ if __name__ == '__main__':
 
         fig = make_subplots(rows=2, cols=6, start_cell="top-left", specs = [ [{ 'type': 'xy' } for _ in range(6)], [{ 'type': 'scene' } for _ in range(6)] ])
 
-        for i, u2 in enumerate(up_2):
+        for i, u2 in enumerate(up_1):
             fig.add_trace(ex.imshow((u2 + 1) * 128).data[0], row=1, col=i+1)
 
         # for i, m1 in enumerate(mesh_1):
diff --git a/downscaling_test_curv.py b/downscaling_test_curv.py
index 0a58101..11e2fe2 100644
--- a/downscaling_test_curv.py
+++ b/downscaling_test_curv.py
@@ -114,10 +114,10 @@ if __name__ == '__main__':
                         #  'uv_shaped_body': (3, 512, 256,),
                         }
 
-    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/stylegan_new/train',
+    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100',
                                   batch_size=1,
                                   signature=signature_dynamic)
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/stylegan_new/train',
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100',
                                    signature=signature_static)
 
 
diff --git a/models/ae_body_model.py b/models/ae_body_model.py
index eaf6924..34e579f 100644
--- a/models/ae_body_model.py
+++ b/models/ae_body_model.py
@@ -1,223 +1,202 @@
 import tensorflow as tf
 from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
-
-
-
-
-class Bottleneck(tf.keras.Model):
-    
-    def __init__(self, dim, latent_dim):
-        super(Bottleneck, self).__init__()
-
-        self.bottleneck = tf.keras.Sequential(
-                [Block(dim=dim, drop_path=0,
-                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
-            )
-        self.flatten = tf.keras.layers.Flatten()
-        self.dense = tf.keras.layers.Dense(latent_dim)
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.bottleneck(x)
-        x = self.flatten(x)
-        x = self.dense(x)
-        
-        return x
-
-
-class Decoder(tf.keras.Model):
-    
-    def __init__(self, bottleneck_size, latent_dim):
-        super(Decoder, self).__init__()
-
-        self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
-        self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
-        self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])))
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
-
-        return x
+from models.base import AEEncoder, Bottleneck, Decoder
 
 
 class AutoEncoder(tf.keras.Model):
     
-    def __init__(self):
+    def __init__(self, size):
         super(AutoEncoder, self).__init__()
 
+        if size == 512:
+            bottleneck_size = 64
+            latent_dim = 128
+            kernel_size = 7
+            encoder_depths = [3, 3, 9, 9, 3]
+            encoder_dims = [32, 64, 128, 256, bottleneck_size]
+            decoder_input = (8, 4, bottleneck_size)
+            decoder_depths = [3, 3, 9, 3, 3]
+            decoder_dims = [bottleneck_size, 256, 128, 64, 16]
+        elif size == 128:
+            raise NotImplementedError
+            bottleneck_size = 512
+            latent_dim = 128
+            kernel_size = 5
+            encoder_depths =[16, 16, 9, 9]
+            encoder_dims = [64, 128, 256, bottleneck_size]
+            decoder_input = (4, 2, bottleneck_size)
+            decoder_depths = [16, 16, 9, 9]
+            decoder_dims = [bottleneck_size, 256, 128, 64]
+        elif size == 32:
+            raise NotImplementedError
+            bottleneck_size = 512
+            latent_dim = 128
+            kernel_size = 3
+            encoder_depths = [27, 27]
+            encoder_dims = [256, bottleneck_size]
+            decoder_input = (4, 2, bottleneck_size)
+            decoder_depths = [27, 27]
+            decoder_dims = [bottleneck_size, 256]
+        else:
+            raise NotImplementedError
 
-        bottleneck_size = 512
-        latent_dim = 128
-
-        self.encoder = ConvNeXt(depths=[3, 3, 9, 9, 3, 3], dims=[32, 64, 128, 256, 256, bottleneck_size], include_top=False)
-        self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
-        self.decoder = Decoder(bottleneck_size, latent_dim)
+        self.encoder : AEEncoder = AEEncoder(latent_dim, encoder_depths, encoder_dims, bottleneck_size, kernel_size)
+        self.decoder = Decoder(decoder_input, decoder_depths, decoder_dims, kernel_size=kernel_size)
 
 
-    def call(self, inputs, training=None):
-        x = inputs
+    def call(self, inputs, training=None, use_mask=True):
+        if use_mask:
+            x, mask = inputs
+        else:
+            x = inputs
 
-        x = tf.transpose(x, [0,2,3,1])
         x = self.encoder(x, training=training)
-        # x = self.pooling(x)
-        # x = self.reshape(x)
-        # x = self.flatten(x)
-
-        x = self.bottleneck(x, training=training)
-
         x = self.decoder(x, training=training)
-        x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-        x = tf.transpose(x, [0,3,1,2])
 
-        # x = x * uv_mask[:,:, tf.newaxis]
-        # x = x * uv_mask_body[tf.newaxis,tf.newaxis,:,:]
+        if use_mask:
+            x = x * mask[tf.newaxis,:,:,tf.newaxis]
 
         return x
 
 
 
-class Bottleneck32(tf.keras.Model):
+# class Bottleneck32(tf.keras.Model):
     
-    def __init__(self, dim, latent_dim):
-        super(Bottleneck32, self).__init__()
+#     def __init__(self, dim, latent_dim):
+#         super(Bottleneck32, self).__init__()
 
-        self.bottleneck = tf.keras.Sequential(
-                [Block(dim=dim, drop_path=0,
-                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
-            )
-        self.flatten = tf.keras.layers.Flatten()
-        self.dense = tf.keras.layers.Dense(latent_dim)
+#         self.bottleneck = tf.keras.Sequential(
+#                 [Block(dim=dim, drop_path=0,
+#                        layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+#             )
+#         self.flatten = tf.keras.layers.Flatten()
+#         self.dense = tf.keras.layers.Dense(latent_dim)
 
-    def call(self, inputs, training=None):
-        x = inputs
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = self.bottleneck(x)
-        x = self.flatten(x)
-        x = self.dense(x)
+#         x = self.bottleneck(x)
+#         x = self.flatten(x)
+#         x = self.dense(x)
         
-        return x
+#         return x
 
 
-class Decoder32(tf.keras.Model):
+# class Decoder32(tf.keras.Model):
     
-    def __init__(self, bottleneck_size, latent_dim):
-        super(Decoder32, self).__init__()
+#     def __init__(self, bottleneck_size, latent_dim):
+#         super(Decoder32, self).__init__()
 
-        self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
-        self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
-        self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])))
+#         self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
+#         self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
+#         self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])))
 
-    def call(self, inputs, training=None):
-        x = inputs
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
+#         x = self.dense(x)
+#         x = self.reshape(x)
+#         x = self.decoder(x)
 
-        return x
+#         return x
 
 
-class AutoEncoder32(tf.keras.Model):
+# class AutoEncoder32(tf.keras.Model):
     
-    def __init__(self):
-        super(AutoEncoder32, self).__init__()
+#     def __init__(self):
+#         super(AutoEncoder32, self).__init__()
 
 
-        bottleneck_size = 512
-        latent_dim = 128
+#         bottleneck_size = 512
+#         latent_dim = 128
 
-        self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
-        self.bottleneck = Bottleneck32(bottleneck_size, latent_dim)
-        self.decoder = Decoder32(bottleneck_size, latent_dim)
+#         self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
+#         self.bottleneck = Bottleneck32(bottleneck_size, latent_dim)
+#         self.decoder = Decoder32(bottleneck_size, latent_dim)
 
 
-    def call(self, inputs, training=None):
-        x = inputs
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = tf.transpose(x, [0,2,3,1])
-        x = self.encoder(x, training=training)
+#         x = tf.transpose(x, [0,2,3,1])
+#         x = self.encoder(x, training=training)
 
-        x = self.bottleneck(x, training=training)
+#         x = self.bottleneck(x, training=training)
 
-        x = self.decoder(x, training=training)
-        # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-        x = tf.transpose(x, [0,3,1,2])
+#         x = self.decoder(x, training=training)
+#         # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+#         x = tf.transpose(x, [0,3,1,2])
 
-        return x
+#         return x
 
 
 
 
-class Bottleneck128(tf.keras.Model):
+# class Bottleneck128(tf.keras.Model):
     
-    def __init__(self, dim, latent_dim):
-        super(Bottleneck128, self).__init__()
+#     def __init__(self, dim, latent_dim):
+#         super(Bottleneck128, self).__init__()
 
-        self.bottleneck = tf.keras.Sequential(
-                [Block(dim=dim, drop_path=0,
-                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
-            )
-        self.flatten = tf.keras.layers.Flatten()
-        self.dense = tf.keras.layers.Dense(latent_dim)
+#         self.bottleneck = tf.keras.Sequential(
+#                 [Block(dim=dim, drop_path=0,
+#                        layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+#             )
+#         self.flatten = tf.keras.layers.Flatten()
+#         self.dense = tf.keras.layers.Dense(latent_dim)
 
-    def call(self, inputs, training=None):
-        x = inputs
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = self.bottleneck(x)
-        x = self.flatten(x)
-        x = self.dense(x)
+#         x = self.bottleneck(x)
+#         x = self.flatten(x)
+#         x = self.dense(x)
         
-        return x
+#         return x
 
 
-class Decoder128(tf.keras.Model):
+# class Decoder128(tf.keras.Model):
     
-    def __init__(self, bottleneck_size, latent_dim):
-        super(Decoder128, self).__init__()
+#     def __init__(self, bottleneck_size, latent_dim):
+#         super(Decoder128, self).__init__()
 
-        self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
-        self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
-        self.decoder = ConvNeXtReverse(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])))
+#         self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
+#         self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
+#         self.decoder = ConvNeXtReverse(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])))
 
-    def call(self, inputs, training=None):
-        x = inputs
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
+#         x = self.dense(x)
+#         x = self.reshape(x)
+#         x = self.decoder(x)
 
-        return x
+#         return x
 
 
-class AutoEncoder128(tf.keras.Model):
+# class AutoEncoder128(tf.keras.Model):
     
-    def __init__(self):
-        super(AutoEncoder128, self).__init__()
+#     def __init__(self):
+#         super(AutoEncoder128, self).__init__()
 
 
-        bottleneck_size = 512
-        latent_dim = 128
+#         bottleneck_size = 512
+#         latent_dim = 128
 
-        self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
-        self.bottleneck = Bottleneck128(bottleneck_size, latent_dim)
-        self.decoder = Decoder128(bottleneck_size, latent_dim)
+#         self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
+#         self.bottleneck = Bottleneck128(bottleneck_size, latent_dim)
+#         self.decoder = Decoder128(bottleneck_size, latent_dim)
 
 
-    def call(self, inputs, training=None):
-        x = inputs
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = tf.transpose(x, [0,2,3,1])
-        x = self.encoder(x, training=training)
+#         x = tf.transpose(x, [0,2,3,1])
+#         x = self.encoder(x, training=training)
 
-        x = self.bottleneck(x, training=training)
+#         x = self.bottleneck(x, training=training)
 
-        x = self.decoder(x, training=training)
-        # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-        x = tf.transpose(x, [0,3,1,2])
+#         x = self.decoder(x, training=training)
+#         # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+#         x = tf.transpose(x, [0,3,1,2])
 
-        return x
\ No newline at end of file
+#         return x
\ No newline at end of file
diff --git a/models/ae_curv_model.py b/models/ae_curv_model.py
index ebb7f5f..3adaa74 100644
--- a/models/ae_curv_model.py
+++ b/models/ae_curv_model.py
@@ -1,63 +1,73 @@
 import tensorflow as tf
 from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverseLabels, Block
+from models.base import AEEncoder, Bottleneck, Decoder
 
 
 
-class Bottleneck(tf.keras.Model):
+# class Bottleneck(tf.keras.Model):
     
-    def __init__(self, dim, latent_dim):
-        super(Bottleneck, self).__init__()
+#     def __init__(self, dim, latent_dim):
+#         super(Bottleneck, self).__init__()
 
-        self.bottleneck = tf.keras.Sequential(
-                [Block(dim=dim, drop_path=0,
-                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
-            )
-        self.flatten = tf.keras.layers.Flatten()
-        self.dense = tf.keras.layers.Dense(latent_dim)
+#         self.bottleneck = tf.keras.Sequential(
+#                 [Block(dim=dim, drop_path=0,
+#                        layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+#             )
+#         self.flatten = tf.keras.layers.Flatten()
+#         self.dense = tf.keras.layers.Dense(latent_dim)
 
-    def call(self, inputs, training=None):
-        x = inputs
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = self.bottleneck(x)
-        x = self.flatten(x)
-        x = self.dense(x)
+#         x = self.bottleneck(x)
+#         x = self.flatten(x)
+#         x = self.dense(x)
         
-        return x
+#         return x
 
 
-class Decoder(tf.keras.Model):
+# class Decoder(tf.keras.Model):
     
-    def __init__(self, bottleneck_size, latent_dim):
-        super(Decoder, self).__init__()
+#     def __init__(self, bottleneck_size, latent_dim):
+#         super(Decoder, self).__init__()
 
-        self.dense = tf.keras.layers.Dense(16 * 8 * bottleneck_size)
-        self.reshape = tf.keras.layers.Reshape((16, 8, bottleneck_size))
-        self.decoder = ConvNeXtReverseLabels(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])), nr_of_classes=7)
+#         self.dense = tf.keras.layers.Dense(16 * 8 * bottleneck_size)
+#         self.reshape = tf.keras.layers.Reshape((16, 8, bottleneck_size))
+#         self.decoder = ConvNeXtReverseLabels(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])), nr_of_classes=7)
 
-    def call(self, inputs, training=None):
-        x = inputs
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
+#         x = self.dense(x)
+#         x = self.reshape(x)
+#         x = self.decoder(x)
 
-        return x
+#         return x
 
         
 class AutoEncoder(tf.keras.Model):
     
-    def __init__(self):
+    def __init__(self, size=512):
         super(AutoEncoder, self).__init__()
 
+        if size == 512:
+            bottleneck_size = 64
+            latent_dim = 128
+            kernel_size = 7
+            encoder_depths = [3, 3, 9, 3, 3]
+            encoder_dims = [32, 64, 128, 256, bottleneck_size]
+            decoder_input = (16, 8, bottleneck_size)
+            decoder_depths = [3, 3, 9, 3, 3]
+            decoder_dims = [bottleneck_size, 256, 128, 64, 16]
+        elif size == 128:
+            raise NotImplementedError
+        elif size == 32:
+            raise NotImplementedError
+        else:
+            raise NotImplementedError
 
-        bottleneck_size = 32
-        latent_dim = 128
-        self.decoder = Decoder(bottleneck_size, latent_dim)
-
-        self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
-
-        self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, bottleneck_size], include_top=False)
-
+        self.encoder : AEEncoder = AEEncoder(latent_dim, encoder_depths, encoder_dims, bottleneck_size, kernel_size)
+        self.decoder = Decoder(decoder_input, decoder_depths, decoder_dims, kernel_size, is_label=True)
 
 
     def call(self, inputs, training=None):
@@ -65,117 +75,113 @@ class AutoEncoder(tf.keras.Model):
 
         # x = tf.transpose(x, [0,2,3,1])
         x = self.encoder(x, training=training)
-
-
-        x = self.bottleneck(x, training=training)
-
         x = self.decoder(x, training=training)
         # x = tf.transpose(x, [0,3,1,2])
 
         # x = x * uv_mask[:,:, tf.newaxis]
         # x = x * uv_mask[:,tf.newaxis,:,:]
-        # x = x * uv_mask[:,:,:,tf.newaxis]
+        x = x * uv_mask[:,:,:,tf.newaxis]
 
         return x
 
 
 
-class Decoder32(tf.keras.Model):
+# class Decoder32(tf.keras.Model):
     
-    def __init__(self, bottleneck_size, latent_dim):
-        super(Decoder32, self).__init__()
+#     def __init__(self, bottleneck_size, latent_dim):
+#         super(Decoder32, self).__init__()
 
-        self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
-        self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
-        self.decoder = ConvNeXtReverseLabels(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])), nr_of_classes=7)
+#         self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
+#         self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
+#         self.decoder = ConvNeXtReverseLabels(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])), nr_of_classes=7)
 
-    def call(self, inputs, training=None):
-        x = inputs
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
+#         x = self.dense(x)
+#         x = self.reshape(x)
+#         x = self.decoder(x)
 
-        return x
+#         return x
 
         
-class AutoEncoder32(tf.keras.Model):
+# class AutoEncoder32(tf.keras.Model):
     
-    def __init__(self):
-        super(AutoEncoder32, self).__init__()
+#     def __init__(self):
+#         super(AutoEncoder32, self).__init__()
 
-        bottleneck_size = 512
-        latent_dim = 128
+#         bottleneck_size = 512
+#         latent_dim = 128
 
-        self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
-        self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
-        self.decoder = Decoder32(bottleneck_size, latent_dim)
+#         self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
+#         self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
+#         self.decoder = Decoder32(bottleneck_size, latent_dim)
 
 
-    def call(self, inputs, training=None):
-        x, uv_mask = inputs
+#     def call(self, inputs, training=None):
+#         x, uv_mask = inputs
 
-        # x = tf.transpose(x, [0,2,3,1])
-        x = self.encoder(x, training=training)
+#         # x = tf.transpose(x, [0,2,3,1])
+#         x = self.encoder(x, training=training)
 
-        x = self.bottleneck(x, training=training)
+#         x = self.bottleneck(x, training=training)
 
-        x = self.decoder(x, training=training)
-        # x = tf.transpose(x, [0,3,1,2])
+#         x = self.decoder(x, training=training)
+#         # x = tf.transpose(x, [0,3,1,2])
 
-        # x = x * uv_mask[:,:, tf.newaxis]
-        # x = x * uv_mask[:,tf.newaxis,:,:]
-        # x = x * uv_mask[:,:,:,tf.newaxis]
+#         # x = x * uv_mask[:,:, tf.newaxis]
+#         # x = x * uv_mask[:,tf.newaxis,:,:]
+#         # x = x * uv_mask[:,:,:,tf.newaxis]
 
-        return x
+#         return x
 
 
 
-class Decoder128(tf.keras.Model):
+# class Decoder128(tf.keras.Model):
     
-    def __init__(self, bottleneck_size, latent_dim):
-        super(Decoder128, self).__init__()
+#     def __init__(self, bottleneck_size, latent_dim):
+#         super(Decoder128, self).__init__()
 
-        self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
-        self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
-        self.decoder = ConvNeXtReverseLabels(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])), nr_of_classes=7)
+#         self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
+#         self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
+#         self.decoder = ConvNeXtReverseLabels(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])), nr_of_classes=7)
 
-    def call(self, inputs, training=None):
-        x = inputs
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
+#         x = self.dense(x)
+#         x = self.reshape(x)
+#         x = self.decoder(x)
 
-        return x
+#         return x
 
         
-class AutoEncoder128(tf.keras.Model):
+# class AutoEncoder128(tf.keras.Model):
     
-    def __init__(self):
-        super(AutoEncoder128, self).__init__()
+#     def __init__(self):
+#         super(AutoEncoder128, self).__init__()
 
-        bottleneck_size = 512
-        latent_dim = 128
+#         bottleneck_size = 512
+#         latent_dim = 128
 
-        self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
-        self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
-        self.decoder = Decoder128(bottleneck_size, latent_dim)
+#         self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
+#         self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
+#         self.decoder = Decoder128(bottleneck_size, latent_dim)
 
 
-    def call(self, inputs, training=None):
-        x, uv_mask = inputs
+#     def call(self, inputs, training=None):
+#         x, uv_mask = inputs
 
-        # x = tf.transpose(x, [0,2,3,1])
-        x = self.encoder(x, training=training)
+#         # x = tf.transpose(x, [0,2,3,1])
+#         x = self.encoder(x, training=training)
 
-        x = self.bottleneck(x, training=training)
+#         x = self.bottleneck(x, training=training)
 
-        x = self.decoder(x, training=training)
-        # x = tf.transpose(x, [0,3,1,2])
+#         x = self.decoder(x, training=training)
+#         # x = tf.transpose(x, [0,3,1,2])
 
-        # x = x * uv_mask[:,:, tf.newaxis]
-        # x = x * uv_mask[:,tf.newaxis,:,:]
-        # x = x * uv_mask[:,:,:,tf.newaxis]
+#         # x = x * uv_mask[:,:, tf.newaxis]
+#         # x = x * uv_mask[:,tf.newaxis,:,:]
+#         # x = x * uv_mask[:,:,:,tf.newaxis]
 
-        return x
+#         return x
diff --git a/models/ae_static_model.py b/models/ae_static_model.py
index bcdf98c..c0c4020 100644
--- a/models/ae_static_model.py
+++ b/models/ae_static_model.py
@@ -1,172 +1,156 @@
 import tensorflow as tf
 from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
+from models.base import AEEncoder, Bottleneck, Decoder
 
 
 
-class Bottleneck(tf.keras.Model):
-    
-    def __init__(self, dim, latent_dim):
-        super(Bottleneck, self).__init__()
-
-        self.bottleneck = tf.keras.Sequential(
-                [Block(dim=dim, drop_path=0,
-                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
-            )
-        self.flatten = tf.keras.layers.Flatten()
-        self.dense = tf.keras.layers.Dense(latent_dim)
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.bottleneck(x)
-        x = self.flatten(x)
-        x = self.dense(x)
-        
-        return x
-
-
-class Decoder(tf.keras.Model):
-    
-    def __init__(self, bottleneck_size, latent_dim):
-        super(Decoder, self).__init__()
-
-        self.dense = tf.keras.layers.Dense(8 * 4 * bottleneck_size)
-        self.reshape = tf.keras.layers.Reshape((8, 4, bottleneck_size))
-        self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3, 3])), dims=list(reversed([bottleneck_size, 128, 128, 64, 16])))
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
-
-        return x
-
-
 class AutoEncoder(tf.keras.Model):
     
-    def __init__(self):
+    def __init__(self, size=512):
         super(AutoEncoder, self).__init__()
 
+        if size == 512:
+            bottleneck_size = 64
+            latent_dim = 128
+            kernel_size = 7
+            encoder_depths = [3, 3, 9, 3, 3]
+            encoder_dims = [32, 64, 128, 256, bottleneck_size]
+            decoder_input = (8, 4, bottleneck_size)
+            decoder_depths = [3, 3, 9, 3, 3]
+            decoder_dims = [bottleneck_size, 256, 128, 64, 16]
+        elif size == 128:
+            bottleneck_size = 512
+            latent_dim = 128
+            kernel_size = 5
+            encoder_depths =[16, 16, 9, 9]
+            encoder_dims = [64, 128, 256, bottleneck_size]
+            decoder_input = (4, 2, bottleneck_size)
+            decoder_depths = [16, 16, 9, 9]
+            decoder_dims = [bottleneck_size, 256, 128, 64]
+        elif size == 32:
+            bottleneck_size = 512
+            latent_dim = 128
+            kernel_size = 3
+            encoder_depths = [27, 27]
+            encoder_dims = [256, bottleneck_size]
+            decoder_input = (4, 2, bottleneck_size)
+            decoder_depths = [27, 27]
+            decoder_dims = [bottleneck_size, 256]
+        else:
+            raise NotImplementedError
 
-        bottleneck_size = 32
-        latent_dim = 128
+        self.encoder : AEEncoder = AEEncoder(latent_dim, encoder_depths, encoder_dims, bottleneck_size, kernel_size)
+        self.decoder : Decoder = Decoder(decoder_input, decoder_depths, decoder_dims, kernel_size)
 
-        self.encoder = ConvNeXt(depths=[3, 3, 9, 3, 3], dims=[32, 64, 128, 128, bottleneck_size], include_top=False)
-        self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
-        self.decoder = Decoder(bottleneck_size, latent_dim)
 
-    def call(self, inputs, training=None):
-        x = inputs
+    def call(self, inputs, training=None, use_mask=True):
+        if use_mask:
+            x, mask = inputs
+        else:
+            x = inputs
 
-        x = tf.transpose(x, [0,2,3,1])
         x = self.encoder(x, training=training)
-
-        x = self.bottleneck(x, training=training)
-
         x = self.decoder(x, training=training)
-        x = tf.transpose(x, [0,3,1,2])
 
-        # x = x * uv_mask[:,:, tf.newaxis]
-        # x = x * uv_mask[:,tf.newaxis,:,:]
+        if use_mask:
+             x = x * mask[:,:,:,tf.newaxis]
 
         return x
 
 
-class Decoder32(tf.keras.Model):
+# class Decoder32(tf.keras.Model):
     
-    def __init__(self, bottleneck_size, latent_dim):
-        super(Decoder32, self).__init__()
+#     def __init__(self, bottleneck_size, latent_dim):
+#         super(Decoder32, self).__init__()
 
-        self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
-        self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
-        self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])))
+#         self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
+#         self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
+#         self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])))
 
-    def call(self, inputs, training=None):
-        x = inputs
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
+#         x = self.dense(x)
+#         x = self.reshape(x)
+#         x = self.decoder(x)
 
-        return x
+#         return x
 
 
-class AutoEncoder32(tf.keras.Model):
+# class AutoEncoder32(tf.keras.Model):
     
-    def __init__(self):
-        super(AutoEncoder32, self).__init__()
+#     def __init__(self):
+#         super(AutoEncoder32, self).__init__()
 
 
-        bottleneck_size = 512
-        latent_dim = 128
+#         bottleneck_size = 512
+#         latent_dim = 128
 
-        self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
-        self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
-        self.decoder = Decoder32(bottleneck_size, latent_dim)
+#         self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
+#         self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
+#         self.decoder = Decoder32(bottleneck_size, latent_dim)
 
-    def call(self, inputs, training=None):
-        x = inputs
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = tf.transpose(x, [0,2,3,1])
-        x = self.encoder(x, training=training)
+#         x = tf.transpose(x, [0,2,3,1])
+#         x = self.encoder(x, training=training)
 
-        x = self.bottleneck(x, training=training)
+#         x = self.bottleneck(x, training=training)
 
-        x = self.decoder(x, training=training)
-        x = tf.transpose(x, [0,3,1,2])
+#         x = self.decoder(x, training=training)
+#         x = tf.transpose(x, [0,3,1,2])
 
-        # x = x * uv_mask[:,:, tf.newaxis]
-        # x = x * uv_mask[:,tf.newaxis,:,:]
+#         # x = x * uv_mask[:,:, tf.newaxis]
+#         # x = x * uv_mask[:,tf.newaxis,:,:]
 
-        return x
+#         return x
 
 
-class Decoder128(tf.keras.Model):
+# class Decoder128(tf.keras.Model):
     
-    def __init__(self, bottleneck_size, latent_dim):
-        super(Decoder128, self).__init__()
+#     def __init__(self, bottleneck_size, latent_dim):
+#         super(Decoder128, self).__init__()
 
-        self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
-        self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
-        self.decoder = ConvNeXtReverse(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])))
+#         self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
+#         self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
+#         self.decoder = ConvNeXtReverse(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])))
 
-    def call(self, inputs, training=None):
-        x = inputs
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
+#         x = self.dense(x)
+#         x = self.reshape(x)
+#         x = self.decoder(x)
 
-        return x
+#         return x
 
 
-class AutoEncoder128(tf.keras.Model):
+# class AutoEncoder128(tf.keras.Model):
     
-    def __init__(self):
-        super(AutoEncoder128, self).__init__()
+#     def __init__(self):
+#         super(AutoEncoder128, self).__init__()
 
 
-        bottleneck_size = 512
-        latent_dim = 128
+#         bottleneck_size = 512
+#         latent_dim = 128
 
-        self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
-        self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
-        self.decoder = Decoder128(bottleneck_size, latent_dim)
+#         self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
+#         self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
+#         self.decoder = Decoder128(bottleneck_size, latent_dim)
 
-    def call(self, inputs, training=None):
-        x = inputs
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = tf.transpose(x, [0,2,3,1])
-        x = self.encoder(x, training=training)
+#         x = tf.transpose(x, [0,2,3,1])
+#         x = self.encoder(x, training=training)
 
-        x = self.bottleneck(x, training=training)
+#         x = self.bottleneck(x, training=training)
 
-        x = self.decoder(x, training=training)
-        x = tf.transpose(x, [0,3,1,2])
+#         x = self.decoder(x, training=training)
+#         x = tf.transpose(x, [0,3,1,2])
 
-        # x = x * uv_mask[:,:, tf.newaxis]
-        # x = x * uv_mask[:,tf.newaxis,:,:]
+#         # x = x * uv_mask[:,:, tf.newaxis]
+#         # x = x * uv_mask[:,tf.newaxis,:,:]
 
-        return x
\ No newline at end of file
+#         return x
\ No newline at end of file
diff --git a/models/base.py b/models/base.py
new file mode 100644
index 0000000..b865077
--- /dev/null
+++ b/models/base.py
@@ -0,0 +1,106 @@
+import os
+import tensorflow as tf
+
+from custom_utils.custom_utils import load_model
+from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels
+
+class Bottleneck(tf.keras.Model):
+    
+    def __init__(self, dim, latent_dim, kernel_size=7):
+        super(Bottleneck, self).__init__()
+
+        self.bottleneck = tf.keras.Sequential(
+                [Block(dim=dim, drop_path=0, kernel_size=kernel_size,
+                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+            )
+        # self.bottleneck = [Block(dim=dim, drop_path=0, kernel_size=kernel_size, layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+            
+        self.flatten = tf.keras.layers.Flatten()
+        self.dense = tf.keras.layers.Dense(latent_dim)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        # for bt in self.bottleneck:
+        #     x = bt(x, training=training)
+        x = self.bottleneck(x)
+        x = self.flatten(x)
+        x = self.dense(x)
+
+        return x
+
+
+class VariationalEncoder(tf.keras.Model):
+    
+    def __init__(self, latent_dim, depths, dims, use_bottleneck=True, kernel_size=7):
+        super(VariationalEncoder, self).__init__()
+
+        self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
+
+        self.flatten = tf.keras.layers.Flatten()
+        self.mean = tf.keras.layers.Dense(latent_dim)
+        self.log_var = tf.keras.layers.Dense(latent_dim)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.encoder(x, training=training)
+        x = self.flatten(x)
+        mean = self.mean(x)
+        log_var = self.log_var(x)
+
+        return mean, log_var
+
+
+class AEEncoder(tf.keras.Model):
+    
+    def __init__(self, latent_dim, depths, dims, bottleneck_size=32, kernel_size=7):
+        super(AEEncoder, self).__init__()
+
+        self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
+        self.bottleneck = Bottleneck(bottleneck_size, latent_dim, kernel_size)
+
+
+    def call(self, inputs, training=None):
+        if len(inputs) == 2:
+            x, _ = inputs
+        else:
+            x = inputs
+
+        x = self.encoder(x, training=training)
+        x = self.bottleneck(x, training=training)
+
+        return x
+
+
+class Decoder(tf.keras.Model):
+    
+    def __init__(self, input_size, depths, dims, kernel_size=7, is_label=False):
+        super(Decoder, self).__init__()
+
+        self.dense = tf.keras.layers.Dense(tf.reduce_prod(input_size))
+        self.reshape = tf.keras.layers.Reshape(input_size)
+        
+        if not is_label:
+            self.decoder = ConvNeXtReverse(depths=list(reversed(depths)), dims=list(reversed(dims)), kernel_size=kernel_size)
+        else:
+            self.decoder = ConvNeXtReverseLabels(depths=list(reversed(depths)), dims=list(reversed(dims)), nr_of_classes=7)
+            
+
+    def call(self, inputs, training=None):
+        if len(inputs) == 2:
+            x, mask = inputs
+        else:
+            x = inputs
+
+        x = self.dense(x)
+        x = self.reshape(x)
+        x = self.decoder(x, training=training)
+
+        if len(inputs) == 2:
+            if len(mask.shape) == 2:
+                x = x * mask[tf.newaxis,:,:,tf.newaxis]
+            else:
+                x = x * mask[:,:,:,tf.newaxis]
+
+        return x
\ No newline at end of file
diff --git a/models/cond_encoder.py b/models/cond_encoder.py
new file mode 100644
index 0000000..cd73c2b
--- /dev/null
+++ b/models/cond_encoder.py
@@ -0,0 +1,20 @@
+import tensorflow as tf
+from models.base import AEEncoder, Decoder
+
+
+class AutoEncoder(tf.keras.Model):
+    
+    def __init__(self, config):
+        super(AutoEncoder, self).__init__()
+
+        self.encoder : AEEncoder = AEEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], config['bottleneck_size'], config['kernel_size'])
+        self.decoder : Decoder= Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
+
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.encoder(x, training=training)
+        x = self.decoder(x, training=training)
+
+        return x
diff --git a/models/cvae_curv_model.py b/models/cvae_curv_model.py
index 07fc82b..ffa4344 100644
--- a/models/cvae_curv_model.py
+++ b/models/cvae_curv_model.py
@@ -80,12 +80,12 @@ class CVAE(tf.keras.Model):
         super(CVAE, self).__init__()
         self.latent_dim = latent_dim
 
-        model : StaticAE = StaticAE()
+        model : StaticAE = StaticAE(size=512)
         load_model(model, static_model)
         self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
         self.static_enc.trainable = False
 
-        model : BodyAE = BodyAE()
+        model : BodyAE = BodyAE(size=512)
         load_model(model, body_model)
         self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
         self.body_enc.trainable = False
diff --git a/models/cvae_model.py b/models/cvae_model.py
index 342f265..c8e48b7 100644
--- a/models/cvae_model.py
+++ b/models/cvae_model.py
@@ -9,119 +9,61 @@ import scipy as sp
 
 from custom_utils.custom_utils import load_model
 from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
-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
+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 Bottleneck(tf.keras.Model):
-    
-    def __init__(self, dim):
-        super(Bottleneck, self).__init__()
-
-        self.bottleneck = tf.keras.Sequential(
-                [Block(dim=dim, drop_path=0,
-                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
-            )
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.bottleneck(x, training=training)
-        
-        return x
-
-
-class Encoder(tf.keras.Model):
-    
-    def __init__(self, latent_dim):
-        super(Encoder, self).__init__()
-
-        self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 256], include_top=False)
-        self.flatten = tf.keras.layers.Flatten()
-        self.mean = tf.keras.layers.Dense(latent_dim)
-        self.log_var = tf.keras.layers.Dense(latent_dim)
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.encoder(x)
-        x = self.flatten(x)
-        mean = self.mean(x)
-        log_var = self.log_var(x)
-
-        return mean, log_var
-
-
-class Decoder(tf.keras.Model):
-    
-    def __init__(self):
-        super(Decoder, self).__init__()
-
-        self.dense = tf.keras.layers.Dense(16 * 8 * 80)
-        self.reshape = tf.keras.layers.Reshape((16, 8, 80))
-        self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([256, 128, 64, 16])))
-
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
-
-        return x
-
 
 class CVAE(tf.keras.Model):
     
-    def __init__(self, latent_dim=128,
-                 avg_garment_path='',
-                 static_model='',
-                 body_model='',
-                 curvature_model='', 
-                 freeze_encoders=False):
+    def __init__(self, config):
         super(CVAE, self).__init__()
-        self.latent_dim = latent_dim
+        self.latent_dim = config['latent_dim']
+        self.size = config['resolution']
 
-        model : StaticAE = StaticAE()
-        if static_model != '':
-            load_model(model, static_model)
-        self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        model : AutoEncoder = AutoEncoder(config['cond_ae'])
+        if config['static_model'] != '':
+            load_model(model, config['static_model'])
+        # self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        self.static_enc = model.encoder
 
-        model : BodyAE = BodyAE()
-        if body_model != '':
-            load_model(model, body_model)
-        self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        model : AutoEncoder = AutoEncoder(config['cond_ae'])
+        if config['body_model'] != '':
+            load_model(model, config['body_model'])
+        self.body_enc = model.encoder
 
-        model : CurvatureAE = CurvatureAE()
-        if curvature_model != '':
-            load_model(model, curvature_model)
-        self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        model : AutoEncoder = AutoEncoder(config['cond_ae'])
+        if config['curvature_model'] != '':
+            load_model(model, config['curvature_model'])
+        self.curv_enc = model.encoder
         
-        if freeze_encoders:
+        if config['freeze_encoders']:
             self.static_enc.trainable = False
             self.body_enc.trainable = False
             self.curv_enc.trainable = False
 
-        self.encoder = Encoder(latent_dim)
-        self.decoder = Decoder()
 
-        if avg_garment_path != '':
-            self.avg_garment = np.load(avg_garment_path)
+        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'])
+
+        if config['avg_garment_path'] != '':
+            self.avg_garment = np.load(config['avg_garment_path'])
         else:
             self.avg_garment = None
 
         self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
                                                  tf.keras.layers.Dense(256, activation='relu'),
                                                  tf.keras.layers.Dense(128, activation='relu'),
-                                                 tf.keras.layers.Dense(latent_dim),])
+                                                 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])]))
 
@@ -212,30 +154,28 @@ class CVAE(tf.keras.Model):
 
 
     def call(self, inputs, training=None, use_mask=False):
+
         if use_mask:
-            uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
+            if self.size == 512:
+                uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
+                uv_template_high, uv_body_posed_high, uv_curv_high = uv_template, uv_body_posed, uv_curv
+            else:
+                uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_template_high, uv_body_posed_high, uv_curv_high, uv_mask_high = inputs
         else:
-            uv_unposed, uv_template, uv_body_posed, uv_curv = inputs
-        uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
+            if self.size == 512:
+                uv_unposed, uv_template, uv_body_posed, uv_curv = inputs
+                uv_template_high, uv_body_posed_high, uv_curv_high = uv_template, uv_body_posed, uv_curv
+            else:
+                uv_unposed, uv_template, uv_body_posed, uv_curv, uv_template_high, uv_body_posed_high, uv_curv_high = inputs
 
         x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
-
-        mean, logvar = self.encoder(x)
+        mean, logvar = self.encoder(x, training=training)
         z = self.reparameterize(mean, logvar)
-
-        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
+        z_cond = self.encode_conds(uv_template_high, uv_body_posed_high, uv_curv_high)
 
         z = tf.concat([z, z_cond], axis=-1)
+        x = self.decoder(z, training=training)
 
-        # z = self.shuffle(z)
-
-        # z = tf.concat([z, z_curv], axis=-1)
-        x = self.decoder(z)
-
-        # x = tf.transpose(x, [0,3,1,2])
-        # x = x * uv_mask[:,tf.newaxis,:,:]
-        # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-        # x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
         if self.avg_garment is not None:
             x = x + self.avg_garment[tf.newaxis,:,:,:]
         if use_mask:
@@ -243,364 +183,3 @@ class CVAE(tf.keras.Model):
 
         return x, mean, logvar, z
 
-
-
-
-class Encoder32(tf.keras.Model):
-    
-    def __init__(self, latent_dim):
-        super(Encoder32, self).__init__()
-
-        self.encoder = ConvNeXt(depths=[27, 27], dims=[256, 256], include_top=False)
-        self.flatten = tf.keras.layers.Flatten()
-        self.mean = tf.keras.layers.Dense(latent_dim)
-        self.log_var = tf.keras.layers.Dense(latent_dim)
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.encoder(x)
-        x = self.flatten(x)
-        mean = self.mean(x)
-        log_var = self.log_var(x)
-
-        return mean, log_var
-
-
-class Decoder32(tf.keras.Model):
-    
-    def __init__(self):
-        super(Decoder32, self).__init__()
-
-        self.dense = tf.keras.layers.Dense(4 * 2 * 256)
-        self.reshape = tf.keras.layers.Reshape((4, 2, 256))
-        self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([256, 256])))
-
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
-
-        return x
-
-
-class CVAE32(tf.keras.Model):
-    
-    def __init__(self, latent_dim=128,
-                 avg_garment_path='',
-                 static_model='',
-                 body_model='',
-                 curvature_model='', 
-                 freeze_encoders=False,
-                 use_full_res_encoders=True):
-        super(CVAE32, self).__init__()
-        self.latent_dim = latent_dim
-
-        if use_full_res_encoders:
-            model : StaticAE = StaticAE()
-        else:
-            model : StaticAE32 = StaticAE32()
-        if static_model != '':
-            load_model(model, static_model)
-        self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
-
-        if use_full_res_encoders:
-            model : BodyAE = BodyAE()
-        else:
-            model : BodyAE32 = BodyAE32()
-        if body_model != '':
-            load_model(model, body_model)
-        self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
-
-        if use_full_res_encoders:
-            model : CurvatureAE = CurvatureAE()
-        else:
-            model : CurvatureAE32 = CurvatureAE32()
-        if curvature_model != '':
-            load_model(model, curvature_model)
-        self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
-        
-        if freeze_encoders:
-            self.static_enc.trainable = False
-            self.body_enc.trainable = False
-            self.curv_enc.trainable = False
-
-        self.encoder = Encoder32(latent_dim)
-        self.decoder = Decoder32()
-
-        if avg_garment_path != '':
-            self.avg_garment = np.load(avg_garment_path)
-        else:
-            self.avg_garment = None
-
-        self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
-                                                 tf.keras.layers.Dense(256, activation='relu'),
-                                                 tf.keras.layers.Dense(128, activation='relu'),
-                                                 tf.keras.layers.Dense(latent_dim),])
-
-
-    # @tf.function
-    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False, zero_latent=False):
-        if not fix_noise:
-            eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
-        else:
-            eps = tf.random.normal(shape=(1, self.latent_dim))
-            eps = tf.tile(eps, [nr_of_samples,1])
-
-        uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
-        uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
-        if len(uv_curv.shape) == 3:
-            uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
-        uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1])
-
-        # uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
-        z_static = self.static_enc(uv_template)
-        z_body = self.body_enc(uv_body_posed)
-        z_curv = self.curv_enc(uv_curv)
-
-        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
-        z_cond = self.cond_encoder(z_cond)
-
-        if zero_latent:
-            eps = tf.zeros_like(eps)
-
-        z = tf.concat([eps, z_cond], axis=-1)
-
-        x = self.decoder(z)
-        x = x + self.avg_garment[tf.newaxis,:,:,:]
-        x = x * uv_mask[:,:,:,tf.newaxis]
-
-        return x
-
-
-    def reparameterize(self, mean, logvar):
-        eps = tf.random.normal(shape=mean.shape)
-        return eps * tf.exp(logvar * .5) + mean
-
-
-    def decode(self, z, apply_sigmoid=False):
-        logits = self.decoder(z)
-        if apply_sigmoid:
-            probs = tf.sigmoid(logits)
-            return probs
-        return logits
-
-
-    def call(self, inputs, training=None, zero_latent=False):
-        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_template_high, uv_body_posed_high, uv_curv_high = inputs
-        uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
-        z_static = self.static_enc(uv_template_high)
-        z_body = self.body_enc(uv_body_posed_high)
-        z_curv = self.curv_enc(uv_curv_high)
-
-        x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
-
-        mean, logvar = self.encoder(x)
-        z = self.reparameterize(mean, logvar)
-        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
-        z_cond = self.cond_encoder(z_cond)
-
-        if zero_latent:
-            z = tf.zeros_like(z)
-
-        z = tf.concat([z, z_cond], axis=-1)
-
-        # z = self.shuffle(z)
-
-        # z = tf.concat([z, z_curv], axis=-1)
-        x = self.decoder(z)
-
-
-        # x = tf.transpose(x, [0,3,1,2])
-        # x = x * uv_mask[:,tf.newaxis,:,:]
-        # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-        # x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
-        if self.avg_garment is not None:
-            x = x + self.avg_garment[tf.newaxis,:,:,:]
-        # x = x * uv_mask[:,:,:,tf.newaxis]
-
-        return x, mean, logvar, z
-
-
-
-
-class Encoder128(tf.keras.Model):
-    
-    def __init__(self, latent_dim):
-        super(Encoder128, self).__init__()
-
-        self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[256, 256, 128, 64], include_top=False)
-        self.flatten = tf.keras.layers.Flatten()
-        self.mean = tf.keras.layers.Dense(latent_dim)
-        self.log_var = tf.keras.layers.Dense(latent_dim)
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.encoder(x)
-        x = self.flatten(x)
-        mean = self.mean(x)
-        log_var = self.log_var(x)
-
-        return mean, log_var
-
-
-class Decoder128(tf.keras.Model):
-    
-    def __init__(self):
-        super(Decoder128, self).__init__()
-
-        self.dense = tf.keras.layers.Dense(4 * 2 * 128)
-        self.reshape = tf.keras.layers.Reshape((4, 2, 128))
-        self.decoder = ConvNeXtReverse(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([256, 256, 128, 64])))
-
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
-
-        return x
-
-
-class CVAE128(tf.keras.Model):
-    
-    def __init__(self, latent_dim=128,
-                 avg_garment_path='',
-                 static_model='',
-                 body_model='',
-                 curvature_model='', 
-                 freeze_encoders=False,
-                 use_full_res_encoders=True):
-        super(CVAE128, self).__init__()
-        self.latent_dim = latent_dim
-
-        if use_full_res_encoders:
-            model : StaticAE = StaticAE()
-        else:
-            model : StaticAE128 = StaticAE128()
-        if static_model != '':
-            load_model(model, static_model)
-        self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
-        
-        if use_full_res_encoders:
-            model : BodyAE = BodyAE()
-        else:
-            model : BodyAE128 = BodyAE128()
-        if body_model != '':
-            load_model(model, body_model)
-        self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
-
-        if use_full_res_encoders:
-            model : CurvatureAE = CurvatureAE()
-        else:
-            model : CurvatureAE128 = CurvatureAE128()
-        if curvature_model != '':
-            load_model(model, curvature_model)
-        self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
-        
-        if freeze_encoders:
-            self.static_enc.trainable = False
-            self.body_enc.trainable = False
-            self.curv_enc.trainable = False
-
-        self.encoder = Encoder128(latent_dim)
-        self.decoder = Decoder128()
-
-        if avg_garment_path != '':
-            self.avg_garment = np.load(avg_garment_path)
-        else:
-            self.avg_garment = None
-
-        self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
-                                                 tf.keras.layers.Dense(256, activation='relu'),
-                                                 tf.keras.layers.Dense(128, activation='relu'),
-                                                 tf.keras.layers.Dense(latent_dim),])
-
-
-    # @tf.function
-    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False, zero_latent=False):
-        if not fix_noise:
-            eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
-        else:
-            eps = tf.random.normal(shape=(1, self.latent_dim))
-            eps = tf.tile(eps, [nr_of_samples,1])
-
-        uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
-        uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
-        if len(uv_curv.shape) == 3:
-            uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
-        uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1])
-
-        # uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
-        z_static = self.static_enc(uv_template)
-        z_body = self.body_enc(uv_body_posed)
-        z_curv = self.curv_enc(uv_curv)
-
-        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
-        z_cond = self.cond_encoder(z_cond)
-
-        if zero_latent:
-            eps = tf.zeros_like(eps)
-
-        z = tf.concat([eps, z_cond], axis=-1)
-
-        x = self.decoder(z)
-        x = x + self.avg_garment[tf.newaxis,:,:,:]
-        x = x * uv_mask[:,:,:,tf.newaxis]
-
-        return x
-
-
-    def reparameterize(self, mean, logvar):
-        eps = tf.random.normal(shape=mean.shape)
-        return eps * tf.exp(logvar * .5) + mean
-
-
-    def decode(self, z, apply_sigmoid=False):
-        logits = self.decoder(z)
-        if apply_sigmoid:
-            probs = tf.sigmoid(logits)
-            return probs
-        return logits
-
-
-    def call(self, inputs, training=None, zero_latent=False):
-        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_template_high, uv_body_posed_high, uv_curv_high = inputs
-        uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
-        z_static = self.static_enc(uv_template_high)
-        z_body = self.body_enc(uv_body_posed_high)
-        z_curv = self.curv_enc(uv_curv_high)
-
-        x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
-
-        mean, logvar = self.encoder(x)
-        z = self.reparameterize(mean, logvar)
-        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
-        z_cond = self.cond_encoder(z_cond)
-
-        if zero_latent:
-            z = tf.zeros_like(z)
-
-        z = tf.concat([z, z_cond], axis=-1)
-
-        # z = self.shuffle(z)
-
-        # z = tf.concat([z, z_curv], axis=-1)
-        x = self.decoder(z)
-
-
-        # x = tf.transpose(x, [0,3,1,2])
-        # x = x * uv_mask[:,tf.newaxis,:,:]
-        # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-        # x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
-        if self.avg_garment is not None:
-            x = x + self.avg_garment[tf.newaxis,:,:,:]
-        # x = x * uv_mask[:,:,:,tf.newaxis]
-
-        return x, mean, logvar, z
\ No newline at end of file
diff --git a/models/cvae_model_old.py b/models/cvae_model_old.py
new file mode 100644
index 0000000..1833bc9
--- /dev/null
+++ b/models/cvae_model_old.py
@@ -0,0 +1,621 @@
+import os
+import time
+import random
+import tensorflow as tf
+import numpy as np
+import scipy as sp
+# import plotly.graph_objects as go
+# import plotly.express as ex
+
+from custom_utils.custom_utils import load_model
+from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
+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 Bottleneck(tf.keras.Model):
+    
+    def __init__(self, dim):
+        super(Bottleneck, self).__init__()
+
+        self.bottleneck = tf.keras.Sequential(
+                [Block(dim=dim, drop_path=0,
+                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+            )
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.bottleneck(x, training=training)
+        
+        return x
+
+
+class Encoder(tf.keras.Model):
+    
+    def __init__(self, latent_dim):
+        super(Encoder, self).__init__()
+
+        self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 256], include_top=False)
+        self.flatten = tf.keras.layers.Flatten()
+        self.mean = tf.keras.layers.Dense(latent_dim)
+        self.log_var = tf.keras.layers.Dense(latent_dim)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.encoder(x)
+        x = self.flatten(x)
+        mean = self.mean(x)
+        log_var = self.log_var(x)
+
+        return mean, log_var
+
+
+class Decoder(tf.keras.Model):
+    
+    def __init__(self):
+        super(Decoder, self).__init__()
+
+        self.dense = tf.keras.layers.Dense(16 * 8 * 80)
+        self.reshape = tf.keras.layers.Reshape((16, 8, 80))
+        self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([256, 128, 64, 16])))
+
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.dense(x)
+        x = self.reshape(x)
+        x = self.decoder(x)
+
+        return x
+
+
+class CVAE(tf.keras.Model):
+    
+    def __init__(self, latent_dim=128,
+                 avg_garment_path='',
+                 static_model='',
+                 body_model='',
+                 curvature_model='', 
+                 freeze_encoders=False):
+        super(CVAE, self).__init__()
+        self.latent_dim = latent_dim
+
+        model : StaticAE = StaticAE()
+        if static_model != '':
+            load_model(model, static_model)
+        self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+
+        model : BodyAE = BodyAE()
+        if body_model != '':
+            load_model(model, body_model)
+        self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+
+        model : CurvatureAE = CurvatureAE()
+        if curvature_model != '':
+            load_model(model, curvature_model)
+        self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        
+        if freeze_encoders:
+            self.static_enc.trainable = False
+            self.body_enc.trainable = False
+            self.curv_enc.trainable = False
+
+        self.encoder = Encoder(latent_dim)
+        self.decoder = Decoder()
+
+        if avg_garment_path != '':
+            self.avg_garment = np.load(avg_garment_path)
+        else:
+            self.avg_garment = None
+
+        self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
+                                                 tf.keras.layers.Dense(256, activation='relu'),
+                                                 tf.keras.layers.Dense(128, activation='relu'),
+                                                 tf.keras.layers.Dense(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):
+        if not fix_noise:
+            eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
+        else:
+            eps = tf.random.normal(shape=(1, self.latent_dim))
+            eps = tf.tile(eps, [nr_of_samples,1])
+
+        uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
+        uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
+        if len(uv_curv.shape) == 3:
+            uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
+        uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1])
+
+        # uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
+        z_static = self.static_enc(uv_template)
+        z_body = self.body_enc(uv_body_posed)
+        z_curv = self.curv_enc(uv_curv)
+
+        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+        z_cond = self.cond_encoder(z_cond)
+
+        z = tf.concat([eps, z_cond], axis=-1)
+
+        x = self.decoder(z)
+        x = x + self.avg_garment[tf.newaxis,:,:,:]
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+
+    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask=None, noise=None):
+        if noise is None:
+            eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
+        else:
+            eps = noise
+
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
+        z = tf.concat([eps, z_cond], axis=-1)
+
+        x = self.decoder(z)
+        
+        if self.avg_garment is not None:
+            x = x + self.avg_garment[tf.newaxis,:,:,:]
+        
+        if uv_mask is not None:
+            x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+
+    def reparameterize(self, mean, logvar, fix_noise=False):
+        if not fix_noise:
+            eps = tf.random.normal(shape=mean.shape)
+        else:
+            eps = tf.random.normal(shape=(1, mean.shape[1]))
+            eps = tf.tile(eps, [mean.shape[0],1])
+        return eps * tf.exp(logvar * .5) + mean
+
+
+    def decode(self, z, apply_sigmoid=False):
+        logits = self.decoder(z)
+        if apply_sigmoid:
+            probs = tf.sigmoid(logits)
+            return probs
+        return logits
+
+
+    def encode_conds(self, uv_template, uv_body_posed, uv_curv):
+        if uv_template.shape[1] == 3:
+            uv_template = tf.transpose(uv_template, [0,2,3,1])
+        if uv_body_posed.shape[1] == 3:
+            uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
+        if uv_curv.shape[1] == 3:
+            uv_curv = tf.transpose(uv_curv, [0,2,3,1])
+
+        z_static = self.static_enc(uv_template)
+        z_body = self.body_enc(uv_body_posed)
+        z_curv = self.curv_enc(uv_curv)
+
+        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+        z_cond = self.cond_encoder(z_cond)
+
+        return z_cond
+
+
+    def call(self, inputs, training=None, use_mask=False):
+        if use_mask:
+            uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
+        else:
+            uv_unposed, uv_template, uv_body_posed, uv_curv = inputs
+        uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
+
+        x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
+
+        mean, logvar = self.encoder(x)
+        z = self.reparameterize(mean, logvar)
+
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
+
+        z = tf.concat([z, z_cond], axis=-1)
+
+        # z = self.shuffle(z)
+
+        # z = tf.concat([z, z_curv], axis=-1)
+        x = self.decoder(z)
+
+        # x = tf.transpose(x, [0,3,1,2])
+        # x = x * uv_mask[:,tf.newaxis,:,:]
+        # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+        # x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
+        if self.avg_garment is not None:
+            x = x + self.avg_garment[tf.newaxis,:,:,:]
+        if use_mask:
+            x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x, mean, logvar, z
+
+
+
+class Encoder32(tf.keras.Model):
+    
+    def __init__(self, latent_dim):
+        super(Encoder32, self).__init__()
+
+        self.encoder = ConvNeXt(depths=[27, 27], dims=[256, 256], include_top=False)
+        self.flatten = tf.keras.layers.Flatten()
+        self.mean = tf.keras.layers.Dense(latent_dim)
+        self.log_var = tf.keras.layers.Dense(latent_dim)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.encoder(x)
+        x = self.flatten(x)
+        mean = self.mean(x)
+        log_var = self.log_var(x)
+
+        return mean, log_var
+
+
+class Decoder32(tf.keras.Model):
+    
+    def __init__(self):
+        super(Decoder32, self).__init__()
+
+        self.dense = tf.keras.layers.Dense(4 * 2 * 256)
+        self.reshape = tf.keras.layers.Reshape((4, 2, 256))
+        self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([256, 256])))
+
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.dense(x)
+        x = self.reshape(x)
+        x = self.decoder(x)
+
+        return x
+
+
+class CVAE32(tf.keras.Model):
+    
+    def __init__(self, latent_dim=128,
+                 avg_garment_path='',
+                 static_model='',
+                 body_model='',
+                 curvature_model='', 
+                 freeze_encoders=False,
+                 use_full_res_encoders=True):
+        super(CVAE32, self).__init__()
+        self.latent_dim = latent_dim
+
+        if use_full_res_encoders:
+            model : StaticAE = StaticAE()
+        else:
+            model : StaticAE32 = StaticAE32()
+        if static_model != '':
+            load_model(model, static_model)
+        self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+
+        if use_full_res_encoders:
+            model : BodyAE = BodyAE()
+        else:
+            model : BodyAE32 = BodyAE32()
+        if body_model != '':
+            load_model(model, body_model)
+        self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+
+        if use_full_res_encoders:
+            model : CurvatureAE = CurvatureAE()
+        else:
+            model : CurvatureAE32 = CurvatureAE32()
+        if curvature_model != '':
+            load_model(model, curvature_model)
+        self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        
+        if freeze_encoders:
+            self.static_enc.trainable = False
+            self.body_enc.trainable = False
+            self.curv_enc.trainable = False
+
+        self.encoder = Encoder32(latent_dim)
+        self.decoder = Decoder32()
+
+        if avg_garment_path != '':
+            self.avg_garment = np.load(avg_garment_path)
+        else:
+            self.avg_garment = None
+
+        self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
+                                                 tf.keras.layers.Dense(256, activation='relu'),
+                                                 tf.keras.layers.Dense(128, activation='relu'),
+                                                 tf.keras.layers.Dense(latent_dim),])
+
+
+    # @tf.function
+    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False, zero_latent=False):
+        if not fix_noise:
+            eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
+        else:
+            eps = tf.random.normal(shape=(1, self.latent_dim))
+            eps = tf.tile(eps, [nr_of_samples,1])
+
+        uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
+        uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
+        if len(uv_curv.shape) == 3:
+            uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
+        uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1])
+
+        # uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
+        z_static = self.static_enc(uv_template)
+        z_body = self.body_enc(uv_body_posed)
+        z_curv = self.curv_enc(uv_curv)
+
+        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+        z_cond = self.cond_encoder(z_cond)
+
+        if zero_latent:
+            eps = tf.zeros_like(eps)
+
+        z = tf.concat([eps, z_cond], axis=-1)
+
+        x = self.decoder(z)
+        x = x + self.avg_garment[tf.newaxis,:,:,:]
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask=None, noise=None):
+        if noise is None:
+            eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
+        else:
+            eps = noise
+
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
+        z = tf.concat([eps, z_cond], axis=-1)
+
+        x = self.decoder(z)
+        
+        if self.avg_garment is not None:
+            x = x + self.avg_garment[tf.newaxis,:,:,:]
+        
+        if uv_mask is not None:
+            x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+    def reparameterize(self, mean, logvar):
+        eps = tf.random.normal(shape=mean.shape)
+        return eps * tf.exp(logvar * .5) + mean
+
+
+    def encode_conds(self, uv_template, uv_body_posed, uv_curv):
+        if uv_template.shape[1] == 3:
+            uv_template = tf.transpose(uv_template, [0,2,3,1])
+        if uv_body_posed.shape[1] == 3:
+            uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
+        if uv_curv.shape[1] == 3:
+            uv_curv = tf.transpose(uv_curv, [0,2,3,1])
+
+        z_static = self.static_enc(uv_template)
+        z_body = self.body_enc(uv_body_posed)
+        z_curv = self.curv_enc(uv_curv)
+
+        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+        z_cond = self.cond_encoder(z_cond)
+
+        return z_cond
+
+    def call(self, inputs, training=None, use_mask=False):
+        if use_mask:
+            uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
+        else:
+            uv_unposed, uv_template, uv_body_posed, uv_curv = inputs
+        uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
+
+
+        x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
+        mean, logvar = self.encoder(x)
+        z = self.reparameterize(mean, logvar)
+
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
+
+        z = tf.concat([z, z_cond], axis=-1)
+        x = self.decoder(z)
+
+        if self.avg_garment is not None:
+            x = x + self.avg_garment[tf.newaxis,:,:,:]
+        if use_mask:
+            x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x, mean, logvar, z
+
+
+
+
+class Encoder128(tf.keras.Model):
+    
+    def __init__(self, latent_dim):
+        super(Encoder128, self).__init__()
+
+        self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[256, 256, 128, 64], include_top=False)
+        self.flatten = tf.keras.layers.Flatten()
+        self.mean = tf.keras.layers.Dense(latent_dim)
+        self.log_var = tf.keras.layers.Dense(latent_dim)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.encoder(x)
+        x = self.flatten(x)
+        mean = self.mean(x)
+        log_var = self.log_var(x)
+
+        return mean, log_var
+
+
+class Decoder128(tf.keras.Model):
+    
+    def __init__(self):
+        super(Decoder128, self).__init__()
+
+        self.dense = tf.keras.layers.Dense(4 * 2 * 128)
+        self.reshape = tf.keras.layers.Reshape((4, 2, 128))
+        self.decoder = ConvNeXtReverse(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([256, 256, 128, 64])))
+
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.dense(x)
+        x = self.reshape(x)
+        x = self.decoder(x)
+
+        return x
+
+
+class CVAE128(tf.keras.Model):
+    
+    def __init__(self, latent_dim=128,
+                 avg_garment_path='',
+                 static_model='',
+                 body_model='',
+                 curvature_model='', 
+                 freeze_encoders=False,
+                 use_full_res_encoders=True):
+        super(CVAE128, self).__init__()
+        self.latent_dim = latent_dim
+
+        if use_full_res_encoders:
+            model : StaticAE = StaticAE()
+        else:
+            model : StaticAE128 = StaticAE128()
+        if static_model != '':
+            load_model(model, static_model)
+        self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        
+        if use_full_res_encoders:
+            model : BodyAE = BodyAE()
+        else:
+            model : BodyAE128 = BodyAE128()
+        if body_model != '':
+            load_model(model, body_model)
+        self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+
+        if use_full_res_encoders:
+            model : CurvatureAE = CurvatureAE()
+        else:
+            model : CurvatureAE128 = CurvatureAE128()
+        if curvature_model != '':
+            load_model(model, curvature_model)
+        self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        
+        if freeze_encoders:
+            self.static_enc.trainable = False
+            self.body_enc.trainable = False
+            self.curv_enc.trainable = False
+
+        self.encoder = Encoder128(latent_dim)
+        self.decoder = Decoder128()
+
+        if avg_garment_path != '':
+            self.avg_garment = np.load(avg_garment_path)
+        else:
+            self.avg_garment = None
+
+        self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
+                                                 tf.keras.layers.Dense(256, activation='relu'),
+                                                 tf.keras.layers.Dense(128, activation='relu'),
+                                                 tf.keras.layers.Dense(latent_dim),])
+
+
+    # @tf.function
+    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False, zero_latent=False):
+        if not fix_noise:
+            eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
+        else:
+            eps = tf.random.normal(shape=(1, self.latent_dim))
+            eps = tf.tile(eps, [nr_of_samples,1])
+
+        uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
+        uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
+        if len(uv_curv.shape) == 3:
+            uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
+        uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1])
+
+        # uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
+        z_static = self.static_enc(uv_template)
+        z_body = self.body_enc(uv_body_posed)
+        z_curv = self.curv_enc(uv_curv)
+
+        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+        z_cond = self.cond_encoder(z_cond)
+
+        if zero_latent:
+            eps = tf.zeros_like(eps)
+
+        z = tf.concat([eps, z_cond], axis=-1)
+
+        x = self.decoder(z)
+        x = x + self.avg_garment[tf.newaxis,:,:,:]
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+
+    def reparameterize(self, mean, logvar):
+        eps = tf.random.normal(shape=mean.shape)
+        return eps * tf.exp(logvar * .5) + mean
+
+
+    def decode(self, z, apply_sigmoid=False):
+        logits = self.decoder(z)
+        if apply_sigmoid:
+            probs = tf.sigmoid(logits)
+            return probs
+        return logits
+
+
+    def call(self, inputs, training=None, zero_latent=False):
+        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_template_high, uv_body_posed_high, uv_curv_high = inputs
+        uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
+        z_static = self.static_enc(uv_template_high)
+        z_body = self.body_enc(uv_body_posed_high)
+        z_curv = self.curv_enc(uv_curv_high)
+
+        x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
+
+        mean, logvar = self.encoder(x)
+        z = self.reparameterize(mean, logvar)
+        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+        z_cond = self.cond_encoder(z_cond)
+
+        if zero_latent:
+            z = tf.zeros_like(z)
+
+        z = tf.concat([z, z_cond], axis=-1)
+
+        # z = self.shuffle(z)
+
+        # z = tf.concat([z, z_curv], axis=-1)
+        x = self.decoder(z)
+
+
+        # x = tf.transpose(x, [0,3,1,2])
+        # x = x * uv_mask[:,tf.newaxis,:,:]
+        # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+        # x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
+        if self.avg_garment is not None:
+            x = x + self.avg_garment[tf.newaxis,:,:,:]
+        # x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x, mean, logvar, z
\ No newline at end of file
diff --git a/models/cvae_pyramid_model.py b/models/cvae_pyramid_model.py
index 922b9fc..4f6d80a 100644
--- a/models/cvae_pyramid_model.py
+++ b/models/cvae_pyramid_model.py
@@ -7,7 +7,7 @@ import numpy as np
 # import plotly.express as ex
 
 from custom_utils.custom_utils import load_model
-from models.cvae_model import CVAE, CVAE128, CVAE32
+from models.cvae_model import CVAE
 from models.upscaler import UpscalerBasic as Upscaler
 
 
@@ -24,26 +24,22 @@ class Pyramid(tf.keras.Model):
                  use_full_res_encoders=True):
         super(Pyramid, self).__init__()
 
-        if low_res == 32:
-            self.model_low = CVAE32(latent_dim=latent_dim,
-                                    static_model=static_model, 
-                                    body_model=body_model, 
-                                    curvature_model=curvature_model, 
-                                    freeze_encoders=freeze_encoders, 
-                                    use_full_res_encoders=use_full_res_encoders)
-        if low_res == 128:
-            self.model_low = CVAE128(latent_dim=latent_dim,
-                                    static_model=static_model, 
-                                    body_model=body_model, 
-                                    curvature_model=curvature_model, 
-                                    freeze_encoders=freeze_encoders, 
-                                    use_full_res_encoders=use_full_res_encoders)
+
+        self.model_low = CVAE(latent_dim=latent_dim,
+                              size=low_res,
+                              static_model=static_model, 
+                              body_model=body_model, 
+                              curvature_model=curvature_model, 
+                              freeze_encoders=freeze_encoders, 
+                              encoder_res=512 if use_full_res_encoders else 32)
 
         self.model_full = CVAE(latent_dim=latent_dim,
+                               size=512,
                                static_model=static_model, 
                                body_model=body_model, 
                                curvature_model=curvature_model, 
-                               freeze_encoders=freeze_encoders)
+                               freeze_encoders=freeze_encoders,
+                               encoder_res=512)
 
         self.upscale_enhancer = None
         if upscaler_model != '':
@@ -64,6 +60,12 @@ class Pyramid(tf.keras.Model):
         else:
             self.avg_garment = None
 
+    def enhance(self, low, mask, training=None):
+        x = self.upscale_enhancer(low, training=training)
+        x = x * (1 - mask[:,:,:,tf.newaxis])
+        low = x + low
+        return low
+
     def call(self, inputs, training=None):
         uv_unposed, uv_static, uv_body_posed, uv_curv, uv_mask, uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low, uv_mask_low = inputs
 
@@ -75,11 +77,9 @@ class Pyramid(tf.keras.Model):
         pred_low = pred_low * uv_mask_low[:,:,:,tf.newaxis]
         
         if self.upscale_enhancer is not None:
-            x = self.upscale_enhancer(pred_low, training=training)
-            x = x * (1 - uv_mask_low[:,:,:,tf.newaxis])
-            pred_low = x + pred_low
+            pred_low = self.enhance(pred_low, uv_mask_low, training=training)
         
-        x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+        x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
         x = x * uv_mask[:,:,:,tf.newaxis]
 
         pred_full, mean_full, logvar_full, _ = self.model_full([uv_unposed, uv_static, uv_body_posed, uv_curv], training=training)
diff --git a/sample_pro.py b/sample_pro.py
index bfe2eba..250beb6 100644
--- a/sample_pro.py
+++ b/sample_pro.py
@@ -1,5 +1,5 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+os.environ["CUDA_VISIBLE_DEVICES"] = "3"
 import tensorflow as tf
 
 gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -27,6 +27,7 @@ from PIL import Image
 from plotly.subplots import make_subplots
 
 from custom_utils.smpl import SMPL
+from custom_utils.custom_utils import load_model
 from cvae_curv import CVAE as CVAECurv
 from models.cvae_model import CVAE as CVAEMain
 
@@ -136,12 +137,6 @@ def plot_everything(batch: dict, type: str):
     return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
 
 
-def load_model(model, checkpoint):
-    ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
-    ckpt.restore(manager.latest_checkpoint)
-
-
 def get_frame(folder, nr):
     pose_folder = os.path.join(folder, 'pose_shape') 
     frames = sorted(os.listdir(pose_folder))
@@ -391,6 +386,7 @@ def sample(folder, variable_condition, fix_noise, nr=4):
         assert False
 
     plot_variable_condition(variable_cond_data, mesh_mask, is_curv=random_curv, is_body=random_body)
+    plot_variable_condition(curv, mesh_mask, is_curv=True, is_body=False)
 
     plot_results(mesh_pred_shaped, mesh_posed, body, mesh_mask)
 
@@ -407,7 +403,7 @@ if __name__ == '__main__':
     args = vars(parser.parse_args())
 
     folder = '/home/workspace/local-data/train_100'
-    folder_main = 'checkpoints/cvae_13'
+    folder_main = 'checkpoints/cvae_18_copy'
     folder_curv = 'checkpoints/cvae_curv_3'
     variable_condition = 'static'
     batch_size = 4
@@ -415,13 +411,17 @@ if __name__ == '__main__':
     smpl = SMPL()
 
     model = CVAEMain(avg_garment_path='custom_data/train_100_no_shape_avg_uv.npy',
-                    latent_dim=128, static_model='checkpoints/ae_static_4',
-                    body_model='checpoints/ae_body_8',
-                    curvature_model='checpoints/ae_curvature_7')
-    curv_model = CVAECurv()
+                    latent_dim=128, static_model='checkpoints/ae_static_6_copy',
+                    body_model='checkpoints/ae_body_8_copy',
+                    curvature_model='checkpoints/ae_curvature_7_copy',
+                    freeze_encoders=True)
+    curv_model = CVAECurv(static_model='checkpoints/ae_static_6_copy',
+                          body_model='checkpoints/ae_body_8_copy',)
 
 
     load_model(model, folder_main)
     load_model(curv_model, folder_curv)
 
     sample(folder, variable_condition, fix_noise=True)
+
+    print()
diff --git a/vae.py b/vae.py
index 35cee89..72366bf 100644
--- a/vae.py
+++ b/vae.py
@@ -22,7 +22,7 @@ from stylegan2_ae.generator import Generator
 from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder, sparse_to_tensor
 from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
 from dataset import DatasetFactory, StaticDataset
-from custom_utils.custom_utils import batch_to_dict
+from custom_utils.custom_utils import batch_to_dict, plot_3d
 from custom_utils.smpl import SMPL
 
 name='vae_11'
@@ -38,53 +38,6 @@ run.log_artifact(artifact)
 
 uv2mesh = UV2Mesh(res=256)
 
-def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
-    mesh_pred_unposed[:,1] += 1
-    collision[:,1] += 2
-    fig = go.Figure(data=[
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=mesh_gt[:,0],
-            y=mesh_gt[:,1],
-            z=mesh_gt[:,2],
-            # i, j and k give the vertices of triangles
-            i = faces[:,0],
-            j = faces[:,1],
-            k = faces[:,2]
-        ),
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=mesh_pred_unposed[:,0],
-            y=mesh_pred_unposed[:,1],
-            z=mesh_pred_unposed[:,2],
-            # i, j and k give the vertices of triangles
-            i = faces[:,0],
-            j = faces[:,1],
-            k = faces[:,2],
-        ),
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=collision[:,0],
-            y=collision[:,1],
-            z=collision[:,2],
-            # i, j and k give the vertices of triangles
-            i = collision_faces[:,0],
-            j = collision_faces[:,1],
-            k = collision_faces[:,2],
-            vertexcolor=collision_color
-        )
-    ])
-
-    fig.update_layout(scene_aspectmode='manual',
-                    scene_aspectratio=dict(x=1, y=2, z=1))
-    fig.update_layout(
-        scene = dict(
-            xaxis = dict(nticks=4, range=[-1,1],),
-                        yaxis = dict(nticks=4, range=[-1,3],),
-                        zaxis = dict(nticks=4, range=[-1,1],),)
-                        )
-
-    return fig
 
 class Bottleneck(tf.keras.Model):
     
diff --git a/verify_upscaler.py b/verify_upscaler.py
index 5167c2c..91abb5c 100644
--- a/verify_upscaler.py
+++ b/verify_upscaler.py
@@ -2,7 +2,7 @@
 import os
 
 from models.upscaler import UpscalerBasic as Upscaler
-os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+os.environ["CUDA_VISIBLE_DEVICES"] = "3"
 
 import tqdm
 import numpy as np