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

Automated PR - 2026-04-13

Browse Source
This commit is contained in:
github-actions[bot]
2026-04-13 14:29:35 +00:00
parent 59ca828d5a
commit d887bbd1e0
29 changed files with 463 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-pipelines/CLAUDE.md
+82
View File
@@ -0,0 +1,82 @@
<!--
MAINTENANCE: When modifying any pipeline class in src/ltx_pipelines/,
update this document to reflect changes to:
- __init__ / __call__ signatures
- sigma handling or step counts
- denoiser types or guidance
- new or removed pipelines
Run: ls src/ltx_pipelines/*.py to check for new pipeline files.
-->
# ltx-pipelines
Inference pipelines for LTX-2 audio-video generation. Depends on `ltx-core` for model definitions, diffusion components, and loading. All pipelines live in `packages/ltx-pipelines/src/ltx_pipelines/`.
## Pipeline selection
| Pipeline | File | Stages | Model | Sampler | Use case |
|----------|------|--------|-------|---------|----------|
| `TI2VidOneStagePipeline` | `ti2vid_one_stage.py` | 1 | Full | Euler | Simple text/image-to-video |
| `TI2VidTwoStagesPipeline` | `ti2vid_two_stages.py` | 2 | Full + distilled LoRA | Euler | Production quality |
| `TI2VidTwoStagesHQPipeline` | `ti2vid_two_stages_hq.py` | 2 | Full + distilled LoRA (both stages) | Res2s | Highest quality, fewer steps |
| `A2VidPipelineTwoStage` | `a2vid_two_stage.py` | 2 | Full + distilled LoRA | Euler | Audio-conditioned video |
| `KeyframeInterpolationPipeline` | `keyframe_interpolation.py` | 2 | Full + distilled LoRA | Euler | Keyframe interpolation |
| `DistilledPipeline` | `distilled.py` | 2 | Distilled only | Euler | Fastest inference |
| `ICLoraPipeline` | `ic_lora.py` | 2 | Distilled only | Euler | Video-to-video with IC-LoRA control |
| `RetakePipeline` | `retake.py` | 1 | Full or distilled | Euler | Video region regeneration |
## Guidance
- **CFG**: Blends conditioned/unconditioned predictions. Defaults: `cfg_scale=3.0` (video), `7.0` (audio).
- **STG**: Perturbs self-attention in transformer blocks. Default `stg_scale=1.0`, `stg_blocks=[28]` (LTX-2.3) / `[29]` (LTX-2). HQ disables STG (`stg_scale=0.0`).
- **Modality guidance**: Cross-modal attention scaling (`modality_scale=3.0`).
- All guidance is stage 1 only. Stage 2 always uses `SimpleDenoiser`.
## Sigma schedules and step counts
- **Scheduler-based** (full model): `self._scheduler = LTX2Scheduler()` with `execute(steps=N)` (HQ also passes `latent=` for token-count-dependent shift). Defaults: 30 steps (LTX-2.3), 40 (LTX-2), 15 (HQ).
- **Distilled**: Fixed 8-step `DISTILLED_SIGMA_VALUES` (9 values). Stage 2 uses 3-step `STAGE_2_DISTILLED_SIGMA_VALUES` (4 values). No `num_inference_steps` param.
- **Retake**: `num_inference_steps=40` default; ignored when `distilled=True` (fixed 8-step).
- **Overrides**: All pipelines accept optional sigma tensors in `__call__`: `sigmas` (one-stage), `stage_1_sigmas` + `stage_2_sigmas` (two-stage).
## LoRA conventions
- No default LoRAs. `loras` param defaults to empty list/tuple. `DEFAULT_LORA_STRENGTH = 1.0`.
- Two-stage non-distilled pipelines require `distilled_lora` (applied to stage 2 only in TI2Vid/A2Vid/Keyframe).
- HQ is unique: applies distilled LoRA to **both** stages with separate `distilled_lora_strength_stage_1` / `_stage_2` params.
## Shared building blocks (`utils/blocks.py`)
- `DiffusionStage` -- owns transformer lifecycle; builds model on call, frees on exit via `gpu_model()` context manager (moves params to meta device to release GPU/CPU memory). Accepts optional `stepper` and `loop` overrides.
- `PromptEncoder` -- Gemma text encoder + embeddings processor (video 4096-dim, audio 2048-dim).
- `ImageConditioner` / `AudioConditioner` -- temporary encoder scope; builds encoder, passes to callable, frees.
- `VideoUpsampler` -- 2x spatial upsampling via encoder + upsampler.
- `VideoDecoder` / `AudioDecoder` -- latent-to-pixel decoding (iterator for video, `Audio` for audio).
### Memory management
- **Model lifecycle**: All blocks build their model on call and free it on exit. `gpu_model()` moves params to `"meta"` device on exit, immediately releasing storage. No model persists between calls.
- **Layer streaming**: When `streaming_prefetch_count` is set, `DiffusionStage` wraps the transformer in `LayerStreamingWrapper`. Layers live on pinned CPU memory; only `1 + prefetch_count` layers are on GPU at a time, with async H2D prefetch on a separate CUDA stream.
- **Batch splitting**: `BatchSplitAdapter` wraps the transformer and splits inputs exceeding `max_batch_size` into sequential chunks. If guidance needs B=4 but `max_batch_size=1`, it runs 4 sequential B=1 passes. Higher `max_batch_size` reduces layer-streaming PCIe transfers at the cost of peak memory.
## Denoisers (`utils/denoisers.py`)
- `SimpleDenoiser` -- single forward pass (B=1), no guidance. Used by distilled pipelines and all stage 2.
- `GuidedDenoiser` -- CFG/STG with static `MultiModalGuider` instances (HQ, A2Vid, Retake non-distilled).
- `FactoryGuidedDenoiser` -- per-step guider creation via factory (OneStageTI2Vid, TwoStagesTI2Vid, Keyframe).
Guided denoisers batch all guidance passes into a **single transformer call**: states are repeated along the batch dimension, contexts concatenated, and a `BatchedPerturbationConfig` controls which attention ops are skipped per sample. Pass count is dynamic: B=2 for CFG-only, up to B=4 with CFG+STG+modality isolation. Results are split back and blended by the guider.
## Per-pipeline unique features
- **HQ**: Res2s second-order sampler for **both** stages, latent-dependent sigma schedule, distilled LoRA on both stages with separate strengths.
- **A2Vid**: Audio frozen in both stages (`frozen=True, noise_scale=0.0`). Returns original audio (not VAE-decoded); no `AudioDecoder`.
- **IC-LoRA**: `VideoConditionByReferenceLatent`, `reference_downscale_factor` from LoRA metadata, `skip_stage_2`, attention mask downsampling. Stage 2 is LoRA-free and uses `combined_image_conditionings` (no IC-LoRA conditioning).
- **Keyframe**: Uses `image_conditionings_by_adding_guiding_latent` in both stages (all frames as keyframe guidance, no replacement) -- unlike TI2Vid which uses `combined_image_conditionings` (frame_idx=0 replaces, others guide).
- **Retake**: `TemporalRegionMask` for selective time-window regeneration. `regenerate_video`/`regenerate_audio` flags. Conditional distilled/full behavior.
- **Distilled**: Single `self.stage` reused for both stages (not `stage_1`/`stage_2`).
## Image conditioning helpers (`utils/helpers.py`)
- `combined_image_conditionings()` -- images with `frame_idx==0` replace latent (`VideoConditionByLatentIndex`), others guide (`VideoConditionByKeyframeIndex`).
- `image_conditionings_by_adding_guiding_latent()` -- all images become keyframe guidance regardless of `frame_idx`.
packages/ltx-pipelines/pyproject.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[project]
name = "ltx-pipelines"
version = "1.1.0"
version = "1.1.1"
description = "Pipelines implementation for Lightricks' LTX-2 model"
readme = "README.md"
requires-python = ">=3.10"
packages/ltx-pipelines/src/ltx_pipelines/a2vid_two_stage.py
+10 -5
View File
@@ -22,7 +22,7 @@ from ltx_pipelines.utils.blocks import (
VideoUpsampler,
)
from ltx_pipelines.utils.constants import (
STAGE_2_DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMAS,
)
from ltx_pipelines.utils.denoisers import GuidedDenoiser, SimpleDenoiser
from ltx_pipelines.utils.helpers import (
@@ -56,6 +56,7 @@ class A2VidPipelineTwoStage:
):
self.device = device or get_device()
self.dtype = torch.bfloat16
self._scheduler = LTX2Scheduler()
self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
@@ -103,6 +104,8 @@ class A2VidPipelineTwoStage:
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
max_batch_size: int = 1,
stage_1_sigmas: torch.Tensor | None = None,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=True)
@@ -148,7 +151,9 @@ class A2VidPipelineTwoStage:
)
)
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
sigmas = (
stage_1_sigmas if stage_1_sigmas is not None else self._scheduler.execute(steps=num_inference_steps)
).to(dtype=torch.float32, device=self.device)
video_state, _ = self.stage_1(
denoiser=GuidedDenoiser(
@@ -185,7 +190,7 @@ class A2VidPipelineTwoStage:
# Stage 2: Upsample and refine the video at higher resolution with distilled LoRA.
upscaled_video_latent = self.upsampler(video_state.latent[:1])
distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
stage_2_sigmas = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = self.image_conditioner(
lambda enc: combined_image_conditionings(
@@ -200,7 +205,7 @@ class A2VidPipelineTwoStage:
video_state, _ = self.stage_2(
denoiser=SimpleDenoiser(v_context_p, a_context_p),
sigmas=distilled_sigmas,
sigmas=stage_2_sigmas,
noiser=noiser,
width=width,
height=height,
@@ -209,7 +214,7 @@ class A2VidPipelineTwoStage:
video=ModalitySpec(
context=v_context_p,
conditionings=stage_2_conditionings,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=upscaled_video_latent,
),
audio=ModalitySpec(
packages/ltx-pipelines/src/ltx_pipelines/distilled.py
+7 -5
View File
@@ -23,8 +23,8 @@ from ltx_pipelines.utils.blocks import (
VideoUpsampler,
)
from ltx_pipelines.utils.constants import (
DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMA_VALUES,
DISTILLED_SIGMAS,
STAGE_2_DISTILLED_SIGMAS,
detect_params,
)
from ltx_pipelines.utils.denoisers import SimpleDenoiser
@@ -77,7 +77,7 @@ class DistilledPipeline:
self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
self.audio_decoder = AudioDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
def __call__(
def __call__( # noqa: PLR0913
self,
prompt: str,
seed: int,
@@ -89,6 +89,8 @@ class DistilledPipeline:
tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=True)
@@ -105,7 +107,7 @@ class DistilledPipeline:
video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
# Stage 1: Initial low resolution video generation.
stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
stage_1_sigmas = stage_1_sigmas.to(dtype=torch.float32, device=self.device)
stage_1_w, stage_1_h = width // 2, height // 2
stage_1_conditionings = self.image_conditioner(
lambda enc: combined_image_conditionings(
@@ -134,7 +136,7 @@ class DistilledPipeline:
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
upscaled_video_latent = self.upsampler(video_state.latent[:1])
stage_2_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
stage_2_sigmas = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
stage_2_conditionings = self.image_conditioner(
lambda enc: combined_image_conditionings(
images=images,
packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py
+9 -7
View File
@@ -32,8 +32,8 @@ from ltx_pipelines.utils.blocks import (
VideoUpsampler,
)
from ltx_pipelines.utils.constants import (
DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMA_VALUES,
DISTILLED_SIGMAS,
STAGE_2_DISTILLED_SIGMAS,
detect_params,
)
from ltx_pipelines.utils.denoisers import SimpleDenoiser
@@ -126,6 +126,8 @@ class ICLoraPipeline:
skip_stage_2: bool = False,
conditioning_attention_mask: torch.Tensor | None = None,
streaming_prefetch_count: int | None = None,
stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
) -> tuple[Iterator[torch.Tensor], Audio]:
"""
Generate video with IC-LoRA conditioning.
@@ -200,7 +202,7 @@ class ICLoraPipeline:
)
)
stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
stage_1_sigmas = stage_1_sigmas.to(dtype=torch.float32, device=self.device)
video_state, audio_state = self.stage_1(
denoiser=SimpleDenoiser(video_context, audio_context),
@@ -230,7 +232,7 @@ class ICLoraPipeline:
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
upscaled_video_latent = self.upsampler(video_state.latent[:1])
distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
stage_2_sigmas = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = self.image_conditioner(
lambda enc: combined_image_conditionings(
@@ -245,7 +247,7 @@ class ICLoraPipeline:
video_state, audio_state = self.stage_2(
denoiser=SimpleDenoiser(video_context, audio_context),
sigmas=distilled_sigmas,
sigmas=stage_2_sigmas,
noiser=noiser,
width=width,
height=height,
@@ -254,12 +256,12 @@ class ICLoraPipeline:
video=ModalitySpec(
context=video_context,
conditionings=stage_2_conditionings,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=upscaled_video_latent,
),
audio=ModalitySpec(
context=audio_context,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=audio_state.latent,
),
streaming_prefetch_count=streaming_prefetch_count,
packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py
+11 -6
View File
@@ -25,7 +25,7 @@ from ltx_pipelines.utils.blocks import (
VideoUpsampler,
)
from ltx_pipelines.utils.constants import (
STAGE_2_DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMAS,
detect_params,
)
from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser, SimpleDenoiser
@@ -62,6 +62,7 @@ class KeyframeInterpolationPipeline:
):
self.device = device or get_device()
self.dtype = torch.bfloat16
self._scheduler = LTX2Scheduler()
self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
@@ -107,6 +108,8 @@ class KeyframeInterpolationPipeline:
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
max_batch_size: int = 1,
stage_1_sigmas: torch.Tensor | None = None,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=True)
@@ -125,7 +128,9 @@ class KeyframeInterpolationPipeline:
v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
# Stage 1: Initial low resolution video generation.
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
sigmas = (
stage_1_sigmas if stage_1_sigmas is not None else self._scheduler.execute(steps=num_inference_steps)
).to(dtype=torch.float32, device=self.device)
stage_1_output_shape = VideoPixelShape(
batch=1,
@@ -181,7 +186,7 @@ class KeyframeInterpolationPipeline:
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
upscaled_video_latent = self.upsampler(video_state.latent[:1])
distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
stage_2_sigmas = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = self.image_conditioner(
lambda enc: image_conditionings_by_adding_guiding_latent(
@@ -196,7 +201,7 @@ class KeyframeInterpolationPipeline:
video_state, audio_state = self.stage_2(
denoiser=SimpleDenoiser(v_context_p, a_context_p),
sigmas=distilled_sigmas,
sigmas=stage_2_sigmas,
noiser=noiser,
width=width,
height=height,
@@ -205,12 +210,12 @@ class KeyframeInterpolationPipeline:
video=ModalitySpec(
context=v_context_p,
conditionings=stage_2_conditionings,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=upscaled_video_latent,
),
audio=ModalitySpec(
context=a_context_p,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=audio_state.latent,
),
streaming_prefetch_count=streaming_prefetch_count,
packages/ltx-pipelines/src/ltx_pipelines/retake.py
+10 -4
View File
@@ -25,7 +25,7 @@ from ltx_pipelines.utils.blocks import (
PromptEncoder,
VideoDecoder,
)
from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES, detect_params
from ltx_pipelines.utils.constants import DISTILLED_SIGMAS, detect_params
from ltx_pipelines.utils.denoisers import GuidedDenoiser, SimpleDenoiser
from ltx_pipelines.utils.helpers import (
audio_latent_from_file,
@@ -78,6 +78,8 @@ class RetakePipeline:
self.device = device or get_device()
self.dtype = torch.bfloat16
self.distilled = distilled
if not distilled:
self._scheduler = LTX2Scheduler()
self.prompt_encoder = PromptEncoder(
checkpoint_path=checkpoint_path,
gemma_root=gemma_root,
@@ -141,6 +143,7 @@ class RetakePipeline:
tiling_config: TilingConfig | None = None,
streaming_prefetch_count: int | None = None,
max_batch_size: int = 1,
sigmas: torch.Tensor | None = None,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
"""Regenerate ``[start_time, end_time]`` of the source video (retake).
Parameters
@@ -227,15 +230,18 @@ class RetakePipeline:
initial_latent=initial_audio_latent,
frozen=initial_audio_latent is not None and not regenerate_audio,
)
# Build denoiser
# Build denoiser and resolve sigma schedule.
if sigmas is None:
sigmas = DISTILLED_SIGMAS if self.distilled else self._scheduler.execute(steps=num_inference_steps)
sigmas = sigmas.to(dtype=torch.float32, device=self.device)
if self.distilled:
sigmas = torch.tensor(DISTILLED_SIGMA_VALUES).to(dtype=torch.float32, device=self.device)
denoiser = SimpleDenoiser(
v_context=v_context_p,
a_context=a_context_p,
)
else:
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
v_context_n, a_context_n = contexts[1].video_encoding, contexts[1].audio_encoding
video_guider = MultiModalGuider(
params=video_guider_params,
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py
+5 -1
View File
@@ -55,6 +55,7 @@ class TI2VidOneStagePipeline:
):
self.dtype = torch.bfloat16
self.device = device or get_device()
self._scheduler = LTX2Scheduler()
self.prompt_encoder = PromptEncoder(
checkpoint_path=checkpoint_path,
gemma_root=gemma_root,
@@ -107,6 +108,7 @@ class TI2VidOneStagePipeline:
streaming_prefetch_count: int | None = None,
tiling_config: TilingConfig | None = None,
max_batch_size: int = 1,
sigmas: torch.Tensor | None = None,
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=False)
@@ -135,7 +137,9 @@ class TI2VidOneStagePipeline:
)
)
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
sigmas = (sigmas if sigmas is not None else self._scheduler.execute(steps=num_inference_steps)).to(
dtype=torch.float32, device=self.device
)
video_guider_factory = create_multimodal_guider_factory(
params=video_guider_params,
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py
+11 -6
View File
@@ -25,7 +25,7 @@ from ltx_pipelines.utils.blocks import (
VideoUpsampler,
)
from ltx_pipelines.utils.constants import (
STAGE_2_DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMAS,
detect_params,
)
from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser, SimpleDenoiser
@@ -61,6 +61,7 @@ class TI2VidTwoStagesPipeline:
):
self.device = device or get_device()
self.dtype = torch.bfloat16
self._scheduler = LTX2Scheduler()
self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
@@ -106,6 +107,8 @@ class TI2VidTwoStagesPipeline:
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
max_batch_size: int = 1,
stage_1_sigmas: torch.Tensor | None = None,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=True)
@@ -142,7 +145,9 @@ class TI2VidTwoStagesPipeline:
)
)
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
sigmas = (
stage_1_sigmas if stage_1_sigmas is not None else self._scheduler.execute(steps=num_inference_steps)
).to(dtype=torch.float32, device=self.device)
video_state, audio_state = self.stage_1(
denoiser=FactoryGuidedDenoiser(
@@ -172,7 +177,7 @@ class TI2VidTwoStagesPipeline:
# Stage 2: Upsample and refine the video at higher resolution with distilled LoRA.
upscaled_video_latent = self.upsampler(video_state.latent[:1])
distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
stage_2_sigmas = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
stage_2_conditionings = self.image_conditioner(
lambda enc: combined_image_conditionings(
images=images,
@@ -186,7 +191,7 @@ class TI2VidTwoStagesPipeline:
video_state, audio_state = self.stage_2(
denoiser=SimpleDenoiser(v_context=v_context_p, a_context=a_context_p),
sigmas=distilled_sigmas,
sigmas=stage_2_sigmas,
noiser=noiser,
width=width,
height=height,
@@ -195,12 +200,12 @@ class TI2VidTwoStagesPipeline:
video=ModalitySpec(
context=v_context_p,
conditionings=stage_2_conditionings,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=upscaled_video_latent,
),
audio=ModalitySpec(
context=a_context_p,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=audio_state.latent,
),
streaming_prefetch_count=streaming_prefetch_count,
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_hq.py
+13 -11
View File
@@ -23,7 +23,7 @@ from ltx_pipelines.utils.blocks import (
)
from ltx_pipelines.utils.constants import (
LTX_2_3_HQ_PARAMS,
STAGE_2_DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMAS,
)
from ltx_pipelines.utils.denoisers import GuidedDenoiser, SimpleDenoiser
from ltx_pipelines.utils.helpers import (
@@ -64,6 +64,7 @@ class TI2VidTwoStagesHQPipeline:
):
self.device = device or get_device()
self.dtype = torch.bfloat16
self._scheduler = LTX2Scheduler()
distilled_lora_stage_1 = LoraPathStrengthAndSDOps(
path=distilled_lora[0].path,
@@ -121,6 +122,8 @@ class TI2VidTwoStagesHQPipeline:
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
max_batch_size: int = 1,
stage_1_sigmas: torch.Tensor | None = None,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=True)
@@ -157,13 +160,12 @@ class TI2VidTwoStagesHQPipeline:
)
)
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)
)
if stage_1_sigmas is None:
empty_latent = torch.empty(VideoLatentShape.from_pixel_shape(stage_1_output_shape).to_torch_shape())
stage_1_sigmas = self._scheduler.execute(latent=empty_latent, steps=num_inference_steps)
sigmas = stage_1_sigmas.to(dtype=torch.float32, device=self.device)
video_state, audio_state = self.stage_1(
denoiser=GuidedDenoiser(
@@ -195,7 +197,7 @@ class TI2VidTwoStagesHQPipeline:
# Stage 2: Upsample and refine the video at higher resolution with distilled LoRA.
upscaled_video_latent = self.upsampler(video_state.latent[:1])
distilled_sigmas = torch.tensor(STAGE_2_DISTILLED_SIGMA_VALUES, device=self.device)
stage_2_sigmas = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = self.image_conditioner(
lambda enc: combined_image_conditionings(
@@ -210,7 +212,7 @@ class TI2VidTwoStagesHQPipeline:
video_state, audio_state = self.stage_2(
denoiser=SimpleDenoiser(v_context=v_context_p, a_context=a_context_p),
sigmas=distilled_sigmas,
sigmas=stage_2_sigmas,
noiser=noiser,
stepper=stepper,
width=width,
@@ -220,12 +222,12 @@ class TI2VidTwoStagesHQPipeline:
video=ModalitySpec(
context=v_context_p,
conditionings=stage_2_conditionings,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=upscaled_video_latent,
),
audio=ModalitySpec(
context=a_context_p,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=audio_state.latent,
),
loop=res2s_audio_video_denoising_loop,
packages/ltx-pipelines/src/ltx_pipelines/utils/constants.py
+4
View File
@@ -1,6 +1,7 @@
import logging
from dataclasses import dataclass, field, replace
import torch
from safetensors import safe_open
from ltx_core.components.guiders import MultiModalGuiderParams
@@ -17,6 +18,9 @@ DISTILLED_SIGMA_VALUES = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725,
# Reduced schedule for super-resolution stage 2 (subset of distilled values)
STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0]
DISTILLED_SIGMAS = torch.tensor(DISTILLED_SIGMA_VALUES)
STAGE_2_DISTILLED_SIGMAS = torch.tensor(STAGE_2_DISTILLED_SIGMA_VALUES)
# =============================================================================
# Pipeline Parameters