Automated PR - 2026-04-23

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

@@ -39,7 +39,7 @@ from ltx_pipelines.utils.constants import (
 from ltx_pipelines.utils.denoisers import SimpleDenoiser
 from ltx_pipelines.utils.helpers import assert_resolution, combined_image_conditionings, get_device
 from ltx_pipelines.utils.media_io import decode_video_by_frame, encode_video, video_preprocess
-from ltx_pipelines.utils.types import ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class ICLoraPipeline:
@@ -63,12 +63,18 @@ class ICLoraPipeline:
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
         torch_compile: bool = False,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
         self.dtype = torch.bfloat16
 
         self.prompt_encoder = PromptEncoder(
-            distilled_checkpoint_path, gemma_root, self.dtype, self.device, registry=registry
+            distilled_checkpoint_path,
+            gemma_root,
+            self.dtype,
+            self.device,
+            registry=registry,
+            offload_mode=offload_mode,
         )
         self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
         self.stage_1 = DiffusionStage(
@@ -79,6 +85,7 @@ class ICLoraPipeline:
             quantization=quantization,
             registry=registry,
             torch_compile=torch_compile,
+            offload_mode=offload_mode,
         )
         self.stage_2 = DiffusionStage(
             distilled_checkpoint_path,
@@ -88,6 +95,7 @@ class ICLoraPipeline:
             quantization=quantization,
             registry=registry,
             torch_compile=torch_compile,
+            offload_mode=offload_mode,
         )
         self.upsampler = VideoUpsampler(
             distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -125,7 +133,6 @@ class ICLoraPipeline:
         conditioning_attention_strength: float = 1.0,
         skip_stage_2: bool = False,
         conditioning_attention_mask: torch.Tensor | None = None,
-        streaming_prefetch_count: int | None = None,
         stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS,
         stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
     ) -> tuple[Iterator[torch.Tensor], Audio]:
@@ -175,7 +182,6 @@ class ICLoraPipeline:
             enhance_first_prompt=enhance_prompt,
             enhance_prompt_image=images[0][0] if len(images) > 0 else None,
             enhance_prompt_seed=seed,
-            streaming_prefetch_count=streaming_prefetch_count,
         )
         video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
 
@@ -219,7 +225,6 @@ class ICLoraPipeline:
             audio=ModalitySpec(
                 context=audio_context,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         if skip_stage_2:
@@ -264,7 +269,6 @@ class ICLoraPipeline:
                 noise_scale=stage_2_sigmas[0].item(),
                 initial_latent=audio_state.latent,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -463,6 +467,7 @@ def main() -> None:
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
         torch_compile=args.compile,
+        offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -479,7 +484,6 @@ def main() -> None:
         conditioning_attention_strength=conditioning_attention_strength,
         skip_stage_2=args.skip_stage_2,
         conditioning_attention_mask=conditioning_attention_mask,
-        streaming_prefetch_count=args.streaming_prefetch_count,
     )
 
     encode_video(