gem/network/gem_diffusion.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
from copy import deepcopy

import torch
import torch.nn as nn
from hydra.utils import instantiate

from gem.diffusion_utils.model_util import create_gaussian_diffusion
from gem.diffusion_utils.resample import create_named_schedule_sampler
from gem.utils.net_utils import length_to_mask
from gem.utils.pylogger import Log

from .gem_cfg_sampler import ClassifierFreeSampleModel


class GEMDiffusion(nn.Module):
    def __init__(
        self,
        model_cfg,
        max_len=120,
        # condition
        cliffcam_dim=3,
        cam_angvel_dim=6,
        cam_t_vel_dim=3,
        imgseq_dim=1024,
        observed_motion_3d_dim=151,
        encoded_music_dim=438,
        encoded_audio_dim=128,
        latent_dim=512,
        dropout=0.1,
        args=None,
        cond_merge_strategy="add",
        cond_exists_dim=512,
        music_mask_prob=0.1,
        img_process_modules=None,
        img_process_modules_enable_grad=None,
        multi_text_module_cfg=None,
        regression_only=False,
        **kwargs,
    ):
        super().__init__()
        self.model_cfg = model_cfg
        self.args = args
        self.max_len = max_len
        if img_process_modules_enable_grad is None:
            img_process_modules_enable_grad = {}
        if multi_text_module_cfg is None:
            multi_text_module_cfg = {}

        self.regression_input_type = self.args.get("regression_input_type", "zero")
        self.regression_only = bool(regression_only or self.model_cfg.get("regression_only", False))
        self.mask_localpose = kwargs.get("mask_localpose", False)
        self.mask_localpose_prob = kwargs.get("mask_localpose_prob", 0.0)

        self.denoiser = instantiate(self.model_cfg.denoiser)
        self.init_diffusion()
        self.text_encoder, self.tokenizer = None, None

    def init_diffusion(self):
        self.train_diffusion = create_gaussian_diffusion(self.model_cfg.diffusion, training=True)
        self.test_diffusion = create_gaussian_diffusion(self.model_cfg.diffusion, training=False)
        gen_only_diffusion = deepcopy(self.model_cfg.diffusion)
        gen_only_diffusion.test_timestep_respacing = self.model_cfg.diffusion.get(
            "gen_only_test_timestep_respacing", "50"
        )
        Log.info(f"Gen only test timestep respacing: {gen_only_diffusion.test_timestep_respacing}")
        self.test_gen_only_diffusion = create_gaussian_diffusion(gen_only_diffusion, training=False)
        self.schedule_sampler = create_named_schedule_sampler(
            self.model_cfg.diffusion.schedule_sampler_type, self.train_diffusion
        )
        return

    def forward_train(self, inputs, mode):
        assert self.training, "forward_train should only be called during training"
        diffusion = self.train_diffusion if self.training else self.test_diffusion
        length = inputs["length"]
        # target_x = inputs["target_x"]
        motion = inputs["motion"]
        f_cond = inputs["f_cond"]
        f_empty = inputs["f_empty"]
        B, L, _ = motion.shape

        vis_mask = length_to_mask(length, L)  # (B, L)
        valid_mask = inputs["mask"]["valid"]
        assert (vis_mask == valid_mask).all()

        denoiser_kwargs = {
            "y": {
                "text": inputs.get("caption", [""] * B),
                "f_cond": f_cond,
                "mask": vis_mask,
                "length": length,
            },
            "inputs": inputs,
            "sample_indices_dict": inputs["sample_indices_dict"],
        }
        if "encoded_text" in inputs:
            denoiser_kwargs["y"]["encoded_text"] = inputs["encoded_text"]
        if "observed_motion_3d" in inputs:
            denoiser_kwargs["observed_motion_3d"] = inputs["observed_motion_3d"]
            denoiser_kwargs["motion_mask_3d"] = inputs["motion_mask_3d"]
            denoiser_kwargs["rm_text_flag"] = inputs["rm_text_flag"]

        if mode == "regression":
            t = (torch.ones(B) * (diffusion.original_num_steps - 1)).long().to(motion.device)
            t_weights = torch.ones(B).to(motion.device)
            x_start = motion
            if self.regression_input_type == "zero":
                x_t = torch.zeros_like(motion)
            elif self.regression_input_type == "normal":
                x_t = torch.randn_like(motion)
            else:
                raise ValueError(f"Unsupported regression_input_type: {self.regression_input_type}")
        elif mode == "diffusion":
            t, t_weights = self.schedule_sampler.sample(motion.shape[0], motion.device)
            if "regression_outputs" in inputs:
                pred_x_start_regression = inputs["regression_outputs"]["model_output"][
                    "pred_x_start"
                ].detach()
            else:
                raise ValueError("No regression outputs found")
                # pred_x_start_regression = torch.zeros_like(motion)
            x_start_reg = pred_x_start_regression.clone()
            x_start = motion.clone()
            x_start[inputs["mask"]["2d_only"]] = x_start_reg[inputs["mask"]["2d_only"]]
            # regression_mask = (
            #     torch.rand(B).to(motion.device) < self.args.use_regression_outputs_prob
            # ).float()
            # if "gen_only" in inputs and self.args.get("use_gt_for_gen_only", True):
            #     regression_mask[inputs["gen_only"]] = 0
            # x_start = x_start_reg * regression_mask[:, None, None] + x_start_gt * (
            #     1 - regression_mask[:, None, None]
            # )
            noise = torch.randn_like(x_start)
            x_t = self.train_diffusion.q_sample(x_start.clone(), t, noise=noise)
            if self.mask_localpose:
                localpose_idx = self.denoiser.denoiser3d.endecoder.obs_indices_dict["body_pose"]
                mask_localpose = torch.rand((B,), device=x_t.device) < self.mask_localpose_prob
                x_t[:, :, localpose_idx] = x_start[:, :, localpose_idx] * mask_localpose[
                    :, None, None
                ].float() + x_t[:, :, localpose_idx] * (1 - mask_localpose[:, None, None].float())
                denoiser_kwargs["y"]["f_cond"] = (
                    f_cond * (1 - mask_localpose[:, None, None].float())
                    + f_empty * mask_localpose[:, None, None].float()
                )

        denoise_out = self.denoiser(
            x_t, diffusion._scale_timesteps(t), return_aux=False, **denoiser_kwargs
        )

        output = {
            "target_x_start": x_start,
            "t_weights": t_weights,
        }
        output.update(denoise_out)
        for x in self.args.out_attr:
            assert x in output, f"Output {x} not found in denoise_out"

        return output

    def forward_test(self, inputs, progress=False):
        assert not self.training, "forward_test should only be called during inference"
        diffusion = self.test_gen_only_diffusion

        denoiser = self.denoiser
        length = inputs["length"]
        B, L = inputs["B"], inputs["L"]

        motion = inputs["motion"]
        f_cond, f_uncond = inputs["f_cond"], inputs["f_uncond"]
        f_empty = inputs["f_empty"]

        vis_mask = length_to_mask(length, L)  # (B, L)

        denoiser_kwargs = {
            "y": {
                "text": inputs.get("caption", [""] * B),
                "f_cond": f_cond,
                "f_uncond": f_uncond,
                "f_empty": f_empty,
                "mask": vis_mask,
                "length": length,
            },
            "inputs": inputs,
            "sample_indices_dict": inputs["sample_indices_dict"],
        }
        if "encoded_text" in inputs:
            denoiser_kwargs["y"]["encoded_text"] = inputs["encoded_text"]
        if "meta" in inputs and "multi_text_data" in inputs["meta"][0]:
            denoiser_kwargs["y"]["multi_text_data"] = inputs["meta"][0]["multi_text_data"]
        if "observed_motion_3d" in inputs:
            denoiser_kwargs["observed_motion_3d"] = inputs["observed_motion_3d"]
            denoiser_kwargs["motion_mask_3d"] = inputs["motion_mask_3d"]
            denoiser_kwargs["rm_text_flag"] = inputs.get("rm_text_flag", None)

        if self.args.get("use_cfg_sampler_for_gen", False):
            denoiser = ClassifierFreeSampleModel(denoiser, mask_localpose=self.mask_localpose)
            denoiser_kwargs["y"]["scale"] = self.model_cfg.diffusion.guidance_param
        diff_sampler = self.model_cfg.diffusion.get("sampler", "ddim")
        if diff_sampler == "ddim":
            sample_fn = diffusion.ddim_sample_loop_with_aux
            kwargs = {"eta": self.model_cfg.diffusion.ddim_eta}
        else:
            raise NotImplementedError(f"Sampler {diff_sampler} not implemented")

        if self.args.get("force_zero_noise", False):
            noise = torch.zeros_like(motion)
        elif self.args.get("force_rand_noise", False):
            noise = torch.randn_like(motion)
        else:
            noise = torch.randn_like(motion)

        if self.args.get("return_mid", False):
            kwargs["return_mid"] = True

        if self.regression_only:
            # Regression-only inference matches the old codepath:
            # directly run the denoiser at the last timestep with zero input.
            x_t = torch.zeros_like(motion)
            t = (
                (torch.ones(B) * (self.train_diffusion.original_num_steps - 1))
                .long()
                .to(motion.device)
            )
            denoise_out = self.denoiser(
                x_t,
                self.train_diffusion._scale_timesteps(t),
                return_aux=False,
                **denoiser_kwargs,
            )
            output = denoise_out.copy()
            if output.get("pred_x", None) is None and "pred_x_start" in output:
                output["pred_x"] = output["pred_x_start"]
        else:
            denoise_out = sample_fn(
                denoiser,
                motion.shape,
                clip_denoised=False,
                model_kwargs=denoiser_kwargs,
                skip_timesteps=0,  # 0 is the default value - i.e. don't skip any step
                init_image=None,
                progress=progress,
                dump_steps=None,
                noise=noise,
                const_noise=False,
                **kwargs,
            )
            output = denoise_out.copy()
            if output.get("pred_x", None) is None and "pred_x_start" in output:
                output["pred_x"] = output["pred_x_start"]

        for x in self.args.out_attr:
            assert x in output, f"Output {x} not found in denoise_out"
        return output

    def forward(
        self,
        inputs,
        train=False,
        postproc=False,
        static_cam=False,
        mode=None,
        test_mode=None,
        normalizer_stats=None,
    ):
        if train:
            return self.forward_train(inputs, mode=mode)
        else:
            return self.forward_test(inputs)