cosmos_predict1/diffusion/training/models/model_image.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.

import functools
from contextlib import contextmanager
from dataclasses import dataclass, fields
from typing import Any, Callable, Dict, Mapping, Optional, Tuple, Type, TypeVar

import numpy as np
import torch
import torch.nn.functional as F
from megatron.core import parallel_state
from torch.distributed.fsdp import FullStateDictConfig
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy, StateDictType
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
from torch.nn.modules.module import _IncompatibleKeys

from cosmos_predict1.diffusion.functional.batch_ops import batch_mul
from cosmos_predict1.diffusion.module.blocks import FourierFeatures
from cosmos_predict1.diffusion.module.pretrained_vae import BaseVAE
from cosmos_predict1.diffusion.modules.denoiser_scaling import EDMScaling
from cosmos_predict1.diffusion.modules.res_sampler import COMMON_SOLVER_OPTIONS, Sampler
from cosmos_predict1.diffusion.training.functional.loss import create_per_sample_loss_mask
from cosmos_predict1.diffusion.training.utils.fsdp_helper import apply_fsdp_checkpointing, hsdp_device_mesh
from cosmos_predict1.diffusion.training.utils.optim_instantiate import get_base_scheduler
from cosmos_predict1.diffusion.types import DenoisePrediction
from cosmos_predict1.utils import distributed, log, misc
from cosmos_predict1.utils.ema import FastEmaModelUpdater
from cosmos_predict1.utils.lazy_config import LazyDict
from cosmos_predict1.utils.lazy_config import instantiate as lazy_instantiate
from cosmos_predict1.utils.model import Model


@dataclass
class CosmosCondition:
    crossattn_emb: torch.Tensor
    crossattn_mask: torch.Tensor
    padding_mask: Optional[torch.Tensor] = None
    scalar_feature: Optional[torch.Tensor] = None

    def to_dict(self) -> Dict[str, Optional[torch.Tensor]]:
        return {f.name: getattr(self, f.name) for f in fields(self)}


class DiffusionModel(Model):
    def __init__(self, config):
        super().__init__()

        self.config = config

        # how many sample have been processed
        self.sample_counter = 0
        self.precision = {
            "float32": torch.float32,
            "float16": torch.float16,
            "bfloat16": torch.bfloat16,
        }[config.precision]
        self.tensor_kwargs = {"device": "cuda", "dtype": self.precision}
        log.warning(f"DiffusionModel: precision {self.precision}")
        # Timer passed to network to detect slow ranks.
        # 1. set data keys and data information
        self.sigma_data = config.sigma_data
        self.state_shape = list(config.latent_shape)
        self.setup_data_key()

        # 2. setup up diffusion processing and scaling~(pre-condition), sampler
        self.sde = lazy_instantiate(config.sde)
        self.sampler = Sampler()
        self.scaling = EDMScaling(self.sigma_data)

        # 3. vae
        with misc.timer("DiffusionModel: set_up_vae"):
            self.vae: BaseVAE = lazy_instantiate(config.vae)
            assert (
                self.vae.latent_ch == self.state_shape[0]
            ), f"latent_ch {self.vae.latent_ch} != state_shape {self.state_shape[0]}"

        # 4. Set up loss options, including loss masking, loss reduce and loss scaling
        self.loss_masking: Optional[Dict] = config.loss_masking
        self.loss_reduce = getattr(config, "loss_reduce", "mean")
        assert self.loss_reduce in ["mean", "sum"]
        self.loss_scale = getattr(config, "loss_scale", 1.0)
        log.critical(f"Using {self.loss_reduce} loss reduce with loss scale {self.loss_scale}")
        log.critical(f"Enable loss masking: {config.loss_mask_enabled}")

        # 5. diffusion neural networks part
        self.set_up_model()

    def setup_data_key(self) -> None:
        self.input_data_key = self.config.input_data_key

    def build_model(self) -> torch.nn.ModuleDict:
        config = self.config
        net = lazy_instantiate(config.net)
        conditioner = lazy_instantiate(config.conditioner)
        logvar = torch.nn.Sequential(
            FourierFeatures(num_channels=128, normalize=True), torch.nn.Linear(128, 1, bias=False)
        )

        return torch.nn.ModuleDict(
            {
                "net": net,
                "conditioner": conditioner,
                "logvar": logvar,
            }
        )

    @misc.timer("DiffusionModel: set_up_model")
    def set_up_model(self):
        config = self.config
        self.model = self.build_model()
        if config.ema.enabled:
            with misc.timer("DiffusionModel: instantiate ema"):
                config.ema.model = self.model
                self.model_ema = lazy_instantiate(config.ema)
                config.ema.model = None
        else:
            self.model_ema = None

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

    @property
    def conditioner(self):
        return self.model.conditioner

    def on_before_zero_grad(
        self, optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler.LRScheduler, iteration: int
    ) -> None:
        """
        update the model_ema
        """
        if self.config.ema.enabled:
            self.model_ema.update_average(self.model, iteration)

    def on_train_start(self, memory_format: torch.memory_format = torch.preserve_format) -> None:
        if self.config.ema.enabled:
            self.model_ema.to(dtype=torch.float32)
        if hasattr(self.vae, "reset_dtype"):
            self.vae.reset_dtype()
        self.model = self.model.to(memory_format=memory_format, **self.tensor_kwargs)

        if hasattr(self.config, "use_torch_compile") and self.config.use_torch_compile:  # compatible with old config
            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."
                )
            # Increasing cache size. It's required because of the model size and dynamic input shapes resulting in
            # multiple different triton kernels. For 28 TransformerBlocks, the cache limit of 256 should be enough for
            # up to 9 different input shapes, as 28*9 < 256. If you have more Blocks or input shapes, and you observe
            # graph breaks at each Block (detectable with torch._dynamo.explain) or warnings about
            # exceeding cache limit, you may want to increase this size.
            # Starting with 24.05 Pytorch container, the default value is 256 anyway.
            # You can read more about it in the comments in Pytorch source code under path torch/_dynamo/cache_size.py.
            torch._dynamo.config.accumulated_cache_size_limit = 256
            # dynamic=False means that a separate kernel is created for each shape. It incurs higher compilation costs
            # at initial iterations, but can result in more specialized and efficient kernels.
            # dynamic=True currently throws errors in pytorch 2.3.
            self.model.net = torch.compile(self.model.net, dynamic=False, disable=not self.config.use_torch_compile)

    def compute_loss_with_epsilon_and_sigma(
        self,
        data_batch: dict[str, torch.Tensor],
        x0_from_data_batch: torch.Tensor,
        x0: torch.Tensor,
        condition: CosmosCondition,
        epsilon: torch.Tensor,
        sigma: torch.Tensor,
    ):
        """
        Compute loss givee epsilon and sigma

        This method is responsible for computing loss give epsilon and sigma. It involves:
        1. Adding noise to the input data using the SDE process.
        2. Passing the noisy data through the network to generate predictions.
        3. Computing the loss based on the difference between the predictions and the original data, \
            considering any configured loss weighting.

        Args:
            data_batch (dict): raw data batch draw from the training data loader.
            x0_from_data_batch: raw image/video
            x0: image/video latent
            condition: text condition
            epsilon: noise
            sigma: noise level

        Returns:
            tuple: A tuple containing four elements:
                - dict: additional data that used to debug / logging / callbacks
                - Tensor 1: kendall loss,
                - Tensor 2: MSE loss,
                - Tensor 3: EDM loss

        Raises:
            AssertionError: If the class is conditional, \
                but no number of classes is specified in the network configuration.

        Notes:
            - The method handles different types of conditioning
            - The method also supports Kendall's loss
        """
        # Get the mean and stand deviation of the marginal probability distribution.
        mean, std = self.sde.marginal_prob(x0, sigma)
        # Generate noisy observations
        xt = mean + batch_mul(std, epsilon)  # corrupted data
        # make prediction
        model_pred = self.denoise(xt, sigma, condition)
        # loss weights for different noise levels
        weights_per_sigma = self.get_per_sigma_loss_weights(sigma=sigma)
        # extra weight for each sample, for example, aesthetic weight, camera weight
        weights_per_sample = self.get_per_sample_weight(data_batch, x0_from_data_batch.shape[0])
        # extra loss mask for each sample, for example, human faces, hands
        loss_mask_per_sample = self.get_per_sample_loss_mask(data_batch, x0_from_data_batch.shape, x0.shape)
        pred_mse = (x0 - model_pred.x0) ** 2 * loss_mask_per_sample
        edm_loss = batch_mul(pred_mse, weights_per_sigma * weights_per_sample)
        if self.config.loss_add_logvar:
            kendall_loss = batch_mul(edm_loss, torch.exp(-model_pred.logvar).view(-1)).flatten(
                start_dim=1
            ) + model_pred.logvar.view(-1, 1)
        else:
            kendall_loss = edm_loss.flatten(start_dim=1)
        output_batch = {
            "x0": x0,
            "xt": xt,
            "sigma": sigma,
            "weights_per_sigma": weights_per_sigma,
            "weights_per_sample": weights_per_sample,
            "loss_mask_per_sample": loss_mask_per_sample,
            "condition": condition,
            "model_pred": model_pred,
            "mse_loss": pred_mse.mean(),
            "edm_loss": edm_loss.mean(),
        }
        return output_batch, kendall_loss, pred_mse, edm_loss

    def training_step(
        self, data_batch: dict[str, torch.Tensor], iteration: int
    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
        """
        Performs a single training step for the diffusion model.

        This method is responsible for executing one iteration of the model's training. It involves:
        1. Adding noise to the input data using the SDE process.
        2. Passing the noisy data through the network to generate predictions.
        3. Computing the loss based on the difference between the predictions and the original data, \
            considering any configured loss weighting.

        Args:
            data_batch (dict): raw data batch draw from the training data loader.
            iteration (int): Current iteration number.

        Returns:
            tuple: A tuple containing two elements:
                - dict: additional data that used to debug / logging / callbacks
                - Tensor: The computed loss for the training step as a PyTorch Tensor.

        Raises:
            AssertionError: If the class is conditional, \
                but no number of classes is specified in the network configuration.

        Notes:
            - The method handles different types of conditioning
            - The method also supports Kendall's loss
        """
        # Get the input data to noise and denoise~(image, video) and the corresponding conditioner.
        x0_from_data_batch, x0, condition = self.get_data_and_condition(data_batch)

        # Sample pertubation noise levels and N(0, 1) noises
        sigma, epsilon = self.draw_training_sigma_and_epsilon(x0.size(), condition)

        output_batch, kendall_loss, pred_mse, edm_loss = self.compute_loss_with_epsilon_and_sigma(
            data_batch, x0_from_data_batch, x0, condition, epsilon, sigma
        )

        if self.loss_reduce == "mean":
            kendall_loss = kendall_loss.mean() * self.loss_scale
        elif self.loss_reduce == "sum":
            kendall_loss = kendall_loss.sum(dim=1).mean() * self.loss_scale
        else:
            raise ValueError(f"Invalid loss_reduce: {self.loss_reduce}")

        return output_batch, kendall_loss

    def denoise(self, xt: torch.Tensor, sigma: torch.Tensor, condition: CosmosCondition) -> DenoisePrediction:
        """
        Performs denoising on the input noise data, noise level, and condition

        Args:
            xt (torch.Tensor): The input noise data.
            sigma (torch.Tensor): The noise level.
            condition (CosmosCondition): conditional information, generated from self.conditioner

        Returns:
            DenoisePrediction: The denoised prediction, it includes clean data predicton (x0), \
                noise prediction (eps_pred) and optional confidence (logvar).
        """

        if getattr(self.config, "use_dummy_temporal_dim", False):
            # When using video DiT model for image, we need to use a dummy temporal dimension.
            xt = xt.unsqueeze(2)

        xt = xt.to(**self.tensor_kwargs)
        sigma = sigma.to(**self.tensor_kwargs)
        # get precondition for the network
        c_skip, c_out, c_in, c_noise = self.scaling(sigma=sigma)

        # forward pass through the network
        net_output = self.net(
            x=batch_mul(c_in, xt),  # Eq. 7 of https://arxiv.org/pdf/2206.00364.pdf
            timesteps=c_noise,  # Eq. 7 of https://arxiv.org/pdf/2206.00364.pdf
            **condition.to_dict(),
        )

        logvar = self.model.logvar(c_noise)
        x0_pred = batch_mul(c_skip, xt) + batch_mul(c_out, net_output)

        # get noise prediction based on sde
        eps_pred = batch_mul(xt - x0_pred, 1.0 / sigma)

        if getattr(self.config, "use_dummy_temporal_dim", False):
            x0_pred = x0_pred.squeeze(2)
            eps_pred = eps_pred.squeeze(2)

        return DenoisePrediction(x0_pred, eps_pred, logvar)

    @torch.no_grad()
    def encode(self, state: torch.Tensor) -> torch.Tensor:
        return self.vae.encode(state) * self.sigma_data

    @torch.no_grad()
    def decode(self, latent: torch.Tensor) -> torch.Tensor:
        return self.vae.decode(latent / self.sigma_data)

    def draw_training_sigma_and_epsilon(self, x0_size: int, condition: Any) -> torch.Tensor:
        del condition
        batch_size = x0_size[0]
        epsilon = torch.randn(x0_size, **self.tensor_kwargs)
        return self.sde.sample_t(batch_size).to(**self.tensor_kwargs), epsilon

    def get_data_and_condition(self, data_batch: dict[str, torch.Tensor]) -> Tuple[torch.Tensor, CosmosCondition]:
        """
        processing data batch draw from data loader and return data and condition that used for denoising task

        Returns:
            raw_state (tensor): the image / video data that feed to vae
            latent_state (tensor): nosie-free state, the vae latent state
            condition (CosmosCondition): condition information for conditional generation. Generated from conditioner
        """
        raw_state = data_batch[self.input_data_key]
        latent_state = self.encode(raw_state)
        condition = self.conditioner(data_batch)
        return raw_state, latent_state, condition

    def get_per_sample_weight(self, data_batch: dict[str, torch.Tensor], batch_size: int):
        r"""
        extra weight for each sample, for example, aesthetic weight
        Args:
            data_batch: raw data batch draw from the training data loader.
            batch_size: int, the batch size of the input data
        """
        aesthetic_cfg = getattr(self.config, "aesthetic_finetuning", None)
        if (aesthetic_cfg is not None) and getattr(aesthetic_cfg, "enabled", False):
            sample_weight = data_batch["aesthetic_weight"]
        else:
            sample_weight = torch.ones(batch_size, **self.tensor_kwargs)

        camera_cfg = getattr(self.config, "camera_sample_weight", None)
        if (camera_cfg is not None) and getattr(camera_cfg, "enabled", False):
            sample_weight *= 1 + (data_batch["camera_attributes"][:, 1:].sum(dim=1) != 0) * (camera_cfg.weight - 1)
        return sample_weight

    def get_per_sample_loss_mask(self, data_batch, raw_x_shape, latent_x_shape):
        """
        extra loss mask for each sample, for example, human faces, hands.

        Args:
            data_batch (dict): raw data batch draw from the training data loader.
            raw_x_shape (tuple): shape of the input data. We need the raw_x_shape for necessary resize operation.
            latent_x_shape (tuple): shape of the latent data
        """
        if self.config.loss_mask_enabled:
            raw_x_shape = [raw_x_shape[0], 1, *raw_x_shape[2:]]
            weights = create_per_sample_loss_mask(
                self.loss_masking, data_batch, raw_x_shape, torch.get_default_dtype(), "cuda"
            )
            return F.interpolate(weights, size=latent_x_shape[2:], mode="bilinear")

        return 1.0

    def get_per_sigma_loss_weights(self, sigma: torch.Tensor):
        """
        Args:
            sigma (tensor): noise level

        Returns:
            loss weights per sigma noise level
        """
        return (sigma**2 + self.sigma_data**2) / (sigma * self.sigma_data) ** 2

    def generate_samples(self, batch_size: int, condition: CosmosCondition) -> torch.Tensor:
        """
        Generate samples with given condition. It is WITHOUT classifier-free-guidance.

        Args:
            batch_size (int):
            condition (CosmosCondition): condition information generated from self.conditioner
        """
        x_sigma_max = torch.randn(batch_size, *self.state_shape, **self.tensor_kwargs) * self.sde.sigma_max

        def x0_fn(x, t):
            return self.denoise(x, t, condition).x0  # ODE function

        return self.sampler(x0_fn, x_sigma_max, sigma_max=self.sde.sigma_max)

    def generate_cfg_samples(
        self, batch_size: int, condition: CosmosCondition, uncondition: CosmosCondition, guidance=1.5
    ) -> torch.Tensor:
        """
        Generate samples with with classifier-free-guidance.

        Args:
            batch_size (int):
            condition (CosmosCondition): condition information generated from self.conditioner
            uncondition (CosmosCondition): uncondition information, possibily generated from self.conditioner
        """
        x_sigma_max = torch.randn(batch_size, *self.state_shape, **self.tensor_kwargs) * self.sde.sigma_max

        def x0_fn(x, t):
            cond_x0 = self.denoise(x, t, condition).x0
            uncond_x0 = self.denoise(x, t, uncondition).x0
            return cond_x0 + guidance * (cond_x0 - uncond_x0)

        return self.sampler(x0_fn, x_sigma_max, sigma_max=self.sde.sigma_max)

    def get_x0_fn_from_batch(
        self,
        data_batch: Dict,
        guidance: float = 1.5,
        is_negative_prompt: bool = False,
    ) -> 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

        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.
        """
        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)

        def x0_fn(noise_x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
            cond_x0 = self.denoise(noise_x, sigma, condition).x0
            uncond_x0 = self.denoise(noise_x, sigma, uncondition).x0
            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: Optional[Tuple] = None,
        n_sample: Optional[int] = None,
        is_negative_prompt: bool = False,
        num_steps: int = 35,
        solver_option: COMMON_SOLVER_OPTIONS = "2ab",
    ) -> torch.Tensor:
        """
        Args:
            data_batch (dict): raw data batch draw from the training data loader.
            iteration (int): Current iteration number.
            guidance (float): guidance weights
            seed (int): random seed
            state_shape (tuple): shape of the state, default to self.state_shape if not provided
            n_sample (int): number of samples to generate
            is_negative_prompt (bool): use negative prompt t5 in uncondition if true
            num_steps (int): number of steps for the diffusion process
            solver_option (str): differential equation solver option, default to "2ab"~(mulitstep solver)
        """
        x0_fn = self.get_x0_fn_from_batch(data_batch, guidance, is_negative_prompt=is_negative_prompt)
        batch_size = n_sample or data_batch[self.input_data_key].shape[0]
        state_shape = state_shape or self.state_shape
        x_sigma_max = (
            misc.arch_invariant_rand(
                (batch_size,) + tuple(state_shape),
                torch.float32,
                self.tensor_kwargs["device"],
                seed,
            )
            * self.sde.sigma_max
        )
        return self.sampler(
            x0_fn, x_sigma_max, sigma_max=self.sde.sigma_max, num_steps=num_steps, solver_option=solver_option
        )

    @torch.no_grad()
    def validation_step(
        self, data: dict[str, torch.Tensor], iteration: int
    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
        """
        Current code does nothing.
        """
        return {}, torch.tensor(0).to(**self.tensor_kwargs)

    @torch.no_grad()
    def forward(self, xt, t, condition: CosmosCondition):
        """
        Performs denoising on the input noise data, noise level, and condition

        Args:
            xt (torch.Tensor): The input noise data.
            sigma (torch.Tensor): The noise level.
            condition (CosmosCondition): conditional information, generated from self.conditioner

        Returns:
            DenoisePrediction: The denoised prediction, it includes clean data predicton (x0), \
                noise prediction (eps_pred) and optional confidence (logvar).
        """
        return self.denoise(xt, t, condition)

    def init_optimizer_scheduler(
        self, optimizer_config: LazyDict, scheduler_config: LazyDict
    ) -> tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LRScheduler]:
        """Creates the optimizer and scheduler for the model.

        Args:
            config_model (ModelConfig): The config object for the model.

        Returns:
            optimizer (torch.optim.Optimizer): The model optimizer.
            scheduler (torch.optim.lr_scheduler.LRScheduler): The optimization scheduler.
        """
        optimizer = lazy_instantiate(optimizer_config, model=self.model)
        scheduler = get_base_scheduler(optimizer, self, scheduler_config)
        return optimizer, scheduler

    def state_dict(self) -> Dict[str, Any]:
        """
        Returns the current state of the model as a dictionary.

        Returns:
            Dict: The current state of the model as a dictionary.
        """
        return {
            "model": self.model.state_dict(),
            "ema": self.model_ema.state_dict() if self.config.ema.enabled else None,
        }

    def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True, assign: bool = False):
        """
        Loads a state dictionary into the model and optionally its EMA counterpart.
        Different from torch strict=False mode, the method will not raise error for unmatched state shape while raise warning.

        Parameters:
            state_dict (Mapping[str, Any]): A dictionary containing separate state dictionaries for the model and
                                            potentially for an EMA version of the model under the keys 'model' and 'ema', respectively.
            strict (bool, optional): If True, the method will enforce that the keys in the state dict match exactly
                                     those in the model and EMA model (if applicable). Defaults to True.
            assign (bool, optional): If True and in strict mode, will assign the state dictionary directly rather than
                                     matching keys one-by-one. This is typically used when loading parts of state dicts
                                     or using customized loading procedures. Defaults to False.
        """
        if strict:
            # the converted tpsp checkpoint has "ema" and it is None
            if self.config.ema.enabled and state_dict["ema"] is not None:
                ema_results: _IncompatibleKeys = self.model_ema.load_state_dict(
                    state_dict["ema"], strict=strict, assign=assign
                )
            reg_results: _IncompatibleKeys = self.model.load_state_dict(
                state_dict["model"], strict=strict, assign=assign
            )
            if self.config.ema.enabled and state_dict["ema"] is not None:
                return _IncompatibleKeys(
                    ema_results.missing_keys + reg_results.missing_keys,
                    ema_results.unexpected_keys + reg_results.unexpected_keys,
                )
            return reg_results
        else:
            from cosmos_predict1.diffusion.training.utils.checkpointer import non_strict_load_model

            log.critical("load model in non-strict mode")
            if "model" in state_dict:
                log.critical(non_strict_load_model(self.model, state_dict["model"]), rank0_only=False)
            else:
                log.critical(non_strict_load_model(self.model, state_dict), rank0_only=False)
            if self.config.ema.enabled and "ema" in state_dict and state_dict["ema"] is not None:
                log.critical("load ema model in non-strict mode")
                log.critical(non_strict_load_model(self.model_ema, state_dict["ema"]), rank0_only=False)

    def get_ckpt_postfix(self) -> Tuple[str, int, int]:
        """Get the checkpoint file postfix.

        Args:
            iteration (int): The current iteration number.

        Returns:
            postfix (str): The postfix of the checkpoint file.
            rank_to_save ema (int), we will not save each ema model in each rank, \
                ema model with same rate will be saved once
            total_ema_num (int)
        """
        total_ema_num = min(self.config.ema.num, distributed.get_world_size())
        rank = distributed.get_rank()
        if rank == 0:
            return "", 0, total_ema_num
        if self.config.ema.enabled:
            if rank < self.config.ema.num:
                return f"_RANK{rank}", rank, total_ema_num
        return "", 0, total_ema_num  # use rank 0 to save the checkpoint

    @contextmanager
    def ema_scope(self, context=None, is_cpu=False):
        if self.config.ema.enabled:
            self.model_ema.cache(self.model.parameters(), is_cpu=is_cpu)
            self.model_ema.copy_to(self.model)
            if context is not None:
                log.info(f"{context}: Switched to EMA weights")
        try:
            yield None
        finally:
            if self.config.ema.enabled:
                self.model_ema.restore(self.model.parameters())
                if context is not None:
                    log.info(f"{context}: Restored training weights")


T = TypeVar("T", bound=DiffusionModel)


def diffusion_fsdp_class_decorator(base_class: Type[T]) -> Type[T]:
    """
    Decorator for the FSDP class for the diffusion model, which handles the FSDP specific logic for the diffusion model.
    """

    class FSDPClass(base_class):
        """
        Handle FSDP specific logic for the diffusion model. Including:
        - FSDP model initialization
        - FSDP model / optimizer save and loading
        - Different from the original DiffusionModel, the impl of multi-rank EMA is a bit hacky. \
            We need to make sure sharded model weights for EMA and regular model are the same.
        """

        def __init__(self, config, fsdp_checkpointer: Any):
            self.fsdp_checkpointer = fsdp_checkpointer
            super().__init__(config)

        def set_up_model(self):
            config = self.config

            # 1. build FSDP sharding strategy and device_mesh
            strategy = {
                "full": ShardingStrategy.FULL_SHARD,
                "hybrid": ShardingStrategy.HYBRID_SHARD,
            }[config.fsdp.sharding_strategy]
            log.critical(f"Using {strategy} sharding strategy for FSDP")

            if config.fsdp.sharding_strategy == "hybrid":
                sharding_group_size = getattr(config.fsdp, "sharding_group_size", 8)
                device_mesh = hsdp_device_mesh(
                    sharding_group_size=sharding_group_size,
                )
                shard_group = device_mesh.get_group(mesh_dim="shard")
                replicate_group = device_mesh.get_group(mesh_dim="replicate")
                fsdp_process_group = (shard_group, replicate_group)
            else:
                device_mesh = hsdp_device_mesh(
                    sharding_group_size=distributed.get_world_size(),
                )
                shard_group = device_mesh.get_group(mesh_dim="shard")
                fsdp_process_group = shard_group

            # We piggyback the `device_mesh` to megatron-core's `parallel_state` for global access.
            # This is not megatron-core's original API.
            parallel_state.fsdp_device_mesh = device_mesh

            def get_wrap_policy(_model):
                if not hasattr(_model.net, "fsdp_wrap_block_cls"):
                    raise ValueError(
                        "Networks does not have fsdp_wrap_block_cls attribute, please check the net definition"
                    )
                fsdp_blocks_cls = _model.net.fsdp_wrap_block_cls
                fsdp_blocks_cls = (
                    list(fsdp_blocks_cls) if isinstance(fsdp_blocks_cls, (list, tuple, set)) else [fsdp_blocks_cls]
                )
                log.critical(f"Using FSDP blocks {fsdp_blocks_cls}")

                log.critical(f"Using wrap policy {config.fsdp.policy}")
                if config.fsdp.policy == "size":
                    min_num_params = getattr(config.fsdp, "min_num_params", 100)
                    log.critical(f"Using {min_num_params} as the minimum number of parameters for auto-wrap policy")
                    wrap_policy = functools.partial(size_based_auto_wrap_policy, min_num_params=min_num_params)
                else:
                    from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy

                    wrap_policy = functools.partial(
                        transformer_auto_wrap_policy,
                        transformer_layer_cls=set(fsdp_blocks_cls),
                    )
                return wrap_policy

            # 2. build naive pytorch model and load weights if exists
            replica_idx, shard_idx = device_mesh.get_coordinate()
            # 2.1 handle ema case first, since float32 is more expensive
            if config.ema.enabled:
                with misc.timer("Creating PyTorch model and loading weights for ema"):
                    model_ema = self.build_model().float()
                    model_ema.cuda().eval().requires_grad_(False)
                    if distributed.get_rank() == 0:
                        # only load model in rank0 to reduce network traffic
                        self.fsdp_checkpointer.load_model_during_init(model_ema, is_ema=True)
                    # sync ema model weights from rank0
                    with misc.timer("Sync model states for EMA model"):
                        #! this is IMPORTANT, see the following comment about regular model for details
                        #! we broadcast the ema model first, since it is fp32 and costs more memory
                        distributed.sync_model_states(model_ema, device_mesh.get_group(mesh_dim="shard"))
                        torch.cuda.empty_cache()
                        distributed.sync_model_states(model_ema, device_mesh.get_group(mesh_dim="replicate"))
                        torch.cuda.empty_cache()
                    # for ema model with dfiferent rate, we download the model when necessary
                    if shard_idx == 0 and replica_idx > 0 and replica_idx < config.ema.num:
                        print("loading ema model in rank", replica_idx)
                        self.fsdp_checkpointer.load_model_during_init(
                            model_ema,
                            is_ema=True,
                            ema_id=replica_idx,
                        )
                        print("finish loading ema model in rank", replica_idx)
                # 2.1.2 create FSDP model for ema model
                with misc.timer("Creating FSDP model for EMA model"):
                    self.model_ema = FSDP(
                        model_ema,
                        sync_module_states=True,  # it can reduce network traffic by only loading model in rank0 and sync
                        process_group=device_mesh.get_group(mesh_dim=1),
                        sharding_strategy=ShardingStrategy.FULL_SHARD,
                        auto_wrap_policy=get_wrap_policy(model_ema),
                        device_id=torch.cuda.current_device(),
                        limit_all_gathers=True,
                    )

                # extra ema model upate logic to the model
                self.model_ema_worker = FastEmaModelUpdater()
                s = 0.1
                replica_idx, shard_idx = device_mesh.get_coordinate()
                divider = 2**replica_idx if replica_idx < config.ema.num else 1
                if replica_idx < config.ema.num:
                    if shard_idx == 0:
                        print(f"EMA: rank {replica_idx}, rate {config.ema.rate / divider}")
                s = config.ema.rate / divider
                self.ema_exp_coefficient = np.roots([1, 7, 16 - s**-2, 12 - s**-2]).real.max()

                torch.cuda.empty_cache()

            # 2.2 handle regular model
            with misc.timer("Creating PyTorch model and loading weights for regular model"):
                model = self.build_model().cuda().to(**self.tensor_kwargs)

                if distributed.get_rank() == 0:
                    # only load model in rank0 to reduce network traffic and sync later
                    self.fsdp_checkpointer.load_model_during_init(model, is_ema=False)

                #! overwrite the forward method so that it will invoke the FSDP-specific pre- and post-forward sharding logic
                model.forward = super().training_step
                #! this is IMPORTANT, though following two lines are identical to sync_module_states=True in FSDP
                #! we do it twice so that following line can warm up and avoid OOM in aws 128+ nodes settings
                #! qsh hypothesize that it is due to overhead of initialization of nccl network communication;
                #! without it, peak mem : reg_model + ema_model + FSDP overhead + nccl communication initialization overhead
                #! with it, peak men: reg_model + ema_model + FSDP overhead
                #! it is tricky, but it works!
                with misc.timer("Sync model states for regular model"):
                    distributed.sync_model_states(model, device_mesh.get_group(mesh_dim="shard"))
                    torch.cuda.empty_cache()
                    distributed.sync_model_states(model, device_mesh.get_group(mesh_dim="replicate"))
                    torch.cuda.empty_cache()

                with misc.timer("Creating FSDP model"):
                    self.model = FSDP(
                        model.to(**self.tensor_kwargs),
                        sync_module_states=True,  # it can reduce network traffic by only loading model in rank0 and sync
                        sharding_strategy=strategy,
                        auto_wrap_policy=get_wrap_policy(model),
                        process_group=fsdp_process_group,
                        limit_all_gathers=True,
                    )

                    if self.config.fsdp.checkpoint:
                        fsdp_blocks_cls = model.net.fsdp_wrap_block_cls
                        fsdp_blocks_cls = (
                            list(fsdp_blocks_cls)
                            if isinstance(fsdp_blocks_cls, (list, tuple, set))
                            else [fsdp_blocks_cls]
                        )
                        log.critical(f"Applying FSDP checkpointing with FSDP blocks: {fsdp_blocks_cls}")
                        apply_fsdp_checkpointing(self.model, list_block_cls=fsdp_blocks_cls)

            torch.cuda.empty_cache()

        def on_before_zero_grad(
            self, optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler.LRScheduler, iteration: int
        ) -> None:
            del scheduler, optimizer

            if self.config.ema.enabled:
                # calculate beta for EMA update
                if iteration == 0:
                    beta = 0.0
                else:
                    i = iteration + 1
                    beta = (1 - 1 / i) ** (self.ema_exp_coefficient + 1)
                self.model_ema_worker.update_average(self.model, self.model_ema, beta=beta)

        def training_step(
            self, data_batch: Dict[str, torch.Tensor], iteration: int
        ) -> Tuple[Dict[str, torch.Tensor] | torch.Tensor]:
            # ! Important!!!
            # ! make sure the training step is the same as the forward method~(training_step in the super class)
            # ! this is necessary to trigger the FSDP-specific pre- and post-forward sharding logic
            return self.model(data_batch, iteration)

        def state_dict(self) -> Dict:
            raise NotImplementedError(
                "FSDPDiffModle does not support state_dict, use state_dict_model and FSDPCheckpointer"
            )

        @misc.timer("FSDP state_dict_model")
        def state_dict_model(self) -> Dict:
            with FSDP.summon_full_params(self.model):
                pass
            with FSDP.state_dict_type(
                self.model, StateDictType.FULL_STATE_DICT, FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
            ):
                model_state = self.model.state_dict()
            if self.config.ema.enabled:
                with FSDP.summon_full_params(self.model_ema):
                    pass
                with FSDP.state_dict_type(
                    self.model_ema,
                    StateDictType.FULL_STATE_DICT,
                    FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
                ):
                    ema_model_state = self.model_ema.state_dict()
            else:
                ema_model_state = None
            return {
                "model": model_state,
                "ema": ema_model_state,
            }

        def load_state_dict(self, state_dict: Dict, strict: bool = True, assign: bool = False) -> None:
            raise NotImplementedError("FSDPDiffModle does not support load_state_dict, using FSDPCheckpointer")

        def init_optimizer_scheduler(
            self, optimizer_config: LazyDict, scheduler_config: LazyDict
        ) -> tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LRScheduler]:
            optimizer, scheduler = super().init_optimizer_scheduler(optimizer_config, scheduler_config)
            self.fsdp_checkpointer.load_optim_scheduler_during_init(
                self.model,
                optimizer,
                scheduler,
            )
            return optimizer, scheduler

        @contextmanager
        def ema_scope(self, context=None, is_cpu=False):
            if self.config.ema.enabled:
                self.model_ema_worker.cache(self.model.parameters(), is_cpu=is_cpu)
                self.model_ema_worker.copy_to(src_model=self.model_ema, tgt_model=self.model)
                if context is not None:
                    log.info(f"{context}: Switched to EMA weights")
            try:
                yield None
            finally:
                if self.config.ema.enabled:
                    self.model_ema_worker.restore(self.model.parameters())
                    if context is not None:
                        log.info(f"{context}: Restored training weights")

        def get_ckpt_postfix(self) -> Tuple[str, int]:
            """Get the checkpoint file postfix. check FSDPCheckpointer for more details

            Args:
                iteration (int): The current iteration number.

            Returns:
                postfix (str): The postfix of the checkpoint file.
                replicate_idx, shard_idx (int), current gpu replicate_idx, shard_idx in FSDP \
                    we will not save each ema model in each GPU, \
                    ema model with same rate will be saved once
                total_ema_num (int)
            """
            mesh_shape = parallel_state.fsdp_device_mesh.shape
            total_ema_num = min(self.config.ema.num, mesh_shape[0])
            replicate_idx, shard_idx = parallel_state.fsdp_device_mesh.get_coordinate()
            if replicate_idx == 0:
                return "", 0, shard_idx, total_ema_num
            if self.config.ema.enabled:
                if replicate_idx < self.config.ema.num:
                    return f"_RANK{replicate_idx}", replicate_idx, shard_idx, total_ema_num
            return "", replicate_idx, shard_idx, total_ema_num

    return FSDPClass


@diffusion_fsdp_class_decorator
class FSDPDiffusionModel(DiffusionModel):
    pass