# 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 Any, TYPE_CHECKING, Optional, Dict, Callable import torch from lipforcing.methods import DMD2Model from lipforcing.utils import basic_utils import lipforcing.utils.logging_utils as logger if TYPE_CHECKING: from lipforcing.networks.network import CausalFastGenNetwork class CausVidModel(DMD2Model): """CausVid implementation""" def _generate_noise_and_time( self, real_data: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Generate random noises and time step Args: real_data: Real data tensor of shape [B, C, T, H, W] Returns: noisy_real_data: Random noise used by the student t_inhom: Inhomogeneous time steps used by the student [B, T] for causal networks t: Homogeneous time step [B] for teacher network eps: Random noise used by a forward process """ assert real_data.ndim == 5, "CausVid only works for video data" batch_size, num_frames = real_data.shape[0], real_data.shape[2] assert hasattr(self.net, "chunk_size"), "net does not have the chunk_size attribute" chunk_size = self.net.chunk_size # Add noise to real image data (for multistep generation) eps_inhom = torch.randn(batch_size, *self.input_shape, device=self.device, dtype=real_data.dtype) assert hasattr( self.net.noise_scheduler, "sample_t_inhom" ), "net.noise_scheduler does not have the sample_t_inhom() method" 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, ) # shape [B, T] t_inhom_expanded = t_inhom[:, None, :, None, None] # shape [B, 1, T, 1, 1] noisy_real_data = self.net.noise_scheduler.forward_process(real_data, eps_inhom, t_inhom_expanded) t = self.net.noise_scheduler.sample_t( batch_size, **basic_utils.convert_cfg_to_dict(self.config.sample_t_cfg), device=self.device, ) eps = torch.randn_like(eps_inhom, device=self.device, dtype=real_data.dtype) return noisy_real_data, t_inhom, t, eps 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) 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, context_noise=getattr(self.config, "context_noise", 0), # Optional context noise ) return {"gen_rand": gen_rand_func, "input_rand": noise, "gen_rand_train": gen_data} @classmethod def _student_sample_loop( cls, net: CausalFastGenNetwork, x: torch.Tensor, t_list: torch.Tensor, condition: Any = None, student_sample_type: str = "sde", context_noise: Optional[float] = 0, **kwargs, ) -> torch.Tensor: """ Sample loop for the student network. Args: net: The FastGenNetwork network x: The latents to start from t_list: Timesteps to sample condition: Optional conditioning information student_sample_type: Type of student multistep sampling Returns: The sampled data """ logger.debug("Using generator_fn in CausVidModel") # cleanup caches before sampling net.clear_caches() batch_size, num_frames = x.shape[0], x.shape[2] chunk_size = net.chunk_size num_chunks = num_frames // chunk_size remaining_size = num_frames % chunk_size # initialize all noise using the first timestep for i in range(max(1, num_chunks)): if num_chunks == 0: # Handle case where num_frames < chunk_size start, end = 0, remaining_size else: # Normal chunking logic start = 0 if i == 0 else chunk_size * i + remaining_size end = chunk_size * (i + 1) + remaining_size x_next = x[:, :, start:end, ...] for step in range(len(t_list) - 1): # denoise t_cur = t_list[step].expand(batch_size) x_cur = x_next x_next = net( x_cur, t_cur, condition=condition, fwd_pred_type="x0", cache_tag="pos", cur_start_frame=start, store_kv=False, is_ar=True, **kwargs, ) # update to the next timestep for forward process t_next = t_list[step + 1] if t_next > 0: t_chunk_next = t_next.expand(batch_size) if student_sample_type == "sde": eps_infer = torch.randn_like(x_next) elif student_sample_type == "ode": eps_infer = net.noise_scheduler.x0_to_eps(xt=x_cur, x0=x_next, t=t_cur) else: raise NotImplementedError( f"student_sample_type must be one of 'sde', 'ode' but got {student_sample_type}" ) x_next = net.noise_scheduler.forward_process(x_next, eps_infer, t_chunk_next) x[:, :, start:end, ...] = x_next # compute and update the KV cache x_cache = x_next t_cache = t_list[-1].expand(batch_size) if context_noise > 0: # Add context noise to denoised frames before caching t_cache = torch.full((batch_size,), context_noise, device=x.device, dtype=x.dtype) x_cache = net.noise_scheduler.forward_process(x_next, torch.randn_like(x_next), t_cache) _ = net( x_cache, t_cache, condition=condition, fwd_pred_type="x0", cache_tag="pos", cur_start_frame=start, store_kv=True, is_ar=True, **kwargs, ) # cleanup caches after full sampling net.clear_caches() return x @classmethod def generator_fn_extrapolation( cls, net: CausalFastGenNetwork, noise: torch.Tensor, condition: Any = None, *, num_segments: int, overlap_frames: int, student_sample_steps: int = 1, student_sample_type: str = "sde", t_list: Optional[torch.Tensor] = None, precision_amp: Optional[torch.dtype] = None, context_noise: Optional[float] = 0, **kwargs, ) -> torch.Tensor: """ Autoregressively generate multiple segments using the student generator_fn stepping, with optional frame-overlap bridging via a VAE. Args: net: The student causal network. noise: Initial latents for a single segment [B, C, T, H, W]. condition: Optional conditioning tensor. num_segments: Number of segments to autoregressively generate (>= 1). overlap_frames: Number of frames to overlap/bridge across segments. Must be divisible by chunk_size. student_sample_steps: Number of denoising steps used by generator_fn. student_sample_type: One of {"sde", "ode"}. t_list: Optional custom t_list; if None, derived from scheduler and student_sample_steps. precision_amp (Optional[torch.dtype]): If not None, uses precision_amp with this dtype for inference. context_noise: Optional context noise scale in [0, 1] for cache prefill. **kwargs: Passed through to network forward calls. Returns: The concatenated video latents across all segments [B, C, num_segments*T - (num_segments-1)*overlap_frames, H, W]. """ logger.debug("Using generator_fn_extrapolation in CausVidModel") with basic_utils.inference_mode(net, precision_amp=precision_amp, device_type=noise.device.type): if num_segments < 1: raise ValueError("num_segments must be >= 1") if overlap_frames > 0 and net.vae is None: raise ValueError("generator_fn_extrapolation requires a VAE instance via `vae` when overlap_frames > 0") batch_size, channels, segment_frames, height, width = noise.shape dtype = noise.dtype device = noise.device chunk_size = net.chunk_size if segment_frames % chunk_size != 0: raise ValueError(f"Segment length {segment_frames} must be divisible by chunk_size {chunk_size}") if overlap_frames < 0 or overlap_frames >= segment_frames: raise ValueError("overlap_frames must be in [0, segment_frames)") if overlap_frames % chunk_size != 0: raise ValueError("overlap_frames must be divisible by chunk_size") # Prepare t_list consistent with generator_fn if t_list is None: t_list = net.noise_scheduler.get_t_list(student_sample_steps, device=device, dtype=torch.float32) else: assert ( len(t_list) - 1 == student_sample_steps ), f"t_list length (excluding zero) != student_sample_steps: {len(t_list) - 1} != {student_sample_steps}" t_list = torch.tensor(t_list, device=device, dtype=torch.float32) assert t_list[-1].item() == 0, "t_list[-1] must be zero" def _prefill_caches(segment_latents: torch.Tensor, frames: int, frame_offset: int = 0) -> None: """Prefill KV caches with overlapping frames from the previous segment. Args: segment_latents: Latents to prefill caches with [B, C, T, H, W] frames: Number of frames to prefill frame_offset: Segment-level offset for extrapolation. For segment N, this is N * (segment_frames - overlap_frames). """ if frames == 0: return start_frame = 0 t_zero = t_list[-1].expand(batch_size) # zero timestep while start_frame < frames: end_frame = min(start_frame + chunk_size, frames) slice_latents = segment_latents[:, :, start_frame:end_frame, ...] _ = net( slice_latents, t_zero, condition=condition, fwd_pred_type="x0", cache_tag="pos", cur_start_frame=start_frame, frame_offset=frame_offset, store_kv=True, is_ar=True, **kwargs, ) start_frame = end_frame def _run_segment(segment_latents: torch.Tensor, prefill_frames: int, frame_offset: int = 0) -> torch.Tensor: """Run a single segment of autoregressive generation. Args: segment_latents: Input latents for this segment [B, C, T, H, W] prefill_frames: Number of frames to prefill from previous segment (for overlap) frame_offset: Segment-level offset for extrapolation. The network computes global frame index as: frame_offset + cur_start_frame. For segment N, this is N * (segment_frames - overlap_frames). """ # Clone to avoid in-place modifications on the input tensor x = segment_latents.clone() # Clear caches before processing a new segment net.clear_caches() # If we have overlapping frames from the previous segment, prefill caches for them # The prefill frames use the same frame_offset since they are the OVERLAP # from the previous segment (their global indices are frame_offset + 0, 1, ...) if prefill_frames > 0: _prefill_caches(x, prefill_frames, frame_offset) # Initialize only the frames we are about to generate using the first timestep sigma if prefill_frames == 0: x = net.noise_scheduler.latents(x, t_init=t_list[0]) else: x[:, :, prefill_frames:, ...] = net.noise_scheduler.latents( x[:, :, prefill_frames:, ...], t_init=t_list[0] ) start_frame = prefill_frames while start_frame < segment_frames: end_frame = min(start_frame + chunk_size, segment_frames) x_next = x[:, :, start_frame:end_frame, ...] for step in range(len(t_list) - 1): # Denoise to x0 using the student network t_cur = t_list[step].expand(batch_size) x_cur = x_next x_next = net( x_cur, t_cur, condition=condition, fwd_pred_type="x0", cache_tag="pos", cur_start_frame=start_frame, frame_offset=frame_offset, store_kv=False, is_ar=True, **kwargs, ) # Move forward in the forward process if not at the final step t_next = t_list[step + 1] if t_next > 0: t_chunk_next = t_next.expand(batch_size) if student_sample_type == "sde": eps_infer = torch.randn_like(x_next) elif student_sample_type == "ode": eps_infer = net.noise_scheduler.x0_to_eps(xt=x_cur, x0=x_next, t=t_cur) else: raise NotImplementedError( f"student_sample_type must be one of 'sde', 'ode' but got {student_sample_type}" ) x_next = net.noise_scheduler.forward_process(x_next, eps_infer, t_chunk_next) # Write the generated slice back x[:, :, start_frame:end_frame, ...] = x_next # Update KV caches with the denoised slice (optionally with context noise) x_cache = x_next t_cache = t_list[-1].expand(batch_size) if context_noise and context_noise > 0: t_cache = torch.full((batch_size,), context_noise, device=device, dtype=dtype) x_cache = net.noise_scheduler.forward_process(x_next, torch.randn_like(x_next), t_cache) _ = net( x_cache, t_cache, condition=condition, fwd_pred_type="x0", cache_tag="pos", cur_start_frame=start_frame, frame_offset=frame_offset, store_kv=True, is_ar=True, **kwargs, ) start_frame = end_frame # Clean up caches after finishing the segment net.clear_caches() return x segments = [] current_latents = noise prefill_frames = 0 for segment_idx in range(num_segments): # Compute the global frame offset for depth/control conditioning # For segment N, the depth conditioning starts at frame N * (segment_frames - overlap_frames) # This ensures: # - Segment 0: local frame i → depth frame i # - Segment 1 with overlap: local frame 0 (overlap from seg 0's tail) → depth frame (seg_frames - overlap) # local frame overlap → depth frame seg_frames # - And so on... frame_offset = segment_idx * (segment_frames - overlap_frames) segment_latents = _run_segment(current_latents, prefill_frames, frame_offset) if segment_idx == 0: segments.append(segment_latents) else: if overlap_frames > 0: segments.append(segment_latents[:, :, overlap_frames:, :, :]) else: segments.append(segment_latents) if segment_idx == num_segments - 1: break # Prepare latents for the next segment if overlap_frames == 0: current_latents = torch.randn_like(noise) prefill_frames = 0 continue # Bridge with VAE: take the last overlap frames from current segment (pixels), encode back to latents decoded_video = net.vae.decode(segment_latents) tail_pixels = decoded_video[:, :, -overlap_frames:, :, :] encoded_tail = net.vae.encode(tail_pixels).to(dtype=dtype, device=device) # Reuse all but the first overlapped latent directly to avoid unnecessary encode/decode if overlap_frames > 1: reused_tail = segment_latents[:, :, -(overlap_frames - 1) :, :, :] encoded_tail = torch.cat([encoded_tail[:, :, :1, :, :], reused_tail], dim=2) # Compose the next segment latents with bridged head and random remainder next_latents = torch.randn_like(segment_latents) next_latents[:, :, :overlap_frames, :, :] = encoded_tail current_latents = next_latents prefill_frames = overlap_frames # Final cleanup and concatenate along the temporal dimension net.clear_caches() return torch.cat(segments, dim=2).to(dtype=noise.dtype)