save

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)