sequential pyramid

models/pyramid_level.py

@@ -155,15 +155,18 @@ class PyramidLevelNew(tf.keras.Model):
         return out
 
 
-    def call(self, inputs, z_cond, training=None):
+    def call(self, inputs, z_cond, prev_latent=None, training=None):
         uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_mask_high = inputs
 
         x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
         mean, logvar = self.encoder(x, training=training)
         z = self.reparameterize(mean, logvar, training)
-        z = self.reshape(z)
-
-        z = tf.concat([z, z_cond], axis=-1)
+        z_latent = self.reshape(z)
+    
+        if prev_latent is not None:
+            z_latent = tf.concat([z_latent, prev_latent], axis=-1)
+            
+        z = tf.concat([z_latent, z_cond], axis=-1)
         x = self.decoder(z, training=training)
 
         if self.avg_garment is not None:
@@ -177,4 +180,4 @@ class PyramidLevelNew(tf.keras.Model):
         
         # x = x * uv_mask_high[:,:,:,tf.newaxis]
 
-        return x, mean, logvar
+        return x, mean, logvar, z_latent