# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ OmniAvatar Self-Forcing model for V2V lip sync distillation. Overrides _prepare_training_data to build OmniAvatar-specific condition dicts with audio, reference frames, spatial mask, masked video, and reference sequence. """ from __future__ import annotations from typing import Any, Dict, TYPE_CHECKING import os import torch from lipforcing.methods.distribution_matching.self_forcing import SelfForcingModel from lipforcing.utils import instantiate from lipforcing.utils.distributed import synchronize import lipforcing.utils.logging_utils as logger if TYPE_CHECKING: from lipforcing.configs.methods.config_self_forcing import ModelConfig class OmniAvatarSelfForcingModel(SelfForcingModel): """Self-Forcing distillation for OmniAvatar V2V audio-driven lip sync. Inherits the full Self-Forcing training loop (rollout_with_gradient, VSD loss, fake_score/discriminator updates). Only overrides data preparation to handle OmniAvatar's condition dict format, and build_model to support a separate fake_score architecture, separate from the teacher. """ def __init__(self, config: ModelConfig): super().__init__(config) def _setup_grad_requirements(self, iteration: int) -> None: """Override parent's grad-toggle to preserve LoRA freeze. The parent dmd2._setup_grad_requirements does ``self.fake_score.train().requires_grad_(True)`` on critic-update steps, which wipes the LoRA freeze and causes the trainable-only checkpoint filter to save the full 14B base. For LoRA mode, only toggle .train() / .eval() — don't touch requires_grad on fake_score. Gated on ``unfreeze_modules`` non-empty so we only override in the explicit "LoRA + selective unfreeze" regime (the 14B convention). 1.3B SF runs construct with merge_lora=False but train fake_score as full-FT (the requires_grad wipes effectively unfreeze the PEFT base) — gating on unfreeze_modules preserves that. Otherwise fall back to parent. """ is_lora_mode = ( hasattr(self.fake_score, "apply_lora_freeze") and not getattr(self.fake_score, "merge_lora", True) and bool(getattr(self.fake_score, "unfreeze_modules", [])) ) if not is_lora_mode: return super()._setup_grad_requirements(iteration) # LoRA mode: only toggle train/eval mode for BN/Dropout; # leave requires_grad as configured by apply_lora_freeze. if iteration % self.config.student_update_freq == 0: # student step self.fake_score.eval() if self.config.gan_loss_weight_gen > 0: self.discriminator.eval().requires_grad_(False) else: # critic-only step self.fake_score.train() if self.config.gan_loss_weight_gen > 0: self.discriminator.train().requires_grad_(True) def single_train_step(self, data: Dict[str, Any], iteration: int): """Combined fake_score + student update on student steps (1:5 ratio). Matches the original Self-Forcing training loop where the critic updates EVERY step, including on the generator (student) step. This gives a true 1:5 ratio (5 critic updates per student update in a 5-step cycle). On non-student steps (iter % freq != 0): delegate to base class (fake_score only). On student steps (iter % freq == 0): run fake_score backward manually (freeing its graph to save memory), then return student loss for the trainer's backward. Both sets of gradients accumulate across grad_accum rounds; the trainer steps both optimizers at the end. """ if iteration % self.config.student_update_freq != 0: # Critic-only step — unchanged from base class return super().single_train_step(data, iteration) # === Combined step: fake_score + student === real_data, condition, neg_condition = self._prepare_training_data(data) grad_accum_rounds = getattr(self.config, "grad_accum_rounds", None) or 1 # --- Step 1: Fake score forward + manual backward (frees graph) --- # Keep self.net.requires_grad_ unchanged (True) — same as the exclusive # pattern. The no_grad() inside _fake_score_discriminator_update_step is # sufficient. Toggling requires_grad on FSDP2 DTensors leaves stale # internal state that breaks gradient checkpointing recomputation. # # For fake_score: only set train() mode (BN/Dropout); leave # requires_grad as configured by build_model.apply_lora_freeze # (LoRA mode) or super().build_model (full FT). Calling # requires_grad_(True) here would wipe the LoRA freeze, making # the trainable-only checkpoint filter save the full 14B base. self.fake_score.train() if self.config.gan_loss_weight_gen > 0: self.discriminator.train().requires_grad_(True) input_fs, t_student_fs, t_fs, eps_fs = self._generate_noise_and_time(real_data) fake_loss_map, _ = self._fake_score_discriminator_update_step( input_fs, t_student_fs, t_fs, eps_fs, real_data, condition=condition, ) # NOTE: the critic backward is intentionally UNSCALED (grad_accum_rounds # is not a declared model-config field, so the divisor below is always 1 # at runtime). The released checkpoints were trained with this behavior; # it is kept as-is to reproduce them. The graph is freed here so it # doesn't overlap with the student forward. (fake_loss_map["total_loss"] / grad_accum_rounds).backward() # Freeze discriminator before the student step so the generator GAN loss # backward doesn't accumulate spurious grads in the discriminator. # Safe: the discriminator is fully FSDP-wrapped (auto-wrap path) and has # no gradient checkpointing, so requires_grad_ toggling doesn't trigger # the DTensor stale-state issue that affects the student's blocks. if self.config.gan_loss_weight_gen > 0: self.discriminator.eval().requires_grad_(False) # --- Step 2: Student forward (returned for trainer's backward) --- self.net.clear_caches() input_student, t_student, t, eps = self._generate_noise_and_time(real_data) student_loss_map, student_outputs = self._student_update_step( input_student, t_student, t, eps, data, condition=condition, neg_condition=neg_condition, ) # Attach Step 1 losses for logging (detached — no gradient) for k, v in fake_loss_map.items(): if k != "total_loss" and torch.is_tensor(v): student_loss_map[k] = v.detach() return student_loss_map, student_outputs def get_optimizers(self, iteration: int) -> list: """On student steps, return all active optimizers (trainer steps all).""" if iteration % self.config.student_update_freq == 0: opts = [self.net_optimizer, self.fake_score_optimizer] if self.config.gan_loss_weight_gen > 0: opts.append(self.discriminator_optimizer) return opts else: if self.config.gan_loss_weight_gen > 0: return [self.fake_score_optimizer, self.discriminator_optimizer] else: return [self.fake_score_optimizer] def get_lr_schedulers(self, iteration: int) -> list: """On student steps, return all active schedulers (trainer steps all).""" if iteration % self.config.student_update_freq == 0: scheds = [self.net_lr_scheduler, self.fake_score_lr_scheduler] if self.config.gan_loss_weight_gen > 0: scheds.append(self.discriminator_lr_scheduler) return scheds else: if self.config.gan_loss_weight_gen > 0: return [self.fake_score_lr_scheduler, self.discriminator_lr_scheduler] else: return [self.fake_score_lr_scheduler] def build_model(self): """Override to instantiate fake_score from config.fake_score if provided. The base DMD2Model.build_model() always creates fake_score from self.teacher_config (= config.teacher), which is 14B. When config.fake_score is set, we use that instead (a separate bidirectional net). """ super().build_model() fake_score_config = getattr(self.config, "fake_score", None) if fake_score_config is not None: logger.info("Re-instantiating fake_score from config.fake_score") with self._get_meta_init_context(): self.fake_score = instantiate(fake_score_config) synchronize() # Restore PEFT-applied freeze after FastGenModel.build_model:260's # `self.net.train().requires_grad_(True)` wipe. Same recovery hook # used in OmniAvatarDiffusionForcingModel. Defensive on fake_score # too — it's not subject to the wipe today (which only touches # self.net), but if a future config sets merge_lora=False on # fake_score with selective unfreeze, this catches drift from any # later mutation. Idempotent and a no-op when merge_lora=True. # # Gated on `unfreeze_modules` being non-empty so we only fire the # freeze in the explicit "LoRA + selective unfreeze" regime (the # 14B convention). 1.3B SF runs construct with merge_lora=False # to preserve PEFT structure for adapter-style ckpt loading, but # train as full-FT — gating on unfreeze_modules preserves that # regime. See `apply_lora_freeze` body which # explicitly relies on unfreeze_modules to re-enable specific # submodules; with no unfreeze_modules the call would freeze # everything except the LoRA adapters (a regime we never use for # 1.3B). if hasattr(self.net, "apply_lora_freeze") and getattr(self.net, "unfreeze_modules", []): self.net.apply_lora_freeze() if ( hasattr(self, "fake_score") and hasattr(self.fake_score, "apply_lora_freeze") and getattr(self.fake_score, "unfreeze_modules", []) ): self.fake_score.apply_lora_freeze() # Load VAE for wandb visual logging (same logic as OmniAvatarDiffusionForcing) 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): """Defensive LoRA freeze re-apply right before optimizer construction. Belt-and-suspenders against any post-build_model code path (e.g., FSDP wrap converting params to DTensors and not preserving requires_grad in some PyTorch versions) that might reset requires_grad on frozen params. apply_lora_freeze is idempotent and a no-op when LoRA isn't in use on the network. Gated on unfreeze_modules — see build_model for rationale. 1.3B SF runs leave fake_score as full-FT after the requires_grad wipes; only the 14B LoRA + selective-unfreeze regime needs the freeze. """ if hasattr(self.net, "apply_lora_freeze") and getattr(self.net, "unfreeze_modules", []): self.net.apply_lora_freeze() if ( hasattr(self, "fake_score") and hasattr(self.fake_score, "apply_lora_freeze") and getattr(self.fake_score, "unfreeze_modules", []) ): self.fake_score.apply_lora_freeze() super().init_optimizers() def _load_vae(self, vae_path: str): """Load WanVideoVAE for visual logging in wandb callback.""" from OmniAvatar.models.wan_video_vae import WanVideoVAE raw_vae = WanVideoVAE(z_dim=16) vae_state = torch.load(vae_path, map_location="cpu", weights_only=False) 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() class VAEWrapper: def __init__(self, vae, device): self._vae = vae self._device = device def decode(self, x): with torch.no_grad(): 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") def validation_step(self, data: Dict[str, Any], iteration: int) -> tuple[dict, dict]: """Validation using CausVid's causal AR inference (chunk-by-chunk with KV cache). Uses CausVidModel._student_sample_loop which does proper AR inference: chunk-by-chunk denoising with KV cache updates, matching inference behavior. No teacher, no fake_score — just the student generating video. """ import time from lipforcing.methods.distribution_matching.causvid import CausVidModel t0 = time.time() real_data, condition, neg_condition = self._prepare_training_data(data) B, C, T, H, W = real_data.shape logger.info(f"[val] Starting CausVid AR inference (B={B}, T={T}, steps={self.config.student_sample_steps})") noise = torch.randn_like(real_data) context_noise = getattr(self.config, "context_noise", 0) with torch.no_grad(): gen_data = CausVidModel.generator_fn( net=self.net, 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, ) t_gen = time.time() - t0 logger.info(f"[val] AR inference done in {t_gen:.1f}s") loss_map = {"total_loss": torch.tensor(0.0, device=self.device)} outputs = {"gen_rand": gen_data} # Already generated, not a callable return loss_map, outputs def _prepare_training_data(self, data: Dict[str, Any]) -> tuple[torch.Tensor, Any, Any]: """Build OmniAvatar condition and neg_condition dicts from dataset output. The OmniAvatar dataset returns: real: [B, 16, 21, 64, 64] — clean video latents masked_video: [B, 16, 21, 64, 64] — mouth-masked video latents audio_emb: [B, 81, 10752] — Wav2Vec2 audio features text_embeds: [B, 1, 512, 4096] — T5 text embedding ref_sequence: [B, 16, 21, 64, 64] — reference sequence latents mask: [B, 64, 64] — spatial mask (LatentSync convention: 1=keep, 0=generate) neg_text_embeds: [B, 1, 512, 4096] — negative text embedding Returns: real_data: [B, 16, 21, 64, 64] condition: dict with all V2V conditioning neg_condition: dict with null audio + negative text """ real_data = data["real"] B = real_data.shape[0] # Reference latent: first frame of clean video ref_latent = real_data[:, :, :1, :, :] # [B, 16, 1, H, W] # Spatial mask — use first sample's mask (same across batch) mask = data["mask"] if mask.dim() == 3: # [B, H, W] from DataLoader batching mask = mask[0] # [H, W] — same for all samples # Positive condition condition = { "text_embeds": data["text_embeds"].squeeze(1) if data["text_embeds"].dim() == 4 else data["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"] # Negative condition: null audio, negative text, same spatial conditioning neg_condition = { "text_embeds": data["neg_text_embeds"].squeeze(1) if data["neg_text_embeds"].dim() == 4 else data["neg_text_embeds"], "audio_emb": torch.zeros_like(data["audio_emb"]), "ref_latent": ref_latent, "mask": mask, "masked_video": data["masked_video"], } if "ref_sequence" in data: neg_condition["ref_sequence"] = data["ref_sequence"] return real_data, condition, neg_condition