Update pipeline.py

pipeline.py

@@ -977,10 +977,16 @@ class AnimateDiffPipeline(DiffusionPipeline, TextualInversionLoaderMixin, IPAdap
                     
                     # select the relevent context from the latents
                     current_context_latents = latents[:, :, current_context_start : current_context_start + context_size, :, :]
+                    
+                    wrap_count = max(current_context_start + context_size - num_frames, 0)
+                    
                     # if context_start + context_size > num_frames: append the remaining frames from the start of the latents
-                    if current_context_start + context_size > num_frames:
-                        print(f"Appending {max(current_context_start + context_size - num_frames, 0)} frames from the start of the latents")    
-                        current_context_latents = torch.cat([current_context_latents, latents[:, :, :max(current_context_start + context_size - num_frames, 0), :, :]], dim=2)
+                    if wrap_count > 0:
+                        print(f"Appending {wrap_count} frames from the start of the latents")    
+                        current_context_latents = torch.cat([current_context_latents, latents[:, :, :wrap_count, :, :]], dim=2)
+                        
+                    # print number of frames in the context
+                    print(f"Number of frames in the context: {current_context_latents.shape[2]}")
                 
                     # expand the latents if we are doing classifier free guidance
                     latent_model_input = torch.cat([current_context_latents] * 2) if do_classifier_free_guidance else current_context_latents
@@ -1004,13 +1010,13 @@ class AnimateDiffPipeline(DiffusionPipeline, TextualInversionLoaderMixin, IPAdap
                     current_context_latents = self.scheduler.step(noise_pred, t, current_context_latents, **extra_step_kwargs).prev_sample
                     
                     # if context_start + context_size > num_frames: remove the appended frames from the end of the current_context_latents
-                    if current_context_start + context_size > num_frames:
+                    if wrap_count > 0:
                         # add the ending frames from current_context_latents to the start of the latent_sum
-                        latent_sum[:, :, -max(current_context_start + context_size - num_frames, 0):, :, :] += current_context_latents[:, :, -max(current_context_start + context_size - num_frames, 0):, :, :]
+                        latent_sum[:, :, 0:wrap_count, :, :] += current_context_latents[:, :, -wrap_count:, :, :]
                         # increase the counter for the ending frames
-                        latent_counter[-max(current_context_start + context_size - num_frames, 0):] += 1
+                        latent_counter[0:wrap_count] += 1
                         # remove the ending frames from current_context_latents
-                        current_context_latents = current_context_latents[:, :, :-max(current_context_start + context_size - num_frames, 0), :, :]
+                        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