| 1 | """ |
| 2 | Bidirectional Audio-Video Trajectory Pipeline for DMD backward simulation. |
| 3 | |
| 4 | This pipeline generates denoising trajectories for backward simulation |
| 5 | in DMD training. It runs the generator through multiple denoising steps |
| 6 | and returns the intermediate states. |
| 7 | """ |
| 8 | |
| 9 | from typing import Tuple, Dict, Any, Optional |
| 10 | import torch |
| 11 | import torch.nn as nn |
| 12 | |
| 13 | |
| 14 | class BidirectionalAVTrajectoryPipeline: |
| 15 | """ |
| 16 | Pipeline for generating audio-video denoising trajectories. |
| 17 | |
| 18 | Used in DMD training for backward simulation: |
| 19 | 1. Start from pure noise |
| 20 | 2. Denoise through multiple steps using the generator |
| 21 | 3. Return trajectory of intermediate states |
| 22 | |
| 23 | The trajectory can be used to sample training inputs at different noise levels. |
| 24 | """ |
| 25 | |
| 26 | def __init__( |
| 27 | self, |
| 28 | generator: nn.Module, |
| 29 | add_noise_fn, |
| 30 | denoising_sigmas: torch.Tensor, |
| 31 | ): |
| 32 | """ |
| 33 | Args: |
| 34 | generator: LTX2DiffusionWrapper instance |
| 35 | add_noise_fn: Callable[[original, noise, sigma], noisy_sample] |
| 36 | Flow matching noise addition: (1-sigma)*x0 + sigma*eps |
| 37 | denoising_sigmas: Tensor of sigma values for denoising steps |
| 38 | """ |
| 39 | self.generator = generator |
| 40 | self.add_noise_fn = add_noise_fn |
| 41 | self.denoising_sigmas = denoising_sigmas |
| 42 | |
| 43 | @torch.no_grad() |
| 44 | def inference_with_trajectory( |
| 45 | self, |
| 46 | video_noise: torch.Tensor, |
| 47 | audio_noise: torch.Tensor, |
| 48 | conditional_dict: Dict[str, Any], |
| 49 | ) -> Tuple[torch.Tensor, torch.Tensor]: |
| 50 | """ |
| 51 | Generate denoising trajectory from noise. |
| 52 | |
| 53 | This implements consistency backward simulation: |
| 54 | At each step, predict x0 and then re-corrupt to the next noise level. |
| 55 | |
| 56 | Args: |
| 57 | video_noise: Initial video noise [B, F_v, C, H, W] |
| 58 | audio_noise: Initial audio noise [B, F_a, C] |
| 59 | conditional_dict: Conditioning dictionary |
| 60 | |
| 61 | Returns: |
| 62 | Tuple of: |
| 63 | - video_trajectory: [B, T, F_v, C, H, W] where T is num steps |
| 64 | - audio_trajectory: [B, T, F_a, C] |
| 65 | """ |
| 66 | B = video_noise.shape[0] |
| 67 | F_v = video_noise.shape[1] |
| 68 | F_a = audio_noise.shape[1] |
| 69 | device = video_noise.device |
| 70 | |
| 71 | video_trajectory = [video_noise] |
| 72 | audio_trajectory = [audio_noise] |
| 73 | |
| 74 | noisy_video = video_noise |
| 75 | noisy_audio = audio_noise |
| 76 | |
| 77 | # Iterate through denoising steps (except the last one which is t=0) |
| 78 | for i, sigma in enumerate(self.denoising_sigmas[:-1]): |
| 79 | # Prepare sigma tensors |
| 80 | video_sigma = sigma * torch.ones([B, F_v], device=device) |
| 81 | audio_sigma = sigma * torch.ones([B, F_a], device=device) |
| 82 | |
| 83 | # Predict x0 |
| 84 | pred_video, pred_audio = self.generator( |
| 85 | noisy_image_or_video=noisy_video, |
| 86 | conditional_dict=conditional_dict, |
| 87 | timestep=video_sigma, |
| 88 | noisy_audio=noisy_audio, |
| 89 | audio_timestep=audio_sigma, |
| 90 | ) |
| 91 | |
| 92 | # Get next sigma |
| 93 | next_sigma = self.denoising_sigmas[i + 1] |
| 94 | |
| 95 | if next_sigma > 0: |
| 96 | # Re-corrupt with next sigma level |
| 97 | # For flow matching: x_t = (1 - sigma) * x_0 + sigma * eps |
| 98 | # We need to add noise at the next sigma level |
| 99 | |
| 100 | # Sample fresh noise |
| 101 | fresh_noise_video = torch.randn_like(video_noise) |
| 102 | fresh_noise_audio = torch.randn_like(audio_noise) |
| 103 | |
| 104 | next_video_sigma = next_sigma * torch.ones([B, F_v], device=device) |
| 105 | next_audio_sigma = next_sigma * torch.ones([B, F_a], device=device) |
| 106 | |
| 107 | noisy_video = self.add_noise_fn( |
| 108 | pred_video.flatten(0, 1), |
| 109 | fresh_noise_video.flatten(0, 1), |
| 110 | next_video_sigma.flatten(0, 1), |
| 111 | ).unflatten(0, (B, F_v)) |
| 112 | |
| 113 | noisy_audio = self.add_noise_fn( |
| 114 | pred_audio, fresh_noise_audio, next_audio_sigma |
| 115 | ) |
| 116 | else: |
| 117 | # At t=0, just use the prediction |
| 118 | noisy_video = pred_video |
| 119 | noisy_audio = pred_audio |
| 120 | |
| 121 | video_trajectory.append(noisy_video) |
| 122 | audio_trajectory.append(noisy_audio) |
| 123 | |
| 124 | # Stack trajectories: [B, T, F, C, H, W] |
| 125 | video_trajectory = torch.stack(video_trajectory, dim=1) |
| 126 | audio_trajectory = torch.stack(audio_trajectory, dim=1) |
| 127 | |
| 128 | return video_trajectory, audio_trajectory |
| 129 | |
| 130 | |
| 131 | class BidirectionalAVInferencePipeline: |
| 132 | """ |
| 133 | Pipeline for few-step bidirectional inference. |
| 134 | |
| 135 | Used for validation after training to generate videos/audio |
| 136 | using the distilled model. |
| 137 | """ |
| 138 | |
| 139 | def __init__( |
| 140 | self, |
| 141 | generator: nn.Module, |
| 142 | add_noise_fn, |
| 143 | denoising_sigmas: torch.Tensor, |
| 144 | ): |
| 145 | """ |
| 146 | Args: |
| 147 | generator: Distilled LTX2DiffusionWrapper |
| 148 | add_noise_fn: Callable[[original, noise, sigma], noisy_sample] |
| 149 | denoising_sigmas: Sigma values for few-step denoising |
| 150 | """ |
| 151 | self.generator = generator |
| 152 | self.add_noise_fn = add_noise_fn |
| 153 | self.denoising_sigmas = denoising_sigmas |
| 154 | |
| 155 | @torch.no_grad() |
| 156 | def generate( |
| 157 | self, |
| 158 | video_shape: Tuple[int, ...], |
| 159 | audio_shape: Tuple[int, ...], |
| 160 | conditional_dict: Dict[str, Any], |
| 161 | seed: Optional[int] = None, |
| 162 | ) -> Tuple[torch.Tensor, torch.Tensor]: |
| 163 | """ |
| 164 | Generate video and audio using few-step denoising. |
| 165 | |
| 166 | Args: |
| 167 | video_shape: (B, F_v, C, H, W) video latent shape |
| 168 | audio_shape: (B, F_a, C) audio latent shape |
| 169 | conditional_dict: Text conditioning |
| 170 | seed: Random seed (optional) |
| 171 | |
| 172 | Returns: |
| 173 | Tuple of (video_latent, audio_latent) |
| 174 | """ |
| 175 | B = video_shape[0] |
| 176 | F_v = video_shape[1] |
| 177 | F_a = audio_shape[1] |
| 178 | |
| 179 | # Set seed if provided |
| 180 | if seed is not None: |
| 181 | torch.manual_seed(seed) |
| 182 | |
| 183 | device = next(self.generator.parameters()).device |
| 184 | dtype = next(self.generator.parameters()).dtype |
| 185 | |
| 186 | # Initialize with noise |
| 187 | video = torch.randn(video_shape, device=device, dtype=dtype) |
| 188 | audio = torch.randn(audio_shape, device=device, dtype=dtype) |
| 189 | |
| 190 | # Few-step denoising |
| 191 | for i, sigma in enumerate(self.denoising_sigmas[:-1]): |
| 192 | video_sigma = sigma * torch.ones([B, F_v], device=device) |
| 193 | audio_sigma = sigma * torch.ones([B, F_a], device=device) |
| 194 | |
| 195 | # Predict x0 |
| 196 | pred_video, pred_audio = self.generator( |
| 197 | noisy_image_or_video=video, |
| 198 | conditional_dict=conditional_dict, |
| 199 | timestep=video_sigma, |
| 200 | noisy_audio=audio, |
| 201 | audio_timestep=audio_sigma, |
| 202 | ) |
| 203 | |
| 204 | # Get next sigma |
| 205 | next_sigma = self.denoising_sigmas[i + 1] |
| 206 | |
| 207 | if next_sigma > 0: |
| 208 | # Euler step or re-corruption |
| 209 | fresh_noise_video = torch.randn_like(video) |
| 210 | fresh_noise_audio = torch.randn_like(audio) |
| 211 | |
| 212 | next_video_sigma = next_sigma * torch.ones([B, F_v], device=device) |
| 213 | next_audio_sigma = next_sigma * torch.ones([B, F_a], device=device) |
| 214 | |
| 215 | video = self.add_noise_fn( |
| 216 | pred_video.flatten(0, 1), |
| 217 | fresh_noise_video.flatten(0, 1), |
| 218 | next_video_sigma.flatten(0, 1), |
| 219 | ).unflatten(0, (B, F_v)) |
| 220 | |
| 221 | audio = self.add_noise_fn( |
| 222 | pred_audio, fresh_noise_audio, next_audio_sigma |
| 223 | ) |
| 224 | else: |
| 225 | video = pred_video |
| 226 | audio = pred_audio |
| 227 | |
| 228 | return video, audio |
| 229 |