cosmos_transfer1/diffusion/training/models/model_ctrl.py

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Callable, Dict, Optional, Tuple, Type, TypeVar, Union

import torch
from einops import rearrange
from megatron.core import parallel_state
from torch import Tensor

from cosmos_transfer1.diffusion.conditioner import CosmosCondition, DataType, VideoConditionerWithCtrl
from cosmos_transfer1.diffusion.diffusion.modules.res_sampler import COMMON_SOLVER_OPTIONS
from cosmos_transfer1.diffusion.inference.inference_utils import (
    merge_patches_into_video,
    non_strict_load_model,
    split_video_into_patches,
)
from cosmos_transfer1.diffusion.module.parallel import cat_outputs_cp, split_inputs_cp
from cosmos_transfer1.diffusion.training.models.extend_model import ExtendDiffusionModel as ExtendVideoDiffusionModel
from cosmos_transfer1.diffusion.training.models.model import DiffusionModel as VideoDiffusionModel
from cosmos_transfer1.diffusion.training.models.model import _broadcast, broadcast_condition
from cosmos_transfer1.diffusion.training.models.model_image import diffusion_fsdp_class_decorator
from cosmos_transfer1.utils import log, misc
from cosmos_transfer1.utils.lazy_config import instantiate

T = TypeVar("T")
IS_PREPROCESSED_KEY = "is_preprocessed"


def ctrlnet_decorator(base_class: Type[T]) -> Type[T]:
    class CtrlNetModel(base_class):
        def __init__(self, config, fsdp_checkpointer=None):
            if fsdp_checkpointer is not None:
                return super().__init__(config, fsdp_checkpointer)
            else:
                return super().__init__(config)

        def build_model(self) -> torch.nn.ModuleDict:
            log.info("Start creating base model")
            base_model = super().build_model()
            # initialize base model
            config = self.config
            self.load_base_model(base_model)
            log.info("Done creating base model")

            log.info("Start creating ctrlnet model")
            net = instantiate(self.config.net_ctrl)
            conditioner = base_model.conditioner
            logvar = base_model.logvar
            # initialize controlnet encoder
            model = torch.nn.ModuleDict({"net": net, "conditioner": conditioner, "logvar": logvar})
            model.load_state_dict(base_model.state_dict(), strict=False)

            model.base_model = base_model
            if not config.finetune_base_model:
                model.base_model.requires_grad_(False)
                log.critical("Only training ctrlnet model and keeping base model frozen")
            else:
                log.critical("Also training base model")
            log.info("Done creating ctrlnet model")

            self.hint_key = self.config.hint_key["hint_key"]
            return model

        @property
        def base_net(self):
            return self.model.base_model.net

        def on_train_start(self, memory_format: torch.memory_format = torch.preserve_format) -> None:
            super().on_train_start(memory_format)
            # self.base_model = self.base_model.to(memory_format=memory_format, **self.tensor_kwargs)
            self.model = self.model.to(memory_format=memory_format, **self.tensor_kwargs)
            if parallel_state.is_initialized() and parallel_state.get_tensor_model_parallel_world_size() > 1:
                if parallel_state.sequence_parallel:
                    self.base_net.enable_sequence_parallel()
            if (
                hasattr(self.config, "use_torch_compile") and self.config.use_torch_compile
            ):  # compatible with old config
                # not tested yet
                if torch.__version__ < "2.3":
                    log.warning(
                        "torch.compile in Pytorch version older than 2.3 doesn't work well with activation checkpointing.\n"
                        "It's very likely there will be no significant speedup from torch.compile.\n"
                        "Please use at least 24.04 Pytorch container, or imaginaire4:v7 container."
                    )
                self.base_net = torch.compile(self.base_net, dynamic=False, disable=not self.config.use_torch_compile)

        def load_base_model(self, base_model) -> None:
            config = self.config
            if config.base_load_from is not None:
                checkpoint_path = config.base_load_from["load_path"]
            else:
                checkpoint_path = ""

            if "*" in checkpoint_path:
                # there might be better ways to decide if it's a converted tp checkpoint
                mp_rank = parallel_state.get_model_parallel_group().rank()
                checkpoint_path = checkpoint_path.replace("*", f"{mp_rank}")

            if checkpoint_path:
                log.info(f"Loading base model checkpoint (local): {checkpoint_path}")
                state_dict = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
                log.success(f"Complete loading base model checkpoint (local): {checkpoint_path}")

                if state_dict.get("ema") is not None:
                    # Copy the base model weights from ema model.
                    log.info("Copying ema to base model")
                    base_state_dict = {k.replace("-", "."): v for k, v in state_dict["ema"].items()}
                elif "model" in state_dict:
                    # Copy the base model weights from reg model.
                    log.warning("Using non-EMA base model")
                    base_state_dict = state_dict["model"]
                else:
                    log.info("Loading from an EMA only model")
                    base_state_dict = state_dict
                try:
                    base_model.load_state_dict(base_state_dict, strict=False)
                except Exception:
                    log.critical("load model in non-strict mode")
                    log.critical(non_strict_load_model(base_model, base_state_dict), rank0_only=False)
            log.info("Done loading the base model checkpoint.")

    return CtrlNetModel


def video_ctrlnet_decorator(base_class: Type[T]) -> Type[T]:
    class VideoDiffusionModelWithCtrlWrapper(base_class):
        def __init__(self, config):
            super().__init__(config)
            if hasattr(config, "pixel_corruptor") and config.pixel_corruptor is not None:
                self.pixel_corruptor = instantiate(config.pixel_corruptor)
                self.pixel_corruptor.to(**self.tensor_kwargs)
            else:
                self.pixel_corruptor = None

        def get_data_and_condition(
            self, data_batch: dict[str, Tensor], **kwargs
        ) -> Tuple[Tensor, VideoConditionerWithCtrl]:
            # process the control input
            hint_key = self.config.hint_key["hint_key"]
            is_image_batch = self.is_image_batch(data_batch)
            _data = {hint_key: data_batch[hint_key]}
            if IS_PREPROCESSED_KEY in data_batch:
                _data[IS_PREPROCESSED_KEY] = data_batch[IS_PREPROCESSED_KEY]
            if not is_image_batch:
                self._normalize_video_databatch_inplace(_data, input_key=hint_key)
            # if it is an image batch, the control input is also image
            if self.input_image_key in data_batch:
                self._augment_image_dim_inplace(_data, input_key=hint_key)
            data_batch[hint_key] = _data[hint_key]
            # else:
            #     raise NotImplementedError(f"{self.config.hint_key} is not implemented.")
            data_batch["hint_key"] = hint_key
            raw_state, latent_state, condition = super().get_data_and_condition(data_batch, **kwargs)
            # if not torch.is_grad_enabled() and all(self.config.hint_mask):
            use_multicontrol = (
                ("control_weight" in data_batch)
                and not isinstance(data_batch["control_weight"], float)
                and data_batch["control_weight"].shape[0] > 1
            )
            if use_multicontrol:  # encode individual conditions separately
                latent_hint = []
                num_conditions = data_batch[data_batch["hint_key"]].size(1) // 3
                for i in range(num_conditions):
                    cond_mask = [False] * num_conditions
                    cond_mask[i] = True
                    latent_hint += [self.encode_latent(data_batch, cond_mask=cond_mask)]
                latent_hint = torch.cat(latent_hint)
            else:
                latent_hint = self.encode_latent(data_batch)
            # copied from model.py
            is_image_batch = self.is_image_batch(data_batch)
            is_video_batch = not is_image_batch
            # VAE has randomness. CP/TP group should have the same encoded output.

            latent_hint = _broadcast(latent_hint, to_tp=True, to_cp=is_video_batch)

            # add extra conditions
            data_batch["latent_hint"] = latent_hint
            setattr(condition, hint_key, latent_hint)
            setattr(condition, "base_model", self.model.base_model)
            return raw_state, latent_state, condition

        def encode_latent(self, data_batch: dict, cond_mask: list = []) -> torch.Tensor:
            x = data_batch[data_batch["hint_key"]]
            if torch.is_grad_enabled() and self.pixel_corruptor is not None:
                x = self.pixel_corruptor(x)
            latent = []
            # control input goes through tokenizer, which always takes 3-input channels
            num_conditions = x.size(1) // 3  # input conditions were concatenated along channel dimension
            if num_conditions > 1 and self.config.hint_dropout_rate > 0:
                if torch.is_grad_enabled():  # during training, randomly dropout some conditions
                    cond_mask = torch.rand(num_conditions) > self.config.hint_dropout_rate
                    if not cond_mask.any():  # make sure at least one condition is present
                        cond_mask[torch.randint(num_conditions, (1,)).item()] = True
                elif not cond_mask:  # during inference, use hint_mask to indicate which conditions are used
                    cond_mask = self.config.hint_mask
            else:
                cond_mask = [True] * num_conditions
            for idx in range(0, x.size(1), 3):
                x_rgb = x[:, idx : idx + 3]
                if self.config.hint_key["grayscale"]:
                    x_rgb = x_rgb.mean(dim=1, keepdim=True).expand_as(x_rgb)
                # if idx == 0:
                #     x_max = x_rgb
                # else:
                #     x_max = torch.maximum(x_rgb, x_max)
                if not cond_mask[idx // 3]:  # if the condition is not selected, replace with a black image
                    x_rgb = torch.zeros_like(x_rgb)
                latent.append(self.encode(x_rgb))
            # latent.append(self.encode(x_max))
            latent = torch.cat(latent, dim=1)
            return latent

        def compute_loss_with_epsilon_and_sigma(
            self,
            data_batch: dict[str, Tensor],
            x0_from_data_batch: Tensor,
            x0: Tensor,
            condition: CosmosCondition,
            epsilon: Tensor,
            sigma: Tensor,
        ):
            if self.is_image_batch(data_batch):
                # Turn off CP
                self.net.disable_context_parallel()
                self.base_net.disable_context_parallel()
            else:
                if parallel_state.get_context_parallel_world_size() > 1:
                    # Turn on CP
                    cp_group = parallel_state.get_context_parallel_group()
                    self.net.enable_context_parallel(cp_group)
                    self.base_net.enable_context_parallel(cp_group)
                    log.debug("[CP] Split hint_input")
                    hint_key = self.config.hint_key["hint_key"]
                    x_hint_raw = getattr(condition, hint_key)
                    x_hint = split_inputs_cp(x=x_hint_raw, seq_dim=2, cp_group=self.net.cp_group)
                    setattr(condition, hint_key, x_hint)
            return super().compute_loss_with_epsilon_and_sigma(
                data_batch, x0_from_data_batch, x0, condition, epsilon, sigma
            )

        def get_x0_fn_from_batch(
            self,
            data_batch: Dict,
            guidance: float = 1.5,
            is_negative_prompt: bool = False,
            condition_latent: torch.Tensor = None,
            num_condition_t: Union[int, None] = None,
            condition_video_augment_sigma_in_inference: float = None,
            seed_inference: int = 1,
        ) -> Callable:
            """
            Generates a callable function `x0_fn` based on the provided data batch and guidance factor.

            This function first processes the input data batch through a conditioning workflow (`conditioner`) to obtain conditioned and unconditioned states. It then defines a nested function `x0_fn` which applies a denoising operation on an input `noise_x` at a given noise level `sigma` using both the conditioned and unconditioned states.

            Args:
            - data_batch (Dict): A batch of data used for conditioning. The format and content of this dictionary should align with the expectations of the `self.conditioner`
            - guidance (float, optional): A scalar value that modulates the influence of the conditioned state relative to the unconditioned state in the output. Defaults to 1.5.
            - is_negative_prompt (bool): use negative prompt t5 in uncondition if true
             condition_latent (torch.Tensor): latent tensor in shape B,C,T,H,W as condition to generate video.
            - num_condition_t (int): number of condition latent T, used in inference to decide the condition region and config.conditioner.video_cond_bool.condition_location == "first_n"
            - condition_video_augment_sigma_in_inference (float): sigma for condition video augmentation in inference

            Returns:
            - Callable: A function `x0_fn(noise_x, sigma)` that takes two arguments, `noise_x` and `sigma`, and return x0 predictoin

            The returned function is suitable for use in scenarios where a denoised state is required based on both conditioned and unconditioned inputs, with an adjustable level of guidance influence.
            """
            # data_batch should be the one processed by self.get_data_and_condition
            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)
            if hasattr(self, "is_extend_model") and self.is_extend_model:
                # Add conditions for long video generation.
                if self.is_image_batch(data_batch):
                    condition.data_type = DataType.IMAGE
                    uncondition.data_type = DataType.IMAGE
                else:
                    if condition_latent is None:
                        condition_latent = torch.zeros(data_batch["latent_hint"].shape, **self.tensor_kwargs)
                        num_condition_t = 0
                        condition_video_augment_sigma_in_inference = 1000

                    condition.video_cond_bool = True
                    condition = self.add_condition_video_indicator_and_video_input_mask(
                        condition_latent, condition, num_condition_t
                    )
                    if self.config.conditioner.video_cond_bool.add_pose_condition:
                        condition = self.add_condition_pose(data_batch, condition)

                    uncondition.video_cond_bool = True  # Not do cfg on condition frames
                    uncondition = self.add_condition_video_indicator_and_video_input_mask(
                        condition_latent, uncondition, num_condition_t
                    )
                    if self.config.conditioner.video_cond_bool.add_pose_condition:
                        uncondition = self.add_condition_pose(data_batch, uncondition)

            # Add extra conditions for ctrlnet.
            latent_hint = data_batch["latent_hint"]
            hint_key = data_batch["hint_key"]
            setattr(condition, hint_key, latent_hint)
            if "use_none_hint" in data_batch and data_batch["use_none_hint"]:
                setattr(uncondition, hint_key, None)
            else:
                setattr(uncondition, hint_key, latent_hint)

            to_cp = self.net.is_context_parallel_enabled
            # For inference, check if parallel_state is initialized
            if parallel_state.is_initialized() and not self.is_image_batch(data_batch):
                condition = broadcast_condition(condition, to_tp=True, to_cp=to_cp)
                uncondition = broadcast_condition(uncondition, to_tp=True, to_cp=to_cp)

                cp_group = parallel_state.get_context_parallel_group()
                latent_hint = getattr(condition, hint_key)
                latent_hint = split_inputs_cp(latent_hint, seq_dim=2, cp_group=cp_group)
                setattr(condition, hint_key, latent_hint)
                if getattr(uncondition, hint_key) is not None:
                    setattr(uncondition, hint_key, latent_hint)
            # else:
            #     assert not to_cp, "parallel_state is not initialized, context parallel should be turned off."

            setattr(condition, "base_model", self.model.base_model)
            setattr(uncondition, "base_model", self.model.base_model)

            def x0_fn(noise_x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
                if self.is_image_batch(data_batch) or not issubclass(base_class, ExtendVideoDiffusionModel):
                    cond_x0 = self.denoise(noise_x, sigma, condition).x0
                    uncond_x0 = self.denoise(noise_x, sigma, uncondition).x0
                else:
                    cond_x0 = self.denoise(
                        noise_x,
                        sigma,
                        condition,
                        condition_video_augment_sigma_in_inference=condition_video_augment_sigma_in_inference,
                        seed_inference=seed_inference,
                    ).x0_pred_replaced
                    uncond_x0 = self.denoise(
                        noise_x,
                        sigma,
                        uncondition,
                        condition_video_augment_sigma_in_inference=condition_video_augment_sigma_in_inference,
                        seed_inference=seed_inference,
                    ).x0_pred_replaced
                return cond_x0 + guidance * (cond_x0 - uncond_x0)

            return x0_fn

        def generate_samples_from_batch(
            self,
            data_batch: Dict,
            guidance: float = 1.5,
            seed: int = 1,
            state_shape: Tuple | None = None,
            n_sample: int | None = None,
            is_negative_prompt: bool = False,
            num_steps: int = 35,
            condition_latent: Union[torch.Tensor, None] = None,
            num_condition_t: Union[int, None] = None,
            condition_video_augment_sigma_in_inference: float = None,
            solver_option: COMMON_SOLVER_OPTIONS = "2ab",
            x_sigma_max: Optional[torch.Tensor] = None,
            sigma_max: float | None = None,
            return_noise: bool = False,
        ) -> Tensor:
            """
            Generate samples from the batch. Based on given batch, it will automatically determine whether to generate image or video samples.
            Different from the base model, this function support condition latent as input, it will create a differnt x0_fn if condition latent is given.
            If this feature is stablized, we could consider to move this function to the base model.

            Args:
                condition_latent (Optional[torch.Tensor]): latent tensor in shape B,C,T,H,W as condition to generate video.
                num_condition_t (Optional[int]): number of condition latent T, if None, will use the whole first half

                return_noise (bool): return the initial noise or not, used for ODE pairs generation. Not used here. Kept for conmpatibility.
            """
            self._normalize_video_databatch_inplace(data_batch)
            self._augment_image_dim_inplace(data_batch)
            is_image_batch = self.is_image_batch(data_batch)
            if is_image_batch:
                log.debug("image batch, call base model generate_samples_from_batch")
                return super().generate_samples_from_batch(
                    data_batch,
                    guidance=guidance,
                    seed=seed,
                    state_shape=state_shape,
                    n_sample=n_sample,
                    is_negative_prompt=is_negative_prompt,
                    num_steps=num_steps,
                )
            if n_sample is None:
                input_key = self.input_image_key if is_image_batch else self.input_data_key
                n_sample = data_batch[input_key].shape[0]
            if state_shape is None:
                if is_image_batch:
                    state_shape = (self.state_shape[0], 1, *self.state_shape[2:])  # C,T,H,W
                else:
                    log.debug(f"Default Video state shape is used. {self.state_shape}")
                    state_shape = self.state_shape

            # assert condition_latent is not None, "condition_latent should be provided"

            # if self.net.is_context_parallel_enabled:
            #     data_batch["latent_hint"] = split_inputs_cp(x=data_batch["latent_hint"], seq_dim=2, cp_group=self.net.cp_group)

            x0_fn = self.get_x0_fn_from_batch(
                data_batch,
                guidance,
                is_negative_prompt=is_negative_prompt,
                condition_latent=condition_latent,
                num_condition_t=num_condition_t,
                condition_video_augment_sigma_in_inference=condition_video_augment_sigma_in_inference,
                seed_inference=seed,
            )

            if sigma_max is None:
                sigma_max = self.sde.sigma_max

            if x_sigma_max is None:
                x_sigma_max = (
                    misc.arch_invariant_rand(
                        (n_sample,) + tuple(state_shape),
                        torch.float32,
                        self.tensor_kwargs["device"],
                        seed,
                    )
                    * sigma_max
                )

            if self.net.is_context_parallel_enabled:
                x_sigma_max = _broadcast(x_sigma_max, to_tp=True, to_cp=True)
                x_sigma_max = split_inputs_cp(x=x_sigma_max, seq_dim=2, cp_group=self.net.cp_group)

            samples = self.sampler(
                x0_fn, x_sigma_max, num_steps=num_steps, sigma_max=sigma_max, solver_option=solver_option
            )
            if self.net.is_context_parallel_enabled:
                samples = cat_outputs_cp(samples, seq_dim=2, cp_group=self.net.cp_group)

            return samples

        def get_patch_based_x0_fn(
            self,
            data_batch: Dict,
            guidance: float = 1.5,
            is_negative_prompt: bool = False,
            condition_latent: torch.Tensor = None,
            num_condition_t: Union[int, None] = None,
            condition_video_augment_sigma_in_inference: float = None,
            target_h: int = 2112,
            target_w: int = 3840,
            patch_h: int = 704,
            patch_w: int = 1280,
            seed_inference: int = 1,
        ) -> Callable:
            """
            Generates a callable function `x0_fn` based on the provided data batch and guidance factor.
            The function will split the input into patches, run inference on each patch, then stitch them together.

            Additional args to original function:
                target_h (int): final stitched video height
                target_w (int): final stitched video width
                patch_h (int): video patch height for each network inference
                patch_w (int): video patch width for each network inference

            Returns:
            - Callable: A function `x0_fn(noise_x, sigma)` that takes two arguments, `noise_x` and `sigma`, and return x0 prediction
            """
            assert patch_h <= target_h and patch_w <= target_w
            # data_batch should be the one processed by self.get_data_and_condition
            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)
            if hasattr(self, "is_extend_model") and self.is_extend_model:
                # Add conditions for long video generation.
                if condition_latent is None:
                    condition_latent = torch.zeros(data_batch["latent_hint"].shape, **self.tensor_kwargs)
                    num_condition_t = 0
                    condition_video_augment_sigma_in_inference = 1000

                condition.video_cond_bool = True
                condition = self.add_condition_video_indicator_and_video_input_mask(
                    condition_latent[:1], condition, num_condition_t
                )
                uncondition.video_cond_bool = True  # Not do cfg on condition frames
                uncondition = self.add_condition_video_indicator_and_video_input_mask(
                    condition_latent[:1], uncondition, num_condition_t
                )
            # Add extra conditions for ctrlnet.
            latent_hint = data_batch["latent_hint"]
            hint_key = data_batch["hint_key"]
            setattr(condition, hint_key, latent_hint)
            if "use_none_hint" in data_batch and data_batch["use_none_hint"]:
                setattr(uncondition, hint_key, None)
            else:
                setattr(uncondition, hint_key, latent_hint)

            # Handle regional prompting information
            if "regional_contexts" in data_batch:
                setattr(condition, "regional_contexts", data_batch["regional_contexts"])
                # For unconditioned generation, we still need the region masks but not the regional contexts
                setattr(uncondition, "regional_contexts", None)  # No regional contexts for unconditioned generation
            original_region_masks = None
            if "region_masks" in data_batch:
                original_region_masks = data_batch["region_masks"]
                setattr(condition, "region_masks", data_batch["region_masks"])
                # For unconditioned generation, we still need the region masks but not the regional contexts
                setattr(uncondition, "region_masks", data_batch["region_masks"])

            to_cp = self.net.is_context_parallel_enabled
            # For inference, check if parallel_state is initialized
            if parallel_state.is_initialized() and not self.is_image_batch(data_batch):
                condition = broadcast_condition(condition, to_tp=True, to_cp=to_cp)
                uncondition = broadcast_condition(uncondition, to_tp=True, to_cp=to_cp)
                cp_group = parallel_state.get_context_parallel_group()
                latent_hint = getattr(condition, hint_key)
                latent_hint = split_inputs_cp(latent_hint, seq_dim=2, cp_group=cp_group)

                if hasattr(condition, "regional_contexts") and getattr(condition, "regional_contexts") is not None:
                    regional_contexts = getattr(condition, "regional_contexts")
                    regional_contexts = split_inputs_cp(regional_contexts, seq_dim=2, cp_group=cp_group)
                    setattr(condition, "regional_contexts", regional_contexts)
                if hasattr(condition, "region_masks") and getattr(condition, "region_masks") is not None:
                    region_masks = getattr(condition, "region_masks")
                    region_masks = split_inputs_cp(region_masks, seq_dim=2, cp_group=cp_group)
                    setattr(condition, "region_masks", region_masks)

            setattr(condition, "base_model", self.model.base_model)
            setattr(uncondition, "base_model", self.model.base_model)
            if hasattr(self, "hint_encoders"):
                self.model.net.hint_encoders = self.hint_encoders

            def x0_fn(noise_x: torch.Tensor, sigma: torch.Tensor):
                w, h = target_w, target_h
                n_img_w = (w - 1) // patch_w + 1
                n_img_h = (h - 1) // patch_h + 1

                overlap_size_w = overlap_size_h = 0
                if n_img_w > 1:
                    overlap_size_w = (n_img_w * patch_w - w) // (n_img_w - 1)
                    assert n_img_w * patch_w - overlap_size_w * (n_img_w - 1) == w
                if n_img_h > 1:
                    overlap_size_h = (n_img_h * patch_h - h) // (n_img_h - 1)
                    assert n_img_h * patch_h - overlap_size_h * (n_img_h - 1) == h

                batch_images = noise_x
                batch_sigma = sigma
                output = []
                for idx, cur_images in enumerate(batch_images):
                    noise_x = cur_images.unsqueeze(0)
                    sigma = batch_sigma[idx : idx + 1]
                    condition.gt_latent = condition_latent[idx : idx + 1]
                    uncondition.gt_latent = condition_latent[idx : idx + 1]
                    setattr(condition, hint_key, latent_hint[idx : idx + 1])
                    if getattr(uncondition, hint_key) is not None:
                        setattr(uncondition, hint_key, latent_hint[idx : idx + 1])

                    if self.is_image_batch(data_batch) or not issubclass(base_class, ExtendVideoDiffusionModel):
                        cond_x0 = self.denoise(noise_x, sigma, condition).x0
                        uncond_x0 = self.denoise(noise_x, sigma, uncondition).x0
                    else:
                        cond_x0 = self.denoise(
                            noise_x,
                            sigma,
                            condition,
                            condition_video_augment_sigma_in_inference=condition_video_augment_sigma_in_inference,
                            seed_inference=seed_inference,
                        ).x0_pred_replaced
                        uncond_x0 = self.denoise(
                            noise_x,
                            sigma,
                            uncondition,
                            condition_video_augment_sigma_in_inference=condition_video_augment_sigma_in_inference,
                            seed_inference=seed_inference,
                        ).x0_pred_replaced
                    x0 = cond_x0 + guidance * (cond_x0 - uncond_x0)
                    output.append(x0)
                output = rearrange(torch.stack(output), "(n t) b ... -> (b n t) ...", n=n_img_h, t=n_img_w)  # 8x3xhxw
                final_output = merge_patches_into_video(output, overlap_size_h, overlap_size_w, n_img_h, n_img_w)
                final_output = split_video_into_patches(final_output, patch_h, patch_w)
                return final_output

            return x0_fn

        def generate_samples_from_patches(
            self,
            data_batch: Dict,
            guidance: float = 1.5,
            seed: int = 1,
            state_shape: Tuple | None = None,
            n_sample: int | None = None,
            is_negative_prompt: bool = False,
            num_steps: int = 35,
            condition_latent: Union[torch.Tensor, None] = None,
            num_condition_t: Union[int, None] = None,
            condition_video_augment_sigma_in_inference: float = None,
            solver_option: COMMON_SOLVER_OPTIONS = "2ab",
            x_sigma_max: Optional[torch.Tensor] = None,
            sigma_max: float | None = None,
            target_h: int = 2112,
            target_w: int = 3840,
            patch_h: int = 704,
            patch_w: int = 1280,
        ) -> Tensor:
            """
            Generate samples from the batch using patch-based inference. During each denoising step, it will denoise each patch
            separately then average the overlapping regions.

            Additional args to original function:
                target_h (int): final stitched video height
                target_w (int): final stitched video width
                patch_h (int): video patch height for each network inference
                patch_w (int): video patch width for each network inference
            """
            assert patch_h <= target_h and patch_w <= target_w
            self._normalize_video_databatch_inplace(data_batch)
            self._augment_image_dim_inplace(data_batch)
            is_image_batch = self.is_image_batch(data_batch)
            if is_image_batch:
                log.debug("image batch, call base model generate_samples_from_batch")
                return super().generate_samples_from_batch(
                    data_batch,
                    guidance=guidance,
                    seed=seed,
                    state_shape=state_shape,
                    n_sample=n_sample,
                    is_negative_prompt=is_negative_prompt,
                    num_steps=num_steps,
                )
            if n_sample is None:
                input_key = self.input_image_key if is_image_batch else self.input_data_key
                n_sample = data_batch[input_key].shape[0]
            if state_shape is None:
                if is_image_batch:
                    state_shape = (self.state_shape[0], 1, *self.state_shape[2:])  # C,T,H,W
                else:
                    log.debug(f"Default Video state shape is used. {self.state_shape}")
                    state_shape = self.state_shape

            x0_fn = self.get_patch_based_x0_fn(
                data_batch,
                guidance,
                is_negative_prompt=is_negative_prompt,
                condition_latent=condition_latent,
                num_condition_t=num_condition_t,
                condition_video_augment_sigma_in_inference=condition_video_augment_sigma_in_inference,
                target_h=target_h,
                target_w=target_w,
                patch_h=patch_h,
                patch_w=patch_w,
                seed_inference=seed,
            )

            if sigma_max is None:
                sigma_max = self.sde.sigma_max

            if x_sigma_max is None:
                x_sigma_max = (
                    misc.arch_invariant_rand(
                        (n_sample,) + tuple(state_shape),
                        torch.float32,
                        self.tensor_kwargs["device"],
                        seed,
                    )
                    * sigma_max
                )

            if self.net.is_context_parallel_enabled:
                x_sigma_max = _broadcast(x_sigma_max, to_tp=True, to_cp=True)
                x_sigma_max = split_inputs_cp(x=x_sigma_max, seq_dim=2, cp_group=self.net.cp_group)

            samples = self.sampler(
                x0_fn, x_sigma_max, num_steps=num_steps, sigma_max=sigma_max, solver_option=solver_option
            )
            if self.net.is_context_parallel_enabled:
                samples = cat_outputs_cp(samples, seq_dim=2, cp_group=self.net.cp_group)

            return samples

        @torch.no_grad()
        def validation_step(
            self, data: dict[str, torch.Tensor], iteration: int
        ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
            """
            save generated videos
            """
            raw_data, x0, condition = self.get_data_and_condition(data)
            guidance = data["guidance"]
            sigma_max = data["sigma_max"]
            is_negative_prompt = data["is_negative_prompt"]
            data = misc.to(data, **self.tensor_kwargs)
            x_sigma_max = None
            if sigma_max is not None:
                x_sigma_max = self.get_x_from_clean(x0, sigma_max)
            sample = self.generate_samples_from_batch(
                data,
                guidance=guidance,
                # make sure no mismatch and also works for cp
                state_shape=x0.shape[1:],
                n_sample=x0.shape[0],
                x_sigma_max=x_sigma_max,
                sigma_max=sigma_max,
                is_negative_prompt=is_negative_prompt,
            )
            sample = self.decode(sample)
            gt = raw_data
            hint = data[data["hint_key"]][:, :3]
            result = torch.cat([hint, sample], dim=3)
            gt = torch.cat([hint, gt], dim=3)
            caption = data["ai_caption"]
            return {"gt": gt, "result": result, "caption": caption}, torch.tensor([0]).to(**self.tensor_kwargs)

    return VideoDiffusionModelWithCtrlWrapper


@video_ctrlnet_decorator
@ctrlnet_decorator
class VideoDiffusionModelWithCtrl(ExtendVideoDiffusionModel):
    pass


@diffusion_fsdp_class_decorator
@video_ctrlnet_decorator
@ctrlnet_decorator
class VideoDiffusionFSDPModelWithCtrl(ExtendVideoDiffusionModel):
    pass


@video_ctrlnet_decorator
@ctrlnet_decorator
class ShortVideoDiffusionModelWithCtrl(VideoDiffusionModel):
    pass


@diffusion_fsdp_class_decorator
@video_ctrlnet_decorator
@ctrlnet_decorator
class ShortVideoDiffusionFSDPModelWithCtrl(VideoDiffusionModel):
    pass