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/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