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

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

@@ -22,6 +22,7 @@ from ltx_pipelines.pipeline_utils import (
     denoise_audio_video,
     encode_text,
     euler_denoising_loop,
+    gradient_estimating_euler_denoising_loop,
     guider_denoising_func,
     simple_denoising_func,
 )
@@ -90,6 +91,10 @@ class TI2VidTwoStagesPipeline:
         cfg_guidance_scale: float,
         images: list[tuple[str, int, float]],
         tiling_config: TilingConfig | None = None,
+        video_context_positive: torch.Tensor | None = None,
+        audio_context_positive: torch.Tensor | None = None,
+        video_context_negative: torch.Tensor | None = None,
+        audio_context_negative: torch.Tensor | None = None,
     ) -> None:
         generator = torch.Generator(device=self.device).manual_seed(seed)
         noiser = GaussianNoiser(generator=generator)
@@ -97,14 +102,23 @@ class TI2VidTwoStagesPipeline:
         cfg_guider = CFGGuider(cfg_guidance_scale)
         dtype = torch.bfloat16
 
-        text_encoder = self.stage_1_model_ledger.text_encoder()
-        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
+        # Use pre-computed embeddings if provided, otherwise encode text
+        if (video_context_positive is None or audio_context_positive is None or
+            video_context_negative is None or audio_context_negative is None):
+            text_encoder = self.stage_1_model_ledger.text_encoder()
+            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
-        utils.cleanup_memory()
+            torch.cuda.synchronize()
+            del text_encoder
+            utils.cleanup_memory()
+        else:
+            # Move pre-computed embeddings to device if needed
+            v_context_p = video_context_positive.to(self.device)
+            a_context_p = audio_context_positive.to(self.device)
+            v_context_n = video_context_negative.to(self.device)
+            a_context_n = audio_context_negative.to(self.device)
 
         # Stage 1: Initial low resolution video generation.
         video_encoder = self.stage_1_model_ledger.video_encoder()
@@ -170,7 +184,18 @@ class TI2VidTwoStagesPipeline:
         def second_stage_denoising_loop(
             sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
         ) -> tuple[LatentState, LatentState]:
-            return euler_denoising_loop(
+            # 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
+            #     ),
+            # )
+            return gradient_estimating_euler_denoising_loop(
                 sigmas=sigmas,
                 video_state=video_state,
                 audio_state=audio_state,
@@ -180,6 +205,7 @@ class TI2VidTwoStagesPipeline:
                     audio_context=a_context_p,
                     transformer=transformer,  # noqa: F821
                 ),
+                ge_gamma=2.0,  # Gradient estimation coefficient
             )
 
         stage_2_output_shape = VideoPixelShape(