| 1 | """Public BF16, FP8 and FP4 R2V DMD inference entrypoint for Echo 1.5.""" |
| 2 | |
| 3 | # ruff: noqa: E402 |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import gc |
| 8 | import json |
| 9 | import os |
| 10 | import sys |
| 11 | import time |
| 12 | from dataclasses import asdict |
| 13 | from datetime import datetime |
| 14 | from glob import glob |
| 15 | from pathlib import Path |
| 16 | from typing import Any |
| 17 | |
| 18 | |
| 19 | REPO_ROOT = Path(__file__).resolve().parent |
| 20 | for _subpath in ("ltx-core/src", "ltx-pipelines/src", "ltx-distillation/src"): |
| 21 | _package_path = str(REPO_ROOT / _subpath) |
| 22 | if _package_path not in sys.path: |
| 23 | sys.path.insert(0, _package_path) |
| 24 | |
| 25 | import torch |
| 26 | import yaml |
| 27 | |
| 28 | from r2v_schema import MAX_MEMORY_SLOTS, R2VRequest, load_r2v_request |
| 29 | from ltx_distillation.audio_voice_filter import VoiceFilterConfig |
| 30 | from ltx_distillation.generator_loader import ( |
| 31 | GeneratorLoadReport, |
| 32 | load_inference_generator, |
| 33 | ) |
| 34 | from ltx_distillation.inference.memory_bidirectional_pipeline import ( |
| 35 | BidirectionalR2VInferencePipeline, |
| 36 | ) |
| 37 | from ltx_distillation.layerwise_offload import DiTLayerwiseOffload |
| 38 | from ltx_distillation.models.vae_wrapper import create_vae_wrappers |
| 39 | from ltx_distillation.r2v_conditioning import ( |
| 40 | R2VConditionBundle, |
| 41 | encode_r2v_requests, |
| 42 | load_r2v_conditioning, |
| 43 | r2v_conditioning_cache_path, |
| 44 | save_r2v_conditioning, |
| 45 | ) |
| 46 | from ltx_distillation.release_checkpoint import ( |
| 47 | ReleaseCheckpoint, |
| 48 | resolve_release_checkpoint, |
| 49 | ) |
| 50 | from ltx_distillation.text_conditioning import artifact_fingerprint |
| 51 | from ltx_distillation.utils import ( |
| 52 | add_noise, |
| 53 | compute_latent_shapes, |
| 54 | decode_generated_sample, |
| 55 | write_generated_media, |
| 56 | ) |
| 57 | from ltx_core.model.video_vae.tiling import ( |
| 58 | SpatialTilingConfig, |
| 59 | TemporalTilingConfig, |
| 60 | TilingConfig, |
| 61 | ) |
| 62 | |
| 63 | |
| 64 | DEFAULT_CONFIG = REPO_ROOT / "configs" / "inference.bf16.yaml" |
| 65 | |
| 66 | |
| 67 | def _load_yaml_config(config_path: Path) -> dict[str, Any]: |
| 68 | with config_path.open("r", encoding="utf-8") as handle: |
| 69 | return yaml.safe_load(handle) or {} |
| 70 | |
| 71 | |
| 72 | def _resolve_path(path_value: str, *, required: bool = True) -> str | None: |
| 73 | value = str(path_value or "").strip() |
| 74 | if not value: |
| 75 | if required: |
| 76 | raise ValueError("required path is empty") |
| 77 | return None |
| 78 | path = Path(value).expanduser() |
| 79 | if not path.is_absolute(): |
| 80 | path = REPO_ROOT / path |
| 81 | return str(path.resolve()) |
| 82 | |
| 83 | |
| 84 | def str_to_bool(value: str | bool) -> bool: |
| 85 | if isinstance(value, bool): |
| 86 | return value |
| 87 | normalized = value.strip().lower() |
| 88 | if normalized in {"1", "true", "yes", "y"}: |
| 89 | return True |
| 90 | if normalized in {"0", "false", "no", "n"}: |
| 91 | return False |
| 92 | raise ValueError(f"invalid boolean value: {value}") |
| 93 | |
| 94 | |
| 95 | def resolve_video_vae_decode_mode( |
| 96 | requested_mode: str, |
| 97 | *, |
| 98 | device_type: str, |
| 99 | platform_name: str, |
| 100 | total_memory_bytes: int | None, |
| 101 | ) -> str: |
| 102 | """Resolve ``auto`` without coupling decode policy to checkpoint precision.""" |
| 103 | mode = requested_mode.strip().lower() |
| 104 | if mode not in {"auto", "tiled", "untiled"}: |
| 105 | raise ValueError("video_vae.decode_mode must be auto, tiled, or untiled") |
| 106 | if mode != "auto": |
| 107 | return mode |
| 108 | if device_type == "cuda" and platform_name == "win32": |
| 109 | return "tiled" |
| 110 | if device_type == "cuda" and ( |
| 111 | total_memory_bytes is None or total_memory_bytes < 48 * 2**30 |
| 112 | ): |
| 113 | return "tiled" |
| 114 | return "untiled" |
| 115 | |
| 116 | |
| 117 | class InferenceConfig: |
| 118 | """Validated YAML configuration with optional CLI overrides.""" |
| 119 | |
| 120 | def __init__(self, config_path: Path, **cli_overrides: Any) -> None: |
| 121 | config = _load_yaml_config(config_path) |
| 122 | paths = config.get("paths", {}) |
| 123 | video = config.get("video", {}) |
| 124 | denoising = config.get("denoising", {}) |
| 125 | memory = config.get("memory", {}) |
| 126 | voice_filter = memory.get("voice_filter", {}) |
| 127 | runtime = config.get("inference", {}) |
| 128 | video_vae = config.get("video_vae", {}) or {} |
| 129 | self.checkpoint = _resolve_path( |
| 130 | paths.get("checkpoint", "checkpoints/echo15_full_dmd") |
| 131 | ) |
| 132 | self.gemma_path = _resolve_path( |
| 133 | paths.get("gemma_path", "checkpoints/gemma-3-12b") |
| 134 | ) |
| 135 | self.requests_dir = _resolve_path( |
| 136 | paths.get("requests_dir", "examples/the_last_visa/requests") |
| 137 | ) |
| 138 | self.requests_glob = str(paths.get("requests_glob", "*.json")) |
| 139 | self.output_root = _resolve_path(paths.get("output_root", "inference_result")) |
| 140 | self.conditioning_cache_dir = _resolve_path( |
| 141 | paths.get("conditioning_cache_dir", ""), required=False |
| 142 | ) |
| 143 | |
| 144 | self.num_frames = int(video.get("num_frames", 241)) |
| 145 | self.video_height = int(video.get("height", 736)) |
| 146 | self.video_width = int(video.get("width", 1280)) |
| 147 | self.video_fps = int(video.get("fps", 25)) |
| 148 | self.seed = int(video.get("seed", 42)) |
| 149 | |
| 150 | self.video_vae_decode_mode = str(video_vae.get("decode_mode", "tiled")).lower() |
| 151 | self.video_vae_tile_size_frames = int(video_vae.get("tile_size_frames", 64)) |
| 152 | self.video_vae_tile_overlap_frames = int( |
| 153 | video_vae.get("tile_overlap_frames", 24) |
| 154 | ) |
| 155 | self.video_vae_tile_size_pixels = int(video_vae.get("tile_size_pixels", 512)) |
| 156 | self.video_vae_tile_overlap_pixels = int( |
| 157 | video_vae.get("tile_overlap_pixels", 64) |
| 158 | ) |
| 159 | |
| 160 | self.denoising_steps = [int(value) for value in denoising.get("steps", [])] |
| 161 | self.denoising_sigmas = [float(value) for value in denoising.get("sigmas", [])] |
| 162 | |
| 163 | self.memory_max_size = int(memory.get("max_size", MAX_MEMORY_SLOTS)) |
| 164 | self.memory_downscale_factor = int(memory.get("downscale_factor", 1)) |
| 165 | self.memory_position_mode = str(memory.get("position_mode", "slot_center")) |
| 166 | self.memory_position_offset = float(memory.get("position_offset", 500.0)) |
| 167 | self.memory_position_slot_stride = float( |
| 168 | memory.get("position_slot_stride", 50.0) |
| 169 | ) |
| 170 | self.enable_audio_memory = bool(memory.get("enable_audio", True)) |
| 171 | self.voice_filter = VoiceFilterConfig( |
| 172 | enabled=bool(voice_filter.get("enabled", True)), |
| 173 | backend=str(voice_filter.get("backend", "msst_speech")), |
| 174 | min_output_rms=float(voice_filter.get("min_output_rms", 0.004)), |
| 175 | msst_dir=str( |
| 176 | _resolve_path(voice_filter.get("msst_dir", "third_party/MSST-WebUI")) |
| 177 | ), |
| 178 | msst_model_path=str( |
| 179 | _resolve_path( |
| 180 | voice_filter.get( |
| 181 | "msst_model_path", |
| 182 | "checkpoints/msst/model_bandit_plus_dnr_sdr_11.47.chpt", |
| 183 | ) |
| 184 | ) |
| 185 | ), |
| 186 | msst_config_path=str( |
| 187 | _resolve_path( |
| 188 | voice_filter.get( |
| 189 | "msst_config_path", |
| 190 | "third_party/MSST-WebUI/configs_backup/multi_stem_models/" |
| 191 | "model_bandit_plus_dnr_sdr_11.47.chpt.yaml", |
| 192 | ) |
| 193 | ) |
| 194 | ), |
| 195 | msst_model_type=str(voice_filter.get("msst_model_type", "bandit")), |
| 196 | msst_sample_rate=int(voice_filter.get("msst_sample_rate", 44100)), |
| 197 | msst_device=str(voice_filter.get("msst_device", "auto")), |
| 198 | msst_local_rank_env=str( |
| 199 | voice_filter.get("msst_local_rank_env", "LOCAL_RANK") |
| 200 | ), |
| 201 | ) |
| 202 | |
| 203 | self.device = str(runtime.get("device", "cuda")) |
| 204 | self.dtype = str(runtime.get("dtype", "bfloat16")).lower() |
| 205 | self.prompt_max_chars = int(runtime.get("prompt_max_chars", 1500)) |
| 206 | self.text_batch_size = int(runtime.get("text_batch_size", 1)) |
| 207 | self.image_batch_size = int(runtime.get("image_batch_size", 1)) |
| 208 | self.audio_batch_size = int(runtime.get("audio_batch_size", 1)) |
| 209 | dit_offload = runtime.get("dit_layerwise_offload", {}) or {} |
| 210 | self.dit_layerwise_offload = bool(dit_offload.get("enabled", False)) |
| 211 | self.dit_resident_blocks = int(dit_offload.get("resident_blocks", 0)) |
| 212 | self.dit_prefetch_blocks = int(dit_offload.get("prefetch_blocks", 1)) |
| 213 | self.dit_pin_memory = bool(dit_offload.get("pin_memory", True)) |
| 214 | |
| 215 | for key, value in cli_overrides.items(): |
| 216 | if value is not None and hasattr(self, key): |
| 217 | setattr(self, key, value) |
| 218 | self.validate() |
| 219 | |
| 220 | def validate(self) -> None: |
| 221 | if self.dtype not in {"bfloat16", "bf16"}: |
| 222 | raise ValueError("Echo 1.5 uses BF16 activations for every checkpoint mode") |
| 223 | if len(self.denoising_sigmas) < 2: |
| 224 | raise ValueError("denoising.sigmas must contain at least two values") |
| 225 | if self.denoising_steps and len(self.denoising_steps) != len( |
| 226 | self.denoising_sigmas |
| 227 | ): |
| 228 | raise ValueError( |
| 229 | "denoising.steps and denoising.sigmas must have equal length" |
| 230 | ) |
| 231 | if self.memory_position_mode != "slot_center": |
| 232 | raise ValueError("Echo 1.5 R2V requires memory.position_mode=slot_center") |
| 233 | if not 0 <= self.memory_max_size <= MAX_MEMORY_SLOTS: |
| 234 | raise ValueError( |
| 235 | f"memory.max_size must be between 0 and {MAX_MEMORY_SLOTS}" |
| 236 | ) |
| 237 | if self.num_frames <= 0 or self.video_height <= 0 or self.video_width <= 0: |
| 238 | raise ValueError("video dimensions and frame count must be positive") |
| 239 | if self.prompt_max_chars <= 0: |
| 240 | raise ValueError("inference.prompt_max_chars must be positive") |
| 241 | if min(self.text_batch_size, self.image_batch_size, self.audio_batch_size) <= 0: |
| 242 | raise ValueError("conditioning batch sizes must be positive") |
| 243 | if self.video_vae_decode_mode not in {"auto", "tiled", "untiled"}: |
| 244 | raise ValueError("video_vae.decode_mode must be auto, tiled, or untiled") |
| 245 | # Reuse the core validators so public configuration follows exactly the |
| 246 | # same alignment constraints as the decoder implementation. |
| 247 | self.video_vae_tiling_config() |
| 248 | if self.dit_resident_blocks < 0: |
| 249 | raise ValueError( |
| 250 | "inference.dit_layerwise_offload.resident_blocks must be non-negative" |
| 251 | ) |
| 252 | if self.dit_prefetch_blocks < 1: |
| 253 | raise ValueError( |
| 254 | "inference.dit_layerwise_offload.prefetch_blocks must be positive" |
| 255 | ) |
| 256 | |
| 257 | def video_vae_tiling_config(self) -> TilingConfig: |
| 258 | return TilingConfig( |
| 259 | spatial_config=SpatialTilingConfig( |
| 260 | tile_size_in_pixels=self.video_vae_tile_size_pixels, |
| 261 | tile_overlap_in_pixels=self.video_vae_tile_overlap_pixels, |
| 262 | ), |
| 263 | temporal_config=TemporalTilingConfig( |
| 264 | tile_size_in_frames=self.video_vae_tile_size_frames, |
| 265 | tile_overlap_in_frames=self.video_vae_tile_overlap_frames, |
| 266 | ), |
| 267 | ) |
| 268 | |
| 269 | |
| 270 | class InferenceEngine: |
| 271 | """Prepare complete R2V conditions first, then run the shared DMD pipeline.""" |
| 272 | |
| 273 | def __init__(self, config: InferenceConfig) -> None: |
| 274 | self.config = config |
| 275 | self.device = torch.device(config.device) |
| 276 | self.dtype = torch.bfloat16 |
| 277 | self.release_checkpoint: ReleaseCheckpoint = resolve_release_checkpoint( |
| 278 | config.checkpoint |
| 279 | ) |
| 280 | self.model_checkpoint = self.release_checkpoint.model_path |
| 281 | self.gemma_path = Path(config.gemma_path) |
| 282 | if not self.gemma_path.exists(): |
| 283 | raise FileNotFoundError(f"Gemma directory not found: {self.gemma_path}") |
| 284 | |
| 285 | self.generator = None |
| 286 | self.video_vae = None |
| 287 | self.audio_vae = None |
| 288 | self.pipeline = None |
| 289 | self.audio_sample_rate: int | None = None |
| 290 | self.generator_load_report: GeneratorLoadReport | None = None |
| 291 | self.dit_offload: DiTLayerwiseOffload | None = None |
| 292 | self.generator_location = "unloaded" |
| 293 | self.decoder_location = "unloaded" |
| 294 | total_memory_bytes = None |
| 295 | if self.device.type == "cuda" and torch.cuda.is_available(): |
| 296 | total_memory_bytes = torch.cuda.get_device_properties( |
| 297 | self.device |
| 298 | ).total_memory |
| 299 | self.video_vae_decode_mode = resolve_video_vae_decode_mode( |
| 300 | config.video_vae_decode_mode, |
| 301 | device_type=self.device.type, |
| 302 | platform_name=sys.platform, |
| 303 | total_memory_bytes=total_memory_bytes, |
| 304 | ) |
| 305 | self.video_vae_tiling_config = ( |
| 306 | config.video_vae_tiling_config() |
| 307 | if self.video_vae_decode_mode == "tiled" |
| 308 | else None |
| 309 | ) |
| 310 | print( |
| 311 | f"[VAE] video decode={self.video_vae_decode_mode} " |
| 312 | f"(configured={config.video_vae_decode_mode})", |
| 313 | flush=True, |
| 314 | ) |
| 315 | |
| 316 | def prepare_conditions( |
| 317 | self, |
| 318 | request_files: list[Path], |
| 319 | requests: list[R2VRequest], |
| 320 | *, |
| 321 | encode_cache_misses: bool = False, |
| 322 | ) -> dict[Path, R2VConditionBundle]: |
| 323 | config = self.config |
| 324 | checkpoint_id = artifact_fingerprint(self.release_checkpoint.root) |
| 325 | gemma_id = artifact_fingerprint(self.gemma_path) |
| 326 | cached: dict[Path, R2VConditionBundle] = {} |
| 327 | pending_files: list[Path] = [] |
| 328 | pending_requests: list[R2VRequest] = [] |
| 329 | if config.conditioning_cache_dir: |
| 330 | print("[Stage 1] Loading precomputed R2V conditioning", flush=True) |
| 331 | for request_file, request in zip(request_files, requests, strict=True): |
| 332 | cache_path = r2v_conditioning_cache_path( |
| 333 | config.conditioning_cache_dir, |
| 334 | config.requests_dir, |
| 335 | request_file, |
| 336 | ) |
| 337 | try: |
| 338 | cached[request_file] = load_r2v_conditioning( |
| 339 | cache_path, |
| 340 | request=request, |
| 341 | checkpoint_fingerprint=checkpoint_id, |
| 342 | gemma_fingerprint=gemma_id, |
| 343 | ) |
| 344 | except (FileNotFoundError, ValueError): |
| 345 | if not encode_cache_misses: |
| 346 | raise |
| 347 | pending_files.append(request_file) |
| 348 | pending_requests.append(request) |
| 349 | if not pending_requests: |
| 350 | return cached |
| 351 | else: |
| 352 | pending_files = list(request_files) |
| 353 | pending_requests = list(requests) |
| 354 | |
| 355 | print( |
| 356 | f"[Stage 1] Encoding {len(pending_requests)} R2V conditions " |
| 357 | f"text_batch={config.text_batch_size} image_batch={config.image_batch_size} " |
| 358 | f"audio_batch={config.audio_batch_size}", |
| 359 | flush=True, |
| 360 | ) |
| 361 | bundles = encode_r2v_requests( |
| 362 | pending_requests, |
| 363 | checkpoint_path=str(self.model_checkpoint), |
| 364 | gemma_path=str(self.gemma_path), |
| 365 | device=self.device, |
| 366 | voice_filter_config=config.voice_filter, |
| 367 | dtype=self.dtype, |
| 368 | text_batch_size=config.text_batch_size, |
| 369 | image_batch_size=config.image_batch_size, |
| 370 | audio_batch_size=config.audio_batch_size, |
| 371 | enable_audio_memory=config.enable_audio_memory, |
| 372 | memory_position_mode=config.memory_position_mode, |
| 373 | memory_position_offset=config.memory_position_offset, |
| 374 | memory_position_slot_stride=config.memory_position_slot_stride, |
| 375 | ) |
| 376 | for request_file, request, bundle in zip( |
| 377 | pending_files, pending_requests, bundles, strict=True |
| 378 | ): |
| 379 | cached[request_file] = bundle |
| 380 | if config.conditioning_cache_dir: |
| 381 | save_r2v_conditioning( |
| 382 | r2v_conditioning_cache_path( |
| 383 | config.conditioning_cache_dir, |
| 384 | config.requests_dir, |
| 385 | request_file, |
| 386 | ), |
| 387 | bundle, |
| 388 | request=request, |
| 389 | checkpoint_fingerprint=checkpoint_id, |
| 390 | gemma_fingerprint=gemma_id, |
| 391 | ) |
| 392 | print("[Stage 1] Complete DiT conditioning ready on CPU", flush=True) |
| 393 | return {request_file: cached[request_file] for request_file in request_files} |
| 394 | |
| 395 | def load_generator(self) -> None: |
| 396 | if self.generator is not None: |
| 397 | return |
| 398 | config = self.config |
| 399 | print( |
| 400 | f"[Stage 2] Loading {self.release_checkpoint.name} " |
| 401 | f"precision={self.release_checkpoint.precision}", |
| 402 | flush=True, |
| 403 | ) |
| 404 | # On Windows, reopening a large BF16 safetensors file for the VAE after |
| 405 | # materializing a large CPU weight pool can crash inside torch_cpu.dll. |
| 406 | # Build the decoder modules before attaching the layerwise offload pool. |
| 407 | precision = self.release_checkpoint.precision |
| 408 | preload_vaes = config.dit_layerwise_offload and precision in {"bf16", "fp4"} |
| 409 | if preload_vaes and (self.video_vae is None or self.audio_vae is None): |
| 410 | print( |
| 411 | f"[Stage 2] Preloading VAEs before {precision.upper()} generator restore", |
| 412 | flush=True, |
| 413 | ) |
| 414 | self.video_vae, self.audio_vae = create_vae_wrappers( |
| 415 | checkpoint_path=str(self.model_checkpoint), |
| 416 | device=torch.device("cpu"), |
| 417 | dtype=self.dtype, |
| 418 | with_video_encoder=False, |
| 419 | with_audio_encoder=False, |
| 420 | decoder_device=torch.device("cpu"), |
| 421 | ) |
| 422 | |
| 423 | self.generator, self.generator_load_report = load_inference_generator( |
| 424 | checkpoint=self.release_checkpoint, |
| 425 | gemma_path=self.gemma_path, |
| 426 | device=self.device, |
| 427 | dtype=self.dtype, |
| 428 | video_height=config.video_height, |
| 429 | video_width=config.video_width, |
| 430 | load_on_cpu=config.dit_layerwise_offload, |
| 431 | ) |
| 432 | if config.dit_layerwise_offload: |
| 433 | self.dit_offload = DiTLayerwiseOffload( |
| 434 | self.generator, |
| 435 | execution_device=self.device, |
| 436 | resident_blocks=config.dit_resident_blocks, |
| 437 | prefetch_blocks=config.dit_prefetch_blocks, |
| 438 | pin_memory=config.dit_pin_memory, |
| 439 | ) |
| 440 | report = self.dit_offload.report |
| 441 | print( |
| 442 | f"[Stage 2] DiT layerwise offload blocks={report.block_count} " |
| 443 | f"resident={report.resident_blocks} prefetch={report.prefetch_blocks} " |
| 444 | f"cpu_weights={report.cpu_weight_bytes / 2**30:.2f}GiB " |
| 445 | f"pinned={report.pinned_weight_bytes / 2**30:.2f}GiB", |
| 446 | flush=True, |
| 447 | ) |
| 448 | print( |
| 449 | f"[Stage 2] format={self.generator_load_report.format} " |
| 450 | f"quantized_modules={self.generator_load_report.quantized_modules} " |
| 451 | f"missing={len(self.generator_load_report.missing_keys)} " |
| 452 | f"unexpected={len(self.generator_load_report.unexpected_keys)}", |
| 453 | flush=True, |
| 454 | ) |
| 455 | |
| 456 | if self.video_vae is None or self.audio_vae is None: |
| 457 | self.video_vae, self.audio_vae = create_vae_wrappers( |
| 458 | checkpoint_path=str(self.model_checkpoint), |
| 459 | device=torch.device("cpu"), |
| 460 | dtype=self.dtype, |
| 461 | with_video_encoder=False, |
| 462 | with_audio_encoder=False, |
| 463 | decoder_device=torch.device("cpu"), |
| 464 | ) |
| 465 | self.video_vae.eval() |
| 466 | self.audio_vae.eval() |
| 467 | sigmas = torch.tensor( |
| 468 | config.denoising_sigmas, device=self.device, dtype=torch.float32 |
| 469 | ) |
| 470 | self.pipeline = BidirectionalR2VInferencePipeline( |
| 471 | self.generator, |
| 472 | add_noise, |
| 473 | sigmas, |
| 474 | memory_downscale_factor=config.memory_downscale_factor, |
| 475 | ) |
| 476 | self.audio_sample_rate = self.audio_vae.get_output_sample_rate() or 24000 |
| 477 | self.generator_location = ( |
| 478 | "cpu" if config.dit_layerwise_offload else str(self.device) |
| 479 | ) |
| 480 | self.decoder_location = "cpu" |
| 481 | print("[Stage 2] Generator and decode VAEs ready", flush=True) |
| 482 | |
| 483 | def unload_generator(self) -> None: |
| 484 | """Release all resident generation/decode weights owned by this engine.""" |
| 485 | |
| 486 | self.release_generator_weights() |
| 487 | self.video_vae = None |
| 488 | self.audio_vae = None |
| 489 | self.audio_sample_rate = None |
| 490 | self.generator_load_report = None |
| 491 | self.decoder_location = "unloaded" |
| 492 | gc.collect() |
| 493 | self._empty_cuda_cache() |
| 494 | |
| 495 | def release_generator_weights(self) -> None: |
| 496 | """Release only DiT/pipeline weights while preserving loaded decoders.""" |
| 497 | |
| 498 | if self.dit_offload is not None: |
| 499 | self.dit_offload.close() |
| 500 | self.pipeline = None |
| 501 | self.dit_offload = None |
| 502 | self.generator = None |
| 503 | self.generator_location = "unloaded" |
| 504 | gc.collect() |
| 505 | self._empty_cuda_cache() |
| 506 | |
| 507 | def generation_storage_bytes(self) -> int: |
| 508 | """Approximate unique tensor storage retained by generation modules.""" |
| 509 | |
| 510 | storages: dict[tuple[str, int], int] = {} |
| 511 | for module in (self.generator, self.video_vae, self.audio_vae): |
| 512 | if module is None: |
| 513 | continue |
| 514 | for tensor in (*module.parameters(), *module.buffers()): |
| 515 | storage = tensor.untyped_storage() |
| 516 | key = (str(tensor.device), int(storage.data_ptr())) |
| 517 | storages.setdefault(key, int(storage.nbytes())) |
| 518 | return sum(storages.values()) |
| 519 | |
| 520 | def stage_generator_for_conditioning(self, policy: str) -> None: |
| 521 | """Place warm generation weights for a conditioning cache miss.""" |
| 522 | |
| 523 | if policy not in {"gpu", "cpu", "release"}: |
| 524 | raise ValueError(f"unsupported conditioning generator policy: {policy}") |
| 525 | if policy == "release": |
| 526 | self.unload_generator() |
| 527 | return |
| 528 | if self.generator is None: |
| 529 | return |
| 530 | if policy == "gpu": |
| 531 | return |
| 532 | if self.dit_offload is not None: |
| 533 | self.dit_offload.deactivate() |
| 534 | else: |
| 535 | self._move(self.generator, "cpu") |
| 536 | if self.video_vae is not None: |
| 537 | self._move(self.video_vae, "cpu") |
| 538 | if self.audio_vae is not None: |
| 539 | self._move(self.audio_vae, "cpu") |
| 540 | self.generator_location = "cpu" |
| 541 | self.decoder_location = "cpu" |
| 542 | self._empty_cuda_cache() |
| 543 | |
| 544 | @staticmethod |
| 545 | def _move(module, target_device: str | torch.device) -> None: |
| 546 | if module is not None: |
| 547 | module.to(target_device) |
| 548 | |
| 549 | def _empty_cuda_cache(self) -> None: |
| 550 | if self.device.type == "cuda": |
| 551 | torch.cuda.empty_cache() |
| 552 | |
| 553 | def _stage_for_denoise(self) -> None: |
| 554 | already_co_resident = self.generator_location.startswith( |
| 555 | "cuda" |
| 556 | ) and self.decoder_location.startswith("cuda") |
| 557 | if not already_co_resident: |
| 558 | self._move(self.video_vae.decoder, "cpu") |
| 559 | self._move(self.audio_vae.decoder, "cpu") |
| 560 | self._move(self.audio_vae.vocoder, "cpu") |
| 561 | self.decoder_location = "cpu" |
| 562 | self._empty_cuda_cache() |
| 563 | if self.dit_offload is not None: |
| 564 | self.dit_offload.activate() |
| 565 | self.generator_location = "layerwise" |
| 566 | else: |
| 567 | self._move(self.generator, self.device) |
| 568 | self.generator_location = str(self.device) |
| 569 | |
| 570 | def _stage_for_decode(self, generator_policy: str = "cpu") -> str: |
| 571 | if generator_policy not in {"gpu", "cpu", "release"}: |
| 572 | raise ValueError(f"unsupported decode generator policy: {generator_policy}") |
| 573 | if generator_policy == "release": |
| 574 | self.release_generator_weights() |
| 575 | elif self.dit_offload is not None: |
| 576 | self.dit_offload.deactivate() |
| 577 | elif generator_policy == "cpu": |
| 578 | self._move(self.generator, "cpu") |
| 579 | self.generator_location = ( |
| 580 | "cpu" |
| 581 | if self.dit_offload is not None or generator_policy == "cpu" |
| 582 | else self.generator_location |
| 583 | ) |
| 584 | self._empty_cuda_cache() |
| 585 | try: |
| 586 | self._move(self.video_vae.decoder, self.device) |
| 587 | self._move(self.audio_vae.decoder, self.device) |
| 588 | self._move(self.audio_vae.vocoder, self.device) |
| 589 | except torch.OutOfMemoryError: |
| 590 | if generator_policy != "gpu": |
| 591 | raise |
| 592 | # Live free-memory sampling is advisory. If another process races |
| 593 | # us, release DiT and retry decoder placement without failing the job. |
| 594 | self._move(self.video_vae, "cpu") |
| 595 | self._move(self.audio_vae, "cpu") |
| 596 | self.release_generator_weights() |
| 597 | self._move(self.video_vae.decoder, self.device) |
| 598 | self._move(self.audio_vae.decoder, self.device) |
| 599 | self._move(self.audio_vae.vocoder, self.device) |
| 600 | generator_policy = "release" |
| 601 | self.decoder_location = str(self.device) |
| 602 | return generator_policy |
| 603 | |
| 604 | def run_request( |
| 605 | self, |
| 606 | request_file: Path, |
| 607 | request: R2VRequest, |
| 608 | output_dir: Path, |
| 609 | bundle: R2VConditionBundle, |
| 610 | *, |
| 611 | stage_callback=None, |
| 612 | decode_generator_policy=None, |
| 613 | ) -> dict[str, Any]: |
| 614 | if ( |
| 615 | self.generator is None |
| 616 | or self.pipeline is None |
| 617 | or self.audio_sample_rate is None |
| 618 | ): |
| 619 | raise RuntimeError("load_generator() must be called before inference") |
| 620 | if len(request.memory_slots) > self.config.memory_max_size: |
| 621 | raise ValueError( |
| 622 | f"request has {len(request.memory_slots)} memory slots; configured maximum is " |
| 623 | f"{self.config.memory_max_size}" |
| 624 | ) |
| 625 | output_dir.mkdir(parents=True, exist_ok=True) |
| 626 | video_shape, audio_shape = compute_latent_shapes( |
| 627 | num_frames=request.num_frames, |
| 628 | video_height=request.height, |
| 629 | video_width=request.width, |
| 630 | batch_size=1, |
| 631 | video_fps=float(self.config.video_fps), |
| 632 | ) |
| 633 | latent_height, latent_width = int(video_shape[-2]), int(video_shape[-1]) |
| 634 | self.generator.latent_height = latent_height |
| 635 | self.generator.latent_width = latent_width |
| 636 | self.generator.video_frame_seqlen = latent_height * latent_width |
| 637 | |
| 638 | condition = { |
| 639 | key: value.to(self.device) if isinstance(value, torch.Tensor) else value |
| 640 | for key, value in bundle.text.items() |
| 641 | } |
| 642 | memory_audio_kwargs = { |
| 643 | key: value.to(self.device) if isinstance(value, torch.Tensor) else value |
| 644 | for key, value in bundle.memory_audio_kwargs.items() |
| 645 | } |
| 646 | if stage_callback is not None: |
| 647 | stage_callback("inferring") |
| 648 | self._stage_for_denoise() |
| 649 | started = time.perf_counter() |
| 650 | denoise_started = time.perf_counter() |
| 651 | fork_devices = ( |
| 652 | [ |
| 653 | self.device.index |
| 654 | if self.device.index is not None |
| 655 | else torch.cuda.current_device() |
| 656 | ] |
| 657 | if self.device.type == "cuda" |
| 658 | else [] |
| 659 | ) |
| 660 | with torch.random.fork_rng(devices=fork_devices): |
| 661 | torch.manual_seed(request.seed) |
| 662 | if self.device.type == "cuda": |
| 663 | torch.cuda.manual_seed(request.seed) |
| 664 | video_latent, audio_latent = self.pipeline.generate( |
| 665 | video_shape=tuple(video_shape), |
| 666 | audio_shape=tuple(audio_shape), |
| 667 | conditional_dict=condition, |
| 668 | memory_video=bundle.memory_video, |
| 669 | first_frame_latent=bundle.first_frame_latent, |
| 670 | seed=request.seed, |
| 671 | **memory_audio_kwargs, |
| 672 | ) |
| 673 | if self.device.type == "cuda": |
| 674 | torch.cuda.synchronize() |
| 675 | denoise_seconds = time.perf_counter() - denoise_started |
| 676 | |
| 677 | if stage_callback is not None: |
| 678 | stage_callback("decoding") |
| 679 | generator_policy = ( |
| 680 | str(decode_generator_policy()) |
| 681 | if decode_generator_policy is not None |
| 682 | else "cpu" |
| 683 | ) |
| 684 | generator_policy = self._stage_for_decode(generator_policy) |
| 685 | decode_started = time.perf_counter() |
| 686 | video_uint8, audio_waveform = decode_generated_sample( |
| 687 | self.video_vae, |
| 688 | self.audio_vae, |
| 689 | video_latent, |
| 690 | audio_latent, |
| 691 | video_tiling_config=self.video_vae_tiling_config, |
| 692 | ) |
| 693 | if self.device.type == "cuda": |
| 694 | torch.cuda.synchronize() |
| 695 | decode_seconds = time.perf_counter() - decode_started |
| 696 | if stage_callback is not None: |
| 697 | stage_callback("writing") |
| 698 | output_path = output_dir / "result.mp4" |
| 699 | write_result = write_generated_media( |
| 700 | output_path=output_path, |
| 701 | video_uint8=video_uint8, |
| 702 | audio_waveform=audio_waveform, |
| 703 | fps=self.config.video_fps, |
| 704 | audio_sr=self.audio_sample_rate, |
| 705 | ) |
| 706 | metadata = { |
| 707 | "schema": "echo15.r2v.result.v1", |
| 708 | "request_file": str(request_file), |
| 709 | "request": request.as_payload(), |
| 710 | "model_checkpoint": str(self.release_checkpoint.root), |
| 711 | "generator": asdict(self.generator_load_report), |
| 712 | "dit_layerwise_offload": ( |
| 713 | asdict(self.dit_offload.report) |
| 714 | if self.dit_offload |
| 715 | else {"enabled": False} |
| 716 | ), |
| 717 | "decode_generator_policy": generator_policy, |
| 718 | "video_vae": { |
| 719 | "decode_mode": self.video_vae_decode_mode, |
| 720 | "tile_size_frames": self.config.video_vae_tile_size_frames, |
| 721 | "tile_overlap_frames": self.config.video_vae_tile_overlap_frames, |
| 722 | "tile_size_pixels": self.config.video_vae_tile_size_pixels, |
| 723 | "tile_overlap_pixels": self.config.video_vae_tile_overlap_pixels, |
| 724 | }, |
| 725 | "conditioning": { |
| 726 | "memory_slots": len(request.memory_slots), |
| 727 | "has_first_frame": bundle.first_frame_latent is not None, |
| 728 | "has_memory_audio": "memory_audio" in bundle.memory_audio_kwargs, |
| 729 | "input_fingerprints": bundle.input_fingerprints, |
| 730 | }, |
| 731 | "output_path": str(output_path), |
| 732 | "audio_latent_shape": list(audio_latent.shape) |
| 733 | if audio_latent is not None |
| 734 | else None, |
| 735 | "audio_stats": write_result["audio_stats"], |
| 736 | "timing": { |
| 737 | "denoise_seconds": round(denoise_seconds, 3), |
| 738 | "decode_seconds": round(decode_seconds, 3), |
| 739 | "total_seconds": round(time.perf_counter() - started, 3), |
| 740 | }, |
| 741 | } |
| 742 | (output_dir / "run_metadata.json").write_text( |
| 743 | json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8" |
| 744 | ) |
| 745 | print( |
| 746 | f"[Inference] {request.shot_id} done denoise={denoise_seconds:.1f}s " |
| 747 | f"decode={decode_seconds:.1f}s output={output_path}", |
| 748 | flush=True, |
| 749 | ) |
| 750 | del video_latent, audio_latent, video_uint8, audio_waveform, condition |
| 751 | self._empty_cuda_cache() |
| 752 | return {"output_path": str(output_path), "metadata": metadata} |
| 753 | |
| 754 | |
| 755 | def load_request_files( |
| 756 | config: InferenceConfig, single_request: str | None = None |
| 757 | ) -> list[Path]: |
| 758 | if single_request: |
| 759 | path = Path(single_request).expanduser().resolve() |
| 760 | if not path.is_file(): |
| 761 | raise FileNotFoundError(f"R2V request not found: {path}") |
| 762 | return [path] |
| 763 | requests_dir = Path(config.requests_dir) |
| 764 | if Path(config.requests_glob).is_absolute(): |
| 765 | files = sorted( |
| 766 | Path(value) for value in glob(config.requests_glob, recursive=True) |
| 767 | ) |
| 768 | else: |
| 769 | files = sorted(requests_dir.glob(config.requests_glob)) |
| 770 | if not files: |
| 771 | raise FileNotFoundError( |
| 772 | f"no R2V requests matched {requests_dir / config.requests_glob}" |
| 773 | ) |
| 774 | return files |
| 775 | |
| 776 | |
| 777 | def load_requests( |
| 778 | config: InferenceConfig, request_files: list[Path] |
| 779 | ) -> list[R2VRequest]: |
| 780 | return [ |
| 781 | load_r2v_request( |
| 782 | path, |
| 783 | default_num_frames=config.num_frames, |
| 784 | default_width=config.video_width, |
| 785 | default_height=config.video_height, |
| 786 | default_seed=config.seed, |
| 787 | prompt_max_chars=config.prompt_max_chars, |
| 788 | ) |
| 789 | for path in request_files |
| 790 | ] |
| 791 | |
| 792 | |
| 793 | def parse_args(): |
| 794 | import argparse |
| 795 | |
| 796 | parser = argparse.ArgumentParser( |
| 797 | description="Echo 1.5 BF16/FP8/FP4 DMD R2V inference", |
| 798 | formatter_class=argparse.ArgumentDefaultsHelpFormatter, |
| 799 | ) |
| 800 | parser.add_argument("--config", default=str(DEFAULT_CONFIG)) |
| 801 | parser.add_argument( |
| 802 | "--request", help="run one production-compatible R2V JSON request" |
| 803 | ) |
| 804 | parser.add_argument("--checkpoint") |
| 805 | parser.add_argument("--gemma-path") |
| 806 | parser.add_argument("--requests-dir") |
| 807 | parser.add_argument("--requests-glob") |
| 808 | parser.add_argument("--output-root") |
| 809 | parser.add_argument("--conditioning-cache-dir") |
| 810 | parser.add_argument( |
| 811 | "--condition-encode", |
| 812 | "--condition_encode", |
| 813 | "--text-encode", |
| 814 | dest="condition_encode", |
| 815 | action="store_true", |
| 816 | help="batch text, condition-image, memory-image and memory-audio encoding, then exit", |
| 817 | ) |
| 818 | parser.add_argument("--overwrite-condition-cache", action="store_true") |
| 819 | parser.add_argument("--seed", type=int) |
| 820 | parser.add_argument("--num-frames", type=int) |
| 821 | parser.add_argument("--video-height", type=int) |
| 822 | parser.add_argument("--video-width", type=int) |
| 823 | parser.add_argument("--video-fps", type=int) |
| 824 | parser.add_argument("--video-vae-decode-mode", choices=("auto", "tiled", "untiled")) |
| 825 | parser.add_argument("--video-vae-tile-size-frames", type=int) |
| 826 | parser.add_argument("--video-vae-tile-overlap-frames", type=int) |
| 827 | parser.add_argument("--video-vae-tile-size-pixels", type=int) |
| 828 | parser.add_argument("--video-vae-tile-overlap-pixels", type=int) |
| 829 | parser.add_argument("--text-batch-size", type=int) |
| 830 | parser.add_argument("--image-batch-size", type=int) |
| 831 | parser.add_argument("--audio-batch-size", type=int) |
| 832 | parser.add_argument("--dit-layerwise-offload", type=str_to_bool) |
| 833 | parser.add_argument("--dit-resident-blocks", type=int) |
| 834 | parser.add_argument("--dit-prefetch-blocks", type=int) |
| 835 | parser.add_argument("--dit-pin-memory", type=str_to_bool) |
| 836 | parser.add_argument("--memory-max-size", type=int) |
| 837 | return parser.parse_args() |
| 838 | |
| 839 | |
| 840 | def main() -> None: |
| 841 | args = parse_args() |
| 842 | config_path = Path(args.config).expanduser().resolve() |
| 843 | if not config_path.is_file(): |
| 844 | raise FileNotFoundError(f"config not found: {config_path}") |
| 845 | |
| 846 | overrides = { |
| 847 | key: value |
| 848 | for key, value in vars(args).items() |
| 849 | if key |
| 850 | not in {"config", "request", "condition_encode", "overwrite_condition_cache"} |
| 851 | and value is not None |
| 852 | } |
| 853 | for path_key in ( |
| 854 | "checkpoint", |
| 855 | "gemma_path", |
| 856 | "requests_dir", |
| 857 | "output_root", |
| 858 | "conditioning_cache_dir", |
| 859 | ): |
| 860 | if path_key in overrides: |
| 861 | overrides[path_key] = str(Path(overrides[path_key]).expanduser().resolve()) |
| 862 | config = InferenceConfig(config_path, **overrides) |
| 863 | request_files = load_request_files(config, args.request) |
| 864 | requests = load_requests(config, request_files) |
| 865 | |
| 866 | if args.condition_encode: |
| 867 | from scripts.precompute_conditioning import main as precompute_main |
| 868 | |
| 869 | precompute_args = ["--config", str(config_path)] |
| 870 | for option, value in ( |
| 871 | ("--request", args.request), |
| 872 | ("--checkpoint", args.checkpoint), |
| 873 | ("--gemma-path", args.gemma_path), |
| 874 | ("--requests-dir", args.requests_dir), |
| 875 | ("--requests-glob", args.requests_glob), |
| 876 | ("--output-dir", args.conditioning_cache_dir), |
| 877 | ("--text-batch-size", args.text_batch_size), |
| 878 | ("--image-batch-size", args.image_batch_size), |
| 879 | ("--audio-batch-size", args.audio_batch_size), |
| 880 | ): |
| 881 | if value is not None: |
| 882 | precompute_args.extend((option, str(value))) |
| 883 | if args.overwrite_condition_cache: |
| 884 | precompute_args.append("--overwrite") |
| 885 | precompute_main(precompute_args) |
| 886 | return |
| 887 | |
| 888 | if int(os.environ.get("WORLD_SIZE", "1")) > 1: |
| 889 | raise RuntimeError( |
| 890 | "multi-process generation is unsupported; use plain Python for inference or " |
| 891 | "--condition-encode with torchrun" |
| 892 | ) |
| 893 | engine = InferenceEngine(config) |
| 894 | bundles = engine.prepare_conditions(request_files, requests) |
| 895 | engine.load_generator() |
| 896 | output_root = Path(config.output_root) |
| 897 | timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") |
| 898 | for request_file, request in zip(request_files, requests, strict=True): |
| 899 | engine.run_request( |
| 900 | request_file, |
| 901 | request, |
| 902 | output_root / request.work_id / request.shot_id / f"inference_{timestamp}", |
| 903 | bundles[request_file], |
| 904 | ) |
| 905 | print(f"[Inference] Processed {len(request_files)} R2V request(s)", flush=True) |
| 906 | |
| 907 | |
| 908 | if __name__ == "__main__": |
| 909 | main() |
| 910 |