| 1 | #!/usr/bin/env python3 |
| 2 | """Standalone multishot inference (release mode - single merged checkpoint).""" |
| 3 | |
| 4 | from __future__ import annotations |
| 5 | |
| 6 | import argparse |
| 7 | import json |
| 8 | from pathlib import Path |
| 9 | from typing import Any, Optional |
| 10 | |
| 11 | import torch |
| 12 | import torchaudio |
| 13 | import yaml |
| 14 | |
| 15 | from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps |
| 16 | from ltx_distillation.inference.bidirectional_pipeline import BidirectionalAVInferencePipeline |
| 17 | from ltx_distillation.inference.memory_bidirectional_pipeline import BidirectionalMemoryAVInferencePipeline |
| 18 | from ltx_distillation.inference.memory_multishot import ( |
| 19 | PairedAudioVideoMemoryBank, |
| 20 | audio_waveform_stats, |
| 21 | build_paired_audio_memory_kwargs, |
| 22 | load_multishot_prompts, |
| 23 | video_uint8_to_pil_frames, |
| 24 | ) |
| 25 | from ltx_distillation.models.ltx_wrapper import create_ltx2_wrapper |
| 26 | from ltx_distillation.models.text_encoder_wrapper import create_text_encoder_wrapper |
| 27 | from ltx_distillation.models.vae_wrapper import create_vae_wrappers |
| 28 | from ltx_distillation.utils import ( |
| 29 | add_noise, |
| 30 | compute_latent_shapes, |
| 31 | concat_shot_audios, |
| 32 | concat_shot_videos, |
| 33 | decode_benchmark_sample, |
| 34 | encode_memory_frames_batch, |
| 35 | save_memory_bank_frames, |
| 36 | write_benchmark_media, |
| 37 | ) |
| 38 | |
| 39 | REPO_ROOT = Path(__file__).resolve().parents[2] |
| 40 | DEFAULT_CONFIG = REPO_ROOT / "configs" / "inference.yaml" |
| 41 | |
| 42 | |
| 43 | def _load_yaml_config(config_path: Path) -> dict[str, Any]: |
| 44 | with open(config_path, "r", encoding="utf-8") as f: |
| 45 | return yaml.safe_load(f) or {} |
| 46 | |
| 47 | |
| 48 | def _resolve_path(path_str: str, repo_root: Path) -> Path: |
| 49 | p = Path(path_str).expanduser() |
| 50 | if not p.is_absolute(): |
| 51 | p = repo_root / p |
| 52 | return p.resolve() |
| 53 | |
| 54 | |
| 55 | def str_to_bool(value: str | bool) -> bool: |
| 56 | if isinstance(value, bool): |
| 57 | return value |
| 58 | normalized = value.strip().lower() |
| 59 | if normalized in {"1", "true", "t", "yes", "y"}: |
| 60 | return True |
| 61 | if normalized in {"0", "false", "f", "no", "n"}: |
| 62 | return False |
| 63 | raise argparse.ArgumentTypeError(f"Invalid boolean value: {value}") |
| 64 | |
| 65 | |
| 66 | def _parse_int_list(value: str) -> list[int]: |
| 67 | items = [item.strip() for item in value.split(",") if item.strip()] |
| 68 | if not items: |
| 69 | raise argparse.ArgumentTypeError("Expected a comma-separated integer list") |
| 70 | return [int(item) for item in items] |
| 71 | |
| 72 | |
| 73 | def _parse_float_list(value: str) -> list[float]: |
| 74 | items = [item.strip() for item in value.split(",") if item.strip()] |
| 75 | if not items: |
| 76 | raise argparse.ArgumentTypeError("Expected a comma-separated float list") |
| 77 | return [float(item) for item in items] |
| 78 | |
| 79 | |
| 80 | def _build_config(config_path: Path, cli_overrides: dict[str, Any]) -> argparse.Namespace: |
| 81 | """Load YAML config and apply CLI overrides on top.""" |
| 82 | cfg = _load_yaml_config(config_path) |
| 83 | |
| 84 | paths_cfg = cfg.get("paths", {}) |
| 85 | video_cfg = cfg.get("video", {}) |
| 86 | denoising_cfg = cfg.get("denoising", {}) |
| 87 | memory_cfg = cfg.get("memory", {}) |
| 88 | audio_cfg = cfg.get("audio_memory", {}) |
| 89 | inference_cfg = cfg.get("inference", {}) |
| 90 | |
| 91 | defaults = argparse.Namespace( |
| 92 | # paths |
| 93 | checkpoint=str(_resolve_path(paths_cfg.get("checkpoint", "checkpoints/echo-longvideo-release.safetensors"), REPO_ROOT)), |
| 94 | gemma_path=str(_resolve_path(paths_cfg.get("gemma_path", "checkpoints/gemma-3-12b"), REPO_ROOT)), |
| 95 | prompts_file=None, |
| 96 | output_dir=None, |
| 97 | # video |
| 98 | num_frames=video_cfg.get("num_frames", 241), |
| 99 | video_height=video_cfg.get("height", 736), |
| 100 | video_width=video_cfg.get("width", 1280), |
| 101 | video_fps=video_cfg.get("fps", 25), |
| 102 | seed=video_cfg.get("seed", 12345), |
| 103 | # denoising |
| 104 | denoising_steps=denoising_cfg.get("steps", [1000, 993, 984, 973, 959, 942, 918, 885, 836, 755, 591, 500, 400, 300, 200, 100, 0]), |
| 105 | denoising_sigmas=denoising_cfg.get("sigmas", [1.0, 0.99256, 0.983633, 0.972721, 0.959082, 0.941546, 0.918165, 0.885432, 0.836333, 0.754505, 0.590859, 0.50, 0.40, 0.30, 0.20, 0.10, 0.0]), |
| 106 | # memory |
| 107 | memory_max_size=memory_cfg.get("max_size", 7), |
| 108 | num_fix_frames=memory_cfg.get("num_fix_frames", 3), |
| 109 | memory_downscale_factor=memory_cfg.get("downscale_factor", 1), |
| 110 | memory_position_mode=memory_cfg.get("position_mode", "reference"), |
| 111 | memory_lora_strength=memory_cfg.get("lora_strength", 1.0), |
| 112 | memory_lora_generator=memory_cfg.get("lora_generator", True), |
| 113 | memory_lora_path=memory_cfg.get("lora_path", "") or None, |
| 114 | save_mode=memory_cfg.get("save_mode", "random_every_shot_frame"), |
| 115 | video_memory_frame_selection_mode=memory_cfg.get("frame_selection_mode", "center"), |
| 116 | video_memory_clip_num_frames=memory_cfg.get("clip_num_frames", 9), |
| 117 | # audio memory |
| 118 | enable_audio_memory=audio_cfg.get("enable", True), |
| 119 | audio_memory_window_size=audio_cfg.get("window_size", 96), |
| 120 | audio_memory_window_selection_mode=audio_cfg.get("window_selection_mode", "max_response"), |
| 121 | audio_memory_sample_rate=audio_cfg.get("sample_rate", 16000), |
| 122 | audio_memory_mel_bins=audio_cfg.get("mel_bins", 128), |
| 123 | audio_memory_mel_hop_length=audio_cfg.get("mel_hop_length", 160), |
| 124 | audio_memory_n_fft=audio_cfg.get("n_fft", 1024), |
| 125 | audio_memory_downsample_factor=audio_cfg.get("downsample_factor", 4), |
| 126 | audio_memory_is_causal=audio_cfg.get("is_causal", True), |
| 127 | # inference |
| 128 | device=inference_cfg.get("device", "cuda"), |
| 129 | dtype=inference_cfg.get("dtype", "bfloat16"), |
| 130 | v2a_grad_scale=inference_cfg.get("v2a_grad_scale", 2.0), |
| 131 | # misc |
| 132 | prompt_max_chars=None, |
| 133 | ) |
| 134 | |
| 135 | for key, value in cli_overrides.items(): |
| 136 | if value is not None: |
| 137 | setattr(defaults, key, value) |
| 138 | |
| 139 | return defaults |
| 140 | |
| 141 | |
| 142 | def parse_args() -> argparse.Namespace: |
| 143 | parser = argparse.ArgumentParser( |
| 144 | description="Multishot inference with merged release checkpoint.", |
| 145 | formatter_class=argparse.ArgumentDefaultsHelpFormatter, |
| 146 | ) |
| 147 | parser.add_argument("--config", type=str, default=str(DEFAULT_CONFIG), help="Path to YAML config file") |
| 148 | parser.add_argument("--checkpoint", type=str, default=None) |
| 149 | parser.add_argument("--gemma-path", type=str, default=None) |
| 150 | parser.add_argument("--prompts-file", type=str, default=None) |
| 151 | parser.add_argument("--output-dir", type=str, default=None) |
| 152 | parser.add_argument("--device", type=str, default=None) |
| 153 | parser.add_argument("--dtype", type=str, default=None, choices=["bfloat16", "float32"]) |
| 154 | parser.add_argument("--seed", type=int, default=None) |
| 155 | parser.add_argument("--num-frames", type=int, default=None) |
| 156 | parser.add_argument("--video-height", type=int, default=None) |
| 157 | parser.add_argument("--video-width", type=int, default=None) |
| 158 | parser.add_argument("--video-fps", type=int, default=None) |
| 159 | parser.add_argument("--save-mode", type=str, default=None) |
| 160 | parser.add_argument("--memory-max-size", type=int, default=None) |
| 161 | parser.add_argument("--num-fix-frames", type=int, default=None) |
| 162 | parser.add_argument("--prompt-max-chars", type=int, default=None) |
| 163 | parser.add_argument("--enable-audio-memory", type=str_to_bool, default=None) |
| 164 | parser.add_argument("--denoising-steps", type=_parse_int_list, default=None) |
| 165 | parser.add_argument("--denoising-sigmas", type=_parse_float_list, default=None) |
| 166 | parser.add_argument("--memory-downscale-factor", type=int, default=None) |
| 167 | parser.add_argument("--v2a-grad-scale", type=float, default=None) |
| 168 | parser.add_argument("--memory-position-mode", type=str, default=None) |
| 169 | parser.add_argument("--audio-memory-window-size", type=int, default=None) |
| 170 | parser.add_argument("--audio-memory-window-selection-mode", type=str, default=None) |
| 171 | parser.add_argument("--video-memory-frame-selection-mode", type=str, default=None) |
| 172 | parser.add_argument("--video-memory-clip-num-frames", type=int, default=None) |
| 173 | parser.add_argument("--audio-memory-sample-rate", type=int, default=None) |
| 174 | parser.add_argument("--audio-memory-mel-bins", type=int, default=None) |
| 175 | parser.add_argument("--audio-memory-mel-hop-length", type=int, default=None) |
| 176 | parser.add_argument("--audio-memory-n-fft", type=int, default=None) |
| 177 | parser.add_argument("--audio-memory-downsample-factor", type=int, default=None) |
| 178 | parser.add_argument("--audio-memory-is-causal", type=str_to_bool, default=None) |
| 179 | |
| 180 | raw_args = parser.parse_args() |
| 181 | |
| 182 | cli_overrides = {k: v for k, v in vars(raw_args).items() if v is not None and k != "config"} |
| 183 | |
| 184 | config_path = Path(raw_args.config).expanduser().resolve() |
| 185 | if not config_path.exists(): |
| 186 | raise FileNotFoundError(f"Config file not found: {config_path}") |
| 187 | |
| 188 | return _build_config(config_path, cli_overrides) |
| 189 | |
| 190 | |
| 191 | def main() -> None: |
| 192 | args = parse_args() |
| 193 | if len(args.denoising_steps) != len(args.denoising_sigmas): |
| 194 | raise ValueError("denoising steps and sigmas must have the same length") |
| 195 | |
| 196 | device = torch.device(args.device) |
| 197 | dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float32 |
| 198 | |
| 199 | checkpoint = Path(args.checkpoint).expanduser().resolve() |
| 200 | gemma_path = Path(args.gemma_path).expanduser().resolve() |
| 201 | prompts_file = Path(args.prompts_file).expanduser().resolve() |
| 202 | output_dir = Path(args.output_dir).expanduser().resolve() |
| 203 | output_dir.mkdir(parents=True, exist_ok=True) |
| 204 | |
| 205 | if not checkpoint.exists(): |
| 206 | raise FileNotFoundError(f"Checkpoint not found: {checkpoint}") |
| 207 | if not gemma_path.exists(): |
| 208 | raise FileNotFoundError(f"Gemma path not found: {gemma_path}") |
| 209 | |
| 210 | prompts = load_multishot_prompts(prompts_file, prompt_max_chars=args.prompt_max_chars) |
| 211 | if not prompts: |
| 212 | raise ValueError(f"No prompts found in {prompts_file}") |
| 213 | |
| 214 | print(f"[Inference] prompts_file={prompts_file}", flush=True) |
| 215 | print(f"[Inference] num_prompts={len(prompts)}", flush=True) |
| 216 | print(f"[Inference] checkpoint={checkpoint}", flush=True) |
| 217 | print( |
| 218 | f"[Inference] frames={args.num_frames} size={args.video_height}x{args.video_width} " |
| 219 | f"fps={args.video_fps} enable_audio_memory={bool(args.enable_audio_memory)}", |
| 220 | flush=True, |
| 221 | ) |
| 222 | |
| 223 | loras: tuple[LoraPathStrengthAndSDOps, ...] = () |
| 224 | if args.memory_lora_path and args.memory_lora_generator: |
| 225 | loras = ( |
| 226 | LoraPathStrengthAndSDOps( |
| 227 | str(Path(args.memory_lora_path).expanduser()), |
| 228 | float(args.memory_lora_strength), |
| 229 | LTXV_LORA_COMFY_RENAMING_MAP, |
| 230 | ), |
| 231 | ) |
| 232 | |
| 233 | # ------------------------------------------------------------------ |
| 234 | # Stage 1: load text encoder, encode all prompts, release encoder. |
| 235 | # ------------------------------------------------------------------ |
| 236 | print(f"[Stage 1] Loading text encoder...", flush=True) |
| 237 | text_encoder = create_text_encoder_wrapper( |
| 238 | checkpoint_path=str(checkpoint), |
| 239 | gemma_path=str(gemma_path), |
| 240 | device=device, |
| 241 | dtype=dtype, |
| 242 | ) |
| 243 | text_encoder.eval() |
| 244 | |
| 245 | print(f"[Stage 1] Encoding {len(prompts)} prompts...", flush=True) |
| 246 | cached_conds: list[dict[str, Any]] = [] |
| 247 | for prompt in prompts: |
| 248 | cond = text_encoder([prompt]) |
| 249 | cached_conds.append( |
| 250 | {k: (v.detach().cpu() if isinstance(v, torch.Tensor) else v) for k, v in cond.items()} |
| 251 | ) |
| 252 | del cond |
| 253 | |
| 254 | del text_encoder |
| 255 | import gc |
| 256 | gc.collect() |
| 257 | if device.type == "cuda": |
| 258 | torch.cuda.empty_cache() |
| 259 | print(f"[Stage 1] Text encoder released.", flush=True) |
| 260 | |
| 261 | # ------------------------------------------------------------------ |
| 262 | # Stage 2: load generator + VAEs. |
| 263 | # ------------------------------------------------------------------ |
| 264 | print(f"[Stage 2] Loading generator + VAEs...", flush=True) |
| 265 | generator = create_ltx2_wrapper( |
| 266 | checkpoint_path=str(checkpoint), |
| 267 | gemma_path=str(gemma_path), |
| 268 | device=device, |
| 269 | dtype=dtype, |
| 270 | video_height=int(args.video_height), |
| 271 | video_width=int(args.video_width), |
| 272 | loras=loras, |
| 273 | ) |
| 274 | generator.eval() |
| 275 | |
| 276 | video_vae, audio_vae = create_vae_wrappers( |
| 277 | checkpoint_path=str(checkpoint), |
| 278 | device=device, |
| 279 | dtype=dtype, |
| 280 | with_video_encoder=True, |
| 281 | with_audio_encoder=True, |
| 282 | decoder_device=device, |
| 283 | ) |
| 284 | video_vae.eval() |
| 285 | audio_vae.eval() |
| 286 | |
| 287 | denoising_sigmas = torch.tensor(list(args.denoising_sigmas), device=device, dtype=torch.float32) |
| 288 | base_pipeline = BidirectionalAVInferencePipeline( |
| 289 | generator=generator, |
| 290 | add_noise_fn=add_noise, |
| 291 | denoising_sigmas=denoising_sigmas, |
| 292 | ) |
| 293 | memory_pipeline = BidirectionalMemoryAVInferencePipeline( |
| 294 | generator=generator, |
| 295 | add_noise_fn=add_noise, |
| 296 | denoising_sigmas=denoising_sigmas, |
| 297 | memory_downscale_factor=int(args.memory_downscale_factor), |
| 298 | ) |
| 299 | |
| 300 | audio_sample_rate = audio_vae.get_output_sample_rate() or 24000 |
| 301 | video_shape, audio_shape = compute_latent_shapes( |
| 302 | num_frames=int(args.num_frames), |
| 303 | video_height=int(args.video_height), |
| 304 | video_width=int(args.video_width), |
| 305 | batch_size=1, |
| 306 | video_fps=float(args.video_fps), |
| 307 | ) |
| 308 | |
| 309 | memory_bank = PairedAudioVideoMemoryBank( |
| 310 | max_size=int(args.memory_max_size), |
| 311 | save_mode=str(args.save_mode), |
| 312 | num_fix_frames=int(args.num_fix_frames), |
| 313 | ) |
| 314 | print(f"[Stage 2] Generator + VAEs ready.", flush=True) |
| 315 | |
| 316 | |
| 317 | shot_paths: list[Path] = [] |
| 318 | shot_audios: list[torch.Tensor] = [] |
| 319 | metadata: dict[str, Any] = { |
| 320 | "checkpoint": str(checkpoint), |
| 321 | "prompts_file": str(prompts_file), |
| 322 | "output_dir": str(output_dir), |
| 323 | "denoising_steps": [int(x) for x in args.denoising_steps], |
| 324 | "denoising_sigmas": [float(x) for x in denoising_sigmas.detach().cpu().tolist()], |
| 325 | "num_prompts": len(prompts), |
| 326 | "save_mode": str(args.save_mode), |
| 327 | "memory_max_size": int(args.memory_max_size), |
| 328 | "num_fix_frames": int(args.num_fix_frames), |
| 329 | "enable_audio_memory": bool(args.enable_audio_memory), |
| 330 | "shots": [], |
| 331 | } |
| 332 | |
| 333 | for shot_idx, prompt in enumerate(prompts): |
| 334 | conditional_dict = { |
| 335 | k: (v.to(device) if isinstance(v, torch.Tensor) else v) |
| 336 | for k, v in cached_conds[shot_idx].items() |
| 337 | } |
| 338 | prompt_seed = int(args.seed) + shot_idx |
| 339 | memory_size_before = len(memory_bank) |
| 340 | |
| 341 | print( |
| 342 | f"[Inference] shot={shot_idx + 1}/{len(prompts)} " |
| 343 | f"memory_size_before={memory_size_before} seed={prompt_seed}", |
| 344 | flush=True, |
| 345 | ) |
| 346 | |
| 347 | memory_video = None |
| 348 | memory_audio_kwargs: dict[str, Any] = {} |
| 349 | |
| 350 | with torch.random.fork_rng(devices=[device]): |
| 351 | torch.manual_seed(prompt_seed) |
| 352 | if device.type == "cuda": |
| 353 | torch.cuda.manual_seed(prompt_seed) |
| 354 | |
| 355 | if len(memory_bank) > 0: |
| 356 | memory_video = encode_memory_frames_batch( |
| 357 | video_vae=video_vae, |
| 358 | batch_memory_frames=[memory_bank.get_memory_frames()], |
| 359 | target_h=int(args.video_height), |
| 360 | target_w=int(args.video_width), |
| 361 | device=device, |
| 362 | dtype=dtype, |
| 363 | ) |
| 364 | memory_audio_kwargs = build_paired_audio_memory_kwargs( |
| 365 | memory_bank, |
| 366 | enable_audio_memory=bool(args.enable_audio_memory), |
| 367 | v2a_grad_scale=float(args.v2a_grad_scale), |
| 368 | memory_position_mode=str(args.memory_position_mode), |
| 369 | ) |
| 370 | |
| 371 | video_latent, audio_latent = memory_pipeline.generate( |
| 372 | video_shape=tuple(video_shape), |
| 373 | audio_shape=tuple(audio_shape), |
| 374 | conditional_dict=conditional_dict, |
| 375 | memory_video=memory_video, |
| 376 | seed=prompt_seed, |
| 377 | **memory_audio_kwargs, |
| 378 | ) |
| 379 | else: |
| 380 | video_latent, audio_latent = base_pipeline.generate( |
| 381 | video_shape=tuple(video_shape), |
| 382 | audio_shape=tuple(audio_shape), |
| 383 | conditional_dict=conditional_dict, |
| 384 | seed=prompt_seed, |
| 385 | ) |
| 386 | |
| 387 | audio_memory_latent = ( |
| 388 | audio_latent.detach().cpu().contiguous() if (args.enable_audio_memory and audio_latent is not None) else None |
| 389 | ) |
| 390 | video_uint8, audio_waveform = decode_benchmark_sample(video_vae, audio_vae, video_latent, audio_latent) |
| 391 | memory_frames_for_bank = video_uint8_to_pil_frames(video_uint8) |
| 392 | |
| 393 | new_memory_metadata: dict[str, Any] = {} |
| 394 | if audio_memory_latent is not None: |
| 395 | new_memory_metadata = memory_bank.save_memory_slot( |
| 396 | memory_frames_for_bank, |
| 397 | audio_memory_latent, |
| 398 | audio_window_size=int(args.audio_memory_window_size), |
| 399 | video_clip_num_frames=int(args.video_memory_clip_num_frames), |
| 400 | audio_waveform=audio_waveform, |
| 401 | audio_sample_rate=int(args.audio_memory_sample_rate), |
| 402 | video_fps=float(args.video_fps), |
| 403 | audio_window_selection_mode=str(args.audio_memory_window_selection_mode), |
| 404 | video_frame_selection_mode=str(args.video_memory_frame_selection_mode), |
| 405 | audio_memory_mel_bins=int(args.audio_memory_mel_bins), |
| 406 | audio_memory_mel_hop_length=int(args.audio_memory_mel_hop_length), |
| 407 | audio_memory_n_fft=int(args.audio_memory_n_fft), |
| 408 | audio_memory_downsample_factor=int(args.audio_memory_downsample_factor), |
| 409 | audio_memory_is_causal=bool(args.audio_memory_is_causal), |
| 410 | ) |
| 411 | |
| 412 | save_memory_bank_frames( |
| 413 | memory_bank.get_memory_frames(), |
| 414 | output_dir / "memory_bank" / f"shot_{shot_idx:03d}", |
| 415 | ) |
| 416 | |
| 417 | shot_path = output_dir / f"shot_{shot_idx:03d}.mp4" |
| 418 | write_result = write_benchmark_media( |
| 419 | output_path=shot_path, |
| 420 | video_uint8=video_uint8, |
| 421 | audio_waveform=audio_waveform, |
| 422 | fps=int(args.video_fps), |
| 423 | audio_sr=int(audio_sample_rate), |
| 424 | ) |
| 425 | shot_paths.append(shot_path) |
| 426 | if audio_waveform is not None: |
| 427 | shot_audios.append(audio_waveform.cpu()) |
| 428 | |
| 429 | metadata["shots"].append( |
| 430 | { |
| 431 | "shot_idx": int(shot_idx), |
| 432 | "prompt": prompt, |
| 433 | "output_path": str(shot_path), |
| 434 | "memory_size_before": int(memory_size_before), |
| 435 | "memory_size_after": int(len(memory_bank)), |
| 436 | "new_memory_entry": new_memory_metadata, |
| 437 | "audio_latent_shape": list(audio_latent.shape) if audio_latent is not None else None, |
| 438 | "wrote_audio_in_mp4": bool(write_result["wrote_audio_in_mp4"]), |
| 439 | "wrote_sidecar_wav": bool(write_result["wrote_sidecar_wav"]), |
| 440 | "audio_stats": write_result["audio_stats"], |
| 441 | "memory_entries": memory_bank.get_memory_metadata(), |
| 442 | } |
| 443 | ) |
| 444 | |
| 445 | del conditional_dict, video_latent, audio_latent, video_uint8, audio_waveform |
| 446 | del audio_memory_latent, memory_frames_for_bank, memory_video, memory_audio_kwargs |
| 447 | if device.type == "cuda": |
| 448 | torch.cuda.empty_cache() |
| 449 | |
| 450 | combined_path = output_dir / "combined_shots.mp4" |
| 451 | concat_shot_videos(shot_paths, combined_path) |
| 452 | combined_audio = concat_shot_audios(shot_audios) |
| 453 | combined_audio_path = None |
| 454 | if combined_audio is not None: |
| 455 | combined_audio_path = output_dir / "combined_shots.wav" |
| 456 | torchaudio.save(str(combined_audio_path), combined_audio, sample_rate=int(audio_sample_rate)) |
| 457 | |
| 458 | metadata["combined_path"] = str(combined_path) |
| 459 | metadata["combined_audio_path"] = str(combined_audio_path) if combined_audio_path is not None else None |
| 460 | metadata["combined_audio_stats"] = audio_waveform_stats(combined_audio) |
| 461 | metadata_path = output_dir / "run_metadata.json" |
| 462 | metadata_path.write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8") |
| 463 | |
| 464 | print(f"[Inference] done -> {combined_path}", flush=True) |
| 465 | |
| 466 | |
| 467 | if __name__ == "__main__": |
| 468 | main() |
| 469 |