update pyramid for ablation

models/pyramid_level.py

@@ -100,7 +100,7 @@ from models.upscaler import Upscaler
 
 class PyramidLevel(nn.Module):
     
-    def __init__(self, config):
+    def __init__(self, config, use_normals=True):
         super(PyramidLevel, self).__init__()
         self.size = config['resolution']
 
@@ -114,6 +114,7 @@ class PyramidLevel(nn.Module):
             load_model(self.upscale_enhancer, config['upscale_enhancer'])
             for param in self.upscale_enhancer.parameters(): param.requires_grad = False
 
+        self.use_normals = use_normals
 
     def enhance(self, low, mask):
         x = self.upscale_enhancer([low, mask])
@@ -152,9 +153,13 @@ class PyramidLevel(nn.Module):
 
 
     def forward(self, inputs, z_cond):
-        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_mask_high = inputs
+        uv_unposed, uv_template, uv_body_posed, uv_normals, uv_mask, uv_mask_high = inputs
 
-        x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], dim=1)
+        if self.use_normals:
+            x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_normals], dim=1)
+        else:   
+            x = torch.concat([uv_unposed, uv_template, uv_body_posed], dim=1)
+            
         mean, logvar = self.encoder(x)
         z = self.reparameterize(mean, logvar)
         z = torch.reshape(z, (mean.shape[0], -1, 8, 4))