| 1 | """Unified inference entrypoint: load models once, process all prompt files.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import json |
| 6 | import sys |
| 7 | import time |
| 8 | from datetime import datetime |
| 9 | from glob import glob |
| 10 | from pathlib import Path |
| 11 | from typing import Any |
| 12 | |
| 13 | # Ensure local packages are importable when running from repo root |
| 14 | _REPO_ROOT = Path(__file__).resolve().parent |
| 15 | for _subpath in ["ltx-core/src", "ltx-pipelines/src", "ltx-distillation/src"]: |
| 16 | _p = str(_REPO_ROOT / _subpath) |
| 17 | if _p not in sys.path: |
| 18 | sys.path.insert(0, _p) |
| 19 | |
| 20 | import torch |
| 21 | import torchaudio |
| 22 | import yaml |
| 23 | |
| 24 | from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps |
| 25 | from ltx_distillation.inference.bidirectional_pipeline import BidirectionalAVInferencePipeline |
| 26 | from ltx_distillation.inference.memory_bidirectional_pipeline import BidirectionalMemoryAVInferencePipeline |
| 27 | from ltx_distillation.inference.memory_multishot import ( |
| 28 | PairedAudioVideoMemoryBank, |
| 29 | audio_waveform_stats, |
| 30 | build_paired_audio_memory_kwargs, |
| 31 | load_multishot_prompts, |
| 32 | video_uint8_to_pil_frames, |
| 33 | ) |
| 34 | from ltx_distillation.models.ltx_wrapper import create_ltx2_wrapper |
| 35 | from ltx_distillation.models.text_encoder_wrapper import create_text_encoder_wrapper |
| 36 | from ltx_distillation.models.vae_wrapper import create_vae_wrappers |
| 37 | from ltx_distillation.utils import ( |
| 38 | add_noise, |
| 39 | compute_latent_shapes, |
| 40 | concat_shot_audios, |
| 41 | concat_shot_videos, |
| 42 | decode_benchmark_sample, |
| 43 | encode_memory_frames_batch, |
| 44 | save_memory_bank_frames, |
| 45 | write_benchmark_media, |
| 46 | ) |
| 47 | |
| 48 | REPO_ROOT = Path(__file__).resolve().parent |
| 49 | DEFAULT_CONFIG = REPO_ROOT / "configs" / "inference.yaml" |
| 50 | |
| 51 | |
| 52 | def _load_yaml_config(config_path: Path) -> dict[str, Any]: |
| 53 | with open(config_path, "r", encoding="utf-8") as f: |
| 54 | return yaml.safe_load(f) or {} |
| 55 | |
| 56 | |
| 57 | def _resolve_path(path_str: str, repo_root: Path) -> Path: |
| 58 | p = Path(path_str).expanduser() |
| 59 | if not p.is_absolute(): |
| 60 | p = repo_root / p |
| 61 | return p.resolve() |
| 62 | |
| 63 | |
| 64 | def str_to_bool(value: str | bool) -> bool: |
| 65 | if isinstance(value, bool): |
| 66 | return value |
| 67 | normalized = value.strip().lower() |
| 68 | if normalized in {"1", "true", "t", "yes", "y"}: |
| 69 | return True |
| 70 | if normalized in {"0", "false", "f", "no", "n"}: |
| 71 | return False |
| 72 | raise ValueError(f"Invalid boolean value: {value}") |
| 73 | |
| 74 | |
| 75 | # --------------------------------------------------------------------------- |
| 76 | # Config |
| 77 | # --------------------------------------------------------------------------- |
| 78 | |
| 79 | |
| 80 | class InferenceConfig: |
| 81 | """Parsed inference configuration from YAML + CLI overrides.""" |
| 82 | |
| 83 | def __init__(self, config_path: Path, **cli_overrides): |
| 84 | cfg = _load_yaml_config(config_path) |
| 85 | |
| 86 | paths_cfg = cfg.get("paths", {}) |
| 87 | video_cfg = cfg.get("video", {}) |
| 88 | denoising_cfg = cfg.get("denoising", {}) |
| 89 | memory_cfg = cfg.get("memory", {}) |
| 90 | audio_cfg = cfg.get("audio_memory", {}) |
| 91 | inference_cfg = cfg.get("inference", {}) |
| 92 | |
| 93 | # Paths |
| 94 | self.checkpoint = str(_resolve_path(paths_cfg.get("checkpoint", "checkpoints/echo-longvideo-release.safetensors"), REPO_ROOT)) |
| 95 | self.gemma_path = str(_resolve_path(paths_cfg.get("gemma_path", "checkpoints/gemma-3-12b"), REPO_ROOT)) |
| 96 | self.prompts_dir = str(_resolve_path(paths_cfg.get("prompts_dir", "prompts"), REPO_ROOT)) |
| 97 | self.prompts_glob = paths_cfg.get("prompts_glob", "*.json") |
| 98 | self.output_root = str(_resolve_path(paths_cfg.get("output_root", "inference_result/dmd"), REPO_ROOT)) |
| 99 | |
| 100 | # Video |
| 101 | self.num_frames = video_cfg.get("num_frames", 241) |
| 102 | self.video_height = video_cfg.get("height", 736) |
| 103 | self.video_width = video_cfg.get("width", 1280) |
| 104 | self.video_fps = video_cfg.get("fps", 25) |
| 105 | self.seed = video_cfg.get("seed", 12345) |
| 106 | |
| 107 | # Denoising |
| 108 | self.denoising_steps = denoising_cfg.get("steps", []) |
| 109 | self.denoising_sigmas = denoising_cfg.get("sigmas", []) |
| 110 | |
| 111 | # Memory |
| 112 | self.memory_max_size = memory_cfg.get("max_size", 7) |
| 113 | self.num_fix_frames = memory_cfg.get("num_fix_frames", 3) |
| 114 | self.memory_downscale_factor = memory_cfg.get("downscale_factor", 1) |
| 115 | self.memory_position_mode = memory_cfg.get("position_mode", "reference") |
| 116 | self.memory_lora_strength = memory_cfg.get("lora_strength", 1.0) |
| 117 | self.memory_lora_generator = memory_cfg.get("lora_generator", True) |
| 118 | self.memory_lora_path = memory_cfg.get("lora_path", "") or None |
| 119 | self.save_mode = memory_cfg.get("save_mode", "random_every_shot_frame") |
| 120 | self.video_memory_frame_selection_mode = memory_cfg.get("frame_selection_mode", "center") |
| 121 | self.video_memory_clip_num_frames = memory_cfg.get("clip_num_frames", 9) |
| 122 | |
| 123 | # Audio memory |
| 124 | self.enable_audio_memory = audio_cfg.get("enable", True) |
| 125 | self.audio_memory_window_size = audio_cfg.get("window_size", 96) |
| 126 | self.audio_memory_window_selection_mode = audio_cfg.get("window_selection_mode", "max_response") |
| 127 | self.audio_memory_sample_rate = audio_cfg.get("sample_rate", 16000) |
| 128 | self.audio_memory_mel_bins = audio_cfg.get("mel_bins", 128) |
| 129 | self.audio_memory_mel_hop_length = audio_cfg.get("mel_hop_length", 160) |
| 130 | self.audio_memory_n_fft = audio_cfg.get("n_fft", 1024) |
| 131 | self.audio_memory_downsample_factor = audio_cfg.get("downsample_factor", 4) |
| 132 | self.audio_memory_is_causal = audio_cfg.get("is_causal", True) |
| 133 | |
| 134 | # Inference |
| 135 | self.device = inference_cfg.get("device", "cuda") |
| 136 | self.dtype = inference_cfg.get("dtype", "bfloat16") |
| 137 | self.v2a_grad_scale = inference_cfg.get("v2a_grad_scale", 2.0) |
| 138 | |
| 139 | # Misc |
| 140 | self.prompt_max_chars = None |
| 141 | |
| 142 | # Apply CLI overrides |
| 143 | for key, value in cli_overrides.items(): |
| 144 | if value is not None and hasattr(self, key): |
| 145 | setattr(self, key, value) |
| 146 | |
| 147 | |
| 148 | # --------------------------------------------------------------------------- |
| 149 | # Engine |
| 150 | # --------------------------------------------------------------------------- |
| 151 | |
| 152 | |
| 153 | class InferenceEngine: |
| 154 | """Two-stage inference engine: encode all prompts first, then load generator. |
| 155 | |
| 156 | This avoids holding the text encoder (~24GB) and the video generator in |
| 157 | memory at the same time. Stage 1 loads only the text encoder, encodes every |
| 158 | prompt, then completely releases it. Stage 2 loads generator + VAEs. |
| 159 | """ |
| 160 | |
| 161 | def __init__(self, cfg: InferenceConfig): |
| 162 | self.cfg = cfg |
| 163 | self.device = torch.device(cfg.device) |
| 164 | self.dtype = torch.bfloat16 if cfg.dtype == "bfloat16" else torch.float32 |
| 165 | |
| 166 | checkpoint = Path(cfg.checkpoint).expanduser().resolve() |
| 167 | gemma_path = Path(cfg.gemma_path).expanduser().resolve() |
| 168 | if not checkpoint.exists(): |
| 169 | raise FileNotFoundError(f"Checkpoint not found: {checkpoint}") |
| 170 | if not gemma_path.exists(): |
| 171 | raise FileNotFoundError(f"Gemma path not found: {gemma_path}") |
| 172 | self._checkpoint = checkpoint |
| 173 | self._gemma_path = gemma_path |
| 174 | |
| 175 | # Stage-2 modules — populated by load_generator(). |
| 176 | self.generator = None |
| 177 | self.video_vae = None |
| 178 | self.audio_vae = None |
| 179 | self.base_pipeline = None |
| 180 | self.memory_pipeline = None |
| 181 | self.audio_sample_rate: int | None = None |
| 182 | |
| 183 | # ------------------------------------------------------------------ |
| 184 | # Stage 1: encode prompts, then free text encoder |
| 185 | # ------------------------------------------------------------------ |
| 186 | |
| 187 | def encode_all_prompts( |
| 188 | self, prompt_files: list[Path] |
| 189 | ) -> dict[Path, list[dict[str, Any]]]: |
| 190 | """Load text encoder, encode every prompt across all files, free encoder. |
| 191 | |
| 192 | Returns: {prompt_file: [cond_dict_on_cpu, ...]} |
| 193 | """ |
| 194 | print(f"[Stage 1] Loading text encoder...", flush=True) |
| 195 | text_encoder = create_text_encoder_wrapper( |
| 196 | checkpoint_path=str(self._checkpoint), |
| 197 | gemma_path=str(self._gemma_path), |
| 198 | device=self.device, |
| 199 | dtype=self.dtype, |
| 200 | ) |
| 201 | text_encoder.eval() |
| 202 | |
| 203 | cached: dict[Path, list[dict[str, Any]]] = {} |
| 204 | for prompts_file in prompt_files: |
| 205 | prompts = load_multishot_prompts(prompts_file, prompt_max_chars=self.cfg.prompt_max_chars) |
| 206 | if not prompts: |
| 207 | print(f"[Stage 1] Skipping empty prompts file: {prompts_file}", flush=True) |
| 208 | cached[prompts_file] = [] |
| 209 | continue |
| 210 | print(f"[Stage 1] Encoding {len(prompts)} prompts from {prompts_file.name}", flush=True) |
| 211 | file_conds: list[dict[str, Any]] = [] |
| 212 | for prompt in prompts: |
| 213 | cond = text_encoder([prompt]) |
| 214 | file_conds.append( |
| 215 | {k: (v.detach().cpu() if isinstance(v, torch.Tensor) else v) for k, v in cond.items()} |
| 216 | ) |
| 217 | del cond |
| 218 | cached[prompts_file] = file_conds |
| 219 | |
| 220 | # Fully release the text encoder (GPU + CPU). |
| 221 | del text_encoder |
| 222 | import gc |
| 223 | gc.collect() |
| 224 | if self.device.type == "cuda": |
| 225 | torch.cuda.empty_cache() |
| 226 | print(f"[Stage 1] Text encoder released.", flush=True) |
| 227 | return cached |
| 228 | |
| 229 | # ------------------------------------------------------------------ |
| 230 | # Stage 2: load generator + VAEs |
| 231 | # ------------------------------------------------------------------ |
| 232 | |
| 233 | def load_generator(self) -> None: |
| 234 | cfg = self.cfg |
| 235 | print(f"[Stage 2] Loading generator + VAEs from {self._checkpoint}", flush=True) |
| 236 | |
| 237 | loras: tuple[LoraPathStrengthAndSDOps, ...] = () |
| 238 | if cfg.memory_lora_path and cfg.memory_lora_generator: |
| 239 | loras = ( |
| 240 | LoraPathStrengthAndSDOps( |
| 241 | str(Path(cfg.memory_lora_path).expanduser()), |
| 242 | float(cfg.memory_lora_strength), |
| 243 | LTXV_LORA_COMFY_RENAMING_MAP, |
| 244 | ), |
| 245 | ) |
| 246 | |
| 247 | self.generator = create_ltx2_wrapper( |
| 248 | checkpoint_path=str(self._checkpoint), |
| 249 | gemma_path=str(self._gemma_path), |
| 250 | device=self.device, |
| 251 | dtype=self.dtype, |
| 252 | video_height=int(cfg.video_height), |
| 253 | video_width=int(cfg.video_width), |
| 254 | loras=loras, |
| 255 | ) |
| 256 | self.generator.eval() |
| 257 | |
| 258 | # Load VAEs to CPU; we hot-swap encoder/decoders per phase to avoid |
| 259 | # holding ~30GB generator and VAE decoders on GPU at the same time. |
| 260 | self.video_vae, self.audio_vae = create_vae_wrappers( |
| 261 | checkpoint_path=str(self._checkpoint), |
| 262 | device=torch.device("cpu"), |
| 263 | dtype=self.dtype, |
| 264 | with_video_encoder=True, |
| 265 | with_audio_encoder=True, |
| 266 | decoder_device=torch.device("cpu"), |
| 267 | ) |
| 268 | self.video_vae.eval() |
| 269 | self.audio_vae.eval() |
| 270 | |
| 271 | denoising_sigmas = torch.tensor(list(cfg.denoising_sigmas), device=self.device, dtype=torch.float32) |
| 272 | self.base_pipeline = BidirectionalAVInferencePipeline( |
| 273 | generator=self.generator, |
| 274 | add_noise_fn=add_noise, |
| 275 | denoising_sigmas=denoising_sigmas, |
| 276 | ) |
| 277 | self.memory_pipeline = BidirectionalMemoryAVInferencePipeline( |
| 278 | generator=self.generator, |
| 279 | add_noise_fn=add_noise, |
| 280 | denoising_sigmas=denoising_sigmas, |
| 281 | memory_downscale_factor=int(cfg.memory_downscale_factor), |
| 282 | ) |
| 283 | |
| 284 | self.audio_sample_rate = self.audio_vae.get_output_sample_rate() or 24000 |
| 285 | print(f"[Stage 2] Generator + VAEs ready.", flush=True) |
| 286 | |
| 287 | # ------------------------------------------------------------------ |
| 288 | # Module hot-swap helpers |
| 289 | # ------------------------------------------------------------------ |
| 290 | |
| 291 | @staticmethod |
| 292 | def _move(module, target_device) -> None: |
| 293 | if module is None: |
| 294 | return |
| 295 | module.to(target_device) |
| 296 | |
| 297 | def _empty(self) -> None: |
| 298 | if self.device.type == "cuda": |
| 299 | torch.cuda.empty_cache() |
| 300 | |
| 301 | def _stage_for_denoise(self) -> None: |
| 302 | """Generator on GPU; all VAE pieces on CPU.""" |
| 303 | self._move(self.video_vae.encoder, "cpu") |
| 304 | self._move(self.video_vae.decoder, "cpu") |
| 305 | self._move(self.audio_vae.encoder, "cpu") |
| 306 | self._move(self.audio_vae.decoder, "cpu") |
| 307 | self._move(self.audio_vae.vocoder, "cpu") |
| 308 | self._move(self.generator, self.device) |
| 309 | self._empty() |
| 310 | |
| 311 | def _stage_for_video_encode(self) -> None: |
| 312 | """Add video VAE encoder onto GPU alongside the generator (brief use).""" |
| 313 | self._move(self.video_vae.encoder, self.device) |
| 314 | |
| 315 | def _stage_after_video_encode(self) -> None: |
| 316 | self._move(self.video_vae.encoder, "cpu") |
| 317 | self._empty() |
| 318 | |
| 319 | def _stage_for_decode(self) -> None: |
| 320 | """Generator off GPU; VAE decoders + vocoder on GPU.""" |
| 321 | self._move(self.generator, "cpu") |
| 322 | self._empty() |
| 323 | self._move(self.video_vae.decoder, self.device) |
| 324 | self._move(self.audio_vae.decoder, self.device) |
| 325 | self._move(self.audio_vae.vocoder, self.device) |
| 326 | |
| 327 | def run_prompt_file( |
| 328 | self, |
| 329 | prompts_file: Path, |
| 330 | output_dir: Path, |
| 331 | cached_conds: list[dict[str, Any]], |
| 332 | ) -> None: |
| 333 | """Run multishot inference for a single prompt file using pre-encoded prompts.""" |
| 334 | if self.generator is None: |
| 335 | raise RuntimeError("call load_generator() before run_prompt_file()") |
| 336 | cfg = self.cfg |
| 337 | device = self.device |
| 338 | dtype = self.dtype |
| 339 | |
| 340 | prompts = load_multishot_prompts(prompts_file, prompt_max_chars=cfg.prompt_max_chars) |
| 341 | if not prompts: |
| 342 | print(f"[Engine] No prompts found in {prompts_file}, skipping.", flush=True) |
| 343 | return |
| 344 | |
| 345 | output_dir.mkdir(parents=True, exist_ok=True) |
| 346 | print(f"[Engine] Processing {prompts_file.name}: {len(prompts)} shots", flush=True) |
| 347 | |
| 348 | if len(cached_conds) != len(prompts): |
| 349 | raise ValueError( |
| 350 | f"cached_conds length ({len(cached_conds)}) does not match prompts count ({len(prompts)})" |
| 351 | ) |
| 352 | |
| 353 | video_shape, audio_shape = compute_latent_shapes( |
| 354 | num_frames=int(cfg.num_frames), |
| 355 | video_height=int(cfg.video_height), |
| 356 | video_width=int(cfg.video_width), |
| 357 | batch_size=1, |
| 358 | video_fps=float(cfg.video_fps), |
| 359 | ) |
| 360 | |
| 361 | memory_bank = PairedAudioVideoMemoryBank( |
| 362 | max_size=int(cfg.memory_max_size), |
| 363 | save_mode=str(cfg.save_mode), |
| 364 | num_fix_frames=int(cfg.num_fix_frames), |
| 365 | ) |
| 366 | |
| 367 | shot_paths: list[Path] = [] |
| 368 | shot_audios: list[torch.Tensor] = [] |
| 369 | metadata: dict[str, Any] = { |
| 370 | "checkpoint": cfg.checkpoint, |
| 371 | "prompts_file": str(prompts_file), |
| 372 | "output_dir": str(output_dir), |
| 373 | "denoising_steps": [int(x) for x in cfg.denoising_steps], |
| 374 | "denoising_sigmas": [float(x) for x in cfg.denoising_sigmas], |
| 375 | "num_prompts": len(prompts), |
| 376 | "save_mode": cfg.save_mode, |
| 377 | "memory_max_size": cfg.memory_max_size, |
| 378 | "num_fix_frames": cfg.num_fix_frames, |
| 379 | "enable_audio_memory": cfg.enable_audio_memory, |
| 380 | "shots": [], |
| 381 | } |
| 382 | |
| 383 | run_started = time.perf_counter() |
| 384 | shot_durations: list[dict[str, float]] = [] |
| 385 | |
| 386 | for shot_idx, prompt in enumerate(prompts): |
| 387 | shot_started = time.perf_counter() |
| 388 | conditional_dict = { |
| 389 | k: (v.to(device) if isinstance(v, torch.Tensor) else v) |
| 390 | for k, v in cached_conds[shot_idx].items() |
| 391 | } |
| 392 | prompt_seed = int(cfg.seed) + shot_idx |
| 393 | memory_size_before = len(memory_bank) |
| 394 | |
| 395 | print( |
| 396 | f"[Engine] shot={shot_idx + 1}/{len(prompts)} " |
| 397 | f"memory_size_before={memory_size_before} seed={prompt_seed}", |
| 398 | flush=True, |
| 399 | ) |
| 400 | |
| 401 | memory_video = None |
| 402 | memory_audio_kwargs: dict[str, Any] = {} |
| 403 | |
| 404 | # Phase A: denoising — generator on GPU, decoders on CPU. |
| 405 | self._stage_for_denoise() |
| 406 | |
| 407 | denoise_started = time.perf_counter() |
| 408 | with torch.random.fork_rng(devices=[device]): |
| 409 | torch.manual_seed(prompt_seed) |
| 410 | if device.type == "cuda": |
| 411 | torch.cuda.manual_seed(prompt_seed) |
| 412 | |
| 413 | if len(memory_bank) > 0: |
| 414 | # Briefly bring video encoder onto GPU. |
| 415 | self._stage_for_video_encode() |
| 416 | memory_video = encode_memory_frames_batch( |
| 417 | video_vae=self.video_vae, |
| 418 | batch_memory_frames=[memory_bank.get_memory_frames()], |
| 419 | target_h=int(cfg.video_height), |
| 420 | target_w=int(cfg.video_width), |
| 421 | device=device, |
| 422 | dtype=dtype, |
| 423 | ) |
| 424 | self._stage_after_video_encode() |
| 425 | |
| 426 | memory_audio_kwargs = build_paired_audio_memory_kwargs( |
| 427 | memory_bank, |
| 428 | enable_audio_memory=bool(cfg.enable_audio_memory), |
| 429 | v2a_grad_scale=float(cfg.v2a_grad_scale), |
| 430 | memory_position_mode=str(cfg.memory_position_mode), |
| 431 | ) |
| 432 | |
| 433 | video_latent, audio_latent = self.memory_pipeline.generate( |
| 434 | video_shape=tuple(video_shape), |
| 435 | audio_shape=tuple(audio_shape), |
| 436 | conditional_dict=conditional_dict, |
| 437 | memory_video=memory_video, |
| 438 | seed=prompt_seed, |
| 439 | **memory_audio_kwargs, |
| 440 | ) |
| 441 | else: |
| 442 | video_latent, audio_latent = self.base_pipeline.generate( |
| 443 | video_shape=tuple(video_shape), |
| 444 | audio_shape=tuple(audio_shape), |
| 445 | conditional_dict=conditional_dict, |
| 446 | seed=prompt_seed, |
| 447 | ) |
| 448 | if device.type == "cuda": |
| 449 | torch.cuda.synchronize() |
| 450 | denoise_elapsed = time.perf_counter() - denoise_started |
| 451 | |
| 452 | # Release intermediates that are no longer needed before the heavy |
| 453 | # decoder swap-in. |
| 454 | del conditional_dict, memory_video, memory_audio_kwargs |
| 455 | memory_video = None |
| 456 | memory_audio_kwargs = {} |
| 457 | |
| 458 | # Phase B: decode — generator off GPU, decoders + vocoder on GPU. |
| 459 | self._stage_for_decode() |
| 460 | |
| 461 | decode_started = time.perf_counter() |
| 462 | audio_memory_latent = ( |
| 463 | audio_latent.detach().cpu().contiguous() |
| 464 | if (cfg.enable_audio_memory and audio_latent is not None) |
| 465 | else None |
| 466 | ) |
| 467 | video_uint8, audio_waveform = decode_benchmark_sample( |
| 468 | self.video_vae, self.audio_vae, video_latent, audio_latent |
| 469 | ) |
| 470 | if device.type == "cuda": |
| 471 | torch.cuda.synchronize() |
| 472 | decode_elapsed = time.perf_counter() - decode_started |
| 473 | memory_frames_for_bank = video_uint8_to_pil_frames(video_uint8) |
| 474 | |
| 475 | new_memory_metadata: dict[str, Any] = {} |
| 476 | if audio_memory_latent is not None: |
| 477 | new_memory_metadata = memory_bank.save_memory_slot( |
| 478 | memory_frames_for_bank, |
| 479 | audio_memory_latent, |
| 480 | audio_window_size=int(cfg.audio_memory_window_size), |
| 481 | video_clip_num_frames=int(cfg.video_memory_clip_num_frames), |
| 482 | audio_waveform=audio_waveform, |
| 483 | audio_sample_rate=int(cfg.audio_memory_sample_rate), |
| 484 | video_fps=float(cfg.video_fps), |
| 485 | audio_window_selection_mode=str(cfg.audio_memory_window_selection_mode), |
| 486 | video_frame_selection_mode=str(cfg.video_memory_frame_selection_mode), |
| 487 | audio_memory_mel_bins=int(cfg.audio_memory_mel_bins), |
| 488 | audio_memory_mel_hop_length=int(cfg.audio_memory_mel_hop_length), |
| 489 | audio_memory_n_fft=int(cfg.audio_memory_n_fft), |
| 490 | audio_memory_downsample_factor=int(cfg.audio_memory_downsample_factor), |
| 491 | audio_memory_is_causal=bool(cfg.audio_memory_is_causal), |
| 492 | ) |
| 493 | |
| 494 | save_memory_bank_frames( |
| 495 | memory_bank.get_memory_frames(), |
| 496 | output_dir / "memory_bank" / f"shot_{shot_idx:03d}", |
| 497 | ) |
| 498 | |
| 499 | shot_path = output_dir / f"shot_{shot_idx:03d}.mp4" |
| 500 | write_result = write_benchmark_media( |
| 501 | output_path=shot_path, |
| 502 | video_uint8=video_uint8, |
| 503 | audio_waveform=audio_waveform, |
| 504 | fps=int(cfg.video_fps), |
| 505 | audio_sr=int(self.audio_sample_rate), |
| 506 | ) |
| 507 | shot_paths.append(shot_path) |
| 508 | if audio_waveform is not None: |
| 509 | shot_audios.append(audio_waveform.cpu()) |
| 510 | |
| 511 | shot_elapsed = time.perf_counter() - shot_started |
| 512 | timing = { |
| 513 | "denoise_sec": round(denoise_elapsed, 3), |
| 514 | "decode_sec": round(decode_elapsed, 3), |
| 515 | "total_sec": round(shot_elapsed, 3), |
| 516 | } |
| 517 | shot_durations.append(timing) |
| 518 | |
| 519 | metadata["shots"].append( |
| 520 | { |
| 521 | "shot_idx": int(shot_idx), |
| 522 | "prompt": prompt, |
| 523 | "output_path": str(shot_path), |
| 524 | "memory_size_before": int(memory_size_before), |
| 525 | "memory_size_after": int(len(memory_bank)), |
| 526 | "new_memory_entry": new_memory_metadata, |
| 527 | "audio_latent_shape": list(audio_latent.shape) if audio_latent is not None else None, |
| 528 | "wrote_audio_in_mp4": bool(write_result["wrote_audio_in_mp4"]), |
| 529 | "wrote_sidecar_wav": bool(write_result["wrote_sidecar_wav"]), |
| 530 | "audio_stats": write_result["audio_stats"], |
| 531 | "memory_entries": memory_bank.get_memory_metadata(), |
| 532 | "timing": timing, |
| 533 | } |
| 534 | ) |
| 535 | |
| 536 | print( |
| 537 | f"[Engine] shot={shot_idx + 1}/{len(prompts)} done " |
| 538 | f"denoise={denoise_elapsed:.1f}s decode={decode_elapsed:.1f}s " |
| 539 | f"total={shot_elapsed:.1f}s", |
| 540 | flush=True, |
| 541 | ) |
| 542 | |
| 543 | del video_latent, audio_latent, video_uint8, audio_waveform |
| 544 | del audio_memory_latent, memory_frames_for_bank |
| 545 | if device.type == "cuda": |
| 546 | torch.cuda.empty_cache() |
| 547 | |
| 548 | run_elapsed = time.perf_counter() - run_started |
| 549 | avg_total = sum(t["total_sec"] for t in shot_durations) / max(len(shot_durations), 1) |
| 550 | avg_denoise = sum(t["denoise_sec"] for t in shot_durations) / max(len(shot_durations), 1) |
| 551 | avg_decode = sum(t["decode_sec"] for t in shot_durations) / max(len(shot_durations), 1) |
| 552 | metadata["timing"] = { |
| 553 | "run_total_sec": round(run_elapsed, 3), |
| 554 | "avg_shot_total_sec": round(avg_total, 3), |
| 555 | "avg_denoise_sec": round(avg_denoise, 3), |
| 556 | "avg_decode_sec": round(avg_decode, 3), |
| 557 | } |
| 558 | print( |
| 559 | f"[Engine] {prompts_file.name} run_total={run_elapsed:.1f}s " |
| 560 | f"avg_shot={avg_total:.1f}s (denoise={avg_denoise:.1f}s decode={avg_decode:.1f}s)", |
| 561 | flush=True, |
| 562 | ) |
| 563 | |
| 564 | combined_path = output_dir / "combined_shots.mp4" |
| 565 | concat_shot_videos(shot_paths, combined_path) |
| 566 | combined_audio = concat_shot_audios(shot_audios) |
| 567 | combined_audio_path = None |
| 568 | if combined_audio is not None: |
| 569 | combined_audio_path = output_dir / "combined_shots.wav" |
| 570 | torchaudio.save(str(combined_audio_path), combined_audio, sample_rate=int(self.audio_sample_rate)) |
| 571 | |
| 572 | metadata["combined_path"] = str(combined_path) |
| 573 | metadata["combined_audio_path"] = str(combined_audio_path) if combined_audio_path else None |
| 574 | metadata["combined_audio_stats"] = audio_waveform_stats(combined_audio) |
| 575 | metadata_path = output_dir / "run_metadata.json" |
| 576 | metadata_path.write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8") |
| 577 | |
| 578 | print(f"[Engine] Done: {prompts_file.name} -> {combined_path}", flush=True) |
| 579 | |
| 580 | |
| 581 | # --------------------------------------------------------------------------- |
| 582 | # CLI |
| 583 | # --------------------------------------------------------------------------- |
| 584 | |
| 585 | |
| 586 | def parse_args(): |
| 587 | import argparse |
| 588 | |
| 589 | parser = argparse.ArgumentParser( |
| 590 | description="Unified inference: load models once, process all prompt files.", |
| 591 | formatter_class=argparse.ArgumentDefaultsHelpFormatter, |
| 592 | ) |
| 593 | parser.add_argument("--config", type=str, default=str(DEFAULT_CONFIG), help="Path to YAML config file") |
| 594 | parser.add_argument("--prompts-dir", type=str, default=None, help="Override prompts directory") |
| 595 | parser.add_argument("--prompts-glob", type=str, default=None, help="Override prompts glob pattern") |
| 596 | parser.add_argument("--output-root", type=str, default=None, help="Override output root directory") |
| 597 | parser.add_argument("--seed", type=int, default=None) |
| 598 | parser.add_argument("--num-frames", type=int, default=None) |
| 599 | parser.add_argument("--video-height", type=int, default=None) |
| 600 | parser.add_argument("--video-width", type=int, default=None) |
| 601 | parser.add_argument("--video-fps", type=int, default=None) |
| 602 | parser.add_argument("--v2a-grad-scale", type=float, default=None) |
| 603 | parser.add_argument("--memory-max-size", type=int, default=None) |
| 604 | parser.add_argument("--num-fix-frames", type=int, default=None) |
| 605 | parser.add_argument("--enable-audio-memory", type=str_to_bool, default=None) |
| 606 | return parser.parse_args() |
| 607 | |
| 608 | |
| 609 | def main() -> None: |
| 610 | args = parse_args() |
| 611 | |
| 612 | config_path = Path(args.config).expanduser().resolve() |
| 613 | if not config_path.exists(): |
| 614 | raise FileNotFoundError(f"Config file not found: {config_path}") |
| 615 | |
| 616 | cli_overrides = {} |
| 617 | for key in ["seed", "num_frames", "video_height", "video_width", "video_fps", |
| 618 | "v2a_grad_scale", "memory_max_size", "num_fix_frames", "enable_audio_memory"]: |
| 619 | val = getattr(args, key, None) |
| 620 | if val is not None: |
| 621 | cli_overrides[key] = val |
| 622 | if args.prompts_dir: |
| 623 | cli_overrides["prompts_dir"] = str(Path(args.prompts_dir).expanduser().resolve()) |
| 624 | if args.prompts_glob: |
| 625 | cli_overrides["prompts_glob"] = args.prompts_glob |
| 626 | if args.output_root: |
| 627 | cli_overrides["output_root"] = str(Path(args.output_root).expanduser().resolve()) |
| 628 | |
| 629 | cfg = InferenceConfig(config_path, **cli_overrides) |
| 630 | |
| 631 | if len(cfg.denoising_steps) != len(cfg.denoising_sigmas): |
| 632 | raise ValueError("denoising steps and sigmas must have the same length") |
| 633 | |
| 634 | engine = InferenceEngine(cfg) |
| 635 | |
| 636 | # Discover prompt files |
| 637 | prompts_dir = Path(cfg.prompts_dir) |
| 638 | prompts_pattern = cfg.prompts_glob |
| 639 | if not prompts_pattern.startswith("/"): |
| 640 | prompt_files = sorted(prompts_dir.glob(prompts_pattern)) |
| 641 | else: |
| 642 | prompt_files = sorted(Path(p) for p in glob(prompts_pattern)) |
| 643 | |
| 644 | if not prompt_files: |
| 645 | raise FileNotFoundError(f"No prompt files matched: {prompts_dir / prompts_pattern}") |
| 646 | |
| 647 | print(f"[Inference] Found {len(prompt_files)} prompt file(s)", flush=True) |
| 648 | |
| 649 | # Stage 1: encode all prompts across all files, then release text encoder. |
| 650 | cached_per_file = engine.encode_all_prompts(prompt_files) |
| 651 | |
| 652 | # Stage 2: now load the generator + VAEs. |
| 653 | engine.load_generator() |
| 654 | |
| 655 | # Stage 3: run inference for each file using the pre-encoded prompts. |
| 656 | output_root = Path(cfg.output_root) / "outputs" |
| 657 | for prompts_file in prompt_files: |
| 658 | cached = cached_per_file.get(prompts_file, []) |
| 659 | if not cached: |
| 660 | continue |
| 661 | prompt_name = prompts_file.stem |
| 662 | timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") |
| 663 | run_output_dir = output_root / prompt_name / f"inference_{timestamp}" |
| 664 | engine.run_prompt_file(prompts_file, run_output_dir, cached) |
| 665 | |
| 666 | print(f"[Inference] All {len(prompt_files)} prompt file(s) processed.", flush=True) |
| 667 | |
| 668 | |
| 669 | if __name__ == "__main__": |
| 670 | main() |
| 671 |