Update pipeline.py

pipeline.py

@@ -1528,13 +1528,13 @@ class AnimateDiffPipeline(DiffusionPipeline, TextualInversionLoaderMixin, IPAdap
                         torch.cuda.synchronize()  # Synchronize GPU before batch addition
                         time_batch_addition_start = time.time()
 
-                        for z, context_index in enumerate(current_context_indexes):
-                            # Perform batch addition
-                            noise_pred_uncond_sum[..., context_index, :, :] += noise_pred_uncond[:, :, z, :, :]
-                            noise_pred_text_sum[..., context_index, :, :] += noise_pred_text[:, :, z, :, :]
+                        current_context_tensor = torch.tensor(current_context_indexes, device=latent_counter.device)
+                        # Perform batch addition
+                        noise_pred_uncond_sum[..., current_context_tensor, :, :] += noise_pred_uncond.transpose(1, 2)
+                        noise_pred_text_sum[..., current_context_tensor, :, :] += noise_pred_text.transpose(1, 2)
 
-                            # Batch increment for latent_counter
-                            latent_counter[context_index] += 1
+                        # Batch increment for latent_counter
+                        latent_counter[current_context_tensor] += 1
 
                         torch.cuda.synchronize()  # Synchronize GPU after batch addition
                         time_batch_addition_end = time.time()