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/__init__.py
+5
View File
@@ -6,21 +6,26 @@ This package provides ready-to-use pipelines for video generation:
- DistilledPipeline: Fast distilled two-stage generation
- ICLoraPipeline: Image/video conditioning with distilled LoRA
- KeyframeInterpolationPipeline: Keyframe-based video interpolation
- RetakePipeline: Regenerate a time region (retake) of an existing video
- ModelLedger: Central coordinator for loading and building models
For more detailed components and utilities, import from specific submodules
like `ltx_pipelines.utils.media_io` or `ltx_pipelines.utils.constants`.
"""
from ltx_pipelines.a2vid_two_stage import A2VidPipelineTwoStage
from ltx_pipelines.distilled import DistilledPipeline
from ltx_pipelines.ic_lora import ICLoraPipeline
from ltx_pipelines.keyframe_interpolation import KeyframeInterpolationPipeline
from ltx_pipelines.retake import RetakePipeline
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
__all__ = [
"A2VidPipelineTwoStage",
"DistilledPipeline",
"ICLoraPipeline",
"KeyframeInterpolationPipeline",
"RetakePipeline",
"TI2VidOneStagePipeline",
"TI2VidTwoStagesPipeline",
]
packages/ltx-pipelines/src/ltx_pipelines/a2vid_two_stage.py
+323
View File
@@ -0,0 +1,323 @@
import logging
from collections.abc import Iterator
import torch
from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.components.schedulers import LTX2Scheduler
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import Audio, AudioLatentShape, LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.args import default_2_stage_arg_parser
from ltx_pipelines.utils.constants import (
STAGE_2_DISTILLED_SIGMA_VALUES,
)
from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
denoise_video_only,
generate_enhanced_prompt,
get_device,
image_conditionings_by_replacing_latent,
multi_modal_guider_denoising_func,
simple_denoising_func,
)
from ltx_pipelines.utils.media_io import decode_audio_from_file, encode_video
from ltx_pipelines.utils.samplers import euler_denoising_loop
from ltx_pipelines.utils.types import PipelineComponents
device = get_device()
class A2VidPipelineTwoStage:
"""
Two-stage audio to video generation pipeline.
Stage 1 generates video at half the target resolution with audio conditioning
(video-only denoising, audio frozen), then Stage 2 upsamples by 2x and refines
both video and audio using a distilled LoRA for higher quality output.
"""
def __init__(
self,
checkpoint_path: str,
distilled_lora: list[LoraPathStrengthAndSDOps],
spatial_upsampler_path: str,
gemma_root: str,
loras: list[LoraPathStrengthAndSDOps],
device: torch.device = device,
quantization: QuantizationPolicy | None = None,
):
self.device = device
self.dtype = torch.bfloat16
self.stage_1_model_ledger = ModelLedger(
dtype=self.dtype,
device=device,
checkpoint_path=checkpoint_path,
gemma_root_path=gemma_root,
spatial_upsampler_path=spatial_upsampler_path,
loras=loras,
quantization=quantization,
)
self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
loras=distilled_lora,
)
self.pipeline_components = PipelineComponents(
dtype=self.dtype,
device=device,
)
def __call__( # noqa: PLR0913
self,
prompt: str,
negative_prompt: str,
seed: int,
height: int,
width: int,
num_frames: int,
frame_rate: float,
num_inference_steps: int,
video_guider_params: MultiModalGuiderParams,
images: list[tuple[str, int, float]],
audio_path: str,
audio_start_time: float = 0.0,
audio_max_duration: float | None = None,
tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False,
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=True)
generator = torch.Generator(device=self.device).manual_seed(seed)
noiser = GaussianNoiser(generator=generator)
stepper = EulerDiffusionStep()
dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(text_encoder, prompt, images[0][0] if len(images) > 0 else None)
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, _ = context_n
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
# Encode audio.
decoded_audio = decode_audio_from_file(audio_path, self.device, audio_start_time, audio_max_duration)
encoded_audio_latent = vae_encode_audio(decoded_audio, self.stage_1_model_ledger.audio_encoder())
audio_shape = AudioLatentShape.from_duration(batch=1, duration=num_frames / frame_rate, channels=8, mel_bins=16)
encoded_audio_latent = encoded_audio_latent[:, :, : audio_shape.frames]
cleanup_memory()
# Stage 1: Initial low resolution video generation with audio conditioning.
transformer = self.stage_1_model_ledger.transformer()
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
def first_stage_denoising_loop(
sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
) -> tuple[LatentState, LatentState]:
return euler_denoising_loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=multi_modal_guider_denoising_func(
video_guider=MultiModalGuider(
params=video_guider_params,
negative_context=v_context_n,
),
audio_guider=MultiModalGuider(
params=MultiModalGuiderParams(),
),
v_context=v_context_p,
a_context=a_context_p,
transformer=transformer, # noqa: F821
),
)
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
video_encoder = self.stage_1_model_ledger.video_encoder()
stage_1_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state = denoise_video_only(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
noiser=noiser,
sigmas=sigmas,
stepper=stepper,
denoising_loop_fn=first_stage_denoising_loop,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
initial_audio_latent=encoded_audio_latent,
)
torch.cuda.synchronize()
del transformer
cleanup_memory()
# Stage 2: Upsample and refine the video at higher resolution with distilled LoRA.
upscaled_video_latent = upsample_video(
latent=video_state.latent[:1],
video_encoder=video_encoder,
upsampler=self.stage_2_model_ledger.spatial_upsampler(),
)
torch.cuda.synchronize()
cleanup_memory()
transformer = self.stage_2_model_ledger.transformer()
distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
def second_stage_denoising_loop(
sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
) -> tuple[LatentState, LatentState]:
return euler_denoising_loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=simple_denoising_func(
video_context=v_context_p,
audio_context=a_context_p,
transformer=transformer, # noqa: F821
),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state = denoise_video_only(
output_shape=stage_2_output_shape,
conditionings=stage_2_conditionings,
noiser=noiser,
sigmas=distilled_sigmas,
stepper=stepper,
denoising_loop_fn=second_stage_denoising_loop,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
noise_scale=distilled_sigmas[0],
initial_video_latent=upscaled_video_latent,
initial_audio_latent=encoded_audio_latent,
)
torch.cuda.synchronize()
del transformer
del video_encoder
cleanup_memory()
decoded_video = vae_decode_video(
video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config, generator
)
# Return the original input audio instead of VAE-decoded audio to preserve fidelity.
# decode_audio_from_file already returns normalised [-1, 1] float values.
original_audio = Audio(waveform=decoded_audio.waveform.squeeze(0), sampling_rate=decoded_audio.sampling_rate)
return decoded_video, original_audio
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
parser = default_2_stage_arg_parser()
parser.add_argument(
"--audio-path",
type=str,
required=True,
help="Path to the audio file to condition the video generation.",
)
parser.add_argument(
"--audio-start-time",
type=float,
default=0.0,
help="Start time in seconds to read audio from (default: 0.0).",
)
parser.add_argument(
"--audio-max-duration",
type=float,
default=None,
help="Maximum audio duration in seconds. Defaults to video duration (num_frames / frame_rate).",
)
args = parser.parse_args()
pipeline = A2VidPipelineTwoStage(
checkpoint_path=args.checkpoint_path,
distilled_lora=args.distilled_lora,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
quantization=args.quantization,
)
tiling_config = TilingConfig.default()
video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
video, audio = pipeline(
prompt=args.prompt,
negative_prompt=args.negative_prompt,
seed=args.seed,
height=args.height,
width=args.width,
num_frames=args.num_frames,
frame_rate=args.frame_rate,
num_inference_steps=args.num_inference_steps,
video_guider_params=MultiModalGuiderParams(
cfg_scale=args.video_cfg_guidance_scale,
stg_scale=args.video_stg_guidance_scale,
rescale_scale=args.video_rescale_scale,
modality_scale=args.a2v_guidance_scale,
skip_step=args.video_skip_step,
stg_blocks=args.video_stg_blocks,
),
images=args.images,
tiling_config=tiling_config,
enhance_prompt=args.enhance_prompt,
audio_path=args.audio_path,
audio_start_time=args.audio_start_time,
audio_max_duration=args.audio_max_duration
if args.audio_max_duration is not None
else args.num_frames / args.frame_rate,
)
encode_video(
video=video,
fps=args.frame_rate,
audio=audio,
output_path=args.output_path,
video_chunks_number=video_chunks_number,
)
if __name__ == "__main__":
main()
packages/ltx-pipelines/src/ltx_pipelines/distilled.py
+17 -13
View File
@@ -13,19 +13,22 @@ from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.args import default_2_stage_distilled_arg_parser
from ltx_core.types import Audio, LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger, euler_denoising_loop
from ltx_pipelines.utils.args import (
ImageConditioningInput,
default_2_stage_distilled_arg_parser,
detect_checkpoint_path,
)
from ltx_pipelines.utils.constants import (
AUDIO_SAMPLE_RATE,
DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMA_VALUES,
detect_params,
)
from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
denoise_audio_video,
euler_denoising_loop,
generate_enhanced_prompt,
get_device,
image_conditionings_by_replacing_latent,
@@ -40,13 +43,13 @@ device = get_device()
class DistilledPipeline:
"""
Two-stage distilled video generation pipeline.
Stage 1 generates video at the target resolution, then Stage 2 upsamples
Stage 1 generates video at half of the target resolution, then Stage 2 upsamples
by 2x and refines with additional denoising steps for higher quality output.
"""
def __init__(
self,
checkpoint_path: str,
distilled_checkpoint_path: str,
gemma_root: str,
spatial_upsampler_path: str,
loras: list[LoraPathStrengthAndSDOps],
@@ -59,7 +62,7 @@ class DistilledPipeline:
self.model_ledger = ModelLedger(
dtype=self.dtype,
device=device,
checkpoint_path=checkpoint_path,
checkpoint_path=distilled_checkpoint_path,
spatial_upsampler_path=spatial_upsampler_path,
gemma_root_path=gemma_root,
loras=loras,
@@ -79,10 +82,10 @@ class DistilledPipeline:
width: int,
num_frames: int,
frame_rate: float,
images: list[tuple[str, int, float]],
images: list[ImageConditioningInput],
tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=True)
generator = torch.Generator(device=self.device).manual_seed(seed)
@@ -198,10 +201,12 @@ class DistilledPipeline:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
parser = default_2_stage_distilled_arg_parser()
checkpoint_path = detect_checkpoint_path(distilled=True)
params = detect_params(checkpoint_path)
parser = default_2_stage_distilled_arg_parser(params=params)
args = parser.parse_args()
pipeline = DistilledPipeline(
checkpoint_path=args.checkpoint_path,
distilled_checkpoint_path=args.distilled_checkpoint_path,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
@@ -225,7 +230,6 @@ def main() -> None:
video=video,
fps=args.frame_rate,
audio=audio,
audio_sample_rate=AUDIO_SAMPLE_RATE,
output_path=args.output_path,
video_chunks_number=video_chunks_number,
)
packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py
+258 -27
View File
@@ -2,12 +2,17 @@ import logging
from collections.abc import Iterator
import torch
from einops import rearrange
from safetensors import safe_open
from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.conditioning import ConditioningItem, VideoConditionByReferenceLatent
from ltx_core.conditioning import (
ConditioningItem,
ConditioningItemAttentionStrengthWrapper,
VideoConditionByReferenceLatent,
)
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
from ltx_core.model.upsampler import upsample_video
@@ -15,15 +20,9 @@ from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunk
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.args import VideoConditioningAction, default_2_stage_distilled_arg_parser
from ltx_pipelines.utils.constants import (
AUDIO_SAMPLE_RATE,
DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMA_VALUES,
)
from ltx_pipelines.utils.helpers import (
from ltx_core.types import Audio, LatentState, VideoLatentShape, VideoPixelShape
from ltx_pipelines.utils import (
ModelLedger,
assert_resolution,
cleanup_memory,
denoise_audio_video,
@@ -33,6 +32,18 @@ from ltx_pipelines.utils.helpers import (
image_conditionings_by_replacing_latent,
simple_denoising_func,
)
from ltx_pipelines.utils.args import (
ImageConditioningInput,
VideoConditioningAction,
VideoMaskConditioningAction,
default_2_stage_distilled_arg_parser,
detect_checkpoint_path,
)
from ltx_pipelines.utils.constants import (
DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMA_VALUES,
detect_params,
)
from ltx_pipelines.utils.media_io import encode_video, load_video_conditioning
from ltx_pipelines.utils.types import PipelineComponents
@@ -45,13 +56,14 @@ class ICLoraPipeline:
Allows conditioning the generated video on control signals such as depth maps,
human pose, or image edges via the video_conditioning parameter.
The specific IC-LoRA model should be provided via the loras parameter.
Stage 1 generates video at the target resolution, then Stage 2 upsamples
Stage 1 generates video at half of the target resolution, then Stage 2 upsamples
by 2x and refines with additional denoising steps for higher quality output.
Both stages use distilled models for efficiency.
"""
def __init__(
self,
checkpoint_path: str,
distilled_checkpoint_path: str,
spatial_upsampler_path: str,
gemma_root: str,
loras: list[LoraPathStrengthAndSDOps],
@@ -62,7 +74,7 @@ class ICLoraPipeline:
self.stage_1_model_ledger = ModelLedger(
dtype=self.dtype,
device=device,
checkpoint_path=checkpoint_path,
checkpoint_path=distilled_checkpoint_path,
spatial_upsampler_path=spatial_upsampler_path,
gemma_root_path=gemma_root,
loras=loras,
@@ -71,7 +83,7 @@ class ICLoraPipeline:
self.stage_2_model_ledger = ModelLedger(
dtype=self.dtype,
device=device,
checkpoint_path=checkpoint_path,
checkpoint_path=distilled_checkpoint_path,
spatial_upsampler_path=spatial_upsampler_path,
gemma_root_path=gemma_root,
loras=[],
@@ -98,8 +110,7 @@ class ICLoraPipeline:
)
self.reference_downscale_factor = scale
@torch.inference_mode()
def __call__(
def __call__( # noqa: PLR0913
self,
prompt: str,
seed: int,
@@ -107,12 +118,51 @@ class ICLoraPipeline:
width: int,
num_frames: int,
frame_rate: float,
images: list[tuple[str, int, float]],
images: list[ImageConditioningInput],
video_conditioning: list[tuple[str, float]],
enhance_prompt: bool = False,
tiling_config: TilingConfig | None = None,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
conditioning_attention_strength: float = 1.0,
skip_stage_2: bool = False,
conditioning_attention_mask: torch.Tensor | None = None,
) -> tuple[Iterator[torch.Tensor], Audio]:
"""
Generate video with IC-LoRA conditioning.
Args:
prompt: Text prompt for video generation.
seed: Random seed for reproducibility.
height: Output video height in pixels (must be divisible by 64).
width: Output video width in pixels (must be divisible by 64).
num_frames: Number of frames to generate.
frame_rate: Output video frame rate.
images: List of (path, frame_idx, strength) tuples for image conditioning.
video_conditioning: List of (path, strength) tuples for IC-LoRA video conditioning.
enhance_prompt: Whether to enhance the prompt using the text encoder.
tiling_config: Optional tiling configuration for VAE decoding.
conditioning_attention_strength: Scale factor for IC-LoRA conditioning attention.
Controls how strongly the conditioning video influences the output.
0.0 = ignore conditioning, 1.0 = full conditioning influence. Default 1.0.
When conditioning_attention_mask is provided, the mask is multiplied by
this strength before being passed to the conditioning items.
skip_stage_2: If True, skip Stage 2 upsampling and refinement. Output will be
at half resolution (height//2, width//2). Default is False.
conditioning_attention_mask: Optional pixel-space attention mask with the same
spatial-temporal dimensions as the input reference video. Shape should be
(B, 1, F, H, W) or (1, 1, F, H, W) where F, H, W match the reference
video's pixel dimensions. Values in [0, 1].
The mask is downsampled to latent space using VAE scale factors (with
causal temporal handling for the first frame), then multiplied by
conditioning_attention_strength.
When None (default): scalar conditioning_attention_strength is used
directly.
Returns:
Tuple of (video_iterator, audio_tensor).
"""
assert_resolution(height=height, width=width, is_two_stage=True)
if not (0.0 <= conditioning_attention_strength <= 1.0):
raise ValueError(
f"conditioning_attention_strength must be in [0.0, 1.0], got {conditioning_attention_strength}"
)
generator = torch.Generator(device=self.device).manual_seed(seed)
noiser = GaussianNoiser(generator=generator)
@@ -158,6 +208,7 @@ class ICLoraPipeline:
height=height // 2,
fps=frame_rate,
)
stage_1_conditionings = self._create_conditionings(
images=images,
video_conditioning=video_conditioning,
@@ -165,7 +216,10 @@ class ICLoraPipeline:
width=stage_1_output_shape.width,
video_encoder=video_encoder,
num_frames=num_frames,
conditioning_attention_strength=conditioning_attention_strength,
conditioning_attention_mask=conditioning_attention_mask,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
@@ -182,6 +236,19 @@ class ICLoraPipeline:
del transformer
cleanup_memory()
if skip_stage_2:
# Skip Stage 2: Decode directly from Stage 1 output at half resolution
logging.info("[IC-LoRA] Skipping Stage 2 (--skip-stage-2 enabled)")
decoded_video = vae_decode_video(
video_state.latent, self.stage_1_model_ledger.video_decoder(), tiling_config, generator
)
decoded_audio = vae_decode_audio(
audio_state.latent, self.stage_1_model_ledger.audio_decoder(), self.stage_1_model_ledger.vocoder()
)
del video_encoder
cleanup_memory()
return decoded_video, decoded_audio
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
upscaled_video_latent = upsample_video(
latent=video_state.latent[:1],
@@ -250,13 +317,29 @@ class ICLoraPipeline:
def _create_conditionings(
self,
images: list[tuple[str, int, float]],
images: list[ImageConditioningInput],
video_conditioning: list[tuple[str, float]],
height: int,
width: int,
num_frames: int,
video_encoder: VideoEncoder,
conditioning_attention_strength: float = 1.0,
conditioning_attention_mask: torch.Tensor | None = None,
) -> list[ConditioningItem]:
"""
Create conditioning items for video generation.
Args:
conditioning_attention_strength: Scalar attention weight in [0, 1].
If conditioning_attention_mask is also provided, the downsampled mask
is multiplied by this strength. Otherwise this scalar is passed
directly as the attention mask.
conditioning_attention_mask: Optional pixel-space attention mask with shape
(B, 1, F_pixel, H_pixel, W_pixel) matching the reference video's
pixel dimensions. Downsampled to latent space with causal temporal
handling, then multiplied by conditioning_attention_strength.
Returns:
List of conditioning items. IC-LoRA conditionings are appended last.
"""
conditionings = image_conditionings_by_replacing_latent(
images=images,
height=height,
@@ -287,21 +370,96 @@ class ICLoraPipeline:
device=self.device,
)
encoded_video = video_encoder(video)
conditionings.append(
VideoConditionByReferenceLatent(
latent=encoded_video,
downscale_factor=scale,
strength=strength,
reference_video_shape = VideoLatentShape.from_torch_shape(encoded_video.shape)
# Build attention_mask for ConditioningItemAttentionStrengthWrapper
if conditioning_attention_mask is not None:
# Downsample pixel-space mask to latent space, then scale by strength
latent_mask = self._downsample_mask_to_latent(
mask=conditioning_attention_mask,
target_latent_shape=reference_video_shape,
)
attn_mask = latent_mask * conditioning_attention_strength
elif conditioning_attention_strength < 1.0:
# Use scalar strength only
attn_mask = conditioning_attention_strength
else:
attn_mask = None
cond = VideoConditionByReferenceLatent(
latent=encoded_video,
downscale_factor=scale,
strength=strength,
)
if attn_mask is not None:
cond = ConditioningItemAttentionStrengthWrapper(cond, attention_mask=attn_mask)
conditionings.append(cond)
if video_conditioning:
logging.info(f"[IC-LoRA] Added {len(video_conditioning)} video conditioning(s)")
return conditionings
@staticmethod
def _downsample_mask_to_latent(
mask: torch.Tensor,
target_latent_shape: VideoLatentShape,
) -> torch.Tensor:
"""
Downsample a pixel-space mask to latent space using VAE scale factors.
Handles causal temporal downsampling: the first frame is kept separately
(temporal scale factor = 1 for the first frame), while the remaining
frames are downsampled by the VAE's temporal scale factor.
Args:
mask: Pixel-space mask of shape (B, 1, F_pixel, H_pixel, W_pixel).
Values in [0, 1].
target_latent_shape: Expected latent shape after VAE encoding.
Used to determine the target (F_latent, H_latent, W_latent).
Returns:
Flattened latent-space mask of shape (B, F_lat * H_lat * W_lat),
matching the patchifier's token ordering (f, h, w).
"""
b = mask.shape[0]
f_lat = target_latent_shape.frames
h_lat = target_latent_shape.height
w_lat = target_latent_shape.width
# Step 1: Spatial downsampling (area interpolation per frame)
f_pix = mask.shape[2]
spatial_down = torch.nn.functional.interpolate(
rearrange(mask, "b 1 f h w -> (b f) 1 h w"),
size=(h_lat, w_lat),
mode="area",
)
spatial_down = rearrange(spatial_down, "(b f) 1 h w -> b 1 f h w", b=b)
# Step 2: Causal temporal downsampling
# First frame: kept as-is (causal VAE encodes first frame independently)
first_frame = spatial_down[:, :, :1, :, :] # (B, 1, 1, H_lat, W_lat)
if f_pix > 1 and f_lat > 1:
# Remaining frames: downsample by temporal factor via group-mean
t = (f_pix - 1) // (f_lat - 1) # temporal downscale factor
assert (f_pix - 1) % (f_lat - 1) == 0, (
f"Pixel frames ({f_pix}) not compatible with latent frames ({f_lat}): "
f"(f_pix - 1) must be divisible by (f_lat - 1)"
)
rest = rearrange(spatial_down[:, :, 1:, :, :], "b 1 (f t) h w -> b 1 f t h w", t=t)
rest = rest.mean(dim=3) # (B, 1, F_lat-1, H_lat, W_lat)
latent_mask = torch.cat([first_frame, rest], dim=2) # (B, 1, F_lat, H_lat, W_lat)
else:
latent_mask = first_frame
# Flatten to (B, F_lat * H_lat * W_lat) matching patchifier token order (f, h, w)
return rearrange(latent_mask, "b 1 f h w -> b (f h w)")
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
parser = default_2_stage_distilled_arg_parser()
checkpoint_path = detect_checkpoint_path(distilled=True)
params = detect_params(checkpoint_path)
parser = default_2_stage_distilled_arg_parser(params=params)
parser.add_argument(
"--video-conditioning",
action=VideoConditioningAction,
@@ -309,9 +467,47 @@ def main() -> None:
metavar=("PATH", "STRENGTH"),
required=True,
)
parser.add_argument(
"--conditioning-attention-mask",
action=VideoMaskConditioningAction,
nargs=2,
metavar=("MASK_PATH", "STRENGTH"),
default=None,
help=(
"Optional spatial attention mask: path to a grayscale mask video and "
"attention strength. The mask video pixel values in [0,1] control "
"per-region conditioning attention strength. The strength scalar is "
"multiplied with the spatial mask. "
"0.0 = ignore IC-LoRA conditioning, 1.0 = full conditioning influence. "
"When not provided, full conditioning strength (1.0) is used. "
"Example: --conditioning-attention-mask path/to/mask.mp4 0.5"
),
)
parser.add_argument(
"--skip-stage-2",
action="store_true",
help=(
"Skip Stage 2 upsampling and refinement. Output will be at half resolution "
"(height//2, width//2). Useful for faster iteration or when GPU memory is limited."
),
)
args = parser.parse_args()
# Load mask video if provided via --conditioning-attention-mask
conditioning_attention_mask = None
conditioning_attention_strength = 1.0
if args.conditioning_attention_mask is not None:
mask_path, mask_strength = args.conditioning_attention_mask
conditioning_attention_strength = mask_strength
conditioning_attention_mask = _load_mask_video(
mask_path=mask_path,
height=args.height // 2, # Stage 1 operates at half resolution
width=args.width // 2,
num_frames=args.num_frames,
)
pipeline = ICLoraPipeline(
checkpoint_path=args.checkpoint_path,
distilled_checkpoint_path=args.distilled_checkpoint_path,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
@@ -329,18 +525,53 @@ def main() -> None:
images=args.images,
video_conditioning=args.video_conditioning,
tiling_config=tiling_config,
conditioning_attention_strength=conditioning_attention_strength,
skip_stage_2=args.skip_stage_2,
conditioning_attention_mask=conditioning_attention_mask,
)
encode_video(
video=video,
fps=args.frame_rate,
audio=audio,
audio_sample_rate=AUDIO_SAMPLE_RATE,
output_path=args.output_path,
video_chunks_number=video_chunks_number,
)
def _load_mask_video(
mask_path: str,
height: int,
width: int,
num_frames: int,
) -> torch.Tensor:
"""Load a mask video and return a pixel-space tensor of shape (1, 1, F, H, W).
The mask video is loaded, resized to (height, width), converted to
grayscale, and normalised to [0, 1].
Args:
mask_path: Path to the mask video file.
height: Target height in pixels.
width: Target width in pixels.
num_frames: Maximum number of frames to load.
Returns:
Tensor of shape ``(1, 1, F, H, W)`` with values in ``[0, 1]``.
"""
mask_video = load_video_conditioning(
video_path=mask_path,
height=height,
width=width,
frame_cap=num_frames,
dtype=torch.bfloat16,
device=device,
)
# mask_video shape: (1, C, F, H, W) — take mean over channels for grayscale
mask = mask_video.mean(dim=1, keepdim=True) # (1, 1, F, H, W)
# Normalise to [0, 1] — load_video_conditioning applies normalize_latent,
# so undo that: values are in [-1, 1], remap to [0, 1]
mask = (mask + 1.0) / 2.0
return mask.clamp(0.0, 1.0)
def _read_lora_reference_downscale_factor(lora_path: str) -> int:
"""Read reference_downscale_factor from LoRA safetensors metadata.
Some IC-LoRA models are trained with reference videos at lower resolution than
packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py
+23 -20
View File
@@ -4,7 +4,11 @@ from collections.abc import Iterator
import torch
from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
from ltx_core.components.guiders import (
MultiModalGuiderFactory,
MultiModalGuiderParams,
create_multimodal_guider_factory,
)
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.components.schedulers import LTX2Scheduler
@@ -15,25 +19,22 @@ from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import LatentState, VideoPixelShape
from ltx_core.types import Audio, LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.args import default_2_stage_arg_parser
from ltx_pipelines.utils.constants import (
AUDIO_SAMPLE_RATE,
STAGE_2_DISTILLED_SIGMA_VALUES,
)
from ltx_pipelines.utils.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path
from ltx_pipelines.utils.constants import STAGE_2_DISTILLED_SIGMA_VALUES, detect_params
from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
denoise_audio_video,
euler_denoising_loop,
generate_enhanced_prompt,
get_device,
image_conditionings_by_adding_guiding_latent,
multi_modal_guider_denoising_func,
multi_modal_guider_factory_denoising_func,
simple_denoising_func,
)
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.samplers import euler_denoising_loop
from ltx_pipelines.utils.types import PipelineComponents
device = get_device()
@@ -43,8 +44,10 @@ class KeyframeInterpolationPipeline:
"""
Keyframe-based Two-stage video interpolation pipeline.
Interpolates between keyframes to generate a video with smoother transitions.
Stage 1 generates video at the target resolution, then Stage 2 upsamples
Stage 1 generates video at half of the target resolution, then Stage 2 upsamples
by 2x and refines with additional denoising steps for higher quality output.
Stage 1 uses full model while Stage 2 uses distilled LORA for efficiency,
as the upsampled video already has good quality and just needs refinement.
"""
def __init__(
@@ -76,7 +79,6 @@ class KeyframeInterpolationPipeline:
device=device,
)
@torch.inference_mode()
def __call__( # noqa: PLR0913
self,
prompt: str,
@@ -87,12 +89,12 @@ class KeyframeInterpolationPipeline:
num_frames: int,
frame_rate: float,
num_inference_steps: int,
video_guider_params: MultiModalGuiderParams,
audio_guider_params: MultiModalGuiderParams,
images: list[tuple[str, int, float]],
video_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
audio_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
images: list[ImageConditioningInput],
tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=True)
generator = torch.Generator(device=self.device).manual_seed(seed)
@@ -126,12 +128,12 @@ class KeyframeInterpolationPipeline:
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=multi_modal_guider_denoising_func(
video_guider=MultiModalGuider(
denoise_fn=multi_modal_guider_factory_denoising_func(
video_guider_factory=create_multimodal_guider_factory(
params=video_guider_params,
negative_context=v_context_n,
),
audio_guider=MultiModalGuider(
audio_guider_factory=create_multimodal_guider_factory(
params=audio_guider_params,
negative_context=a_context_n,
),
@@ -241,7 +243,9 @@ class KeyframeInterpolationPipeline:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
parser = default_2_stage_arg_parser()
checkpoint_path = detect_checkpoint_path()
params = detect_params(checkpoint_path)
parser = default_2_stage_arg_parser(params=params)
args = parser.parse_args()
pipeline = KeyframeInterpolationPipeline(
checkpoint_path=args.checkpoint_path,
@@ -286,7 +290,6 @@ def main() -> None:
video=video,
fps=args.frame_rate,
audio=audio,
audio_sample_rate=AUDIO_SAMPLE_RATE,
output_path=args.output_path,
video_chunks_number=video_chunks_number,
)
packages/ltx-pipelines/src/ltx_pipelines/retake.py
+476
View File
@@ -0,0 +1,476 @@
from __future__ import annotations
import argparse
import logging
from collections.abc import Iterator
from dataclasses import dataclass
import torch
from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.patchifiers import get_pixel_coords
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.components.schedulers import LTX2Scheduler
from ltx_core.conditioning import ConditioningItem
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.tools import LatentTools
from ltx_core.types import (
Audio,
AudioLatentShape,
LatentState,
SpatioTemporalScaleFactors,
VideoPixelShape,
)
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES
from ltx_pipelines.utils.helpers import (
cleanup_memory,
generate_enhanced_prompt,
get_device,
multi_modal_guider_denoising_func,
noise_audio_state,
noise_video_state,
simple_denoising_func,
)
from ltx_pipelines.utils.media_io import (
decode_audio_from_file,
encode_video,
get_videostream_metadata,
load_video_conditioning,
)
from ltx_pipelines.utils.samplers import euler_denoising_loop
from ltx_pipelines.utils.types import PipelineComponents
device = get_device()
def _encode_video_for_retake(
video_encoder: torch.nn.Module,
video_path: str,
output_shape: VideoPixelShape,
dtype: torch.dtype,
device: torch.device,
) -> torch.Tensor:
"""Load video and encode to latents."""
pixel_video = load_video_conditioning(
video_path=video_path,
height=output_shape.height,
width=output_shape.width,
frame_cap=output_shape.frames,
dtype=dtype,
device=device,
) # (1, C, F, H, W)
return video_encoder(pixel_video)
def _encode_audio_for_retake(
audio_encoder: torch.nn.Module,
waveform: torch.Tensor,
waveform_sr: int,
output_shape: VideoPixelShape,
dtype: torch.dtype,
) -> torch.Tensor:
"""Encode audio to latents and trim/pad to match output_shape."""
waveform_batch = waveform.unsqueeze(0) if waveform.dim() == 2 else waveform
initial_audio_latent = vae_encode_audio(
Audio(waveform=waveform_batch.to(dtype), sampling_rate=waveform_sr), audio_encoder, None
)
expected_audio_shape = AudioLatentShape.from_video_pixel_shape(output_shape)
expected_frames = expected_audio_shape.frames
actual_frames = initial_audio_latent.shape[2]
if actual_frames > expected_frames:
initial_audio_latent = initial_audio_latent[:, :, :expected_frames, :]
elif actual_frames < expected_frames:
pad = torch.zeros(
initial_audio_latent.shape[0],
initial_audio_latent.shape[1],
expected_frames - actual_frames,
initial_audio_latent.shape[3],
device=initial_audio_latent.device,
dtype=initial_audio_latent.dtype,
)
initial_audio_latent = torch.cat([initial_audio_latent, pad], dim=2)
return initial_audio_latent
# ---------------------------------------------------------------------------
# Custom conditioning item: temporal region mask
# ---------------------------------------------------------------------------
@dataclass(frozen=True)
class TemporalRegionMask:
"""Conditioning item that sets ``denoise_mask = 0`` outside a time range
and ``1`` inside, so only the specified temporal region is regenerated.
Uses ``start_time`` and ``end_time`` in seconds. Works in *patchified*
(token) space using the patchifier's ``get_patch_grid_bounds``: for video
coords are latent frame indices (converted from seconds via ``fps``), for
audio coords are already in seconds.
"""
start_time: float # seconds, inclusive
end_time: float # seconds, exclusive
fps: float
def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
coords = latent_tools.patchifier.get_patch_grid_bounds(
latent_tools.target_shape, device=latent_state.denoise_mask.device
)
# coords: [B, 3, N, 2] (video) or [B, 1, N, 2] (audio); temporal dim is index 0
if coords.shape[1] == 1:
# Audio: patchifier returns seconds
t_start = coords[:, 0, :, 0] # [B, N]
t_end = coords[:, 0, :, 1] # [B, N]
in_region = (t_end > self.start_time) & (t_start < self.end_time)
else:
# Video: get pixel bounds per patch, find patches for start/end frame, read latent from coords.
scale_factors = getattr(latent_tools, "scale_factors", SpatioTemporalScaleFactors.default())
pixel_bounds = get_pixel_coords(coords, scale_factors, causal_fix=getattr(latent_tools, "causal_fix", True))
timestamp_bounds = pixel_bounds[0, 0] / self.fps
t_start, t_end = timestamp_bounds.unbind(dim=-1)
in_region = (t_end > self.start_time) & (t_start < self.end_time)
state = latent_state.clone()
mask_val = in_region.to(state.denoise_mask.dtype)
if state.denoise_mask.dim() == 3:
mask_val = mask_val.unsqueeze(-1)
state.denoise_mask.copy_(mask_val)
return state
# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------
class RetakePipeline:
"""Regenerate a time region (retake) of an existing video.
Given a source video file and a time window ``[start_time, end_time]``
(in seconds), this pipeline keeps the video/audio outside that window
unchanged and *regenerates* the content inside the window from a text
prompt using the LTX-2 diffusion model.
Parameters
----------
checkpoint_path : str
Path to the LTX-2 model checkpoint.
gemma_root : str
Root directory containing Gemma text-encoder weights.
loras : list[LoraPathStrengthAndSDOps]
Optional LoRA configs applied to the transformer.
device : torch.device
Target device (default: CUDA if available).
quantization : QuantizationPolicy | None
Optional quantization policy for the transformer.
"""
def __init__(
self,
checkpoint_path: str,
gemma_root: str,
loras: list[LoraPathStrengthAndSDOps],
device: torch.device = device,
quantization: QuantizationPolicy | None = None,
):
self.device = device
self.dtype = torch.bfloat16
self.model_ledger = ModelLedger(
dtype=self.dtype,
device=device,
checkpoint_path=checkpoint_path,
gemma_root_path=gemma_root,
loras=loras,
quantization=quantization,
)
self.pipeline_components = PipelineComponents(
dtype=self.dtype,
device=device,
)
# --------------------------------------------------------------------- #
# Public entry point #
# --------------------------------------------------------------------- #
@torch.inference_mode()
def __call__( # noqa: PLR0913, PLR0915
self,
video_path: str,
prompt: str,
start_time: float,
end_time: float,
seed: int,
*,
negative_prompt: str = "",
num_inference_steps: int = 40,
video_guider_params: MultiModalGuiderParams | None = None,
audio_guider_params: MultiModalGuiderParams | None = None,
regenerate_video: bool = True,
regenerate_audio: bool = True,
enhance_prompt: bool = False,
distilled: bool = False,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
"""Regenerate ``[start_time, end_time]`` of the source video (retake).
Parameters
----------
video_path : str
Path to the source video file (must contain video; audio is optional).
prompt : str
Text prompt describing the *regenerated* section.
start_time, end_time : float
Time window (in seconds) of the section to regenerate.
seed : int
Random seed for reproducibility.
negative_prompt : str
Negative prompt for CFG guidance (ignored in distilled mode).
num_inference_steps : int
Number of Euler denoising steps (ignored in distilled mode which
uses a fixed 8-step schedule).
video_guider_params, audio_guider_params : MultiModalGuiderParams | None
Guidance parameters for video and audio modalities. Ignored in
distilled mode.
regenerate_video : bool
If ``True`` (default), preserve video outside ``[start_time, end_time]``
and only regenerate the masked region. If ``False``, fully regenerate
all video frames (the encoded video is still used as the initial latent
but with ``denoise_mask = 1`` everywhere).
regenerate_audio : bool
If True, regenerate audio in the [start_time, end_time] window; if False,
audio is preserved as-is (no regeneration).
enhance_prompt : bool
Whether to enhance the prompt via the text encoder.
distilled : bool
If ``True``, use the distilled sigma schedule
(``DISTILLED_SIGMA_VALUES``) and a simple (non-guided) denoising
function. The model checkpoint must be the distilled variant.
Returns
-------
tuple[Iterator[torch.Tensor], torch.Tensor]
``(video_frames_iterator, audio_waveform)``
"""
if start_time >= end_time:
raise ValueError(f"start_time ({start_time}) must be less than end_time ({end_time})")
effective_seed = torch.randint(0, 2**31, (1,), device=self.device).item() if seed < 0 else seed
generator = torch.Generator(device=self.device).manual_seed(effective_seed)
noiser = GaussianNoiser(generator=generator)
stepper = EulerDiffusionStep()
dtype = self.dtype
video_encoder = self.model_ledger.video_encoder()
# Use av to get metadata
fps, num_pixel_frames, src_width, src_height = get_videostream_metadata(video_path)
output_shape = VideoPixelShape(
batch=1,
frames=num_pixel_frames,
width=src_width,
height=src_height,
fps=fps,
)
initial_video_latent = _encode_video_for_retake(
video_encoder=video_encoder,
video_path=video_path,
output_shape=output_shape,
dtype=dtype,
device=self.device,
)
video_conditionings: list[ConditioningItem] = [
TemporalRegionMask(
start_time=start_time if regenerate_video else 0.0,
end_time=end_time if regenerate_video else 0.0,
fps=fps,
)
]
del video_encoder
cleanup_memory()
initial_audio_latent: torch.Tensor | None = None
audio_conditionings: list[ConditioningItem] = []
audio_in = decode_audio_from_file(video_path, self.device)
audio_encoder = self.model_ledger.audio_encoder()
if audio_in is not None:
waveform = audio_in.waveform.squeeze(0)
waveform_sr = audio_in.sampling_rate
else:
waveform, waveform_sr = None, None
if waveform is not None:
initial_audio_latent = _encode_audio_for_retake(
audio_encoder=audio_encoder,
waveform=waveform,
waveform_sr=waveform_sr,
output_shape=output_shape,
dtype=dtype,
)
audio_conditionings = [
TemporalRegionMask(
start_time=start_time if regenerate_audio else 0.0,
end_time=end_time if regenerate_audio else 0.0,
fps=fps,
)
]
del audio_encoder
cleanup_memory()
text_encoder = self.model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(text_encoder, prompt, None, seed=effective_seed)
if distilled:
# Distilled mode: single prompt, no negative
context_p = encode_text(text_encoder, prompts=[prompt])[0]
v_context_p, a_context_p = context_p
else:
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, a_context_n = context_n
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
transformer = self.model_ledger.transformer()
sigmas = (
torch.tensor(DISTILLED_SIGMA_VALUES) if distilled else LTX2Scheduler().execute(steps=num_inference_steps)
).to(dtype=torch.float32, device=self.device)
if distilled:
denoise_fn = simple_denoising_func(
video_context=v_context_p,
audio_context=a_context_p,
transformer=transformer,
)
else:
video_guider = MultiModalGuider(
params=video_guider_params,
negative_context=v_context_n,
)
audio_guider = MultiModalGuider(
params=audio_guider_params,
negative_context=a_context_n,
)
denoise_fn = multi_modal_guider_denoising_func(
video_guider,
audio_guider,
v_context=v_context_p,
a_context=a_context_p,
transformer=transformer,
)
def denoising_loop(
sigmas: torch.Tensor,
video_state: LatentState,
audio_state: LatentState,
stepper: DiffusionStepProtocol,
) -> tuple[LatentState, LatentState]:
return euler_denoising_loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=denoise_fn,
)
# Build noised states with the encoded latents as initial values and
# the temporal masks applied via conditionings.
video_state, video_tools = noise_video_state(
output_shape=output_shape,
noiser=noiser,
conditionings=video_conditionings,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
initial_latent=initial_video_latent,
)
audio_state, audio_tools = noise_audio_state(
output_shape=output_shape,
noiser=noiser,
conditionings=audio_conditionings,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
initial_latent=initial_audio_latent,
)
video_state, audio_state = denoising_loop(sigmas, video_state, audio_state, stepper)
video_state = video_tools.clear_conditioning(video_state)
video_state = video_tools.unpatchify(video_state)
audio_state = audio_tools.clear_conditioning(audio_state)
audio_state = audio_tools.unpatchify(audio_state)
torch.cuda.synchronize()
del transformer
cleanup_memory()
decoded_video = vae_decode_video(video_state.latent, self.model_ledger.video_decoder(), generator=generator)
decoded_audio = vae_decode_audio(
audio_state.latent, self.model_ledger.audio_decoder(), self.model_ledger.vocoder()
)
return decoded_video, decoded_audio
def main() -> None:
"""CLI entry point for retake (regenerate a time region)."""
logging.getLogger().setLevel(logging.INFO)
parser = argparse.ArgumentParser(description="Retake: regenerate a time region of a video with LTX-2.")
parser.add_argument("--video-path", type=str, required=True, help="Path to the source video.")
parser.add_argument("--prompt", type=str, required=True, help="Text prompt for the regenerated region.")
parser.add_argument("--start-time", type=float, required=True, help="Start time of the region to regenerate (s).")
parser.add_argument("--end-time", type=float, required=True, help="End time of the region to regenerate (s).")
parser.add_argument("--output-path", type=str, required=True, help="Path for the output video.")
parser.add_argument("--checkpoint-path", type=str, required=True, help="Path to the LTX-2 checkpoint.")
parser.add_argument("--gemma-root", type=str, required=True, help="Path to Gemma text encoder weights.")
parser.add_argument("--seed", type=int, default=42, help="Random seed. Use -1 for a random seed.")
parser.add_argument("--loras", nargs="*", default=[], help="LoRA paths (optional).")
args = parser.parse_args()
if args.start_time >= args.end_time:
raise ValueError("start_time must be less than end_time")
# Validate frame count (8k+1) and resolution (multiples of 32) at CLI stage
video_scale = SpatioTemporalScaleFactors.default()
fps, num_frames, width, height = get_videostream_metadata(args.video_path)
if (num_frames - 1) % video_scale.time != 0:
snapped = ((num_frames - 1) // video_scale.time) * video_scale.time + 1
raise ValueError(
f"Video frame count must satisfy 8k+1 (e.g. 97, 193). Got {num_frames}; use a video with {snapped} frames."
)
if width % 32 != 0 or height % 32 != 0:
raise ValueError(f"Video width and height must be multiples of 32. Got {width}x{height}.")
pipeline = RetakePipeline(
checkpoint_path=args.checkpoint_path,
gemma_root=args.gemma_root,
loras=args.loras or [],
)
video_iter, audio = pipeline(
video_path=args.video_path,
prompt=args.prompt,
start_time=args.start_time,
end_time=args.end_time,
seed=args.seed,
)
tiling_config = TilingConfig.default()
video_chunks_number = get_video_chunks_number(num_frames, tiling_config)
encode_video(
video=video_iter,
fps=int(fps),
audio=audio,
output_path=args.output_path,
video_chunks_number=video_chunks_number,
)
if __name__ == "__main__":
main()
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py
+31 -23
View File
@@ -4,7 +4,11 @@ from collections.abc import Iterator
import torch
from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
from ltx_core.components.guiders import (
MultiModalGuiderFactory,
MultiModalGuiderParams,
create_multimodal_guider_factory,
)
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.components.schedulers import LTX2Scheduler
@@ -13,11 +17,9 @@ from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.args import default_1_stage_arg_parser
from ltx_pipelines.utils.constants import AUDIO_SAMPLE_RATE
from ltx_pipelines.utils.helpers import (
from ltx_core.types import Audio, LatentState, VideoPixelShape
from ltx_pipelines.utils import (
ModelLedger,
assert_resolution,
cleanup_memory,
denoise_audio_video,
@@ -25,8 +27,10 @@ from ltx_pipelines.utils.helpers import (
generate_enhanced_prompt,
get_device,
image_conditionings_by_replacing_latent,
multi_modal_guider_denoising_func,
multi_modal_guider_factory_denoising_func,
)
from ltx_pipelines.utils.args import ImageConditioningInput, default_1_stage_arg_parser, detect_checkpoint_path
from ltx_pipelines.utils.constants import detect_params
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import PipelineComponents
@@ -39,6 +43,7 @@ class TI2VidOneStagePipeline:
Generates video at the target resolution in a single diffusion pass with
classifier-free guidance (CFG). Supports optional image conditioning via
the images parameter.
Assumes full non distilled model is provided in the checkpoint_path.
"""
def __init__(
@@ -74,11 +79,11 @@ class TI2VidOneStagePipeline:
num_frames: int,
frame_rate: float,
num_inference_steps: int,
video_guider_params: MultiModalGuiderParams,
audio_guider_params: MultiModalGuiderParams,
images: list[tuple[str, int, float]],
video_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
audio_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
images: list[ImageConditioningInput],
enhance_prompt: bool = False,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=False)
generator = torch.Generator(device=self.device).manual_seed(seed)
@@ -104,6 +109,15 @@ class TI2VidOneStagePipeline:
transformer = self.model_ledger.transformer()
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
video_guider_factory = create_multimodal_guider_factory(
params=video_guider_params,
negative_context=v_context_n,
)
audio_guider_factory = create_multimodal_guider_factory(
params=audio_guider_params,
negative_context=a_context_n,
)
def first_stage_denoising_loop(
sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
) -> tuple[LatentState, LatentState]:
@@ -112,15 +126,9 @@ class TI2VidOneStagePipeline:
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=multi_modal_guider_denoising_func(
video_guider=MultiModalGuider(
params=video_guider_params,
negative_context=v_context_n,
),
audio_guider=MultiModalGuider(
params=audio_guider_params,
negative_context=a_context_n,
),
denoise_fn=multi_modal_guider_factory_denoising_func(
video_guider_factory=video_guider_factory,
audio_guider_factory=audio_guider_factory,
v_context=v_context_p,
a_context=a_context_p,
transformer=transformer, # noqa: F821
@@ -157,14 +165,15 @@ class TI2VidOneStagePipeline:
decoded_audio = vae_decode_audio(
audio_state.latent, self.model_ledger.audio_decoder(), self.model_ledger.vocoder()
)
return decoded_video, decoded_audio
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
parser = default_1_stage_arg_parser()
checkpoint_path = detect_checkpoint_path()
params = detect_params(checkpoint_path)
parser = default_1_stage_arg_parser(params=params)
args = parser.parse_args()
pipeline = TI2VidOneStagePipeline(
checkpoint_path=args.checkpoint_path,
@@ -204,7 +213,6 @@ def main() -> None:
video=video,
fps=args.frame_rate,
audio=audio,
audio_sample_rate=AUDIO_SAMPLE_RATE,
output_path=args.output_path,
video_chunks_number=1,
)
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py
+26 -25
View File
@@ -4,7 +4,11 @@ from collections.abc import Iterator
import torch
from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
from ltx_core.components.guiders import (
MultiModalGuiderFactory,
MultiModalGuiderParams,
create_multimodal_guider_factory,
)
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.components.schedulers import LTX2Scheduler
@@ -15,14 +19,9 @@ from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.args import default_2_stage_arg_parser
from ltx_pipelines.utils.constants import (
AUDIO_SAMPLE_RATE,
STAGE_2_DISTILLED_SIGMA_VALUES,
)
from ltx_pipelines.utils.helpers import (
from ltx_core.types import Audio, LatentState, VideoPixelShape
from ltx_pipelines.utils import (
ModelLedger,
assert_resolution,
cleanup_memory,
denoise_audio_video,
@@ -30,9 +29,11 @@ from ltx_pipelines.utils.helpers import (
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.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path
from ltx_pipelines.utils.constants import STAGE_2_DISTILLED_SIGMA_VALUES, detect_params
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import PipelineComponents
@@ -42,9 +43,10 @@ device = get_device()
class TI2VidTwoStagesPipeline:
"""
Two-stage text/image-to-video generation pipeline.
Stage 1 generates video at the target resolution with CFG guidance, then
Stage 2 upsamples by 2x and refines using a distilled LoRA for higher
quality output. Supports optional image conditioning via the images parameter.
Stage 1 generates video at half of the target resolution with CFG guidance (assuming
full model is used), then Stage 2 upsamples by 2x and refines using a distilled
LoRA for higher quality output. Supports optional image conditioning via the
images parameter.
"""
def __init__(
@@ -54,7 +56,7 @@ class TI2VidTwoStagesPipeline:
spatial_upsampler_path: str,
gemma_root: str,
loras: list[LoraPathStrengthAndSDOps],
device: str = device,
device: torch.device = device,
quantization: QuantizationPolicy | None = None,
):
self.device = device
@@ -78,7 +80,6 @@ class TI2VidTwoStagesPipeline:
device=device,
)
@torch.inference_mode()
def __call__( # noqa: PLR0913
self,
prompt: str,
@@ -89,12 +90,12 @@ class TI2VidTwoStagesPipeline:
num_frames: int,
frame_rate: float,
num_inference_steps: int,
video_guider_params: MultiModalGuiderParams,
audio_guider_params: MultiModalGuiderParams,
images: list[tuple[str, int, float]],
video_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
audio_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
images: list[ImageConditioningInput],
tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=True)
generator = torch.Generator(device=self.device).manual_seed(seed)
@@ -128,12 +129,12 @@ class TI2VidTwoStagesPipeline:
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=multi_modal_guider_denoising_func(
video_guider=MultiModalGuider(
denoise_fn=multi_modal_guider_factory_denoising_func(
video_guider_factory=create_multimodal_guider_factory(
params=video_guider_params,
negative_context=v_context_n,
),
audio_guider=MultiModalGuider(
audio_guider_factory=create_multimodal_guider_factory(
params=audio_guider_params,
negative_context=a_context_n,
),
@@ -237,14 +238,15 @@ class TI2VidTwoStagesPipeline:
decoded_audio = vae_decode_audio(
audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder()
)
return decoded_video, decoded_audio
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
parser = default_2_stage_arg_parser()
checkpoint_path = detect_checkpoint_path()
params = detect_params(checkpoint_path)
parser = default_2_stage_arg_parser(params=params)
args = parser.parse_args()
pipeline = TI2VidTwoStagesPipeline(
checkpoint_path=args.checkpoint_path,
@@ -289,7 +291,6 @@ def main() -> None:
video=video,
fps=args.frame_rate,
audio=audio,
audio_sample_rate=AUDIO_SAMPLE_RATE,
output_path=args.output_path,
video_chunks_number=video_chunks_number,
)
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_res2s.py
+307
View File
@@ -0,0 +1,307 @@
import logging
from collections.abc import Iterator
import torch
from ltx_core.components.diffusion_steps import Res2sDiffusionStep
from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.components.schedulers import LTX2Scheduler
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.tools import VideoLatentShape
from ltx_core.types import Audio, LatentState, VideoPixelShape
from ltx_pipelines.utils import (
ModelLedger,
assert_resolution,
cleanup_memory,
denoise_audio_video,
generate_enhanced_prompt,
get_device,
image_conditionings_by_replacing_latent,
multi_modal_guider_denoising_func,
res2s_audio_video_denoising_loop,
simple_denoising_func,
)
from ltx_pipelines.utils.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path
from ltx_pipelines.utils.constants import STAGE_2_DISTILLED_SIGMA_VALUES, detect_params
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import PipelineComponents
device = get_device()
class TI2VidTwoStagesRes2sPipeline:
"""
Two-stage text/image-to-video generation pipeline using the res_2s sampler.
Same structure as :class:`TI2VidTwoStagesPipeline`: stage 1 generates video at
half of the target resolution with CFG guidance (assuming full model is used),
then Stage 2 upsamples by 2x and refines using a distilled LoRA for higher
quality output.
Uses the res_2s second-order sampler instead of Euler, allowing fewer
steps for comparable quality. Supports optional image conditioning via
the images parameter.
"""
def __init__(
self,
checkpoint_path: str,
distilled_lora: list[LoraPathStrengthAndSDOps],
spatial_upsampler_path: str,
gemma_root: str,
loras: list[LoraPathStrengthAndSDOps],
device: str = device,
quantization: QuantizationPolicy | None = None,
):
self.device = device
self.dtype = torch.bfloat16
self.stage_1_model_ledger = ModelLedger(
dtype=self.dtype,
device=device,
checkpoint_path=checkpoint_path,
gemma_root_path=gemma_root,
spatial_upsampler_path=spatial_upsampler_path,
loras=loras,
quantization=quantization,
)
self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
loras=distilled_lora,
)
self.pipeline_components = PipelineComponents(
dtype=self.dtype,
device=device,
)
@torch.inference_mode()
def __call__( # noqa: PLR0913
self,
prompt: str,
negative_prompt: str,
seed: int,
height: int,
width: int,
num_frames: int,
frame_rate: float,
num_inference_steps: int,
video_guider_params: MultiModalGuiderParams,
audio_guider_params: MultiModalGuiderParams,
images: list[ImageConditioningInput],
tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False,
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=True)
generator = torch.Generator(device=self.device).manual_seed(seed)
noiser = GaussianNoiser(generator=generator)
dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(
text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
)
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, a_context_n = context_n
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
# Stage 1: Initial low resolution video generation.
video_encoder = self.stage_1_model_ledger.video_encoder()
transformer = self.stage_1_model_ledger.transformer()
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
empty_latent = torch.empty(VideoLatentShape.from_pixel_shape(stage_1_output_shape).to_torch_shape())
stepper = Res2sDiffusionStep()
sigmas = (
LTX2Scheduler()
.execute(latent=empty_latent, steps=num_inference_steps)
.to(dtype=torch.float32, device=self.device)
)
def first_stage_denoising_loop(
sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
) -> tuple[LatentState, LatentState]:
return res2s_audio_video_denoising_loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=multi_modal_guider_denoising_func(
video_guider=MultiModalGuider(
params=video_guider_params,
negative_context=v_context_n,
),
audio_guider=MultiModalGuider(
params=audio_guider_params,
negative_context=a_context_n,
),
v_context=v_context_p,
a_context=a_context_p,
transformer=transformer, # noqa: F821
),
)
stage_1_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
noiser=noiser,
sigmas=sigmas,
stepper=stepper,
denoising_loop_fn=first_stage_denoising_loop,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
)
torch.cuda.synchronize()
del transformer
cleanup_memory()
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
upscaled_video_latent = upsample_video(
latent=video_state.latent[:1],
video_encoder=video_encoder,
upsampler=self.stage_2_model_ledger.spatial_upsampler(),
)
torch.cuda.synchronize()
cleanup_memory()
transformer = self.stage_2_model_ledger.transformer()
distilled_sigmas = torch.tensor(STAGE_2_DISTILLED_SIGMA_VALUES, device=self.device)
def second_stage_denoising_loop(
sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
) -> tuple[LatentState, LatentState]:
return res2s_audio_video_denoising_loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=simple_denoising_func(
video_context=v_context_p,
audio_context=a_context_p,
transformer=transformer, # noqa: F821
),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_2_output_shape,
conditionings=stage_2_conditionings,
noiser=noiser,
sigmas=distilled_sigmas,
stepper=stepper,
denoising_loop_fn=second_stage_denoising_loop,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
noise_scale=distilled_sigmas[0],
initial_video_latent=upscaled_video_latent,
initial_audio_latent=audio_state.latent,
)
torch.cuda.synchronize()
del transformer
del video_encoder
cleanup_memory()
decoded_video = vae_decode_video(
video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config, generator
)
decoded_audio = vae_decode_audio(
audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder()
)
return decoded_video, decoded_audio
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
checkpoint_path = detect_checkpoint_path()
params = detect_params(checkpoint_path)
parser = default_2_stage_arg_parser(params=params)
args = parser.parse_args()
pipeline = TI2VidTwoStagesRes2sPipeline(
checkpoint_path=args.checkpoint_path,
distilled_lora=args.distilled_lora,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
quantization=args.quantization,
)
tiling_config = TilingConfig.default()
video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
video, audio = pipeline(
prompt=args.prompt,
negative_prompt=args.negative_prompt,
seed=args.seed,
height=args.height,
width=args.width,
num_frames=args.num_frames,
frame_rate=args.frame_rate,
num_inference_steps=args.num_inference_steps,
video_guider_params=MultiModalGuiderParams(
cfg_scale=args.video_cfg_guidance_scale,
stg_scale=args.video_stg_guidance_scale,
rescale_scale=args.video_rescale_scale,
modality_scale=args.a2v_guidance_scale,
skip_step=args.video_skip_step,
stg_blocks=args.video_stg_blocks,
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=args.audio_cfg_guidance_scale,
stg_scale=args.audio_stg_guidance_scale,
rescale_scale=args.audio_rescale_scale,
modality_scale=args.v2a_guidance_scale,
skip_step=args.audio_skip_step,
stg_blocks=args.audio_stg_blocks,
),
images=args.images,
tiling_config=tiling_config,
)
encode_video(
video=video,
fps=args.frame_rate,
audio=audio,
output_path=args.output_path,
video_chunks_number=video_chunks_number,
)
if __name__ == "__main__":
main()
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