Update pipeline.py

pipeline.py

@@ -1006,22 +1006,33 @@ class AnimateDiffPipeline(DiffusionPipeline, TextualInversionLoaderMixin, IPAdap
         
         # divide the initial latents into context groups
         
-        def context_scheduler(context_size, overlap, offset, num_frames, num_timesteps):
-            num_context_groups = (num_frames // (context_size-overlap))+1
-            context_indexes = []
-            for t in range(num_timesteps):
-                context_groups = []
-                for context_group_index in range(num_context_groups):
-                    context_group = []
-                    for i in range(context_size):
-                        # calculate the frame index
-                        frame_index = ((t+1) * context_group_index * (context_size-overlap)) + (offset * t) + i
-                        # wrap around at the end
-                        if frame_index >= num_frames:frame_index = frame_index % num_frames
-                        context_group.append(frame_index)
-                    context_groups.append(context_groups)
-                context_indexes.append(context_groups)
-            return context_indexes
+        def context_scheduler(context_size, overlap, offset, total_frames, total_timesteps):
+            # Calculate the number of context groups based on frame count and context size
+            number_of_context_groups = (total_frames // (context_size - overlap)) + 1
+            # Initialize a list to store context indexes for all timesteps
+            all_context_indexes = []
+            # Iterate over each timestep
+            for timestep in range(total_timesteps):
+                # Initialize a list to store groups of context indexes for this timestep
+                timestep_context_groups = []
+                # Iterate over each context group
+                for group_index in range(number_of_context_groups):
+                    # Initialize a list to store context indexes for this group
+                    context_group_indexes = []
+                    # Iterate over each index in the context group
+                    for index in range(context_size):
+                        # Calculate the frame index
+                        frame_index = ((timestep + 1) * group_index * (context_size - overlap)) + (offset * timestep) + index
+                        # If frame index exceeds total frames, wrap around
+                        if frame_index >= total_frames:
+                            frame_index %= total_frames
+                        # Add the frame index to the group's list
+                        context_group_indexes.append(frame_index)
+                    # Add the group's indexes to the timestep's list
+                    timestep_context_groups.append(context_group_indexes)
+                # Add the timestep's context groups to the overall list
+                all_context_indexes.append(timestep_context_groups)
+            return all_context_indexes
                 
         context_indexes = context_scheduler(context_size, overlap, step, num_frames, len(timesteps))