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

Automated PR - 2026-03-04

Browse Source
This commit is contained in:
sync-bot
2026-03-04 19:34:46 +00:00
parent 28c3c73fe5
commit 822ce3c4b1
73 changed files with 4984 additions and 1220 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-pipelines/src/ltx_pipelines/utils/__init__.py
+28
View File
@@ -1,5 +1,33 @@
from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
denoise_audio_video,
generate_enhanced_prompt,
get_device,
image_conditionings_by_replacing_latent,
multi_modal_guider_denoising_func,
multi_modal_guider_factory_denoising_func,
simple_denoising_func,
)
from ltx_pipelines.utils.model_ledger import ModelLedger
from ltx_pipelines.utils.samplers import (
euler_denoising_loop,
gradient_estimating_euler_denoising_loop,
res2s_audio_video_denoising_loop,
)
__all__ = [
"ModelLedger",
"assert_resolution",
"cleanup_memory",
"denoise_audio_video",
"euler_denoising_loop",
"generate_enhanced_prompt",
"get_device",
"gradient_estimating_euler_denoising_loop",
"image_conditionings_by_replacing_latent",
"multi_modal_guider_denoising_func",
"multi_modal_guider_factory_denoising_func",
"res2s_audio_video_denoising_loop",
"simple_denoising_func",
]
packages/ltx-pipelines/src/ltx_pipelines/utils/args.py
+132 -80
View File
@@ -1,24 +1,25 @@
import argparse
from pathlib import Path
from typing import NamedTuple
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
from ltx_core.quantization import QuantizationPolicy
from ltx_pipelines.utils.constants import (
DEFAULT_1_STAGE_HEIGHT,
DEFAULT_1_STAGE_WIDTH,
DEFAULT_2_STAGE_HEIGHT,
DEFAULT_2_STAGE_WIDTH,
DEFAULT_AUDIO_GUIDER_PARAMS,
DEFAULT_FRAME_RATE,
DEFAULT_IMAGE_CRF,
DEFAULT_LORA_STRENGTH,
DEFAULT_NEGATIVE_PROMPT,
DEFAULT_NUM_FRAMES,
DEFAULT_NUM_INFERENCE_STEPS,
DEFAULT_SEED,
DEFAULT_VIDEO_GUIDER_PARAMS,
LTX_2_3_PARAMS,
PipelineParams,
)
class ImageConditioningInput(NamedTuple):
path: str
frame_idx: int
strength: float
crf: int = DEFAULT_IMAGE_CRF
class VideoConditioningAction(argparse.Action):
def __call__(
self,
@@ -35,20 +36,50 @@ class VideoConditioningAction(argparse.Action):
setattr(namespace, self.dest, current)
class VideoMaskConditioningAction(argparse.Action):
"""Parse ``--conditioning-attention-mask PATH STRENGTH``.
Stores a ``(mask_path, strength)`` tuple on the namespace. The mask video
should be grayscale with pixel values in [0, 1] controlling per-region
conditioning attention strength. The scalar *STRENGTH* is multiplied with
the spatial mask before it is applied.
"""
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} requires exactly 2 arguments (MASK_PATH STRENGTH), got {len(values)}"
raise argparse.ArgumentError(self, msg)
mask_path = resolve_path(values[0])
strength = float(values[1])
setattr(namespace, self.dest, (mask_path, strength))
class ImageAction(argparse.Action):
def __call__(
self,
parser: argparse.ArgumentParser, # noqa: ARG002
namespace: argparse.Namespace,
values: list[str],
option_string: str | None = None, # noqa: ARG002
option_string: str | None = None,
) -> None:
path, frame_idx, strength_str = values
resolved_path = resolve_path(path)
frame_idx = int(frame_idx)
strength = float(strength_str)
if len(values) not in (3, 4):
msg = f"{option_string} requires 3 or 4 arguments (PATH FRAME_IDX STRENGTH [CRF]), got {len(values)}"
raise argparse.ArgumentError(self, msg)
conditioning = ImageConditioningInput(
path=resolve_path(values[0]),
frame_idx=int(values[1]),
strength=float(values[2]),
crf=int(values[3]) if len(values) > 3 else DEFAULT_IMAGE_CRF,
)
current = getattr(namespace, self.dest) or []
current.append((resolved_path, frame_idx, strength))
current.append(conditioning)
setattr(namespace, self.dest, current)
@@ -113,14 +144,34 @@ class QuantizationAction(argparse.Action):
setattr(namespace, self.dest, policy)
def basic_arg_parser() -> argparse.ArgumentParser:
def detect_checkpoint_path(distilled: bool = False) -> str:
"""Pre-parse argv to extract the checkpoint path before building the full parser."""
pre = argparse.ArgumentParser(add_help=False)
flag = "--distilled-checkpoint-path" if distilled else "--checkpoint-path"
pre.add_argument(flag, type=resolve_path, required=True)
known, _ = pre.parse_known_args()
return known.distilled_checkpoint_path if distilled else known.checkpoint_path
def basic_arg_parser(
params: PipelineParams = LTX_2_3_PARAMS,
distilled: bool = False,
) -> argparse.ArgumentParser:
parser = argparse.ArgumentParser()
parser.add_argument(
"--checkpoint-path",
type=resolve_path,
required=True,
help="Path to LTX-2 model checkpoint (.safetensors file).",
)
if distilled:
parser.add_argument(
"--distilled-checkpoint-path",
type=resolve_path,
required=True,
help="Path to LTX-2 distilled model checkpoint (.safetensors file).",
)
else:
parser.add_argument(
"--checkpoint-path",
type=resolve_path,
required=True,
help="Path to LTX-2 model checkpoint (.safetensors file).",
)
parser.add_argument(
"--gemma-root",
type=resolve_path,
@@ -142,58 +193,57 @@ def basic_arg_parser() -> argparse.ArgumentParser:
parser.add_argument(
"--seed",
type=int,
default=DEFAULT_SEED,
help=(
f"Random seed value used to initialize the noise tensor for "
f"reproducible generation (default: {DEFAULT_SEED})."
),
default=params.seed,
help=f"Random seed for reproducible generation (default: {params.seed}).",
)
parser.add_argument(
"--height",
type=int,
default=DEFAULT_1_STAGE_HEIGHT,
help=f"Height of the generated video in pixels, should be divisible by 32 (default: {DEFAULT_1_STAGE_HEIGHT}).",
default=params.stage_1_height,
help=f"Video height in pixels, divisible by 32 (default: {params.stage_1_height}).",
)
parser.add_argument(
"--width",
type=int,
default=DEFAULT_1_STAGE_WIDTH,
help=f"Width of the generated video in pixels, should be divisible by 32 (default: {DEFAULT_1_STAGE_WIDTH}).",
default=params.stage_1_width,
help=f"Width of the generated video in pixels, should be divisible by 32 (default: {params.stage_1_width}).",
)
parser.add_argument(
"--num-frames",
type=int,
default=DEFAULT_NUM_FRAMES,
default=params.num_frames,
help=f"Number of frames to generate in the output video sequence, num-frames = (8 x K) + 1, "
f"where k is a non-negative integer (default: {DEFAULT_NUM_FRAMES}).",
f"where k is a non-negative integer (default: {params.num_frames}).",
)
parser.add_argument(
"--frame-rate",
type=float,
default=DEFAULT_FRAME_RATE,
help=f"Frame rate of the generated video (fps) (default: {DEFAULT_FRAME_RATE}).",
default=params.frame_rate,
help=f"Frame rate of the generated video (fps) (default: {params.frame_rate}).",
)
parser.add_argument(
"--num-inference-steps",
type=int,
default=DEFAULT_NUM_INFERENCE_STEPS,
default=params.num_inference_steps,
help=(
f"Number of denoising steps in the diffusion sampling process. "
f"Higher values improve quality but increase generation time (default: {DEFAULT_NUM_INFERENCE_STEPS})."
f"Higher values improve quality but increase generation time (default: {params.num_inference_steps})."
),
)
parser.add_argument(
"--image",
dest="images",
action=ImageAction,
nargs=3,
metavar=("PATH", "FRAME_IDX", "STRENGTH"),
nargs="+",
metavar="ARG",
default=[],
help=(
"Image conditioning input: path to image file, target frame index, "
"and conditioning strength (all three required). Default: empty list [] (no image conditioning). "
"Image conditioning input: PATH FRAME_IDX STRENGTH [CRF]. "
"PATH is the image file, FRAME_IDX is the target frame index, "
"STRENGTH is the conditioning strength (all three required). "
f"CRF is the optional H.264 compression quality (0=lossless, default: {DEFAULT_IMAGE_CRF}). "
"Can be specified multiple times. Example: --image path/to/image1.jpg 0 0.8 "
"--image path/to/image2.jpg 160 0.9"
"--image path/to/image2.jpg 160 0.9 0"
),
)
parser.add_argument(
@@ -228,8 +278,10 @@ def basic_arg_parser() -> argparse.ArgumentParser:
return parser
def default_1_stage_arg_parser() -> argparse.ArgumentParser:
parser = basic_arg_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
parser = basic_arg_parser(params=params)
parser.add_argument(
"--negative-prompt",
type=str,
@@ -243,139 +295,139 @@ def default_1_stage_arg_parser() -> argparse.ArgumentParser:
parser.add_argument(
"--video-cfg-guidance-scale",
type=float,
default=DEFAULT_VIDEO_GUIDER_PARAMS.cfg_scale,
default=video_guider.cfg_scale,
help=(
f"Classifier-free guidance (CFG) scale controlling how strongly "
f"the model adheres to the video prompt. Higher values increase prompt "
"adherence but may reduce diversity. 1.0 means no effect "
f"(default: {DEFAULT_VIDEO_GUIDER_PARAMS.cfg_scale})."
f"adherence but may reduce diversity. 1.0 means no effect "
f"(default: {video_guider.cfg_scale})."
),
)
parser.add_argument(
"--video-stg-guidance-scale",
type=float,
default=DEFAULT_VIDEO_GUIDER_PARAMS.stg_scale,
default=video_guider.stg_scale,
help=(
f"STG (Spatio-Temporal Guidance) scale controlling how strongly "
f"the model reacts to the perturbation of the video modality. Higher values increase "
f"the effect but may reduce quality. 0.0 means no effect "
f"(default: {DEFAULT_VIDEO_GUIDER_PARAMS.stg_scale})."
f"(default: {video_guider.stg_scale})."
),
)
parser.add_argument(
"--video-rescale-scale",
type=float,
default=DEFAULT_VIDEO_GUIDER_PARAMS.rescale_scale,
default=video_guider.rescale_scale,
help=(
f"Rescale scale controlling how strongly "
f"the model rescales the video modality after applying other guidance. Higher values tend to decrease "
f"oversaturation effects. 0.0 means no effect (default: {DEFAULT_VIDEO_GUIDER_PARAMS.rescale_scale})."
f"oversaturation effects. 0.0 means no effect (default: {video_guider.rescale_scale})."
),
)
parser.add_argument(
"--video-stg-blocks",
type=int,
nargs="*",
default=DEFAULT_VIDEO_GUIDER_PARAMS.stg_blocks,
help=(f"Which transformer blocks to perturb for STG. Default: {DEFAULT_VIDEO_GUIDER_PARAMS.stg_blocks}."),
default=video_guider.stg_blocks,
help=(f"Which transformer blocks to perturb for STG. Default: {video_guider.stg_blocks}."),
)
parser.add_argument(
"--a2v-guidance-scale",
type=float,
default=DEFAULT_VIDEO_GUIDER_PARAMS.modality_scale,
default=video_guider.modality_scale,
help=(
f"A2V (Audio-to-Video) guidance scale controlling how strongly "
f"the model reacts to the perturbation of the audio-to-video cross-attention. Higher values may increase "
f"lipsync quality. 1.0 means no effect (default: {DEFAULT_VIDEO_GUIDER_PARAMS.modality_scale})."
f"lipsync quality. 1.0 means no effect (default: {video_guider.modality_scale})."
),
)
parser.add_argument(
"--video-skip-step",
type=int,
default=DEFAULT_VIDEO_GUIDER_PARAMS.skip_step,
default=video_guider.skip_step,
help=(
"Video skip step N controls periodic skipping during the video diffusion process: "
"only steps where step_index % (N + 1) == 0 are processed, all others are skipped "
f"(e.g., 0 = no skipping; 1 = skip every other step; 2 = skip 2 of every 3 steps; "
f"default: {DEFAULT_VIDEO_GUIDER_PARAMS.skip_step})."
f"default: {video_guider.skip_step})."
),
)
parser.add_argument(
"--audio-cfg-guidance-scale",
type=float,
default=DEFAULT_AUDIO_GUIDER_PARAMS.cfg_scale,
default=audio_guider.cfg_scale,
help=(
f"Audio CFG (Classifier-free guidance) scale controlling how strongly "
f"the model adheres to the audio prompt. Higher values increase prompt "
f"adherence but may reduce diversity. 1.0 means no effect "
f"(default: {DEFAULT_AUDIO_GUIDER_PARAMS.cfg_scale})."
f"(default: {audio_guider.cfg_scale})."
),
)
parser.add_argument(
"--audio-stg-guidance-scale",
type=float,
default=DEFAULT_AUDIO_GUIDER_PARAMS.stg_scale,
default=audio_guider.stg_scale,
help=(
f"Audio STG (Spatio-Temporal Guidance) scale controlling how strongly "
f"the model reacts to the perturbation of the audio modality. Higher values increase "
f"the effect but may reduce quality. 0.0 means no effect "
f"(default: {DEFAULT_AUDIO_GUIDER_PARAMS.stg_scale})."
f"(default: {audio_guider.stg_scale})."
),
)
parser.add_argument(
"--audio-rescale-scale",
type=float,
default=DEFAULT_AUDIO_GUIDER_PARAMS.rescale_scale,
default=audio_guider.rescale_scale,
help=(
f"Audio rescale scale controlling how strongly "
f"the model rescales the audio modality after applying other guidance. "
f"Experimental. 0.0 means no effect (default: {DEFAULT_AUDIO_GUIDER_PARAMS.rescale_scale})."
f"Experimental. 0.0 means no effect (default: {audio_guider.rescale_scale})."
),
)
parser.add_argument(
"--audio-stg-blocks",
type=int,
nargs="*",
default=DEFAULT_AUDIO_GUIDER_PARAMS.stg_blocks,
help=(f"Which transformer blocks to perturb for Audio STG. Default: {DEFAULT_AUDIO_GUIDER_PARAMS.stg_blocks}."),
default=audio_guider.stg_blocks,
help=(f"Which transformer blocks to perturb for Audio STG. Default: {audio_guider.stg_blocks}."),
)
parser.add_argument(
"--v2a-guidance-scale",
type=float,
default=DEFAULT_AUDIO_GUIDER_PARAMS.modality_scale,
default=audio_guider.modality_scale,
help=(
f"V2A (Video-to-Audio) guidance scale controlling how strongly "
f"the model reacts to the perturbation of the video-to-audio cross-attention. Higher values may increase "
f"lipsync quality. 1.0 means no effect (default: {DEFAULT_AUDIO_GUIDER_PARAMS.modality_scale})."
f"lipsync quality. 1.0 means no effect (default: {audio_guider.modality_scale})."
),
)
parser.add_argument(
"--audio-skip-step",
type=int,
default=DEFAULT_AUDIO_GUIDER_PARAMS.skip_step,
default=audio_guider.skip_step,
help=(
"Audio skip step N controls periodic skipping during the audio diffusion process: "
"only steps where step_index % (N + 1) == 0 are processed, all others are skipped "
f"(e.g., 0 = no skipping; 1 = skip every other step; 2 = skip 2 of every 3 steps; "
f"default: {DEFAULT_AUDIO_GUIDER_PARAMS.skip_step})."
f"default: {audio_guider.skip_step})."
),
)
return parser
def default_2_stage_arg_parser() -> argparse.ArgumentParser:
parser = default_1_stage_arg_parser()
parser.set_defaults(height=DEFAULT_2_STAGE_HEIGHT, width=DEFAULT_2_STAGE_WIDTH)
def default_2_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
parser = default_1_stage_arg_parser(params=params)
parser.set_defaults(height=params.stage_2_height, width=params.stage_2_width)
# Update help text to reflect 2-stage defaults
for action in parser._actions:
if "--height" in action.option_strings:
action.help = (
f"Height of the generated video in pixels, should be divisible by 64 "
f"(default: {DEFAULT_2_STAGE_HEIGHT})."
f"(default: {params.stage_2_height})."
)
if "--width" in action.option_strings:
action.help = (
f"Width of the generated video in pixels, should be divisible by 64 (default: {DEFAULT_2_STAGE_WIDTH})."
f"Width of the generated video in pixels, should be divisible by 64 (default: {params.stage_2_width})."
)
parser.add_argument(
"--distilled-lora",
@@ -405,19 +457,19 @@ def default_2_stage_arg_parser() -> argparse.ArgumentParser:
return parser
def default_2_stage_distilled_arg_parser() -> argparse.ArgumentParser:
parser = basic_arg_parser()
parser.set_defaults(height=DEFAULT_2_STAGE_HEIGHT, width=DEFAULT_2_STAGE_WIDTH)
def default_2_stage_distilled_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
parser = basic_arg_parser(params=params, distilled=True)
parser.set_defaults(height=params.stage_2_height, width=params.stage_2_width)
# Update help text to reflect 2-stage defaults
for action in parser._actions:
if "--height" in action.option_strings:
action.help = (
f"Height of the generated video in pixels, should be divisible by 64 "
f"(default: {DEFAULT_2_STAGE_HEIGHT})."
f"(default: {params.stage_2_height})."
)
if "--width" in action.option_strings:
action.help = (
f"Width of the generated video in pixels, should be divisible by 64 (default: {DEFAULT_2_STAGE_WIDTH})."
f"Width of the generated video in pixels, should be divisible by 64 (default: {params.stage_2_width})."
)
parser.add_argument(
"--spatial-upsampler-path",
packages/ltx-pipelines/src/ltx_pipelines/utils/constants.py
+80 -51
View File
@@ -1,12 +1,17 @@
import logging
from dataclasses import dataclass, field, replace
from safetensors import safe_open
from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_core.types import SpatioTemporalScaleFactors
# =============================================================================
# Diffusion Schedule
# =============================================================================
# Noise schedule for the distilled pipeline. These sigma values control noise
# levels at each denoising step and were tuned to match the distillation process.
from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_core.types import SpatioTemporalScaleFactors
DISTILLED_SIGMA_VALUES = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0]
# Reduced schedule for super-resolution stage 2 (subset of distilled values)
@@ -14,64 +19,88 @@ STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0]
# =============================================================================
# Video Generation Defaults
# Pipeline Parameters
# =============================================================================
DEFAULT_SEED = 10
DEFAULT_1_STAGE_HEIGHT = 512
DEFAULT_1_STAGE_WIDTH = 768
DEFAULT_2_STAGE_HEIGHT = DEFAULT_1_STAGE_HEIGHT * 2
DEFAULT_2_STAGE_WIDTH = DEFAULT_1_STAGE_WIDTH * 2
DEFAULT_NUM_FRAMES = 121
DEFAULT_FRAME_RATE = 24.0
DEFAULT_NUM_INFERENCE_STEPS = 40
DEFAULT_VIDEO_GUIDER_PARAMS = MultiModalGuiderParams(
cfg_scale=3.0,
stg_scale=1.0,
rescale_scale=0.7,
modality_scale=3.0,
skip_step=0,
stg_blocks=[29],
@dataclass(frozen=True)
class PipelineParams:
seed: int = 10
stage_1_height: int = 512
stage_1_width: int = 768
num_frames: int = 121
frame_rate: float = 24.0
num_inference_steps: int = 40
video_guider_params: MultiModalGuiderParams = field(
default_factory=lambda: MultiModalGuiderParams(
cfg_scale=3.0,
stg_scale=1.0,
rescale_scale=0.7,
modality_scale=3.0,
skip_step=0,
stg_blocks=[29],
)
)
audio_guider_params: MultiModalGuiderParams = field(
default_factory=lambda: MultiModalGuiderParams(
cfg_scale=7.0,
stg_scale=1.0,
rescale_scale=0.7,
modality_scale=3.0,
skip_step=0,
stg_blocks=[29],
)
)
@property
def stage_2_height(self) -> int:
return int(self.stage_1_height * 2)
@property
def stage_2_width(self) -> int:
return int(self.stage_1_width * 2)
# Default params for LTX-2.0 non-distilled models. These can be overridden by detecting from checkpoint metadata.
LTX_2_PARAMS = PipelineParams()
# Default params for LTX-2.3 non-distilled models. These override some of the LTX-2.0 defaults.
LTX_2_3_PARAMS = replace(
LTX_2_PARAMS,
num_inference_steps=30,
video_guider_params=replace(LTX_2_PARAMS.video_guider_params, stg_blocks=[28]),
audio_guider_params=replace(LTX_2_PARAMS.audio_guider_params, stg_blocks=[28]),
)
# =============================================================================
# Audio
# =============================================================================
DEFAULT_AUDIO_GUIDER_PARAMS = MultiModalGuiderParams(
cfg_scale=7.0,
stg_scale=1.0,
rescale_scale=0.7,
modality_scale=3.0,
skip_step=0,
stg_blocks=[29],
)
AUDIO_SAMPLE_RATE = 24000
# =============================================================================
# LoRA
# =============================================================================
DEFAULT_LORA_STRENGTH = 1.0
# =============================================================================
# Video VAE Architecture
# =============================================================================
DEFAULT_IMAGE_CRF = 33
VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
VIDEO_LATENT_CHANNELS = 128
_LTX_2_3_MODEL_VERSION_PREFIX = "2.3"
# =============================================================================
# Image Preprocessing
# =============================================================================
# CRF (Constant Rate Factor) for H.264 encoding used in image conditioning.
# Lower = higher quality, 0 = lossless. This mimics compression artifacts.
DEFAULT_IMAGE_CRF = 33
def detect_params(checkpoint_path: str) -> PipelineParams:
"""Detect pipeline params from checkpoint metadata.
Reads the ``model_version`` field from the safetensors config metadata.
Returns ``LTX_2_3_PARAMS`` when the version starts with "2.3",
otherwise falls back to ``LTX_2_PARAMS``.
"""
logger = logging.getLogger(__name__)
try:
with safe_open(checkpoint_path, framework="pt") as f:
metadata = f.metadata() or {}
version = metadata.get("model_version", "")
except Exception:
logger.warning("Could not read checkpoint metadata from %s, using LTX-2 defaults", checkpoint_path)
return LTX_2_PARAMS
if version.startswith(_LTX_2_3_MODEL_VERSION_PREFIX):
return LTX_2_3_PARAMS
logger.info("Using LTX_2_PARAMS for checkpoint (version=%s)", version or "unknown")
return LTX_2_PARAMS
# =============================================================================
packages/ltx-pipelines/src/ltx_pipelines/utils/helpers.py
+112 -128
View File
@@ -3,9 +3,8 @@ import logging
from dataclasses import replace
import torch
from tqdm import tqdm
from ltx_core.components.guiders import MultiModalGuider
from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderFactory
from ltx_core.components.noisers import Noiser
from ltx_core.components.protocols import DiffusionStepProtocol, GuiderProtocol
from ltx_core.conditioning import (
@@ -21,10 +20,10 @@ from ltx_core.guidance.perturbations import (
)
from ltx_core.model.transformer import Modality, X0Model
from ltx_core.model.video_vae import VideoEncoder
from ltx_core.text_encoders.gemma import GemmaTextEncoderModelBase
from ltx_core.text_encoders.gemma import GemmaTextEncoder
from ltx_core.tools import AudioLatentTools, LatentTools, VideoLatentTools
from ltx_core.types import AudioLatentShape, LatentState, VideoLatentShape, VideoPixelShape
from ltx_core.utils import to_denoised, to_velocity
from ltx_pipelines.utils.args import ImageConditioningInput
from ltx_pipelines.utils.media_io import decode_image, load_image_conditioning, resize_aspect_ratio_preserving
from ltx_pipelines.utils.types import (
DenoisingFunc,
@@ -46,7 +45,7 @@ def cleanup_memory() -> None:
def image_conditionings_by_replacing_latent(
images: list[tuple[str, int, float]],
images: list[ImageConditioningInput],
height: int,
width: int,
video_encoder: VideoEncoder,
@@ -54,20 +53,21 @@ def image_conditionings_by_replacing_latent(
device: torch.device,
) -> list[ConditioningItem]:
conditionings = []
for image_path, frame_idx, strength in images:
for img in images:
image = load_image_conditioning(
image_path=image_path,
image_path=img.path,
height=height,
width=width,
dtype=dtype,
device=device,
crf=img.crf,
)
encoded_image = video_encoder(image)
conditionings.append(
VideoConditionByLatentIndex(
latent=encoded_image,
strength=strength,
latent_idx=frame_idx,
strength=img.strength,
latent_idx=img.frame_idx,
)
)
@@ -75,7 +75,7 @@ def image_conditionings_by_replacing_latent(
def image_conditionings_by_adding_guiding_latent(
images: list[tuple[str, int, float]],
images: list[ImageConditioningInput],
height: int,
width: int,
video_encoder: VideoEncoder,
@@ -83,131 +83,22 @@ def image_conditionings_by_adding_guiding_latent(
device: torch.device,
) -> list[ConditioningItem]:
conditionings = []
for image_path, frame_idx, strength in images:
for img in images:
image = load_image_conditioning(
image_path=image_path,
image_path=img.path,
height=height,
width=width,
dtype=dtype,
device=device,
crf=img.crf,
)
encoded_image = video_encoder(image)
conditionings.append(
VideoConditionByKeyframeIndex(keyframes=encoded_image, frame_idx=frame_idx, strength=strength)
VideoConditionByKeyframeIndex(keyframes=encoded_image, frame_idx=img.frame_idx, strength=img.strength)
)
return conditionings
def euler_denoising_loop(
sigmas: torch.Tensor,
video_state: LatentState,
audio_state: LatentState,
stepper: DiffusionStepProtocol,
denoise_fn: DenoisingFunc,
) -> tuple[LatentState, LatentState]:
"""
Perform the joint audio-video denoising loop over a diffusion schedule.
This function iterates over all but the final value in ``sigmas`` and, at
each diffusion step, calls ``denoise_fn`` to obtain denoised video and
audio latents. The denoised latents are post-processed with their
respective denoise masks and clean latents, then passed to ``stepper`` to
advance the noisy latents one step along the diffusion schedule.
### Parameters
sigmas:
A 1D tensor of noise levels (diffusion sigmas) defining the sampling
schedule. All steps except the last element are iterated over.
video_state:
The current video :class:`LatentState`, containing the noisy latent,
its clean reference latent, and the denoising mask.
audio_state:
The current audio :class:`LatentState`, analogous to ``video_state``
but for the audio modality.
stepper:
An implementation of :class:`DiffusionStepProtocol` that updates a
latent given the current latent, its denoised estimate, the full
``sigmas`` schedule, and the current step index.
denoise_fn:
A callable implementing :class:`DenoisingFunc`. It is invoked as
``denoise_fn(video_state, audio_state, sigmas, step_index)`` and must
return a tuple ``(denoised_video, denoised_audio)``, where each element
is a tensor with the same shape as the corresponding latent.
### Returns
tuple[LatentState, LatentState]
A pair ``(video_state, audio_state)`` containing the final video and
audio latent states after completing the denoising loop.
"""
for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
denoised_video, denoised_audio = denoise_fn(video_state, audio_state, sigmas, step_idx)
denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
denoised_audio = post_process_latent(denoised_audio, audio_state.denoise_mask, audio_state.clean_latent)
video_state = replace(video_state, latent=stepper.step(video_state.latent, denoised_video, sigmas, step_idx))
audio_state = replace(audio_state, latent=stepper.step(audio_state.latent, denoised_audio, sigmas, step_idx))
return (video_state, audio_state)
def gradient_estimating_euler_denoising_loop(
sigmas: torch.Tensor,
video_state: LatentState,
audio_state: LatentState,
stepper: DiffusionStepProtocol,
denoise_fn: DenoisingFunc,
ge_gamma: float = 2.0,
) -> tuple[LatentState, LatentState]:
"""
Perform the joint audio-video denoising loop using gradient-estimation sampling.
This function is similar to :func:`euler_denoising_loop`, but applies
gradient estimation to improve the denoised estimates by tracking velocity
changes across steps. See the referenced function for detailed parameter
documentation.
### Parameters
ge_gamma:
Gradient estimation coefficient controlling the velocity correction term.
Default is 2.0. Paper: https://openreview.net/pdf?id=o2ND9v0CeK
sigmas, video_state, audio_state, stepper, denoise_fn:
See :func:`euler_denoising_loop` for parameter descriptions.
### Returns
tuple[LatentState, LatentState]
See :func:`euler_denoising_loop` for return value description.
"""
previous_audio_velocity = None
previous_video_velocity = None
def update_velocity_and_sample(
noisy_sample: torch.Tensor, denoised_sample: torch.Tensor, sigma: float, previous_velocity: torch.Tensor | None
) -> tuple[torch.Tensor, torch.Tensor]:
current_velocity = to_velocity(noisy_sample, sigma, denoised_sample)
if previous_velocity is not None:
delta_v = current_velocity - previous_velocity
total_velocity = ge_gamma * delta_v + previous_velocity
denoised_sample = to_denoised(noisy_sample, total_velocity, sigma)
return current_velocity, denoised_sample
for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
denoised_video, denoised_audio = denoise_fn(video_state, audio_state, sigmas, step_idx)
denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
denoised_audio = post_process_latent(denoised_audio, audio_state.denoise_mask, audio_state.clean_latent)
if sigmas[step_idx + 1] == 0:
return replace(video_state, latent=denoised_video), replace(audio_state, latent=denoised_audio)
previous_video_velocity, denoised_video = update_velocity_and_sample(
video_state.latent, denoised_video, sigmas[step_idx], previous_video_velocity
)
previous_audio_velocity, denoised_audio = update_velocity_and_sample(
audio_state.latent, denoised_audio, sigmas[step_idx], previous_audio_velocity
)
video_state = replace(video_state, latent=stepper.step(video_state.latent, denoised_video, sigmas, step_idx))
audio_state = replace(audio_state, latent=stepper.step(audio_state.latent, denoised_audio, sigmas, step_idx))
return (video_state, audio_state)
def noise_video_state(
output_shape: VideoPixelShape,
noiser: Noiser,
@@ -313,7 +204,10 @@ def post_process_latent(denoised: torch.Tensor, denoise_mask: torch.Tensor, clea
def modality_from_latent_state(
state: LatentState, context: torch.Tensor, sigma: float | torch.Tensor, enabled: bool = True
state: LatentState,
context: torch.Tensor,
sigma: torch.Tensor,
enabled: bool = True,
) -> Modality:
"""Create a Modality from a latent state.
Constructs a Modality object with the latent state's data, timesteps derived
@@ -322,10 +216,12 @@ def modality_from_latent_state(
return Modality(
enabled=enabled,
latent=state.latent,
sigma=sigma,
timesteps=timesteps_from_mask(state.denoise_mask, sigma),
positions=state.positions,
context=context,
context_mask=None,
attention_mask=state.attention_mask,
)
@@ -389,10 +285,10 @@ def multi_modal_guider_denoising_func(
v_context: torch.Tensor,
a_context: torch.Tensor,
transformer: X0Model,
*,
last_denoised_video: torch.Tensor | None = None,
last_denoised_audio: torch.Tensor | None = None,
) -> DenoisingFunc:
last_denoised_video = None
last_denoised_audio = None
def guider_denoising_step(
video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -490,6 +386,43 @@ def multi_modal_guider_denoising_func(
return guider_denoising_step
def multi_modal_guider_factory_denoising_func(
video_guider_factory: MultiModalGuiderFactory,
audio_guider_factory: MultiModalGuiderFactory | None,
v_context: torch.Tensor,
a_context: torch.Tensor,
transformer: X0Model,
) -> DenoisingFunc:
"""Resolve guiders per step via factory.build_from_sigma, then multi_modal_guider_denoising_func."""
last_denoised_video: torch.Tensor | None = None
last_denoised_audio: torch.Tensor | None = None
sigma_vals_cached: list[float] | None = None
def guider_denoising_step(
video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
) -> tuple[torch.Tensor, torch.Tensor]:
nonlocal last_denoised_video, last_denoised_audio, sigma_vals_cached
if sigma_vals_cached is None:
sigma_vals_cached = sigmas.detach().cpu().tolist()
sigma_val = sigma_vals_cached[step_index]
video_guider = video_guider_factory.build_from_sigma(sigma_val)
audio_guider = (audio_guider_factory or video_guider_factory).build_from_sigma(sigma_val)
denoise_fn = multi_modal_guider_denoising_func(
video_guider,
audio_guider,
v_context,
a_context,
transformer,
last_denoised_video=last_denoised_video,
last_denoised_audio=last_denoised_audio,
)
denoised_video, denoised_audio = denoise_fn(video_state, audio_state, sigmas, step_index)
last_denoised_video, last_denoised_audio = denoised_video, denoised_audio
return denoised_video, denoised_audio
return guider_denoising_step
def denoise_audio_video( # noqa: PLR0913
output_shape: VideoPixelShape,
conditionings: list[ConditioningItem],
@@ -540,6 +473,57 @@ def denoise_audio_video( # noqa: PLR0913
return video_state, audio_state
def denoise_video_only( # noqa: PLR0913
output_shape: VideoPixelShape,
conditionings: list[ConditioningItem],
noiser: Noiser,
sigmas: torch.Tensor,
stepper: DiffusionStepProtocol,
denoising_loop_fn: DenoisingLoopFunc,
components: PipelineComponents,
dtype: torch.dtype,
device: torch.device,
noise_scale: float = 1.0,
initial_video_latent: torch.Tensor | None = None,
initial_audio_latent: torch.Tensor | None = None,
) -> LatentState:
video_state, video_tools = noise_video_state(
output_shape=output_shape,
noiser=noiser,
conditionings=conditionings,
components=components,
dtype=dtype,
device=device,
noise_scale=noise_scale,
initial_latent=initial_video_latent,
)
audio_state, _ = noise_audio_state(
output_shape=output_shape,
noiser=noiser,
conditionings=[],
components=components,
dtype=dtype,
device=device,
noise_scale=0.0,
initial_latent=initial_audio_latent,
)
audio_state = replace(audio_state, denoise_mask=torch.zeros_like(audio_state.denoise_mask))
video_state, audio_state = denoising_loop_fn(
sigmas,
video_state,
audio_state,
stepper,
)
video_state = video_tools.clear_conditioning(video_state)
video_state = video_tools.unpatchify(video_state)
return video_state
_UNICODE_REPLACEMENTS = str.maketrans("\u2018\u2019\u201c\u201d\u2014\u2013\u00a0\u2032\u2212", "''\"\"-- '-")
@@ -555,7 +539,7 @@ def clean_response(text: str) -> str:
def generate_enhanced_prompt(
text_encoder: GemmaTextEncoderModelBase,
text_encoder: GemmaTextEncoder,
prompt: str,
image_path: str | None = None,
image_long_side: int = 896,
packages/ltx-pipelines/src/ltx_pipelines/utils/media_io.py
+111 -23
View File
@@ -12,6 +12,7 @@ from PIL import Image
from torch._prims_common import DeviceLikeType
from tqdm import tqdm
from ltx_core.types import Audio
from ltx_pipelines.utils.constants import DEFAULT_IMAGE_CRF
logger = logging.getLogger(__name__)
@@ -79,14 +80,19 @@ def normalize_latent(latent: torch.Tensor, device: torch.device, dtype: torch.dt
def load_image_conditioning(
image_path: str, height: int, width: int, dtype: torch.dtype, device: torch.device
image_path: str,
height: int,
width: int,
dtype: torch.dtype,
device: torch.device,
crf: int = DEFAULT_IMAGE_CRF,
) -> torch.Tensor:
"""
Loads an image from a path and preprocesses it for conditioning.
Note: The image is resized to the nearest multiple of 2 for compatibility with video codecs.
"""
image = decode_image(image_path=image_path)
image = preprocess(image=image)
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)
@@ -115,9 +121,8 @@ def decode_image(image_path: str) -> np.ndarray:
return np_array
def _write_audio(
container: av.container.Container, audio_stream: av.audio.AudioStream, samples: torch.Tensor, audio_sample_rate: int
) -> None:
def _write_audio(container: av.container.Container, audio_stream: av.audio.AudioStream, audio: Audio) -> None:
samples = audio.waveform
if samples.ndim == 1:
samples = samples[:, None]
@@ -137,7 +142,7 @@ def _write_audio(
format="s16",
layout="stereo",
)
frame_in.sample_rate = audio_sample_rate
frame_in.sample_rate = audio.sampling_rate
_resample_audio(container, audio_stream, frame_in)
@@ -185,8 +190,7 @@ def _resample_audio(
def encode_video(
video: torch.Tensor | Iterator[torch.Tensor],
fps: int,
audio: torch.Tensor | None,
audio_sample_rate: int | None,
audio: Audio | None,
output_path: str,
video_chunks_number: int,
) -> None:
@@ -204,10 +208,7 @@ def encode_video(
stream.pix_fmt = "yuv420p"
if audio is not None:
if audio_sample_rate is None:
raise ValueError("audio_sample_rate is required when audio is provided")
audio_stream = _prepare_audio_stream(container, audio_sample_rate)
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]
@@ -227,27 +228,114 @@ def encode_video(
container.mux(packet)
if audio is not None:
_write_audio(container, audio_stream, audio, audio_sample_rate)
_write_audio(container, audio_stream, audio)
container.close()
logger.info(f"Video saved to {output_path}")
def decode_audio_from_file(path: str, device: torch.device) -> torch.Tensor | None:
_INT_FORMAT_MAX: dict[str, float] = {
"u8": 128.0,
"u8p": 128.0,
"s16": 32768.0,
"s16p": 32768.0,
"s32": 2147483648.0,
"s32p": 2147483648.0,
}
def _audio_frame_to_float(frame: av.AudioFrame) -> np.ndarray:
"""Convert an audio frame to a float32 ndarray with values in [-1, 1] and shape (channels, samples)."""
fmt = frame.format.name
arr = frame.to_ndarray().astype(np.float32)
if fmt in _INT_FORMAT_MAX:
arr = arr / _INT_FORMAT_MAX[fmt]
if not frame.format.is_planar:
# Interleaved formats have shape (1, samples * channels) — reshape to (channels, samples).
channels = len(frame.layout.channels)
arr = arr.reshape(-1, channels).T
return arr
def get_videostream_metadata(path: str) -> tuple[float, int, int, int]:
"""Read video stream metadata: (fps, num_frames, width, height).
If frame count is missing in the container, decodes the stream to count frames.
"""
container = av.open(path)
try:
audio = []
audio_stream = next(s for s in container.streams if s.type == "audio")
for frame in container.decode(audio_stream):
audio.append(torch.tensor(frame.to_ndarray(), dtype=torch.float32, device=device).unsqueeze(0))
container.close()
audio = torch.cat(audio)
except StopIteration:
audio = None
video_stream = next(s for s in container.streams if s.type == "video")
fps = float(video_stream.average_rate)
num_frames = video_stream.frames or 0
if num_frames == 0:
num_frames = sum(1 for _ in container.decode(video_stream))
width = video_stream.codec_context.width
height = video_stream.codec_context.height
return fps, num_frames, width, height
finally:
container.close()
return audio
def decode_audio_from_file(
path: str, device: torch.device, start_time: float = 0.0, max_duration: float | None = None
) -> Audio | None:
"""Decodes audio from a file, optionally seeking to a start time and limiting duration.
Args:
path: Path to the audio/video file containing an audio stream.
device: Device to place the resulting tensor on.
start_time: Start time in seconds to begin reading audio from.
max_duration: Maximum audio duration in seconds. If None, reads to end of stream.
Returns:
An Audio object with waveform of shape (1, channels, samples), or None if no audio stream.
"""
container = av.open(path)
try:
audio_stream = next(s for s in container.streams if s.type == "audio")
except StopIteration:
container.close()
return None
sample_rate = audio_stream.rate
start_pts = int(start_time / audio_stream.time_base)
end_time = start_time + max_duration if max_duration else audio_stream.duration * audio_stream.time_base
container.seek(start_pts, stream=audio_stream)
samples = []
first_frame_time = None
for frame in container.decode(audio=0):
if frame.pts is None:
continue
frame_time = float(frame.pts * audio_stream.time_base)
frame_end = frame_time + frame.samples / frame.sample_rate
if frame_end < start_time:
continue
if frame_time > end_time:
break
if first_frame_time is None:
first_frame_time = frame_time
samples.append(_audio_frame_to_float(frame))
container.close()
if not samples:
return None
audio = np.concatenate(samples, axis=-1)
# Trim samples that fall outside the requested [start_time, start_time + max_duration] window.
# Audio codecs decode in fixed-size frames whose boundaries may not align with the requested
# time range, so the first frame can start before start_time and the last frame can end after
# start_time + max_duration.
skip_samples = round((start_time - first_frame_time) * sample_rate)
if skip_samples > 0:
audio = audio[..., skip_samples:]
if max_duration is not None:
max_samples = round(max_duration * sample_rate)
audio = audio[..., :max_samples]
waveform = torch.from_numpy(audio).to(device).unsqueeze(0)
return Audio(waveform=waveform, sampling_rate=sample_rate)
def decode_video_from_file(path: str, frame_cap: int, device: DeviceLikeType) -> Generator[torch.Tensor]:
packages/ltx-pipelines/src/ltx_pipelines/utils/model_ledger.py
+23 -5
View File
@@ -8,9 +8,12 @@ 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 (
AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
VOCODER_COMFY_KEYS_FILTER,
AudioDecoder,
AudioDecoderConfigurator,
AudioEncoder,
AudioEncoderConfigurator,
Vocoder,
VocoderConfigurator,
)
@@ -31,11 +34,11 @@ from ltx_core.model.video_vae import (
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import (
AV_GEMMA_TEXT_ENCODER_KEY_OPS,
AVGemmaTextEncoderModel,
AVGemmaTextEncoderModelConfigurator,
GEMMA_MODEL_OPS,
GemmaTextEncoder,
GemmaTextEncoderConfigurator,
module_ops_from_gemma_root,
)
from ltx_core.text_encoders.gemma.encoders.av_encoder import GEMMA_MODEL_OPS
from ltx_core.utils import find_matching_file
@@ -131,6 +134,13 @@ class ModelLedger:
registry=self.registry,
)
self.audio_encoder_builder = Builder[AudioEncoder](
model_path=self.checkpoint_path,
model_class_configurator=AudioEncoderConfigurator,
model_sd_ops=AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
registry=self.registry,
)
self.audio_decoder_builder = Builder(
model_path=self.checkpoint_path,
model_class_configurator=AudioDecoderConfigurator,
@@ -152,7 +162,7 @@ class ModelLedger:
self.text_encoder_builder = Builder(
model_path=(str(self.checkpoint_path), *weight_paths),
model_class_configurator=AVGemmaTextEncoderModelConfigurator,
model_class_configurator=GemmaTextEncoderConfigurator,
model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS,
registry=self.registry,
module_ops=(GEMMA_MODEL_OPS, *module_ops),
@@ -225,7 +235,7 @@ class ModelLedger:
return self.vae_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
def text_encoder(self) -> AVGemmaTextEncoderModel:
def text_encoder(self) -> GemmaTextEncoder:
if not hasattr(self, "text_encoder_builder"):
raise ValueError(
"Text encoder not initialized. Please provide a checkpoint path and gemma root path to the "
@@ -234,6 +244,14 @@ class ModelLedger:
return self.text_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
def audio_encoder(self) -> AudioEncoder:
if not hasattr(self, "audio_encoder_builder"):
raise ValueError(
"Audio encoder not initialized. Please provide a checkpoint path to the ModelLedger constructor."
)
return self.audio_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
def audio_decoder(self) -> AudioDecoder:
if not hasattr(self, "audio_decoder_builder"):
raise ValueError(
packages/ltx-pipelines/src/ltx_pipelines/utils/res2s.py
+62
View File
@@ -0,0 +1,62 @@
import math
def phi(j: int, neg_h: float) -> float:
"""
Compute φⱼ(z) where z = -h (negative step size in log-space)
φ₁(z) = (e^z - 1) / z
φ₂(z) = (e^z - 1 - z) / z²
φⱼ(z) = (e^z - Σₖ₌₀^(j-1) zᵏ/k!) / zʲ
These functions naturally appear when solving:
dx/dt = A*x + g(x,t) (linear drift + nonlinear part)
"""
if abs(neg_h) < 1e-10:
# Taylor series for small h to avoid division by zero
# φⱼ(0) = 1/j!
return 1.0 / math.factorial(j)
# Compute the "remainder" sum: Σₖ₌₀^(j-1) z^k/k!
remainder = sum(neg_h**k / math.factorial(k) for k in range(j))
# φⱼ(z) = (e^z - remainder) / z^j
return (math.exp(neg_h) - remainder) / (neg_h**j)
def get_res2s_coefficients(h: float, phi_cache: dict, c2: float = 0.5) -> tuple[float, float, float]:
"""
Compute res_2s Runge-Kutta coefficients for a given step size.
Args:
h: Step size in log-space = log(sigma / sigma_next)
phi_cache: Dictionary to cache phi function results. Cache key: (j, neg_h)
c2: Substep position (default 0.5 = midpoint)
Returns:
a21: Coefficient for computing intermediate x
b1, b2: Coefficients for final combination
"""
def get_phi(j: int, neg_h: float) -> float:
"""Get phi value with caching."""
cache_key = (j, neg_h)
if cache_key in phi_cache:
return phi_cache[cache_key]
result = phi(j, neg_h)
phi_cache[cache_key] = result
return result
# Substep coefficient: how much of ε₁ to use for intermediate point
# a21 = c2 * φ₁(-h*c2)
neg_h_c2 = -h * c2
phi_1_c2 = get_phi(1, neg_h_c2)
a21 = c2 * phi_1_c2
# Final combination weights
# b2 = φ₂(-h) / c2
neg_h_full = -h
phi_2_full = get_phi(2, neg_h_full)
b2 = phi_2_full / c2
# b1 = φ₁(-h) - b2
phi_1_full = get_phi(1, neg_h_full)
b1 = phi_1_full - b2
return a21, b1, b2
packages/ltx-pipelines/src/ltx_pipelines/utils/samplers.py
+363
View File
@@ -0,0 +1,363 @@
import logging
from dataclasses import replace
from functools import partial
from typing import Callable
import torch
from tqdm import tqdm
from ltx_core.components.diffusion_steps import Res2sDiffusionStep
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.utils import to_denoised, to_velocity
from ltx_pipelines.utils.helpers import post_process_latent, timesteps_from_mask
from ltx_pipelines.utils.res2s import get_res2s_coefficients
from ltx_pipelines.utils.types import DenoisingFunc, LatentState
logger = logging.getLogger(__name__)
def euler_denoising_loop(
sigmas: torch.Tensor,
video_state: LatentState,
audio_state: LatentState,
stepper: DiffusionStepProtocol,
denoise_fn: DenoisingFunc,
) -> tuple[LatentState, LatentState]:
"""
Perform the joint audio-video denoising loop over a diffusion schedule.
This function iterates over all but the final value in ``sigmas`` and, at
each diffusion step, calls ``denoise_fn`` to obtain denoised video and
audio latents. The denoised latents are post-processed with their
respective denoise masks and clean latents, then passed to ``stepper`` to
advance the noisy latents one step along the diffusion schedule.
### Parameters
sigmas:
A 1D tensor of noise levels (diffusion sigmas) defining the sampling
schedule. All steps except the last element are iterated over.
video_state:
The current video :class:`LatentState`, containing the noisy latent,
its clean reference latent, and the denoising mask.
audio_state:
The current audio :class:`LatentState`, analogous to ``video_state``
but for the audio modality.
stepper:
An implementation of :class:`DiffusionStepProtocol` that updates a
latent given the current latent, its denoised estimate, the full
``sigmas`` schedule, and the current step index.
denoise_fn:
A callable implementing :class:`DenoisingFunc`. It is invoked as
``denoise_fn(video_state, audio_state, sigmas, step_index)`` and must
return a tuple ``(denoised_video, denoised_audio)``, where each element
is a tensor with the same shape as the corresponding latent.
### Returns
tuple[LatentState, LatentState]
A pair ``(video_state, audio_state)`` containing the final video and
audio latent states after completing the denoising loop.
"""
for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
denoised_video, denoised_audio = denoise_fn(video_state, audio_state, sigmas, step_idx)
denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
denoised_audio = post_process_latent(denoised_audio, audio_state.denoise_mask, audio_state.clean_latent)
video_state = replace(video_state, latent=stepper.step(video_state.latent, denoised_video, sigmas, step_idx))
audio_state = replace(audio_state, latent=stepper.step(audio_state.latent, denoised_audio, sigmas, step_idx))
return (video_state, audio_state)
def gradient_estimating_euler_denoising_loop(
sigmas: torch.Tensor,
video_state: LatentState,
audio_state: LatentState,
stepper: DiffusionStepProtocol,
denoise_fn: DenoisingFunc,
ge_gamma: float = 2.0,
) -> tuple[LatentState, LatentState]:
"""
Perform the joint audio-video denoising loop using gradient-estimation sampling.
This function is similar to :func:`euler_denoising_loop`, but applies
gradient estimation to improve the denoised estimates by tracking velocity
changes across steps. See the referenced function for detailed parameter
documentation.
### Parameters
ge_gamma:
Gradient estimation coefficient controlling the velocity correction term.
Default is 2.0. Paper: https://openreview.net/pdf?id=o2ND9v0CeK
sigmas, video_state, audio_state, stepper, denoise_fn:
See :func:`euler_denoising_loop` for parameter descriptions.
### Returns
tuple[LatentState, LatentState]
See :func:`euler_denoising_loop` for return value description.
"""
previous_audio_velocity = None
previous_video_velocity = None
def update_velocity_and_sample(
noisy_sample: torch.Tensor, denoised_sample: torch.Tensor, sigma: float, previous_velocity: torch.Tensor | None
) -> tuple[torch.Tensor, torch.Tensor]:
current_velocity = to_velocity(noisy_sample, sigma, denoised_sample)
if previous_velocity is not None:
delta_v = current_velocity - previous_velocity
total_velocity = ge_gamma * delta_v + previous_velocity
denoised_sample = to_denoised(noisy_sample, total_velocity, sigma)
return current_velocity, denoised_sample
for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
denoised_video, denoised_audio = denoise_fn(video_state, audio_state, sigmas, step_idx)
denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
denoised_audio = post_process_latent(denoised_audio, audio_state.denoise_mask, audio_state.clean_latent)
if sigmas[step_idx + 1] == 0:
return replace(video_state, latent=denoised_video), replace(audio_state, latent=denoised_audio)
previous_video_velocity, denoised_video = update_velocity_and_sample(
video_state.latent, denoised_video, sigmas[step_idx], previous_video_velocity
)
previous_audio_velocity, denoised_audio = update_velocity_and_sample(
audio_state.latent, denoised_audio, sigmas[step_idx], previous_audio_velocity
)
video_state = replace(video_state, latent=stepper.step(video_state.latent, denoised_video, sigmas, step_idx))
audio_state = replace(audio_state, latent=stepper.step(audio_state.latent, denoised_audio, sigmas, step_idx))
return (video_state, audio_state)
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))
def _get_new_noise(x: torch.Tensor, generator: torch.Generator) -> torch.Tensor:
noise = torch.randn(x.shape, generator=generator, dtype=torch.float64, device=generator.device)
noise = (noise - noise.mean()) / noise.std()
return _channelwise_normalize(noise)
def _inject_sde_noise(
state: LatentState,
sample: torch.Tensor,
denoised_sample: torch.Tensor,
step_noise_generator: torch.Generator,
new_noise_fn: Callable[[torch.Tensor, torch.Generator], torch.Tensor],
stepper: DiffusionStepProtocol,
sigmas: torch.Tensor,
step_idx: int,
legacy_mode: bool = False,
) -> torch.Tensor:
sigmas_copy = sigmas.clone()
new_noise = new_noise_fn(state.latent, step_noise_generator)
if not legacy_mode:
timesteps = timesteps_from_mask(state.denoise_mask.double(), sigmas_copy[step_idx].double())
next_timesteps = timesteps_from_mask(state.denoise_mask.double(), sigmas_copy[step_idx + 1].double())
sigmas = torch.stack([timesteps, next_timesteps])
step_idx = 0
x_next = stepper.step(
sample=sample,
denoised_sample=denoised_sample,
sigmas=sigmas,
step_index=step_idx,
noise=new_noise,
)
if legacy_mode:
x_next = post_process_latent(x_next, state.denoise_mask, state.clean_latent)
return x_next
def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915
sigmas: torch.Tensor,
video_state: LatentState,
audio_state: LatentState,
stepper: DiffusionStepProtocol,
denoise_fn: DenoisingFunc,
noise_seed: int = -1,
noise_seed_substep: int | None = None,
bongmath: bool = True,
bongmath_max_iter: int = 100,
new_noise_fn: Callable[[torch.Tensor, torch.Generator], torch.Tensor] = _get_new_noise,
model_dtype: torch.dtype = torch.bfloat16,
legacy_mode: bool = True,
) -> tuple[LatentState, LatentState]:
"""
Joint audio-video denoising loop using the res_2s second-order sampler.
Iterates over the diffusion schedule with a two-stage Runge-Kutta step:
evaluates the denoiser at the current point and at a midpoint (with SDE
noise), then combines both with RK coefficients. Supports anchor-point
refinement (bong iteration) and optional SDE noise injection. Requires
:class:`Res2sDiffusionStep` as ``stepper``.
### Parameters
sigmas:
A 1D tensor of noise levels defining the sampling schedule.
video_state:
Current video :class:`LatentState` (noisy latent, clean reference, mask).
audio_state:
Current audio :class:`LatentState`, same structure as ``video_state``.
stepper:
Must be an instance of :class:`Res2sDiffusionStep`; performs SDE step
with noise injection.
denoise_fn:
Callable ``(video_state, audio_state, sigmas, step_index)`` returning
``(denoised_video, denoised_audio)``.
noise_seed:
Seed for step-level SDE noise; substep seed defaults to ``noise_seed + 10000``.
noise_seed_substep:
Optional seed for substep SDE noise; if None, derived from ``noise_seed``.
bongmath:
Whether to run iterative anchor refinement (bong iteration) when step size is small.
bongmath_max_iter:
Max iterations for bong refinement when enabled.
new_noise_fn:
Callable ``(latent, generator) -> noise`` for SDE injection; default
uses normalized channel-wise Gaussian noise.
model_dtype:
Dtype for latent state updates (e.g. bfloat16).
### Returns
tuple[LatentState, LatentState]
Final ``(video_state, audio_state)`` after the denoising loop.
"""
# Initialize noise generators with different seeds
if noise_seed_substep is None:
noise_seed_substep = noise_seed + 10000 # Offset to ensure different seeds
step_noise_generator = torch.Generator(device=video_state.latent.device).manual_seed(noise_seed)
substep_noise_generator = torch.Generator(device=video_state.latent.device).manual_seed(noise_seed_substep)
sde_noise_injecting_fn = partial(
_inject_sde_noise, stepper=stepper, new_noise_fn=new_noise_fn, legacy_mode=legacy_mode
)
step_noise_injecting_fn = partial(sde_noise_injecting_fn, step_noise_generator=step_noise_generator)
substep_noise_injecting_fn = partial(sde_noise_injecting_fn, step_noise_generator=substep_noise_generator)
if not isinstance(stepper, Res2sDiffusionStep):
raise ValueError("stepper must be an instance of Res2sDiffusionStep")
n_full_steps = len(sigmas) - 1
# inject minimal sigma value to avoid division by zero
if sigmas[-1] == 0:
sigmas = torch.cat([sigmas[:-1], torch.tensor([0.0011, 0.0], device=sigmas.device)], dim=0)
# Compute step sizes in hyperbolic space
hs = -torch.log(sigmas[1:].double().cpu() / (sigmas[:-1].double().cpu()))
# Initialize phi cache for reuse across loop iterations
# Cache key: (j, neg_h) where j is phi order and neg_h is negative step value
phi_cache = {}
c2 = 0.5 # Midpoint for res_2s
# Progress bar shows only full two-stage steps; final (sigma_next==0) step is done silently
for step_idx in tqdm(range(n_full_steps)):
sigma = sigmas[step_idx].double()
sigma_next = sigmas[step_idx + 1].double()
# Initialize anchor point
x_anchor_video = video_state.latent.clone().double()
x_anchor_audio = audio_state.latent.clone().double()
# ====================================================================
# STAGE 1: Evaluate at current point
# ====================================================================
denoised_video_1, denoised_audio_1 = denoise_fn(video_state, audio_state, sigmas, step_idx)
denoised_video_1 = post_process_latent(denoised_video_1, video_state.denoise_mask, video_state.clean_latent)
denoised_audio_1 = post_process_latent(denoised_audio_1, audio_state.denoise_mask, audio_state.clean_latent)
h = hs[step_idx].item()
# Compute RK coefficients (pass phi_cache for caching)
a21, b1, b2 = get_res2s_coefficients(h, phi_cache, c2)
# Compute substep sigma, sqrt is a hardcode for c2 = 0.5
sub_sigma = torch.sqrt(sigma * sigma_next)
# ====================================================================
# Compute substep x using RK coefficient a21
# ====================================================================
eps_1_video = denoised_video_1.double() - x_anchor_video
eps_1_audio = denoised_audio_1.double() - x_anchor_audio
x_mid_video = x_anchor_video.double() + h * a21 * eps_1_video
x_mid_audio = x_anchor_audio.double() + h * a21 * eps_1_audio
# ====================================================================
# SDE noise injection at substep
# ====================================================================
x_mid_video = substep_noise_injecting_fn(
state=video_state,
sample=x_anchor_video,
denoised_sample=x_mid_video,
sigmas=torch.stack([sigma, sub_sigma]),
step_idx=0,
)
x_mid_audio = substep_noise_injecting_fn(
state=audio_state,
sample=x_anchor_audio,
denoised_sample=x_mid_audio,
sigmas=torch.stack([sigma, sub_sigma]),
step_idx=0,
)
# ====================================================================
# ITERATIVE REFINEMENT (Bong Iteration) - Stabilize anchor point
# ====================================================================
if bongmath and h < 0.5 and sigma > 0.03:
for _ in range(bongmath_max_iter):
x_anchor_video = x_mid_video - h * a21 * eps_1_video
eps_1_video = denoised_video_1.double() - x_anchor_video
x_anchor_audio = x_mid_audio - h * a21 * eps_1_audio
eps_1_audio = denoised_audio_1.double() - x_anchor_audio
# ====================================================================
# STAGE 2: Evaluate at substep point (WITH NOISE)
# ====================================================================
mid_video_state = replace(video_state, latent=x_mid_video.to(model_dtype))
mid_audio_state = replace(audio_state, latent=x_mid_audio.to(model_dtype))
denoised_video_2, denoised_audio_2 = denoise_fn(
video_state=mid_video_state,
audio_state=mid_audio_state,
sigmas=torch.stack([sub_sigma]).to(sigmas.device),
step_index=0,
)
denoised_video_2 = post_process_latent(denoised_video_2, video_state.denoise_mask, video_state.clean_latent)
denoised_audio_2 = post_process_latent(denoised_audio_2, audio_state.denoise_mask, audio_state.clean_latent)
# ====================================================================
# FINAL COMBINATION: Compute x_next using RK coefficients
# ====================================================================
eps_2_video = denoised_video_2.double() - x_anchor_video
eps_2_audio = denoised_audio_2.double() - x_anchor_audio
x_next_video = x_anchor_video + h * (b1 * eps_1_video + b2 * eps_2_video)
x_next_audio = x_anchor_audio + h * (b1 * eps_1_audio + b2 * eps_2_audio)
# ====================================================================
# SDE NOISE INJECTION AT STEP LEVEL
# ====================================================================
x_next_video = step_noise_injecting_fn(
state=video_state,
sample=x_anchor_video,
denoised_sample=x_next_video,
sigmas=sigmas,
step_idx=step_idx,
)
x_next_audio = step_noise_injecting_fn(
state=audio_state,
sample=x_anchor_audio,
denoised_sample=x_next_audio,
sigmas=sigmas,
step_idx=step_idx,
)
# Update states
video_state = replace(video_state, latent=x_next_video.to(model_dtype))
audio_state = replace(audio_state, latent=x_next_audio.to(model_dtype))
# Final step if we need to fully remove the noise
if sigmas[-1] == 0:
denoised_video_1, denoised_audio_1 = denoise_fn(video_state, audio_state, sigmas, n_full_steps)
denoised_video_1 = post_process_latent(denoised_video_1, video_state.denoise_mask, video_state.clean_latent)
denoised_audio_1 = post_process_latent(denoised_audio_1, audio_state.denoise_mask, audio_state.clean_latent)
video_state = replace(video_state, latent=denoised_video_1.to(model_dtype))
audio_state = replace(audio_state, latent=denoised_audio_1.to(model_dtype))
return video_state, audio_state