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

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

@@ -75,6 +75,11 @@ class TI2VidTwoStagesPipeline:
             device=device,
         )
 
+        # Cached models to avoid reloading
+        self._video_encoder = None
+        self._stage_1_transformer = None
+        self._stage_2_transformer = None
+
     @torch.inference_mode()
     def __call__(  # noqa: PLR0913
         self,
@@ -90,6 +95,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,18 +106,33 @@ 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()
-        transformer = self.stage_1_model_ledger.transformer()
+        # Load models only if not already cached
+        if self._video_encoder is None:
+            self._video_encoder = self.stage_1_model_ledger.video_encoder()
+        video_encoder = self._video_encoder
+
+        if self._stage_1_transformer is None:
+            self._stage_1_transformer = self.stage_1_model_ledger.transformer()
+        transformer = self._stage_1_transformer
         sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
 
         def first_stage_denoising_loop(
@@ -151,8 +175,8 @@ class TI2VidTwoStagesPipeline:
         )
 
         torch.cuda.synchronize()
-        del transformer
-        utils.cleanup_memory()
+        # del transformer
+        # utils.cleanup_memory()
 
         # Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
         upscaled_video_latent = utils.upsample_video(
@@ -162,9 +186,12 @@ class TI2VidTwoStagesPipeline:
         )
 
         torch.cuda.synchronize()
-        utils.cleanup_memory()
+        # utils.cleanup_memory()
 
-        transformer = self.stage_2_model_ledger.transformer()
+        # Load stage 2 transformer only if not already cached
+        if self._stage_2_transformer is None:
+            self._stage_2_transformer = self.stage_2_model_ledger.transformer()
+        transformer = self._stage_2_transformer
         distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
 
         def second_stage_denoising_loop(
@@ -209,9 +236,9 @@ class TI2VidTwoStagesPipeline:
         )
 
         torch.cuda.synchronize()
-        del transformer
-        del video_encoder
-        utils.cleanup_memory()
+        # del transformer
+        # del video_encoder
+        # utils.cleanup_memory()
 
         decoded_video = vae_decode_video(video_state, self.stage_2_model_ledger.video_decoder(), tiling_config)