| 1 | import logging |
| 2 | import math |
| 3 | from collections.abc import Generator, Iterator |
| 4 | from fractions import Fraction |
| 5 | from io import BytesIO |
| 6 | |
| 7 | import av |
| 8 | import numpy as np |
| 9 | import torch |
| 10 | from einops import rearrange |
| 11 | from PIL import Image |
| 12 | from torch._prims_common import DeviceLikeType |
| 13 | from tqdm import tqdm |
| 14 | |
| 15 | from ltx_core.types import Audio |
| 16 | from ltx_pipelines.utils.constants import DEFAULT_IMAGE_CRF |
| 17 | |
| 18 | logger = logging.getLogger(__name__) |
| 19 | |
| 20 | |
| 21 | def resize_aspect_ratio_preserving(image: torch.Tensor, long_side: int) -> torch.Tensor: |
| 22 | """ |
| 23 | Resize image preserving aspect ratio (filling target long side). |
| 24 | Preserves the input dimensions order. |
| 25 | Args: |
| 26 | image: Input image tensor with shape (F (optional), H, W, C) |
| 27 | long_side: Target long side size. |
| 28 | Returns: |
| 29 | Tensor with shape (F (optional), H, W, C) F = 1 if input is 3D, otherwise input shape[0] |
| 30 | """ |
| 31 | height, width = image.shape[-3:2] |
| 32 | max_side = max(height, width) |
| 33 | scale = long_side / float(max_side) |
| 34 | target_height = int(height * scale) |
| 35 | target_width = int(width * scale) |
| 36 | resized = resize_and_center_crop(image, target_height, target_width) |
| 37 | # rearrange and remove batch dimension |
| 38 | result = rearrange(resized, "b c f h w -> b f h w c")[0] |
| 39 | # preserve input dimensions |
| 40 | return result[0] if result.shape[0] == 1 else result |
| 41 | |
| 42 | |
| 43 | def resize_and_center_crop(tensor: torch.Tensor, height: int, width: int) -> torch.Tensor: |
| 44 | """ |
| 45 | Resize tensor preserving aspect ratio (filling target), then center crop to exact dimensions. |
| 46 | Args: |
| 47 | latent: Input tensor with shape (H, W, C) or (F, H, W, C) |
| 48 | height: Target height |
| 49 | width: Target width |
| 50 | Returns: |
| 51 | Tensor with shape (1, C, 1, height, width) for 3D input or (1, C, F, height, width) for 4D input |
| 52 | """ |
| 53 | if tensor.ndim == 3: |
| 54 | tensor = rearrange(tensor, "h w c -> 1 c h w") |
| 55 | elif tensor.ndim == 4: |
| 56 | tensor = rearrange(tensor, "f h w c -> f c h w") |
| 57 | else: |
| 58 | raise ValueError(f"Expected input with 3 or 4 dimensions; got shape {tensor.shape}.") |
| 59 | |
| 60 | _, _, src_h, src_w = tensor.shape |
| 61 | |
| 62 | scale = max(height / src_h, width / src_w) |
| 63 | # Use ceil to avoid floating-point rounding causing new_h/new_w to be |
| 64 | # slightly smaller than target, which would result in negative crop offsets. |
| 65 | new_h = math.ceil(src_h * scale) |
| 66 | new_w = math.ceil(src_w * scale) |
| 67 | |
| 68 | tensor = torch.nn.functional.interpolate(tensor, size=(new_h, new_w), mode="bilinear", align_corners=False) |
| 69 | |
| 70 | crop_top = (new_h - height) // 2 |
| 71 | crop_left = (new_w - width) // 2 |
| 72 | tensor = tensor[:, :, crop_top : crop_top + height, crop_left : crop_left + width] |
| 73 | |
| 74 | tensor = rearrange(tensor, "f c h w -> 1 c f h w") |
| 75 | return tensor |
| 76 | |
| 77 | |
| 78 | def normalize_latent(latent: torch.Tensor, device: torch.device, dtype: torch.dtype) -> torch.Tensor: |
| 79 | return (latent / 127.5 - 1.0).to(device=device, dtype=dtype) |
| 80 | |
| 81 | |
| 82 | def load_image_conditioning( |
| 83 | image_path: str, |
| 84 | height: int, |
| 85 | width: int, |
| 86 | dtype: torch.dtype, |
| 87 | device: torch.device, |
| 88 | crf: int = DEFAULT_IMAGE_CRF, |
| 89 | ) -> torch.Tensor: |
| 90 | """ |
| 91 | Loads an image from a path and preprocesses it for conditioning. |
| 92 | Note: The image is resized to the nearest multiple of 2 for compatibility with video codecs. |
| 93 | """ |
| 94 | image = decode_image(image_path=image_path) |
| 95 | image = preprocess(image=image, crf=crf) |
| 96 | image = torch.tensor(image, dtype=torch.float32, device=device) |
| 97 | image = resize_and_center_crop(image, height, width) |
| 98 | image = normalize_latent(image, device, dtype) |
| 99 | return image |
| 100 | |
| 101 | |
| 102 | def load_video_conditioning( |
| 103 | video_path: str, height: int, width: int, frame_cap: int, dtype: torch.dtype, device: torch.device |
| 104 | ) -> torch.Tensor: |
| 105 | """ |
| 106 | Loads a video from a path and preprocesses it for conditioning. |
| 107 | Note: The video is resized to the nearest multiple of 2 for compatibility with video codecs. |
| 108 | """ |
| 109 | frames = decode_video_from_file(path=video_path, frame_cap=frame_cap, device=device) |
| 110 | result = None |
| 111 | for f in frames: |
| 112 | frame = resize_and_center_crop(f.to(torch.float32), height, width) |
| 113 | frame = normalize_latent(frame, device, dtype) |
| 114 | result = frame if result is None else torch.cat([result, frame], dim=2) |
| 115 | return result |
| 116 | |
| 117 | |
| 118 | def decode_image(image_path: str) -> np.ndarray: |
| 119 | image = Image.open(image_path) |
| 120 | np_array = np.array(image)[..., :3] |
| 121 | return np_array |
| 122 | |
| 123 | |
| 124 | def _write_audio(container: av.container.Container, audio_stream: av.audio.AudioStream, audio: Audio) -> None: |
| 125 | samples = audio.waveform |
| 126 | if samples.ndim == 1: |
| 127 | samples = samples[:, None] |
| 128 | |
| 129 | if samples.shape[1] != 2 and samples.shape[0] == 2: |
| 130 | samples = samples.T |
| 131 | |
| 132 | if samples.shape[1] != 2: |
| 133 | raise ValueError(f"Expected samples with 2 channels; got shape {samples.shape}.") |
| 134 | |
| 135 | # Convert to int16 packed for ingestion; resampler converts to encoder fmt. |
| 136 | if samples.dtype != torch.int16: |
| 137 | samples = torch.clip(samples, -1.0, 1.0) |
| 138 | samples = (samples * 32767.0).to(torch.int16) |
| 139 | |
| 140 | frame_in = av.AudioFrame.from_ndarray( |
| 141 | samples.contiguous().reshape(1, -1).cpu().numpy(), |
| 142 | format="s16", |
| 143 | layout="stereo", |
| 144 | ) |
| 145 | frame_in.sample_rate = audio.sampling_rate |
| 146 | |
| 147 | _resample_audio(container, audio_stream, frame_in) |
| 148 | |
| 149 | |
| 150 | def _prepare_audio_stream(container: av.container.Container, audio_sample_rate: int) -> av.audio.AudioStream: |
| 151 | """ |
| 152 | Prepare the audio stream for writing. |
| 153 | """ |
| 154 | audio_stream = container.add_stream("aac", rate=audio_sample_rate) |
| 155 | audio_stream.codec_context.sample_rate = audio_sample_rate |
| 156 | audio_stream.codec_context.layout = "stereo" |
| 157 | audio_stream.codec_context.time_base = Fraction(1, audio_sample_rate) |
| 158 | return audio_stream |
| 159 | |
| 160 | |
| 161 | def _resample_audio( |
| 162 | container: av.container.Container, audio_stream: av.audio.AudioStream, frame_in: av.AudioFrame |
| 163 | ) -> None: |
| 164 | cc = audio_stream.codec_context |
| 165 | |
| 166 | # Use the encoder's format/layout/rate as the *target* |
| 167 | target_format = cc.format or "fltp" # AAC → usually fltp |
| 168 | target_layout = cc.layout or "stereo" |
| 169 | target_rate = cc.sample_rate or frame_in.sample_rate |
| 170 | |
| 171 | audio_resampler = av.audio.resampler.AudioResampler( |
| 172 | format=target_format, |
| 173 | layout=target_layout, |
| 174 | rate=target_rate, |
| 175 | ) |
| 176 | |
| 177 | audio_next_pts = 0 |
| 178 | for rframe in audio_resampler.resample(frame_in): |
| 179 | if rframe.pts is None: |
| 180 | rframe.pts = audio_next_pts |
| 181 | audio_next_pts += rframe.samples |
| 182 | rframe.sample_rate = frame_in.sample_rate |
| 183 | container.mux(audio_stream.encode(rframe)) |
| 184 | |
| 185 | # flush audio encoder |
| 186 | for packet in audio_stream.encode(): |
| 187 | container.mux(packet) |
| 188 | |
| 189 | |
| 190 | def encode_video( |
| 191 | video: torch.Tensor | Iterator[torch.Tensor], |
| 192 | fps: int, |
| 193 | audio: Audio | None, |
| 194 | output_path: str, |
| 195 | video_chunks_number: int, |
| 196 | ) -> None: |
| 197 | if isinstance(video, torch.Tensor): |
| 198 | video = iter([video]) |
| 199 | |
| 200 | first_chunk = next(video) |
| 201 | |
| 202 | _, height, width, _ = first_chunk.shape |
| 203 | |
| 204 | container = av.open(output_path, mode="w") |
| 205 | stream = container.add_stream("libx264", rate=int(fps)) |
| 206 | stream.width = width |
| 207 | stream.height = height |
| 208 | stream.pix_fmt = "yuv420p" |
| 209 | |
| 210 | if audio is not None: |
| 211 | audio_stream = _prepare_audio_stream(container, audio.sampling_rate) |
| 212 | |
| 213 | def all_tiles( |
| 214 | first_chunk: torch.Tensor, tiles_generator: Generator[tuple[torch.Tensor, int], None, None] |
| 215 | ) -> Generator[tuple[torch.Tensor, int], None, None]: |
| 216 | yield first_chunk |
| 217 | yield from tiles_generator |
| 218 | |
| 219 | for video_chunk in tqdm(all_tiles(first_chunk, video), total=video_chunks_number): |
| 220 | video_chunk_cpu = video_chunk.to("cpu").numpy() |
| 221 | for frame_array in video_chunk_cpu: |
| 222 | frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24") |
| 223 | for packet in stream.encode(frame): |
| 224 | container.mux(packet) |
| 225 | |
| 226 | # Flush encoder |
| 227 | for packet in stream.encode(): |
| 228 | container.mux(packet) |
| 229 | |
| 230 | if audio is not None: |
| 231 | _write_audio(container, audio_stream, audio) |
| 232 | |
| 233 | container.close() |
| 234 | logger.info(f"Video saved to {output_path}") |
| 235 | |
| 236 | |
| 237 | _INT_FORMAT_MAX: dict[str, float] = { |
| 238 | "u8": 128.0, |
| 239 | "u8p": 128.0, |
| 240 | "s16": 32768.0, |
| 241 | "s16p": 32768.0, |
| 242 | "s32": 2147483648.0, |
| 243 | "s32p": 2147483648.0, |
| 244 | } |
| 245 | |
| 246 | |
| 247 | def _audio_frame_to_float(frame: av.AudioFrame) -> np.ndarray: |
| 248 | """Convert an audio frame to a float32 ndarray with values in [-1, 1] and shape (channels, samples).""" |
| 249 | fmt = frame.format.name |
| 250 | arr = frame.to_ndarray().astype(np.float32) |
| 251 | if fmt in _INT_FORMAT_MAX: |
| 252 | arr = arr / _INT_FORMAT_MAX[fmt] |
| 253 | if not frame.format.is_planar: |
| 254 | # Interleaved formats have shape (1, samples * channels) — reshape to (channels, samples). |
| 255 | channels = len(frame.layout.channels) |
| 256 | arr = arr.reshape(-1, channels).T |
| 257 | return arr |
| 258 | |
| 259 | |
| 260 | def get_videostream_metadata(path: str) -> tuple[float, int, int, int]: |
| 261 | """Read video stream metadata: (fps, num_frames, width, height). |
| 262 | If frame count is missing in the container, decodes the stream to count frames. |
| 263 | """ |
| 264 | container = av.open(path) |
| 265 | try: |
| 266 | video_stream = next(s for s in container.streams if s.type == "video") |
| 267 | fps = float(video_stream.average_rate) |
| 268 | num_frames = video_stream.frames or 0 |
| 269 | if num_frames == 0: |
| 270 | num_frames = sum(1 for _ in container.decode(video_stream)) |
| 271 | width = video_stream.codec_context.width |
| 272 | height = video_stream.codec_context.height |
| 273 | return fps, num_frames, width, height |
| 274 | finally: |
| 275 | container.close() |
| 276 | |
| 277 | |
| 278 | def decode_audio_from_file( |
| 279 | path: str, device: torch.device, start_time: float = 0.0, max_duration: float | None = None |
| 280 | ) -> Audio | None: |
| 281 | """Decodes audio from a file, optionally seeking to a start time and limiting duration. |
| 282 | Args: |
| 283 | path: Path to the audio/video file containing an audio stream. |
| 284 | device: Device to place the resulting tensor on. |
| 285 | start_time: Start time in seconds to begin reading audio from. |
| 286 | max_duration: Maximum audio duration in seconds. If None, reads to end of stream. |
| 287 | Returns: |
| 288 | An Audio object with waveform of shape (1, channels, samples), or None if no audio stream. |
| 289 | """ |
| 290 | container = av.open(path) |
| 291 | try: |
| 292 | audio_stream = next(s for s in container.streams if s.type == "audio") |
| 293 | except StopIteration: |
| 294 | container.close() |
| 295 | return None |
| 296 | |
| 297 | sample_rate = audio_stream.rate |
| 298 | start_pts = int(start_time / audio_stream.time_base) |
| 299 | end_time = start_time + max_duration if max_duration else audio_stream.duration * audio_stream.time_base |
| 300 | container.seek(start_pts, stream=audio_stream) |
| 301 | |
| 302 | samples = [] |
| 303 | first_frame_time = None |
| 304 | for frame in container.decode(audio=0): |
| 305 | if frame.pts is None: |
| 306 | continue |
| 307 | frame_time = float(frame.pts * audio_stream.time_base) |
| 308 | frame_end = frame_time + frame.samples / frame.sample_rate |
| 309 | if frame_end < start_time: |
| 310 | continue |
| 311 | if frame_time > end_time: |
| 312 | break |
| 313 | if first_frame_time is None: |
| 314 | first_frame_time = frame_time |
| 315 | samples.append(_audio_frame_to_float(frame)) |
| 316 | |
| 317 | container.close() |
| 318 | |
| 319 | if not samples: |
| 320 | return None |
| 321 | |
| 322 | audio = np.concatenate(samples, axis=-1) |
| 323 | |
| 324 | # Trim samples that fall outside the requested [start_time, start_time + max_duration] window. |
| 325 | # Audio codecs decode in fixed-size frames whose boundaries may not align with the requested |
| 326 | # time range, so the first frame can start before start_time and the last frame can end after |
| 327 | # start_time + max_duration. |
| 328 | skip_samples = round((start_time - first_frame_time) * sample_rate) |
| 329 | if skip_samples > 0: |
| 330 | audio = audio[..., skip_samples:] |
| 331 | |
| 332 | if max_duration is not None: |
| 333 | max_samples = round(max_duration * sample_rate) |
| 334 | audio = audio[..., :max_samples] |
| 335 | |
| 336 | waveform = torch.from_numpy(audio).to(device).unsqueeze(0) |
| 337 | |
| 338 | return Audio(waveform=waveform, sampling_rate=sample_rate) |
| 339 | |
| 340 | |
| 341 | def decode_video_from_file(path: str, frame_cap: int, device: DeviceLikeType) -> Generator[torch.Tensor]: |
| 342 | container = av.open(path) |
| 343 | try: |
| 344 | video_stream = next(s for s in container.streams if s.type == "video") |
| 345 | for frame in container.decode(video_stream): |
| 346 | tensor = torch.tensor(frame.to_rgb().to_ndarray(), dtype=torch.uint8, device=device).unsqueeze(0) |
| 347 | yield tensor |
| 348 | frame_cap = frame_cap - 1 |
| 349 | if frame_cap == 0: |
| 350 | break |
| 351 | finally: |
| 352 | container.close() |
| 353 | |
| 354 | |
| 355 | def encode_single_frame(output_file: str, image_array: np.ndarray, crf: float) -> None: |
| 356 | container = av.open(output_file, "w", format="mp4") |
| 357 | try: |
| 358 | stream = container.add_stream("libx264", rate=1, options={"crf": str(crf), "preset": "veryfast"}) |
| 359 | # Round to nearest multiple of 2 for compatibility with video codecs |
| 360 | height = image_array.shape[0] // 2 * 2 |
| 361 | width = image_array.shape[1] // 2 * 2 |
| 362 | image_array = image_array[:height, :width] |
| 363 | stream.height = height |
| 364 | stream.width = width |
| 365 | av_frame = av.VideoFrame.from_ndarray(image_array, format="rgb24").reformat(format="yuv420p") |
| 366 | container.mux(stream.encode(av_frame)) |
| 367 | container.mux(stream.encode()) |
| 368 | finally: |
| 369 | container.close() |
| 370 | |
| 371 | |
| 372 | def decode_single_frame(video_file: str) -> np.array: |
| 373 | container = av.open(video_file) |
| 374 | try: |
| 375 | stream = next(s for s in container.streams if s.type == "video") |
| 376 | frame = next(container.decode(stream)) |
| 377 | finally: |
| 378 | container.close() |
| 379 | return frame.to_ndarray(format="rgb24") |
| 380 | |
| 381 | |
| 382 | def preprocess(image: np.array, crf: float = DEFAULT_IMAGE_CRF) -> np.array: |
| 383 | if crf == 0: |
| 384 | return image |
| 385 | |
| 386 | with BytesIO() as output_file: |
| 387 | encode_single_frame(output_file, image, crf) |
| 388 | video_bytes = output_file.getvalue() |
| 389 | with BytesIO(video_bytes) as video_file: |
| 390 | image_array = decode_single_frame(video_file) |
| 391 | return image_array |
| 392 |