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

Automated PR - 2026-01-05

Browse Source
This commit is contained in:
sync-bot
2026-01-05 20:10:38 +00:00
parent fc3b319d34
commit 9ce438b353
153 changed files with 28100 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-pipelines/src/ltx_pipelines/utils/__init__.py
+5
View File
@@ -0,0 +1,5 @@
from ltx_pipelines.utils.model_ledger import ModelLedger
__all__ = [
"ModelLedger",
]
packages/ltx-pipelines/src/ltx_pipelines/utils/args.py
+277
View File
@@ -0,0 +1,277 @@
import argparse
from pathlib import Path
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
from ltx_pipelines.utils.constants import (
DEFAULT_1_STAGE_HEIGHT,
DEFAULT_1_STAGE_WIDTH,
DEFAULT_2_STAGE_HEIGHT,
DEFAULT_2_STAGE_WIDTH,
DEFAULT_CFG_GUIDANCE_SCALE,
DEFAULT_FRAME_RATE,
DEFAULT_LORA_STRENGTH,
DEFAULT_NEGATIVE_PROMPT,
DEFAULT_NUM_FRAMES,
DEFAULT_NUM_INFERENCE_STEPS,
DEFAULT_SEED,
)
class VideoConditioningAction(argparse.Action):
def __call__(
self,
parser: argparse.ArgumentParser, # noqa: ARG002
namespace: argparse.Namespace,
values: list[str],
option_string: str | None = None, # noqa: ARG002
) -> None:
path, strength_str = values
resolved_path = resolve_path(path)
strength = float(strength_str)
current = getattr(namespace, self.dest) or []
current.append((resolved_path, strength))
setattr(namespace, self.dest, current)
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
) -> None:
path, frame_idx, strength_str = values
resolved_path = resolve_path(path)
frame_idx = int(frame_idx)
strength = float(strength_str)
current = getattr(namespace, self.dest) or []
current.append((resolved_path, frame_idx, strength))
setattr(namespace, self.dest, current)
class LoraAction(argparse.Action):
def __call__(
self,
parser: argparse.ArgumentParser, # noqa: ARG002
namespace: argparse.Namespace,
values: list[str],
option_string: str | None = None,
) -> None:
if len(values) > 2:
msg = f"{option_string} accepts at most 2 arguments (PATH and optional STRENGTH), got {len(values)} values"
raise argparse.ArgumentError(self, msg)
path = values[0]
strength_str = values[1] if len(values) > 1 else str(DEFAULT_LORA_STRENGTH)
resolved_path = resolve_path(path)
strength = float(strength_str)
current = getattr(namespace, self.dest) or []
current.append(LoraPathStrengthAndSDOps(resolved_path, strength, LTXV_LORA_COMFY_RENAMING_MAP))
setattr(namespace, self.dest, current)
def resolve_path(path: str) -> str:
return str(Path(path).expanduser().resolve().as_posix())
def basic_arg_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser()
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,
required=True,
help="Path to the root directory containing the Gemma text encoder model files.",
)
parser.add_argument(
"--prompt",
type=str,
required=True,
help="Text prompt describing the desired video content to be generated by the model.",
)
parser.add_argument(
"--output-path",
type=resolve_path,
required=True,
help="Path to the output video file (MP4 format).",
)
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})."
),
)
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}).",
)
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}).",
)
parser.add_argument(
"--num-frames",
type=int,
default=DEFAULT_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}).",
)
parser.add_argument(
"--frame-rate",
type=float,
default=DEFAULT_FRAME_RATE,
help=f"Frame rate of the generated video (fps) (default: {DEFAULT_FRAME_RATE}).",
)
parser.add_argument(
"--num-inference-steps",
type=int,
default=DEFAULT_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})."
),
)
parser.add_argument(
"--image",
dest="images",
action=ImageAction,
nargs=3,
metavar=("PATH", "FRAME_IDX", "STRENGTH"),
default=[],
help=(
"Image conditioning input: path to image file, target frame index, "
"and conditioning strength (all three required). Default: empty list [] (no image conditioning). "
"Can be specified multiple times. Example: --image path/to/image1.jpg 0 0.8 "
"--image path/to/image2.jpg 160 0.9"
),
)
parser.add_argument(
"--lora",
dest="lora",
action=LoraAction,
nargs="+", # Accept 1-2 arguments per use (path and optional strength); validation is handled in LoraAction
metavar=("PATH", "STRENGTH"),
default=[],
help=(
"LoRA (Low-Rank Adaptation) model: path to model file and optional strength "
f"(default strength: {DEFAULT_LORA_STRENGTH}). Can be specified multiple times. "
"Example: --lora path/to/lora1.safetensors 0.8 --lora path/to/lora2.safetensors"
),
)
parser.add_argument(
"--enable-fp8",
action="store_true",
help="Enable FP8 mode to reduce memory footprint by keeping model in lower precision. "
"Note that calculations are still performed in bfloat16 precision.",
)
parser.add_argument("--enhance-prompt", action="store_true")
return parser
def default_1_stage_arg_parser() -> argparse.ArgumentParser:
parser = basic_arg_parser()
parser.add_argument(
"--cfg-guidance-scale",
type=float,
default=DEFAULT_CFG_GUIDANCE_SCALE,
help=(
f"Classifier-free guidance (CFG) scale controlling how strongly "
f"the model adheres to the prompt. Higher values increase prompt "
f"adherence but may reduce diversity (default: {DEFAULT_CFG_GUIDANCE_SCALE})."
),
)
parser.add_argument(
"--negative-prompt",
type=str,
default=DEFAULT_NEGATIVE_PROMPT,
help=(
"Negative prompt describing what should not appear in the generated video, "
"used to guide the diffusion process away from unwanted content. "
"Default: a comprehensive negative prompt covering common artifacts and quality issues."
),
)
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)
# 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})."
)
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})."
)
parser.add_argument(
"--distilled-lora",
dest="distilled_lora",
action=LoraAction,
nargs="+", # Accept 1-2 arguments per use (path and optional strength); validation is handled in LoraAction
metavar=("PATH", "STRENGTH"),
required=True,
help=(
"Distilled LoRA (Low-Rank Adaptation) model used in the second stage (upscaling and refinement): "
f"path to model file and optional strength (default strength: {DEFAULT_LORA_STRENGTH}). "
"The second stage upsamples the video by 2x resolution and refines it using a distilled "
"denoising schedule (fewer steps, no CFG). The distilled LoRA is specifically trained "
"for this refinement process to improve quality at higher resolutions. "
"Example: --distilled-lora path/to/distilled_lora.safetensors 0.8"
),
)
parser.add_argument(
"--spatial-upsampler-path",
type=resolve_path,
required=True,
help=(
"Path to the spatial upsampler model used to increase the resolution "
"of the generated video in the latent space."
),
)
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)
# 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})."
)
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})."
)
parser.add_argument(
"--spatial-upsampler-path",
type=resolve_path,
required=True,
help=(
"Path to the spatial upsampler model used to increase the resolution "
"of the generated video in the latent space."
),
)
return parser
packages/ltx-pipelines/src/ltx_pipelines/utils/constants.py
+77
View File
@@ -0,0 +1,77 @@
# =============================================================================
# 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.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)
STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0]
# =============================================================================
# Video Generation Defaults
# =============================================================================
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_CFG_GUIDANCE_SCALE = 4.0
# =============================================================================
# Audio
# =============================================================================
AUDIO_SAMPLE_RATE = 24000
# =============================================================================
# LoRA
# =============================================================================
DEFAULT_LORA_STRENGTH = 1.0
# =============================================================================
# Video VAE Architecture
# =============================================================================
VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
VIDEO_LATENT_CHANNELS = 128
# =============================================================================
# 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
# =============================================================================
# Prompts
# =============================================================================
DEFAULT_NEGATIVE_PROMPT = (
"blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, "
"grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, "
"deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, "
"wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of "
"field, background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent "
"lighting direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny "
"valley effect, incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, "
"mismatched lip sync, silent or muted audio, distorted voice, robotic voice, echo, background noise, "
"off-sync audio, incorrect dialogue, added dialogue, repetitive speech, jittery movement, awkward "
"pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, flat lighting, "
"inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts."
)
packages/ltx-pipelines/src/ltx_pipelines/utils/helpers.py
+474
View File
@@ -0,0 +1,474 @@
import gc
import logging
from dataclasses import replace
import torch
from tqdm import tqdm
from ltx_core.components.noisers import Noiser
from ltx_core.components.protocols import DiffusionStepProtocol, GuiderProtocol
from ltx_core.conditioning import (
ConditioningItem,
VideoConditionByKeyframeIndex,
VideoConditionByLatentIndex,
)
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.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.media_io import decode_image, load_image_conditioning, resize_aspect_ratio_preserving
from ltx_pipelines.utils.types import (
DenoisingFunc,
DenoisingLoopFunc,
PipelineComponents,
)
def get_device() -> torch.device:
if torch.cuda.is_available():
return torch.device("cuda")
return torch.device("cpu")
def cleanup_memory() -> None:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.synchronize()
def image_conditionings_by_replacing_latent(
images: list[tuple[str, int, float]],
height: int,
width: int,
video_encoder: VideoEncoder,
dtype: torch.dtype,
device: torch.device,
) -> list[ConditioningItem]:
conditionings = []
for image_path, frame_idx, strength in images:
image = load_image_conditioning(
image_path=image_path,
height=height,
width=width,
dtype=dtype,
device=device,
)
encoded_image = video_encoder(image)
conditionings.append(
VideoConditionByLatentIndex(
latent=encoded_image,
strength=strength,
latent_idx=frame_idx,
)
)
return conditionings
def image_conditionings_by_adding_guiding_latent(
images: list[tuple[str, int, float]],
height: int,
width: int,
video_encoder: VideoEncoder,
dtype: torch.dtype,
device: torch.device,
) -> list[ConditioningItem]:
conditionings = []
for image_path, frame_idx, strength in images:
image = load_image_conditioning(
image_path=image_path,
height=height,
width=width,
dtype=dtype,
device=device,
)
encoded_image = video_encoder(image)
conditionings.append(
VideoConditionByKeyframeIndex(keyframes=encoded_image, frame_idx=frame_idx, strength=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,
conditionings: list[ConditioningItem],
components: PipelineComponents,
dtype: torch.dtype,
device: torch.device,
noise_scale: float = 1.0,
initial_latent: torch.Tensor | None = None,
) -> tuple[LatentState, VideoLatentTools]:
"""Initialize and noise a video latent state for the diffusion pipeline.
Creates a video latent state from the output shape, applies conditionings,
and adds noise using the provided noiser. Returns the noised state and
video latent tools for further processing. If initial_latent is provided, it will be used to create the initial
state, otherwise an empty initial state will be created.
"""
video_latent_shape = VideoLatentShape.from_pixel_shape(
shape=output_shape,
latent_channels=components.video_latent_channels,
scale_factors=components.video_scale_factors,
)
video_tools = VideoLatentTools(components.video_patchifier, video_latent_shape, output_shape.fps)
video_state = create_noised_state(
tools=video_tools,
conditionings=conditionings,
noiser=noiser,
dtype=dtype,
device=device,
noise_scale=noise_scale,
initial_latent=initial_latent,
)
return video_state, video_tools
def noise_audio_state(
output_shape: VideoPixelShape,
noiser: Noiser,
conditionings: list[ConditioningItem],
components: PipelineComponents,
dtype: torch.dtype,
device: torch.device,
noise_scale: float = 1.0,
initial_latent: torch.Tensor | None = None,
) -> tuple[LatentState, AudioLatentTools]:
"""Initialize and noise an audio latent state for the diffusion pipeline.
Creates an audio latent state from the output shape, applies conditionings,
and adds noise using the provided noiser. Returns the noised state and
audio latent tools for further processing. If initial_latent is provided, it will be used to create the initial
state, otherwise an empty initial state will be created.
"""
audio_latent_shape = AudioLatentShape.from_video_pixel_shape(output_shape)
audio_tools = AudioLatentTools(components.audio_patchifier, audio_latent_shape)
audio_state = create_noised_state(
tools=audio_tools,
conditionings=conditionings,
noiser=noiser,
dtype=dtype,
device=device,
noise_scale=noise_scale,
initial_latent=initial_latent,
)
return audio_state, audio_tools
def create_noised_state(
tools: LatentTools,
conditionings: list[ConditioningItem],
noiser: Noiser,
dtype: torch.dtype,
device: torch.device,
noise_scale: float = 1.0,
initial_latent: torch.Tensor | None = None,
) -> LatentState:
"""Create a noised latent state from empty state, conditionings, and noiser.
Creates an empty latent state, applies conditionings, and then adds noise
using the provided noiser. Returns the final noised state ready for diffusion.
"""
state = tools.create_initial_state(device, dtype, initial_latent)
state = state_with_conditionings(state, conditionings, tools)
state = noiser(state, noise_scale)
return state
def state_with_conditionings(
latent_state: LatentState, conditioning_items: list[ConditioningItem], latent_tools: LatentTools
) -> LatentState:
"""Apply a list of conditionings to a latent state.
Iterates through the conditioning items and applies each one to the latent
state in sequence. Returns the modified state with all conditionings applied.
"""
for conditioning in conditioning_items:
latent_state = conditioning.apply_to(latent_state=latent_state, latent_tools=latent_tools)
return latent_state
def post_process_latent(denoised: torch.Tensor, denoise_mask: torch.Tensor, clean: torch.Tensor) -> torch.Tensor:
"""Blend denoised output with clean state based on mask."""
return (denoised * denoise_mask + clean.float() * (1 - denoise_mask)).to(denoised.dtype)
def modality_from_latent_state(
state: LatentState, context: torch.Tensor, sigma: float | 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
from the denoise mask and sigma, positions, and the provided context.
"""
return Modality(
enabled=enabled,
latent=state.latent,
timesteps=timesteps_from_mask(state.denoise_mask, sigma),
positions=state.positions,
context=context,
context_mask=None,
)
def timesteps_from_mask(denoise_mask: torch.Tensor, sigma: float | torch.Tensor) -> torch.Tensor:
"""Compute timesteps from a denoise mask and sigma value.
Multiplies the denoise mask by sigma to produce timesteps for each position
in the latent state. Areas where the mask is 0 will have zero timesteps.
"""
return denoise_mask * sigma
def simple_denoising_func(
video_context: torch.Tensor, audio_context: torch.Tensor, transformer: X0Model
) -> DenoisingFunc:
def simple_denoising_step(
video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
) -> tuple[torch.Tensor, torch.Tensor]:
sigma = sigmas[step_index]
pos_video = modality_from_latent_state(video_state, video_context, sigma)
pos_audio = modality_from_latent_state(audio_state, audio_context, sigma)
denoised_video, denoised_audio = transformer(video=pos_video, audio=pos_audio, perturbations=None)
return denoised_video, denoised_audio
return simple_denoising_step
def guider_denoising_func(
guider: GuiderProtocol,
v_context_p: torch.Tensor,
v_context_n: torch.Tensor,
a_context_p: torch.Tensor,
a_context_n: torch.Tensor,
transformer: X0Model,
) -> DenoisingFunc:
def guider_denoising_step(
video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
) -> tuple[torch.Tensor, torch.Tensor]:
sigma = sigmas[step_index]
pos_video = modality_from_latent_state(video_state, v_context_p, sigma)
pos_audio = modality_from_latent_state(audio_state, a_context_p, sigma)
denoised_video, denoised_audio = transformer(video=pos_video, audio=pos_audio, perturbations=None)
if guider.enabled():
neg_video = modality_from_latent_state(video_state, v_context_n, sigma)
neg_audio = modality_from_latent_state(audio_state, a_context_n, sigma)
neg_denoised_video, neg_denoised_audio = transformer(video=neg_video, audio=neg_audio, perturbations=None)
denoised_video = denoised_video + guider.delta(denoised_video, neg_denoised_video)
denoised_audio = denoised_audio + guider.delta(denoised_audio, neg_denoised_audio)
return denoised_video, denoised_audio
return guider_denoising_step
def denoise_audio_video( # 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,
) -> tuple[LatentState, 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, audio_tools = noise_audio_state(
output_shape=output_shape,
noiser=noiser,
conditionings=[],
components=components,
dtype=dtype,
device=device,
noise_scale=noise_scale,
initial_latent=initial_audio_latent,
)
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)
audio_state = audio_tools.clear_conditioning(audio_state)
audio_state = audio_tools.unpatchify(audio_state)
return video_state, audio_state
_UNICODE_REPLACEMENTS = str.maketrans("\u2018\u2019\u201c\u201d\u2014\u2013\u00a0\u2032\u2212", "''\"\"-- '-")
def clean_response(text: str) -> str:
"""Clean a response from curly quotes and leading non-letter characters which Gemma tends to insert."""
text = text.translate(_UNICODE_REPLACEMENTS)
# Remove leading non-letter characters
for i, char in enumerate(text):
if char.isalpha():
return text[i:]
return text
def generate_enhanced_prompt(
text_encoder: GemmaTextEncoderModelBase,
prompt: str,
image_path: str | None = None,
image_long_side: int = 896,
seed: int = 42,
) -> str:
"""Generate an enhanced prompt from a text encoder and a prompt."""
image = None
if image_path:
image = decode_image(image_path=image_path)
image = torch.tensor(image)
image = resize_aspect_ratio_preserving(image, image_long_side).to(torch.uint8)
prompt = text_encoder.enhance_i2v(prompt, image, seed=seed)
else:
prompt = text_encoder.enhance_t2v(prompt, seed=seed)
logging.info(f"Enhanced prompt: {prompt}")
return clean_response(prompt)
def assert_resolution(height: int, width: int, is_two_stage: bool) -> None:
"""Assert that the resolution is divisible by the required divisor.
For two-stage pipelines, the resolution must be divisible by 64.
For one-stage pipelines, the resolution must be divisible by 32.
"""
divisor = 64 if is_two_stage else 32
if height % divisor != 0 or width % divisor != 0:
raise ValueError(
f"Resolution ({height}x{width}) is not divisible by {divisor}. "
f"For {'two-stage' if is_two_stage else 'one-stage'} pipelines, "
f"height and width must be multiples of {divisor}."
)
packages/ltx-pipelines/src/ltx_pipelines/utils/media_io.py
+299
View File
@@ -0,0 +1,299 @@
import math
from collections.abc import Generator, Iterator
from fractions import Fraction
from io import BytesIO
import av
import numpy as np
import torch
from einops import rearrange
from PIL import Image
from torch._prims_common import DeviceLikeType
from tqdm import tqdm
from ltx_pipelines.utils.constants import DEFAULT_IMAGE_CRF
def resize_aspect_ratio_preserving(image: torch.Tensor, long_side: int) -> torch.Tensor:
"""
Resize image preserving aspect ratio (filling target long side).
Preserves the input dimensions order.
Args:
image: Input image tensor with shape (F (optional), H, W, C)
long_side: Target long side size.
Returns:
Tensor with shape (F (optional), H, W, C) F = 1 if input is 3D, otherwise input shape[0]
"""
height, width = image.shape[-3:2]
max_side = max(height, width)
scale = long_side / float(max_side)
target_height = int(height * scale)
target_width = int(width * scale)
resized = resize_and_center_crop(image, target_height, target_width)
# rearrange and remove batch dimension
result = rearrange(resized, "b c f h w -> b f h w c")[0]
# preserve input dimensions
return result[0] if result.shape[0] == 1 else result
def resize_and_center_crop(tensor: torch.Tensor, height: int, width: int) -> torch.Tensor:
"""
Resize tensor preserving aspect ratio (filling target), then center crop to exact dimensions.
Args:
latent: Input tensor with shape (H, W, C) or (F, H, W, C)
height: Target height
width: Target width
Returns:
Tensor with shape (1, C, 1, height, width) for 3D input or (1, C, F, height, width) for 4D input
"""
if tensor.ndim == 3:
tensor = rearrange(tensor, "h w c -> 1 c h w")
elif tensor.ndim == 4:
tensor = rearrange(tensor, "f h w c -> f c h w")
else:
raise ValueError(f"Expected input with 3 or 4 dimensions; got shape {tensor.shape}.")
_, _, src_h, src_w = tensor.shape
scale = max(height / src_h, width / src_w)
# Use ceil to avoid floating-point rounding causing new_h/new_w to be
# slightly smaller than target, which would result in negative crop offsets.
new_h = math.ceil(src_h * scale)
new_w = math.ceil(src_w * scale)
tensor = torch.nn.functional.interpolate(tensor, size=(new_h, new_w), mode="bilinear", align_corners=False)
crop_top = (new_h - height) // 2
crop_left = (new_w - width) // 2
tensor = tensor[:, :, crop_top : crop_top + height, crop_left : crop_left + width]
tensor = rearrange(tensor, "f c h w -> 1 c f h w")
return tensor
def normalize_latent(latent: torch.Tensor, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
return (latent / 127.5 - 1.0).to(device=device, dtype=dtype)
def load_image_conditioning(
image_path: str, height: int, width: int, dtype: torch.dtype, device: torch.device
) -> 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 = torch.tensor(image, dtype=torch.float32, device=device)
image = resize_and_center_crop(image, height, width)
image = normalize_latent(image, device, dtype)
return image
def load_video_conditioning(
video_path: str, height: int, width: int, frame_cap: int, dtype: torch.dtype, device: torch.device
) -> torch.Tensor:
"""
Loads a video from a path and preprocesses it for conditioning.
Note: The video is resized to the nearest multiple of 2 for compatibility with video codecs.
"""
frames = decode_video_from_file(path=video_path, frame_cap=frame_cap, device=device)
result = None
for f in frames:
frame = resize_and_center_crop(f.to(torch.float32), height, width)
frame = normalize_latent(frame, device, dtype)
result = frame if result is None else torch.cat([result, frame], dim=2)
return result
def decode_image(image_path: str) -> np.ndarray:
image = Image.open(image_path)
np_array = np.array(image)[..., :3]
return np_array
def _write_audio(
container: av.container.Container, audio_stream: av.audio.AudioStream, samples: torch.Tensor, audio_sample_rate: int
) -> None:
if samples.ndim == 1:
samples = samples[:, None]
if samples.shape[1] != 2 and samples.shape[0] == 2:
samples = samples.T
if samples.shape[1] != 2:
raise ValueError(f"Expected samples with 2 channels; got shape {samples.shape}.")
# Convert to int16 packed for ingestion; resampler converts to encoder fmt.
if samples.dtype != torch.int16:
samples = torch.clip(samples, -1.0, 1.0)
samples = (samples * 32767.0).to(torch.int16)
frame_in = av.AudioFrame.from_ndarray(
samples.contiguous().reshape(1, -1).cpu().numpy(),
format="s16",
layout="stereo",
)
frame_in.sample_rate = audio_sample_rate
_resample_audio(container, audio_stream, frame_in)
def _prepare_audio_stream(container: av.container.Container, audio_sample_rate: int) -> av.audio.AudioStream:
"""
Prepare the audio stream for writing.
"""
audio_stream = container.add_stream("aac", rate=audio_sample_rate)
audio_stream.codec_context.sample_rate = audio_sample_rate
audio_stream.codec_context.layout = "stereo"
audio_stream.codec_context.time_base = Fraction(1, audio_sample_rate)
return audio_stream
def _resample_audio(
container: av.container.Container, audio_stream: av.audio.AudioStream, frame_in: av.AudioFrame
) -> None:
cc = audio_stream.codec_context
# Use the encoder's format/layout/rate as the *target*
target_format = cc.format or "fltp" # AAC → usually fltp
target_layout = cc.layout or "stereo"
target_rate = cc.sample_rate or frame_in.sample_rate
audio_resampler = av.audio.resampler.AudioResampler(
format=target_format,
layout=target_layout,
rate=target_rate,
)
audio_next_pts = 0
for rframe in audio_resampler.resample(frame_in):
if rframe.pts is None:
rframe.pts = audio_next_pts
audio_next_pts += rframe.samples
rframe.sample_rate = frame_in.sample_rate
container.mux(audio_stream.encode(rframe))
# flush audio encoder
for packet in audio_stream.encode():
container.mux(packet)
def encode_video(
video: torch.Tensor | Iterator[torch.Tensor],
fps: int,
audio: torch.Tensor | None,
audio_sample_rate: int | None,
output_path: str,
video_chunks_number: int,
) -> None:
if isinstance(video, torch.Tensor):
video = iter([video])
first_chunk = next(video)
_, height, width, _ = first_chunk.shape
container = av.open(output_path, mode="w")
stream = container.add_stream("libx264", rate=int(fps))
stream.width = width
stream.height = height
stream.pix_fmt = "yuv420p"
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)
def all_tiles(
first_chunk: torch.Tensor, tiles_generator: Generator[tuple[torch.Tensor, int], None, None]
) -> Generator[tuple[torch.Tensor, int], None, None]:
yield first_chunk
yield from tiles_generator
for video_chunk in tqdm(all_tiles(first_chunk, video), total=video_chunks_number):
video_chunk_cpu = video_chunk.to("cpu").numpy()
for frame_array in video_chunk_cpu:
frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24")
for packet in stream.encode(frame):
container.mux(packet)
# Flush encoder
for packet in stream.encode():
container.mux(packet)
if audio is not None:
_write_audio(container, audio_stream, audio, audio_sample_rate)
container.close()
def decode_audio_from_file(path: str, device: torch.device) -> torch.Tensor | None:
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
finally:
container.close()
return audio
def decode_video_from_file(path: str, frame_cap: int, device: DeviceLikeType) -> Generator[torch.Tensor]:
container = av.open(path)
try:
video_stream = next(s for s in container.streams if s.type == "video")
for frame in container.decode(video_stream):
tensor = torch.tensor(frame.to_rgb().to_ndarray(), dtype=torch.uint8, device=device).unsqueeze(0)
yield tensor
frame_cap = frame_cap - 1
if frame_cap == 0:
break
finally:
container.close()
def encode_single_frame(output_file: str, image_array: np.ndarray, crf: float) -> None:
container = av.open(output_file, "w", format="mp4")
try:
stream = container.add_stream("libx264", rate=1, options={"crf": str(crf), "preset": "veryfast"})
# Round to nearest multiple of 2 for compatibility with video codecs
height = image_array.shape[0] // 2 * 2
width = image_array.shape[1] // 2 * 2
image_array = image_array[:height, :width]
stream.height = height
stream.width = width
av_frame = av.VideoFrame.from_ndarray(image_array, format="rgb24").reformat(format="yuv420p")
container.mux(stream.encode(av_frame))
container.mux(stream.encode())
finally:
container.close()
def decode_single_frame(video_file: str) -> np.array:
container = av.open(video_file)
try:
stream = next(s for s in container.streams if s.type == "video")
frame = next(container.decode(stream))
finally:
container.close()
return frame.to_ndarray(format="rgb24")
def preprocess(image: np.array, crf: float = DEFAULT_IMAGE_CRF) -> np.array:
if crf == 0:
return image
with BytesIO() as output_file:
encode_single_frame(output_file, image, crf)
video_bytes = output_file.getvalue()
with BytesIO(video_bytes) as video_file:
image_array = decode_single_frame(video_file)
return image_array
packages/ltx-pipelines/src/ltx_pipelines/utils/model_ledger.py
+243
View File
@@ -0,0 +1,243 @@
from dataclasses import replace
import torch
from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
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,
VOCODER_COMFY_KEYS_FILTER,
AudioDecoder,
AudioDecoderConfigurator,
Vocoder,
VocoderConfigurator,
)
from ltx_core.model.transformer import (
LTXV_MODEL_COMFY_RENAMING_MAP,
LTXV_MODEL_COMFY_RENAMING_WITH_TRANSFORMER_LINEAR_DOWNCAST_MAP,
UPCAST_DURING_INFERENCE,
LTXModelConfigurator,
X0Model,
)
from ltx_core.model.upsampler import LatentUpsampler, LatentUpsamplerConfigurator
from ltx_core.model.video_vae import (
VAE_DECODER_COMFY_KEYS_FILTER,
VAE_ENCODER_COMFY_KEYS_FILTER,
VideoDecoder,
VideoDecoderConfigurator,
VideoEncoder,
VideoEncoderConfigurator,
)
from ltx_core.text_encoders.gemma import (
AV_GEMMA_TEXT_ENCODER_KEY_OPS,
AVGemmaTextEncoderModel,
AVGemmaTextEncoderModelConfigurator,
module_ops_from_gemma_root,
)
class ModelLedger:
"""
Central coordinator for loading and building models used in an LTX pipeline.
The ledger wires together multiple model builders (transformer, video VAE encoder/decoder,
audio VAE decoder, vocoder, text encoder, and optional latent upsampler) and exposes
factory methods for constructing model instances.
### Model Building
Each model method (e.g. :meth:`transformer`, :meth:`video_decoder`, :meth:`text_encoder`)
constructs a new model instance on each call. The builder uses the
:class:`~ltx_core.loader.registry.Registry` to load weights from the checkpoint,
instantiates the model with the configured ``dtype``, and moves it to ``self.device``.
.. note::
Models are **not cached**. Each call to a model method creates a new instance.
Callers are responsible for storing references to models they wish to reuse
and for freeing GPU memory (e.g. by deleting references and calling
``torch.cuda.empty_cache()``).
### Constructor parameters
dtype:
Torch dtype used when constructing all models (e.g. ``torch.bfloat16``).
device:
Target device to which models are moved after construction (e.g. ``torch.device("cuda")``).
checkpoint_path:
Path to a checkpoint directory or file containing the core model weights
(transformer, video VAE, audio VAE, text encoder, vocoder). If ``None``, the
corresponding builders are not created and calling those methods will raise
a :class:`ValueError`.
gemma_root_path:
Base path to Gemma-compatible CLIP/text encoder weights. Required to
initialize the text encoder builder; if omitted, :meth:`text_encoder` cannot be used.
spatial_upsampler_path:
Optional path to a latent upsampler checkpoint. If provided, the
:meth:`spatial_upsampler` method becomes available; otherwise calling it raises
a :class:`ValueError`.
loras:
Optional collection of LoRA configurations (paths, strengths, and key operations)
that are applied on top of the base transformer weights when building the model.
registry:
Optional :class:`Registry` instance for weight caching across builders.
Defaults to :class:`DummyRegistry` which performs no cross-builder caching.
fp8transformer:
If ``True``, builds the transformer with FP8 quantization and upcasting during inference.
### Creating Variants
Use :meth:`with_loras` to create a new ``ModelLedger`` instance that includes
additional LoRA configurations while sharing the same registry for weight caching.
"""
def __init__(
self,
dtype: torch.dtype,
device: torch.device,
checkpoint_path: str | None = None,
gemma_root_path: str | None = None,
spatial_upsampler_path: str | None = None,
loras: LoraPathStrengthAndSDOps | None = None,
registry: Registry | None = None,
fp8transformer: bool = False,
):
self.dtype = dtype
self.device = device
self.checkpoint_path = checkpoint_path
self.gemma_root_path = gemma_root_path
self.spatial_upsampler_path = spatial_upsampler_path
self.loras = loras or ()
self.registry = registry or DummyRegistry()
self.fp8transformer = fp8transformer
self.build_model_builders()
def build_model_builders(self) -> None:
if self.checkpoint_path is not None:
self.transformer_builder = Builder(
model_path=self.checkpoint_path,
model_class_configurator=LTXModelConfigurator,
model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
loras=tuple(self.loras),
registry=self.registry,
)
self.vae_decoder_builder = Builder(
model_path=self.checkpoint_path,
model_class_configurator=VideoDecoderConfigurator,
model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
registry=self.registry,
)
self.vae_encoder_builder = Builder(
model_path=self.checkpoint_path,
model_class_configurator=VideoEncoderConfigurator,
model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
registry=self.registry,
)
self.audio_decoder_builder = Builder(
model_path=self.checkpoint_path,
model_class_configurator=AudioDecoderConfigurator,
model_sd_ops=AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
registry=self.registry,
)
self.vocoder_builder = Builder(
model_path=self.checkpoint_path,
model_class_configurator=VocoderConfigurator,
model_sd_ops=VOCODER_COMFY_KEYS_FILTER,
registry=self.registry,
)
if self.gemma_root_path is not None:
self.text_encoder_builder = Builder(
model_path=self.checkpoint_path,
model_class_configurator=AVGemmaTextEncoderModelConfigurator,
model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS,
registry=self.registry,
module_ops=module_ops_from_gemma_root(self.gemma_root_path),
)
if self.spatial_upsampler_path is not None:
self.upsampler_builder = Builder(
model_path=self.spatial_upsampler_path,
model_class_configurator=LatentUpsamplerConfigurator,
registry=self.registry,
)
def _target_device(self) -> torch.device:
if isinstance(self.registry, DummyRegistry) or self.registry is None:
return self.device
else:
return torch.device("cpu")
def with_loras(self, loras: LoraPathStrengthAndSDOps) -> "ModelLedger":
return ModelLedger(
dtype=self.dtype,
device=self.device,
checkpoint_path=self.checkpoint_path,
gemma_root_path=self.gemma_root_path,
spatial_upsampler_path=self.spatial_upsampler_path,
loras=(*self.loras, *loras),
registry=self.registry,
fp8transformer=self.fp8transformer,
)
def transformer(self) -> X0Model:
if not hasattr(self, "transformer_builder"):
raise ValueError(
"Transformer not initialized. Please provide a checkpoint path to the ModelLedger constructor."
)
if self.fp8transformer:
fp8_builder = replace(
self.transformer_builder,
module_ops=(UPCAST_DURING_INFERENCE,),
model_sd_ops=LTXV_MODEL_COMFY_RENAMING_WITH_TRANSFORMER_LINEAR_DOWNCAST_MAP,
)
return X0Model(fp8_builder.build(device=self._target_device())).to(self.device).eval()
else:
return (
X0Model(self.transformer_builder.build(device=self._target_device(), dtype=self.dtype))
.to(self.device)
.eval()
)
def video_decoder(self) -> VideoDecoder:
if not hasattr(self, "vae_decoder_builder"):
raise ValueError(
"Video decoder not initialized. Please provide a checkpoint path to the ModelLedger constructor."
)
return self.vae_decoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
def video_encoder(self) -> VideoEncoder:
if not hasattr(self, "vae_encoder_builder"):
raise ValueError(
"Video encoder not initialized. Please provide a checkpoint path to the ModelLedger constructor."
)
return self.vae_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
def text_encoder(self) -> AVGemmaTextEncoderModel:
if not hasattr(self, "text_encoder_builder"):
raise ValueError(
"Text encoder not initialized. Please provide a checkpoint path and gemma root path to the "
"ModelLedger constructor."
)
return self.text_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(
"Audio decoder not initialized. Please provide a checkpoint path to the ModelLedger constructor."
)
return self.audio_decoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
def vocoder(self) -> Vocoder:
if not hasattr(self, "vocoder_builder"):
raise ValueError(
"Vocoder not initialized. Please provide a checkpoint path to the ModelLedger constructor."
)
return self.vocoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
def spatial_upsampler(self) -> LatentUpsampler:
if not hasattr(self, "upsampler_builder"):
raise ValueError("Upsampler not initialized. Please provide upsampler path to the ModelLedger constructor.")
return self.upsampler_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
packages/ltx-pipelines/src/ltx_pipelines/utils/types.py
+73
View File
@@ -0,0 +1,73 @@
from typing import Protocol
import torch
from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.types import LatentState
from ltx_pipelines.utils.constants import VIDEO_LATENT_CHANNELS, VIDEO_SCALE_FACTORS
class PipelineComponents:
"""
Container class for pipeline components used throughout the LTX pipelines.
Attributes:
dtype (torch.dtype): Default torch dtype for tensors in the pipeline.
device (torch.device): Target device to place tensors and modules on.
video_scale_factors (SpatioTemporalScaleFactors): Scale factors (T, H, W) for VAE latent space.
video_latent_channels (int): Number of channels in the video latent representation.
video_patchifier (VideoLatentPatchifier): Patchifier instance for video latents.
audio_patchifier (AudioPatchifier): Patchifier instance for audio latents.
"""
def __init__(
self,
dtype: torch.dtype,
device: torch.device,
):
self.dtype = dtype
self.device = device
self.video_scale_factors = VIDEO_SCALE_FACTORS
self.video_latent_channels = VIDEO_LATENT_CHANNELS
self.video_patchifier = VideoLatentPatchifier(patch_size=1)
self.audio_patchifier = AudioPatchifier(patch_size=1)
class DenoisingFunc(Protocol):
"""
Protocol for a denoising function used in the LTX pipeline.
Args:
video_state (LatentState): The current latent state for video.
audio_state (LatentState): The current latent state for audio.
sigmas (torch.Tensor): A 1D tensor of sigma values for each diffusion step.
step_index (int): Index of the current denoising step.
Returns:
tuple[torch.Tensor, torch.Tensor]: The denoised video and audio tensors.
"""
def __call__(
self, video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
) -> tuple[torch.Tensor, torch.Tensor]: ...
class DenoisingLoopFunc(Protocol):
"""
Protocol for a denoising loop function used in the LTX pipeline.
Args:
sigmas (torch.Tensor): A 1D tensor of sigma values for each diffusion step.
video_state (LatentState): The current latent state for video.
audio_state (LatentState): The current latent state for audio.
stepper (DiffusionStepProtocol): The diffusion step protocol to use.
Returns:
tuple[LatentState, LatentState]: The denoised video and audio latent states.
"""
def __call__(
self,
sigmas: torch.Tensor,
video_state: LatentState,
audio_state: LatentState,
stepper: DiffusionStepProtocol,
) -> tuple[torch.Tensor, torch.Tensor]: ...