| 1 | """Causal denoising schedule and aligned audio-video block layout.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | from itertools import pairwise |
| 6 | |
| 7 | from ltx_core.components.schedulers import LTX2Scheduler |
| 8 | |
| 9 | DEFAULT_CAUSAL_TIMESTEPS = (1000, 750, 500, 250) |
| 10 | CAUSAL_VIDEO_CHUNK_SIZE = 3 |
| 11 | AUDIO_PREFIX_FRAMES = 2 |
| 12 | AUDIO_FRAMES_PER_VIDEO_BLOCK = 25 |
| 13 | |
| 14 | |
| 15 | def resolve_causal_sigmas( |
| 16 | timesteps: tuple[int, ...] | list[int] = DEFAULT_CAUSAL_TIMESTEPS, |
| 17 | *, |
| 18 | num_train_timesteps: int = 1000, |
| 19 | ) -> list[float]: |
| 20 | """Map distilled student timesteps to model sigmas without appending zero.""" |
| 21 | if not timesteps: |
| 22 | raise ValueError("at least one causal timestep is required") |
| 23 | schedule = LTX2Scheduler().execute(steps=num_train_timesteps) |
| 24 | result = [] |
| 25 | for timestep in timesteps: |
| 26 | index = num_train_timesteps - int(timestep) |
| 27 | if not 0 <= index < len(schedule): |
| 28 | raise ValueError(f"causal timestep {timestep} is outside [0, {num_train_timesteps}]") |
| 29 | result.append(float(schedule[index])) |
| 30 | if any(current <= following for current, following in pairwise(result)): |
| 31 | raise ValueError(f"causal sigmas must be strictly descending, got {result}") |
| 32 | return result |
| 33 | |
| 34 | |
| 35 | def _causal_block_count(video_frames: int, chunk_size: int) -> int: |
| 36 | if chunk_size != CAUSAL_VIDEO_CHUNK_SIZE: |
| 37 | raise ValueError( |
| 38 | f"Echo-WM Flash requires video_chunk_size={CAUSAL_VIDEO_CHUNK_SIZE}, " |
| 39 | f"got {chunk_size}" |
| 40 | ) |
| 41 | if video_frames < 1 or (video_frames - 1) % chunk_size: |
| 42 | raise ValueError(f"latent video length must be 1 + n * chunk_size, got {video_frames}") |
| 43 | return (video_frames - 1) // chunk_size |
| 44 | |
| 45 | |
| 46 | def causal_video_blocks( |
| 47 | video_frames: int, |
| 48 | chunk_size: int = CAUSAL_VIDEO_CHUNK_SIZE, |
| 49 | ) -> list[tuple[int, int]]: |
| 50 | """Return ``[0, 1]`` followed by fixed-size causal generation blocks.""" |
| 51 | _causal_block_count(video_frames, chunk_size) |
| 52 | return [(0, 1), *[(start, start + chunk_size) for start in range(1, video_frames, chunk_size)]] |
| 53 | |
| 54 | |
| 55 | def causal_audio_frames( |
| 56 | video_frames: int, |
| 57 | chunk_size: int = CAUSAL_VIDEO_CHUNK_SIZE, |
| 58 | ) -> int: |
| 59 | """Map video latent frames to the aligned audio layout.""" |
| 60 | block_count = _causal_block_count(video_frames, chunk_size) |
| 61 | return AUDIO_PREFIX_FRAMES + block_count * AUDIO_FRAMES_PER_VIDEO_BLOCK |
| 62 | |
| 63 | |
| 64 | def causal_audio_blocks( |
| 65 | video_frames: int, |
| 66 | chunk_size: int = CAUSAL_VIDEO_CHUNK_SIZE, |
| 67 | ) -> list[tuple[int, int]]: |
| 68 | """Return the audio blocks paired with :func:`causal_video_blocks`.""" |
| 69 | total = causal_audio_frames(video_frames, chunk_size) |
| 70 | return [ |
| 71 | (0, AUDIO_PREFIX_FRAMES), |
| 72 | *[ |
| 73 | (start, min(start + AUDIO_FRAMES_PER_VIDEO_BLOCK, total)) |
| 74 | for start in range(AUDIO_PREFIX_FRAMES, total, AUDIO_FRAMES_PER_VIDEO_BLOCK) |
| 75 | ], |
| 76 | ] |
| 77 |