'testing'

pyramid_dit/pyramid_dit_for_video_gen_pipeline.py

@@ -558,18 +558,32 @@ class PyramidDiTForVideoGeneration:
                     generator,
                     is_first_frame=True,
                 )
-                generated_latents_list.append(intermed_latents[-1])
+                # Ensure 
+                if not intermed_latents or len(intermed_latents) == 0:
+                    raise ValueError("First frame generation failed")
+                    generated_latents_list.append(intermed_latents[-1].clone()) 
             else:
+
                 
-                # check if we have previous frames to condition on 
-                if not generated_latents_list or len(generated_latents_list) == 0:
-                    raise ValueError("No latent tensors generated - check previous steps ")
 
                 
+                # Subsequent frames
+                if len(generated_latents_list) == 0:
+                    raise ValueError("No previous frames available for conditioning")
+                    
+                # Debug print
+                print(f"Conditioning on {len(generated_latents_list)} existing frames")
                 
-                # prepare the condition latents
-                past_condition_latents = []
-                clean_latents_list = self.get_pyramid_latent(torch.cat(generated_latents_list, dim=2), len(stages) - 1)
+                # Get pyramid latents from existing frames
+                try:
+                    clean_latents_list = self.get_pyramid_latent(
+                        torch.cat(generated_latents_list, dim=2),
+                        len(stages) - 1
+                    )
+                except Exception as e:
+                    print(f"Error concatenating latents: {e}")
+                    print(f"Shapes: {[x.shape for x in generated_latents_list]}")
+                    raise
                 
                 for i_s in range(len(stages)):
                     last_cond_latent = clean_latents_list[i_s][:,:,-(self.frame_per_unit):]