Update pipeline.py

pipeline.py

@@ -1407,6 +1407,8 @@ class AnimateDiffPipeline(DiffusionPipeline, TextualInversionLoaderMixin, IPAdap
         num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
         with self.progress_bar(total=len(timesteps)) as progress_bar:
             for i, t in enumerate(timesteps):
+                print("i", i)
+                print("t", t)
                 noise_pred_uncond_sum = torch.zeros_like(latents).to(device).to(dtype=torch.float16)
                 noise_pred_text_sum = torch.zeros_like(latents).to(device).to(dtype=torch.float16)
                 latent_counter = torch.zeros(num_frames).to(device).to(dtype=torch.float16)
@@ -1424,7 +1426,7 @@ class AnimateDiffPipeline(DiffusionPipeline, TextualInversionLoaderMixin, IPAdap
                     latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
                     
                     
-                    if self.controlnet != None or i > 2:
+                    if self.controlnet != None and i < 4:
                         contorl_start = time.time()
                         
                         current_context_conditioning_frames = conditioning_frames[current_context_indexes, :, :, :] 
@@ -1467,7 +1469,7 @@ class AnimateDiffPipeline(DiffusionPipeline, TextualInversionLoaderMixin, IPAdap
                         )
                         print("controlnet time", time.time() - contorl_start)
                         
-                        
+                        unet_start = time.time()
                         # predict the noise residual with the added controlnet residuals
                         noise_pred = self.unet(
                             latent_model_input,
@@ -1478,8 +1480,10 @@ class AnimateDiffPipeline(DiffusionPipeline, TextualInversionLoaderMixin, IPAdap
                             down_block_additional_residuals=down_block_res_samples,
                             mid_block_additional_residual=mid_block_res_sample,
                         ).sample
+                        print("unet time", time.time() - unet_start)
                     
                     else:
+                        unet_start = time.time()
                         # predict the noise residual without contorlnet
                         noise_pred = self.unet(
                             latent_model_input,
@@ -1488,6 +1492,7 @@ class AnimateDiffPipeline(DiffusionPipeline, TextualInversionLoaderMixin, IPAdap
                             cross_attention_kwargs=cross_attention_kwargs,
                             added_cond_kwargs=added_cond_kwargs,
                         ).sample
+                        print("unet time", time.time() - unet_start)
 
                     # sum the noise predictions for the unconditional and text conditioned noise
                     if do_classifier_free_guidance: