indigo/LTX-2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Automated PR - 2026-04-23

Browse Source
This commit is contained in:
github-actions[bot]
2026-04-23 12:43:54 +00:00
parent a2c3f24078
commit b604d3fab3
49 changed files with 2664 additions and 568 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-pipelines/src/ltx_pipelines/utils/args.py
+12 -8
View File
@@ -12,6 +12,7 @@ from ltx_pipelines.utils.constants import (
LTX_2_3_PARAMS,
PipelineParams,
)
from ltx_pipelines.utils.types import OffloadMode
class ImageConditioningInput(NamedTuple):
@@ -231,16 +232,19 @@ def basic_arg_parser(
except ValueError as e:
raise argparse.ArgumentTypeError(f"must be an integer, got {value}") from e
# Layer streaming
# Weight offloading
parser.add_argument(
"--streaming-prefetch-count",
type=_positive_int,
default=None,
metavar="N",
"--offload",
dest="offload_mode",
type=OffloadMode,
default=OffloadMode.NONE,
choices=list(OffloadMode),
help=(
"Enable layer streaming prefetching N layers ahead. "
"At most 1 + N layers reside on GPU at once. "
"Must be >= 1. Example: --streaming-prefetch-count 2"
"Weight offloading strategy. "
"'none' keeps all weights on GPU (default). "
"'cpu' pins weights in CPU RAM, streams to GPU per layer. "
"'disk' reads weights from disk on demand (lowest memory). "
"Example: --offload cpu"
),
)
packages/ltx-pipelines/src/ltx_pipelines/utils/blocks.py
+128 -83
View File
@@ -15,11 +15,11 @@ from typing import Callable, TypeVar
import torch
from ltx_core.batch_split import BatchSplitAdapter
from ltx_core.block_streaming import DISK_CPU_SLOTS, StreamingModelBuilder
from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.noisers import Noiser
from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.layer_streaming import LayerStreamingWrapper
from ltx_core.loader import SDOps
from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import DummyRegistry, Registry
@@ -70,7 +70,7 @@ from ltx_pipelines.utils.helpers import (
generate_enhanced_prompt,
)
from ltx_pipelines.utils.samplers import euler_denoising_loop
from ltx_pipelines.utils.types import Denoiser, ModalitySpec
from ltx_pipelines.utils.types import Denoiser, ModalitySpec, OffloadMode
logger = logging.getLogger(__name__)
@@ -85,34 +85,24 @@ _M = TypeVar("_M", bound=torch.nn.Module)
@contextmanager
def _streaming_model(
model: _M,
layers_attr: str,
builder: StreamingModelBuilder,
offload_mode: OffloadMode,
target_device: torch.device,
prefetch_count: int,
) -> Iterator[_M]:
"""Wrap *model* with :class:`LayerStreamingWrapper`, yield it, then tear down."""
wrapped = LayerStreamingWrapper(
model,
layers_attr=layers_attr,
dtype: torch.dtype,
) -> Iterator:
"""Build a streaming wrapper, yield it, then tear down and free memory."""
cpu_slots_count = DISK_CPU_SLOTS if offload_mode == OffloadMode.DISK else None
wrapped = builder.build(
target_device=target_device,
prefetch_count=prefetch_count,
dtype=dtype,
cpu_slots_count=cpu_slots_count,
)
try:
yield wrapped # type: ignore[misc]
yield wrapped
finally:
wrapped.teardown()
wrapped.to("meta")
cleanup_memory()
# Flush the host (pinned) memory cache so that freed pinned pages are
# returned to the OS. Without this, sequential streaming models
# (e.g. text encoder then transformer) exhaust host memory because the
# CachingHostAllocator keeps freed blocks cached indefinitely.
torch.cuda.synchronize(device=target_device)
try:
if hasattr(torch._C, "_host_emptyCache"):
torch._C._host_emptyCache()
except Exception:
logger.warning("Host empty cache cleanup failed; ignoring.", exc_info=True)
def _build_state(
@@ -163,11 +153,30 @@ class DiffusionStage:
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
offload_mode: OffloadMode = OffloadMode.NONE,
) -> None:
if offload_mode != OffloadMode.NONE:
if torch_compile:
raise ValueError("torch.compile is not supported with layer streaming")
if quantization is not None:
raise ValueError("quantization is not supported with layer streaming")
self._streaming_builder = StreamingModelBuilder(
model_class_configurator=LTXModelConfigurator,
model_path=checkpoint_path,
model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
loras=tuple(loras),
registry=registry or DummyRegistry(),
blocks_attr="velocity_model.transformer_blocks",
blocks_prefix="transformer_blocks",
state_dict_prefix="velocity_model.",
model_wrapper=lambda m: X0Model(m).eval(),
)
self._dtype = dtype
self._device = device
self._quantization = quantization
self._torch_compile = torch_compile
self._offload_mode = offload_mode
self._transformer_builder = Builder(
model_path=checkpoint_path,
model_class_configurator=LTXModelConfigurator,
@@ -205,22 +214,21 @@ class DiffusionStage:
builder = self._transformer_builder.with_module_ops(module_ops).with_sd_ops(sd_ops).with_loras(loras)
return X0Model(builder.build(device=target, **kwargs)).to(target).eval()
def _transformer_ctx(
self,
streaming_prefetch_count: int | None,
**kwargs: object,
) -> AbstractContextManager:
if streaming_prefetch_count is not None:
return _streaming_model(
self._build_transformer(device=torch.device("cpu"), **kwargs),
layers_attr="velocity_model.transformer_blocks",
target_device=self._device,
prefetch_count=streaming_prefetch_count,
)
def _transformer_ctx(self, **kwargs: object) -> AbstractContextManager:
if self._offload_mode != OffloadMode.NONE:
return _streaming_model(self._streaming_builder, self._offload_mode, self._device, self._dtype)
return gpu_model(self._build_transformer(**kwargs))
def __call__( # noqa: PLR0913
def model_context(self, **kwargs: object) -> AbstractContextManager:
"""Build the transformer, yield it, then free its memory on exit.
Keyword arguments are forwarded to the underlying builder (e.g.
``video_tools`` required by ``TiledDataParallelBuilder``).
"""
return self._transformer_ctx(**kwargs)
def run( # noqa: PLR0913
self,
transformer: object,
denoiser: Denoiser,
sigmas: torch.Tensor,
noiser: Noiser,
@@ -232,27 +240,14 @@ class DiffusionStage:
audio: ModalitySpec | None = None,
stepper: DiffusionStepProtocol | None = None,
loop: Callable[..., tuple[LatentState | None, LatentState | None]] | None = None,
streaming_prefetch_count: int | None = None,
max_batch_size: int = 1,
) -> tuple[LatentState | None, LatentState | None]:
"""Build transformer → run denoising loop → free transformer.
Args:
width: Output width in pixels.
height: Output height in pixels.
frames: Number of output frames.
fps: Frame rate.
loop: Denoising loop function. Must accept
``(sigmas, video_state, audio_state, stepper, transformer, denoiser)``
as the first six positional arguments. When ``None``, resolves to
:func:`euler_denoising_loop` at call time.
streaming_prefetch_count: When set, build the transformer on CPU and
wrap with :class:`LayerStreamingWrapper` for memory-efficient
inference, prefetching this many layers ahead.
max_batch_size: Maximum batch size per transformer forward pass.
Guided denoisers make up to 4 transformer calls per step.
When set to a value > 1, the transformer batches multiple
calls together, reducing layer-streaming PCIe transfers.
Default ``1`` preserves sequential behavior.
"""Run denoising with a pre-built transformer.
Same semantics as ``__call__`` but accepts a pre-built transformer so
the model can be shared across multiple calls (e.g. tiled inference
inside a single ``model_context()`` block). Audio supports
``ModalitySpec(frozen=True)`` to keep the latent unchanged throughout
denoising while still providing cross-modal context to the transformer.
Returns ``(video_state | None, audio_state | None)`` with cleared
conditionings and unpatchified latents for present modalities.
"""
@@ -261,7 +256,6 @@ class DiffusionStage:
if loop is None:
loop = euler_denoising_loop
if stepper is None:
stepper = EulerDiffusionStep()
@@ -281,28 +275,70 @@ class DiffusionStage:
audio_tools = AudioLatentTools(AudioPatchifier(patch_size=1), a_shape)
audio_state = _build_state(audio, audio_tools, noiser, self._dtype, self._device)
with self._transformer_ctx(streaming_prefetch_count, video_tools=video_tools) as base_transformer:
transformer = BatchSplitAdapter(base_transformer, max_batch_size=max_batch_size)
video_state, audio_state = loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
transformer=transformer,
denoiser=denoiser,
)
wrapped = BatchSplitAdapter(transformer, max_batch_size=max_batch_size) # type: ignore[arg-type]
video_state, audio_state = loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
transformer=wrapped,
denoiser=denoiser,
)
# Post-process: clear conditionings and unpatchify
if video_state is not None and video_tools is not None:
video_state = video_tools.clear_conditioning(video_state)
video_state = video_tools.unpatchify(video_state)
if audio_state is not None and audio_tools is not None:
audio_state = audio_tools.clear_conditioning(audio_state)
audio_state = audio_tools.unpatchify(audio_state)
return video_state, audio_state
def __call__( # noqa: PLR0913
self,
denoiser: Denoiser,
sigmas: torch.Tensor,
noiser: Noiser,
width: int,
height: int,
frames: int,
fps: float,
video: ModalitySpec | None = None,
audio: ModalitySpec | None = None,
stepper: DiffusionStepProtocol | None = None,
loop: Callable[..., tuple[LatentState | None, LatentState | None]] | None = None,
max_batch_size: int = 1,
) -> tuple[LatentState | None, LatentState | None]:
"""Build transformer -> run denoising loop -> free transformer.
Returns ``(video_state | None, audio_state | None)`` with cleared
conditionings and unpatchified latents for present modalities.
"""
# Build video_tools up front so it can be forwarded to the transformer
# context (required by TiledDataParallelBuilder in multi-GPU mode).
# `run()` rebuilds its own tools internally; the duplication is cheap.
video_tools: LatentTools | None = None
if video is not None:
pixel_shape = VideoPixelShape(batch=1, frames=frames, height=height, width=width, fps=fps)
v_shape = VideoLatentShape.from_pixel_shape(pixel_shape)
video_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape, fps)
with self._transformer_ctx(video_tools=video_tools) as transformer:
return self.run(
transformer,
denoiser,
sigmas,
noiser,
width,
height,
frames,
fps,
video,
audio,
stepper,
loop,
max_batch_size,
)
# ---------------------------------------------------------------------------
# PromptEncoder
@@ -322,9 +358,11 @@ class PromptEncoder:
dtype: torch.dtype,
device: torch.device,
registry: Registry | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
) -> None:
self._dtype = dtype
self._device = device
self._offload_mode = offload_mode
module_ops = module_ops_from_gemma_root(gemma_root)
model_folder = find_matching_file(gemma_root, "model*.safetensors").parent
@@ -337,6 +375,15 @@ class PromptEncoder:
module_ops=(GEMMA_MODEL_OPS, *module_ops),
registry=registry or DummyRegistry(),
)
self._streaming_text_encoder_builder = StreamingModelBuilder(
model_path=tuple(weight_paths),
model_class_configurator=GemmaTextEncoderConfigurator,
model_sd_ops=GEMMA_LLM_KEY_OPS,
module_ops=(GEMMA_MODEL_OPS, *module_ops),
registry=registry or DummyRegistry(),
blocks_attr="model.model.language_model.layers",
blocks_prefix="model.model.language_model.layers",
)
self._embeddings_processor_builder = Builder(
model_path=checkpoint_path,
model_class_configurator=EmbeddingsProcessorConfigurator,
@@ -344,17 +391,9 @@ class PromptEncoder:
registry=registry or DummyRegistry(),
)
def _text_encoder_ctx(
self,
streaming_prefetch_count: int | None,
) -> AbstractContextManager:
if streaming_prefetch_count is not None:
return _streaming_model(
self._text_encoder_builder.build(device=torch.device("cpu"), dtype=self._dtype).eval(),
layers_attr="model.model.language_model.layers",
target_device=self._device,
prefetch_count=streaming_prefetch_count,
)
def _text_encoder_ctx(self) -> AbstractContextManager:
if self._offload_mode != OffloadMode.NONE:
return _streaming_model(self._streaming_text_encoder_builder, self._offload_mode, self._device, self._dtype)
return gpu_model(self._text_encoder_builder.build(device=self._device, dtype=self._dtype).eval())
def __call__(
@@ -364,10 +403,9 @@ class PromptEncoder:
enhance_first_prompt: bool = False,
enhance_prompt_image: str | None = None,
enhance_prompt_seed: int = 42,
streaming_prefetch_count: int | None = None,
) -> list[EmbeddingsProcessorOutput]:
"""Encode *prompts* through Gemma → embeddings processor, freeing each model after use."""
with self._text_encoder_ctx(streaming_prefetch_count) as text_encoder:
"""Encode *prompts* through Gemma -> embeddings processor, freeing each model after use."""
with self._text_encoder_ctx() as text_encoder:
if enhance_first_prompt:
prompts = list(prompts)
prompts[0] = generate_enhanced_prompt(
@@ -490,10 +528,17 @@ class VideoDecoder:
latent: torch.Tensor,
tiling_config: TilingConfig | None = None,
generator: torch.Generator | None = None,
*,
output_dtype: torch.dtype = torch.uint8,
) -> Iterator[torch.Tensor]:
"""Decode *latent* to pixel-space video chunks. Decoder freed after exhaustion."""
"""Decode *latent* to pixel-space video chunks. Decoder freed after exhaustion.
Args:
output_dtype: Target dtype for output tensors. ``torch.uint8``
(default) maps to ``[0, 255]``. Any floating dtype returns
``[0, 1]`` cast to that dtype.
"""
decoder = self._decoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
return _cleanup_iter(decoder.decode_video(latent, tiling_config, generator), decoder)
return _cleanup_iter(decoder.decode_video(latent, tiling_config, generator, output_dtype=output_dtype), decoder)
# ---------------------------------------------------------------------------
packages/ltx-pipelines/src/ltx_pipelines/utils/helpers.py
+5
View File
@@ -37,6 +37,11 @@ def cleanup_memory() -> None:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.synchronize()
try:
if hasattr(torch._C, "_host_emptyCache"):
torch._C._host_emptyCache()
except Exception:
logging.warning("Host empty cache cleanup failed; ignoring.", exc_info=True)
def _conform_latent_length(latent: torch.Tensor, expected_frames_count: int) -> torch.Tensor:
packages/ltx-pipelines/src/ltx_pipelines/utils/media_io.py
+205
View File
@@ -1,23 +1,34 @@
import enum
import logging
import math
from collections.abc import Generator, Iterator
from fractions import Fraction
from io import BytesIO
from pathlib import Path
import av
import numpy as np
import OpenImageIO
import torch
from einops import rearrange
from PIL import Image
from torch._prims_common import DeviceLikeType
from tqdm import tqdm
from ltx_core.hdr import LogC3
from ltx_core.types import Audio, VideoPixelShape
from ltx_pipelines.utils.constants import DEFAULT_IMAGE_CRF
logger = logging.getLogger(__name__)
class ResizeMode(enum.Enum):
"""How to fit a conditioning video to the target resolution."""
CENTER_CROP = "center_crop"
REFLECT_PAD = "reflect_pad"
def resize_aspect_ratio_preserving(image: torch.Tensor, long_side: int) -> torch.Tensor:
"""
Resize image preserving aspect ratio (filling target long side).
@@ -79,6 +90,16 @@ def normalize_latent(latent: torch.Tensor, device: torch.device, dtype: torch.dt
return (latent / 127.5 - 1.0).to(device=device, dtype=dtype)
def to_vae_range(x: torch.Tensor) -> torch.Tensor:
"""Map [0, 1] to [-1, 1] (VAE input convention)."""
return torch.clamp(x, 0.0, 1.0) * 2.0 - 1.0
def from_vae_range(z: torch.Tensor) -> torch.Tensor:
"""Map [-1, 1] (VAE output convention) to [0, 1]."""
return torch.clamp((z + 1.0) / 2.0, 0.0, 1.0)
def load_image_and_preprocess(
image_path: str,
height: int,
@@ -124,6 +145,108 @@ def video_preprocess(
return result
def align_resolution(
width: int,
height: int,
resize_mode: ResizeMode,
divisor: int = 64,
) -> tuple[int, int, int, int]:
"""Compute aligned generation dimensions and crop-back size.
Args:
width: Source video width (need not be aligned).
height: Source video height (need not be aligned).
resize_mode: CENTER_CROP rounds down; REFLECT_PAD rounds up.
divisor: Alignment divisor (default 64 for two-stage pipelines).
Returns:
``(gen_width, gen_height, crop_width, crop_height)`` where
``gen_*`` are multiples of *divisor* and ``crop_*`` are the
original dimensions to trim back to after decoding. When no
cropping is needed ``crop_*`` equals ``gen_*``.
"""
if resize_mode is ResizeMode.REFLECT_PAD:
gen_w = ((width + divisor - 1) // divisor) * divisor
gen_h = ((height + divisor - 1) // divisor) * divisor
else:
gen_w = (width // divisor) * divisor
gen_h = (height // divisor) * divisor
crop_w = width if gen_w != width else gen_w
crop_h = height if gen_h != height else gen_h
return gen_w, gen_h, crop_w, crop_h
def resize_and_reflect_pad(tensor: torch.Tensor, height: int, width: int) -> torch.Tensor:
"""Resize tensor to fit within target, then reflect-pad to exact dimensions.
Unlike resize_and_center_crop which stretches and crops, this preserves the
original aspect ratio and pads the shorter dimension with reflected pixels.
When the target is already >= the source in both dimensions, interpolation
is skipped entirely to preserve original pixels.
Args:
tensor: Input with shape (H, W, C) or (F, H, W, C)
height: Target height
width: Target width
Returns:
Tensor with shape (1, C, 1, height, width) for 3D or (1, C, F, height, width) for 4D
"""
if tensor.ndim == 3:
tensor = rearrange(tensor, "h w c -> 1 c h w")
elif tensor.ndim == 4:
tensor = rearrange(tensor, "f h w c -> f c h w")
else:
raise ValueError(f"Expected input with 3 or 4 dimensions; got shape {tensor.shape}.")
_, _, src_h, src_w = tensor.shape
if height >= src_h and width >= src_w:
new_h, new_w = src_h, src_w
else:
scale = min(height / src_h, width / src_w)
new_h = round(src_h * scale)
new_w = round(src_w * scale)
tensor = torch.nn.functional.interpolate(tensor, size=(new_h, new_w), mode="bilinear", align_corners=False)
pad_bottom = height - new_h
pad_right = width - new_w
if pad_bottom > 0 or pad_right > 0:
pad_mode = "reflect" if pad_bottom < new_h and pad_right < new_w else "replicate"
tensor = torch.nn.functional.pad(tensor, (0, pad_right, 0, pad_bottom), mode=pad_mode)
tensor = rearrange(tensor, "f c h w -> 1 c f h w")
return tensor
def load_video_conditioning_hdr(
video_path: str,
height: int,
width: int,
frame_cap: int,
dtype: torch.dtype,
device: torch.device,
hdr_transform: str = "logc3",
resize_mode: ResizeMode = ResizeMode.CENTER_CROP,
) -> Iterator[torch.Tensor]:
"""Load a video and yield preprocessed frames for HDR IC-LoRA conditioning.
Decodes through the standard path and applies the LDR compression that
matches training. Callers are responsible for providing Rec.709 SDR
input — the HDR IC-LoRA was trained on that color space.
Args:
hdr_transform: LDR-compression name (currently only ``logc3``).
resize_mode: How to fit the video to the target resolution.
Yields:
Per-frame tensors of shape ``(1, C, 1, height, width)``.
"""
if hdr_transform != "logc3":
raise ValueError(f"Unsupported HDR transform: {hdr_transform}")
resize_fn = resize_and_reflect_pad if resize_mode is ResizeMode.REFLECT_PAD else resize_and_center_crop
for f in decode_video_by_frame(path=video_path, frame_cap=frame_cap, device=device):
frame = resize_fn(f.to(torch.float32), height, width)
ldr = (frame / 255.0).clamp(0.0, 1.0)
compressed = LogC3().compress_ldr(ldr)
yield to_vae_range(compressed).to(device=device, dtype=dtype)
def decode_image(image_path: str) -> np.ndarray:
image = Image.open(image_path)
np_array = np.array(image)[..., :3]
@@ -481,3 +604,85 @@ def preprocess(image: np.array, crf: float = DEFAULT_IMAGE_CRF) -> np.array:
with BytesIO(video_bytes) as video_file:
image_array = decode_single_frame(video_file)
return image_array
def save_exr_tensor(tensor: torch.Tensor, file_path: str | Path, half: bool = False) -> None:
"""Save a single tensor frame as EXR with linear sRGB colorspace metadata.
Args:
tensor: ``[H, W, C]`` or ``[C, H, W]`` float tensor.
file_path: Output path (e.g. ``frame_0000.exr``).
half: Force float16 output with ZIP compression.
"""
if tensor.dim() == 3 and tensor.shape[0] == 3:
tensor = tensor.permute(1, 2, 0)
use_half = half or tensor.dtype in (torch.float16, torch.half)
img_np = np.ascontiguousarray(tensor.cpu().numpy().astype(np.float32))
file_path = str(file_path)
h, w = img_np.shape[:2]
fmt = OpenImageIO.HALF if use_half else OpenImageIO.FLOAT
spec = OpenImageIO.ImageSpec(w, h, 3, fmt)
spec.channelnames = ("R", "G", "B")
spec.attribute("compression", "zip")
spec.attribute("chromaticities", "float[8]", (0.64, 0.33, 0.30, 0.60, 0.15, 0.06, 0.3127, 0.3290))
spec.attribute("colorSpace", "sRGB")
out = OpenImageIO.ImageOutput.create(file_path)
if out is None:
raise RuntimeError(
f"Failed to create EXR writer for '{file_path}'. Ensure OpenImageIO is built with OpenEXR support."
)
try:
if not out.open(file_path, spec):
raise RuntimeError(f"Failed to open EXR file '{file_path}': {out.geterror()}")
if not out.write_image(img_np):
raise RuntimeError(f"Failed to write EXR image '{file_path}': {out.geterror()}")
finally:
out.close()
def _linear_to_srgb(x: np.ndarray) -> np.ndarray:
"""Linear -> sRGB OETF per IEC 61966-2-1. Input assumed in [0, 1]."""
x = np.clip(x, 0.0, 1.0)
return np.where(x <= 0.0031308, x * 12.92, 1.055 * np.power(x, 1.0 / 2.4) - 0.055)
def encode_exr_sequence_to_mp4(exr_dir: Path, output_mp4: Path, frame_rate: float) -> None:
"""Convert a linear EXR frame sequence to sRGB and encode to H.264 .mp4 via PyAV.
Exposure is fixed at EV=0 (no gain). Each EXR frame is clamped to [0, 1],
passed through the sRGB OETF, quantised to 8-bit BGR, and fed to a libx264
stream (crf 18, yuv420p). ``frame_rate`` is the original source video's
frame rate so playback matches the input timing.
"""
import os # noqa: PLC0415
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
import cv2 # noqa: PLC0415
exr_files = sorted(exr_dir.glob("frame_*.exr"))
if not exr_files:
raise FileNotFoundError(f"No EXR frames found in {exr_dir}")
container = av.open(str(output_mp4), mode="w")
stream = container.add_stream("libx264", rate=Fraction(frame_rate).limit_denominator(1000))
stream.pix_fmt = "yuv420p"
stream.options = {"crf": "18", "movflags": "+faststart"}
try:
for i, exr_path in enumerate(exr_files):
hdr = cv2.imread(str(exr_path), cv2.IMREAD_UNCHANGED).astype(np.float32)
sdr = _linear_to_srgb(np.maximum(hdr, 0.0))
bgr8 = (sdr * 255.0 + 0.5).astype(np.uint8)
if i == 0:
stream.height = bgr8.shape[0]
stream.width = bgr8.shape[1]
frame = av.VideoFrame.from_ndarray(bgr8, format="bgr24")
for packet in stream.encode(frame):
container.mux(packet)
for packet in stream.encode():
container.mux(packet)
finally:
container.close()
packages/ltx-pipelines/src/ltx_pipelines/utils/types.py
+21
View File
@@ -1,4 +1,7 @@
from __future__ import annotations
from dataclasses import dataclass, field
from enum import Enum
from typing import Protocol
import torch
@@ -74,3 +77,21 @@ class ModalitySpec:
noise_scale: float = 1.0
frozen: bool = False
initial_latent: torch.Tensor | None = None
class OffloadMode(Enum):
"""Weight offloading strategy.
Controls where model weights reside during inference:
- ``NONE``: All weights on GPU (no streaming). Fastest inference,
requires enough VRAM for the full model (~28 GB for LTX-2).
- ``CPU``: Weights pinned in CPU RAM, streamed layer-by-layer to a
small GPU buffer. First pass reads from disk; subsequent passes
reuse the CPU cache. Requires ~36 GB RAM + ~5 GB VRAM.
- ``DISK``: Weights read from disk on demand through a small CPU
buffer, then streamed to GPU. Every pass re-reads from disk.
Lowest memory: ~5 GB RAM + ~5 GB VRAM.
"""
NONE = "none"
CPU = "cpu"
DISK = "disk"