Update pipeline.py

pipeline.py

@@ -1497,28 +1497,16 @@ class AnimateDiffPipeline(DiffusionPipeline, TextualInversionLoaderMixin, IPAdap
                         
                     if do_classifier_free_guidance:
                         # Split tensor along its first dimension
+                        time_chunk = time.time()
                         noise_pred_uncond, noise_pred_text = torch.chunk(noise_pred, 2, dim=0)
-
-                        # Reshape or expand noise_pred_uncond and noise_pred_text to match the dimensions of the sum tensors
-                        # This step depends on the dimensions of your tensors and how they need to align for the operation
-                        expanded_noise_pred_uncond = noise_pred_uncond.expand_as(noise_pred_uncond_sum[..., current_context_indexes, :, :])
-                        expanded_noise_pred_text = noise_pred_text.expand_as(noise_pred_text_sum[..., current_context_indexes, :, :])
-
-                        # Perform batch addition
-                        noise_pred_uncond_sum[..., current_context_indexes, :, :] += expanded_noise_pred_uncond
-                        noise_pred_text_sum[..., current_context_indexes, :, :] += expanded_noise_pred_text
-                        
-                        # print devices and shapes for everything
-                        print("noise_pred_uncond_sum", noise_pred_uncond_sum.device, noise_pred_uncond_sum.shape)
-                        print("noise_pred_text_sum", noise_pred_text_sum.device, noise_pred_text_sum.shape)
-                        print("expanded_noise_pred_uncond", expanded_noise_pred_uncond.device, expanded_noise_pred_uncond.shape)
-                        print("expanded_noise_pred_text", expanded_noise_pred_text.device, expanded_noise_pred_text.shape)
-                        print("current_context_latents", current_context_latents.device, current_context_latents.shape)
-                        print("latent_counter", latent_counter.device, latent_counter.shape)
-                        
-
-                        # Batch increment for latent_counter
-                        latent_counter[current_context_indexes] += 1
+                        print("chunk time", time.time() - time_chunk)
+                        for i,context_index in range(len(current_context_indexes)):
+                            # Perform batch addition
+                            noise_pred_uncond_sum[..., context_index, :, :] += noise_pred_uncond[ :, :,i, :, :]
+                            noise_pred_text_sum[..., context_index, :, :] += noise_pred_text[ :, :,i, :, :]
+
+                            # Batch increment for latent_counter
+                            latent_counter[context_index] += 1
                         
                     print("guidance time", time.time() - start_guidance_time)