Update pipeline.py

pipeline.py

@@ -1011,7 +1011,9 @@ 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):
-                latent_sum = torch.zeros_like(latents).to(device).to(dtype=torch.float16)
+                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_sum = torch.zeros_like(latents).to(device).to(dtype=torch.float16)
                 latent_counter = torch.zeros(num_frames).to(device).to(dtype=torch.float16)
         
                 # foreach context group seperately denoise the current timestep
@@ -1045,37 +1047,45 @@ class AnimateDiffPipeline(DiffusionPipeline, TextualInversionLoaderMixin, IPAdap
                     # perform guidance
                     if do_classifier_free_guidance:
                         noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-                        
+                        noise_pred_uncond_sum += noise_pred_uncond
+                        noise_pred_text_sum += noise_pred_text
                     
                     # set the step index to the current batch
-                    # self.scheduler._step_index = i
-                    
-                    # compute the previous noisy sample x_t -> x_t-1
-                    current_context_latents = self.scheduler.step(noise_pred, t, current_context_latents, **extra_step_kwargs).prev_sample
+                    self.scheduler._step_index = i
                     
                     # if context_start + context_size > num_frames: remove the appended frames from the end of the current_context_latents
-                    if wrap_count > 0:
-                        # add the ending frames from current_context_latents to the start of the latent_sum
-                        latent_sum[:, :, 0:wrap_count, :, :] += current_context_latents[:, :, -wrap_count:, :, :]
-                        # increase the counter for the ending frames
-                        latent_counter[0:wrap_count] += 1
-                        # remove the ending frames from current_context_latents
-                        current_context_latents = current_context_latents[:, :, :-wrap_count, :, :]
+                    # if wrap_count > 0:
+                    #     # add the ending frames from current_context_latents to the start of the latent_sum
+                    #     latent_sum[:, :, 0:wrap_count, :, :] += current_context_latents[:, :, -wrap_count:, :, :]
+                    #     # increase the counter for the ending frames
+                    #     latent_counter[0:wrap_count] += 1
+                    #     # remove the ending frames from current_context_latents
+                    #     current_context_latents = current_context_latents[:, :, :-wrap_count, :, :]
 
                     #add the context current_context_latents back to the latent sum starting from the current context start
-                    latent_sum[:, :, current_context_start : current_context_start + context_size, :, :] += current_context_latents    
+                    # latent_sum[:, :, current_context_start : current_context_start + context_size, :, :] += current_context_latents    
+                    
                     # add one to the counter for each timestep in the context
                     latent_counter[current_context_start : current_context_start + context_size] += 1
                     
+                # perform guidance
+                if do_classifier_free_guidance:
+                    latent_counter = latent_counter.reshape(1, 1, num_frames, 1, 1)
+                    noise_pred_uncond = noise_pred_uncond_sum / latent_counter
+                    noise_pred_text = noise_pred_text_sum / latent_counter
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+                
+                # compute the previous noisy sample x_t -> x_t-1
+                current_context_latents = self.scheduler.step(noise_pred, t, current_context_latents, **extra_step_kwargs).prev_sample
+                
                 # call the callback, if provided
                 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                     progress_bar.update()
                     if callback is not None and i % callback_steps == 0:
                         callback(i, t, None)
                     
-                latent_counter = latent_counter.reshape(1, 1, num_frames, 1, 1)
-                latents = latent_sum / latent_counter
+
+                # latents = latent_sum / latent_counter
                 
                 # shuffle rotate latent images by step places, wrapping around the last 2 to the start
                 latents = torch.cat([latents[:, :, -step:, :, :], latents[:, :, :-step, :, :]], dim=2)