diff --git a/cvae_pyramid.py b/cvae_pyramid.py
index 81b8717..e4ba70b 100644
--- a/cvae_pyramid.py
+++ b/cvae_pyramid.py
@@ -16,12 +16,13 @@ if gpus:
         print(e)
 
 import time
+import tqdm
 import argparse
 import wandb
 import random
 import numpy as np
 import scipy as sp
-from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces
+from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, get_summary_str
 import plotly.graph_objects as go
 import plotly.express as ex
 from plotly.subplots import make_subplots
@@ -33,17 +34,17 @@ else:
     from losses import edge_loss
 
 from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
-from losses import uv_mask_loss, bend_loss, kl_loss
+from losses import uv_mask_loss, bend_loss, kl_loss, collision_ratio
 # 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.pyramid_model import Pyramid
+from models.pyramid_model_optimized import Pyramid
 
 
 def init_wandb():
-    run = wandb.init(project='cvae', name=config['name'], notes=config['notes'])
+    run = wandb.init(project='cvae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
 
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae_pyramid.py')
@@ -60,7 +61,7 @@ def init_wandb():
     artifact.add_file('convnext/models/convnext_tf.py')
     run.log_artifact(artifact)
 
-uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'])
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
 
 def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, collision_faces, collision_color):
     fig = make_subplots(rows=1, cols=1+len(mesh_pred_unposed_levels)+1+1, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(1+len(mesh_pred_unposed_levels)+1+1)] ])
@@ -215,9 +216,9 @@ def group_inputs_for_levels(batch):
 
 @tf.function
 def train_step(batch):
+    levels_input, conds = group_inputs_for_levels(batch)
 
     with tf.GradientTape() as tape:
-        levels_input, conds = group_inputs_for_levels(batch)
         out_levels = model([levels_input, conds], training=True)
 
         shape = batch['pose_shape'][:, 72:]
@@ -246,9 +247,10 @@ def train_step(batch):
         l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
         # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
         ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
+        l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
         kl = kl_loss(mean, logvar)
 
-        reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
+        reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
         # reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
         kl_total += kl
 
@@ -260,7 +262,7 @@ def train_step(batch):
         output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric, 
                                     'train/ms_ssim' + res_modifier: ms_ssim, 'train/mean' + res_modifier: mean, 'train/logvar' + res_modifier: logvar, 'train/kl_loss' + res_modifier: kl,
                                     'train/var' + res_modifier: np.e ** logvar, 'train/l2_loss' + res_modifier: l2_loss,
-                                    'train/euclid_error': metric } }
+                                    'train/euclid_error': metric, 'train/l2_mesh_loss': l2_mesh_loss } }
 
         loss = reconstr_loss_total + balancing_coeff * kl_total
 
@@ -300,9 +302,10 @@ def valid_step(batch):
     l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
     # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
     ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
+    l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
     kl = kl_loss(mean, logvar)
 
-    reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
+    reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
     # reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
     kl_total += kl
 
@@ -314,13 +317,68 @@ def valid_step(batch):
     output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss,'valid/euclid_error' + res_modifier: metric, 
                                 'valid/ms_ssim' + res_modifier: ms_ssim, 'valid/mean' + res_modifier: mean, 'valid/logvar' + res_modifier: logvar, 'valid/kl_loss' + res_modifier: kl,
                                 'valid/var' + res_modifier: np.e ** logvar, 'valid/l2_loss' + res_modifier: l2_loss,
-                                'valid/euclid_error': metric } }
+                                'valid/euclid_error': metric, 'valid/l2_mesh_loss': l2_mesh_loss } }
 
     loss = reconstr_loss_total + balancing_coeff * kl_total
 
     return { **{ 'valid/loss': loss, 'valid/reconstr_loss_total': reconstr_loss_total, 'valid/kl_total': kl_total }, **output }
 
 
+@tf.function
+def test_step(batch):
+    levels_input, conds = group_inputs_for_levels(batch)
+    out_levels = model([levels_input, conds], training=False)
+
+    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_orig = batch['mesh_unposed']
+
+    output = dict()
+    reconstr_loss_total = 0
+    kl_total = 0
+
+    #for i, res in enumerate(config['pyramid']['resolutions']):
+    i = -1
+    res = config['pyramid']['resolutions'][-1]
+    res_modifier = '_' + str(res) if res != 512 else ''
+    pred, mean, logvar = out_levels[i]
+
+    mesh_pred = uv2mesh(pred)
+    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)
+
+    # is the output already masked?
+
+    e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+    c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
+    c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
+    l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
+    # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
+    ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
+    kl = kl_loss(mean, logvar)
+
+    reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
+    # reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
+    kl_total += kl
+
+    mean = tf.reduce_mean(mean)
+    logvar = tf.reduce_mean(logvar)
+
+    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    euclid_error_per_garment = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1)), axis=1) / tf.reduce_sum(batch['mesh_mask'], axis=1)
+
+    output = { **output, **{ 'test/edge_loss' + res_modifier: e_loss, 'test/collision_loss' + res_modifier: c_loss,'test/euclid_error' + res_modifier: metric, 
+                                'test/ms_ssim' + res_modifier: ms_ssim, 'test/mean' + res_modifier: mean, 'test/logvar' + res_modifier: logvar, 'test/kl_loss' + res_modifier: kl,
+                                'test/var' + res_modifier: np.e ** logvar, 'test/l2_loss' + res_modifier: l2_loss,
+                                'test/euclid_error': metric, 'test/collision_ratio' + res_modifier: c_ratio } }
+
+    loss = reconstr_loss_total + balancing_coeff * kl_total
+
+    return { **{ 'test/loss': loss, 'test/reconstr_loss_total': reconstr_loss_total, 'test/kl_total': kl_total }, **output }, euclid_error_per_garment
+
+
+
 def plot_everything(batch: dict, type: str):
     levels_input = group_inputs_for_levels(batch)
     output_levels = model(levels_input, training=False)
@@ -343,7 +401,7 @@ def plot_everything(batch: dict, type: str):
     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])]).numpy()
 
-    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy")
+    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
     mask = batch['mesh_mask'].numpy().astype(bool)
 
     _, faces = maskFacesBatch(np.ones_like(mesh_orig), F, mask)
@@ -357,16 +415,22 @@ def plot_everything(batch: dict, type: str):
             mesh_unposed_shaped_levels[k].append(mesh_unposed_shaped[k])
     
     # full res version
-    mesh_unposed_shaped_full = np.zeros((mask.shape[0], 14475, 3), dtype=np.float32)
+    mesh_unposed_shaped_full = np.zeros((mask.shape[0], config['smpl_res'], 3), dtype=np.float32)
     for i in range(mask.shape[0]):
         mesh_unposed_shaped_full[i, mask[i], :] = mesh_unposed_shaped[i]
     mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_unposed_shaped_full).numpy()
     mesh_posed = [mesh_posed[j, mask[j], :] for j in range(mask.shape[0])]
 
     collision = [np.concatenate((body[i], mesh_posed[i]), axis=0) for i in range(len(mesh_posed))]
-    collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(mesh_posed))]
+    collision_faces = [np.concatenate((F, faces[i] + config['smpl_res']), axis=0) for i in range(len(mesh_posed))]
     collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(mesh_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(mesh_posed))]
 
+    # body_2 = [b.copy() for b in body]
+    # for i in range(len(gt_posed)):
+    #     body_2[i][:, 1] = body[i][:,2].copy()
+    #     body_2[i][:, 2] = body[i][:,1].copy()
+    # gt_posed_with_body = [np.concatenate((body_2[i], gt_posed[i]), axis=0) for i in range(len(gt_posed))]
+
     figs = { type + "/garment_" + str(i): plot_3d(gt[i], mesh_unposed_shaped_levels[i], faces[i], gt_posed[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
 
     return { **figs, **{ type + '/images': imgs, } }
@@ -494,6 +558,68 @@ def sample(folder):
         print()
 
 
+def test():
+    garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item()
+    str_split = config['name'].split('_')
+    test_idx = str_split.index('test')
+    checkpoint_name = '_'.join(str_split[:test_idx])
+    load_model(model, 'checkpoints/' + checkpoint_name)
+
+    dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=False, do_repeat=False)
+    static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=False)
+   
+    step = 0
+    losses = []
+
+    if config['is_skirt']:
+        per_garment_errors = { 'Tshirt': [], 'Dress': [], 'Skirt': [], 'Top': [] }
+    else:
+        per_garment_errors = { 'Tshirt': [], 'Jumpsuit': [], 'Trousers': [], 'Top': [] }
+
+    for batch in tqdm.tqdm(dataset_test):
+        batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
+
+        test_loss, euclid_error_per_garment = test_step(batch)
+        losses.append(test_loss)
+
+        for i, id in enumerate(batch['id']):
+            id_str = id.numpy().decode("utf-8")
+            per_garment_errors[garment_type_table[id_str]].append(euclid_error_per_garment[i].numpy()) 
+
+        to_log = { **{ 'current_step': step }, **test_loss }
+        if step % 1000 == 0:
+            figs = plot_everything(batch, 'test')
+            to_log = { **to_log, **figs }
+
+        wandb.log(to_log, step = step)
+        
+        # print(step)
+        step += 1
+   
+    # averages = dict()
+    path_to_plotly_html = "temp/plotly_figure.html"
+
+    histograms = dict()
+    for key in test_loss.keys():
+        accumulated = [losses[i][key].numpy() for i in range(len(losses))]
+        fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean()))
+        fig.write_html(path_to_plotly_html, auto_play = False)
+
+        histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html)
+        # averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean()
+
+    for key in per_garment_errors.keys():
+        fig = ex.histogram(per_garment_errors[key], title='Average: ' + str(np.asarray(per_garment_errors[key]).mean()))
+        fig.write_html(path_to_plotly_html, auto_play = False)
+        histograms['all/' + key + 'euclid_error_hist'] = wandb.Html(path_to_plotly_html)
+    
+
+    wandb.log({ **histograms }, step = step)
+    wandb.finish()
+    pass
+
+
+
 if __name__ == '__main__':
     init_wandb()
     allow_memory_growth()
@@ -501,17 +627,17 @@ if __name__ == '__main__':
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())
 
-    smpl = SMPL()
+    smpl = SMPL(orig_smpl=config['is_cape'])
     skirt_modifier = '_skirt' if config['is_skirt'] else ''
     dataset_resolutions = config['pyramid']['resolutions'].copy()
     if 512 not in dataset_resolutions:
         dataset_resolutions.append(512)
 
     factory = DatasetFactory()
-    dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'])
-    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions)
-    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'])
-    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions)
+    dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
+    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
+    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
+    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
 
     model: Pyramid = Pyramid(config)
 
@@ -523,15 +649,16 @@ if __name__ == '__main__':
     balancing_activated_step = 0
 
     # sample('checkpoints/cvae_13')
+    # test()
 
     ckpt = tf.train.Checkpoint(model=model)
     manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20)
 
-    load_model(model, config['pyramid']['pretrained_pyramid'])
+    # load_model(model, config['pyramid']['pretrained_pyramid'])
     model.static_enc.trainable = False
     model.body_enc.trainable = False
     model.curv_enc.trainable = False
-    model.cond_encoder.trainable = False
+    # model.cond_encoder.trainable = False
 
 
     # load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer'])