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

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

@@ -64,6 +64,10 @@ class DistilledPipeline:
             device=device,
         )
 
+        # Cached models to avoid reloading
+        self._video_encoder = None
+        self._transformer = None
+
     @torch.inference_mode()
     def __call__(
         self,
@@ -76,23 +80,37 @@ class DistilledPipeline:
         frame_rate: float,
         images: list[tuple[str, int, float]],
         tiling_config: TilingConfig | None = None,
+        video_context: torch.Tensor | None = None,
+        audio_context: torch.Tensor | None = None,
     ) -> None:
         generator = torch.Generator(device=self.device).manual_seed(seed)
         noiser = GaussianNoiser(generator=generator)
         stepper = EulerDiffusionStep()
         dtype = torch.bfloat16
 
-        text_encoder = self.model_ledger.text_encoder()
-        context_p = encode_text(text_encoder, prompts=[prompt])[0]
-        video_context, audio_context = context_p
+        # 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()
+            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.
-        video_encoder = self.model_ledger.video_encoder()
-        transformer = self.model_ledger.transformer()
+        # 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(
@@ -168,9 +186,9 @@ class DistilledPipeline:
         )
 
         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.model_ledger.video_decoder(), tiling_config)
 
@@ -214,4 +232,4 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    main()