'testing'

pyramid_dit/pyramid_dit_for_video_gen_pipeline.py

@@ -533,7 +533,7 @@ class PyramidDiTForVideoGeneration:
         stages = self.stages
 
         generated_latents_list = []    # The generated results
-        last_generated_latents = None
+        # last_generated_latents = None
 
         for unit_index in tqdm(range(num_units)):
             if use_linear_guidance:
@@ -542,92 +542,114 @@ class PyramidDiTForVideoGeneration:
 
             if unit_index == 0:
                 # Generate first frame
-                past_condition_latents = [[] for _ in range(len(stages))]
-                intermed_latents = self.generate_one_unit(
-                    latents[:,:,:1],
-                    past_condition_latents,
-                    prompt_embeds,
-                    prompt_attention_mask,
-                    pooled_prompt_embeds,
-                    num_inference_steps,
-                    height,
-                    width,
-                    1,
-                    device,
-                    dtype,
-                    generator,
-                    is_first_frame=True,
-                )
-                # Ensure 
-                if not intermed_latents or len(intermed_latents) == 0:
-                    raise ValueError("First frame generation failed")
+                try:
+                    
+                    past_condition_latents = [[] for _ in range(len(stages))]
+                    intermed_latents = self.generate_one_unit(
+                        latents[:,:,:1],
+                        past_condition_latents,
+                        prompt_embeds,
+                        prompt_attention_mask,
+                        pooled_prompt_embeds,
+                        num_inference_steps,
+                        height,
+                        width,
+                        1,
+                        device,
+                        dtype,
+                        generator,
+                        is_first_frame=True,
+                    )
+                    # Ensure 
+                    if not intermed_latents:
+                        raise ValueError("First frame generation failed")
+                    
                     generated_latents_list.append(intermed_latents[-1].clone()) 
+                    print(f"successfully generated first frame. shape: {generated_latents_list[-1].shape}")
+
+                except Exception as e:
+                    print(f"First frame generation failed: {str(e)}")
+                    raise ValueError("Could not generate initial frame") from e 
             else:
 
-                
+                if not generated_latents_list:
+                    raise ValueError("No previous frames available for conditioning (this should never happen)")
 
+                try:
+
+                    # prepare conditioning from existing frames 
+                    concatenated_latents = torch.cat(generated_latents_list, dim=2)
+                    print(f"Conditioning on {len(generated_latents_list)} frame. Concatenated shape: {concatenated_latents.shape}")
+
+                    clean_latents_list = self.get_pyramid_latent(concatenated_latents, len(stages) - 1)
+
+                    # prepare past conditions
+                    past_condition_latents = []
+
+
+                    for i_s in range(len(stages)):
+                        last_cond_latent = clean_latents_list[i_s][:,:,-(self.frame_per_unit):]
+
+                        stage_input = [torch.cat([last_cond_latent] * 2) if self.do_classifier_free_guidance else last_cond_latent]
                 
-                # Subsequent frames
-                if len(generated_latents_list) == 0:
-                    raise ValueError("No previous frames available for conditioning")
+                        # pad the past clean latents
+                        cur_unit_num = unit_index
+                        cur_stage = i_s
+                        cur_unit_ptx = 1
+
+                        while cur_unit_ptx < cur_unit_num:
+                            cur_stage = max(cur_stage - 1, 0)
+                            if cur_stage == 0:
+                                break
+                            cur_unit_ptx += 1
+                            cond_latents = clean_latents_list[cur_stage][:, :, -(cur_unit_ptx * self.frame_per_unit) : -((cur_unit_ptx - 1) * self.frame_per_unit)]
+                            stage_input.append(torch.cat([cond_latents] * 2) if self.do_classifier_free_guidance else cond_latents)
+
+                        if cur_stage == 0 and cur_unit_ptx < cur_unit_num:
+                            cond_latents = clean_latents_list[0][:, :, :-(cur_unit_ptx * self.frame_per_unit)]
+                            stage_input.append(torch.cat([cond_latents] * 2) if self.do_classifier_free_guidance else cond_latents)
                     
-                # Debug print
-                print(f"Conditioning on {len(generated_latents_list)} existing frames")
-                
-                # Get pyramid latents from existing frames
-                try:
-                    clean_latents_list = self.get_pyramid_latent(
-                        torch.cat(generated_latents_list, dim=2),
-                        len(stages) - 1
+                        stage_input = list(reversed(stage_input))
+                        past_condition_latents.append(stage_input)
+
+                    # Generate current frame unit 
+                    frame_slice = slice(
+                        1 + (unit_index - 1) * self.frame_per_unit,
+                        1 + unit_index * self.frame_per_unit
                     )
-                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):]
 
-                    stage_input = [torch.cat([last_cond_latent] * 2) if self.do_classifier_free_guidance else last_cond_latent]
-            
-                    # pad the past clean latents
-                    cur_unit_num = unit_index
-                    cur_stage = i_s
-                    cur_unit_ptx = 1
-
-                    while cur_unit_ptx < cur_unit_num:
-                        cur_stage = max(cur_stage - 1, 0)
-                        if cur_stage == 0:
-                            break
-                        cur_unit_ptx += 1
-                        cond_latents = clean_latents_list[cur_stage][:, :, -(cur_unit_ptx * self.frame_per_unit) : -((cur_unit_ptx - 1) * self.frame_per_unit)]
-                        stage_input.append(torch.cat([cond_latents] * 2) if self.do_classifier_free_guidance else cond_latents)
-
-                    if cur_stage == 0 and cur_unit_ptx < cur_unit_num:
-                        cond_latents = clean_latents_list[0][:, :, :-(cur_unit_ptx * self.frame_per_unit)]
-                        stage_input.append(torch.cat([cond_latents] * 2) if self.do_classifier_free_guidance else cond_latents)
-                
-                    stage_input = list(reversed(stage_input))
-                    past_condition_latents.append(stage_input)
-
-                intermed_latents = self.generate_one_unit(
-                    latents[:,:, 1 + (unit_index - 1) * self.frame_per_unit:1 + unit_index * self.frame_per_unit],
-                    past_condition_latents,
-                    prompt_embeds,
-                    prompt_attention_mask,
-                    pooled_prompt_embeds,
-                    video_num_inference_steps,
-                    height,
-                    width,
-                    self.frame_per_unit,
-                    device,
-                    dtype,
-                    generator,
-                    is_first_frame=False,
-                )
+                    intermed_latents = self.generate_one_unit(
+                        latents[:,:, frame_slice],
+                        past_condition_latents,
+                        prompt_embeds,
+                        prompt_attention_mask,
+                        pooled_prompt_embeds,
+                        video_num_inference_steps,
+                        height,
+                        width,
+                        self.frame_per_unit,
+                        device,
+                        dtype,
+                        generator,
+                        is_first_frame=False,
+                    )
 
-            generated_latents_list.append(intermed_latents[-1])
-            last_generated_latents = intermed_latents
+                    if not intermed_latents:
+                        raise ValueError(f"Frame generation failed for unit {unit_index}")
+                    
+
+                    generated_latents_list.append(intermed_latents[-1].clone())
+                    print(f"Successfully generated frame unit {unit_index}. Shape: {generated_latents[-1].shape}")
+                    # last_generated_latents = intermed_latents
+                
+                except Exception as e:
+                    print(f"Frame generation faild for unit {unit_index}: {str(e)}")
+                    print(f"Current frames : {len(generated_latents_list)}")
+                    raise ValueError(f"Could not generate frame unit {unit_index}") from e 
+                
+        # final processing 
+        if not generated_latents_list:
+            raise ValueError("No frames were generated")
 
         generated_latents = torch.cat(generated_latents_list, dim=2)