| 1 | """Shared utilities for inference: latent computation, noise, media I/O, video concat.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import shutil |
| 6 | import subprocess |
| 7 | import tempfile |
| 8 | from pathlib import Path |
| 9 | from typing import Any, Optional |
| 10 | |
| 11 | import torch |
| 12 | import torchaudio |
| 13 | |
| 14 | from ltx_distillation.inference.memory_multishot import ( |
| 15 | audio_waveform_stats, |
| 16 | normalize_audio_waveform_for_media, |
| 17 | ) |
| 18 | |
| 19 | |
| 20 | def _write_video(*args, **kwargs) -> None: |
| 21 | """Import the pinned torchvision video backend only when media is written.""" |
| 22 | |
| 23 | from torchvision.io import write_video |
| 24 | |
| 25 | write_video(*args, **kwargs) |
| 26 | |
| 27 | |
| 28 | def compute_latent_shapes( |
| 29 | *, |
| 30 | num_frames: int, |
| 31 | video_height: int, |
| 32 | video_width: int, |
| 33 | batch_size: int = 1, |
| 34 | latent_channels: int = 128, |
| 35 | vae_temporal_compression: int = 8, |
| 36 | vae_spatial_compression: int = 32, |
| 37 | video_fps: float = 24.0, |
| 38 | audio_sample_rate: int = 16000, |
| 39 | audio_hop_length: int = 160, |
| 40 | audio_latent_downsample: int = 4, |
| 41 | ) -> tuple[list[int], list[int]]: |
| 42 | if (num_frames - 1) % vae_temporal_compression != 0: |
| 43 | raise ValueError(f"num_frames must be 1 + 8*k, got {num_frames}") |
| 44 | |
| 45 | latent_frames = 1 + (num_frames - 1) // vae_temporal_compression |
| 46 | latent_h = video_height // vae_spatial_compression |
| 47 | latent_w = video_width // vae_spatial_compression |
| 48 | |
| 49 | video_duration = float(num_frames) / float(video_fps) |
| 50 | audio_latent_fps = ( |
| 51 | float(audio_sample_rate) |
| 52 | / float(audio_hop_length) |
| 53 | / float(audio_latent_downsample) |
| 54 | ) |
| 55 | audio_frames = round(video_duration * audio_latent_fps) |
| 56 | |
| 57 | return ( |
| 58 | [batch_size, latent_frames, latent_channels, latent_h, latent_w], |
| 59 | [batch_size, audio_frames, latent_channels], |
| 60 | ) |
| 61 | |
| 62 | |
| 63 | def add_noise( |
| 64 | original: torch.Tensor, noise: torch.Tensor, sigma: torch.Tensor |
| 65 | ) -> torch.Tensor: |
| 66 | sigma = sigma.to(device=original.device, dtype=original.dtype) |
| 67 | if sigma.dim() == 1: |
| 68 | sigma = sigma.reshape(-1, *[1] * (original.dim() - 1)) |
| 69 | elif sigma.dim() == 2: |
| 70 | sigma = sigma.reshape(*sigma.shape, *[1] * (original.dim() - 2)) |
| 71 | return (1 - sigma) * original + sigma * noise |
| 72 | |
| 73 | |
| 74 | @torch.no_grad() |
| 75 | def decode_generated_sample( |
| 76 | video_vae, |
| 77 | audio_vae, |
| 78 | video_latent, |
| 79 | audio_latent, |
| 80 | *, |
| 81 | video_tiling_config=None, |
| 82 | ): |
| 83 | if video_tiling_config is None: |
| 84 | video_pixel = video_vae.decode_to_pixel(video_latent) |
| 85 | video_uint8 = video_pixel[0] |
| 86 | if video_uint8.shape[0] == 3: |
| 87 | video_uint8 = video_uint8.permute(1, 0, 2, 3) |
| 88 | video_uint8 = video_uint8.permute(0, 2, 3, 1) |
| 89 | video_uint8 = (video_uint8.clamp(0, 1) * 255).cpu().to(torch.uint8).contiguous() |
| 90 | else: |
| 91 | video_chunks = list( |
| 92 | video_vae.decode_to_uint8_chunks(video_latent, video_tiling_config) |
| 93 | ) |
| 94 | if not video_chunks: |
| 95 | raise RuntimeError("tiled video VAE decode produced no frames") |
| 96 | video_uint8 = torch.cat(video_chunks, dim=0) |
| 97 | |
| 98 | audio_waveform = ( |
| 99 | audio_vae.decode_to_waveform(audio_latent) if audio_latent is not None else None |
| 100 | ) |
| 101 | |
| 102 | audio_float = normalize_audio_waveform_for_media(audio_waveform) |
| 103 | return video_uint8, audio_float |
| 104 | |
| 105 | |
| 106 | def fit_audio_to_video_frames( |
| 107 | audio: torch.Tensor, |
| 108 | *, |
| 109 | sample_rate: int, |
| 110 | video_frames: int, |
| 111 | video_fps: float, |
| 112 | ) -> torch.Tensor: |
| 113 | """Trim or zero-pad audio to the exact encoded video-frame duration.""" |
| 114 | |
| 115 | if sample_rate <= 0 or video_frames <= 0 or video_fps <= 0: |
| 116 | return audio |
| 117 | target_samples = max(1, round(video_frames * sample_rate / video_fps)) |
| 118 | current_samples = int(audio.shape[-1]) |
| 119 | if current_samples > target_samples: |
| 120 | return audio[..., :target_samples].contiguous() |
| 121 | if current_samples < target_samples: |
| 122 | return torch.nn.functional.pad(audio, (0, target_samples - current_samples)) |
| 123 | return audio |
| 124 | |
| 125 | |
| 126 | def _write_video_with_aligned_audio( |
| 127 | *, |
| 128 | video_uint8: torch.Tensor, |
| 129 | output_path: Path, |
| 130 | audio_path: Path, |
| 131 | fps: float, |
| 132 | ) -> None: |
| 133 | """Mux aligned audio without ``-shortest``, matching production.""" |
| 134 | |
| 135 | frame_count = int(video_uint8.shape[0]) |
| 136 | if frame_count <= 0: |
| 137 | raise ValueError("cannot write an empty video") |
| 138 | if fps <= 0: |
| 139 | raise ValueError(f"invalid frame rate: {fps}") |
| 140 | |
| 141 | duration_arg = f"{frame_count / float(fps):.9f}" |
| 142 | silent_path = output_path.with_name( |
| 143 | f"{output_path.stem}_silent{output_path.suffix}" |
| 144 | ) |
| 145 | _write_video(str(silent_path), video_uint8, fps=int(fps)) |
| 146 | |
| 147 | ffmpeg = shutil.which("ffmpeg") |
| 148 | if ffmpeg is None: |
| 149 | silent_path.unlink(missing_ok=True) |
| 150 | raise RuntimeError("ffmpeg not found") |
| 151 | |
| 152 | command = [ |
| 153 | ffmpeg, |
| 154 | "-y", |
| 155 | "-i", |
| 156 | str(silent_path), |
| 157 | "-i", |
| 158 | str(audio_path), |
| 159 | "-map", |
| 160 | "0:v:0", |
| 161 | "-map", |
| 162 | "1:a:0", |
| 163 | "-c:v", |
| 164 | "copy", |
| 165 | "-c:a", |
| 166 | "aac", |
| 167 | "-af", |
| 168 | f"apad,atrim=duration={duration_arg},asetpts=N/SR/TB", |
| 169 | "-t", |
| 170 | duration_arg, |
| 171 | str(output_path), |
| 172 | ] |
| 173 | try: |
| 174 | subprocess.run( |
| 175 | command, |
| 176 | check=True, |
| 177 | capture_output=True, |
| 178 | text=True, |
| 179 | timeout=300, |
| 180 | ) |
| 181 | finally: |
| 182 | silent_path.unlink(missing_ok=True) |
| 183 | |
| 184 | |
| 185 | def write_generated_media( |
| 186 | *, |
| 187 | output_path: Path, |
| 188 | video_uint8: torch.Tensor, |
| 189 | audio_waveform: Optional[torch.Tensor], |
| 190 | fps: int, |
| 191 | audio_sr: int, |
| 192 | ) -> dict[str, Any]: |
| 193 | output_path.parent.mkdir(parents=True, exist_ok=True) |
| 194 | audio_waveform = normalize_audio_waveform_for_media(audio_waveform) |
| 195 | if audio_waveform is not None: |
| 196 | audio_waveform = fit_audio_to_video_frames( |
| 197 | audio_waveform, |
| 198 | sample_rate=audio_sr, |
| 199 | video_frames=int(video_uint8.shape[0]), |
| 200 | video_fps=float(fps), |
| 201 | ) |
| 202 | stats = audio_waveform_stats(audio_waveform) |
| 203 | |
| 204 | wrote_with_audio = False |
| 205 | wrote_sidecar_wav = False |
| 206 | if audio_waveform is not None: |
| 207 | audio_path = output_path.with_suffix(".wav") |
| 208 | torchaudio.save(str(audio_path), audio_waveform, audio_sr) |
| 209 | wrote_sidecar_wav = True |
| 210 | try: |
| 211 | _write_video_with_aligned_audio( |
| 212 | video_uint8=video_uint8, |
| 213 | output_path=output_path, |
| 214 | audio_path=audio_path, |
| 215 | fps=float(fps), |
| 216 | ) |
| 217 | wrote_with_audio = True |
| 218 | except Exception as exc: |
| 219 | print( |
| 220 | f"[warn] aligned audio mux failed for {output_path}: {exc}; " |
| 221 | "retrying direct write_video", |
| 222 | flush=True, |
| 223 | ) |
| 224 | try: |
| 225 | _write_video( |
| 226 | str(output_path), |
| 227 | video_uint8, |
| 228 | fps=fps, |
| 229 | audio_array=audio_waveform, |
| 230 | audio_fps=audio_sr, |
| 231 | audio_codec="aac", |
| 232 | ) |
| 233 | wrote_with_audio = True |
| 234 | except Exception as fallback_exc: |
| 235 | print( |
| 236 | f"[warn] direct audio mux failed for {output_path}: " |
| 237 | f"{fallback_exc}; writing video only; audio_stats={stats}", |
| 238 | flush=True, |
| 239 | ) |
| 240 | |
| 241 | if not wrote_with_audio: |
| 242 | _write_video(str(output_path), video_uint8, fps=fps) |
| 243 | |
| 244 | return { |
| 245 | "wrote_audio_in_mp4": wrote_with_audio, |
| 246 | "wrote_sidecar_wav": wrote_sidecar_wav, |
| 247 | "audio_stats": stats, |
| 248 | } |
| 249 | |
| 250 | |
| 251 | def concat_shot_videos(shot_paths: list[Path], output_path: Path) -> None: |
| 252 | if not shot_paths: |
| 253 | raise ValueError("No shot videos provided for concatenation") |
| 254 | |
| 255 | output_path.parent.mkdir(parents=True, exist_ok=True) |
| 256 | with tempfile.NamedTemporaryFile( |
| 257 | "w", suffix=".txt", delete=False, encoding="utf-8" |
| 258 | ) as fp: |
| 259 | concat_file = Path(fp.name) |
| 260 | for shot_path in shot_paths: |
| 261 | fp.write(f"file '{shot_path.resolve().as_posix()}'\n") |
| 262 | |
| 263 | try: |
| 264 | cmd = [ |
| 265 | "ffmpeg", |
| 266 | "-y", |
| 267 | "-f", |
| 268 | "concat", |
| 269 | "-safe", |
| 270 | "0", |
| 271 | "-i", |
| 272 | str(concat_file), |
| 273 | "-c", |
| 274 | "copy", |
| 275 | str(output_path), |
| 276 | ] |
| 277 | result = subprocess.run(cmd, capture_output=True, text=True) |
| 278 | if result.returncode != 0: |
| 279 | fallback_cmd = [ |
| 280 | "ffmpeg", |
| 281 | "-y", |
| 282 | "-f", |
| 283 | "concat", |
| 284 | "-safe", |
| 285 | "0", |
| 286 | "-i", |
| 287 | str(concat_file), |
| 288 | "-c:v", |
| 289 | "libx264", |
| 290 | "-preset", |
| 291 | "medium", |
| 292 | "-crf", |
| 293 | "18", |
| 294 | "-c:a", |
| 295 | "aac", |
| 296 | "-b:a", |
| 297 | "192k", |
| 298 | str(output_path), |
| 299 | ] |
| 300 | fallback_result = subprocess.run( |
| 301 | fallback_cmd, capture_output=True, text=True |
| 302 | ) |
| 303 | if fallback_result.returncode != 0: |
| 304 | raise RuntimeError( |
| 305 | "Failed to concatenate shot videos with ffmpeg.\n" |
| 306 | f"copy stderr:\n{result.stderr}\n" |
| 307 | f"reencode stderr:\n{fallback_result.stderr}" |
| 308 | ) |
| 309 | finally: |
| 310 | concat_file.unlink(missing_ok=True) |
| 311 | |
| 312 | |
| 313 | def concat_shot_audios(audios: list[torch.Tensor]) -> Optional[torch.Tensor]: |
| 314 | if not audios: |
| 315 | return None |
| 316 | audio = audios[0] |
| 317 | if audio.ndim == 1: |
| 318 | sample_dim = 0 |
| 319 | elif audio.ndim == 2: |
| 320 | sample_dim = 1 if audio.shape[0] <= audio.shape[1] else 0 |
| 321 | else: |
| 322 | raise ValueError( |
| 323 | f"Expected audio tensor with 1 or 2 dims, got shape={tuple(audio.shape)}" |
| 324 | ) |
| 325 | return torch.cat([a.contiguous() for a in audios], dim=sample_dim).contiguous() |
| 326 |