Update pipeline.py

pipeline.py

@@ -1498,14 +1498,28 @@ class AnimateDiffPipeline(DiffusionPipeline, TextualInversionLoaderMixin, IPAdap
 
                     # sum the noise predictions for the unconditional and text conditioned noise
                     start_guidance_time = time.time()
+                    
+                    # if do_classifier_free_guidance:
+                    #     noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+
+                    #     # add the ending frames from noise_pred_uncond to the start of the noise_pred_uncond_sum
+                    #     noise_pred_uncond_sum[:, :,current_context_indexes, :, :] += noise_pred_uncond
+                    #     noise_pred_text_sum[:, :,current_context_indexes, :, :] += noise_pred_text
+                    #     #increase the counter for the ending frames
+                    #     latent_counter[current_context_indexes] += 1
+                        
+                        
                     if do_classifier_free_guidance:
-                        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                        # Split tensor along its first dimension
+                        noise_pred_uncond, noise_pred_text = torch.chunk(noise_pred, 2, dim=0)
 
-                        # add the ending frames from noise_pred_uncond to the start of the noise_pred_uncond_sum
-                        noise_pred_uncond_sum[:, :,current_context_indexes, :, :] += noise_pred_uncond
-                        noise_pred_text_sum[:, :,current_context_indexes, :, :] += noise_pred_text
-                        #increase the counter for the ending frames
+                        # Efficient in-place addition using advanced indexing
+                        noise_pred_uncond_sum[..., current_context_indexes, :, :] += noise_pred_uncond
+                        noise_pred_text_sum[..., current_context_indexes, :, :] += noise_pred_text
+
+                        # Efficient in-place increment for latent_counter
                         latent_counter[current_context_indexes] += 1
+                        
                     print("guidance time", time.time() - start_guidance_time)
 
                     # set the step index to the current batch