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

Automated PR - 2026-05-11

Browse Source
This commit is contained in:
github-actions[bot]
2026-05-11 13:14:05 +00:00
parent 41d9243716
commit 7df34dfa83
72 changed files with 3299 additions and 911 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-pipelines/src/ltx_pipelines/utils/__init__.py
+4 -1
View File
@@ -16,15 +16,17 @@ from ltx_pipelines.utils.helpers import (
image_conditionings_by_adding_guiding_latent,
)
from ltx_pipelines.utils.samplers import (
euler_cfg_pp_denoising_loop,
euler_denoising_loop,
gradient_estimating_euler_denoising_loop,
res2s_audio_video_denoising_loop,
)
from ltx_pipelines.utils.types import Denoiser, ModalitySpec
from ltx_pipelines.utils.types import DenoisedLatentResult, Denoiser, ModalitySpec
__all__ = [
"AudioConditioner",
"AudioDecoder",
"DenoisedLatentResult",
"Denoiser",
"DiffusionStage",
"FactoryGuidedDenoiser",
@@ -38,6 +40,7 @@ __all__ = [
"assert_resolution",
"cleanup_memory",
"combined_image_conditionings",
"euler_cfg_pp_denoising_loop",
"euler_denoising_loop",
"get_device",
"gradient_estimating_euler_denoising_loop",
packages/ltx-pipelines/src/ltx_pipelines/utils/args.py
+78 -33
View File
@@ -1,4 +1,5 @@
import argparse
from collections.abc import Sequence
from pathlib import Path
from typing import NamedTuple
@@ -115,35 +116,34 @@ def resolve_path(path: str) -> str:
QUANTIZATION_POLICIES = ("fp8-cast", "fp8-scaled-mm")
class QuantizationAction(argparse.Action):
def __call__(
self,
parser: argparse.ArgumentParser, # noqa: ARG002
namespace: argparse.Namespace,
values: list[str],
option_string: str | None = None,
) -> None:
if len(values) > 2:
msg = (
f"{option_string} accepts at most 2 arguments (POLICY and optional AMAX_PATH), got {len(values)} values"
def _resolve_quantization(namespace: argparse.Namespace) -> None:
# Resolution is deferred until after parse_args because fp8-scaled-mm needs the
# checkpoint path, which isn't on the namespace when the --quantization argument
# is parsed.
name = getattr(namespace, "quantization", None)
if name is None or isinstance(name, QuantizationPolicy):
return
if name == "fp8-cast":
namespace.quantization = QuantizationPolicy.fp8_cast()
return
if name == "fp8-scaled-mm":
ckpt = getattr(namespace, "checkpoint_path", None) or getattr(namespace, "distilled_checkpoint_path", None)
if ckpt is None:
raise SystemExit(
"--quantization fp8-scaled-mm requires --checkpoint-path (or --distilled-checkpoint-path)."
)
raise argparse.ArgumentError(self, msg)
namespace.quantization = QuantizationPolicy.fp8_scaled_mm(ckpt)
policy_name = values[0]
if policy_name not in QUANTIZATION_POLICIES:
msg = f"Unknown quantization policy '{policy_name}'. Choose from: {', '.join(QUANTIZATION_POLICIES)}"
raise argparse.ArgumentError(self, msg)
if policy_name == "fp8-cast":
if len(values) > 1:
msg = f"{option_string} fp8-cast does not accept additional arguments"
raise argparse.ArgumentError(self, msg)
policy = QuantizationPolicy.fp8_cast()
elif policy_name == "fp8-scaled-mm":
amax_path = resolve_path(values[1]) if len(values) > 1 else None
policy = QuantizationPolicy.fp8_scaled_mm(amax_path)
setattr(namespace, self.dest, policy)
class _PipelineArgumentParser(argparse.ArgumentParser):
def parse_args( # type: ignore[override]
self,
args: Sequence[str] | None = None,
namespace: argparse.Namespace | None = None,
) -> argparse.Namespace:
ns = super().parse_args(args, namespace)
_resolve_quantization(ns)
return ns
def detect_checkpoint_path(distilled: bool = False) -> str:
@@ -159,7 +159,7 @@ def basic_arg_parser(
params: PipelineParams = LTX_2_3_PARAMS,
distilled: bool = False,
) -> argparse.ArgumentParser:
parser = argparse.ArgumentParser()
parser = _PipelineArgumentParser()
if distilled:
parser.add_argument(
"--distilled-checkpoint-path",
@@ -264,16 +264,14 @@ def basic_arg_parser(
parser.add_argument(
"--quantization",
dest="quantization",
action=QuantizationAction,
nargs="+",
metavar=("POLICY", "AMAX_PATH"),
choices=QUANTIZATION_POLICIES,
default=None,
help=(
f"Quantization policy: {', '.join(QUANTIZATION_POLICIES)}. "
"fp8-cast uses FP8 casting with upcasting during inference. "
"fp8-scaled-mm uses FP8 scaled matrix multiplication (optionally provide amax calibration file path). "
"Example: --quantization fp8-cast or --quantization fp8-scaled-mm /path/to/amax.json"
"fp8-scaled-mm uses FP8 scaled matrix multiplication; the layer set is auto-discovered "
"from the checkpoint's .weight_scale tensors. "
"Example: --quantization fp8-cast or --quantization fp8-scaled-mm"
),
)
parser.add_argument(
@@ -348,6 +346,53 @@ def video_editing_arg_parser(
return parser
def lipdub_arg_parser(
params: PipelineParams = LTX_2_3_PARAMS,
) -> argparse.ArgumentParser:
"""Argument parser for the lip-dub pipeline.
Frame count and frame rate are derived from the reference video at runtime (the frame count
is silently snapped down to the nearest 8k+1), so this parser intentionally omits
--num-frames, --frame-rate, and --image. Distilled checkpoint only.
"""
parser = basic_arg_parser(params=params, distilled=True)
parser.add_argument(
"--height",
type=int,
default=params.stage_2_height,
help=(
f"Height of the generated video in pixels, should be divisible by 64 (default: {params.stage_2_height})."
),
)
parser.add_argument(
"--width",
type=int,
default=params.stage_2_width,
help=f"Width of the generated video in pixels, should be divisible by 64 (default: {params.stage_2_width}).",
)
parser.add_argument(
"--spatial-upsampler-path",
type=resolve_path,
required=True,
help=(
"Path to the spatial upsampler model used to increase the resolution "
"of the generated video in the latent space."
),
)
parser.add_argument(
"--reference-video",
type=resolve_path,
required=True,
help="Reference video file (video + audio track used for IC-LoRA and audio identity).",
)
parser.add_argument(
"--reference-strength",
type=float,
default=1.0,
help="Strength for IC-LoRA video reference conditioning (default: 1.0).",
)
return parser
def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
video_guider = params.video_guider_params
audio_guider = params.audio_guider_params
packages/ltx-pipelines/src/ltx_pipelines/utils/blocks.py
+104 -64
View File
@@ -21,7 +21,8 @@ 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.loader import SDOps
from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.primitives import BuilderProtocol, LoraPathStrengthAndSDOps, ModelBuilderProtocol
from ltx_core.loader.registry import DummyRegistry, Registry
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
from ltx_core.model.audio_vae import (
@@ -37,12 +38,14 @@ from ltx_core.model.audio_vae import (
)
from ltx_core.model.transformer import (
LTXV_MODEL_COMFY_RENAMING_MAP,
LTXModel,
LTXModelConfigurator,
X0Model,
)
from ltx_core.model.transformer.compiling import COMPILE_TRANSFORMER, modify_sd_ops_for_compilation
from ltx_core.model.upsampler import LatentUpsamplerConfigurator, upsample_video
from ltx_core.model.video_vae import (
MEMORY_EFFICIENT_DECODE,
VAE_DECODER_COMFY_KEYS_FILTER,
VAE_ENCODER_COMFY_KEYS_FILTER,
TilingConfig,
@@ -59,10 +62,11 @@ from ltx_core.text_encoders.gemma import (
GemmaTextEncoderConfigurator,
module_ops_from_gemma_root,
)
from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessorOutput
from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessor, EmbeddingsProcessorOutput
from ltx_core.tools import AudioLatentTools, LatentTools, VideoLatentTools
from ltx_core.types import Audio, AudioLatentShape, LatentState, VideoLatentShape, VideoPixelShape
from ltx_core.utils import find_matching_file
from ltx_pipelines.multigpu.delegating_builder import DelegatingBuilder
from ltx_pipelines.utils.gpu_model import gpu_model
from ltx_pipelines.utils.helpers import (
cleanup_memory,
@@ -83,6 +87,20 @@ _M = TypeVar("_M", bound=torch.nn.Module)
# ---------------------------------------------------------------------------
def _chain_quantization(
sd_ops: SDOps,
module_ops: tuple[ModuleOps, ...],
quantization: QuantizationPolicy,
) -> tuple[SDOps, tuple[ModuleOps, ...]]:
chained_sd_ops = sd_ops
if quantization.sd_ops is not None:
chained_sd_ops = SDOps(
name=f"sd_ops_chain_{sd_ops.name}+{quantization.sd_ops.name}",
mapping=(*sd_ops.mapping, *quantization.sd_ops.mapping),
)
return chained_sd_ops, (*module_ops, *quantization.module_ops)
@contextmanager
def _streaming_model(
builder: StreamingModelBuilder,
@@ -154,16 +172,43 @@ class DiffusionStage:
registry: Registry | None = None,
torch_compile: bool = False,
offload_mode: OffloadMode = OffloadMode.NONE,
transformer_builder: ModelBuilderProtocol[LTXModel] | DelegatingBuilder[LTXModel] | None = None,
) -> None:
self._dtype = dtype
self._device = device
self._quantization = quantization
self._torch_compile = torch_compile
self._offload_mode = offload_mode
if transformer_builder is not None:
self._transformer_builder = transformer_builder
else:
self._transformer_builder = Builder(
model_path=checkpoint_path,
model_class_configurator=LTXModelConfigurator,
model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
loras=tuple(loras),
registry=registry or DummyRegistry(),
)
if offload_mode != OffloadMode.NONE:
if torch_compile:
raise ValueError("torch.compile is not supported with layer streaming")
streaming_sd_ops: SDOps = LTXV_MODEL_COMFY_RENAMING_MAP
streaming_module_ops: tuple[ModuleOps, ...] = ()
if quantization is not None:
raise ValueError("quantization is not supported with layer streaming")
if quantization.kind != QuantizationPolicy.Kind.FP8_CAST:
raise ValueError(
f"Layer streaming supports only QuantizationPolicy.fp8_cast(); "
f"got kind={quantization.kind!r} which produces heterogeneous block layouts."
)
streaming_sd_ops, streaming_module_ops = _chain_quantization(
streaming_sd_ops, streaming_module_ops, quantization
)
self._streaming_builder = StreamingModelBuilder(
model_class_configurator=LTXModelConfigurator,
model_path=checkpoint_path,
model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
model_sd_ops=streaming_sd_ops,
module_ops=streaming_module_ops,
loras=tuple(loras),
registry=registry or DummyRegistry(),
blocks_attr="velocity_model.transformer_blocks",
@@ -172,19 +217,6 @@ class DiffusionStage:
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,
model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
loras=tuple(loras),
registry=registry or DummyRegistry(),
)
def _build_transformer(self, *, device: torch.device | None = None, **kwargs: object) -> X0Model:
target = device or self._device
sd_ops = self._transformer_builder.model_sd_ops
@@ -198,18 +230,12 @@ class DiffusionStage:
LoraPathStrengthAndSDOps(
lora.path,
lora.strength,
modify_sd_ops_for_compilation(
lora.sd_ops if lora.sd_ops is not None else SDOps(name="identity"), number_of_layers
),
modify_sd_ops_for_compilation(lora.sd_ops, number_of_layers),
)
for lora in loras
)
if self._quantization is not None:
module_ops = (*module_ops, *self._quantization.module_ops)
sd_ops = SDOps(
name=f"sd_ops_chain_{sd_ops.name}+{self._quantization.sd_ops.name}",
mapping=(*sd_ops.mapping, *self._quantization.sd_ops.mapping),
)
sd_ops, module_ops = _chain_quantization(sd_ops, module_ops, self._quantization)
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()
@@ -359,31 +385,40 @@ class PromptEncoder:
device: torch.device,
registry: Registry | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
text_encoder_builder: BuilderProtocol | None = 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
weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")]
self._text_encoder_builder = Builder(
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(),
)
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",
)
if text_encoder_builder is not None:
if offload_mode != OffloadMode.NONE:
raise ValueError(
"text_encoder_builder cannot be used with offload_mode != OffloadMode.NONE "
"because no streaming text encoder builder is available."
)
self._text_encoder_builder = text_encoder_builder
self._streaming_text_encoder_builder = None
else:
module_ops = module_ops_from_gemma_root(gemma_root)
model_folder = find_matching_file(gemma_root, "model*.safetensors").parent
weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")]
self._text_encoder_builder = Builder(
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(),
)
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,
@@ -391,10 +426,18 @@ class PromptEncoder:
registry=registry or DummyRegistry(),
)
def _build_text_encoder(self) -> torch.nn.Module:
"""Build the Gemma text encoder (non-streaming path)."""
return self._text_encoder_builder.build(device=self._device, dtype=self._dtype).eval()
def _build_embeddings_processor(self) -> EmbeddingsProcessor:
"""Build the embeddings processor on the target device."""
return self._embeddings_processor_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
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())
return gpu_model(self._build_text_encoder())
def __call__(
self,
@@ -413,9 +456,7 @@ class PromptEncoder:
)
raw_outputs = [text_encoder.encode(p) for p in prompts]
with gpu_model(
self._embeddings_processor_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
) as embeddings_processor:
with gpu_model(self._build_embeddings_processor()) as embeddings_processor:
return [embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs]
@@ -513,32 +554,31 @@ class VideoDecoder:
dtype: torch.dtype,
device: torch.device,
registry: Registry | None = None,
memory_efficient: bool = True,
decoder_builder: BuilderProtocol | None = None,
) -> None:
self._dtype = dtype
self._device = device
self._decoder_builder = Builder(
model_path=checkpoint_path,
model_class_configurator=VideoDecoderConfigurator,
model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
registry=registry or DummyRegistry(),
)
if decoder_builder is not None:
self._decoder_builder = decoder_builder
else:
self._decoder_builder = Builder(
model_path=checkpoint_path,
model_class_configurator=VideoDecoderConfigurator,
model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
registry=registry or DummyRegistry(),
module_ops=(MEMORY_EFFICIENT_DECODE,) if memory_efficient else (),
)
def __call__(
self,
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.
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.
"""
"""Decode *latent* to pixel-space video chunks. Decoder freed after exhaustion."""
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, output_dtype=output_dtype), decoder)
return _cleanup_iter(decoder.decode_video(latent, tiling_config, generator), decoder)
# ---------------------------------------------------------------------------
packages/ltx-pipelines/src/ltx_pipelines/utils/color_conversion.py
+224
View File
@@ -0,0 +1,224 @@
"""Color space conversion utilities for video encoding.
Provides GPU-accelerated RGB to YUV420 conversion that runs between the
VAE decoder (which yields float RGB chunks) and ``encode_video``, bypassing
pyav's CPU-side libswscale conversion. The ``FrameConverter`` also carries
the codec metadata (pixel format, colour space, colour range) that
``encode_video`` needs to tag the output stream.
"""
from __future__ import annotations
import enum
from collections.abc import Callable
from dataclasses import dataclass, field
import torch
class ColorSpace(enum.Enum):
"""YUV color space standard."""
BT_709 = "bt709"
BT_2020_NCL = "bt2020ncl"
@property
def av_colorspace(self) -> int:
"""FFmpeg ``AVCOL_SPC_*`` constant for ``codec_context.colorspace``."""
return _AV_COLORSPACE[self]
class ColorRange(enum.Enum):
"""YUV color range."""
MPEG = "mpeg"
JPEG = "jpeg"
@property
def av_color_range(self) -> int:
"""FFmpeg ``AVCOL_RANGE_*`` constant for ``codec_context.color_range``."""
return _AV_COLOR_RANGE[self]
class PixelFormat(enum.Enum):
"""Pixel format for video frames."""
RGB24 = "rgb24"
YUV420P = "yuv420p"
@property
def av_format(self) -> str:
"""PyAV format string for ``VideoFrame.from_ndarray``."""
return self.value
_AV_COLORSPACE = {
ColorSpace.BT_709: 1, # AVCOL_SPC_BT709
ColorSpace.BT_2020_NCL: 9, # AVCOL_SPC_BT2020_NCL
}
_AV_COLOR_RANGE = {
ColorRange.MPEG: 1, # AVCOL_RANGE_MPEG (limited)
ColorRange.JPEG: 2, # AVCOL_RANGE_JPEG (full)
}
# BT.709 RGB->YUV matrix (row-major: each row produces one of Y, U, V)
_BT709_MATRIX = torch.tensor(
[
[0.2126, 0.7152, 0.0722],
[-0.1146, -0.3854, 0.5],
[0.5, -0.4542, -0.0458],
],
dtype=torch.float32,
)
# BT.2020 NCL RGB->YUV matrix
_KR_2020 = 0.2627
_KG_2020 = 0.6780
_KB_2020 = 0.0593
_BT2020_MATRIX = torch.tensor(
[
[_KR_2020, _KG_2020, _KB_2020],
[-_KR_2020 / 1.8814, -_KG_2020 / 1.8814, 0.5],
[0.5, -_KG_2020 / 1.4746, -_KB_2020 / 1.4746],
],
dtype=torch.float32,
)
_COLOR_SPACE_MATRICES = {
ColorSpace.BT_709: _BT709_MATRIX,
ColorSpace.BT_2020_NCL: _BT2020_MATRIX,
}
@dataclass(frozen=True)
class FrameConverter:
"""Converts ``[*, C, H, W]`` float ``[0, 1]`` frames to uint8.
Carries encoding metadata so ``encode_video`` can derive pixel format,
color space, and color range from the converter itself.
The ``fn_`` callable **may mutate its input** (PyTorch trailing-underscore
convention). Callers that need to keep the original ``frames`` afterwards
must pass ``frames.clone()``. Inside ``encode_video``'s per-chunk
generator each chunk is consumed once, so direct passthrough is safe.
"""
pixel_format: PixelFormat
fn_: Callable[[torch.Tensor], torch.Tensor] = field(repr=False)
color_space: ColorSpace | None = None
color_range: ColorRange | None = None
def __call__(self, frames: torch.Tensor) -> torch.Tensor:
return self.fn_(frames)
def rgb_to_yuv(image: torch.Tensor, color_space: ColorSpace) -> torch.Tensor:
"""Convert an RGB image to YUV.
The image data is assumed to be in the range of ``[0, 1]``.
Uses a single matrix multiply for better memory locality.
Args:
image: RGB image with shape ``(*, 3, H, W)``.
color_space: Color space standard for the conversion matrix.
Returns:
YUV image with shape ``(*, 3, H, W)``.
"""
if len(image.shape) < 3 or image.shape[-3] != 3:
raise ValueError(f"Input size must have a shape of (*, 3, H, W). Got {image.shape}")
mat = _COLOR_SPACE_MATRICES[color_space].to(device=image.device, dtype=image.dtype)
# [*, 3, H, W] -> [*, H, W, 3] @ [3, 3]^T -> [*, H, W, 3] -> [*, 3, H, W]
pixels = image.movedim(-3, -1) # [*, H, W, 3]
yuv = pixels @ mat.T # [*, H, W, 3]
return yuv.movedim(-1, -3) # [*, 3, H, W]
def apply_color_range_(y: torch.Tensor, uv: torch.Tensor, color_range: ColorRange) -> tuple[torch.Tensor, torch.Tensor]:
"""Scale Y and UV planes to the specified color range, in-place.
Args:
y: Luma plane in ``[0, 1]``.
uv: Chroma planes centered at 0.
color_range: Target color range.
Returns:
Scaled ``(Y, UV)`` tensors (modified in-place).
"""
if color_range == ColorRange.MPEG:
y.mul_(219).add_(16)
uv.mul_(224).add_(128)
elif color_range == ColorRange.JPEG:
y.mul_(255)
uv.add_(0.5).mul_(255)
else:
raise ValueError(f"Unsupported color range: {color_range}")
return y, uv
def rgb_to_yuv420(
image: torch.Tensor, color_space: ColorSpace, color_range: ColorRange
) -> tuple[torch.Tensor, torch.Tensor]:
"""Convert an RGB image to YUV 4:2:0 with chroma subsampling.
Chroma is subsampled by averaging 2x2 pixel blocks (chroma siting
``(128, 128)``).
Args:
image: RGB image with shape ``(*, 3, H, W)`` in ``[0, 1]``.
H and W must be divisible by 2.
color_space: Color space standard.
color_range: Color range for the output.
Returns:
``(Y, UV)`` where Y has shape ``(*, 1, H, W)`` and UV has shape
``(*, 2, H//2, W//2)``.
"""
if len(image.shape) < 3 or image.shape[-3] != 3:
raise ValueError(f"Input size must have a shape of (*, 3, H, W). Got {image.shape}")
if image.shape[-2] % 2 != 0 or image.shape[-1] % 2 != 0:
raise ValueError(f"Input H and W must be divisible by 2. Got {image.shape}")
yuv = rgb_to_yuv(image, color_space)
y = yuv[..., :1, :, :]
# Subsample chroma: average 2x2 blocks via avg_pool2d (contiguous, fused kernel)
uv_full = yuv[..., 1:3, :, :].contiguous()
# Flatten leading dims for avg_pool2d which expects [N, C, H, W]
lead = uv_full.shape[:-3]
uv_flat = uv_full.reshape(-1, 2, uv_full.shape[-2], uv_full.shape[-1])
uv = torch.nn.functional.avg_pool2d(uv_flat, kernel_size=2, stride=2)
uv = uv.reshape(*lead, 2, uv.shape[-2], uv.shape[-1])
return apply_color_range_(y, uv, color_range)
def pack_i420(y: torch.Tensor, uv: torch.Tensor) -> torch.Tensor:
"""Pack Y and UV planes into I420 layout for pyav.
I420 packs the three planes into a single 2D array of height ``H * 3 // 2``
and width ``W``. The Y plane occupies the first ``H`` rows. The UV tensor
``(*, 2, H//2, W//2)`` is reshaped to ``(*, H//2, W)`` -- U rows packed
two-by-two followed by V rows packed two-by-two -- and appended below.
Args:
y: Luma with shape ``(*, 1, H, W)``.
uv: Chroma with shape ``(*, 2, H//2, W//2)``.
Returns:
Packed tensor with shape ``(*, H*3//2, W)`` uint8.
"""
y_plane = y[..., 0, :, :] # [*, H, W]
uv_packed = uv.reshape(*uv.shape[:-3], uv.shape[-2], uv.shape[-1] * 2) # [*, H//2, W]
packed = torch.cat([y_plane, uv_packed], dim=-2) # [*, H*3//2, W]
return packed.clamp_(0, 255).to(torch.uint8)
def _rgb_uint8_fn_(frames: torch.Tensor) -> torch.Tensor:
"""In-place: mutates ``frames`` via ``clamp_`` + ``mul_``, returns a uint8 view."""
return frames.clamp_(0.0, 1.0).mul_(255.0).to(torch.uint8).movedim(-3, -1)
rgb_uint8_converter_ = FrameConverter(pixel_format=PixelFormat.RGB24, fn_=_rgb_uint8_fn_)
"""``(*, 3, H, W)`` float ``[0, 1]`` to ``(*, H, W, 3)`` uint8. Mutates input."""
def _yuv420p_bt709_fn_(frames: torch.Tensor) -> torch.Tensor:
y, uv = rgb_to_yuv420(frames, ColorSpace.BT_709, ColorRange.MPEG)
return pack_i420(y, uv)
yuv420p_bt709_converter_ = FrameConverter(
pixel_format=PixelFormat.YUV420P,
fn_=_yuv420p_bt709_fn_,
color_space=ColorSpace.BT_709,
color_range=ColorRange.MPEG,
)
"""``(*, 3, H, W)`` float ``[0, 1]`` to ``(*, H*3//2, W)`` uint8 YUV420p BT.709 MPEG."""
packages/ltx-pipelines/src/ltx_pipelines/utils/denoisers.py
+42 -19
View File
@@ -20,6 +20,7 @@ from ltx_core.guidance.perturbations import (
from ltx_core.model.transformer import X0Model
from ltx_core.types import LatentState
from ltx_pipelines.utils.helpers import modality_from_latent_state
from ltx_pipelines.utils.types import DenoisedLatentResult
_POSITIVE_ONLY_GUIDER = MultiModalGuider(
params=MultiModalGuiderParams(cfg_scale=1.0, stg_scale=0.0, modality_scale=1.0),
@@ -53,7 +54,7 @@ def _repeat_state(state: LatentState, n: int) -> LatentState:
)
def _guided_denoise( # noqa: PLR0913
def _guided_denoise( # noqa: PLR0913,PLR0915
transformer: X0Model,
video_state: LatentState | None,
audio_state: LatentState | None,
@@ -66,7 +67,8 @@ def _guided_denoise( # noqa: PLR0913
last_denoised_video: torch.Tensor | None,
last_denoised_audio: torch.Tensor | None,
step_index: int,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
force_uncond_pass: bool = False,
) -> tuple[DenoisedLatentResult | None, DenoisedLatentResult | None]:
"""Core guided denoising — batches all guidance passes into one transformer call.
Collects per-pass contexts first, then builds a single batched Modality
per present modality via :func:`modality_from_latent_state`. When wrapped
@@ -80,7 +82,9 @@ def _guided_denoise( # noqa: PLR0913
a_skip = audio_guider.should_skip_step(step_index)
if v_skip and a_skip:
return last_denoised_video, last_denoised_audio
video_result = DenoisedLatentResult.result_or_none(denoised=last_denoised_video)
audio_result = DenoisedLatentResult.result_or_none(denoised=last_denoised_audio)
return video_result, audio_result
if video_state is not None and v_context is None:
raise ValueError("v_context is required when video_state is provided")
@@ -91,10 +95,12 @@ def _guided_denoise( # noqa: PLR0913
_pass = tuple[str, torch.Tensor | None, torch.Tensor | None, PerturbationConfig]
passes: list[_pass] = [("cond", v_context, a_context, PerturbationConfig.empty())]
if video_guider.do_unconditional_generation() or audio_guider.do_unconditional_generation():
if video_guider.do_unconditional_generation() and video_guider.negative_context is None:
v_needs_neg = video_guider.do_unconditional_generation() or (force_uncond_pass and video_state is not None)
a_needs_neg = audio_guider.do_unconditional_generation() or (force_uncond_pass and audio_state is not None)
if v_needs_neg or a_needs_neg:
if v_needs_neg and video_guider.negative_context is None:
raise ValueError("Negative context is required for unconditioned denoising")
if audio_guider.do_unconditional_generation() and audio_guider.negative_context is None:
if a_needs_neg and audio_guider.negative_context is None:
raise ValueError("Negative context is required for unconditioned denoising")
v_neg = video_guider.negative_context if video_guider.negative_context is not None else v_context
a_neg = audio_guider.negative_context if audio_guider.negative_context is not None else a_context
@@ -172,7 +178,14 @@ def _guided_denoise( # noqa: PLR0913
denoised_video = last_denoised_video if v_skip else video_guider.calculate(cond_v, uncond_v, ptb_v, mod_v)
denoised_audio = last_denoised_audio if a_skip else audio_guider.calculate(cond_a, uncond_a, ptb_a, mod_a)
return denoised_video, denoised_audio
return (
DenoisedLatentResult.result_or_none(
denoised=denoised_video, uncond=uncond_v, cond=cond_v, ptb=ptb_v, mod=mod_v
),
DenoisedLatentResult.result_or_none(
denoised=denoised_audio, uncond=uncond_a, cond=cond_a, ptb=ptb_a, mod=mod_a
),
)
class SimpleDenoiser:
@@ -195,11 +208,15 @@ class SimpleDenoiser:
audio_state: LatentState | None,
sigmas: torch.Tensor,
step_index: int,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
) -> tuple[DenoisedLatentResult | None, DenoisedLatentResult | None]:
sigma = sigmas[step_index]
pos_video = modality_from_latent_state(video_state, self.v_context, sigma) if video_state is not None else None
pos_audio = modality_from_latent_state(audio_state, self.a_context, sigma) if audio_state is not None else None
return transformer(video=pos_video, audio=pos_audio, perturbations=None)
denoised_video, denoised_audio = transformer(video=pos_video, audio=pos_audio, perturbations=None)
return (
DenoisedLatentResult.result_or_none(denoised=denoised_video),
DenoisedLatentResult.result_or_none(denoised=denoised_audio),
)
class GuidedDenoiser:
@@ -214,11 +231,13 @@ class GuidedDenoiser:
a_context: torch.Tensor | None,
video_guider: MultiModalGuider | None = None,
audio_guider: MultiModalGuider | None = None,
force_uncond_pass: bool = False,
) -> None:
self.v_context = v_context
self.a_context = a_context
self.video_guider = video_guider
self.audio_guider = audio_guider
self.force_uncond_pass = force_uncond_pass
self._last_denoised_video: torch.Tensor | None = None
self._last_denoised_audio: torch.Tensor | None = None
@@ -229,8 +248,8 @@ class GuidedDenoiser:
audio_state: LatentState | None,
sigmas: torch.Tensor,
step_index: int,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
denoised_video, denoised_audio = _guided_denoise(
) -> tuple[DenoisedLatentResult | None, DenoisedLatentResult | None]:
guided_denoise_result_v, guided_denoise_result_a = _guided_denoise(
transformer=transformer,
video_state=video_state,
audio_state=audio_state,
@@ -242,10 +261,11 @@ class GuidedDenoiser:
last_denoised_video=self._last_denoised_video,
last_denoised_audio=self._last_denoised_audio,
step_index=step_index,
force_uncond_pass=self.force_uncond_pass,
)
self._last_denoised_video = denoised_video
self._last_denoised_audio = denoised_audio
return denoised_video, denoised_audio
self._last_denoised_video = guided_denoise_result_v.denoised
self._last_denoised_audio = guided_denoise_result_a.denoised
return guided_denoise_result_v, guided_denoise_result_a
class FactoryGuidedDenoiser:
@@ -257,11 +277,13 @@ class FactoryGuidedDenoiser:
a_context: torch.Tensor | None,
video_guider_factory: MultiModalGuiderFactory | None = None,
audio_guider_factory: MultiModalGuiderFactory | None = None,
force_uncond_pass: bool = False,
) -> None:
self.v_context = v_context
self.a_context = a_context
self.video_guider_factory = video_guider_factory
self.audio_guider_factory = audio_guider_factory
self.force_uncond_pass = force_uncond_pass
self._last_denoised_video: torch.Tensor | None = None
self._last_denoised_audio: torch.Tensor | None = None
self._sigma_vals_cached: list[float] | None = None
@@ -273,7 +295,7 @@ class FactoryGuidedDenoiser:
audio_state: LatentState | None,
sigmas: torch.Tensor,
step_index: int,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
) -> tuple[DenoisedLatentResult | None, DenoisedLatentResult | None]:
if self._sigma_vals_cached is None:
self._sigma_vals_cached = sigmas.detach().cpu().tolist()
sigma_val = self._sigma_vals_cached[step_index]
@@ -287,7 +309,7 @@ class FactoryGuidedDenoiser:
else None
)
denoised_video, denoised_audio = _guided_denoise(
guided_denoise_result_v, guided_denoise_result_a = _guided_denoise(
transformer=transformer,
video_state=video_state,
audio_state=audio_state,
@@ -299,7 +321,8 @@ class FactoryGuidedDenoiser:
last_denoised_video=self._last_denoised_video,
last_denoised_audio=self._last_denoised_audio,
step_index=step_index,
force_uncond_pass=self.force_uncond_pass,
)
self._last_denoised_video = denoised_video
self._last_denoised_audio = denoised_audio
return denoised_video, denoised_audio
self._last_denoised_video = guided_denoise_result_v.denoised
self._last_denoised_audio = guided_denoise_result_a.denoised
return guided_denoise_result_v, guided_denoise_result_a
packages/ltx-pipelines/src/ltx_pipelines/utils/media_io.py
+109 -31
View File
@@ -1,10 +1,12 @@
import enum
import logging
import math
import threading
from collections.abc import Generator, Iterator
from fractions import Fraction
from io import BytesIO
from pathlib import Path
from queue import Queue
import av
import numpy as np
@@ -17,6 +19,7 @@ from tqdm import tqdm
from ltx_core.hdr import LogC3
from ltx_core.types import Audio, VideoPixelShape
from ltx_pipelines.utils.color_conversion import FrameConverter, PixelFormat, yuv420p_bt709_converter_
from ltx_pipelines.utils.constants import DEFAULT_IMAGE_CRF
logger = logging.getLogger(__name__)
@@ -86,8 +89,8 @@ def resize_and_center_crop(tensor: torch.Tensor, height: int, width: int) -> tor
return tensor
def normalize_latent(latent: torch.Tensor, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
return (latent / 127.5 - 1.0).to(device=device, dtype=dtype)
def normalize_images(images: torch.Tensor, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
return (images / 127.5 - 1.0).to(device=device, dtype=dtype)
def to_vae_range(x: torch.Tensor) -> torch.Tensor:
@@ -116,7 +119,7 @@ def load_image_and_preprocess(
image = preprocess(image=image, crf=crf)
image = torch.tensor(image, dtype=torch.float32, device=device)
image = resize_and_center_crop(image, height, width)
image = normalize_latent(image, device, dtype)
image = normalize_images(image, device, dtype)
return image
@@ -137,11 +140,13 @@ def video_preprocess(
Returns:
Tensor of shape (1, C, F, height, width) with values in [-1, 1].
"""
result = None
result: torch.Tensor | None = None
for f in frames:
frame = resize_and_center_crop(f.to(torch.float32), height, width)
frame = normalize_latent(frame, device, dtype)
frame = normalize_images(frame, device, dtype)
result = frame if result is None else torch.cat([result, frame], dim=2)
if result is None:
raise ValueError("video_preprocess received an empty frame generator; no frames were decoded from the source.")
return result
@@ -325,47 +330,120 @@ def encode_video(
audio: Audio | None,
output_path: str,
video_chunks_number: int,
frame_converter: FrameConverter = yuv420p_bt709_converter_,
crf: int = 19,
preset: str = "veryfast",
thread_count: int = 0,
) -> None:
if isinstance(video, torch.Tensor):
video = iter([video])
first_chunk = next(video)
def convert(chunk: torch.Tensor) -> torch.Tensor:
return frame_converter(chunk.movedim(-1, -3))
_, height, width, _ = first_chunk.shape
first_chunk = convert(next(video))
if frame_converter.pixel_format == PixelFormat.RGB24:
height, width = first_chunk.shape[-3], first_chunk.shape[-2]
else:
height = first_chunk.shape[-2] * 2 // 3
width = first_chunk.shape[-1]
container = av.open(output_path, mode="w")
stream = container.add_stream("libx264", rate=int(fps))
stream.width = width
stream.height = height
stream.pix_fmt = "yuv420p"
success = False
try:
stream = container.add_stream("libx264", rate=int(fps), options={"crf": str(crf), "preset": preset})
stream.width = width
stream.height = height
stream.pix_fmt = "yuv420p"
stream.codec_context.thread_count = thread_count
stream.codec_context.thread_type = "FRAME"
if frame_converter.color_space is not None:
stream.codec_context.colorspace = frame_converter.color_space.av_colorspace
if frame_converter.color_range is not None:
stream.codec_context.color_range = frame_converter.color_range.av_color_range
if audio is not None:
audio_stream = _prepare_audio_stream(container, audio.sampling_rate)
if audio is not None:
audio_stream = _prepare_audio_stream(container, audio.sampling_rate)
def all_tiles(
first_chunk: torch.Tensor, tiles_generator: Generator[tuple[torch.Tensor, int], None, None]
) -> Generator[tuple[torch.Tensor, int], None, None]:
yield first_chunk
yield from tiles_generator
av_format = frame_converter.pixel_format.av_format
for video_chunk in tqdm(all_tiles(first_chunk, video), total=video_chunks_number):
video_chunk_cpu = video_chunk.to("cpu").numpy()
for frame_array in video_chunk_cpu:
frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24")
for packet in stream.encode(frame):
container.mux(packet)
def cpu_chunks() -> Generator[np.ndarray, None, None]:
yield first_chunk.to("cpu").numpy()
for chunk in video:
yield convert(chunk).to("cpu").numpy()
# Flush encoder
for packet in stream.encode():
container.mux(packet)
_encode_chunks_threaded(
container=container,
stream=stream,
av_format=av_format,
chunks=cpu_chunks(),
progress_total=video_chunks_number,
)
if audio is not None:
_write_audio(container, audio_stream, audio)
container.close()
if audio is not None:
_write_audio(container, audio_stream, audio)
success = True
finally:
container.close()
if not success:
Path(output_path).unlink(missing_ok=True)
logger.info(f"Video saved to {output_path}")
def _encode_chunks_threaded(
container: av.container.Container,
stream: av.video.stream.VideoStream,
av_format: str,
chunks: Iterator[np.ndarray],
progress_total: int,
) -> None:
"""Run libx264 frame.encode + container.mux on a background thread while
the caller produces numpy chunks on the current thread. The 1-slot queue
lets the producer get one chunk ahead (so the next VAE/gather chunk
overlaps with libx264 encoding the previous chunk) without buffering more
than one chunk in CPU memory.
"""
chunk_queue: Queue[np.ndarray | None] = Queue(maxsize=1)
encoder_error: list[BaseException] = []
def encoder_worker() -> None:
error: BaseException | None = None
while True:
arr = chunk_queue.get()
if arr is None:
break
if error is not None:
continue
try:
for frame_array in arr:
frame = av.VideoFrame.from_ndarray(frame_array, format=av_format)
for packet in stream.encode(frame):
container.mux(packet)
except Exception as e:
error = e
if error is None:
try:
for packet in stream.encode():
container.mux(packet)
except Exception as e:
error = e
if error is not None:
encoder_error.append(error)
encoder_thread = threading.Thread(target=encoder_worker, name="h264-encoder")
encoder_thread.start()
try:
for arr in tqdm(chunks, total=progress_total):
chunk_queue.put(arr)
finally:
chunk_queue.put(None)
encoder_thread.join()
if encoder_error:
raise encoder_error[0]
_INT_FORMAT_MAX: dict[str, float] = {
"u8": 128.0,
"u8p": 128.0,
packages/ltx-pipelines/src/ltx_pipelines/utils/samplers.py
+140 -7
View File
@@ -6,7 +6,7 @@ from typing import Callable
import torch
from tqdm import tqdm
from ltx_core.components.diffusion_steps import Res2sDiffusionStep
from ltx_core.components.diffusion_steps import EulerCfgPpDiffusionStep, Res2sDiffusionStep
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.model.transformer import X0Model
from ltx_core.utils import to_denoised, to_velocity
@@ -60,13 +60,15 @@ def euler_denoising_loop(
denoiser:
A callable implementing :class:`Denoiser`. It is invoked as
``denoiser(transformer, video_state, audio_state, sigmas, step_index)``
and must return ``(denoised_video, denoised_audio)``.
and must return a :class:`~ltx_pipelines.utils.types.DenoisedLatentResult`.
### Returns
tuple[LatentState | None, LatentState | None]
Final ``(video_state, audio_state)`` after the denoising loop.
"""
for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
denoised_video, denoised_audio = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
video_result, audio_result = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
denoised_video = video_result.denoised if video_result is not None else None
denoised_audio = audio_result.denoised if audio_result is not None else None
video_state = _step_state(video_state, denoised_video, stepper, sigmas, step_idx)
audio_state = _step_state(audio_state, denoised_audio, stepper, sigmas, step_idx)
@@ -110,7 +112,9 @@ def gradient_estimating_euler_denoising_loop(
return current_velocity, denoised_sample
for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
denoised_video, denoised_audio = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
video_result, audio_result = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
denoised_video = video_result.denoised if video_result is not None else None
denoised_audio = audio_result.denoised if audio_result is not None else None
if video_state is not None and denoised_video is not None:
denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
@@ -143,6 +147,11 @@ def gradient_estimating_euler_denoising_loop(
return (video_state, audio_state)
def _get_plain_noise(x: torch.Tensor, generator: torch.Generator) -> torch.Tensor:
"""Draw standard Gaussian noise matching the shape, dtype, and device of ``x``."""
return torch.randn(x.shape, generator=generator, dtype=x.dtype, device=x.device)
def _channelwise_normalize(x: torch.Tensor) -> torch.Tensor:
return x.sub_(x.mean(dim=(-2, -1), keepdim=True)).div_(x.std(dim=(-2, -1), keepdim=True))
@@ -278,7 +287,9 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915,PLR0912
# ====================================================================
# STAGE 1: Evaluate at current point
# ====================================================================
denoised_video_1, denoised_audio_1 = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
video_result, audio_result = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
denoised_video_1 = video_result.denoised if video_result is not None else None
denoised_audio_1 = audio_result.denoised if audio_result is not None else None
if video_state is not None and denoised_video_1 is not None:
denoised_video_1 = post_process_latent(denoised_video_1, video_state.denoise_mask, video_state.clean_latent)
if audio_state is not None and denoised_audio_1 is not None:
@@ -355,13 +366,15 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915,PLR0912
else None
)
denoised_video_2, denoised_audio_2 = denoiser(
video_result_2, audio_result_2 = denoiser(
transformer,
video_state=mid_video_state,
audio_state=mid_audio_state,
sigmas=torch.stack([sub_sigma]).to(sigmas.device),
step_index=0,
)
denoised_video_2 = video_result_2.denoised if video_result_2 is not None else None
denoised_audio_2 = audio_result_2.denoised if audio_result_2 is not None else None
if video_state is not None and denoised_video_2 is not None:
denoised_video_2 = post_process_latent(denoised_video_2, video_state.denoise_mask, video_state.clean_latent)
if audio_state is not None and denoised_audio_2 is not None:
@@ -410,7 +423,9 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915,PLR0912
# Final step if we need to fully remove the noise
if sigmas[-1] == 0:
denoised_video_1, denoised_audio_1 = denoiser(transformer, video_state, audio_state, sigmas, n_full_steps)
video_result_final, audio_result_final = denoiser(transformer, video_state, audio_state, sigmas, n_full_steps)
denoised_video_1 = video_result_final.denoised if video_result_final is not None else None
denoised_audio_1 = audio_result_final.denoised if audio_result_final is not None else None
if video_state is not None and denoised_video_1 is not None:
denoised_video_1 = post_process_latent(denoised_video_1, video_state.denoise_mask, video_state.clean_latent)
video_state = replace(video_state, latent=denoised_video_1.to(model_dtype))
@@ -419,3 +434,121 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915,PLR0912
audio_state = replace(audio_state, latent=denoised_audio_1.to(model_dtype))
return video_state, audio_state
def euler_cfg_pp_denoising_loop(
sigmas: torch.Tensor,
video_state: LatentState | None,
audio_state: LatentState | None,
stepper: EulerCfgPpDiffusionStep,
transformer: X0Model,
denoiser: Denoiser,
noise_seed: int = -1,
new_noise_fn: Callable[[torch.Tensor, torch.Generator], torch.Tensor] = _get_plain_noise,
model_dtype: torch.dtype = torch.bfloat16,
) -> tuple[LatentState | None, LatentState | None]:
"""
Joint audio-video denoising loop using the CFG++ corrected Euler sampler.
Applies the CFG++ update rule at each step: the ODE derivative is computed
from the unconditioned denoised prediction rather than the standard velocity,
and an ancestral DDIM noise injection is applied in the rescaled sigma space.
Requires a guided denoiser whose :class:`~ltx_pipelines.utils.types.DenoisedLatentResult`
carries ``uncond`` tensors (i.e. CFG must be enabled).
Either ``video_state`` or ``audio_state`` may be ``None`` for absent modalities.
When both are present, noise is drawn from the same seeded generator (video
first, audio second) to produce a consistent random sequence.
### Parameters
sigmas:
1-D tensor of noise levels defining the sampling schedule.
video_state:
Current video :class:`~ltx_core.types.LatentState`, or ``None``.
audio_state:
Current audio :class:`~ltx_core.types.LatentState`, or ``None``.
stepper:
:class:`~ltx_core.components.diffusion_steps.EulerCfgPpDiffusionStep`
instance carrying ``eta`` and ``s_noise`` parameters.
transformer:
The diffusion model passed to the denoiser at each step.
denoiser:
Callable implementing :class:`~ltx_pipelines.utils.types.Denoiser`.
noise_seed:
Integer seed for the noise generator. Default ``-1``.
new_noise_fn:
``(latent, generator) -> noise`` callable. Defaults to plain
``torch.randn`` (no channel-wise normalization). Pass
:func:`_get_new_noise` for the normalized variant used in res2s.
model_dtype:
Dtype for latent state updates. Default ``bfloat16``.
### Returns
tuple[LatentState | None, LatentState | None]
Final ``(video_state, audio_state)`` after the denoising loop.
"""
if not isinstance(stepper, EulerCfgPpDiffusionStep):
raise ValueError(f"stepper must be an instance of EulerCfgPpDiffusionStep, got {type(stepper).__name__}")
present_state = video_state or audio_state
if present_state is None:
raise ValueError("At least one of video_state or audio_state must be provided")
generator = torch.Generator(device=present_state.latent.device).manual_seed(noise_seed)
draw_noise = stepper.eta > 0 and stepper.s_noise > 0
for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
video_result, audio_result = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
denoised_video = video_result.denoised if video_result is not None else None
denoised_audio = audio_result.denoised if audio_result is not None else None
uncond_video = video_result.uncond if video_result is not None else None
uncond_audio = audio_result.uncond if audio_result is not None else None
if video_state is not None and not isinstance(uncond_video, torch.Tensor):
raise ValueError(
"euler_cfg_pp_denoising_loop requires video DenoisedLatentResult.uncond to be a tensor. "
"Use GuidedDenoiser or FactoryGuidedDenoiser with cfg_scale != 1 "
"or force_uncond_pass=True and a negative_context."
)
if audio_state is not None and not isinstance(uncond_audio, torch.Tensor):
raise ValueError(
"euler_cfg_pp_denoising_loop requires audio DenoisedLatentResult.uncond to be a tensor. "
"Use GuidedDenoiser or FactoryGuidedDenoiser with cfg_scale != 1 "
"or force_uncond_pass=True and a negative_context."
)
if video_state is not None and denoised_video is not None:
denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
if audio_state is not None and denoised_audio is not None:
denoised_audio = post_process_latent(denoised_audio, audio_state.denoise_mask, audio_state.clean_latent)
if sigmas[step_idx + 1] == 0:
if video_state is not None and denoised_video is not None:
video_state = replace(video_state, latent=denoised_video.to(model_dtype))
if audio_state is not None and denoised_audio is not None:
audio_state = replace(audio_state, latent=denoised_audio.to(model_dtype))
return video_state, audio_state
# Draw noise consecutively from the same generator: video first, audio second.
noise_video = new_noise_fn(video_state.latent, generator) if (video_state is not None and draw_noise) else None
noise_audio = new_noise_fn(audio_state.latent, generator) if (audio_state is not None and draw_noise) else None
if video_state is not None and denoised_video is not None:
x_next = stepper.step(
sample=video_state.latent,
denoised_sample=denoised_video,
sigmas=sigmas,
step_index=step_idx,
uncond_denoised=uncond_video,
noise=noise_video,
)
video_state = replace(video_state, latent=x_next.to(model_dtype))
if audio_state is not None and denoised_audio is not None:
x_next = stepper.step(
sample=audio_state.latent,
denoised_sample=denoised_audio,
sigmas=sigmas,
step_index=step_idx,
uncond_denoised=uncond_audio,
noise=noise_audio,
)
audio_state = replace(audio_state, latent=x_next.to(model_dtype))
return video_state, audio_state
packages/ltx-pipelines/src/ltx_pipelines/utils/types.py
+33 -2
View File
@@ -40,6 +40,36 @@ class PipelineComponents:
self.audio_patchifier = AudioPatchifier(patch_size=1)
@dataclass(frozen=True)
class DenoisedLatentResult:
"""Output of one denoiser call for a single modality.
``denoised`` is the final blended prediction for this modality.
The remaining fields carry the per-pass raw outputs from ``_guided_denoise``
(all ``None`` for ``SimpleDenoiser``). Denoisers return a
``(video_result, audio_result)`` tuple; either element may be ``None``
for absent modalities.
"""
denoised: torch.Tensor
uncond: torch.Tensor | None = None
cond: torch.Tensor | None = None
ptb: torch.Tensor | None = None
mod: torch.Tensor | None = None
@classmethod
def result_or_none(
cls,
denoised: torch.Tensor | None,
uncond: torch.Tensor | None = None,
cond: torch.Tensor | None = None,
ptb: torch.Tensor | None = None,
mod: torch.Tensor | None = None,
) -> DenoisedLatentResult | None:
if denoised is None:
return None
return cls(denoised=denoised, uncond=uncond, cond=cond, ptb=ptb, mod=mod)
class Denoiser(Protocol):
"""Protocol for a denoiser that receives the transformer at call time.
The transformer is not stored — it is passed as the first argument so the
@@ -51,7 +81,8 @@ class Denoiser(Protocol):
sigmas: 1-D tensor of sigma values for each diffusion step.
step_index: Index of the current denoising step.
Returns:
``(denoised_video, denoised_audio)`` tensors (either may be ``None``).
A ``(video_result, audio_result)`` tuple of :class:`DenoisedLatentResult`,
either may be ``None`` for absent modalities.
"""
def __call__(
@@ -61,7 +92,7 @@ class Denoiser(Protocol):
audio_state: LatentState | None,
sigmas: torch.Tensor,
step_index: int,
) -> tuple[torch.Tensor | None, torch.Tensor | None]: ...
) -> tuple[DenoisedLatentResult | None, DenoisedLatentResult | None]: ...
@dataclass(frozen=True)