diff --git a/data/dataset.py b/data/dataset.py
index 1b38abc..2c0ab84 100644
--- a/data/dataset.py
+++ b/data/dataset.py
@@ -131,7 +131,7 @@ def read_frames_file(frames_file, shuffle=False):
     with open(frames_file, 'r') as f:
         lines = f.readlines()
         lines = [line.strip() for line in lines]
-    frames = [Frame(line.split()[0], int(line.split()[1]), line.split()[2]) for line in lines]
+    frames = [Frame(frame_str=line) for line in lines]
 
     if shuffle:
         random.shuffle(frames)
diff --git a/data/dataset_cape.py b/data/dataset_cape.py
index 1680741..36b5fc3 100644
--- a/data/dataset_cape.py
+++ b/data/dataset_cape.py
@@ -44,7 +44,10 @@ def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=
         with open(os.path.join(folder, frames_file), 'r') as f:
             lines = f.readlines()
             lines = [line.strip() for line in lines]
-        frames = [Frame(line.split('_')[0], line.split('_')[1]) for line in lines]
+        if not is_cape:
+            frames = [Frame(line.split('_')[0], line.split('_')[1]) for line in lines]
+        else:
+            frames = [Frame(line.split('.')[0], line.split('.')[1], 'cape') for line in lines]
 
     if shuffle:
         random.shuffle(frames)
@@ -160,6 +163,7 @@ class DatasetCape(torch.utils.data.Dataset):
             item['garment'] = 'cape'
         if 'pose_shape' in self.signature:
             pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy')).astype(np.float32)
+            item['global_orient'] = pose_shape[:3].copy()
             pose_shape[:3] = 0
             item['pose_shape'] = pose_shape
         if 'mesh_body_posed' in self.signature:
diff --git a/models/cvae_model.py b/models/cvae_model.py
index 788f31e..c00e6cd 100644
--- a/models/cvae_model.py
+++ b/models/cvae_model.py
@@ -55,28 +55,24 @@ class CVAE(nn.Module):
         if self.avg_garment is not None:
             assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'
 
-    # @tf.function
-    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
+    
+    def sample(self, uv_template, uv_body_posed, uv_normals, uv_mask, noise=None):
         if noise is None:
-            eps = tf.random.normal(shape=(uv_template.shape[0], 8, 4, self.latent_dim))
+            eps = torch.randn(uv_template.shape[0], self.latent_shape[0], self.latent_shape[1], self.latent_shape[2]).to(uv_template.device)
         else:
             eps = noise
 
-        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
-        if add_noise:
-            z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
-        
-        z = tf.concat([eps, z_cond], axis=-1)
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_normals)
+        z = torch.cat([eps, z_cond], dim=1)
+        x = self.decoder(z)
 
-        x = self.decoder(z, training=False)
-        
         if self.avg_garment is not None:
-            x = x + self.avg_garment[tf.newaxis,:,:,:]
-        
+            x = x + self.avg_garment[None,:,:,:]
+    
         if self.use_offsets_over_body:
             x = x + uv_body_posed
 
-        x = x * uv_mask[:,:,:,tf.newaxis]
+        x = x * uv_mask[:,None,:,:]
 
         return x
 
@@ -94,11 +90,11 @@ class CVAE(nn.Module):
         return out
 
 
-    def encode_conds(self, uv_template, uv_body_posed, uv_curv):
+    def encode_conds(self, uv_template, uv_body_posed, uv_normals):
         z_body = self.body_enc(uv_body_posed)
         z_static = self.static_enc(uv_template)
         if self.use_normals:
-            z_norm = self.norm_enc(uv_curv)
+            z_norm = self.norm_enc(uv_normals)
             z_cond = torch.concat([z_static, z_body, z_norm], dim=1)
         else:
             z_cond = torch.concat([z_static, z_body], dim=1)
diff --git a/models/pyramid_level.py b/models/pyramid_level.py
index d5cfda9..482f3e1 100644
--- a/models/pyramid_level.py
+++ b/models/pyramid_level.py
@@ -39,22 +39,16 @@ class PyramidLevel(nn.Module):
 
     def sample(self, z_cond, uv_mask, noise=None):
         if noise is None:
-            eps = torch.normal(shape=(uv_mask.shape[0], self.latent_dim))
+            eps = torch.randn(z_cond.shape).to(uv_mask.device)
         else:
             eps = noise
 
-        z = torch.concat([eps, z_cond], dim=-1)
-
+        z = torch.cat([eps, z_cond], dim=1)
         x = self.decoder(z)
-        
-        if self.avg_garment is not None:
-            x = x + self.avg_garment[None,:,:,:]
-        
-        x = x * uv_mask[:,None:,:]
+        x = x * uv_mask[:,None,:,:]
 
         return x
 
-
     def reparameterize(self, mean, logvar, fix_noise=False):
         if self.training == True:
             if not fix_noise:
diff --git a/models/pyramid_model.py b/models/pyramid_model.py
index 1996ecd..6f81627 100644
--- a/models/pyramid_model.py
+++ b/models/pyramid_model.py
@@ -42,6 +42,28 @@ class Pyramid(nn.Module):
 
         self.use_normals = config['use_normals']
 
+
+    def sample(self, uv_template, uv_body_posed, uv_normals, uv_mask, noise=None):
+
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_normals)
+
+        assert len(self.levels) == 4
+        out = torch.zeros((uv_template.shape[0], 3, 64, 32)).to(uv_template.device)
+
+        for i in range(len(self.levels)):
+            level_out = self.levels[i].sample(z_cond, uv_mask[self.levels[i].size], noise=noise)
+            if i != len(self.levels) -  1:
+                out = out + level_out
+                out = self.levels[i].enhance(out, uv_mask[self.levels[i].size])
+                out = self.resizers[self.levels[i+1].size](out)
+                out = out * uv_mask[self.levels[i+1].size][:,None,:,:]
+            else:
+                out = out + level_out
+
+        out = out * uv_mask[512][:,None,:,:]
+
+        return out
+    
     def encode_conds(self, uv_template, uv_body_posed, uv_norm):
         z_static = self.static_enc(uv_template)
         z_body = self.body_enc(uv_body_posed)
diff --git a/scripts/calculate_error.py b/scripts/calculate_error.py
index 83b1b84..4e977b1 100644
--- a/scripts/calculate_error.py
+++ b/scripts/calculate_error.py
@@ -9,11 +9,11 @@ parser.add_argument('--name', type=str)
 parser.add_argument('--config', type=str)
 parser.add_argument('--debug', action='store_true')
 parser.add_argument('--iteration', type=str, default='')
-parser.add_argument('-g', '--gpu', type=str)
+parser.add_argument('-g', '--gpu', type=str, default='1')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64_128_cape_f_2 import config_1
+    from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_1
 else:
     config_module = importlib.import_module(args['config'])
     config_1 = config_module.config
@@ -65,19 +65,32 @@ def calculate():
         # fig.update_layout(scene_aspectmode='cube')
         # fig.show()
 
-        pred = np.load(f'{out_folder}/{frame}.npz')['arr_0']
+        pred = np.load(f'{out_folder}/{frame}.npz', allow_pickle=True)['arr_0']
+
+        try:
+            pred.item()['lvl3']
+            all_levels = True
+        except:
+            all_levels = False
+
         gt = readPC2Frame(f'{config_1["test_dir"]}/mesh_posed_no_orient/{frame.id}_{frame.garment}.pc16', frame.frame_nr, True)
-        # mask = np.load(f'{config["test_dir"]}/mesh_mask/{frame.id}_{frame.garment}.npy')
 
         fix_data = np.load(f'{config_1["test_dir"]}/mesh_fixed/{frame.id}_{frame.garment}.npz')
         mask_convert = fix_data['mas_convert']
         new_mask = fix_data['mask']
-
-        pred = pred[new_mask]
         gt = gt[mask_convert]
 
-        error = np.linalg.norm(pred - gt, axis=1).mean()
-        errors.append((str(frame), error))
+        if not all_levels:
+            pred = pred[new_mask]
+
+            error = np.linalg.norm(pred - gt, axis=1).mean()
+            errors.append((str(frame), error))
+        else:
+            for i in range(4):
+                _pred = pred.item()[f'lvl{i}'][new_mask]
+                error = np.linalg.norm(_pred - gt, axis=1).mean()
+                errors.append((f'{frame}_lvl{i}', error))
+
 
     # save errors to a csv file
     with open(f'{results_folder}/errors/{name}.csv', 'w') as f:
diff --git a/scripts/cvae.py b/scripts/cvae.py
index 3bfcb35..20d4140 100644
--- a/scripts/cvae.py
+++ b/scripts/cvae.py
@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.abl_all_4 import config
+    from configs.iccvw.abl_all_5_var import config
 else:
     config_module = importlib.import_module(args['config'])
     config = config_module.config
@@ -30,6 +30,7 @@ import numpy as np
 import plotly.graph_objects as go
 import plotly.express as ex
 from PIL import Image
+from pathlib import Path
 
 if config['is_skirt']:
     from losses.losses import edge_loss_skirt as edge_loss
@@ -277,171 +278,80 @@ def plot_everything(batch: dict, type: str):
     return { **figs, **figs_coll, **{ f'{type}/batch_info': batch_info }, **{ 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)
+
+
+def sample():
+    from data.utils import Frame
+    from utils.utils import save_stl
+
+    def create_temp_frames(folder: str, frame: Frame):
+        with open(Path(folder) / 'frames_temp.txt', 'w') as f:
+            f.writelines([f'{frame}\n' for _ in range(4)])
+
+    model.eval()
+    load_model(model, f'checkpoints/iccvw/{config["name"]}')
     count = 0
-    for batch in dataset_train:
-        batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
-        uv_normals = tf.one_hot(batch['uv_normals'], 7)
-        nr_of_samples=2
 
-        pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
-        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(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]
-
-        body = batch['mesh_body_posed'][0].numpy()
-        body_color = np.ones_like(body, dtype=np.int32)*100
-
-        F = np.load("custom_utils/smpl_faces.npy")
-        mask = batch['mesh_mask'].numpy().astype(bool)[0]
-        _, faces = _maskFaces(np.ones((config['smpl_res'], 3)), F, mask)
-
-        pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
-        mesh_pred_zero = uv2mesh(pred_zero)
-        mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
+    frames = ['01852_70_Tshirt']
+    frames = [Frame(frame_str=frame) for frame in frames]
 
 
-        meshes = []
-        meshes.append(go.Mesh3d(
-                    # 8 vertices of a cube
-                    x=mesh_orig_reposed[mask,0],
-                    y=mesh_orig_reposed[mask,2],
-                    z=mesh_orig_reposed[mask,1],
-                    # i, j and k give the vertices of triangles
-                    i = faces[:,0],
-                    j = faces[:,1],
-                    k = faces[:,2]
-                    ))
 
-        for i in range(nr_of_samples):
-            x = mesh_posed[i][mask, :].copy()
-            # x[:, 2] += 1 * (i + 1)
-            x[:, 2] += 1
-            meshes.append(go.Mesh3d(
-                        # 8 vertices of a cube
-                        x=x[:,0],
-                        y=x[:,2],
-                        z=x[:,1],
-                        # i, j and k give the vertices of triangles
-                        i = faces[:,0],
-                        j = faces[:,1],
-                        k = faces[:,2]
-                        ))
+    for frame in frames:
+        create_temp_frames(config['train_dir'], frame)
+        dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[512], is_cape=config['is_cape'])
+        dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
+        batch = next(iter(dataloader))
 
+        uv_static = batch['uv_static']
+        uv_body_posed = batch['uv_body_posed']
+        uv_mask = batch['uv_mask']
+        uv_normals = batch['uv_normals']
+
+        batch_size = uv_static.size(0)
+
+
+        for i in range(batch_size):
+            counter = 0
+            frame = Frame(batch['id'][i], batch['frame_nr'][i], batch['garment'][i])
+            fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
+            mask = fix_data['mask']
+            faces = fix_data['faces']
+
+            gt = batch['mesh_posed'][i].cpu().numpy()[mask]
+            save_stl(Path('temp') / 'sampling' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces)
+
+            uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1)
+            uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1)
+            uv_mask_tiled = uv_mask[i][None,:,:].repeat(4, 1, 1)
+            uv_normals_tiled = uv_normals[i][None,:,:,:].repeat(4, 1, 1, 1)
+            pose_shape = batch['pose_shape'][i][None,:].repeat(4, 1)
         
-        for i in range(nr_of_samples):
-            x = mesh_posed_zero[i][mask, :].copy()
-            # x[:, 2] += 1 * (i + 1)
-            x[:, 2] += 2
-            meshes.append(go.Mesh3d(
-                        # 8 vertices of a cube
-                        x=x[:,0],
-                        y=x[:,2],
-                        z=x[:,1],
-                        # i, j and k give the vertices of triangles
-                        i = faces[:,0],
-                        j = faces[:,1],
-                        k = faces[:,2]
-                        ))
+            for _ in range(4):
+                with torch.inference_mode():
+                    pred = model.sample(uv_static_tiled, uv_body_posed_tiled, uv_normals_tiled, uv_mask_tiled)
+                
+                mesh_pred = uv2mesh(pred)
 
-        # for i in range(5):
-        #     x = body.copy()
-        #     x[:, 2] += 1 * i
-        #     meshes.append(go.Mesh3d(
-        #                 # 8 vertices of a cube
-        #                 x=x[:,0],
-        #                 y=x[:,2],
-        #                 z=x[:,1],
-        #                 # i, j and k give the vertices of triangles
-        #                 i = F[:,0],
-        #                 j = F[:,1],
-        #                 k = F[:,2],
-        #                 vertexcolor=body_color
-        #                 ))
+                shape = pose_shape[:, 72:]
+                gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
+                shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+
+                mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+                mesh_pred_posed = smpl(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped)[0]
+
+                for k in range(4):
+                    posed = mesh_pred_posed.cpu().numpy()[k, mask, :]
+                    save_stl(Path('temp') / 'sampling' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}.stl', posed, faces)
+                    counter += 1
+
+            pass
 
 
-        fig = go.Figure(data=meshes)
-
-        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1))
-        # fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1))
-        fig.update_layout(
-            scene = dict(
-                xaxis = dict(nticks=4, range=[-1,1],),
-                            # yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],),
-                            yaxis = dict(nticks=4, range=[-1,3],),
-                            zaxis = dict(nticks=4, range=[-1,1],),)
-                            )
-
-        # wandb.log({'plot_' + str(count): fig})
-        fig.show()
-
-        pred = pred.numpy()
-        frames = [Image.fromarray(((pred[i] + 1) * 128).astype(np.int8), mode='RGB') for i in range(pred.shape[0])]
-        frames[0].save("temp/garment_1.gif", format="GIF", append_images=frames[1:],
-               save_all=True, duration=500, loop=0)
-
-
-        if count > 10:
+        if count > 100:
             break
 
-        print(count)
-        count += 1
 
-        print()
-
-
-
-def test():
-    str_split = config['name'].split('_')
-    test_idx = str_split.index('test')
-    checkpoint_name = '_'.join(str_split[:test_idx])
-    load_model(model, 'checkpoints/' + checkpoint_name)
-
-    dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False, do_repeat=False)
-    static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False)
-   
-    step = 0
-    losses = []
-
-    for batch in tqdm.tqdm(dataset_test):
-        batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
-
-        test_loss = test_step(batch)
-        losses.append(test_loss)
-
-        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()
-
-    
-
-    wandb.log({ **histograms }, step = step)
-    wandb.finish()
-    pass
 
 
 
@@ -500,6 +410,8 @@ if __name__ == '__main__':
 
     model = CVAE(config)
 
+    sample()
+    exit()
 
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
diff --git a/scripts/cvae_pyramid.py b/scripts/cvae_pyramid.py
index 4b36b5c..bd3c471 100644
--- a/scripts/cvae_pyramid.py
+++ b/scripts/cvae_pyramid.py
@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64_cape_f_skin import config
+    from configs.iccvw.pyr_64_128_256_512 import config
 else:
     config_module = importlib.import_module(args['config'])
     config = config_module.config
@@ -501,6 +501,86 @@ def speed_step(levels_input, conds):
     return model.sample([levels_input, conds], training=False)
 
 
+def sample():
+    from data.utils import Frame
+    from utils.utils import save_stl
+    from pathlib import Path
+    import trimesh as tm
+
+    def create_temp_frames(folder: str, frame: Frame):
+        with open(Path(folder) / 'frames_temp.txt', 'w') as f:
+            f.writelines([f'{frame}\n' for _ in range(4)])
+
+    model.eval()
+    load_model(model, f'checkpoints/iccvw/{config["name"]}')
+    count = 0
+
+    frames = ['01852_70_Tshirt']
+    frames = [Frame(frame_str=frame) for frame in frames]
+
+
+
+    for frame in frames:
+        create_temp_frames(config['train_dir'], frame)
+        dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape'])
+        dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
+        batch = next(iter(dataloader))
+
+        uv_static = batch['uv_static']
+        uv_body_posed = batch['uv_body_posed']
+        uv_mask = { 512: batch['uv_mask'], 256: batch['uv_mask_256'], 128: batch['uv_mask_128'], 64: batch['uv_mask_64'] }
+        uv_normals = batch['uv_normals']
+
+        batch_size = uv_static.size(0)
+
+
+        for i in range(batch_size):
+            counter = 0
+            frame = Frame(batch['id'][i], batch['frame_nr'][i], batch['garment'][i])
+            fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
+            mask = fix_data['mask']
+            faces = fix_data['faces']
+
+            gt = batch['mesh_posed'][i].cpu().numpy()[mask]
+            save_stl(Path('temp') / 'sampling_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces)
+
+            uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1)
+            uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1)
+            uv_mask_tiled = { 512: uv_mask[512][i][None,:,:].repeat(4, 1, 1), 256: uv_mask[256][i][None,:,:].repeat(4, 1, 1), 128: uv_mask[128][i][None,:,:].repeat(4, 1, 1), 64: uv_mask[64][i][None,:,:].repeat(4, 1, 1)}
+            uv_normals_tiled = uv_normals[i][None,:,:,:].repeat(4, 1, 1, 1)
+            pose_shape = batch['pose_shape'][i][None,:].repeat(4, 1)
+            body_posed = batch['mesh_body_posed'][i][None,:,:].repeat(4, 1, 1)
+            mask_tiled = torch.from_numpy(mask[None,:]).to(body_posed.device).repeat(4, 1)
+        
+            for _ in range(4):
+                with torch.inference_mode():
+                    pred = model.sample(uv_static_tiled, uv_body_posed_tiled, uv_normals_tiled, uv_mask_tiled)
+                
+                mesh_pred = uv2mesh(pred)
+
+                shape = pose_shape[:, 72:]
+                gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
+                shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+
+                mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+                mesh_pred_posed = smpl(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped)[0]
+
+                mesh_posed_fixed = fix_collision(mesh_pred_posed, body_posed, mask_tiled).cpu().numpy()
+
+                for k in range(4):
+                    _mesh_posed_fixed = mesh_posed_fixed[k, mask, :]
+                    m = tm.Trimesh(_mesh_posed_fixed, faces, process=False, validate=False)
+                    tm.smoothing.filter_humphrey(m, iterations=2)
+                    verts = m.vertices
+                    save_stl(Path('temp') / 'sampling_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}.stl', verts, faces)
+                    counter += 1
+
+            pass
+
+
+        if count > 100:
+            break
+
 
 def train():
     global g_balancing_coeff
@@ -560,6 +640,9 @@ if __name__ == '__main__':
 
     model: Pyramid = Pyramid(config)
 
+    sample()
+    exit()
+
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
     ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=50)
diff --git a/scripts/cvae_pyramid_e2e.py b/scripts/cvae_pyramid_e2e.py
index a7d7215..5b08e35 100644
--- a/scripts/cvae_pyramid_e2e.py
+++ b/scripts/cvae_pyramid_e2e.py
@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64_128_256_512_cape_m_e2e import config
+    from configs.iccvw.pyr_64_128_256_512_reg_e2e import config
 else:
     config_module = importlib.import_module(args['config'])
     config = config_module.config
@@ -246,34 +246,49 @@ def train_step(batch):
         if res != 512:
             pred = resizer(pred) * batch['uv_mask'][:,None,:,:]
 
-        l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2])) # levels_input[i][0] = uv_orig
+        l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2]))
         kl = kl_loss(mean, logvar)
 
-        mesh_pred = uv2mesh(pred)
+        if config['use_extra_losses']:
+            mesh_pred = uv2mesh(pred)
 
-        if config['use_skinning']:
-            shape = batch['pose_shape'][:, 72:]
-            gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
-            shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+            if config['use_skinning']:
+                shape = batch['pose_shape'][:, 72:]
+                gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
+                shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
 
-            mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_fix']['mesh_mask'][:,:,None]
-            mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+                mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_fix']['mesh_mask'][:,:,None]
+                mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+            else:
+                mesh_pred_posed = mesh_pred
+
+            e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_fix']['mesh_mask']))
+            c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_fix']['mesh_mask'])
+            c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
+            n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_fix']['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+            l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1))
+            
+            reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
+            if not config['is_cape']:
+                reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff'] 
+                if config['use_norm_loss'] and step > 2000:
+                    reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff']
+
+            output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss, 
+                            '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/normal_loss': n_loss,
+                            'train/l2_mesh_loss': l2_mesh_loss, 'train/collision_ratio': c_ratio } }
         else:
-            mesh_pred_posed = mesh_pred
+            reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
 
-        e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_fix']['mesh_mask']))
-        c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_fix']['mesh_mask'])
-        c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
-        n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_fix']['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
-        l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1))
-        
-        reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
-        if not config['is_cape']:
-            reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff'] 
-            if config['use_norm_loss'] and step > 2000:
-                reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * n_loss * config['reg_loss_coeff']
-        
         kl_total += kl
+
+        model.eval()
+        with torch.inference_mode():
+            mean = torch.mean(mean)
+            logvar = torch.mean(logvar)
+
+    
     loss: torch.Tensor = reconstr_loss_total + g_balancing_coeff * kl_total
 
     optimizer.zero_grad()
@@ -282,17 +297,45 @@ def train_step(batch):
 
     model.eval()
     with torch.inference_mode():
-        mean = torch.mean(mean)
-        logvar = torch.mean(logvar)
+        for i, res in enumerate(config['pyramid']['resolutions']):
+            # i = -1
+            res = config['pyramid']['resolutions'][i]
+            res_modifier = '_' + str(res) if res != 512 else ''
+            pred, mean, logvar = out_levels[i]
+            if res != 512:
+                pred = resizer(pred) * batch['uv_mask'][:,None,:,:]
 
-        ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1))
-        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask'])
+            mean = torch.mean(mean)
+            logvar = torch.mean(logvar)
 
-    output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric, 
-                                '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/normal_loss': n_loss,
-                                'train/euclid_error': metric, 'train/l2_mesh_loss': l2_mesh_loss, 'train/ms_ssim' + res_modifier: ms_ssim,
-                                'train/collision_ratio': c_ratio } } # 
+            if not config['use_extra_losses']:
+                mesh_pred = uv2mesh(pred)
+
+                if config['use_skinning']:
+                    shape = batch['pose_shape'][:, 72:]
+                    gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
+                    shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+
+                    mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_fix']['mesh_mask'][:,:,None]
+                    mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+                else:
+                    mesh_pred_posed = mesh_pred
+
+                e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_fix']['mesh_mask']))
+                c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_fix']['mesh_mask'])
+                c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
+                l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1))
+                n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_fix']['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+                
+                output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,
+                                            '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/normal_loss': n_loss,
+                                            'train/l2_mesh_loss': l2_mesh_loss, 'train/collision_ratio': c_ratio } }
+
+            ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1))
+            metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask'])
+
+            output = { **output, **{ 'train/euclid_error' + res_modifier: metric, 'train/ms_ssim' + res_modifier: ms_ssim } }
 
 
     return { **{ 'train/loss': loss, 'train/reconstr_loss_total': reconstr_loss_total, 'train/kl_total': kl_total }, **output }
@@ -356,18 +399,21 @@ def valid_step(batch):
                 reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
 
             kl_total += kl
+
+            mean = torch.mean(mean)
+            logvar = torch.mean(logvar)
+
+            ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1))
+            metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask'])
+
+            output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss, 
+                                        '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/normal_loss': n_loss,
+                                        'valid/euclid_error' + res_modifier: metric, 'valid/l2_mesh_loss': l2_mesh_loss, 'valid/ms_ssim' + res_modifier: ms_ssim, } } # 
+
+
         loss = reconstr_loss_total + g_balancing_coeff * kl_total
 
-        mean = torch.mean(mean)
-        logvar = torch.mean(logvar)
-
-        ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1))
-        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask'])
-
-        output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss,'valid/euclid_error' + res_modifier: metric, 
-                                    '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/normal_loss': n_loss,
-                                    'valid/euclid_error': metric, 'valid/l2_mesh_loss': l2_mesh_loss, 'valid/ms_ssim' + res_modifier: ms_ssim, } } # 
 
 
 
diff --git a/scripts/generate_results.py b/scripts/generate_results.py
index 7e7e6f4..30d7b90 100644
--- a/scripts/generate_results.py
+++ b/scripts/generate_results.py
@@ -10,6 +10,7 @@ parser.add_argument('--config', type=str)
 parser.add_argument('--iteration', type=str, default='')
 parser.add_argument('--debug', action='store_true')
 parser.add_argument('-g', '--gpu', type=str)
+parser.add_argument('-l', '--level', type=int, default=-1)
 args = vars(parser.parse_args())
 
 if args['debug']:
@@ -62,7 +63,7 @@ def step_pyramid(batch):
         levels_input, conds = group_inputs_for_levels(batch, config_1)
 
         out_levels = model([levels_input, conds])
-        pred, _, _ = out_levels[-1]
+        pred, _, _ = out_levels[args['level']]
 
         mesh_pred = uv2mesh(pred)
 
@@ -167,7 +168,10 @@ if __name__ == '__main__':
     model.eval()
     # it = load_model(model, 'checkpoints/iccvw/' + config['name'])
 
-    name = config_1['name'] + '_' + str(it)
+    if args['level'] == -1:
+        name = config_1['name'] + '_' + str(it)
+    else:
+        name = config_1['name'] + '_lvl' + str(args['level']) + '_' + str(it)
 
     folder = '/data-local/127_2/results'
     out_folder = f'{folder}/{name}'
diff --git a/scripts/get_error.py b/scripts/get_error.py
index bf36b24..1a0bbfd 100644
--- a/scripts/get_error.py
+++ b/scripts/get_error.py
@@ -8,12 +8,13 @@ parser = argparse.ArgumentParser(description='')
 parser.add_argument('--name', type=str)
 parser.add_argument('--config', type=str)
 parser.add_argument('--debug', action='store_true')
-parser.add_argument('--iteration', type=str, default='')
+parser.add_argument('--iteration', type=str, default='190000')
 parser.add_argument('-g', '--gpu', type=str)
+parser.add_argument('-a', '--all', action='store_true')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64_128_cape_f_2 import config_1
+    from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_1
 else:
     config_module = importlib.import_module(args['config'])
     config_1 = config_module.config
@@ -32,8 +33,14 @@ from data.utils import Frame
 
 
 def calculate():
-    errors = dict()
-    errors_mean = dict()
+    all_levels = args['all']
+    if all_levels:
+        errors = {0: dict(), 1: dict(), 2: dict(), 3: dict()}
+        errors_mean = {0: dict(), 1: dict(), 2: dict(), 3: dict()}
+    else:
+        errors = dict()
+        errors_mean = dict()
+
 
     with open(f'{results_folder}/errors/{name}.csv', 'r') as f:
         lines = f.readlines()
@@ -41,18 +48,31 @@ def calculate():
     for line in lines:
         line = line.split(',')
         category = line[2]
-        if category not in errors:
-            errors[category] = []
-        errors[category].append(float(line[3]))
+        if all_levels:
+            level = int(line[3][-1])
+            if category not in errors:
+                errors[level][category] = []
+            errors[level][category].append(float(line[4]))
+        else:
+            if category not in errors:
+                errors[category] = []
+            errors[category].append(float(line[3]))
 
-    for category in errors:
-        errors_mean[category] = np.mean(errors[category])
+    if all_levels:
+        for i in range(4):
+            for category in errors[i]:
+                errors_mean[i][category] = round(np.mean(errors[i][category]), 5)
+            print(f'Level {i}: {errors_mean[i]}')
 
-    print(errors_mean) 
+            mean = np.mean(list(errors_mean[i].values()))
+            print(f'Final error level {i}: {mean}')
+    else:
+        for category in errors:
+            errors_mean[category] = round(np.mean(errors[category]), 5)
+        print(errors_mean) 
 
-    mean = np.mean(list(errors_mean.values()))
-
-    print(f'Final error: {mean}')
+        mean = np.mean(list(errors_mean.values()))
+        print(f'Final error: {mean}')
 
 
 
diff --git a/scripts/get_final_cape.py b/scripts/get_final_cape.py
index edcdb2b..f81a1ff 100644
--- a/scripts/get_final_cape.py
+++ b/scripts/get_final_cape.py
@@ -13,10 +13,10 @@ parser.add_argument('-g', '--gpu', type=str, default='1')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_1
+    from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_f
 else:
     config_module = importlib.import_module(args['config'])
-    config_1 = config_module.config
+    config_f = config_module.config
 
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = args['gpu']
@@ -54,13 +54,13 @@ from models.cvae_model import CVAE
 from data.utils import Frame
 
 
-uv2mesh = UV2Mesh(res=256, is_skirt=config_1['is_skirt'], is_smpl_orig=config_1['is_cape'])
+uv2mesh = UV2Mesh(res=256, is_skirt=config_f['is_skirt'], is_smpl_orig=config_f['is_cape'])
 
 
 def step_pyramid(batch):
     model.eval()
     with torch.inference_mode():
-        levels_input, conds = group_inputs_for_levels(batch, config_1)
+        levels_input, conds = group_inputs_for_levels(batch, config_f)
 
         out = []
 
@@ -70,16 +70,13 @@ def step_pyramid(batch):
 
             mesh_pred = uv2mesh(pred)
 
-            if config_1['use_skinning']:
+            if config_f['use_skinning']:
                 shape = batch['pose_shape'][:, 72:]
                 gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
                 shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
 
                 mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
-                if not config_1['is_cape']:
-                    mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
-                else:
-                    mesh_pred_posed = mesh_pred_shaped
+                mesh_pred_posed = mesh_pred_shaped
             else:
                 mesh_pred_posed = mesh_pred
 
@@ -94,8 +91,8 @@ cape_mask = np.zeros((6890), bool)
 cape_mask[clothing_verts_idx] = True
 
 def get_error(pred, frame: Frame):
-    gt = np.load(f'{config_1["test_dir"]}/mesh_unposed_no_shape/{frame.id}.{frame.frame_nr}.npz')['mesh_unposed_no_shape']
-    pose_shape = np.load(f'{config_1["test_dir"]}/pose_shape/{frame.id}.{frame.frame_nr}.npy').astype(np.float32)
+    gt = np.load(f'{config_f["test_dir"]}/mesh_unposed_no_shape/{frame.id}.{frame.frame_nr}.npz')['mesh_unposed_no_shape']
+    pose_shape = np.load(f'{config_f["test_dir"]}/pose_shape/{frame.id}.{frame.frame_nr}.npy').astype(np.float32)
     pose, shape = pose_shape[:72], pose_shape[72:]
     gender = int(shape[10] >= 0.5)
     shape_offsets = torch.tensordot(torch.from_numpy(shape[:10][None, :]), smpl.shapedirs[gender].to('cpu'), dims=[[1], [2]])[0,smpl.to_keep,:].numpy()
@@ -155,30 +152,30 @@ def generate():
 
 if __name__ == '__main__':
 
-    smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device'])
-    skirt_modifier = '_skirt' if config_1['is_skirt'] else ''
+    smpl = SMPL(orig_smpl=config_f['is_cape'], path='smpl', device=config_f['device'])
+    skirt_modifier = '_skirt' if config_f['is_skirt'] else ''
     dataset_resolutions = []
-    if 'pyramid' in config_1:
-        dataset_resolutions = config_1['pyramid']['resolutions'].copy()
+    if 'pyramid' in config_f:
+        dataset_resolutions = config_f['pyramid']['resolutions'].copy()
     if 512 not in dataset_resolutions:
         dataset_resolutions.append(512)
 
     # remove all values from config['signature] that start with mesh_
-    config_1['signature'] = set([x for x in config_1['signature'] if not x.startswith('mesh_')])
+    config_f['signature'] = set([x for x in config_f['signature'] if not x.startswith('mesh_')])
 
-    dataset_test = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file=config_1['test_frames'], use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
+    dataset_test = Dataset(config_f['test_dir'], config_f['signature'], device=config_f['device'], frames_file=config_f['test_frames'], use_cache=config_f['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_f['is_cape'])
     # dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
     test_dataloader = DataLoader(dataset_test, batch_size=4, drop_last=True, shuffle=True, )
 
-    is_pyramid = config_1['name'].startswith('pyr')
+    is_pyramid = config_f['name'].startswith('pyr')
 
-    model: Pyramid = Pyramid(config_1)
+    model: Pyramid = Pyramid(config_f)
 
-    it = load_model(model, 'checkpoints/iccvw/' + config_1['name'], iteration=args['iteration'])
+    it = load_model(model, 'checkpoints/iccvw/' + config_f['name'], iteration=args['iteration'])
     model.eval()
     # it = load_model(model, 'checkpoints/iccvw/' + config['name'])
 
-    name = config_1['name'] + '_' + str(it)
+    name = config_f['name'] + '_' + str(it)
 
     folder = '/data-local/127_2/results'
     out_folder = f'{folder}/{name}'
diff --git a/scripts/plot_ablation.py b/scripts/plot_ablation.py
index 26df6dd..bb9fe27 100644
--- a/scripts/plot_ablation.py
+++ b/scripts/plot_ablation.py
@@ -85,7 +85,7 @@ if __name__ == '__main__':
 
     experiments = {
                    'final': { 'nonskirt': ('pyr_64_128_256_512', 100000), 'skirt': ('pyr_64_128_256_512_skirt', 200000) },
-                   'e2e': { 'nonskirt': ('pyr_64_128_256_512_e2e', 165000), 'skirt': ('pyr_64_128_256_512_skirt_e2e', 175000) },
+                   'e2e': { 'nonskirt': ('pyr_64_128_256_512_e2e', 195000), 'skirt': ('pyr_64_128_256_512_skirt_e2e', 200000) },
                    'skin': { 'nonskirt': ('pyr_64_128_256_512_skin', 225000), 'skirt': ('pyr_64_128_256_512_skin_skirt', 200000) },
                    'norm': { 'nonskirt': ('pyr_64_128_256_512_norm', 200000), 'skirt': ('pyr_64_128_256_512_norm_skirt', 200000) },
                    'reg': { 'nonskirt': ('pyr_64_128_256_512_reg', 200000), 'skirt': ('pyr_64_128_256_512_reg_skirt', 200000) },
@@ -110,7 +110,8 @@ if __name__ == '__main__':
         models[experiment]['skirt'].eval()
 
 
-    frames = ['00054_1_Dress', '00005_1_Jumpsuit']
+    # frames = ['00054_1_Dress', '00005_1_Jumpsuit']
+    frames = ['03849_27_Dress', '00054_1_Dress', '01852_70_Tshirt']
     frames = [Frame(frame_str=frame) for frame in frames]
     body_faces = smpl.f
 
@@ -145,13 +146,13 @@ if __name__ == '__main__':
             save_stl(out_folder / f'{frame}_{experiment}_body.stl', body, body_faces)
             save_stl(out_folder / f'{frame}_{experiment}_mesh.stl', verts, faces)
 
-            fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
-                                  go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
-                                  ])
-            fig.update_layout(scene=dict(aspectmode='cube'))
-            fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
-            fig.update_layout(title_text=f'{frame} {experiment}')
-            fig.show()
+            # fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
+            #                       go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
+            #                       ])
+            # fig.update_layout(scene=dict(aspectmode='cube'))
+            # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
+            # fig.update_layout(title_text=f'{frame} {experiment}')
+            # fig.show()
             pass
 
         pass
diff --git a/scripts/plot_ablation_pyramid.py b/scripts/plot_ablation_pyramid.py
index e8d4d75..58168af 100644
--- a/scripts/plot_ablation_pyramid.py
+++ b/scripts/plot_ablation_pyramid.py
@@ -8,7 +8,7 @@ os.environ["CUDA_VISIBLE_DEVICES"] = '1'
 import time
 import argparse
 import wandb
-import tqdm
+from tqdm import tqdm
 import random
 import numpy as np
 import plotly.graph_objects as go
@@ -39,30 +39,32 @@ import trimesh as tm
 from zipfile import ZipFile
 
 
-from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_1
+from configs.iccvw.pyr_64_128_256_512_e2e import config as config_1
+from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_2
+
+uv2mesh = dict()
+uv2mesh['nonskirt'] = UV2Mesh(res=256, is_skirt=False, is_smpl_orig=config_1['is_cape'])
+uv2mesh['skirt'] = UV2Mesh(res=256, is_skirt=True, is_smpl_orig=config_1['is_cape'])
 
 
-uv2mesh = UV2Mesh(res=256, is_skirt=config_1['is_skirt'], is_smpl_orig=config_1['is_cape'])
-
-
-def step_pyramid(model, batch):
+def step_pyramid(model, batch, config):
     with torch.inference_mode():
-        levels_input, conds = group_inputs_for_levels(batch, config_1)
+        levels_input, conds = group_inputs_for_levels(batch, config)
 
         out = []
         out_levels = model([levels_input, conds])
         for i in range(len(out_levels)):
             pred, _, _ = out_levels[i]
 
-            mesh_pred = uv2mesh(pred)
+            mesh_pred = uv2mesh[picker](pred)
 
-            if config_1['use_skinning']:
+            if config['use_skinning']:
                 shape = batch['pose_shape'][:, 72:]
                 gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
                 shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
 
                 mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
-                if not config_1['is_cape']:
+                if not config['is_cape']:
                     mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
                 else:
                     mesh_pred_posed = mesh_pred_shaped
@@ -82,31 +84,39 @@ def create_temp_frames(folder: str, frame: Frame):
 if __name__ == '__main__':
     results_folder = Path('/data-local/127_2/results')
     out_folder = Path('temp/pyramid_ablation')
-    iteration = '100000'
+    iteration = {'nonskirt': 190000, 'skirt': 200000 }
 
     smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device'])
-    skirt_modifier = '_skirt' if config_1['is_skirt'] else ''
     dataset_resolutions = []
     if 'pyramid' in config_1:
         dataset_resolutions = config_1['pyramid']['resolutions'].copy()
     if 512 not in dataset_resolutions:
         dataset_resolutions.append(512)
 
-    model: Pyramid = Pyramid(config_1)
-    it = load_model(model, 'checkpoints/iccvw/' + config_1['name'], iteration=iteration)
-    model.eval()
+    model = { 'nonskirt': Pyramid(config_1), 'skirt': Pyramid(config_2) }
+    load_model(model['nonskirt'], 'checkpoints/iccvw/' + config_1['name'], iteration=iteration['nonskirt'])
+    load_model(model['skirt'], 'checkpoints/iccvw/' + config_2['name'], iteration=iteration['skirt'])
+    model['nonskirt'].eval()
+    model['skirt'].eval()
 
-    frames = ['00054_1_Dress']
+    frames = ['00054_1_Dress', '01070_13_Jumpsuit', '03849_27_Dress', '06057_152_Dress',
+              '03179_274_Dress', '04247_102_Tshirt', '01852_70_Tshirt']
+
+    frames_corrupt = ['00300_69_Jumpsuit', '07608_77_Jumpsuit', '06955_117_Dress']
     frames = [Frame(frame_str=frame) for frame in frames]
     body_faces = smpl.f
 
-    for frame in frames:
+    config = { 'nonskirt': config_1, 'skirt': config_2 }
+
+    for frame in tqdm(frames):
         create_temp_frames(config_1['train_dir'], frame)
-        dataset = Dataset(config_1['train_dir'], config_1['signature'], device=config_1['device'], frames_file='frames_temp.txt', use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
+        dataset = Dataset(config_1['train_dir'], config_1['signature'], device=config_1['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
         dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
 
+        picker = 'skirt' if frame.garment in ['Skirt', 'Dress'] else 'nonskirt'
+
         batch = next(iter(dataloader))
-        pred = step_pyramid(model, batch)
+        pred = step_pyramid(model[picker], batch, config[picker])
         
         fix_data = np.load(Path(dataset.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
         mask = fix_data['mask']
@@ -124,19 +134,18 @@ if __name__ == '__main__':
             mesh_posed_fixed = mesh_posed_fixed[0, mask, :]
             m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False)
             tm.smoothing.filter_humphrey(m, iterations=2)
-
             verts = m.vertices
 
             save_stl(out_folder / f'{frame}_lvl{i}_body.stl', body, body_faces)
             save_stl(out_folder / f'{frame}_lvl{i}_mesh.stl', verts, faces)
 
-
             # fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
             #                       go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
             #                       ])
             # fig.update_layout(scene=dict(aspectmode='cube'))
             # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
             # fig.show()
+            pass
 
         pass
 
diff --git a/scripts/plot_cape.py b/scripts/plot_cape.py
index fad3046..5e912d3 100644
--- a/scripts/plot_cape.py
+++ b/scripts/plot_cape.py
@@ -134,7 +134,7 @@ if __name__ == '__main__':
     dataset_resolutions = [64, 128, 256, 512]
 
     # cape_exp = { 'cape_f': ('pyr_64_128_256_512_cape', 100000), 'cape_m': ('pyr_64_128_256_512_skirt', 200000) }
-    ours_exp = { 'cape_f': ('pyr_iccvw_64_128_256_512_cape_f_e2e_fine_3', 170000), 'cape_m': ('pyr_iccvw_64_128_256_512_cape_m_e2e', 230000) }
+    ours_exp = { 'cape_f': ('pyr_iccvw_64_128_256_512_cape_f_e2e_fine_3_lvl2', 170000), 'cape_m': ('pyr_iccvw_64_128_256_512_cape_m_e2e_lvl2', 230000) }
 
     picker = 'cape_m'
     config = dict()
@@ -142,8 +142,8 @@ if __name__ == '__main__':
     config['cape_m'] = config_m
 
     frames_override = ['00122_shortshort_ballerina_spin_000153', '00127_shortlong_ballerina_spin_000141',
-                       '00127_shortlong_soccer_000306', '00134_longlong_climb_trial2_000040',
-                       '00134_shortlong_ski_trial2_000230', '00134_shortlong_soccer_trial2_000059']
+                       '00127_shortlong_soccer_000306', 
+                       '00134_shortlong_ski_trial2_000230',] #'00134_longlong_climb_trial2_000040', '00134_shortlong_soccer_trial2_000059'
 
     _frames = frames_override if frames_override is not None else frames[picker]
     for _frame in tqdm(_frames):
diff --git a/utils/utils.py b/utils/utils.py
index 5b5fe02..394c549 100644
--- a/utils/utils.py
+++ b/utils/utils.py
@@ -2,6 +2,7 @@ import os
 import sys
 import numpy as np
 from torch.utils.data import DataLoader
+from stl import mesh as mesh_stl
 
 
 
@@ -48,4 +49,14 @@ def text_to_html(lines):
     html = ''
     for line in lines:
         html += "<pre>" + line + "</pre>"
-    return html
\ No newline at end of file
+    return html
+
+
+
+def save_stl(file, mesh, faces):
+    obj = mesh_stl.Mesh(np.zeros(faces.shape[0], dtype=mesh_stl.Mesh.dtype))
+    for i, f in enumerate(faces):
+        for j in range(3):
+            obj.vectors[i][j] = mesh[f[j],:]    
+
+    obj.save(file)
\ No newline at end of file