Automated PR - 2026-04-13

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

@@ -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,