add chunking

app.py

@@ -55,6 +55,7 @@ import random
 from io import BytesIO
 
 import torch
+
 from cosmos_transfer1.checkpoints import (
     BASE_7B_CHECKPOINT_AV_SAMPLE_PATH,
     BASE_7B_CHECKPOINT_PATH,
@@ -70,14 +71,13 @@ from cosmos_transfer1.diffusion.inference.world_generation_pipeline import (
 )
 from cosmos_transfer1.utils import log, misc
 from cosmos_transfer1.utils.io import read_prompts_from_file, save_video
-
 from helper import parse_arguments
 
 torch.enable_grad(False)
 torch.serialization.add_safe_globals([BytesIO])
 
 
-def inference(cfg, control_inputs) -> Tuple[List[str], List[str]]:
+def inference(cfg, control_inputs, chunking) -> Tuple[List[str], List[str]]:
     video_paths = []
     prompt_paths = []
 
@@ -87,9 +87,10 @@ def inference(cfg, control_inputs) -> Tuple[List[str], List[str]]:
     device_rank = 0
     process_group = None
     if cfg.num_gpus > 1:
-        from cosmos_transfer1.utils import distributed
         from megatron.core import parallel_state
 
+        from cosmos_transfer1.utils import distributed
+
         distributed.init()
         parallel_state.initialize_model_parallel(context_parallel_size=cfg.num_gpus)
         process_group = parallel_state.get_context_parallel_group()
@@ -142,6 +143,7 @@ def inference(cfg, control_inputs) -> Tuple[List[str], List[str]]:
             upsample_prompt=cfg.upsample_prompt,
             offload_prompt_upsampler=cfg.offload_prompt_upsampler,
             process_group=process_group,
+            chunking=chunking,
         )
 
     if cfg.batch_input_path:
@@ -278,6 +280,7 @@ def generate_video(
     negative_prompt="The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality.",  # noqa: E501
     seed=42,
     randomize_seed=False,
+    chunking=False,
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
@@ -315,7 +318,7 @@ def generate_video(
     watcher = watch_gpu_memory(10)
 
     # start inference
-    videos, prompts = inference(args, control_inputs)
+    videos, prompts = inference(args, control_inputs, chunking)
 
     # print the generation time
     end_time = time.time()
@@ -361,6 +364,7 @@ with gr.Blocks() as demo:
                 randomize_seed_checkbox = gr.Checkbox(label="Randomize Seed", value=False)
                 seed_input = gr.Slider(minimum=0, maximum=1000000, value=1, step=1, label="Seed")
 
+            chunking_checkbox = gr.Checkbox(label="Chunking", value=True)
             generate_button = gr.Button("Generate Image")
 
         with gr.Column():
@@ -369,7 +373,16 @@ with gr.Blocks() as demo:
 
     generate_button.click(
         fn=generate_video,
-        inputs=[rgb_video_input, hdmap_input, lidar_input, prompt_input, negative_prompt_input, seed_input, randomize_seed_checkbox],
+        inputs=[
+            rgb_video_input,
+            hdmap_input,
+            lidar_input,
+            prompt_input,
+            negative_prompt_input,
+            seed_input,
+            randomize_seed_checkbox,
+            chunking_checkbox,
+        ],
         outputs=[output_video, output_file, seed_input],
     )
 

cosmos_transfer1/diffusion/inference/inference_utils.py

@@ -710,6 +710,7 @@ def generate_world_from_control(
     x_sigma_max=None,
     augment_sigma=None,
     use_batch_processing: bool = True,
+    chunking: bool = False,
 ) -> Tuple[np.array, list, list]:
     """Generate video using a conditioning video/image input.
 
@@ -723,6 +724,7 @@ def generate_world_from_control(
         seed (int): Random seed for generation
         condition_latent (torch.Tensor): Latent tensor from conditioning video/image file
         num_input_frames (int): Number of input frames
+        chunking: Whether to use the chunking method in generation pipeline
 
     Returns:
         np.array: Generated video frames in shape [T,H,W,C], range [0,255]
@@ -761,6 +763,7 @@ def generate_world_from_control(
         patch_h=h,
         patch_w=w,
         use_batch_processing=use_batch_processing,
+        chunking=chunking,
     )
     return sample
 

cosmos_transfer1/diffusion/inference/world_generation_pipeline.py

@@ -151,6 +151,7 @@ class DiffusionControl2WorldGenerationPipeline(BaseWorldGenerationPipeline):
         regional_prompts: List[str] = None,
         region_definitions: Union[List[List[float]], torch.Tensor] = None,
         waymo_example: bool = False,
+        chunking: bool = False,
     ):
         """Initialize diffusion world generation pipeline.
 
@@ -178,6 +179,7 @@ class DiffusionControl2WorldGenerationPipeline(BaseWorldGenerationPipeline):
             offload_prompt_upsampler: Whether to offload prompt upsampler after use
             process_group: Process group for distributed training
             waymo_example: Whether to use the waymo example post-training checkpoint
+            chunking: Whether to use the chunking method in generation pipeline
         """
         self.num_input_frames = num_input_frames
         self.control_inputs = control_inputs
@@ -201,6 +203,7 @@ class DiffusionControl2WorldGenerationPipeline(BaseWorldGenerationPipeline):
         self.seed = seed
         self.regional_prompts = regional_prompts
         self.region_definitions = region_definitions
+        self.chunking = chunking
 
         super().__init__(
             checkpoint_dir=checkpoint_dir,
@@ -621,6 +624,7 @@ class DiffusionControl2WorldGenerationPipeline(BaseWorldGenerationPipeline):
                 sigma_max=self.sigma_max if x_sigma_max is not None else None,
                 x_sigma_max=x_sigma_max,
                 use_batch_processing=False if is_upscale_case else True,
+                chunking=self.chunking,
             )
             log.info("Completed diffusion sampling")
             log.info("Starting VAE decode")

cosmos_transfer1/diffusion/model/model_v2w.py

@@ -16,6 +16,7 @@
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
+import numpy as np
 import torch
 from megatron.core import parallel_state
 from torch import Tensor
@@ -167,6 +168,39 @@ class DiffusionV2WModel(DiffusionT2WModel):
             x0_pred_replaced=x0_pred_replaced,
         )
 
+    CHUNKING_SIZE = 4
+    CHUNKING_MODE = "rand_order"  # ["shuffle", "in_order", "rand_order"]
+    IS_STAGGERED = True
+
+    def get_chunks_indices(self, total_flen) -> List[torch.Tensor]:
+        chunks_indices = []
+        if self.CHUNKING_MODE == "shuffle":
+            for index in torch.arange(0, total_flen, 1).split(self.CHUNKING_SIZE):
+                chunks_indices.append(index)
+            np.random.shuffle(chunks_indices)
+        else:
+            first_chunk_end = (
+                int(torch.randint(low=0, high=self.CHUNKING_SIZE, size=(1,)) + 1) if self.IS_STAGGERED else self.CHUNKING_SIZE
+            )
+
+            if first_chunk_end >= total_flen:
+                chunks_indices.append(torch.arange(total_flen))
+            else:
+                chunks_indices.append(torch.arange(first_chunk_end))
+
+                for index in torch.arange(first_chunk_end, total_flen, 1).split(self.CHUNKING_SIZE):
+                    chunks_indices.append(index)
+
+                if self.CHUNKING_MODE == "in_order":
+                    pass
+                elif self.CHUNKING_MODE == "rand_order":
+                    if np.random.rand() > 0.5:
+                        chunks_indices = chunks_indices[::-1]
+                else:
+                    raise NotImplementedError(f"{self.CHUNKING_MODE} mode not implemented!!")
+
+        return chunks_indices
+
     def generate_samples_from_batch(
         self,
         data_batch: Dict,
@@ -182,6 +216,7 @@ class DiffusionV2WModel(DiffusionT2WModel):
         add_input_frames_guidance: bool = False,
         x_sigma_max: Optional[torch.Tensor] = None,
         sigma_max: Optional[float] = None,
+        chunking: bool = False,
         **kwargs,
     ) -> Tensor:
         """Generates video samples conditioned on input frames.
@@ -199,6 +234,7 @@ class DiffusionV2WModel(DiffusionT2WModel):
             condition_video_augment_sigma_in_inference: Noise level for condition augmentation
             add_input_frames_guidance: Whether to apply guidance to input frames
             x_sigma_max: Maximum noise level tensor
+            chunking: Whether to use the chunking method in generation pipeline
 
         Returns:
             Generated video samples tensor
@@ -213,6 +249,7 @@ class DiffusionV2WModel(DiffusionT2WModel):
 
         assert condition_latent is not None, "condition_latent should be provided"
 
+        # try to add chunking here !!!
         x0_fn = self.get_x0_fn_from_batch_with_condition_latent(
             data_batch,
             guidance,
@@ -222,6 +259,7 @@ class DiffusionV2WModel(DiffusionT2WModel):
             condition_video_augment_sigma_in_inference=condition_video_augment_sigma_in_inference,
             add_input_frames_guidance=add_input_frames_guidance,
             seed=seed,
+            chunking=chunking,
         )
         if sigma_max is None:
             sigma_max = self.sde.sigma_max
@@ -256,6 +294,7 @@ class DiffusionV2WModel(DiffusionT2WModel):
         condition_video_augment_sigma_in_inference: float = None,
         add_input_frames_guidance: bool = False,
         seed: int = 1,
+        chunking: bool = False,
     ) -> Callable:
         """Creates denoising function for conditional video generation.
 
@@ -268,44 +307,92 @@ class DiffusionV2WModel(DiffusionT2WModel):
             condition_video_augment_sigma_in_inference: Noise level for condition augmentation
             add_input_frames_guidance: Whether to apply guidance to input frames
             seed: Random seed for reproducibility
+            chunking: Whether to use the chunking method in generation pipeline
 
         Returns:
             Function that takes noisy input and noise level and returns denoised prediction
         """
-        if is_negative_prompt:
-            condition, uncondition = self.conditioner.get_condition_with_negative_prompt(data_batch)
-        else:
-            condition, uncondition = self.conditioner.get_condition_uncondition(data_batch)
-
-        condition.video_cond_bool = True
-        condition = self.add_condition_video_indicator_and_video_input_mask(
-            condition_latent, condition, num_condition_t
-        )
+        if not chunking:
+            if is_negative_prompt:
+                condition, uncondition = self.conditioner.get_condition_with_negative_prompt(data_batch)
+            else:
+                condition, uncondition = self.conditioner.get_condition_uncondition(data_batch)
+
+            condition.video_cond_bool = True
+            condition = self.add_condition_video_indicator_and_video_input_mask(
+                condition_latent, condition, num_condition_t
+            )
 
-        uncondition.video_cond_bool = False if add_input_frames_guidance else True
-        uncondition = self.add_condition_video_indicator_and_video_input_mask(
-            condition_latent, uncondition, num_condition_t
-        )
+            uncondition.video_cond_bool = False if add_input_frames_guidance else True
+            uncondition = self.add_condition_video_indicator_and_video_input_mask(
+                condition_latent, uncondition, num_condition_t
+            )
 
-        def x0_fn(noise_x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
-            cond_x0 = self.denoise(
-                noise_x,
-                sigma,
-                condition,
-                condition_video_augment_sigma_in_inference=condition_video_augment_sigma_in_inference,
-                seed=seed,
-            ).x0_pred_replaced
-            uncond_x0 = self.denoise(
-                noise_x,
-                sigma,
-                uncondition,
-                condition_video_augment_sigma_in_inference=condition_video_augment_sigma_in_inference,
-                seed=seed,
-            ).x0_pred_replaced
-
-            return cond_x0 + guidance * (cond_x0 - uncond_x0)
-
-        return x0_fn
+            def x0_fn(noise_x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
+                cond_x0 = self.denoise(
+                    noise_x,
+                    sigma,
+                    condition,
+                    condition_video_augment_sigma_in_inference=condition_video_augment_sigma_in_inference,
+                    seed=seed,
+                ).x0_pred_replaced
+                uncond_x0 = self.denoise(
+                    noise_x,
+                    sigma,
+                    uncondition,
+                    condition_video_augment_sigma_in_inference=condition_video_augment_sigma_in_inference,
+                    seed=seed,
+                ).x0_pred_replaced
+
+                return cond_x0 + guidance * (cond_x0 - uncond_x0)
+
+            return x0_fn
+        else:
+            log.critical("GO CHUNKING !!!")
+
+            def x0_fn(noise_x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
+                if is_negative_prompt:
+                    condition, uncondition = self.conditioner.get_condition_with_negative_prompt(data_batch)
+                else:
+                    condition, uncondition = self.conditioner.get_condition_uncondition(data_batch)
+
+                noises = torch.zeros_like(condition_latent)
+                T = condition_latent.shape[2]
+                for chunk_idx in self.get_chunks_indices(T):
+                    latents_ = condition_latent[:, :, chunk_idx, :, :]
+                    log.info(f"chunk_idx: {chunk_idx}, chunk shape: {latents_.shape}")
+                    # controlnet_cond_ = self.controlnet_data[:, chunk_idx]
+
+                    condition.video_cond_bool = True
+                    condition = self.add_condition_video_indicator_and_video_input_mask(
+                        latents_, condition, num_condition_t
+                    )
+
+                    uncondition.video_cond_bool = False if add_input_frames_guidance else True
+                    uncondition = self.add_condition_video_indicator_and_video_input_mask(
+                        latents_, uncondition, num_condition_t
+                    )
+
+                    cond_x0 = self.denoise(
+                        noise_x,
+                        sigma,
+                        condition,
+                        condition_video_augment_sigma_in_inference=condition_video_augment_sigma_in_inference,
+                        seed=seed,
+                    ).x0_pred_replaced
+                    uncond_x0 = self.denoise(
+                        noise_x,
+                        sigma,
+                        uncondition,
+                        condition_video_augment_sigma_in_inference=condition_video_augment_sigma_in_inference,
+                        seed=seed,
+                    ).x0_pred_replaced
+
+                    noises[:, :, chunk_idx, :, :] = cond_x0 + guidance * (cond_x0 - uncond_x0)
+
+                # TODO: need scheduler ?
+                return noises
+            return x0_fn
 
     def add_condition_video_indicator_and_video_input_mask(
         self, latent_state: torch.Tensor, condition: VideoExtendCondition, num_condition_t: Union[int, None] = None