| 1 | #!/usr/bin/env python3 |
| 2 | """Precompute complete Echo R2V conditioning with independent torchrun workers.""" |
| 3 | |
| 4 | # ruff: noqa: E402 |
| 5 | |
| 6 | from __future__ import annotations |
| 7 | |
| 8 | import argparse |
| 9 | import os |
| 10 | import sys |
| 11 | from pathlib import Path |
| 12 | |
| 13 | |
| 14 | REPO_ROOT = Path(__file__).resolve().parents[1] |
| 15 | if str(REPO_ROOT) not in sys.path: |
| 16 | sys.path.insert(0, str(REPO_ROOT)) |
| 17 | for _subpath in ("ltx-core/src", "ltx-pipelines/src", "ltx-distillation/src"): |
| 18 | _package_path = str(REPO_ROOT / _subpath) |
| 19 | if _package_path not in sys.path: |
| 20 | sys.path.insert(0, _package_path) |
| 21 | |
| 22 | import torch |
| 23 | |
| 24 | from inference import InferenceConfig, load_request_files, load_requests |
| 25 | from r2v_schema import R2VRequest |
| 26 | from ltx_distillation.r2v_conditioning import ( |
| 27 | encode_r2v_requests, |
| 28 | load_r2v_conditioning, |
| 29 | r2v_conditioning_cache_path, |
| 30 | save_r2v_conditioning, |
| 31 | ) |
| 32 | from ltx_distillation.release_checkpoint import resolve_release_checkpoint |
| 33 | from ltx_distillation.text_conditioning import artifact_fingerprint |
| 34 | |
| 35 | |
| 36 | def parse_args(argv: list[str] | None = None) -> argparse.Namespace: |
| 37 | parser = argparse.ArgumentParser(description="Precompute complete Echo 1.5 R2V conditions") |
| 38 | parser.add_argument( |
| 39 | "--config", default=str(REPO_ROOT / "configs" / "inference.bf16.yaml") |
| 40 | ) |
| 41 | parser.add_argument("--request") |
| 42 | parser.add_argument("--checkpoint") |
| 43 | parser.add_argument("--gemma-path") |
| 44 | parser.add_argument("--requests-dir") |
| 45 | parser.add_argument("--requests-glob") |
| 46 | parser.add_argument("--output-dir", "--conditioning-cache-dir", dest="output_dir") |
| 47 | parser.add_argument("--text-batch-size", type=int) |
| 48 | parser.add_argument("--image-batch-size", type=int) |
| 49 | parser.add_argument("--audio-batch-size", type=int) |
| 50 | parser.add_argument("--overwrite", action="store_true") |
| 51 | return parser.parse_args(argv) |
| 52 | |
| 53 | |
| 54 | def main(argv: list[str] | None = None) -> None: |
| 55 | args = parse_args(argv) |
| 56 | overrides = { |
| 57 | key: value |
| 58 | for key, value in { |
| 59 | "checkpoint": args.checkpoint, |
| 60 | "gemma_path": args.gemma_path, |
| 61 | "requests_dir": args.requests_dir, |
| 62 | "requests_glob": args.requests_glob, |
| 63 | "conditioning_cache_dir": args.output_dir, |
| 64 | "text_batch_size": args.text_batch_size, |
| 65 | "image_batch_size": args.image_batch_size, |
| 66 | "audio_batch_size": args.audio_batch_size, |
| 67 | }.items() |
| 68 | if value is not None |
| 69 | } |
| 70 | for path_key in ("checkpoint", "gemma_path", "requests_dir", "conditioning_cache_dir"): |
| 71 | if path_key in overrides: |
| 72 | overrides[path_key] = str(Path(overrides[path_key]).expanduser().resolve()) |
| 73 | config = InferenceConfig(Path(args.config).expanduser().resolve(), **overrides) |
| 74 | if not config.conditioning_cache_dir: |
| 75 | raise ValueError( |
| 76 | "paths.conditioning_cache_dir is required for --condition-encode " |
| 77 | "(or override it with --conditioning-cache-dir)" |
| 78 | ) |
| 79 | request_files = load_request_files(config, args.request) |
| 80 | requests = load_requests(config, request_files) |
| 81 | rank = int(os.environ.get("RANK", "0")) |
| 82 | local_rank = int(os.environ.get("LOCAL_RANK", "0")) |
| 83 | world_size = int(os.environ.get("WORLD_SIZE", "1")) |
| 84 | if not torch.cuda.is_available(): |
| 85 | raise RuntimeError("CUDA is required for R2V conditioning precomputation") |
| 86 | if local_rank >= torch.cuda.device_count(): |
| 87 | raise ValueError( |
| 88 | f"LOCAL_RANK={local_rank} is outside {torch.cuda.device_count()} visible CUDA devices" |
| 89 | ) |
| 90 | torch.cuda.set_device(local_rank) |
| 91 | device = torch.device(f"cuda:{local_rank}") |
| 92 | |
| 93 | release = resolve_release_checkpoint(config.checkpoint) |
| 94 | checkpoint_id = artifact_fingerprint(release.root) |
| 95 | gemma_id = artifact_fingerprint(config.gemma_path) |
| 96 | assigned = list(zip(request_files, requests, strict=True))[rank::world_size] |
| 97 | pending: list[tuple[Path, R2VRequest]] = [] |
| 98 | for request_file, request in assigned: |
| 99 | cache_path = r2v_conditioning_cache_path( |
| 100 | config.conditioning_cache_dir, config.requests_dir, request_file |
| 101 | ) |
| 102 | if cache_path.is_file() and not args.overwrite: |
| 103 | try: |
| 104 | load_r2v_conditioning( |
| 105 | cache_path, |
| 106 | request=request, |
| 107 | checkpoint_fingerprint=checkpoint_id, |
| 108 | gemma_fingerprint=gemma_id, |
| 109 | ) |
| 110 | print(f"[rank {rank}] cached {cache_path}", flush=True) |
| 111 | continue |
| 112 | except ValueError: |
| 113 | pass |
| 114 | pending.append((request_file, request)) |
| 115 | if not pending: |
| 116 | print(f"[rank {rank}] nothing to encode", flush=True) |
| 117 | return |
| 118 | |
| 119 | print( |
| 120 | f"[rank {rank}/{world_size}] encoding {len(pending)} of {len(assigned)} requests on {device}", |
| 121 | flush=True, |
| 122 | ) |
| 123 | bundles = encode_r2v_requests( |
| 124 | [request for _, request in pending], |
| 125 | checkpoint_path=str(release.model_path), |
| 126 | gemma_path=str(config.gemma_path), |
| 127 | device=device, |
| 128 | voice_filter_config=config.voice_filter, |
| 129 | dtype=torch.bfloat16, |
| 130 | text_batch_size=config.text_batch_size, |
| 131 | image_batch_size=config.image_batch_size, |
| 132 | audio_batch_size=config.audio_batch_size, |
| 133 | enable_audio_memory=config.enable_audio_memory, |
| 134 | memory_position_mode=config.memory_position_mode, |
| 135 | memory_position_offset=config.memory_position_offset, |
| 136 | memory_position_slot_stride=config.memory_position_slot_stride, |
| 137 | ) |
| 138 | for (request_file, request), bundle in zip(pending, bundles, strict=True): |
| 139 | cache_path = r2v_conditioning_cache_path( |
| 140 | config.conditioning_cache_dir, config.requests_dir, request_file |
| 141 | ) |
| 142 | save_r2v_conditioning( |
| 143 | cache_path, |
| 144 | bundle, |
| 145 | request=request, |
| 146 | checkpoint_fingerprint=checkpoint_id, |
| 147 | gemma_fingerprint=gemma_id, |
| 148 | ) |
| 149 | print(f"[rank {rank}] wrote {cache_path}", flush=True) |
| 150 | print(f"[rank {rank}] complete", flush=True) |
| 151 | |
| 152 | |
| 153 | if __name__ == "__main__": |
| 154 | main() |
| 155 |