| 1 | """Batch text conditioning and safe, portable cache serialization.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import gc |
| 6 | import hashlib |
| 7 | import os |
| 8 | from pathlib import Path |
| 9 | from typing import TypeAlias |
| 10 | |
| 11 | import torch |
| 12 | from safetensors import safe_open |
| 13 | from safetensors.torch import save_file |
| 14 | |
| 15 | from ltx_distillation.models.text_encoder_wrapper import ( |
| 16 | create_language_only_text_encoder, |
| 17 | create_text_embeddings_processor, |
| 18 | ) |
| 19 | |
| 20 | |
| 21 | TextCondition: TypeAlias = dict[str, torch.Tensor | None] |
| 22 | RawTextCondition: TypeAlias = tuple[tuple[torch.Tensor, ...], torch.Tensor] |
| 23 | |
| 24 | _CACHE_SCHEMA_VERSION = "1" |
| 25 | _REQUIRED_TENSORS = {"video_context", "attention_mask"} |
| 26 | |
| 27 | |
| 28 | def _release_cuda(device: torch.device) -> None: |
| 29 | gc.collect() |
| 30 | if device.type == "cuda": |
| 31 | torch.cuda.empty_cache() |
| 32 | |
| 33 | |
| 34 | def prompt_fingerprint(prompt: str) -> str: |
| 35 | return hashlib.sha256(prompt.encode("utf-8")).hexdigest() |
| 36 | |
| 37 | |
| 38 | def artifact_fingerprint(path: str | Path) -> str: |
| 39 | """Create a cheap, path-independent fingerprint from artifact names and sizes.""" |
| 40 | root = Path(path).resolve() |
| 41 | if root.is_file(): |
| 42 | records = [(root.name, root.stat().st_size)] |
| 43 | elif root.is_dir(): |
| 44 | relevant_names = { |
| 45 | "checkpoint.json", |
| 46 | "config.json", |
| 47 | "preprocessor_config.json", |
| 48 | "tokenizer.model", |
| 49 | "tokenizer_config.json", |
| 50 | } |
| 51 | files = sorted( |
| 52 | item |
| 53 | for item in root.rglob("*") |
| 54 | if item.is_file() |
| 55 | and (item.suffix in {".safetensors", ".pt"} or item.name in relevant_names) |
| 56 | ) |
| 57 | records = [(item.relative_to(root).as_posix(), item.stat().st_size) for item in files] |
| 58 | else: |
| 59 | raise FileNotFoundError(f"model artifact not found: {root}") |
| 60 | payload = "\n".join(f"{name}:{size}" for name, size in records) |
| 61 | return hashlib.sha256(payload.encode("utf-8")).hexdigest() |
| 62 | |
| 63 | |
| 64 | def conditioning_cache_path( |
| 65 | cache_root: str | Path, |
| 66 | prompts_root: str | Path, |
| 67 | prompt_file: str | Path, |
| 68 | shot_index: int, |
| 69 | ) -> Path: |
| 70 | prompts_root = Path(prompts_root).resolve() |
| 71 | prompt_file = Path(prompt_file).resolve() |
| 72 | try: |
| 73 | relative_prompt = prompt_file.relative_to(prompts_root) |
| 74 | except ValueError: |
| 75 | relative_prompt = Path(prompt_file.name) |
| 76 | return Path(cache_root) / relative_prompt.with_suffix("") / f"shot_{shot_index:04d}.safetensors" |
| 77 | |
| 78 | |
| 79 | def save_text_conditioning( |
| 80 | path: str | Path, |
| 81 | condition: TextCondition, |
| 82 | *, |
| 83 | prompt: str, |
| 84 | checkpoint_fingerprint: str, |
| 85 | gemma_fingerprint: str, |
| 86 | ) -> None: |
| 87 | destination = Path(path) |
| 88 | destination.parent.mkdir(parents=True, exist_ok=True) |
| 89 | tensors = { |
| 90 | key: value.detach().cpu().contiguous() |
| 91 | for key, value in condition.items() |
| 92 | if isinstance(value, torch.Tensor) |
| 93 | } |
| 94 | missing = _REQUIRED_TENSORS - tensors.keys() |
| 95 | if missing: |
| 96 | raise ValueError(f"conditioning is missing required tensors: {sorted(missing)}") |
| 97 | metadata = { |
| 98 | "schema_version": _CACHE_SCHEMA_VERSION, |
| 99 | "prompt_sha256": prompt_fingerprint(prompt), |
| 100 | "checkpoint_fingerprint": checkpoint_fingerprint, |
| 101 | "gemma_fingerprint": gemma_fingerprint, |
| 102 | "has_audio_context": str("audio_context" in tensors).lower(), |
| 103 | } |
| 104 | temporary = destination.with_name(f".{destination.name}.{os.getpid()}.tmp") |
| 105 | save_file(tensors, str(temporary), metadata=metadata) |
| 106 | os.replace(temporary, destination) |
| 107 | |
| 108 | |
| 109 | def load_text_conditioning( |
| 110 | path: str | Path, |
| 111 | *, |
| 112 | prompt: str, |
| 113 | checkpoint_fingerprint: str, |
| 114 | gemma_fingerprint: str, |
| 115 | ) -> TextCondition: |
| 116 | source = Path(path) |
| 117 | if not source.is_file(): |
| 118 | raise FileNotFoundError(f"text conditioning cache not found: {source}") |
| 119 | with safe_open(str(source), framework="pt", device="cpu") as handle: |
| 120 | metadata = handle.metadata() or {} |
| 121 | expected_metadata = { |
| 122 | "schema_version": _CACHE_SCHEMA_VERSION, |
| 123 | "prompt_sha256": prompt_fingerprint(prompt), |
| 124 | "checkpoint_fingerprint": checkpoint_fingerprint, |
| 125 | "gemma_fingerprint": gemma_fingerprint, |
| 126 | } |
| 127 | mismatches = { |
| 128 | key: (metadata.get(key), expected) |
| 129 | for key, expected in expected_metadata.items() |
| 130 | if metadata.get(key) != expected |
| 131 | } |
| 132 | if mismatches: |
| 133 | details = ", ".join( |
| 134 | f"{key}={actual!r} (expected {expected!r})" |
| 135 | for key, (actual, expected) in mismatches.items() |
| 136 | ) |
| 137 | raise ValueError(f"stale or incompatible text conditioning cache {source}: {details}") |
| 138 | tensors = {key: handle.get_tensor(key) for key in handle.keys()} |
| 139 | |
| 140 | missing = _REQUIRED_TENSORS - tensors.keys() |
| 141 | if missing: |
| 142 | raise ValueError(f"conditioning cache {source} is missing tensors: {sorted(missing)}") |
| 143 | for key, tensor in tensors.items(): |
| 144 | if tensor.shape[0] != 1: |
| 145 | raise ValueError( |
| 146 | f"conditioning cache {source} has invalid {key} batch shape {tuple(tensor.shape)}" |
| 147 | ) |
| 148 | return { |
| 149 | "video_context": tensors["video_context"], |
| 150 | "audio_context": tensors.get("audio_context"), |
| 151 | "attention_mask": tensors["attention_mask"], |
| 152 | } |
| 153 | |
| 154 | |
| 155 | def encode_prompts_two_stage( |
| 156 | prompts: list[str], |
| 157 | *, |
| 158 | checkpoint_path: str, |
| 159 | gemma_path: str, |
| 160 | device: torch.device, |
| 161 | dtype: torch.dtype = torch.bfloat16, |
| 162 | batch_size: int = 1, |
| 163 | ) -> list[TextCondition]: |
| 164 | """Encode prompts in batches while keeping Gemma and Echo connectors disjoint.""" |
| 165 | if batch_size <= 0: |
| 166 | raise ValueError("batch_size must be positive") |
| 167 | if not prompts: |
| 168 | return [] |
| 169 | |
| 170 | text_encoder = create_language_only_text_encoder( |
| 171 | checkpoint_path=checkpoint_path, |
| 172 | gemma_path=gemma_path, |
| 173 | device=device, |
| 174 | dtype=dtype, |
| 175 | ) |
| 176 | raw_conditions: list[RawTextCondition] = [] |
| 177 | try: |
| 178 | with torch.inference_mode(): |
| 179 | for offset in range(0, len(prompts), batch_size): |
| 180 | prompt_batch = prompts[offset : offset + batch_size] |
| 181 | hidden_states, attention_mask = text_encoder.encode_batch(prompt_batch) |
| 182 | for batch_index in range(len(prompt_batch)): |
| 183 | raw_conditions.append( |
| 184 | ( |
| 185 | tuple( |
| 186 | hidden[batch_index : batch_index + 1].detach().cpu() |
| 187 | for hidden in hidden_states |
| 188 | ), |
| 189 | attention_mask[batch_index : batch_index + 1].detach().cpu(), |
| 190 | ) |
| 191 | ) |
| 192 | del hidden_states, attention_mask |
| 193 | finally: |
| 194 | del text_encoder |
| 195 | _release_cuda(device) |
| 196 | |
| 197 | embeddings_processor = create_text_embeddings_processor( |
| 198 | checkpoint_path=checkpoint_path, |
| 199 | device=device, |
| 200 | dtype=dtype, |
| 201 | ) |
| 202 | conditions: list[TextCondition] = [] |
| 203 | try: |
| 204 | with torch.inference_mode(): |
| 205 | for offset in range(0, len(raw_conditions), batch_size): |
| 206 | raw_batch = raw_conditions[offset : offset + batch_size] |
| 207 | hidden_states = tuple( |
| 208 | torch.cat( |
| 209 | [condition[0][layer_index] for condition in raw_batch], |
| 210 | dim=0, |
| 211 | ).to(device) |
| 212 | for layer_index in range(len(raw_batch[0][0])) |
| 213 | ) |
| 214 | attention_mask = torch.cat([condition[1] for condition in raw_batch], dim=0).to( |
| 215 | device |
| 216 | ) |
| 217 | output = embeddings_processor.process_hidden_states( |
| 218 | hidden_states, |
| 219 | attention_mask, |
| 220 | padding_side="left", |
| 221 | ) |
| 222 | for batch_index in range(len(raw_batch)): |
| 223 | conditions.append( |
| 224 | { |
| 225 | "video_context": output.video_encoding[batch_index : batch_index + 1] |
| 226 | .detach() |
| 227 | .cpu(), |
| 228 | "audio_context": ( |
| 229 | output.audio_encoding[batch_index : batch_index + 1].detach().cpu() |
| 230 | if output.audio_encoding is not None |
| 231 | else None |
| 232 | ), |
| 233 | "attention_mask": output.attention_mask[batch_index : batch_index + 1] |
| 234 | .detach() |
| 235 | .cpu(), |
| 236 | } |
| 237 | ) |
| 238 | del hidden_states, attention_mask, output |
| 239 | raw_batch.clear() |
| 240 | finally: |
| 241 | del embeddings_processor, raw_conditions |
| 242 | _release_cuda(device) |
| 243 | return conditions |
| 244 |