| 1 | """Batch R2V conditioning that stops exactly at the DiT input boundary.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import base64 |
| 6 | import binascii |
| 7 | import gc |
| 8 | import hashlib |
| 9 | import io |
| 10 | import json |
| 11 | import os |
| 12 | from collections import defaultdict |
| 13 | from dataclasses import dataclass |
| 14 | from pathlib import Path |
| 15 | from typing import Any, TypeAlias |
| 16 | from urllib.request import Request, urlopen |
| 17 | |
| 18 | import torch |
| 19 | import torchaudio |
| 20 | from PIL import Image |
| 21 | from safetensors import safe_open |
| 22 | from safetensors.torch import save_file |
| 23 | from torchvision.transforms import functional as TVF |
| 24 | |
| 25 | from ltx_distillation.audio_voice_filter import VoiceFilterConfig, filter_voice_only |
| 26 | from ltx_distillation.models.vae_wrapper import create_vae_wrappers |
| 27 | from ltx_distillation.text_conditioning import TextCondition, encode_prompts_two_stage |
| 28 | from r2v_schema import R2VRequest |
| 29 | |
| 30 | CONDITIONING_SCHEMA_VERSION = "echo15.r2v.conditioning.v1" |
| 31 | MAX_IMAGE_BYTES = 20 * 1024 * 1024 |
| 32 | MAX_AUDIO_BYTES = 100 * 1024 * 1024 |
| 33 | TensorDict: TypeAlias = dict[str, torch.Tensor | None] |
| 34 | |
| 35 | |
| 36 | @dataclass |
| 37 | class R2VConditionBundle: |
| 38 | """All CPU tensors and scalar arguments required immediately before DiT.""" |
| 39 | |
| 40 | text: TextCondition |
| 41 | first_frame_latent: torch.Tensor | None |
| 42 | memory_video: torch.Tensor | None |
| 43 | memory_audio_kwargs: dict[str, Any] |
| 44 | input_fingerprints: dict[str, str] |
| 45 | |
| 46 | |
| 47 | def _validate_bundle(bundle: R2VConditionBundle, request: R2VRequest) -> None: |
| 48 | """Reject incomplete caches before they reach the generator.""" |
| 49 | |
| 50 | for key in ("video_context", "attention_mask"): |
| 51 | value = bundle.text.get(key) |
| 52 | if not isinstance(value, torch.Tensor) or value.shape[0] != 1: |
| 53 | raise ValueError(f"R2V conditioning requires batched text.{key}") |
| 54 | audio_context = bundle.text.get("audio_context") |
| 55 | if audio_context is not None and ( |
| 56 | not isinstance(audio_context, torch.Tensor) or audio_context.shape[0] != 1 |
| 57 | ): |
| 58 | raise ValueError("R2V conditioning text.audio_context must have batch size 1") |
| 59 | |
| 60 | first_frame = bundle.first_frame_latent |
| 61 | if (first_frame is not None) != (request.condition_img is not None): |
| 62 | raise ValueError("condition_img and first_frame_latent must either both exist or both be absent") |
| 63 | if first_frame is not None and ( |
| 64 | first_frame.ndim != 5 or first_frame.shape[0] != 1 or first_frame.shape[1] != 1 |
| 65 | ): |
| 66 | raise ValueError( |
| 67 | "first_frame_latent must have shape [1, 1, C, H, W], got " |
| 68 | f"{tuple(first_frame.shape)}" |
| 69 | ) |
| 70 | |
| 71 | slot_count = len(request.memory_slots) |
| 72 | memory_video = bundle.memory_video |
| 73 | if (memory_video is not None) != bool(slot_count): |
| 74 | raise ValueError("memory slots and memory_video must either both exist or both be absent") |
| 75 | if memory_video is not None and ( |
| 76 | memory_video.ndim != 5 |
| 77 | or memory_video.shape[0] != 1 |
| 78 | or memory_video.shape[1] != slot_count |
| 79 | ): |
| 80 | raise ValueError( |
| 81 | f"memory_video must have shape [1, {slot_count}, C, H, W], got " |
| 82 | f"{tuple(memory_video.shape)}" |
| 83 | ) |
| 84 | |
| 85 | audio = bundle.memory_audio_kwargs.get("memory_audio") |
| 86 | timestep = bundle.memory_audio_kwargs.get("memory_audio_timestep") |
| 87 | lengths = bundle.memory_audio_kwargs.get("memory_audio_segment_lengths") |
| 88 | if audio is None: |
| 89 | if timestep is not None or lengths is not None: |
| 90 | raise ValueError("memory audio timestep/segments require memory_audio") |
| 91 | return |
| 92 | if not isinstance(audio, torch.Tensor) or audio.ndim != 3 or audio.shape[0] != 1: |
| 93 | raise ValueError("memory_audio must have shape [1, T, C]") |
| 94 | if not isinstance(timestep, torch.Tensor) or tuple(timestep.shape) != tuple(audio.shape[:2]): |
| 95 | raise ValueError("memory_audio_timestep must match memory_audio batch/time dimensions") |
| 96 | if ( |
| 97 | not isinstance(lengths, tuple) |
| 98 | or len(lengths) != 1 |
| 99 | or len(lengths[0]) != slot_count |
| 100 | or sum(int(value) for value in lengths[0]) != audio.shape[1] |
| 101 | ): |
| 102 | raise ValueError("memory_audio_segment_lengths must align one-to-one with memory slots") |
| 103 | |
| 104 | |
| 105 | def r2v_conditioning_cache_path( |
| 106 | cache_root: str | Path, |
| 107 | requests_root: str | Path, |
| 108 | request_file: str | Path, |
| 109 | ) -> Path: |
| 110 | requests_root = Path(requests_root).resolve() |
| 111 | request_file = Path(request_file).resolve() |
| 112 | try: |
| 113 | relative = request_file.relative_to(requests_root) |
| 114 | except ValueError: |
| 115 | relative = Path(request_file.name) |
| 116 | return Path(cache_root) / relative.with_suffix(".safetensors") |
| 117 | |
| 118 | |
| 119 | class _ResourceStore: |
| 120 | def __init__(self) -> None: |
| 121 | self._bytes: dict[str, bytes] = {} |
| 122 | self.fingerprints: dict[str, str] = {} |
| 123 | |
| 124 | def read(self, source: str, *, kind: str, max_bytes: int) -> bytes: |
| 125 | if source in self._bytes: |
| 126 | return self._bytes[source] |
| 127 | if source.startswith("data:"): |
| 128 | header, separator, encoded = source.partition(",") |
| 129 | if not separator or ";base64" not in header.lower(): |
| 130 | raise ValueError(f"{kind} has an invalid data URL") |
| 131 | try: |
| 132 | data = base64.b64decode(encoded, validate=True) |
| 133 | except (binascii.Error, ValueError) as exc: |
| 134 | raise ValueError(f"{kind} has invalid base64 data") from exc |
| 135 | elif source.startswith(("http://", "https://")): |
| 136 | request = Request(source, headers={"User-Agent": "JoyAI-Echo15/1.0"}) |
| 137 | with urlopen(request, timeout=30) as response: |
| 138 | length = response.headers.get("Content-Length") |
| 139 | if length and int(length) > max_bytes: |
| 140 | raise ValueError(f"{kind} exceeds {max_bytes} bytes: {source}") |
| 141 | data = response.read(max_bytes + 1) |
| 142 | else: |
| 143 | path = Path(source) |
| 144 | if not path.is_file(): |
| 145 | raise FileNotFoundError(f"{kind} not found: {path}") |
| 146 | if path.stat().st_size > max_bytes: |
| 147 | raise ValueError(f"{kind} exceeds {max_bytes} bytes: {path}") |
| 148 | data = path.read_bytes() |
| 149 | if not data or len(data) > max_bytes: |
| 150 | raise ValueError(f"{kind} must contain 1..{max_bytes} bytes: {source}") |
| 151 | self._bytes[source] = data |
| 152 | self.fingerprints[source] = hashlib.sha256(data).hexdigest() |
| 153 | return data |
| 154 | |
| 155 | def image(self, source: str) -> Image.Image: |
| 156 | raw = self.read(source, kind="R2V image", max_bytes=MAX_IMAGE_BYTES) |
| 157 | with Image.open(io.BytesIO(raw)) as image: |
| 158 | image.load() |
| 159 | return image.convert("RGB") |
| 160 | |
| 161 | def audio(self, source: str) -> tuple[torch.Tensor, int]: |
| 162 | raw = self.read(source, kind="R2V audio", max_bytes=MAX_AUDIO_BYTES) |
| 163 | waveform, sample_rate = torchaudio.load(io.BytesIO(raw)) |
| 164 | return waveform.detach().cpu().float().contiguous(), int(sample_rate) |
| 165 | |
| 166 | |
| 167 | def _release_cuda(device: torch.device) -> None: |
| 168 | gc.collect() |
| 169 | if device.type == "cuda": |
| 170 | torch.cuda.empty_cache() |
| 171 | |
| 172 | |
| 173 | def _request_cache_fingerprint(request: R2VRequest) -> str: |
| 174 | if request.request_sha256: |
| 175 | return request.request_sha256 |
| 176 | payload = request.as_payload() |
| 177 | return hashlib.sha256( |
| 178 | json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode( |
| 179 | "utf-8" |
| 180 | ) |
| 181 | ).hexdigest() |
| 182 | |
| 183 | |
| 184 | def _image_tensor(image: Image.Image, *, height: int, width: int) -> torch.Tensor: |
| 185 | if image.size != (width, height): |
| 186 | image = image.resize((width, height), Image.Resampling.BICUBIC) |
| 187 | return (TVF.to_tensor(image) * 2.0 - 1.0).unsqueeze(1).contiguous() |
| 188 | |
| 189 | |
| 190 | def _normalize_audio(waveform: torch.Tensor) -> torch.Tensor: |
| 191 | value = torch.as_tensor(waveform).detach().cpu().float() |
| 192 | while value.ndim > 2 and value.shape[0] == 1: |
| 193 | value = value.squeeze(0) |
| 194 | if value.ndim == 1: |
| 195 | value = value.unsqueeze(0) |
| 196 | elif value.ndim > 2: |
| 197 | value = value.reshape(value.shape[-2], value.shape[-1]) |
| 198 | if value.ndim != 2 or value.shape[-1] <= 1: |
| 199 | raise ValueError(f"R2V audio has no usable samples: shape={tuple(value.shape)}") |
| 200 | if value.shape[0] == 1: |
| 201 | value = value.repeat(2, 1) |
| 202 | elif value.shape[0] > 2: |
| 203 | value = value[:2] |
| 204 | return value.contiguous() |
| 205 | |
| 206 | |
| 207 | def _chunked(values: list[Any], size: int): |
| 208 | for offset in range(0, len(values), size): |
| 209 | yield values[offset : offset + size] |
| 210 | |
| 211 | |
| 212 | def _encode_unique_images( |
| 213 | requests: list[R2VRequest], |
| 214 | *, |
| 215 | video_vae, |
| 216 | resources: _ResourceStore, |
| 217 | device: torch.device, |
| 218 | dtype: torch.dtype, |
| 219 | batch_size: int, |
| 220 | ) -> dict[tuple[str, int, int], torch.Tensor]: |
| 221 | grouped: dict[tuple[int, int], list[tuple[str, int, int]]] = defaultdict(list) |
| 222 | seen: set[tuple[str, int, int]] = set() |
| 223 | for request in requests: |
| 224 | sources = [request.condition_img] + [slot.image_url for slot in request.memory_slots] |
| 225 | for source in sources: |
| 226 | if not source: |
| 227 | continue |
| 228 | key = (source, request.height, request.width) |
| 229 | if key not in seen: |
| 230 | seen.add(key) |
| 231 | grouped[(request.height, request.width)].append(key) |
| 232 | |
| 233 | encoded: dict[tuple[str, int, int], torch.Tensor] = {} |
| 234 | if not grouped: |
| 235 | return encoded |
| 236 | video_vae.encoder.to(device=device, dtype=dtype) |
| 237 | with torch.inference_mode(): |
| 238 | for (height, width), keys in grouped.items(): |
| 239 | for batch in _chunked(keys, batch_size): |
| 240 | pixels = torch.stack( |
| 241 | [ |
| 242 | _image_tensor(resources.image(source), height=height, width=width) |
| 243 | for source, _, _ in batch |
| 244 | ], |
| 245 | dim=0, |
| 246 | ).to(device=device, dtype=dtype) |
| 247 | latents = video_vae.encode(pixels).permute(0, 2, 1, 3, 4) |
| 248 | for index, key in enumerate(batch): |
| 249 | encoded[key] = latents[index : index + 1].detach().cpu().contiguous() |
| 250 | del pixels, latents |
| 251 | video_vae.encoder.to("cpu") |
| 252 | _release_cuda(device) |
| 253 | return encoded |
| 254 | |
| 255 | |
| 256 | def _encode_unique_audio( |
| 257 | requests: list[R2VRequest], |
| 258 | *, |
| 259 | audio_vae, |
| 260 | resources: _ResourceStore, |
| 261 | voice_filter_config: VoiceFilterConfig, |
| 262 | device: torch.device, |
| 263 | batch_size: int, |
| 264 | enabled: bool, |
| 265 | ) -> dict[str, torch.Tensor]: |
| 266 | if not enabled: |
| 267 | return {} |
| 268 | sources = list( |
| 269 | dict.fromkeys( |
| 270 | slot.audio_url |
| 271 | for request in requests |
| 272 | for slot in request.memory_slots |
| 273 | if slot.audio_url |
| 274 | ) |
| 275 | ) |
| 276 | if not sources: |
| 277 | return {} |
| 278 | prepared: dict[str, tuple[torch.Tensor, int]] = {} |
| 279 | for source in sources: |
| 280 | waveform, sample_rate = resources.audio(source) |
| 281 | filtered = filter_voice_only(waveform, sample_rate, voice_filter_config) |
| 282 | if filtered is None: |
| 283 | raise ValueError(f"voice filter unexpectedly removed R2V audio: {source}") |
| 284 | prepared[source] = (_normalize_audio(filtered), sample_rate) |
| 285 | |
| 286 | groups: dict[tuple[int, tuple[int, ...]], list[str]] = defaultdict(list) |
| 287 | for source, (waveform, sample_rate) in prepared.items(): |
| 288 | groups[(sample_rate, tuple(waveform.shape))].append(source) |
| 289 | |
| 290 | encoded: dict[str, torch.Tensor] = {} |
| 291 | audio_vae.encoder.to(device=device, dtype=torch.float32) |
| 292 | with torch.inference_mode(): |
| 293 | for (sample_rate, _shape), group_sources in groups.items(): |
| 294 | for batch in _chunked(group_sources, batch_size): |
| 295 | waveforms = torch.stack([prepared[source][0] for source in batch], dim=0) |
| 296 | latents = audio_vae.encode(waveforms, sample_rate) |
| 297 | for index, source in enumerate(batch): |
| 298 | encoded[source] = latents[index : index + 1].detach().cpu().contiguous() |
| 299 | del waveforms, latents |
| 300 | audio_vae.encoder.to("cpu") |
| 301 | _release_cuda(device) |
| 302 | return encoded |
| 303 | |
| 304 | |
| 305 | def encode_r2v_requests( |
| 306 | requests: list[R2VRequest], |
| 307 | *, |
| 308 | checkpoint_path: str, |
| 309 | gemma_path: str, |
| 310 | device: torch.device, |
| 311 | voice_filter_config: VoiceFilterConfig, |
| 312 | dtype: torch.dtype = torch.bfloat16, |
| 313 | text_batch_size: int = 1, |
| 314 | image_batch_size: int = 1, |
| 315 | audio_batch_size: int = 1, |
| 316 | enable_audio_memory: bool = True, |
| 317 | memory_position_mode: str = "slot_center", |
| 318 | memory_position_offset: float = 500.0, |
| 319 | memory_position_slot_stride: float = 50.0, |
| 320 | ) -> list[R2VConditionBundle]: |
| 321 | """Batch all modalities, then return one fully assembled DiT input per request.""" |
| 322 | |
| 323 | if min(text_batch_size, image_batch_size, audio_batch_size) <= 0: |
| 324 | raise ValueError("conditioning batch sizes must be positive") |
| 325 | if not requests: |
| 326 | return [] |
| 327 | for request in requests: |
| 328 | unresolved = [slot.shot_id for slot in request.memory_slots if slot.shot_id] |
| 329 | if unresolved: |
| 330 | raise ValueError( |
| 331 | "offline inference cannot resolve shot_id memory slots; provide image_url/audio_url " |
| 332 | f"instead (shot={request.shot_id}, references={unresolved})" |
| 333 | ) |
| 334 | |
| 335 | unique_prompts = list(dict.fromkeys(request.prompt for request in requests)) |
| 336 | unique_text_conditions = encode_prompts_two_stage( |
| 337 | unique_prompts, |
| 338 | checkpoint_path=checkpoint_path, |
| 339 | gemma_path=gemma_path, |
| 340 | device=device, |
| 341 | dtype=dtype, |
| 342 | batch_size=text_batch_size, |
| 343 | ) |
| 344 | text_by_prompt = dict(zip(unique_prompts, unique_text_conditions, strict=True)) |
| 345 | resources = _ResourceStore() |
| 346 | video_vae, audio_vae = create_vae_wrappers( |
| 347 | checkpoint_path=checkpoint_path, |
| 348 | device=torch.device("cpu"), |
| 349 | dtype=dtype, |
| 350 | with_video_encoder=True, |
| 351 | with_audio_encoder=True, |
| 352 | with_decoders=False, |
| 353 | ) |
| 354 | try: |
| 355 | images = _encode_unique_images( |
| 356 | requests, |
| 357 | video_vae=video_vae, |
| 358 | resources=resources, |
| 359 | device=device, |
| 360 | dtype=dtype, |
| 361 | batch_size=image_batch_size, |
| 362 | ) |
| 363 | audios = _encode_unique_audio( |
| 364 | requests, |
| 365 | audio_vae=audio_vae, |
| 366 | resources=resources, |
| 367 | voice_filter_config=voice_filter_config, |
| 368 | device=device, |
| 369 | batch_size=audio_batch_size, |
| 370 | enabled=enable_audio_memory, |
| 371 | ) |
| 372 | finally: |
| 373 | del video_vae, audio_vae |
| 374 | _release_cuda(device) |
| 375 | |
| 376 | bundles: list[R2VConditionBundle] = [] |
| 377 | for request in requests: |
| 378 | text_condition = text_by_prompt[request.prompt] |
| 379 | first_frame = ( |
| 380 | images[(request.condition_img, request.height, request.width)] |
| 381 | if request.condition_img |
| 382 | else None |
| 383 | ) |
| 384 | video_slices = [ |
| 385 | images[(slot.image_url, request.height, request.width)] |
| 386 | for slot in request.memory_slots |
| 387 | if slot.image_url |
| 388 | ] |
| 389 | if len(video_slices) != len(request.memory_slots): |
| 390 | raise ValueError(f"every offline R2V memory slot needs image_url: {request.shot_id}") |
| 391 | memory_video = torch.cat(video_slices, dim=1).contiguous() if video_slices else None |
| 392 | |
| 393 | audio_slices = [ |
| 394 | audios.get(slot.audio_url) if slot.audio_url else None |
| 395 | for slot in request.memory_slots |
| 396 | ] |
| 397 | template = next((item for item in audio_slices if item is not None), None) |
| 398 | memory_audio_kwargs: dict[str, Any] = ( |
| 399 | { |
| 400 | "memory_position_mode": str(memory_position_mode), |
| 401 | "memory_position_offset": float(memory_position_offset), |
| 402 | "memory_position_slot_stride": float(memory_position_slot_stride), |
| 403 | } |
| 404 | if memory_video is not None |
| 405 | else {} |
| 406 | ) |
| 407 | if template is not None: |
| 408 | aligned = [item if item is not None else torch.zeros_like(template) for item in audio_slices] |
| 409 | memory_audio = torch.cat(aligned, dim=1).contiguous() |
| 410 | memory_audio_kwargs.update({ |
| 411 | "memory_audio": memory_audio, |
| 412 | "memory_audio_timestep": torch.zeros(memory_audio.shape[:2], dtype=torch.float32), |
| 413 | "memory_audio_segment_lengths": ( |
| 414 | tuple(int(item.shape[1]) for item in aligned), |
| 415 | ), |
| 416 | }) |
| 417 | bundle = R2VConditionBundle( |
| 418 | text=text_condition, |
| 419 | first_frame_latent=first_frame, |
| 420 | memory_video=memory_video, |
| 421 | memory_audio_kwargs=memory_audio_kwargs, |
| 422 | input_fingerprints=dict(resources.fingerprints), |
| 423 | ) |
| 424 | _validate_bundle(bundle, request) |
| 425 | bundles.append(bundle) |
| 426 | return bundles |
| 427 | |
| 428 | |
| 429 | def save_r2v_conditioning( |
| 430 | path: str | Path, |
| 431 | bundle: R2VConditionBundle, |
| 432 | *, |
| 433 | request: R2VRequest, |
| 434 | checkpoint_fingerprint: str, |
| 435 | gemma_fingerprint: str, |
| 436 | ) -> None: |
| 437 | destination = Path(path) |
| 438 | _validate_bundle(bundle, request) |
| 439 | destination.parent.mkdir(parents=True, exist_ok=True) |
| 440 | tensors: dict[str, torch.Tensor] = {} |
| 441 | for key, value in bundle.text.items(): |
| 442 | if isinstance(value, torch.Tensor): |
| 443 | tensors[f"text.{key}"] = value.detach().cpu().contiguous() |
| 444 | if bundle.first_frame_latent is not None: |
| 445 | tensors["first_frame_latent"] = bundle.first_frame_latent.detach().cpu().contiguous() |
| 446 | if bundle.memory_video is not None: |
| 447 | tensors["memory_video"] = bundle.memory_video.detach().cpu().contiguous() |
| 448 | for key in ("memory_audio", "memory_audio_timestep"): |
| 449 | value = bundle.memory_audio_kwargs.get(key) |
| 450 | if isinstance(value, torch.Tensor): |
| 451 | tensors[key] = value.detach().cpu().contiguous() |
| 452 | |
| 453 | scalar_audio_kwargs = { |
| 454 | key: value |
| 455 | for key, value in bundle.memory_audio_kwargs.items() |
| 456 | if not isinstance(value, torch.Tensor) |
| 457 | } |
| 458 | metadata = { |
| 459 | "schema_version": CONDITIONING_SCHEMA_VERSION, |
| 460 | "request_sha256": _request_cache_fingerprint(request), |
| 461 | "checkpoint_fingerprint": checkpoint_fingerprint, |
| 462 | "gemma_fingerprint": gemma_fingerprint, |
| 463 | "shot_id": request.shot_id, |
| 464 | "memory_slot_count": str(len(request.memory_slots)), |
| 465 | "memory_audio_kwargs": json.dumps(scalar_audio_kwargs, separators=(",", ":")), |
| 466 | "input_fingerprints": json.dumps( |
| 467 | bundle.input_fingerprints, sort_keys=True, separators=(",", ":") |
| 468 | ), |
| 469 | } |
| 470 | temporary = destination.with_name(f".{destination.name}.{os.getpid()}.tmp") |
| 471 | save_file(tensors, str(temporary), metadata=metadata) |
| 472 | os.replace(temporary, destination) |
| 473 | |
| 474 | |
| 475 | def load_r2v_conditioning( |
| 476 | path: str | Path, |
| 477 | *, |
| 478 | request: R2VRequest, |
| 479 | checkpoint_fingerprint: str, |
| 480 | gemma_fingerprint: str, |
| 481 | ) -> R2VConditionBundle: |
| 482 | source = Path(path) |
| 483 | if not source.is_file(): |
| 484 | raise FileNotFoundError(f"R2V conditioning cache not found: {source}") |
| 485 | with safe_open(str(source), framework="pt", device="cpu") as handle: |
| 486 | metadata = handle.metadata() or {} |
| 487 | expected = { |
| 488 | "schema_version": CONDITIONING_SCHEMA_VERSION, |
| 489 | "request_sha256": _request_cache_fingerprint(request), |
| 490 | "checkpoint_fingerprint": checkpoint_fingerprint, |
| 491 | "gemma_fingerprint": gemma_fingerprint, |
| 492 | "shot_id": request.shot_id, |
| 493 | "memory_slot_count": str(len(request.memory_slots)), |
| 494 | } |
| 495 | mismatches = { |
| 496 | key: (metadata.get(key), value) |
| 497 | for key, value in expected.items() |
| 498 | if metadata.get(key) != value |
| 499 | } |
| 500 | if mismatches: |
| 501 | details = ", ".join( |
| 502 | f"{key}={actual!r} (expected {wanted!r})" |
| 503 | for key, (actual, wanted) in mismatches.items() |
| 504 | ) |
| 505 | raise ValueError(f"stale or incompatible R2V cache {source}: {details}") |
| 506 | tensors = {key: handle.get_tensor(key) for key in handle.keys()} |
| 507 | |
| 508 | required_text = {"text.video_context", "text.attention_mask"} |
| 509 | missing = required_text - tensors.keys() |
| 510 | if missing: |
| 511 | raise ValueError(f"R2V cache {source} is missing tensors: {sorted(missing)}") |
| 512 | audio_kwargs = json.loads(metadata.get("memory_audio_kwargs", "{}")) |
| 513 | segment_lengths = audio_kwargs.get("memory_audio_segment_lengths") |
| 514 | if segment_lengths is not None: |
| 515 | audio_kwargs["memory_audio_segment_lengths"] = tuple( |
| 516 | tuple(int(value) for value in row) for row in segment_lengths |
| 517 | ) |
| 518 | for key in ("memory_audio", "memory_audio_timestep"): |
| 519 | if key in tensors: |
| 520 | audio_kwargs[key] = tensors[key] |
| 521 | bundle = R2VConditionBundle( |
| 522 | text={ |
| 523 | "video_context": tensors["text.video_context"], |
| 524 | "audio_context": tensors.get("text.audio_context"), |
| 525 | "attention_mask": tensors["text.attention_mask"], |
| 526 | }, |
| 527 | first_frame_latent=tensors.get("first_frame_latent"), |
| 528 | memory_video=tensors.get("memory_video"), |
| 529 | memory_audio_kwargs=audio_kwargs, |
| 530 | input_fingerprints=json.loads(metadata.get("input_fingerprints", "{}")), |
| 531 | ) |
| 532 | _validate_bundle(bundle, request) |
| 533 | return bundle |
| 534 |