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

Automated PR - 2026-03-30

Browse Source
This commit is contained in:
github-actions[bot]
2026-03-30 17:59:34 +00:00
parent ae855f8538
commit f4d0c1ec0e
48 changed files with 8429 additions and 6644 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-pipelines/src/ltx_pipelines/distilled.py
+87 -107
View File
@@ -3,40 +3,38 @@ from collections.abc import Iterator
import torch
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.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.loader.registry import Registry
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.types import Audio, LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger, euler_denoising_loop
from ltx_core.types import Audio
from ltx_pipelines.utils.args import (
ImageConditioningInput,
default_2_stage_distilled_arg_parser,
detect_checkpoint_path,
)
from ltx_pipelines.utils.blocks import (
AudioDecoder,
DiffusionStage,
ImageConditioner,
PromptEncoder,
VideoDecoder,
VideoUpsampler,
)
from ltx_pipelines.utils.constants import (
DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMA_VALUES,
detect_params,
)
from ltx_pipelines.utils.denoisers import SimpleDenoiser
from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
combined_image_conditionings,
denoise_audio_video,
encode_prompts,
get_device,
simple_denoising_func,
)
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import PipelineComponents
device = get_device()
from ltx_pipelines.utils.types import ModalitySpec
class DistilledPipeline:
@@ -52,26 +50,32 @@ class DistilledPipeline:
gemma_root: str,
spatial_upsampler_path: str,
loras: list[LoraPathStrengthAndSDOps],
device: torch.device = device,
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
):
self.device = device
self.device = device or get_device()
self.dtype = torch.bfloat16
self.model_ledger = ModelLedger(
dtype=self.dtype,
device=device,
checkpoint_path=distilled_checkpoint_path,
spatial_upsampler_path=spatial_upsampler_path,
gemma_root_path=gemma_root,
loras=loras,
self.prompt_encoder = PromptEncoder(
distilled_checkpoint_path, gemma_root, self.dtype, self.device, registry=registry
)
self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
self.stage = DiffusionStage(
distilled_checkpoint_path,
self.dtype,
self.device,
loras=tuple(loras),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
)
self.pipeline_components = PipelineComponents(
dtype=self.dtype,
device=device,
self.upsampler = VideoUpsampler(
distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
)
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__(
self,
@@ -84,114 +88,88 @@ class DistilledPipeline:
images: list[ImageConditioningInput],
tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
) -> 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
(ctx_p,) = encode_prompts(
(ctx_p,) = self.prompt_encoder(
[prompt],
self.model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
streaming_prefetch_count=streaming_prefetch_count,
)
video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
# Stage 1: Initial low resolution video generation.
video_encoder = self.model_ledger.video_encoder()
transformer = self.model_ledger.transformer()
stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
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=simple_denoising_func(
video_context=video_context,
audio_context=audio_context,
transformer=transformer, # noqa: F821
),
stage_1_w, stage_1_h = width // 2, height // 2
stage_1_conditionings = self.image_conditioner(
lambda enc: combined_image_conditionings(
images=images,
height=stage_1_h,
width=stage_1_w,
video_encoder=enc,
dtype=dtype,
device=self.device,
)
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
stage_1_conditionings = combined_image_conditionings(
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,
video_state, audio_state = self.stage(
denoiser=SimpleDenoiser(video_context, audio_context),
sigmas=stage_1_sigmas,
stepper=stepper,
denoising_loop_fn=denoising_loop,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
noiser=noiser,
width=stage_1_w,
height=stage_1_h,
frames=num_frames,
fps=frame_rate,
video=ModalitySpec(context=video_context, conditionings=stage_1_conditionings),
audio=ModalitySpec(context=audio_context),
streaming_prefetch_count=streaming_prefetch_count,
)
# 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.model_ledger.spatial_upsampler()
)
torch.cuda.synchronize()
cleanup_memory()
upscaled_video_latent = self.upsampler(video_state.latent[:1])
stage_2_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = combined_image_conditionings(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
stage_2_conditionings = self.image_conditioner(
lambda enc: combined_image_conditionings(
images=images,
height=height,
width=width,
video_encoder=enc,
dtype=dtype,
device=self.device,
)
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_2_output_shape,
conditionings=stage_2_conditionings,
noiser=noiser,
video_state, audio_state = self.stage(
denoiser=SimpleDenoiser(video_context, audio_context),
sigmas=stage_2_sigmas,
stepper=stepper,
denoising_loop_fn=denoising_loop,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
noise_scale=stage_2_sigmas[0],
initial_video_latent=upscaled_video_latent,
initial_audio_latent=audio_state.latent,
noiser=noiser,
width=width,
height=height,
frames=num_frames,
fps=frame_rate,
video=ModalitySpec(
context=video_context,
conditionings=stage_2_conditionings,
noise_scale=stage_2_sigmas[0].item(),
initial_latent=upscaled_video_latent,
),
audio=ModalitySpec(
context=audio_context,
noise_scale=stage_2_sigmas[0].item(),
initial_latent=audio_state.latent,
),
streaming_prefetch_count=streaming_prefetch_count,
)
torch.cuda.synchronize()
del transformer
del video_encoder
cleanup_memory()
decoded_video = vae_decode_video(
video_state.latent, self.model_ledger.video_decoder(), tiling_config, generator
)
decoded_audio = vae_decode_audio(
audio_state.latent, self.model_ledger.audio_decoder(), self.model_ledger.vocoder()
)
decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
decoded_audio = self.audio_decoder(audio_state.latent)
return decoded_video, decoded_audio
@@ -208,6 +186,7 @@ def main() -> None:
gemma_root=args.gemma_root,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
)
tiling_config = TilingConfig.default()
video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -221,6 +200,7 @@ def main() -> None:
images=args.images,
tiling_config=tiling_config,
enhance_prompt=args.enhance_prompt,
streaming_prefetch_count=args.streaming_prefetch_count,
)
encode_video(