Upload 6 files

packages/ltx-pipelines/src/ltx_pipelines/__init__.py

@@ -0,0 +1,26 @@
+"""
+LTX-2 Pipelines: High-level video generation pipelines and utilities.
+This package provides ready-to-use pipelines for video generation:
+- TI2VidOneStagePipeline: Text/image-to-video in a single stage
+- TI2VidTwoStagesPipeline: Two-stage generation with upsampling
+- DistilledPipeline: Fast distilled two-stage generation
+- ICLoraPipeline: Image/video conditioning with distilled LoRA
+- KeyframeInterpolationPipeline: Keyframe-based video interpolation
+- ModelLedger: Central coordinator for loading and building models
+For more detailed components and utilities, import from specific submodules
+like `ltx_pipelines.utils.media_io` or `ltx_pipelines.utils.constants`.
+"""
+
+from ltx_pipelines.distilled import DistilledPipeline
+from ltx_pipelines.ic_lora import ICLoraPipeline
+from ltx_pipelines.keyframe_interpolation import KeyframeInterpolationPipeline
+from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
+from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
+
+__all__ = [
+    "DistilledPipeline",
+    "ICLoraPipeline",
+    "KeyframeInterpolationPipeline",
+    "TI2VidOneStagePipeline",
+    "TI2VidTwoStagesPipeline",
+]

packages/ltx-pipelines/src/ltx_pipelines/distilled.py

@@ -0,0 +1,475 @@
+# Copyright (c) 2025 Lightricks. All rights reserved.
+# Created by Amit Pintz.
+
+
+import torch
+
+from ltx_core.components.diffusion_steps import EulerDiffusionStep
+from ltx_core.loader.sd_ops import LTXV_LORA_COMFY_RENAMING_MAP
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.protocols import DiffusionStepProtocol
+from ltx_core.conditioning import ConditioningItem, VideoConditionByKeyframeIndex, ConditioningError
+from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
+from ltx_core.model.upsampler import upsample_video
+from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
+from ltx_core.model.video_vae import decode_video as vae_decode_video
+from ltx_core.text_encoders.gemma import encode_text
+from ltx_core.types import LatentState, VideoPixelShape
+from ltx_core.tools import LatentTools
+from ltx_pipelines import utils
+from ltx_pipelines.utils import ModelLedger
+from ltx_pipelines.utils.args import default_2_stage_distilled_arg_parser
+from ltx_pipelines.utils.constants import (
+    AUDIO_SAMPLE_RATE,
+    DEFAULT_LORA_STRENGTH,
+    DISTILLED_SIGMA_VALUES,
+    STAGE_2_DISTILLED_SIGMA_VALUES,
+)
+from ltx_pipelines.utils.helpers import (
+    assert_resolution,
+    cleanup_memory,
+    denoise_audio_video,
+    euler_denoising_loop,
+    generate_enhanced_prompt,
+    get_device,
+    image_conditionings_by_replacing_latent,
+    simple_denoising_func,
+)
+from ltx_pipelines.utils.media_io import encode_video, load_video_conditioning
+from ltx_pipelines.utils.types import PipelineComponents
+
+import torchaudio
+from ltx_core.model.audio_vae import AudioProcessor
+from ltx_core.types import AudioLatentShape, VideoPixelShape
+
+class AudioConditionByLatent(ConditioningItem):
+    """
+    Conditions audio generation by injecting a full latent sequence.
+    Replaces tokens in the latent state with the provided audio latents,
+    and sets denoise strength according to the strength parameter.
+    """
+
+    def __init__(self, latent: torch.Tensor, strength: float):
+        self.latent = latent
+        self.strength = strength
+
+    def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
+        if not isinstance(latent_tools.target_shape, AudioLatentShape):
+            raise ConditioningError("Audio conditioning requires an audio latent target shape.")
+
+        cond_batch, cond_channels, cond_frames, cond_bins = self.latent.shape
+        tgt_batch, tgt_channels, tgt_frames, tgt_bins = latent_tools.target_shape.to_torch_shape()
+
+        if (cond_batch, cond_channels, cond_frames, cond_bins) != (tgt_batch, tgt_channels, tgt_frames, tgt_bins):
+            raise ConditioningError(
+                f"Can't apply audio conditioning item to latent with shape {latent_tools.target_shape}, expected "
+                f"shape is ({tgt_batch}, {tgt_channels}, {tgt_frames}, {tgt_bins})."
+            )
+
+        tokens = latent_tools.patchifier.patchify(self.latent)
+        latent_state = latent_state.clone()
+        latent_state.latent[:, : tokens.shape[1]] = tokens
+        latent_state.clean_latent[:, : tokens.shape[1]] = tokens
+        latent_state.denoise_mask[:, : tokens.shape[1]] = 1.0 - self.strength
+
+        return latent_state
+
+device = get_device()
+
+
+class DistilledPipeline:
+    def __init__(
+        self,
+        checkpoint_path: str,
+        gemma_root: str,
+        spatial_upsampler_path: str,
+        loras: list[LoraPathStrengthAndSDOps],
+        device: torch.device = device,
+        fp8transformer: bool = False,
+        local_files_only: bool = True,
+    ):
+        self.device = device
+        self.dtype = torch.bfloat16
+
+        self.model_ledger = ModelLedger(
+            dtype=self.dtype,
+            device=device,
+            checkpoint_path=checkpoint_path,
+            spatial_upsampler_path=spatial_upsampler_path,
+            gemma_root_path=gemma_root,
+            loras=loras,
+            fp8transformer=fp8transformer,
+            local_files_only=local_files_only
+        )
+
+        self.pipeline_components = PipelineComponents(
+            dtype=self.dtype,
+            device=device,
+        )
+
+        # Cached models to avoid reloading
+        self._video_encoder = None
+        self._transformer = None
+        
+    def _build_audio_conditionings_from_waveform(
+        self,
+        input_waveform: torch.Tensor,
+        input_sample_rate: int,
+        num_frames: int,
+        fps: float,
+        strength: float,
+    ) -> list[AudioConditionByLatent] | None:
+        strength = float(strength)
+        if strength <= 0.0:
+            return None
+
+        # Expect waveform as:
+        #  - (T,) or (C,T) or (B,C,T). Convert to (B,C,T)
+        waveform = input_waveform
+        if waveform.ndim == 1:
+            waveform = waveform.unsqueeze(0).unsqueeze(0)
+        elif waveform.ndim == 2:
+            waveform = waveform.unsqueeze(0)
+        elif waveform.ndim != 3:
+            raise ValueError(f"input_waveform must be 1D/2D/3D, got shape {tuple(waveform.shape)}")
+
+        # Get audio encoder + its config
+        audio_encoder = self.model_ledger.audio_encoder()  # assumes ledger exposes it
+        # If you want to cache it like video_encoder/transformer, you can.
+        target_sr = int(getattr(audio_encoder, "sample_rate"))
+        target_channels = int(getattr(audio_encoder, "in_channels", waveform.shape[1]))
+        mel_bins = int(getattr(audio_encoder, "mel_bins"))
+        mel_hop = int(getattr(audio_encoder, "mel_hop_length"))
+        n_fft = int(getattr(audio_encoder, "n_fft"))
+
+        # Match channels
+        if waveform.shape[1] != target_channels:
+            if waveform.shape[1] == 1 and target_channels > 1:
+                waveform = waveform.repeat(1, target_channels, 1)
+            elif target_channels == 1:
+                waveform = waveform.mean(dim=1, keepdim=True)
+            else:
+                waveform = waveform[:, :target_channels, :]
+                if waveform.shape[1] < target_channels:
+                    pad_ch = target_channels - waveform.shape[1]
+                    pad = torch.zeros((waveform.shape[0], pad_ch, waveform.shape[2]), dtype=waveform.dtype)
+                    waveform = torch.cat([waveform, pad], dim=1)
+
+        # Resample if needed (CPU float32 is safest for torchaudio)
+        waveform = waveform.to(device="cpu", dtype=torch.float32)
+        if int(input_sample_rate) != target_sr:
+            waveform = torchaudio.functional.resample(waveform, int(input_sample_rate), target_sr)
+
+        # Waveform -> Mel
+        audio_processor = AudioProcessor(
+            sample_rate=target_sr,
+            mel_bins=mel_bins,
+            mel_hop_length=mel_hop,
+            n_fft=n_fft,
+        ).to(waveform.device)
+
+        mel = audio_processor.waveform_to_mel(waveform, target_sr)
+
+        # Mel -> latent (run encoder on its own device/dtype)
+        audio_params = next(audio_encoder.parameters(), None)
+        enc_device = audio_params.device if audio_params is not None else self.device
+        enc_dtype = audio_params.dtype if audio_params is not None else self.dtype
+
+        mel = mel.to(device=enc_device, dtype=enc_dtype)
+        with torch.inference_mode():
+            audio_latent = audio_encoder(mel)
+
+        # Pad/trim latent to match the target video duration
+        audio_downsample = getattr(getattr(audio_encoder, "patchifier", None), "audio_latent_downsample_factor", 4)
+        target_shape = AudioLatentShape.from_video_pixel_shape(
+            VideoPixelShape(batch=audio_latent.shape[0], frames=int(num_frames), width=1, height=1, fps=float(fps)),
+            channels=audio_latent.shape[1],
+            mel_bins=audio_latent.shape[3],
+            sample_rate=target_sr,
+            hop_length=mel_hop,
+            audio_latent_downsample_factor=audio_downsample,
+        )
+        target_frames = int(target_shape.frames)
+
+        if audio_latent.shape[2] < target_frames:
+            pad_frames = target_frames - audio_latent.shape[2]
+            pad = torch.zeros(
+                (audio_latent.shape[0], audio_latent.shape[1], pad_frames, audio_latent.shape[3]),
+                device=audio_latent.device,
+                dtype=audio_latent.dtype,
+            )
+            audio_latent = torch.cat([audio_latent, pad], dim=2)
+        elif audio_latent.shape[2] > target_frames:
+            audio_latent = audio_latent[:, :, :target_frames, :]
+
+        # Move latent to pipeline device/dtype for conditioning object
+        audio_latent = audio_latent.to(device=self.device, dtype=self.dtype)
+
+        return [AudioConditionByLatent(audio_latent, strength)]
+
+    def _prepare_output_waveform(
+        self,
+        input_waveform: torch.Tensor,
+        input_sample_rate: int,
+        target_sample_rate: int,
+        num_frames: int,
+        fps: float,
+    ) -> torch.Tensor:
+        """
+        Returns waveform on CPU, float32, resampled to target_sample_rate and
+        trimmed/padded to match video duration.
+        Output shape: (T,) for mono or (C, T) for multi-channel.
+        """
+        wav = input_waveform
+
+        # Accept (T,), (C,T), (B,C,T)
+        if wav.ndim == 3:
+            wav = wav[0]
+        elif wav.ndim == 2:
+            pass
+        elif wav.ndim == 1:
+            wav = wav.unsqueeze(0)
+        else:
+            raise ValueError(f"input_waveform must be 1D/2D/3D, got {tuple(wav.shape)}")
+
+        # Now wav is (C, T)
+        wav = wav.detach().to("cpu", dtype=torch.float32)
+
+        # Resample if needed
+        if int(input_sample_rate) != int(target_sample_rate):
+            wav = torchaudio.functional.resample(wav, int(input_sample_rate), int(target_sample_rate))
+
+        # Match video duration
+        duration_sec = float(num_frames) / float(fps)
+        target_len = int(round(duration_sec * float(target_sample_rate)))
+
+        cur_len = int(wav.shape[-1])
+        if cur_len > target_len:
+            wav = wav[..., :target_len]
+        elif cur_len < target_len:
+            pad = target_len - cur_len
+            wav = torch.nn.functional.pad(wav, (0, pad))
+
+        # If mono, return (T,) for convenience
+        if wav.shape[0] == 1:
+            return wav[0]
+        return wav
+
+
+    @torch.inference_mode()
+    def __call__(
+        self,
+        prompt: str,
+        output_path: str,
+        seed: int,
+        height: int,
+        width: int,
+        num_frames: int,
+        frame_rate: float,
+        images: list[tuple[str, int, float]],
+        video_conditioning: list[tuple[str, float]] | None = None,
+        video_conditioning_frame_idx: int = 0,
+        apply_video_conditioning_to_stage2: bool = False,
+        tiling_config: TilingConfig | None = None,
+        video_context: torch.Tensor | None = None,
+        audio_context: torch.Tensor | None = None,
+        input_waveform: torch.Tensor | None = None,
+        input_waveform_sample_rate: int | None = None,
+        audio_strength: float = 1.0,   # or audio_scale, your naming
+    ) -> None:
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        noiser = GaussianNoiser(generator=generator)
+        stepper = EulerDiffusionStep()
+        dtype = torch.bfloat16
+
+        audio_conditionings = None
+        if input_waveform is not None:
+            if input_waveform_sample_rate is None:
+                raise ValueError("input_waveform_sample_rate must be provided when input_waveform is set.")
+            audio_conditionings = self._build_audio_conditionings_from_waveform(
+                input_waveform=input_waveform,
+                input_sample_rate=int(input_waveform_sample_rate),
+                num_frames=num_frames,
+                fps=frame_rate,
+                strength=audio_strength,
+            )
+
+        # Use pre-computed embeddings if provided, otherwise encode text
+        if video_context is None or audio_context is None:
+            text_encoder = self.model_ledger.text_encoder()
+            context_p = encode_text(text_encoder, prompts=[prompt])[0]
+            video_context, audio_context = context_p
+
+            torch.cuda.synchronize()
+            del text_encoder
+            utils.cleanup_memory()
+        else:
+            # Move pre-computed embeddings to device if needed
+            video_context = video_context.to(self.device)
+            audio_context = audio_context.to(self.device)
+
+        # Stage 1: Initial low resolution video generation.
+        # Load models only if not already cached
+        if self._video_encoder is None:
+            self._video_encoder = self.model_ledger.video_encoder()
+        video_encoder = self._video_encoder
+
+        if self._transformer is None:
+            self._transformer = self.model_ledger.transformer()
+        transformer = self._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_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
+        stage_1_conditionings = self._create_conditionings(
+            images=images,
+            video_conditioning=video_conditioning,
+            height=stage_1_output_shape.height,
+            width=stage_1_output_shape.width,
+            num_frames=num_frames,
+            video_encoder=video_encoder,
+            video_conditioning_frame_idx=video_conditioning_frame_idx,
+            dtype=dtype,
+        )
+
+        video_state, audio_state = denoise_audio_video(
+            output_shape=stage_1_output_shape,
+            conditionings=stage_1_conditionings,
+            audio_conditionings=audio_conditionings,
+            noiser=noiser,
+            sigmas=stage_1_sigmas,
+            stepper=stepper,
+            denoising_loop_fn=denoising_loop,
+            components=self.pipeline_components,
+            dtype=dtype,
+            device=self.device,
+        )
+
+        # 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()
+
+        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 * 2, height=height * 2, fps=frame_rate
+        )
+        if apply_video_conditioning_to_stage2:
+            stage_2_conditionings = self._create_conditionings(
+                images=images,
+                video_conditioning=video_conditioning,
+                height=stage_2_output_shape.height,
+                width=stage_2_output_shape.width,
+                num_frames=num_frames,
+                video_encoder=video_encoder,
+                video_conditioning_frame_idx=video_conditioning_frame_idx,
+            )
+        else:
+            stage_2_conditionings = image_conditionings_by_replacing_latent(
+                images=images,
+                height=stage_2_output_shape.height,
+                width=stage_2_output_shape.width,
+                video_encoder=video_encoder,
+                dtype=dtype,
+                device=self.device,
+            )
+        video_chunks_number = get_video_chunks_number(num_frames, tiling_config)
+        video_state, audio_state = denoise_audio_video(
+            output_shape=stage_2_output_shape,
+            conditionings=stage_2_conditionings,
+            audio_conditionings=audio_conditionings,
+            noiser=noiser,
+            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,
+        )
+
+        torch.cuda.synchronize()
+        # del transformer
+        # del video_encoder
+        # utils.cleanup_memory()
+
+        decoded_video = vae_decode_video(video_state.latent, self.model_ledger.video_decoder(), tiling_config)
+        decoded_audio = vae_decode_audio(audio_state.latent, self.model_ledger.audio_decoder(), self.model_ledger.vocoder())
+
+        encode_video(
+            video=decoded_video,
+            fps=frame_rate,
+            audio=decoded_audio,
+            audio_sample_rate=AUDIO_SAMPLE_RATE,
+            output_path=output_path,
+            video_chunks_number=video_chunks_number,
+        )
+
+
+    def _create_conditionings(
+        self,
+        images: list[tuple[str, int, float]],
+        video_conditioning: list[tuple[str, float]] | None,
+        height: int,
+        width: int,
+        num_frames: int,
+        video_encoder,
+        video_conditioning_frame_idx: int,
+        dtype: torch.dtype,
+    ):
+        # 1) Keep ORIGINAL behavior: image conditioning by replacing latent
+        conditionings = image_conditionings_by_replacing_latent(
+            images=images,
+            height=height,
+            width=width,
+            video_encoder=video_encoder,
+            dtype=dtype,
+            device=self.device,
+        )
+
+        # 2) Optional: add video conditioning (IC-LoRA style)
+        if not video_conditioning:
+            return conditionings
+
+        for video_path, strength in video_conditioning:
+            video = load_video_conditioning(
+                video_path=video_path,
+                height=height,
+                width=width,
+                frame_cap=num_frames,   # ✅ correct kwarg name
+                dtype=dtype,
+                device=self.device,
+            )
+
+            encoded_video = video_encoder(video)
+
+            # ✅ match IC-LoRA: append the conditioning object directly
+            conditionings.append(
+                VideoConditionByKeyframeIndex(
+                    keyframes=encoded_video,
+                    frame_idx=video_conditioning_frame_idx,
+                    strength=strength,
+                )
+            )
+
+        return conditionings

packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py

@@ -0,0 +1,309 @@
+import logging
+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.conditioning import ConditioningItem, VideoConditionByKeyframeIndex
+from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
+from ltx_core.model.upsampler import upsample_video
+from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
+from ltx_core.model.video_vae import decode_video as vae_decode_video
+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 (
+    assert_resolution,
+    cleanup_memory,
+    denoise_audio_video,
+    euler_denoising_loop,
+    generate_enhanced_prompt,
+    get_device,
+    image_conditionings_by_replacing_latent,
+    simple_denoising_func,
+)
+from ltx_pipelines.utils.media_io import encode_video, load_video_conditioning
+from ltx_pipelines.utils.types import PipelineComponents
+
+device = get_device()
+
+
+class ICLoraPipeline:
+    """
+    Two-stage video generation pipeline with In-Context (IC) LoRA support.
+    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
+    by 2x and refines with additional denoising steps for higher quality output.
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        spatial_upsampler_path: str,
+        gemma_root: str,
+        loras: list[LoraPathStrengthAndSDOps],
+        device: torch.device = device,
+        fp8transformer: bool = False,
+    ):
+        self.dtype = torch.bfloat16
+        self.stage_1_model_ledger = ModelLedger(
+            dtype=self.dtype,
+            device=device,
+            checkpoint_path=checkpoint_path,
+            spatial_upsampler_path=spatial_upsampler_path,
+            gemma_root_path=gemma_root,
+            loras=loras,
+            fp8transformer=fp8transformer,
+        )
+        self.stage_2_model_ledger = ModelLedger(
+            dtype=self.dtype,
+            device=device,
+            checkpoint_path=checkpoint_path,
+            spatial_upsampler_path=spatial_upsampler_path,
+            gemma_root_path=gemma_root,
+            loras=[],
+            fp8transformer=fp8transformer,
+        )
+        self.pipeline_components = PipelineComponents(
+            dtype=self.dtype,
+            device=device,
+        )
+        self.device = device
+
+    @torch.inference_mode()
+    def __call__(
+        self,
+        prompt: str,
+        seed: int,
+        height: int,
+        width: int,
+        num_frames: int,
+        frame_rate: float,
+        images: list[tuple[str, int, float]],
+        video_conditioning: list[tuple[str, float]],
+        enhance_prompt: bool = False,
+        tiling_config: TilingConfig | None = None,
+    ) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
+        assert_resolution(height=height, width=width, is_two_stage=True)
+
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        noiser = GaussianNoiser(generator=generator)
+        stepper = EulerDiffusionStep()
+        dtype = torch.bfloat16
+
+        text_encoder = self.stage_1_model_ledger.text_encoder()
+
+        if enhance_prompt:
+            prompt = generate_enhanced_prompt(
+                text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
+            )
+        video_context, audio_context = encode_text(text_encoder, prompts=[prompt])[0]
+
+        torch.cuda.synchronize()
+        del text_encoder
+        cleanup_memory()
+
+        # Stage 1: Initial low resolution video generation.
+        video_encoder = self.stage_1_model_ledger.video_encoder()
+        transformer = self.stage_1_model_ledger.transformer()
+        stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
+
+        def first_stage_denoising_loop(
+            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
+        ) -> tuple[LatentState, LatentState]:
+            return euler_denoising_loop(
+                sigmas=sigmas,
+                video_state=video_state,
+                audio_state=audio_state,
+                stepper=stepper,
+                denoise_fn=simple_denoising_func(
+                    video_context=video_context,
+                    audio_context=audio_context,
+                    transformer=transformer,  # noqa: F821
+                ),
+            )
+
+        stage_1_output_shape = VideoPixelShape(
+            batch=1,
+            frames=num_frames,
+            width=width // 2,
+            height=height // 2,
+            fps=frame_rate,
+        )
+        stage_1_conditionings = self._create_conditionings(
+            images=images,
+            video_conditioning=video_conditioning,
+            height=stage_1_output_shape.height,
+            width=stage_1_output_shape.width,
+            video_encoder=video_encoder,
+            num_frames=num_frames,
+        )
+        video_state, audio_state = denoise_audio_video(
+            output_shape=stage_1_output_shape,
+            conditionings=stage_1_conditionings,
+            noiser=noiser,
+            sigmas=stage_1_sigmas,
+            stepper=stepper,
+            denoising_loop_fn=first_stage_denoising_loop,
+            components=self.pipeline_components,
+            dtype=dtype,
+            device=self.device,
+        )
+
+        torch.cuda.synchronize()
+        del transformer
+        cleanup_memory()
+
+        # Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
+        upscaled_video_latent = upsample_video(
+            latent=video_state.latent[:1],
+            video_encoder=video_encoder,
+            upsampler=self.stage_2_model_ledger.spatial_upsampler(),
+        )
+
+        torch.cuda.synchronize()
+        cleanup_memory()
+
+        transformer = self.stage_2_model_ledger.transformer()
+        distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
+
+        def second_stage_denoising_loop(
+            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
+        ) -> tuple[LatentState, LatentState]:
+            return euler_denoising_loop(
+                sigmas=sigmas,
+                video_state=video_state,
+                audio_state=audio_state,
+                stepper=stepper,
+                denoise_fn=simple_denoising_func(
+                    video_context=video_context,
+                    audio_context=audio_context,
+                    transformer=transformer,  # noqa: F821
+                ),
+            )
+
+        stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
+        stage_2_conditionings = image_conditionings_by_replacing_latent(
+            images=images,
+            height=stage_2_output_shape.height,
+            width=stage_2_output_shape.width,
+            video_encoder=video_encoder,
+            dtype=self.dtype,
+            device=self.device,
+        )
+
+        video_state, audio_state = denoise_audio_video(
+            output_shape=stage_2_output_shape,
+            conditionings=stage_2_conditionings,
+            noiser=noiser,
+            sigmas=distilled_sigmas,
+            stepper=stepper,
+            denoising_loop_fn=second_stage_denoising_loop,
+            components=self.pipeline_components,
+            dtype=dtype,
+            device=self.device,
+            noise_scale=distilled_sigmas[0],
+            initial_video_latent=upscaled_video_latent,
+            initial_audio_latent=audio_state.latent,
+        )
+
+        torch.cuda.synchronize()
+        del transformer
+        del video_encoder
+        cleanup_memory()
+
+        decoded_video = vae_decode_video(video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config)
+        decoded_audio = vae_decode_audio(
+            audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder()
+        )
+        return decoded_video, decoded_audio
+
+    def _create_conditionings(
+        self,
+        images: list[tuple[str, int, float]],
+        video_conditioning: list[tuple[str, float]],
+        height: int,
+        width: int,
+        num_frames: int,
+        video_encoder: VideoEncoder,
+    ) -> list[ConditioningItem]:
+        conditionings = image_conditionings_by_replacing_latent(
+            images=images,
+            height=height,
+            width=width,
+            video_encoder=video_encoder,
+            dtype=self.dtype,
+            device=self.device,
+        )
+
+        for video_path, strength in video_conditioning:
+            video = load_video_conditioning(
+                video_path=video_path,
+                height=height,
+                width=width,
+                frame_cap=num_frames,
+                dtype=self.dtype,
+                device=self.device,
+            )
+            encoded_video = video_encoder(video)
+            conditionings.append(VideoConditionByKeyframeIndex(keyframes=encoded_video, frame_idx=0, strength=strength))
+
+        return conditionings
+
+
+@torch.inference_mode()
+def main() -> None:
+    logging.getLogger().setLevel(logging.INFO)
+    parser = default_2_stage_distilled_arg_parser()
+    parser.add_argument(
+        "--video-conditioning",
+        action=VideoConditioningAction,
+        nargs=2,
+        metavar=("PATH", "STRENGTH"),
+        required=True,
+    )
+    args = parser.parse_args()
+    pipeline = ICLoraPipeline(
+        checkpoint_path=args.checkpoint_path,
+        spatial_upsampler_path=args.spatial_upsampler_path,
+        gemma_root=args.gemma_root,
+        loras=args.lora,
+        fp8transformer=args.enable_fp8,
+    )
+    tiling_config = TilingConfig.default()
+    video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
+    video, audio = pipeline(
+        prompt=args.prompt,
+        seed=args.seed,
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        frame_rate=args.frame_rate,
+        images=args.images,
+        video_conditioning=args.video_conditioning,
+        tiling_config=tiling_config,
+    )
+
+    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,
+    )
+
+
+if __name__ == "__main__":
+    main()

packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py

@@ -0,0 +1,273 @@
+import logging
+from collections.abc import Iterator
+
+import torch
+
+from ltx_core.components.diffusion_steps import EulerDiffusionStep
+from ltx_core.components.guiders import CFGGuider
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.protocols import DiffusionStepProtocol
+from ltx_core.components.schedulers import LTX2Scheduler
+from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
+from ltx_core.model.upsampler import upsample_video
+from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
+from ltx_core.model.video_vae import decode_video as vae_decode_video
+from ltx_core.text_encoders.gemma import encode_text
+from ltx_core.types import LatentState, VideoPixelShape
+from ltx_pipelines.utils import ModelLedger
+from ltx_pipelines.utils.args import default_2_stage_arg_parser
+from ltx_pipelines.utils.constants import (
+    AUDIO_SAMPLE_RATE,
+    STAGE_2_DISTILLED_SIGMA_VALUES,
+)
+from ltx_pipelines.utils.helpers import (
+    assert_resolution,
+    cleanup_memory,
+    denoise_audio_video,
+    euler_denoising_loop,
+    generate_enhanced_prompt,
+    get_device,
+    guider_denoising_func,
+    image_conditionings_by_adding_guiding_latent,
+    simple_denoising_func,
+)
+from ltx_pipelines.utils.media_io import encode_video
+from ltx_pipelines.utils.types import PipelineComponents
+
+device = get_device()
+
+
+class KeyframeInterpolationPipeline:
+    """
+    Keyframe-based Two-stage video interpolation pipeline.
+    Interpolates between keyframes to generate a video with smoother transitions.
+    Stage 1 generates video at the target resolution, then Stage 2 upsamples
+    by 2x and refines with additional denoising steps for higher quality output.
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        distilled_lora: list[LoraPathStrengthAndSDOps],
+        spatial_upsampler_path: str,
+        gemma_root: str,
+        loras: list[LoraPathStrengthAndSDOps],
+        device: torch.device = device,
+        fp8transformer: bool = False,
+    ):
+        self.device = device
+        self.dtype = torch.bfloat16
+        self.stage_1_model_ledger = ModelLedger(
+            dtype=self.dtype,
+            device=device,
+            checkpoint_path=checkpoint_path,
+            spatial_upsampler_path=spatial_upsampler_path,
+            gemma_root_path=gemma_root,
+            loras=loras,
+            fp8transformer=fp8transformer,
+        )
+        self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
+            loras=distilled_lora,
+        )
+        self.pipeline_components = PipelineComponents(
+            dtype=self.dtype,
+            device=device,
+        )
+
+    @torch.inference_mode()
+    def __call__(  # noqa: PLR0913
+        self,
+        prompt: str,
+        negative_prompt: str,
+        seed: int,
+        height: int,
+        width: int,
+        num_frames: int,
+        frame_rate: float,
+        num_inference_steps: int,
+        cfg_guidance_scale: float,
+        images: list[tuple[str, int, float]],
+        tiling_config: TilingConfig | None = None,
+        enhance_prompt: bool = False,
+    ) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
+        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()
+        cfg_guider = CFGGuider(cfg_guidance_scale)
+        dtype = torch.bfloat16
+
+        text_encoder = self.stage_1_model_ledger.text_encoder()
+        if enhance_prompt:
+            prompt = generate_enhanced_prompt(
+                text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
+            )
+        context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
+        v_context_p, a_context_p = context_p
+        v_context_n, a_context_n = context_n
+
+        torch.cuda.synchronize()
+        del text_encoder
+        cleanup_memory()
+
+        # Stage 1: Initial low resolution video generation.
+        video_encoder = self.stage_1_model_ledger.video_encoder()
+        transformer = self.stage_1_model_ledger.transformer()
+        sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
+
+        def first_stage_denoising_loop(
+            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
+        ) -> tuple[LatentState, LatentState]:
+            return euler_denoising_loop(
+                sigmas=sigmas,
+                video_state=video_state,
+                audio_state=audio_state,
+                stepper=stepper,
+                denoise_fn=guider_denoising_func(
+                    cfg_guider,
+                    v_context_p,
+                    v_context_n,
+                    a_context_p,
+                    a_context_n,
+                    transformer=transformer,  # noqa: F821
+                ),
+            )
+
+        stage_1_output_shape = VideoPixelShape(
+            batch=1,
+            frames=num_frames,
+            width=width // 2,
+            height=height // 2,
+            fps=frame_rate,
+        )
+        stage_1_conditionings = image_conditionings_by_adding_guiding_latent(
+            images=images,
+            height=stage_1_output_shape.height,
+            width=stage_1_output_shape.width,
+            video_encoder=video_encoder,
+            dtype=dtype,
+            device=self.device,
+        )
+        video_state, audio_state = denoise_audio_video(
+            output_shape=stage_1_output_shape,
+            conditionings=stage_1_conditionings,
+            noiser=noiser,
+            sigmas=sigmas,
+            stepper=stepper,
+            denoising_loop_fn=first_stage_denoising_loop,
+            components=self.pipeline_components,
+            dtype=dtype,
+            device=self.device,
+        )
+
+        torch.cuda.synchronize()
+        del transformer
+        cleanup_memory()
+
+        # Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
+        upscaled_video_latent = upsample_video(
+            latent=video_state.latent[:1],
+            video_encoder=video_encoder,
+            upsampler=self.stage_2_model_ledger.spatial_upsampler(),
+        )
+
+        torch.cuda.synchronize()
+        cleanup_memory()
+
+        transformer = self.stage_2_model_ledger.transformer()
+        distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
+
+        def second_stage_denoising_loop(
+            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
+        ) -> tuple[LatentState, LatentState]:
+            return euler_denoising_loop(
+                sigmas=sigmas,
+                video_state=video_state,
+                audio_state=audio_state,
+                stepper=stepper,
+                denoise_fn=simple_denoising_func(
+                    video_context=v_context_p,
+                    audio_context=a_context_p,
+                    transformer=transformer,  # noqa: F821
+                ),
+            )
+
+        stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
+        stage_2_conditionings = image_conditionings_by_adding_guiding_latent(
+            images=images,
+            height=stage_2_output_shape.height,
+            width=stage_2_output_shape.width,
+            video_encoder=video_encoder,
+            dtype=dtype,
+            device=self.device,
+        )
+        video_state, audio_state = denoise_audio_video(
+            output_shape=stage_2_output_shape,
+            conditionings=stage_2_conditionings,
+            noiser=noiser,
+            sigmas=distilled_sigmas,
+            stepper=stepper,
+            denoising_loop_fn=second_stage_denoising_loop,
+            components=self.pipeline_components,
+            dtype=dtype,
+            device=self.device,
+            noise_scale=distilled_sigmas[0],
+            initial_video_latent=upscaled_video_latent,
+            initial_audio_latent=audio_state.latent,
+        )
+
+        torch.cuda.synchronize()
+        del transformer
+        del video_encoder
+        cleanup_memory()
+
+        decoded_video = vae_decode_video(video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config)
+        decoded_audio = vae_decode_audio(
+            audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder()
+        )
+        return decoded_video, decoded_audio
+
+
+@torch.inference_mode()
+def main() -> None:
+    logging.getLogger().setLevel(logging.INFO)
+    parser = default_2_stage_arg_parser()
+    args = parser.parse_args()
+    pipeline = KeyframeInterpolationPipeline(
+        checkpoint_path=args.checkpoint_path,
+        distilled_lora=args.distilled_lora,
+        spatial_upsampler_path=args.spatial_upsampler_path,
+        gemma_root=args.gemma_root,
+        loras=args.lora,
+        fp8transformer=args.enable_fp8,
+    )
+    tiling_config = TilingConfig.default()
+    video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
+    video, audio = pipeline(
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        seed=args.seed,
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        frame_rate=args.frame_rate,
+        num_inference_steps=args.num_inference_steps,
+        cfg_guidance_scale=args.cfg_guidance_scale,
+        images=args.images,
+        tiling_config=tiling_config,
+    )
+
+    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,
+    )
+
+
+if __name__ == "__main__":
+    main()

packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py

@@ -0,0 +1,193 @@
+import logging
+from collections.abc import Iterator
+
+import torch
+
+from ltx_core.components.diffusion_steps import EulerDiffusionStep
+from ltx_core.components.guiders import CFGGuider
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.protocols import DiffusionStepProtocol
+from ltx_core.components.schedulers import LTX2Scheduler
+from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
+from ltx_core.model.video_vae import decode_video as vae_decode_video
+from ltx_core.text_encoders.gemma import encode_text
+from ltx_core.types import LatentState, VideoPixelShape
+from ltx_pipelines.utils import ModelLedger
+from ltx_pipelines.utils.args import default_1_stage_arg_parser
+from ltx_pipelines.utils.constants import AUDIO_SAMPLE_RATE
+from ltx_pipelines.utils.helpers import (
+    assert_resolution,
+    cleanup_memory,
+    denoise_audio_video,
+    euler_denoising_loop,
+    generate_enhanced_prompt,
+    get_device,
+    guider_denoising_func,
+    image_conditionings_by_replacing_latent,
+)
+from ltx_pipelines.utils.media_io import encode_video
+from ltx_pipelines.utils.types import PipelineComponents
+
+device = get_device()
+
+
+class TI2VidOneStagePipeline:
+    """
+    Single-stage text/image-to-video generation pipeline.
+    Generates video at the target resolution in a single diffusion pass with
+    classifier-free guidance (CFG). Supports optional image conditioning via
+    the images parameter.
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        gemma_root: str,
+        loras: list[LoraPathStrengthAndSDOps],
+        device: torch.device = device,
+        fp8transformer: bool = False,
+    ):
+        self.dtype = torch.bfloat16
+        self.device = device
+        self.model_ledger = ModelLedger(
+            dtype=self.dtype,
+            device=device,
+            checkpoint_path=checkpoint_path,
+            gemma_root_path=gemma_root,
+            loras=loras,
+            fp8transformer=fp8transformer,
+        )
+        self.pipeline_components = PipelineComponents(
+            dtype=self.dtype,
+            device=device,
+        )
+
+    def __call__(  # noqa: PLR0913
+        self,
+        prompt: str,
+        negative_prompt: str,
+        seed: int,
+        height: int,
+        width: int,
+        num_frames: int,
+        frame_rate: float,
+        num_inference_steps: int,
+        cfg_guidance_scale: float,
+        images: list[tuple[str, int, float]],
+        enhance_prompt: bool = False,
+    ) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
+        assert_resolution(height=height, width=width, is_two_stage=False)
+
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        noiser = GaussianNoiser(generator=generator)
+        stepper = EulerDiffusionStep()
+        cfg_guider = CFGGuider(cfg_guidance_scale)
+        dtype = torch.bfloat16
+
+        text_encoder = self.model_ledger.text_encoder()
+        if enhance_prompt:
+            prompt = generate_enhanced_prompt(
+                text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
+            )
+        context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
+        v_context_p, a_context_p = context_p
+        v_context_n, a_context_n = context_n
+
+        torch.cuda.synchronize()
+        del text_encoder
+        cleanup_memory()
+
+        # Stage 1: Initial low resolution video generation.
+        video_encoder = self.model_ledger.video_encoder()
+        transformer = self.model_ledger.transformer()
+        sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
+
+        def first_stage_denoising_loop(
+            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
+        ) -> tuple[LatentState, LatentState]:
+            return euler_denoising_loop(
+                sigmas=sigmas,
+                video_state=video_state,
+                audio_state=audio_state,
+                stepper=stepper,
+                denoise_fn=guider_denoising_func(
+                    cfg_guider,
+                    v_context_p,
+                    v_context_n,
+                    a_context_p,
+                    a_context_n,
+                    transformer=transformer,  # noqa: F821
+                ),
+            )
+
+        stage_1_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
+        stage_1_conditionings = image_conditionings_by_replacing_latent(
+            images=images,
+            height=stage_1_output_shape.height,
+            width=stage_1_output_shape.width,
+            video_encoder=video_encoder,
+            dtype=dtype,
+            device=self.device,
+        )
+
+        video_state, audio_state = denoise_audio_video(
+            output_shape=stage_1_output_shape,
+            conditionings=stage_1_conditionings,
+            noiser=noiser,
+            sigmas=sigmas,
+            stepper=stepper,
+            denoising_loop_fn=first_stage_denoising_loop,
+            components=self.pipeline_components,
+            dtype=dtype,
+            device=self.device,
+        )
+
+        torch.cuda.synchronize()
+        del transformer
+        cleanup_memory()
+
+        decoded_video = vae_decode_video(video_state.latent, self.model_ledger.video_decoder())
+        decoded_audio = vae_decode_audio(
+            audio_state.latent, self.model_ledger.audio_decoder(), self.model_ledger.vocoder()
+        )
+
+        return decoded_video, decoded_audio
+
+
+@torch.inference_mode()
+def main() -> None:
+    logging.getLogger().setLevel(logging.INFO)
+    parser = default_1_stage_arg_parser()
+    args = parser.parse_args()
+    pipeline = TI2VidOneStagePipeline(
+        checkpoint_path=args.checkpoint_path,
+        gemma_root=args.gemma_root,
+        loras=args.lora,
+        fp8transformer=args.enable_fp8,
+    )
+    video, audio = pipeline(
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        seed=args.seed,
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        frame_rate=args.frame_rate,
+        num_inference_steps=args.num_inference_steps,
+        cfg_guidance_scale=args.cfg_guidance_scale,
+        images=args.images,
+    )
+
+    encode_video(
+        video=video,
+        fps=args.frame_rate,
+        audio=audio,
+        audio_sample_rate=AUDIO_SAMPLE_RATE,
+        output_path=args.output_path,
+        video_chunks_number=1,
+    )
+
+
+if __name__ == "__main__":
+    main()

packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py

@@ -0,0 +1,276 @@
+import logging
+from collections.abc import Iterator
+
+import torch
+
+from ltx_core.components.diffusion_steps import EulerDiffusionStep
+from ltx_core.components.guiders import CFGGuider
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.protocols import DiffusionStepProtocol
+from ltx_core.components.schedulers import LTX2Scheduler
+from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
+from ltx_core.model.upsampler import upsample_video
+from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
+from ltx_core.model.video_vae import decode_video as vae_decode_video
+from ltx_core.text_encoders.gemma import encode_text
+from ltx_core.types import LatentState, VideoPixelShape
+from ltx_pipelines.utils import ModelLedger
+from ltx_pipelines.utils.args import default_2_stage_arg_parser
+from ltx_pipelines.utils.constants import (
+    AUDIO_SAMPLE_RATE,
+    STAGE_2_DISTILLED_SIGMA_VALUES,
+)
+from ltx_pipelines.utils.helpers import (
+    assert_resolution,
+    cleanup_memory,
+    denoise_audio_video,
+    euler_denoising_loop,
+    generate_enhanced_prompt,
+    get_device,
+    guider_denoising_func,
+    image_conditionings_by_replacing_latent,
+    simple_denoising_func,
+)
+from ltx_pipelines.utils.media_io import encode_video
+from ltx_pipelines.utils.types import PipelineComponents
+
+device = get_device()
+
+
+class TI2VidTwoStagesPipeline:
+    """
+    Two-stage text/image-to-video generation pipeline.
+    Stage 1 generates video at the target resolution with CFG guidance, then
+    Stage 2 upsamples by 2x and refines using a distilled LoRA for higher
+    quality output. Supports optional image conditioning via the images parameter.
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        distilled_lora: list[LoraPathStrengthAndSDOps],
+        spatial_upsampler_path: str,
+        gemma_root: str,
+        loras: list[LoraPathStrengthAndSDOps],
+        device: str = device,
+        fp8transformer: bool = False,
+    ):
+        self.device = device
+        self.dtype = torch.bfloat16
+        self.stage_1_model_ledger = ModelLedger(
+            dtype=self.dtype,
+            device=device,
+            checkpoint_path=checkpoint_path,
+            gemma_root_path=gemma_root,
+            spatial_upsampler_path=spatial_upsampler_path,
+            loras=loras,
+            fp8transformer=fp8transformer,
+        )
+
+        self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
+            loras=distilled_lora,
+        )
+
+        self.pipeline_components = PipelineComponents(
+            dtype=self.dtype,
+            device=device,
+        )
+
+    @torch.inference_mode()
+    def __call__(  # noqa: PLR0913
+        self,
+        prompt: str,
+        negative_prompt: str,
+        seed: int,
+        height: int,
+        width: int,
+        num_frames: int,
+        frame_rate: float,
+        num_inference_steps: int,
+        cfg_guidance_scale: float,
+        images: list[tuple[str, int, float]],
+        tiling_config: TilingConfig | None = None,
+        enhance_prompt: bool = False,
+    ) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
+        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()
+        cfg_guider = CFGGuider(cfg_guidance_scale)
+        dtype = torch.bfloat16
+
+        text_encoder = self.stage_1_model_ledger.text_encoder()
+        if enhance_prompt:
+            prompt = generate_enhanced_prompt(
+                text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
+            )
+        context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
+        v_context_p, a_context_p = context_p
+        v_context_n, a_context_n = context_n
+
+        torch.cuda.synchronize()
+        del text_encoder
+        cleanup_memory()
+
+        # Stage 1: Initial low resolution video generation.
+        video_encoder = self.stage_1_model_ledger.video_encoder()
+        transformer = self.stage_1_model_ledger.transformer()
+        sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
+
+        def first_stage_denoising_loop(
+            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
+        ) -> tuple[LatentState, LatentState]:
+            return euler_denoising_loop(
+                sigmas=sigmas,
+                video_state=video_state,
+                audio_state=audio_state,
+                stepper=stepper,
+                denoise_fn=guider_denoising_func(
+                    cfg_guider,
+                    v_context_p,
+                    v_context_n,
+                    a_context_p,
+                    a_context_n,
+                    transformer=transformer,  # noqa: F821
+                ),
+            )
+
+        stage_1_output_shape = VideoPixelShape(
+            batch=1,
+            frames=num_frames,
+            width=width // 2,
+            height=height // 2,
+            fps=frame_rate,
+        )
+        stage_1_conditionings = image_conditionings_by_replacing_latent(
+            images=images,
+            height=stage_1_output_shape.height,
+            width=stage_1_output_shape.width,
+            video_encoder=video_encoder,
+            dtype=dtype,
+            device=self.device,
+        )
+        video_state, audio_state = denoise_audio_video(
+            output_shape=stage_1_output_shape,
+            conditionings=stage_1_conditionings,
+            noiser=noiser,
+            sigmas=sigmas,
+            stepper=stepper,
+            denoising_loop_fn=first_stage_denoising_loop,
+            components=self.pipeline_components,
+            dtype=dtype,
+            device=self.device,
+        )
+
+        torch.cuda.synchronize()
+        del transformer
+        cleanup_memory()
+
+        # Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
+        upscaled_video_latent = upsample_video(
+            latent=video_state.latent[:1],
+            video_encoder=video_encoder,
+            upsampler=self.stage_2_model_ledger.spatial_upsampler(),
+        )
+
+        torch.cuda.synchronize()
+        cleanup_memory()
+
+        transformer = self.stage_2_model_ledger.transformer()
+        distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
+
+        def second_stage_denoising_loop(
+            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
+        ) -> tuple[LatentState, LatentState]:
+            return euler_denoising_loop(
+                sigmas=sigmas,
+                video_state=video_state,
+                audio_state=audio_state,
+                stepper=stepper,
+                denoise_fn=simple_denoising_func(
+                    video_context=v_context_p,
+                    audio_context=a_context_p,
+                    transformer=transformer,  # noqa: F821
+                ),
+            )
+
+        stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
+        stage_2_conditionings = image_conditionings_by_replacing_latent(
+            images=images,
+            height=stage_2_output_shape.height,
+            width=stage_2_output_shape.width,
+            video_encoder=video_encoder,
+            dtype=dtype,
+            device=self.device,
+        )
+        video_state, audio_state = denoise_audio_video(
+            output_shape=stage_2_output_shape,
+            conditionings=stage_2_conditionings,
+            noiser=noiser,
+            sigmas=distilled_sigmas,
+            stepper=stepper,
+            denoising_loop_fn=second_stage_denoising_loop,
+            components=self.pipeline_components,
+            dtype=dtype,
+            device=self.device,
+            noise_scale=distilled_sigmas[0],
+            initial_video_latent=upscaled_video_latent,
+            initial_audio_latent=audio_state.latent,
+        )
+
+        torch.cuda.synchronize()
+        del transformer
+        del video_encoder
+        cleanup_memory()
+
+        decoded_video = vae_decode_video(video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config)
+        decoded_audio = vae_decode_audio(
+            audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder()
+        )
+
+        return decoded_video, decoded_audio
+
+
+@torch.inference_mode()
+def main() -> None:
+    logging.getLogger().setLevel(logging.INFO)
+    parser = default_2_stage_arg_parser()
+    args = parser.parse_args()
+    pipeline = TI2VidTwoStagesPipeline(
+        checkpoint_path=args.checkpoint_path,
+        distilled_lora=args.distilled_lora,
+        spatial_upsampler_path=args.spatial_upsampler_path,
+        gemma_root=args.gemma_root,
+        loras=args.lora,
+        fp8transformer=args.enable_fp8,
+    )
+    tiling_config = TilingConfig.default()
+    video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
+    video, audio = pipeline(
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        seed=args.seed,
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        frame_rate=args.frame_rate,
+        num_inference_steps=args.num_inference_steps,
+        cfg_guidance_scale=args.cfg_guidance_scale,
+        images=args.images,
+        tiling_config=tiling_config,
+    )
+
+    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,
+    )
+
+
+if __name__ == "__main__":
+    main()