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/a2vid_two_stage.py
+8 -7
View File
@@ -31,7 +31,7 @@ from ltx_pipelines.utils.helpers import (
get_device,
)
from ltx_pipelines.utils.media_io import decode_audio_from_file, encode_video
from ltx_pipelines.utils.types import ModalitySpec
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
class A2VidPipelineTwoStage:
@@ -53,12 +53,15 @@ class A2VidPipelineTwoStage:
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
self.dtype = torch.bfloat16
self._scheduler = LTX2Scheduler()
self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
self.prompt_encoder = PromptEncoder(
checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
)
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
self.audio_conditioner = AudioConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
self.stage_1 = DiffusionStage(
@@ -69,6 +72,7 @@ class A2VidPipelineTwoStage:
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
offload_mode=offload_mode,
)
stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
self.stage_2 = DiffusionStage(
@@ -79,6 +83,7 @@ class A2VidPipelineTwoStage:
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
offload_mode=offload_mode,
)
self.upsampler = VideoUpsampler(
checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -102,7 +107,6 @@ class A2VidPipelineTwoStage:
audio_max_duration: float | None = None,
tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
max_batch_size: int = 1,
stage_1_sigmas: torch.Tensor | None = None,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
@@ -117,7 +121,6 @@ class A2VidPipelineTwoStage:
[prompt, negative_prompt],
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
streaming_prefetch_count=streaming_prefetch_count,
)
v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
v_context_n, _ = ctx_n.video_encoding, ctx_n.audio_encoding
@@ -183,7 +186,6 @@ class A2VidPipelineTwoStage:
noise_scale=0.0,
initial_latent=encoded_audio_latent,
),
streaming_prefetch_count=streaming_prefetch_count,
max_batch_size=max_batch_size,
)
@@ -223,7 +225,6 @@ class A2VidPipelineTwoStage:
noise_scale=0.0,
initial_latent=encoded_audio_latent,
),
streaming_prefetch_count=streaming_prefetch_count,
)
decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -266,6 +267,7 @@ def main() -> None:
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -294,7 +296,6 @@ def main() -> None:
audio_max_duration=args.audio_max_duration
if args.audio_max_duration is not None
else args.num_frames / args.frame_rate,
streaming_prefetch_count=args.streaming_prefetch_count,
max_batch_size=args.max_batch_size,
)
packages/ltx-pipelines/src/ltx_pipelines/distilled.py
+10 -7
View File
@@ -34,7 +34,7 @@ from ltx_pipelines.utils.helpers import (
get_device,
)
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import ModalitySpec
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
class DistilledPipeline:
@@ -54,12 +54,18 @@ class DistilledPipeline:
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
self.dtype = torch.bfloat16
self.prompt_encoder = PromptEncoder(
distilled_checkpoint_path, gemma_root, self.dtype, self.device, registry=registry
distilled_checkpoint_path,
gemma_root,
self.dtype,
self.device,
registry=registry,
offload_mode=offload_mode,
)
self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
self.stage = DiffusionStage(
@@ -70,6 +76,7 @@ class DistilledPipeline:
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
offload_mode=offload_mode,
)
self.upsampler = VideoUpsampler(
distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -88,7 +95,6 @@ class DistilledPipeline:
images: list[ImageConditioningInput],
tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
) -> tuple[Iterator[torch.Tensor], Audio]:
@@ -102,7 +108,6 @@ class DistilledPipeline:
[prompt],
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
streaming_prefetch_count=streaming_prefetch_count,
)
video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
@@ -130,7 +135,6 @@ class DistilledPipeline:
fps=frame_rate,
video=ModalitySpec(context=video_context, conditionings=stage_1_conditionings),
audio=ModalitySpec(context=audio_context),
streaming_prefetch_count=streaming_prefetch_count,
)
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
@@ -167,7 +171,6 @@ class DistilledPipeline:
noise_scale=stage_2_sigmas[0].item(),
initial_latent=audio_state.latent,
),
streaming_prefetch_count=streaming_prefetch_count,
)
decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -189,6 +192,7 @@ def main() -> None:
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -202,7 +206,6 @@ def main() -> None:
images=args.images,
tiling_config=tiling_config,
enhance_prompt=args.enhance_prompt,
streaming_prefetch_count=args.streaming_prefetch_count,
)
encode_video(
packages/ltx-pipelines/src/ltx_pipelines/hdr_ic_lora.py
+886
View File
@@ -0,0 +1,886 @@
"""HDR IC-LoRA pipeline: two-stage video generation with HDR output.
Extends the standard IC-LoRA pipeline with HDR decode via LogC3 inverse
transform. ``__call__`` returns a **linear HDR float** tensor
``[f, h, w, c]``; tonemapping and EXR saving are the caller's
responsibility.
Text embeddings must be pre-computed externally (e.g. using
``PromptEncoder`` from ``ltx_pipelines.utils.blocks`` with a Gemma text
encoder) and saved as a ``.safetensors`` file with ``video_context``
and ``audio_context`` tensors (via ``safetensors.torch.save_file``).
The path is passed via ``text_embeddings_path``.
Run as a script for batch inference::
python -m ltx_pipelines.hdr_ic_lora \\
--input ./videos/ \\
--output-dir ./hdr-output \\
--distilled-checkpoint-path /models/ltx-2.3-22b-distilled.safetensors \\
--spatial-upsampler-path /models/ltx-2.3-spatial-upscaler-x2-1.0.safetensors \\
--hdr-lora /path/to/hdr_lora.safetensors \\
--text-embeddings /path/to/hdr_scene_emb.safetensors \\
--num-frames 161
Supports resolutions up to 4K (3840x2160 @ 121 frames on 80 GB,
49 frames on 48 GB). The caller is responsible for choosing a resolution
and frame count that fits in GPU memory. See ``--help`` for a reference
table, or use ``ltx_pipelines.utils.vram_budget.max_frames_for_resolution``
to query your specific configuration.
"""
import dataclasses
import logging
from dataclasses import replace
from pathlib import Path
import torch
from einops import rearrange
from safetensors import safe_open
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.patchifiers import VideoLatentPatchifier
from ltx_core.conditioning import (
ConditioningItem,
VideoConditionByReferenceLatent,
)
from ltx_core.hdr import apply_hdr_decode_postprocess
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import Registry
from ltx_core.loader.sd_ops import LTXV_LORA_COMFY_RENAMING_MAP
from ltx_core.modality_tiling import VideoModalityTilingHelper
from ltx_core.model.video_vae import TilingConfig, VideoEncoder
from ltx_core.quantization import QuantizationPolicy
from ltx_core.tiling import DimensionTilingConfig, TileCountConfig
from ltx_core.tools import VideoLatentTools
from ltx_core.types import VideoLatentShape
from ltx_pipelines.utils.blocks import (
DiffusionStage,
ImageConditioner,
VideoDecoder,
VideoUpsampler,
)
from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
from ltx_pipelines.utils.denoisers import SimpleDenoiser
from ltx_pipelines.utils.helpers import get_device, modality_from_latent_state
from ltx_pipelines.utils.media_io import ResizeMode, align_resolution, load_video_conditioning_hdr
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
DEFAULT_NUM_FRAMES = 161
MIN_RESOLUTION = 64
ALIGNMENT_DIVISOR = 64
# Conditioning videos whose spatial resolution (H x W) exceeds this value are
# encoded with the tiled encoder. The default (512 x 768) is suitable for
# H100-80GB. On lower-VRAM GPUs pass tiled_vae_encode_pixel_threshold=256*256
# to the pipeline constructor.
TILED_VAE_ENCODE_PIXEL_THRESHOLD = 512 * 768
_DEFAULT_QUANTIZATION = QuantizationPolicy.fp8_cast()
# Default stage-2 configuration: one refinement phase with modest 2-way tiling
# in every dimension and a short 2-step distilled sigma schedule.
_S2 = STAGE_2_DISTILLED_SIGMA_VALUES
_TILED_2F2H2W_OV8_6 = TileCountConfig(
frames=DimensionTilingConfig(2, 8),
height=DimensionTilingConfig(2, 6),
width=DimensionTilingConfig(2, 6),
)
STAGE2_TILINGS = [_TILED_2F2H2W_OV8_6]
STAGE2_SIGMAS = [[_S2[0], _S2[1], 0.0]]
STAGE2_USE_IC_LORA = [True]
def _clamp_dim_tiling(cfg: DimensionTilingConfig, dim_size: int, axis: str) -> DimensionTilingConfig:
"""Clamp a single dim's tile count and overlap to the latent's extent.
``split_by_count`` requires ``overlap < tile_size``; with
``tile_size = (dim_size + overlap*(n-1)) // n`` this reduces to
``overlap <= dim_size - n``. When the configured overlap exceeds this
bound it is clamped; if the latent is too small to hold ``n`` tiles
at all, tiling falls back to a single tile on this axis.
"""
n = cfg.num_tiles
if n <= 1:
return cfg
if dim_size < n:
logger.warning(
"%s tiling: dim_size=%d < num_tiles=%d; falling back to 1 tile on this axis.",
axis,
dim_size,
n,
)
return DimensionTilingConfig(1, 0)
max_overlap = dim_size - n
if cfg.overlap <= max_overlap:
return cfg
logger.warning(
"%s tiling: overlap=%d exceeds latent bound (%d); clamping to %d.",
axis,
cfg.overlap,
max_overlap,
max_overlap,
)
return DimensionTilingConfig(n, max_overlap)
def _clamp_tile_to_latent(tiling: TileCountConfig, latent_shape: tuple[int, int, int]) -> TileCountConfig:
"""Clamp frame, height, and width tilings to the latent's extents.
``latent_shape`` is ``(F, H, W)`` in latent units.
"""
f, h, w = latent_shape
return replace(
tiling,
frames=_clamp_dim_tiling(tiling.frames, f, "Frame"),
height=_clamp_dim_tiling(tiling.height, h, "Height"),
width=_clamp_dim_tiling(tiling.width, w, "Width"),
)
# Default tiling config (spatial tile 1280 px, overlap 256 px; temporal 32
# frames, overlap 16). On GPUs with < 80 GB VRAM you may need to shrink
# the spatial tile size (e.g. 768) to avoid OOM during VAE decode.
DEFAULT_SPATIAL_TILE = 1280
DEFAULT_SPATIAL_OVERLAP = 256
DEFAULT_TEMPORAL_TILE = 32
DEFAULT_TEMPORAL_OVERLAP = 16
# ---------------------------------------------------------------------------
# HDR LoRA config
# ---------------------------------------------------------------------------
@dataclasses.dataclass(frozen=True)
class HdrLoraConfig:
"""Explicit HDR LoRA parameters.
Read from LoRA safetensors metadata by :func:`read_hdr_lora_config`, or
constructed manually for testing.
"""
hdr_transform: str = "logc3"
reference_downscale_factor: int = 1
def read_hdr_lora_config(lora_path: str) -> HdrLoraConfig | None:
"""Read HDR config from LoRA safetensors metadata.
Returns ``None`` when the LoRA has no HDR metadata.
"""
try:
with safe_open(lora_path, framework="pt") as f:
metadata = f.metadata() or {}
except (OSError, ValueError) as e:
logger.warning("Failed to read metadata from LoRA file '%s': %s", lora_path, e)
return None
hdr_transform = metadata.get("hdr_transform", "")
has_hdr = bool(hdr_transform or metadata.get("use_hdr_transform"))
if not has_hdr:
return None
transform = hdr_transform if hdr_transform and hdr_transform != "true" else "logc3"
scale = int(metadata.get("reference_downscale_factor", 1))
return HdrLoraConfig(hdr_transform=transform, reference_downscale_factor=scale)
# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------
class HDRICLoraPipeline:
"""Two-stage IC-LoRA pipeline with HDR support.
Same two-stage architecture as ICLoraPipeline (half-res generation + 2x
upscale refinement), with HDR decode via LogC3 inverse.
``__call__`` returns a **linear HDR float** tensor ``[f, h, w, c]``.
Tonemapping and EXR saving are the caller's responsibility.
"""
def __init__(
self,
distilled_checkpoint_path: str,
spatial_upsampler_path: str,
hdr_lora: str | Path,
text_embeddings_path: str | Path,
device: torch.device | None = None,
quantization: QuantizationPolicy = _DEFAULT_QUANTIZATION,
registry: Registry | None = None,
hdr_lora_config: HdrLoraConfig | None = None,
tiled_vae_encode_pixel_threshold: int = TILED_VAE_ENCODE_PIXEL_THRESHOLD,
offload_mode: OffloadMode = OffloadMode.NONE,
):
"""
Args:
distilled_checkpoint_path: Path to the distilled model checkpoint.
spatial_upsampler_path: Path to the spatial upsampler checkpoint.
hdr_lora: Path to the HDR IC-LoRA ``.safetensors`` file.
text_embeddings_path: Path to pre-computed text embeddings
(``.safetensors`` file with ``video_context`` and
``audio_context`` tensors).
device: Target device. Auto-detected when ``None``.
quantization: Quantization policy. Defaults to ``fp8_cast``.
registry: Optional model registry for caching loaded components.
hdr_lora_config: Explicit HDR LoRA config override. When ``None``,
auto-detected from LoRA safetensors metadata.
tiled_vae_encode_pixel_threshold: Conditioning videos whose spatial
area (H x W) exceeds this value are encoded with the tiled
encoder. Default ``512 * 768`` is suitable for 80 GB GPUs.
Use ``256 * 256`` on GPUs with less VRAM.
offload_mode: Weight offloading strategy for diffusion stages.
"""
self.device = device or get_device()
self._tiled_vae_encode_threshold = tiled_vae_encode_pixel_threshold
if offload_mode != OffloadMode.NONE and quantization is not None:
logger.info("Offload mode enabled — disabling quantization (not supported with layer streaming).")
quantization = None
self.dtype = torch.bfloat16
lora_path = str(Path(hdr_lora).resolve())
loras = (LoraPathStrengthAndSDOps(lora_path, 1.0, LTXV_LORA_COMFY_RENAMING_MAP),)
# Load pre-computed text embeddings from safetensors.
emb_path = Path(text_embeddings_path)
logger.info("Loading text embeddings from %s", emb_path)
with safe_open(emb_path, framework="pt", device=str(self.device)) as f:
self.text_embeddings: tuple[torch.Tensor, torch.Tensor] = (
f.get_tensor("video_context"),
f.get_tensor("audio_context"),
)
self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
self.stage_1 = DiffusionStage(
distilled_checkpoint_path,
self.dtype,
self.device,
loras=loras,
quantization=quantization,
registry=registry,
offload_mode=offload_mode,
)
self.stage_2 = DiffusionStage(
distilled_checkpoint_path,
self.dtype,
self.device,
loras=loras,
quantization=quantization,
registry=registry,
offload_mode=offload_mode,
)
self.upsampler = VideoUpsampler(
distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
)
self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
# HDR config: explicit override, or auto-detect from LoRA metadata.
if hdr_lora_config is not None:
self._hdr_config: HdrLoraConfig | None = hdr_lora_config
else:
self._hdr_config = read_hdr_lora_config(lora_path)
if self._hdr_config is not None:
logger.info("[HDR IC-LoRA] HDR mode enabled (%s decode)", self._hdr_config.hdr_transform)
@property
def hdr_transform(self) -> str:
"""Active HDR transform name (defaults to 'logc3')."""
return self._hdr_config.hdr_transform if self._hdr_config is not None else "logc3"
@property
def reference_downscale_factor(self) -> int:
"""Reference video downscale factor from HDR LoRA config."""
return self._hdr_config.reference_downscale_factor if self._hdr_config is not None else 1
def __call__( # noqa: PLR0913
self,
seed: int,
height: int,
width: int,
num_frames: int,
frame_rate: float,
video_conditioning: list[tuple[str, float]],
tiling_config: TilingConfig | None = None,
high_quality_hdr: bool = False,
stage2_tilings: list[TileCountConfig] | None = None,
stage2_sigmas: list[list[float]] | None = None,
stage2_use_ic_lora: list[bool] | None = None,
) -> torch.Tensor:
"""Generate video with IC-LoRA conditioning and HDR output.
Returns a linear HDR float tensor ``[f, h, w, c]``.
Args:
seed: Random seed for reproducibility.
height: Desired output video height in pixels. Aligned internally
to the nearest multiple of 64 (rounded up). Decoded output is
cropped back to this size.
width: Desired output video width in pixels. Same alignment rules
as *height*.
num_frames: Number of frames to generate.
frame_rate: Output video frame rate.
video_conditioning: List of (path, strength) tuples for IC-LoRA video conditioning.
high_quality_hdr: High-quality HDR mode. Duplicates each conditioning
frame and generates at 2x frame count, then keeps every other
output frame. Reduces temporal artifacts at the cost of ~2x
generation time.
Returns:
Linear HDR float tensor ``[f, h, w, c]``.
"""
# In high-quality HDR mode, generate 2*N - 1 frames internally
# (satisfies (n-1)%8==0 when N itself does), then keep every other frame.
if high_quality_hdr:
gen_num_frames = 2 * num_frames - 1
logger.info("[HDR IC-LoRA] High-quality HDR: %d -> %d internal frames", num_frames, gen_num_frames)
else:
gen_num_frames = num_frames
gen_w, gen_h, crop_w, crop_h = align_resolution(
width, height, ResizeMode.REFLECT_PAD, divisor=ALIGNMENT_DIVISOR
)
if gen_h < MIN_RESOLUTION or gen_w < MIN_RESOLUTION:
raise ValueError(
f"Resolution ({width}x{height}) is too small after alignment "
f"(got {gen_w}x{gen_h}, need at least {MIN_RESOLUTION}x{MIN_RESOLUTION})."
)
needs_crop = crop_w != gen_w or crop_h != gen_h
if needs_crop:
logger.info(
"[HDR IC-LoRA] Aligned %dx%d -> %dx%d, will crop to %dx%d",
width,
height,
gen_w,
gen_h,
crop_w,
crop_h,
)
generator = torch.Generator(device=self.device).manual_seed(seed)
noiser = GaussianNoiser(generator=generator)
video_context, _ = self.text_embeddings
# Stage 1: Initial low resolution video generation.
s1_w, s1_h = gen_w // 2, gen_h // 2
stage_1_conditionings = self.image_conditioner(
lambda enc: self._create_conditionings(
video_conditioning=video_conditioning,
height=s1_h,
width=s1_w,
video_encoder=enc,
num_frames=gen_num_frames,
tiling_config=tiling_config,
high_quality_hdr=high_quality_hdr,
)
)
stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
# HDR is video-only: skip the audio stream to avoid denoising 5B audio params.
video_state, _ = self.stage_1(
denoiser=SimpleDenoiser(video_context, None),
sigmas=stage_1_sigmas,
noiser=noiser,
width=s1_w,
height=s1_h,
frames=gen_num_frames,
fps=frame_rate,
video=ModalitySpec(
context=video_context,
conditionings=stage_1_conditionings,
),
)
if stage2_tilings is None:
stage2_tilings = list(STAGE2_TILINGS)
if stage2_sigmas is None:
stage2_sigmas = [list(s) for s in STAGE2_SIGMAS]
if stage2_use_ic_lora is None:
stage2_use_ic_lora = list(STAGE2_USE_IC_LORA)
if not (len(stage2_tilings) == len(stage2_sigmas) == len(stage2_use_ic_lora)):
raise ValueError("stage2_tilings, stage2_sigmas, and stage2_use_ic_lora must have equal length")
# Stage 2: Upsample and refine at full resolution.
upscaled_video_latent = self.upsampler(video_state.latent[:1])
stage_2_conditionings = self.image_conditioner(
lambda enc: self._create_conditionings(
video_conditioning=video_conditioning,
height=gen_h,
width=gen_w,
video_encoder=enc,
num_frames=gen_num_frames,
tiling_config=tiling_config,
high_quality_hdr=high_quality_hdr,
)
)
with self.stage_2.model_context() as transformer:
phase_latent = upscaled_video_latent
for phase_idx, (tiling, sigmas_list, use_ic) in enumerate(
zip(stage2_tilings, stage2_sigmas, stage2_use_ic_lora, strict=True)
):
diffusion_tiling = _clamp_tile_to_latent(tiling, tuple(phase_latent.shape[2:5]))
conditionings = stage_2_conditionings if use_ic else []
sigma_t = torch.tensor(sigmas_list, dtype=torch.float32, device=self.device)
logger.info(
"[Stage 2 / phase %d] sigmas=%s ic_lora=%s tiling_h=%s tiling_w=%s",
phase_idx,
sigmas_list,
use_ic,
diffusion_tiling.height,
diffusion_tiling.width,
)
phase_latent = self._run_stage2_phase(
transformer=transformer,
latent=phase_latent,
conditionings=conditionings,
tiling=diffusion_tiling,
sigmas=sigma_t,
v_ctx=video_context,
frame_rate=frame_rate,
seed=seed,
)
final_video_latent = phase_latent
crop_size = (crop_w, crop_h) if needs_crop else None
return self._decode_video(
final_video_latent,
tiling_config,
generator,
crop_size,
high_quality_hdr=high_quality_hdr,
)
def _run_stage2_phase(
self,
transformer: object,
latent: torch.Tensor,
conditionings: list[ConditioningItem],
tiling: TileCountConfig,
sigmas: torch.Tensor,
v_ctx: torch.Tensor,
frame_rate: float,
seed: int,
) -> torch.Tensor:
"""Run one stage-2 denoising phase with optional IC-LoRA conditioning.
Each tile calls ``stage_2.run()`` with a tile-sized ``ModalitySpec`` for
video only (audio is omitted entirely for HDR). IC-LoRA conditionings
are sliced spatially to match each tile's extent.
"""
batch, n_channels, n_frames, n_height, n_width = latent.shape
full_shape = VideoLatentShape(batch=batch, channels=n_channels, frames=n_frames, height=n_height, width=n_width)
full_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), full_shape, frame_rate)
helper = VideoModalityTilingHelper(tiling, full_tools)
ref_initial = full_tools.create_initial_state(device=self.device, dtype=self.dtype)
ref_modality = modality_from_latent_state(ref_initial, v_ctx, sigmas[0])
n_gen = full_tools.target_shape.token_count()
blend_output = torch.zeros(batch, n_gen, n_channels, device=self.device, dtype=self.dtype)
patchifier = VideoLatentPatchifier(patch_size=1)
df = self.reference_downscale_factor
for tile_idx, tile in enumerate(helper.tiles):
_, ctx = helper.tile_modality(ref_modality, tile, normalize_positions=True)
frame_s, height_s, width_s = tile.in_coords
tile_h = height_s.stop - height_s.start
tile_w = width_s.stop - width_s.start
tile_f = frame_s.stop - frame_s.start
tile_conditionings = [
VideoConditionByReferenceLatent(
latent=cond.latent[
:,
:,
frame_s,
slice(height_s.start // df, height_s.stop // df),
slice(width_s.start // df, width_s.stop // df),
].to(device=self.device, dtype=self.dtype),
downscale_factor=cond.downscale_factor,
strength=cond.strength,
)
for cond in conditionings
]
tile_video_state, _ = self.stage_2.run(
transformer=transformer,
denoiser=SimpleDenoiser(v_ctx, None),
sigmas=sigmas,
noiser=GaussianNoiser(generator=torch.Generator(device=self.device).manual_seed(seed + tile_idx)),
width=tile_w * 32,
height=tile_h * 32,
frames=(tile_f - 1) * 8 + 1,
fps=frame_rate,
video=ModalitySpec(
context=v_ctx,
conditionings=tile_conditionings,
noise_scale=sigmas[0].item(),
initial_latent=latent[:, :, frame_s, height_s, width_s].to(device=self.device, dtype=self.dtype),
),
)
tile_tokens = patchifier.patchify(tile_video_state.latent)
blend_output = helper.blend(tile_tokens, tile, ctx, blend_output)
return full_tools.unpatchify(replace(ref_initial, latent=blend_output)).latent
def _decode_video(
self,
latent: torch.Tensor,
tiling_config: TilingConfig | None,
generator: torch.Generator,
crop_size: tuple[int, int] | None = None,
*,
high_quality_hdr: bool = False,
) -> torch.Tensor:
"""Decode latent to HDR video, optionally cropping to target size.
Args:
crop_size: ``(width, height)`` to crop decoded frames to, or
``None`` to skip cropping.
high_quality_hdr: When True, keep only every other frame (undoes the
2x generation applied during high-quality HDR mode).
Returns:
Linear HDR float tensor ``[f, h, w, c]``.
"""
# Cast to float32 so tiled-decode accumulation buffers and blending
# masks run in full precision, avoiding bfloat16 seam artifacts.
# Request float32 [0, 1] output — apply_hdr_decode_postprocess expects it.
latent = latent.float()
decoded = torch.cat(
list(self.video_decoder(latent, tiling_config, generator, output_dtype=torch.float32)),
dim=0,
)
decoded = rearrange(decoded, "f h w c -> 1 c f h w")
hdr = apply_hdr_decode_postprocess(decoded, transform=self.hdr_transform)
del decoded
out = rearrange(hdr[0], "c f h w -> f h w c")
if crop_size is not None:
out = out[:, : crop_size[1], : crop_size[0], :]
if high_quality_hdr:
out = out[::2]
return out
def _create_conditionings(
self,
video_conditioning: list[tuple[str, float]],
height: int,
width: int,
num_frames: int,
video_encoder: VideoEncoder,
tiling_config: TilingConfig | None = None,
high_quality_hdr: bool = False,
) -> list[ConditioningItem]:
"""Create conditioning items for video generation."""
conditionings: list[ConditioningItem] = []
scale = self.reference_downscale_factor
if scale != 1 and (height % scale != 0 or width % scale != 0):
raise ValueError(
f"Output dimensions ({height}x{width}) must be divisible by reference_downscale_factor ({scale})"
)
ref_height = height // scale
ref_width = width // scale
# In high-quality HDR mode, load half the frames then duplicate each one.
load_frame_cap = (num_frames + 1) // 2 if high_quality_hdr else num_frames
for video_path, strength in video_conditioning:
video = torch.cat(
list(
load_video_conditioning_hdr(
video_path=video_path,
height=ref_height,
width=ref_width,
frame_cap=load_frame_cap,
dtype=self.dtype,
device=self.device,
hdr_transform=self.hdr_transform,
resize_mode=ResizeMode.REFLECT_PAD,
)
),
dim=2,
)
if high_quality_hdr:
video = video.repeat_interleave(2, dim=2)[:, :, :num_frames, :, :]
if tiling_config is not None and ref_height * ref_width > self._tiled_vae_encode_threshold:
encoded_video = video_encoder.tiled_encode(video, tiling_config)
else:
encoded_video = video_encoder(video)
cond = VideoConditionByReferenceLatent(
latent=encoded_video,
downscale_factor=scale,
strength=strength,
)
conditionings.append(cond)
if video_conditioning:
logger.info("[HDR IC-LoRA] Added %d video conditioning(s)", len(video_conditioning))
return conditionings
# ---------------------------------------------------------------------------
# CLI helpers
# ---------------------------------------------------------------------------
def _make_tiling_config(
spatial_tile: int = DEFAULT_SPATIAL_TILE,
spatial_overlap: int = DEFAULT_SPATIAL_OVERLAP,
temporal_tile: int = DEFAULT_TEMPORAL_TILE,
temporal_overlap: int = DEFAULT_TEMPORAL_OVERLAP,
) -> TilingConfig:
"""Build a TilingConfig from explicit sizes.
The defaults (1280 px spatial tile, 256 px overlap; 32 temporal frames,
16 overlap) are suitable for H100-80 GB. On GPUs with less VRAM,
reduce the spatial tile size (e.g. ``spatial_tile=768``).
"""
from ltx_core.model.video_vae.tiling import SpatialTilingConfig, TemporalTilingConfig # noqa: PLC0415
return TilingConfig(
spatial_config=SpatialTilingConfig(tile_size_in_pixels=spatial_tile, tile_overlap_in_pixels=spatial_overlap),
temporal_config=TemporalTilingConfig(
tile_size_in_frames=temporal_tile,
tile_overlap_in_frames=temporal_overlap,
),
)
_VIDEO_SUFFIXES = {".mp4", ".mov"}
def _collect_videos(input_path: Path) -> list[Path]:
"""Return a list of .mp4/.mov files from *input_path* (file or directory)."""
if input_path.is_file():
return [input_path]
if input_path.is_dir():
return sorted(p for p in input_path.iterdir() if p.is_file() and p.suffix.lower() in _VIDEO_SUFFIXES)
logger.error("Input %s is not a file or directory", input_path)
return []
def _process_single_video( # noqa: PLR0913
pipeline: HDRICLoraPipeline,
video_path: Path,
vid_w: int,
vid_h: int,
num_frames: int,
frame_rate: float,
output_dir: Path,
tiling_config: TilingConfig,
seed: int,
skip_mp4: bool,
exr_half: bool,
exr_executor: "ThreadPoolExecutor", # noqa: F821
exr_futures: list,
high_quality_hdr: bool = False,
) -> None:
"""Run inference on a single video: generate EXR frames + optional H.264 .mp4 preview."""
import gc # noqa: PLC0415
import time # noqa: PLC0415
from ltx_pipelines.utils.media_io import encode_exr_sequence_to_mp4, save_exr_tensor # noqa: PLC0415
output_mp4 = output_dir / f"{video_path.stem}.mp4"
exr_dir = output_dir / f"{video_path.stem}_exr"
t0 = time.time()
hdr_video = pipeline(
seed=seed,
height=vid_h,
width=vid_w,
num_frames=num_frames,
frame_rate=frame_rate,
video_conditioning=[(str(video_path), 1.0)],
tiling_config=tiling_config,
high_quality_hdr=high_quality_hdr,
)
exr_dir.mkdir(parents=True, exist_ok=True)
for j in range(hdr_video.shape[0]):
frame_cpu = hdr_video[j].cpu().clone()
path = exr_dir / f"frame_{j:05d}.exr"
exr_futures.append(exr_executor.submit(save_exr_tensor, frame_cpu, str(path), exr_half))
del hdr_video
gc.collect()
torch.cuda.empty_cache()
if not skip_mp4:
# Wait for EXR saves to finish before encoding.
for fut in exr_futures:
fut.result()
logger.info("Encoding H.264 sRGB preview: %s", video_path.name)
encode_exr_sequence_to_mp4(exr_dir, output_mp4, frame_rate)
elapsed = time.time() - t0
logger.info("Decode + encode: %.1fs | %s", elapsed, output_mp4)
# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------
def _build_arg_parser() -> "argparse.ArgumentParser": # noqa: F821
"""Build the argument parser for HDR IC-LoRA batch inference."""
import argparse # noqa: PLC0415
parser = argparse.ArgumentParser(
description="HDR IC-LoRA inference: EXR frames + tonemapped ProRes .mov.",
epilog="""\
Resolution & frame constraints
------------------------------
* Width and height must each be divisible by 32.
* Frame count must satisfy (frames - 1) %% 8 == 0.
Valid counts: 1, 9, 17, 25, ..., 121, 129, 137, 145, 153, 161.
Max frames by resolution (fp8_cast, bfloat16 VAE, tiled decode)
---------------------------------------------------------------
Resolution 80 GB (H100) 48 GB (A6000)
------------------------------------------------
720p 1280x720 161+ frames 161+ frames
1080p 1920x1080 161+ frames 161+ frames
2K 2048x1080 161+ frames 161+ frames
1440p 2560x1440 161+ frames 137 frames
4K 3840x2160 121 frames 49 frames
4K 4096x2160 105 frames 49 frames
Estimates from ltx_pipelines.utils.vram_budget. Run
python -c "from ltx_pipelines.utils.vram_budget import \\
max_frames_for_resolution as mf; print(mf(W, H, vram_gb=GB))"
to check your specific resolution and GPU.
* The tiled-encode threshold (%(tiled_threshold)s px) and the default
tiling config (%(stile)s px spatial tile) are tuned for 80 GB.
On lower-VRAM GPUs pass --spatial-tile 768 (or smaller).
"""
% {"tiled_threshold": TILED_VAE_ENCODE_PIXEL_THRESHOLD, "stile": DEFAULT_SPATIAL_TILE},
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument("--input", required=True, help="Single .mp4 or directory of .mp4 videos.")
parser.add_argument("--output-dir", required=True, help="Directory for .mov and EXR folders.")
parser.add_argument("--hdr-lora", required=True, help="HDR IC-LoRA .safetensors file.")
parser.add_argument("--text-embeddings", required=True, help="Pre-computed text embeddings (.safetensors file).")
parser.add_argument("--distilled-checkpoint-path", required=True, help="Distilled model checkpoint (.safetensors).")
parser.add_argument("--spatial-upsampler-path", required=True, help="Spatial upsampler (.safetensors).")
parser.add_argument(
"--num-frames",
type=int,
default=DEFAULT_NUM_FRAMES,
help=f"Number of output frames. Must satisfy (n-1) %% 8 == 0 (default: {DEFAULT_NUM_FRAMES}).",
)
parser.add_argument(
"--spatial-tile",
type=int,
default=DEFAULT_SPATIAL_TILE,
help=f"Spatial tile size in pixels for tiled VAE decode (default: {DEFAULT_SPATIAL_TILE}). "
"Reduce on lower-VRAM GPUs (e.g. 768 for 48 GB).",
)
parser.add_argument("--skip-mp4", action="store_true", help="Skip H.264 MP4 encoding, only produce EXR.")
parser.add_argument("--exr-half", action="store_true", help="Save EXR as float16.")
parser.add_argument("--seed", type=int, default=10, help="Random seed (default: 10).")
parser.add_argument(
"--offload",
dest="offload_mode",
type=OffloadMode,
default=OffloadMode.NONE,
choices=list(OffloadMode),
help=(
"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"
),
)
parser.add_argument(
"--high-quality",
action="store_true",
help="High-quality HDR mode. Generates at 2x frame count internally "
"and keeps every other frame for smoother output. ~2x slower.",
)
return parser
@torch.inference_mode()
def main() -> None:
"""Batch HDR IC-LoRA inference: per-frame EXR + tonemapped ProRes .mov."""
import time # noqa: PLC0415
from concurrent.futures import ThreadPoolExecutor # noqa: PLC0415
from ltx_pipelines.utils.media_io import get_videostream_metadata # noqa: PLC0415
logging.basicConfig(level=logging.INFO)
args = _build_arg_parser().parse_args()
high_quality = args.high_quality
num_frames = args.num_frames
tiling_config = _make_tiling_config(spatial_tile=args.spatial_tile)
input_path = Path(args.input)
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
videos = _collect_videos(input_path)
if not videos:
logger.error("No valid videos to process.")
return
logger.info("Found %d video(s), generating %d frames each", len(videos), num_frames)
logger.info("Loading pipeline...")
pipeline = HDRICLoraPipeline(
distilled_checkpoint_path=args.distilled_checkpoint_path,
spatial_upsampler_path=args.spatial_upsampler_path,
hdr_lora=args.hdr_lora,
text_embeddings_path=args.text_embeddings,
offload_mode=args.offload_mode,
)
logger.info("Pipeline loaded.")
exr_executor = ThreadPoolExecutor(max_workers=4)
exr_futures: list = []
total_t0 = time.time()
successes = 0
for i, video_path in enumerate(videos, 1):
meta = get_videostream_metadata(str(video_path))
vid_w, vid_h = meta.width, meta.height
logger.info("%s", "=" * 60)
logger.info("[%d/%d] %s (%dx%d, %df)", i, len(videos), video_path.name, vid_w, vid_h, num_frames)
_process_single_video(
pipeline=pipeline,
video_path=video_path,
vid_w=vid_w,
vid_h=vid_h,
num_frames=num_frames,
frame_rate=meta.fps,
output_dir=output_dir,
tiling_config=tiling_config,
seed=args.seed,
skip_mp4=args.skip_mp4,
exr_half=args.exr_half,
exr_executor=exr_executor,
exr_futures=exr_futures,
high_quality_hdr=high_quality,
)
successes += 1
infer_elapsed = time.time() - total_t0
logger.info("%s", "=" * 60)
logger.info("All inference done in %.0fs (%d/%d OK)", infer_elapsed, successes, len(videos))
if exr_futures:
t0 = time.time()
logger.info("Waiting for %d EXR saves...", len(exr_futures))
for fut in exr_futures:
fut.result()
exr_wait = time.time() - t0
if exr_wait > 0.1:
logger.info("EXR save wait: %.1fs", exr_wait)
logger.info("Total wall time: %.0fs", time.time() - total_t0)
if __name__ == "__main__":
main()
packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py
+11 -7
View File
@@ -39,7 +39,7 @@ from ltx_pipelines.utils.constants import (
from ltx_pipelines.utils.denoisers import SimpleDenoiser
from ltx_pipelines.utils.helpers import assert_resolution, combined_image_conditionings, get_device
from ltx_pipelines.utils.media_io import decode_video_by_frame, encode_video, video_preprocess
from ltx_pipelines.utils.types import ModalitySpec
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
class ICLoraPipeline:
@@ -63,12 +63,18 @@ class ICLoraPipeline:
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
self.dtype = torch.bfloat16
self.prompt_encoder = PromptEncoder(
distilled_checkpoint_path, gemma_root, self.dtype, self.device, registry=registry
distilled_checkpoint_path,
gemma_root,
self.dtype,
self.device,
registry=registry,
offload_mode=offload_mode,
)
self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
self.stage_1 = DiffusionStage(
@@ -79,6 +85,7 @@ class ICLoraPipeline:
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
offload_mode=offload_mode,
)
self.stage_2 = DiffusionStage(
distilled_checkpoint_path,
@@ -88,6 +95,7 @@ class ICLoraPipeline:
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
offload_mode=offload_mode,
)
self.upsampler = VideoUpsampler(
distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -125,7 +133,6 @@ class ICLoraPipeline:
conditioning_attention_strength: float = 1.0,
skip_stage_2: bool = False,
conditioning_attention_mask: torch.Tensor | None = None,
streaming_prefetch_count: int | None = None,
stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
) -> tuple[Iterator[torch.Tensor], Audio]:
@@ -175,7 +182,6 @@ class ICLoraPipeline:
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
enhance_prompt_seed=seed,
streaming_prefetch_count=streaming_prefetch_count,
)
video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
@@ -219,7 +225,6 @@ class ICLoraPipeline:
audio=ModalitySpec(
context=audio_context,
),
streaming_prefetch_count=streaming_prefetch_count,
)
if skip_stage_2:
@@ -264,7 +269,6 @@ class ICLoraPipeline:
noise_scale=stage_2_sigmas[0].item(),
initial_latent=audio_state.latent,
),
streaming_prefetch_count=streaming_prefetch_count,
)
decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -463,6 +467,7 @@ def main() -> None:
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -479,7 +484,6 @@ def main() -> None:
conditioning_attention_strength=conditioning_attention_strength,
skip_stage_2=args.skip_stage_2,
conditioning_attention_mask=conditioning_attention_mask,
streaming_prefetch_count=args.streaming_prefetch_count,
)
encode_video(
packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py
+13 -8
View File
@@ -15,7 +15,11 @@ from ltx_core.loader.registry import Registry
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, VideoPixelShape
from ltx_pipelines.utils.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path
from ltx_pipelines.utils.args import (
ImageConditioningInput,
default_2_stage_arg_parser,
detect_checkpoint_path,
)
from ltx_pipelines.utils.blocks import (
AudioDecoder,
DiffusionStage,
@@ -35,7 +39,7 @@ from ltx_pipelines.utils.helpers import (
image_conditionings_by_adding_guiding_latent,
)
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import ModalitySpec
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
class KeyframeInterpolationPipeline:
@@ -59,12 +63,15 @@ class KeyframeInterpolationPipeline:
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
self.dtype = torch.bfloat16
self._scheduler = LTX2Scheduler()
self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
self.prompt_encoder = PromptEncoder(
checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
)
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
self.stage_1 = DiffusionStage(
checkpoint_path,
@@ -74,6 +81,7 @@ class KeyframeInterpolationPipeline:
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
offload_mode=offload_mode,
)
stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
self.stage_2 = DiffusionStage(
@@ -84,6 +92,7 @@ class KeyframeInterpolationPipeline:
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
offload_mode=offload_mode,
)
self.upsampler = VideoUpsampler(
checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -106,7 +115,6 @@ class KeyframeInterpolationPipeline:
images: list[ImageConditioningInput],
tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
max_batch_size: int = 1,
stage_1_sigmas: torch.Tensor | None = None,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
@@ -122,7 +130,6 @@ class KeyframeInterpolationPipeline:
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
enhance_prompt_seed=seed,
streaming_prefetch_count=streaming_prefetch_count,
)
v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
@@ -179,7 +186,6 @@ class KeyframeInterpolationPipeline:
audio=ModalitySpec(
context=a_context_p,
),
streaming_prefetch_count=streaming_prefetch_count,
max_batch_size=max_batch_size,
)
@@ -218,7 +224,6 @@ class KeyframeInterpolationPipeline:
noise_scale=stage_2_sigmas[0].item(),
initial_latent=audio_state.latent,
),
streaming_prefetch_count=streaming_prefetch_count,
)
decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -241,6 +246,7 @@ def main() -> None:
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -271,7 +277,6 @@ def main() -> None:
),
images=args.images,
tiling_config=tiling_config,
streaming_prefetch_count=args.streaming_prefetch_count,
max_batch_size=args.max_batch_size,
)
packages/ltx-pipelines/src/ltx_pipelines/retake.py
+5 -5
View File
@@ -36,7 +36,7 @@ from ltx_pipelines.utils.media_io import (
encode_video,
get_videostream_metadata,
)
from ltx_pipelines.utils.types import ModalitySpec
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
class RetakePipeline:
@@ -74,6 +74,7 @@ class RetakePipeline:
registry: Registry | None = None,
distilled: bool = True,
torch_compile: bool = False,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
self.dtype = torch.bfloat16
@@ -86,6 +87,7 @@ class RetakePipeline:
dtype=self.dtype,
device=self.device,
registry=registry,
offload_mode=offload_mode,
)
self.image_conditioner = ImageConditioner(
checkpoint_path=checkpoint_path,
@@ -107,6 +109,7 @@ class RetakePipeline:
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
offload_mode=offload_mode,
)
self.video_decoder = VideoDecoder(
checkpoint_path=checkpoint_path,
@@ -141,7 +144,6 @@ class RetakePipeline:
regenerate_audio: bool = True,
enhance_prompt: bool = False,
tiling_config: TilingConfig | None = None,
streaming_prefetch_count: int | None = None,
max_batch_size: int = 1,
sigmas: torch.Tensor | None = None,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
@@ -210,7 +212,6 @@ class RetakePipeline:
prompts_to_encode,
enhance_first_prompt=enhance_prompt,
enhance_prompt_seed=seed,
streaming_prefetch_count=streaming_prefetch_count,
)
v_context_p, a_context_p = contexts[0].video_encoding, contexts[0].audio_encoding
@@ -269,7 +270,6 @@ class RetakePipeline:
fps=output_shape.fps,
video=video_modality_spec,
audio=audio_modality_spec,
streaming_prefetch_count=streaming_prefetch_count,
max_batch_size=max_batch_size,
)
@@ -309,6 +309,7 @@ def main() -> None:
quantization=args.quantization,
distilled=args.distilled,
torch_compile=args.compile,
offload_mode=args.offload_mode,
)
params = detect_params(args.distilled_checkpoint_path)
tiling_config = TilingConfig.default()
@@ -321,7 +322,6 @@ def main() -> None:
video_guider_params=params.video_guider_params,
audio_guider_params=params.audio_guider_params,
tiling_config=tiling_config,
streaming_prefetch_count=args.streaming_prefetch_count,
max_batch_size=args.max_batch_size,
)
video_chunks_number = get_video_chunks_number(src.frames, tiling_config)
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py
+10 -6
View File
@@ -20,7 +20,11 @@ from ltx_pipelines.utils import (
combined_image_conditionings,
get_device,
)
from ltx_pipelines.utils.args import ImageConditioningInput, default_1_stage_arg_parser, detect_checkpoint_path
from ltx_pipelines.utils.args import (
ImageConditioningInput,
default_1_stage_arg_parser,
detect_checkpoint_path,
)
from ltx_pipelines.utils.blocks import (
AudioDecoder,
DiffusionStage,
@@ -31,7 +35,7 @@ from ltx_pipelines.utils.blocks import (
from ltx_pipelines.utils.constants import detect_params
from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import ModalitySpec
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
class TI2VidOneStagePipeline:
@@ -52,6 +56,7 @@ class TI2VidOneStagePipeline:
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.dtype = torch.bfloat16
self.device = device or get_device()
@@ -62,6 +67,7 @@ class TI2VidOneStagePipeline:
dtype=self.dtype,
device=self.device,
registry=registry,
offload_mode=offload_mode,
)
self.image_conditioner = ImageConditioner(
checkpoint_path=checkpoint_path,
@@ -77,6 +83,7 @@ class TI2VidOneStagePipeline:
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
offload_mode=offload_mode,
)
self.video_decoder = VideoDecoder(
checkpoint_path=checkpoint_path,
@@ -105,7 +112,6 @@ class TI2VidOneStagePipeline:
audio_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
images: list[ImageConditioningInput],
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
tiling_config: TilingConfig | None = None,
max_batch_size: int = 1,
sigmas: torch.Tensor | None = None,
@@ -121,7 +127,6 @@ class TI2VidOneStagePipeline:
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
enhance_prompt_seed=seed,
streaming_prefetch_count=streaming_prefetch_count,
)
v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
@@ -170,7 +175,6 @@ class TI2VidOneStagePipeline:
audio=ModalitySpec(
context=a_context_p,
),
streaming_prefetch_count=streaming_prefetch_count,
max_batch_size=max_batch_size,
)
@@ -192,6 +196,7 @@ def main() -> None:
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
offload_mode=args.offload_mode,
)
video, audio = pipeline(
prompt=args.prompt,
@@ -219,7 +224,6 @@ def main() -> None:
stg_blocks=args.audio_stg_blocks,
),
images=args.images,
streaming_prefetch_count=args.streaming_prefetch_count,
max_batch_size=args.max_batch_size,
)
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py
+13 -8
View File
@@ -15,7 +15,11 @@ from ltx_core.loader.registry import Registry
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, VideoPixelShape
from ltx_pipelines.utils.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path
from ltx_pipelines.utils.args import (
ImageConditioningInput,
default_2_stage_arg_parser,
detect_checkpoint_path,
)
from ltx_pipelines.utils.blocks import (
AudioDecoder,
DiffusionStage,
@@ -35,7 +39,7 @@ from ltx_pipelines.utils.helpers import (
get_device,
)
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import ModalitySpec
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
class TI2VidTwoStagesPipeline:
@@ -58,12 +62,15 @@ class TI2VidTwoStagesPipeline:
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
self.dtype = torch.bfloat16
self._scheduler = LTX2Scheduler()
self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
self.prompt_encoder = PromptEncoder(
checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
)
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
self.upsampler = VideoUpsampler(
checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -79,6 +86,7 @@ class TI2VidTwoStagesPipeline:
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
offload_mode=offload_mode,
)
self.stage_2 = DiffusionStage(
checkpoint_path,
@@ -88,6 +96,7 @@ class TI2VidTwoStagesPipeline:
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
offload_mode=offload_mode,
)
def __call__( # noqa: PLR0913
@@ -105,7 +114,6 @@ class TI2VidTwoStagesPipeline:
images: list[ImageConditioningInput],
tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
max_batch_size: int = 1,
stage_1_sigmas: torch.Tensor | None = None,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
@@ -121,7 +129,6 @@ class TI2VidTwoStagesPipeline:
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
enhance_prompt_seed=seed,
streaming_prefetch_count=streaming_prefetch_count,
)
v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
@@ -170,7 +177,6 @@ class TI2VidTwoStagesPipeline:
fps=frame_rate,
video=ModalitySpec(context=v_context_p, conditionings=stage_1_conditionings),
audio=ModalitySpec(context=a_context_p),
streaming_prefetch_count=streaming_prefetch_count,
max_batch_size=max_batch_size,
)
@@ -208,7 +214,6 @@ class TI2VidTwoStagesPipeline:
noise_scale=stage_2_sigmas[0].item(),
initial_latent=audio_state.latent,
),
streaming_prefetch_count=streaming_prefetch_count,
)
decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -231,6 +236,7 @@ def main() -> None:
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -261,7 +267,6 @@ def main() -> None:
),
images=args.images,
tiling_config=tiling_config,
streaming_prefetch_count=args.streaming_prefetch_count,
max_batch_size=args.max_batch_size,
)
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_hq.py
+8 -7
View File
@@ -33,7 +33,7 @@ from ltx_pipelines.utils.helpers import (
)
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.samplers import res2s_audio_video_denoising_loop
from ltx_pipelines.utils.types import ModalitySpec
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
class TI2VidTwoStagesHQPipeline:
@@ -61,6 +61,7 @@ class TI2VidTwoStagesHQPipeline:
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
self.dtype = torch.bfloat16
@@ -77,7 +78,9 @@ class TI2VidTwoStagesHQPipeline:
sd_ops=distilled_lora[0].sd_ops,
)
self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
self.prompt_encoder = PromptEncoder(
checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
)
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
self.upsampler = VideoUpsampler(
checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -93,6 +96,7 @@ class TI2VidTwoStagesHQPipeline:
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
offload_mode=offload_mode,
)
self.stage_2 = DiffusionStage(
checkpoint_path,
@@ -102,6 +106,7 @@ class TI2VidTwoStagesHQPipeline:
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
offload_mode=offload_mode,
)
@torch.inference_mode()
@@ -120,7 +125,6 @@ class TI2VidTwoStagesHQPipeline:
images: list[ImageConditioningInput],
tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
max_batch_size: int = 1,
stage_1_sigmas: torch.Tensor | None = None,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
@@ -136,7 +140,6 @@ class TI2VidTwoStagesHQPipeline:
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
enhance_prompt_seed=seed,
streaming_prefetch_count=streaming_prefetch_count,
)
v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
@@ -190,7 +193,6 @@ class TI2VidTwoStagesHQPipeline:
video=ModalitySpec(context=v_context_p, conditionings=stage_1_conditionings),
audio=ModalitySpec(context=a_context_p),
loop=res2s_audio_video_denoising_loop,
streaming_prefetch_count=streaming_prefetch_count,
max_batch_size=max_batch_size,
)
@@ -231,7 +233,6 @@ class TI2VidTwoStagesHQPipeline:
initial_latent=audio_state.latent,
),
loop=res2s_audio_video_denoising_loop,
streaming_prefetch_count=streaming_prefetch_count,
)
decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -254,6 +255,7 @@ def main() -> None:
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -284,7 +286,6 @@ def main() -> None:
),
images=args.images,
tiling_config=tiling_config,
streaming_prefetch_count=args.streaming_prefetch_count,
max_batch_size=args.max_batch_size,
)
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"