lipforcing/methods/omniavatar_diffusion_forcing.py

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

"""
OmniAvatar Diffusion Forcing model for Stage 1 initialization.

Alternative to ODE-based KD (CausalKDModel). Instead of pre-computing
ODE trajectories from the teacher, this adds Gaussian noise to real data at
inhomogeneous block-wise timesteps and trains the student to denoise with L2 loss.
No teacher model or ODE generation needed.
"""

from __future__ import annotations

import os
from typing import Any, Dict, TYPE_CHECKING, Callable
from functools import partial

import torch
import torch.nn.functional as F

from lipforcing.methods.knowledge_distillation.KD import KDModel
from lipforcing.methods.distribution_matching.causvid import CausVidModel
import lipforcing.utils.logging_utils as logger

if TYPE_CHECKING:
    from lipforcing.configs.config import BaseModelConfig as ModelConfig


class OmniAvatarDiffusionForcingModel(KDModel):
    """Diffusion Forcing on real data — alternative to ODE KD for Stage 1.

    Adds noise to real data at inhomogeneous block-wise timesteps.
    Student denoises -> L2 loss vs clean data. No teacher ODE needed.

    Inheritance: OmniAvatarDiffusionForcingModel -> KDModel -> FastGenModel
    """

    def __init__(self, config: ModelConfig):
        super().__init__(config)

    def build_model(self):
        """Build model and optionally load VAE for visual logging."""
        super().build_model()
        # FastGenModel.build_model (model.py:260) unconditionally does
        # `self.net.train().requires_grad_(True)` after instantiating the
        # network, which destroys the LoRA freeze that PEFT's
        # inject_adapter_in_model set up inside _load_weights.  Re-apply
        # the freeze here so the optimizer (constructed downstream by
        # init_optimizers, which filters on requires_grad) only sees the
        # intended trainable subset (LoRA A/B + user-listed unfreeze
        # modules).  No-op when merge_lora=True or when PEFT injection
        # didn't actually run (apply_lora_freeze handles both).
        if hasattr(self.net, "apply_lora_freeze"):
            self.net.apply_lora_freeze()

        vae_path = getattr(self.config, "vae_path", "") or ""
        if vae_path and os.path.exists(vae_path):
            self._load_vae(vae_path)

    def init_optimizers(self):
        """Initialize optimizers, with a defensive LoRA freeze re-apply.

        Belt-and-suspenders against any post-build_model code path that
        might also reset requires_grad (e.g., FSDP wrap converting params
        to DTensors and not preserving requires_grad in some PyTorch
        versions).  apply_lora_freeze is idempotent.
        """
        if hasattr(self.net, "apply_lora_freeze"):
            self.net.apply_lora_freeze()
        super().init_optimizers()

    def _load_vae(self, vae_path: str):
        """Load WanVideoVAE for decoding generated samples in wandb visual logging.

        The wandb callback calls model.net.vae.decode(tensor) with a single [B,C,T,H,W] tensor.
        WanVideoVAE.decode expects (hidden_states_list, device). We wrap it for compatibility.
        """
        from OmniAvatar.models.wan_video_vae import WanVideoVAE

        # Load VAE weights — checkpoint keys lack "model." prefix, use converter
        raw_vae = WanVideoVAE(z_dim=16)
        vae_state = torch.load(vae_path, map_location="cpu", weights_only=False)
        # Add "model." prefix to match WanVideoVAE's self.model attribute
        if any(k.startswith("encoder.") for k in vae_state):
            vae_state = {f"model.{k}": v for k, v in vae_state.items()}
        raw_vae.load_state_dict(vae_state)
        device_str = f"cuda:{self.device}" if isinstance(self.device, int) else str(self.device)
        raw_vae = raw_vae.to(device_str).eval()

        # Wrap decode to match wandb callback's expected interface: decode(tensor) -> tensor
        class VAEWrapper:
            def __init__(self, vae, device):
                self._vae = vae
                self._device = device

            def decode(self, x):
                """Decode [B, C, T, H, W] latent to [B, 3, T*4, H*8, W*8] pixel video."""
                with torch.no_grad():
                    # WanVideoVAE.decode expects list of [C,T,H,W] tensors in float32
                    return self._vae.decode([xi.float() for xi in x], self._device)

            def to(self, *args, **kwargs):
                return self

        self.net.vae = VAEWrapper(raw_vae, device_str)
        logger.info(f"Loaded WanVideoVAE from {vae_path} for visual logging")

    # Use CausVidModel's AR sample loop for visualization (chunk-by-chunk with KV cache).
    # Without this, FastGenModel._student_sample_loop processes the entire video as one
    # bidirectional pass, which doesn't reflect actual AR inference behavior.
    _student_sample_loop = CausVidModel._student_sample_loop

    def _build_condition(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """Build OmniAvatar condition dict from data batch.

        Expected shapes (after collation, with batch dim):
            text_embeds:    [B, 1, 512, 4096] or [B, 512, 4096]
            audio_emb:      [B, 81, audio_dim]
            mask:           [B, H, W] or [H, W]
            masked_video:   [B, 16, T, H, W]
            ref_sequence:   [B, 16, T, H, W] (optional)

        Args:
            data: Batch from OmniAvatarDataset.

        Returns:
            Condition dict for OmniAvatar networks.
        """
        for key in ("real", "text_embeds", "audio_emb", "mask", "masked_video"):
            assert key in data, f"Missing required key '{key}' in data batch"

        real_data = data["real"]
        ref_latent = real_data[:, :, :1, :, :]  # [B, 16, 1, H, W]

        mask = data["mask"]
        if mask.dim() == 3:
            mask = mask[0]

        text_embeds = data["text_embeds"]
        if text_embeds.dim() == 4:
            assert text_embeds.shape[1] == 1, (
                f"text_embeds dim 1 must be 1 for squeeze, got shape {list(text_embeds.shape)}"
            )
            text_embeds = text_embeds.squeeze(1)

        condition = {
            "text_embeds": text_embeds,
            "audio_emb": data["audio_emb"],
            "ref_latent": ref_latent,
            "mask": mask,
            "masked_video": data["masked_video"],
        }
        if "ref_sequence" in data:
            condition["ref_sequence"] = data["ref_sequence"]

        return condition

    def _get_outputs(
        self,
        gen_data: torch.Tensor,
        input_student: torch.Tensor = None,
        condition: Any = None,
    ) -> Dict[str, torch.Tensor | Callable]:
        has_vae = hasattr(self.net, "vae")
        if not has_vae:
            logger.debug("No VAE loaded on net — visual logging disabled")
        if has_vae and condition is not None:
            noise = torch.randn_like(gen_data, dtype=self.precision)
            context_noise = getattr(self.config, "context_noise", 0)
            gen_rand_func = partial(
                CausVidModel.generator_fn,
                net=self.net_inference,
                noise=noise,
                condition=condition,
                student_sample_steps=self.config.student_sample_steps,
                student_sample_type=self.config.student_sample_type,
                t_list=self.config.sample_t_cfg.t_list,
                context_noise=context_noise,
                precision_amp=self.precision_amp_infer,
            )
            return {"gen_rand": gen_rand_func, "input_rand": noise, "gen_rand_train": gen_data.detach()}
        return {"gen_rand_train": gen_data.detach()}

    def single_train_step(
        self, data: Dict[str, Any], iteration: int
    ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor | Callable]]:
        """Single training step using diffusion forcing on real data.

        Instead of gathering from pre-computed ODE trajectories (as in CausalKDModel),
        this adds Gaussian noise to real data at inhomogeneous block-wise timesteps.
        """
        real_data = data["real"]  # [B, 16, 21, 64, 64]
        condition = self._build_condition(data)

        batch_size, num_frames = real_data.shape[0], real_data.shape[2]
        chunk_size = self.net.chunk_size

        # Sample inhomogeneous block-wise timesteps
        t_inhom, _ = self.net.noise_scheduler.sample_t_inhom(
            batch_size,
            num_frames,
            chunk_size,
            sample_steps=self.config.student_sample_steps,
            t_list=self.config.sample_t_cfg.t_list,
            device=self.device,
            dtype=real_data.dtype,
        )  # [B, T]

        # Diffusion forcing: add noise to real data at sampled timesteps
        eps = torch.randn_like(real_data)
        t_inhom_expanded = t_inhom[:, None, :, None, None]  # [B, 1, T, 1, 1]
        noisy_data = self.net.noise_scheduler.forward_process(real_data, eps, t_inhom_expanded)

        # Student denoise
        gen_data = self.gen_data_from_net(noisy_data, t_inhom, condition=condition)

        # L2 loss
        loss = 0.5 * F.mse_loss(gen_data, real_data, reduction="mean")

        # Outputs for logging (detached to avoid holding autograd references)
        outputs = self._get_outputs(gen_data.detach(), condition=condition)

        loss_map = {"total_loss": loss, "recon_loss": loss.detach()}
        return loss_map, outputs