# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations from functools import partial from typing import Dict, Any, TYPE_CHECKING, Callable import torch import torch.nn.functional as F from lipforcing.methods import FastGenModel, CausVidModel from lipforcing.utils import expand_like if TYPE_CHECKING: from lipforcing.configs.config import BaseModelConfig as ModelConfig class KDModel(FastGenModel): def __init__(self, config: ModelConfig): """ Args: config (ModelConfig): The configuration for the knowledge distillation model. This model directly learns the pre-constructed ODE pairs from the teacher model. """ super().__init__(config) self.config = config def build_model(self): super().build_model() self.load_student_weights_and_ema() def _get_outputs( self, gen_data: torch.Tensor, input_student: torch.Tensor = None, condition: Any = None, ) -> Dict[str, torch.Tensor | Callable]: if self.config.student_sample_steps == 1: assert input_student is not None, "input_student must be provided for KDModel" return {"gen_rand": gen_data, "input_rand": input_student} else: noise = torch.randn_like(gen_data, dtype=self.precision) gen_rand_func = partial( self.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, precision_amp=self.precision_amp_infer, ) return {"gen_rand": gen_rand_func, "input_rand": noise, "gen_rand_train": gen_data} def single_train_step( self, data: Dict[str, Any], iteration: int ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor | Callable]]: """ Single training step for knowledge distillation model. Important! For multistep KD distillation, t_list must be aligned with the `path`'s timesteps: 1) Ensure t_list corresponds exactly to path timesteps 2) Please check the `path_timesteps` item in index.json of the paired dataset 3) num_inference_steps in denoise path must be 4 4) student_sample_steps must be either 2 or 4 5) Current approach assumes uniform spacing: t_list=[t1, t3] → path indices [0, 2] Args: data (Dict[str, Any]): Data dict for the current iteration. iteration (int): Current training iteration Returns: loss_map (dict[str, torch.Tensor]): Dictionary containing the loss values outputs (dict[str, torch.Tensor]): Dictionary containing the network output """ denoised_data = data["real"] condition = data["condition"] batch_size = denoised_data.shape[0] if self.config.student_sample_steps == 1: # perform single-step distillation if "noise" in data: input_student = data["noise"] elif "path" in data: input_student = data["path"][:, 0, ...] # the first step is noise else: raise ValueError("Noise or path must be provided for KDModel") t_student = torch.full( (batch_size,), self.net.noise_scheduler.max_t, device=self.device, dtype=self.net.noise_scheduler.t_precision, ) else: # perform multiple-step distillation assert "path" in data, "path must be provided for KDModel" denoise_path = data["path"] # [batch_size, num_inf_steps, C, num_frames, H, W] assert denoise_path.shape[1] == 4, "num_inference_steps in denoise path must be 4" assert ( denoise_path.shape[1] % self.config.student_sample_steps == 0 ), f"student_sample_steps must be either 2 or 4, but got {self.config.student_sample_steps}" t_student, t_list_ids = self.net.noise_scheduler.sample_from_t_list( batch_size, sample_steps=self.config.student_sample_steps, t_list=self.config.sample_t_cfg.t_list, return_ids=True, device=self.device, ) # Important: Ensure t_list corresponds exactly to path timesteps # Current approach assumes uniform spacing: t_list=[t1, t3] → path indices [0, 2] path_indices = t_list_ids * (denoise_path.shape[1] // self.config.student_sample_steps) path_indices = expand_like(path_indices, denoise_path).expand( -1, -1, *denoise_path.shape[2:] ) # [batch_size, 1, C, num_frames, H, W] input_student = torch.gather(denoise_path, 1, path_indices).squeeze(1) # [batch_size, C, num_frames, H, W] gen_data = self.gen_data_from_net(input_student, t_student, condition=condition) # Compute the l2 loss between the generated data and the denoised data loss = 0.5 * F.mse_loss(gen_data, denoised_data, reduction="mean") # Build output dictionaries loss_map = { "total_loss": loss, "recon_loss": loss, } outputs = self._get_outputs(gen_data, input_student, condition=condition) return loss_map, outputs class CausalKDModel(KDModel): def _get_outputs( self, gen_data: torch.Tensor, input_student: torch.Tensor = None, condition: Any = None, ) -> Dict[str, torch.Tensor | Callable]: noise = torch.randn_like(gen_data, dtype=self.precision) context_noise = getattr(self.config, "context_noise", 0) # Reuse CausVidModel's autoregressive generation logic gen_rand_func = partial( CausVidModel.generator_fn, net=self.net_inference, noise=noise, condition=condition, student_sample_steps=self.config.student_sample_steps, 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} def single_train_step( self, data: Dict[str, Any], iteration: int ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor | Callable]]: """ Single training step for knowledge distillation model. Important! t_list must be the same with the `path`'s timesteps. Please check the `path_timesteps` item in index.json of the paired dataset. Args: data (Dict[str, Any]): Data dict for the current iteration. iteration (int): Current training iteration Returns: loss_map (dict[str, torch.Tensor]): Dictionary containing the loss values outputs (dict[str, torch.Tensor]): Dictionary containing the network output """ denoise_path = data["path"] # shape is [batch_size, num_inf_steps, C, num_frames, H, W] denoised_data = data["real"] # [batch_size, C, num_frames, H, W] condition = data["condition"] batch_size, num_frames = denoise_path.shape[0], denoise_path.shape[3] chunk_size = self.net.chunk_size # add noise t_inhom, ids = 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, # Note t_list to be aligned the `path`'s timesteps device=self.device, dtype=denoise_path.dtype, ) # [batch_size, num_frames] expand_shape = [ids.shape[0], 1, 1, ids.shape[1]] + [1] * max(0, denoise_path.ndim - 4) ids = ids.view(expand_shape).expand(-1, -1, *denoise_path.shape[2:]) # [batch_size, 1, C, num_frames, H, W] denoise_path_all = torch.cat([denoise_path, denoised_data.unsqueeze(1)], dim=1) # gather clean data noisy_data = torch.gather(denoise_path_all, 1, ids).squeeze(1) # [batch_size, C, num_frames, H, W] # generate data gen_data = self.gen_data_from_net(noisy_data, t_inhom, condition=condition) # Compute the l2 loss between the generated data and the denoised data loss = 0.5 * F.mse_loss(gen_data, denoised_data, reduction="mean") # Build output dictionaries loss_map = { "total_loss": loss, "recon_loss": loss, } outputs = self._get_outputs(gen_data, condition=condition) return loss_map, outputs