Update models/pipeline_controlvideo.py

models/pipeline_controlvideo.py

@@ -652,7 +652,7 @@ class ControlVideoPipeline(DiffusionPipeline):
         prev_timestep = timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
         alpha_prod_t_prev = self.scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.scheduler.final_alpha_cumprod
         return alpha_prod_t_prev
-
+    """
     def get_slide_window_indices(self, video_length, window_size):
         assert window_size >=3 
         key_frame_indices = np.arange(0, video_length, window_size-1).tolist()
@@ -668,7 +668,45 @@ class ControlVideoPipeline(DiffusionPipeline):
                 continue
             inter_frame_list.append(s[1:].tolist())
         return key_frame_indices, inter_frame_list
-        
+    """
+    def get_slide_window_indices(self, video_length, window_size):
+    assert window_size >= 3
+
+    # Define the chunk size for processing
+    chunk_size = 4
+
+    # Calculate the number of chunks
+    num_chunks = (video_length - 1) // chunk_size + 1
+
+    # Initialize the lists to store the results
+    key_frame_indices = []
+    inter_frame_list = []
+
+    for chunk_index in range(num_chunks):
+        # Calculate the start and end indices for the current chunk
+        start_index = chunk_index * chunk_size
+        end_index = min((chunk_index + 1) * chunk_size, video_length)
+
+        # Generate key frame indices for the current chunk
+        chunk_key_frame_indices = np.arange(start_index, end_index, window_size - 1).tolist()
+
+        # Append the last index if it's not already included
+        if chunk_key_frame_indices[-1] != (end_index - 1):
+            chunk_key_frame_indices.append(end_index - 1)
+
+        # Append the key frame indices of the current chunk to the overall list
+        key_frame_indices.extend(chunk_key_frame_indices)
+
+        # Generate slices for the current chunk
+        chunk_slices = np.split(np.arange(start_index, end_index), chunk_key_frame_indices)
+
+        # Process each slice in the current chunk
+        for s in chunk_slices:
+            if len(s) < 2:
+                continue
+            inter_frame_list.append(s[1:].tolist())
+
+    return key_frame_indices, inter_frame_list
     @torch.no_grad()
     def __call__(
         self,