| 1 | """Shared utilities for inference: latent computation, noise, media I/O, video concat.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import subprocess |
| 6 | import tempfile |
| 7 | from pathlib import Path |
| 8 | from typing import Any, Optional |
| 9 | |
| 10 | import torch |
| 11 | import torchaudio |
| 12 | from torchvision.io import write_video |
| 13 | from torchvision.transforms import functional as TVF |
| 14 | |
| 15 | from ltx_distillation.inference.memory_multishot import ( |
| 16 | audio_waveform_stats, |
| 17 | normalize_audio_waveform_for_media, |
| 18 | ) |
| 19 | |
| 20 | |
| 21 | def compute_latent_shapes( |
| 22 | *, |
| 23 | num_frames: int, |
| 24 | video_height: int, |
| 25 | video_width: int, |
| 26 | batch_size: int = 1, |
| 27 | latent_channels: int = 128, |
| 28 | vae_temporal_compression: int = 8, |
| 29 | vae_spatial_compression: int = 32, |
| 30 | video_fps: float = 24.0, |
| 31 | audio_sample_rate: int = 16000, |
| 32 | audio_hop_length: int = 160, |
| 33 | audio_latent_downsample: int = 4, |
| 34 | ) -> tuple[list[int], list[int]]: |
| 35 | if (num_frames - 1) % vae_temporal_compression != 0: |
| 36 | raise ValueError(f"num_frames must be 1 + 8*k, got {num_frames}") |
| 37 | |
| 38 | latent_frames = 1 + (num_frames - 1) // vae_temporal_compression |
| 39 | latent_h = video_height // vae_spatial_compression |
| 40 | latent_w = video_width // vae_spatial_compression |
| 41 | |
| 42 | video_duration = float(num_frames) / float(video_fps) |
| 43 | audio_latent_fps = float(audio_sample_rate) / float(audio_hop_length) / float(audio_latent_downsample) |
| 44 | audio_frames = round(video_duration * audio_latent_fps) |
| 45 | |
| 46 | return ( |
| 47 | [batch_size, latent_frames, latent_channels, latent_h, latent_w], |
| 48 | [batch_size, audio_frames, latent_channels], |
| 49 | ) |
| 50 | |
| 51 | |
| 52 | def add_noise(original: torch.Tensor, noise: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: |
| 53 | sigma = sigma.to(device=original.device, dtype=original.dtype) |
| 54 | if sigma.dim() == 1: |
| 55 | sigma = sigma.reshape(-1, *[1] * (original.dim() - 1)) |
| 56 | elif sigma.dim() == 2: |
| 57 | sigma = sigma.reshape(*sigma.shape, *[1] * (original.dim() - 2)) |
| 58 | return (1 - sigma) * original + sigma * noise |
| 59 | |
| 60 | |
| 61 | def frames_to_video_tensor(frames, target_h: int, target_w: int) -> torch.Tensor: |
| 62 | tensors = [] |
| 63 | for idx, image in enumerate(frames): |
| 64 | if image.size != (target_w, target_h): |
| 65 | raise ValueError( |
| 66 | f"Frame size mismatch at index {idx}: got={image.size}, expected={(target_w, target_h)}" |
| 67 | ) |
| 68 | tensor = TVF.to_tensor(image) |
| 69 | tensors.append(tensor * 2.0 - 1.0) |
| 70 | return torch.stack(tensors, dim=1).contiguous() |
| 71 | |
| 72 | |
| 73 | @torch.no_grad() |
| 74 | def encode_memory_frames_batch( |
| 75 | *, |
| 76 | video_vae, |
| 77 | batch_memory_frames, |
| 78 | target_h: int, |
| 79 | target_w: int, |
| 80 | device: torch.device, |
| 81 | dtype: torch.dtype, |
| 82 | ) -> torch.Tensor: |
| 83 | if getattr(video_vae, "encoder", None) is None: |
| 84 | raise RuntimeError("video VAE encoder is not initialized for memory encoding") |
| 85 | |
| 86 | latents = [] |
| 87 | for memory_frames in batch_memory_frames: |
| 88 | if not memory_frames: |
| 89 | raise ValueError("memory_frames cannot be empty when encoding memory video") |
| 90 | per_frame_latents = [] |
| 91 | for memory_item in memory_frames: |
| 92 | is_clip_memory = isinstance(memory_item, list) |
| 93 | frame_video = frames_to_video_tensor( |
| 94 | memory_item if is_clip_memory else [memory_item], |
| 95 | target_h, |
| 96 | target_w, |
| 97 | ).unsqueeze(0).to(device=device, dtype=dtype) |
| 98 | latent = video_vae.encode(frame_video) |
| 99 | del frame_video |
| 100 | latent = latent.permute(0, 2, 1, 3, 4).to(dtype=dtype) |
| 101 | if is_clip_memory: |
| 102 | latent = latent[:, -1:, :, :, :].contiguous() |
| 103 | per_frame_latents.append(latent) |
| 104 | latents.append(torch.cat(per_frame_latents, dim=1)) |
| 105 | del per_frame_latents |
| 106 | return torch.cat(latents, dim=0) |
| 107 | |
| 108 | |
| 109 | @torch.no_grad() |
| 110 | def decode_benchmark_sample(video_vae, audio_vae, video_latent, audio_latent): |
| 111 | video_pixel = video_vae.decode_to_pixel(video_latent) |
| 112 | audio_waveform = audio_vae.decode_to_waveform(audio_latent) if audio_latent is not None else None |
| 113 | |
| 114 | video_uint8 = video_pixel[0] |
| 115 | if video_uint8.shape[0] == 3: |
| 116 | video_uint8 = video_uint8.permute(1, 0, 2, 3) |
| 117 | video_uint8 = video_uint8.permute(0, 2, 3, 1) |
| 118 | video_uint8 = (video_uint8.clamp(0, 1) * 255).cpu().to(torch.uint8).contiguous() |
| 119 | |
| 120 | audio_float = normalize_audio_waveform_for_media(audio_waveform) |
| 121 | return video_uint8, audio_float |
| 122 | |
| 123 | |
| 124 | def write_benchmark_media( |
| 125 | *, |
| 126 | output_path: Path, |
| 127 | video_uint8: torch.Tensor, |
| 128 | audio_waveform: Optional[torch.Tensor], |
| 129 | fps: int, |
| 130 | audio_sr: int, |
| 131 | ) -> dict[str, Any]: |
| 132 | output_path.parent.mkdir(parents=True, exist_ok=True) |
| 133 | audio_waveform = normalize_audio_waveform_for_media(audio_waveform) |
| 134 | stats = audio_waveform_stats(audio_waveform) |
| 135 | |
| 136 | wrote_with_audio = False |
| 137 | wrote_sidecar_wav = False |
| 138 | if audio_waveform is not None: |
| 139 | try: |
| 140 | write_video( |
| 141 | str(output_path), |
| 142 | video_uint8, |
| 143 | fps=fps, |
| 144 | audio_array=audio_waveform, |
| 145 | audio_fps=audio_sr, |
| 146 | audio_codec="aac", |
| 147 | ) |
| 148 | wrote_with_audio = True |
| 149 | except Exception as exc: |
| 150 | print(f"[warn] write_video with audio failed for {output_path}: {exc}; audio_stats={stats}", flush=True) |
| 151 | |
| 152 | if not wrote_with_audio: |
| 153 | write_video(str(output_path), video_uint8, fps=fps) |
| 154 | if audio_waveform is not None: |
| 155 | try: |
| 156 | torchaudio.save(str(output_path.with_suffix(".wav")), audio_waveform, audio_sr) |
| 157 | wrote_sidecar_wav = True |
| 158 | except Exception as exc: |
| 159 | print(f"[warn] torchaudio.save failed for {output_path}: {exc}; audio_stats={stats}", flush=True) |
| 160 | |
| 161 | return { |
| 162 | "wrote_audio_in_mp4": wrote_with_audio, |
| 163 | "wrote_sidecar_wav": wrote_sidecar_wav, |
| 164 | "audio_stats": stats, |
| 165 | } |
| 166 | |
| 167 | |
| 168 | def save_memory_bank_frames(memory_frames: list[Any], save_dir: Path) -> None: |
| 169 | save_dir.mkdir(parents=True, exist_ok=True) |
| 170 | for old_file in save_dir.glob("*.jpg"): |
| 171 | old_file.unlink() |
| 172 | for idx, frame in enumerate(memory_frames): |
| 173 | if isinstance(frame, list): |
| 174 | frame = frame[len(frame) // 2] |
| 175 | frame.convert("RGB").save(save_dir / f"memory_{idx:03d}.jpg") |
| 176 | |
| 177 | |
| 178 | def concat_shot_videos(shot_paths: list[Path], output_path: Path) -> None: |
| 179 | if not shot_paths: |
| 180 | raise ValueError("No shot videos provided for concatenation") |
| 181 | |
| 182 | output_path.parent.mkdir(parents=True, exist_ok=True) |
| 183 | with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False, encoding="utf-8") as fp: |
| 184 | concat_file = Path(fp.name) |
| 185 | for shot_path in shot_paths: |
| 186 | fp.write(f"file '{shot_path.resolve().as_posix()}'\n") |
| 187 | |
| 188 | try: |
| 189 | cmd = [ |
| 190 | "ffmpeg", "-y", "-f", "concat", "-safe", "0", |
| 191 | "-i", str(concat_file), "-c", "copy", str(output_path), |
| 192 | ] |
| 193 | result = subprocess.run(cmd, capture_output=True, text=True) |
| 194 | if result.returncode != 0: |
| 195 | fallback_cmd = [ |
| 196 | "ffmpeg", "-y", "-f", "concat", "-safe", "0", |
| 197 | "-i", str(concat_file), |
| 198 | "-c:v", "libx264", "-preset", "medium", "-crf", "18", |
| 199 | "-c:a", "aac", "-b:a", "192k", |
| 200 | str(output_path), |
| 201 | ] |
| 202 | fallback_result = subprocess.run(fallback_cmd, capture_output=True, text=True) |
| 203 | if fallback_result.returncode != 0: |
| 204 | raise RuntimeError( |
| 205 | "Failed to concatenate shot videos with ffmpeg.\n" |
| 206 | f"copy stderr:\n{result.stderr}\n" |
| 207 | f"reencode stderr:\n{fallback_result.stderr}" |
| 208 | ) |
| 209 | finally: |
| 210 | concat_file.unlink(missing_ok=True) |
| 211 | |
| 212 | |
| 213 | def concat_shot_audios(audios: list[torch.Tensor]) -> Optional[torch.Tensor]: |
| 214 | if not audios: |
| 215 | return None |
| 216 | audio = audios[0] |
| 217 | if audio.ndim == 1: |
| 218 | sample_dim = 0 |
| 219 | elif audio.ndim == 2: |
| 220 | sample_dim = 1 if audio.shape[0] <= audio.shape[1] else 0 |
| 221 | else: |
| 222 | raise ValueError(f"Expected audio tensor with 1 or 2 dims, got shape={tuple(audio.shape)}") |
| 223 | return torch.cat([a.contiguous() for a in audios], dim=sample_dim).contiguous() |
| 224 |