Automated PR - 2026-03-30

packages/ltx-pipelines/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ltx-pipelines"
-version = "1.0.0"
+version = "1.1.0"
 description = "Pipelines implementation for Lightricks' LTX-2 model"
 readme = "README.md"
 requires-python = ">=3.10"

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

@@ -7,7 +7,6 @@ This package provides ready-to-use pipelines for video generation:
 - ICLoraPipeline: Image/video conditioning with distilled LoRA
 - KeyframeInterpolationPipeline: Keyframe-based video interpolation
 - RetakePipeline: Regenerate a time region (retake) of an existing video
-- ModelLedger: Central coordinator for loading and building models
 For more detailed components and utilities, import from specific submodules
 like `ltx_pipelines.utils.media_io` or `ltx_pipelines.utils.constants`.
 """

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

@@ -3,38 +3,35 @@ from collections.abc import Iterator
 
 import torch
 
-from ltx_core.components.diffusion_steps import EulerDiffusionStep
 from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
 from ltx_core.components.noisers import GaussianNoiser
-from ltx_core.components.protocols import DiffusionStepProtocol
 from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.loader.registry import Registry
 from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
-from ltx_core.model.upsampler import upsample_video
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
-from ltx_core.model.video_vae import decode_video as vae_decode_video
 from ltx_core.quantization import QuantizationPolicy
-from ltx_core.types import Audio, AudioLatentShape, LatentState, VideoPixelShape
-from ltx_pipelines.utils import ModelLedger
+from ltx_core.types import Audio, AudioLatentShape, VideoPixelShape
 from ltx_pipelines.utils.args import default_2_stage_arg_parser
+from ltx_pipelines.utils.blocks import (
+    AudioConditioner,
+    DiffusionStage,
+    ImageConditioner,
+    PromptEncoder,
+    VideoDecoder,
+    VideoUpsampler,
+)
 from ltx_pipelines.utils.constants import (
     STAGE_2_DISTILLED_SIGMA_VALUES,
 )
+from ltx_pipelines.utils.denoisers import GuidedDenoiser, SimpleDenoiser
 from ltx_pipelines.utils.helpers import (
     assert_resolution,
-    cleanup_memory,
     combined_image_conditionings,
-    denoise_video_only,
-    encode_prompts,
     get_device,
-    multi_modal_guider_denoising_func,
-    simple_denoising_func,
 )
 from ltx_pipelines.utils.media_io import decode_audio_from_file, encode_video
-from ltx_pipelines.utils.samplers import euler_denoising_loop
-from ltx_pipelines.utils.types import PipelineComponents
-
-device = get_device()
+from ltx_pipelines.utils.types import ModalitySpec
 
 
 class A2VidPipelineTwoStage:
@@ -52,30 +49,40 @@ class A2VidPipelineTwoStage:
         spatial_upsampler_path: str,
         gemma_root: str,
         loras: list[LoraPathStrengthAndSDOps],
-        device: torch.device = device,
+        device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
+        registry: Registry | None = None,
+        torch_compile: bool = False,
     ):
-        self.device = device
+        self.device = device or get_device()
         self.dtype = torch.bfloat16
 
-        self.stage_1_model_ledger = ModelLedger(
-            dtype=self.dtype,
-            device=device,
-            checkpoint_path=checkpoint_path,
-            gemma_root_path=gemma_root,
-            spatial_upsampler_path=spatial_upsampler_path,
-            loras=loras,
+        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)
+        self.audio_conditioner = AudioConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
+        self.stage_1 = DiffusionStage(
+            checkpoint_path,
+            self.dtype,
+            self.device,
+            loras=tuple(loras),
             quantization=quantization,
+            registry=registry,
+            torch_compile=torch_compile,
         )
-
-        self.stage_2_model_ledger = self.stage_1_model_ledger.with_additional_loras(
-            loras=distilled_lora,
+        stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
+        self.stage_2 = DiffusionStage(
+            checkpoint_path,
+            self.dtype,
+            self.device,
+            loras=stage_2_loras,
+            quantization=quantization,
+            registry=registry,
+            torch_compile=torch_compile,
         )
-
-        self.pipeline_components = PipelineComponents(
-            dtype=self.dtype,
-            device=device,
+        self.upsampler = VideoUpsampler(
+            checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
         )
+        self.video_decoder = VideoDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
 
     def __call__(  # noqa: PLR0913
         self,
@@ -94,31 +101,35 @@ class A2VidPipelineTwoStage:
         audio_max_duration: float | None = None,
         tiling_config: TilingConfig | None = None,
         enhance_prompt: bool = False,
+        streaming_prefetch_count: int | None = None,
+        max_batch_size: int = 1,
     ) -> tuple[Iterator[torch.Tensor], Audio]:
         assert_resolution(height=height, width=width, is_two_stage=True)
 
         generator = torch.Generator(device=self.device).manual_seed(seed)
         noiser = GaussianNoiser(generator=generator)
-        stepper = EulerDiffusionStep()
         dtype = torch.bfloat16
 
-        ctx_p, ctx_n = encode_prompts(
+        ctx_p, ctx_n = self.prompt_encoder(
             [prompt, negative_prompt],
-            self.stage_1_model_ledger,
             enhance_first_prompt=enhance_prompt,
             enhance_prompt_image=images[0][0] if len(images) > 0 else None,
+            streaming_prefetch_count=streaming_prefetch_count,
         )
         v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
         v_context_n, _ = ctx_n.video_encoding, ctx_n.audio_encoding
 
         # Encode audio.
         decoded_audio = decode_audio_from_file(audio_path, self.device, audio_start_time, audio_max_duration)
-        encoded_audio_latent = vae_encode_audio(decoded_audio, self.stage_1_model_ledger.audio_encoder())
+        if decoded_audio is None:
+            raise ValueError(f"Failed to decode audio from {audio_path}. Please check the file and try again.")
+
+        encoded_audio_latent = self.audio_conditioner(lambda enc: vae_encode_audio(decoded_audio, enc, None))
         audio_shape = AudioLatentShape.from_duration(batch=1, duration=num_frames / frame_rate, channels=8, mel_bins=16)
         encoded_audio_latent = encoded_audio_latent[:, :, : audio_shape.frames]
 
-        # Stage 1: encode image conditionings with the VAE encoder, then free it
-        # before loading the transformer to reduce peak VRAM.
+        # Stage 1: encode image conditionings with the VAE encoder, then denoise
+        # video-only (audio frozen).
         stage_1_output_shape = VideoPixelShape(
             batch=1,
             frames=num_frames,
@@ -126,122 +137,91 @@ class A2VidPipelineTwoStage:
             height=height // 2,
             fps=frame_rate,
         )
-        video_encoder = self.stage_1_model_ledger.video_encoder()
-        stage_1_conditionings = combined_image_conditionings(
-            images=images,
-            height=stage_1_output_shape.height,
-            width=stage_1_output_shape.width,
-            video_encoder=video_encoder,
-            dtype=dtype,
-            device=self.device,
+        stage_1_conditionings = self.image_conditioner(
+            lambda enc: combined_image_conditionings(
+                images=images,
+                height=stage_1_output_shape.height,
+                width=stage_1_output_shape.width,
+                video_encoder=enc,
+                dtype=dtype,
+                device=self.device,
+            )
         )
-        torch.cuda.synchronize()
-        del video_encoder
-        cleanup_memory()
 
-        transformer = self.stage_1_model_ledger.transformer()
         sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
 
-        def first_stage_denoising_loop(
-            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
-        ) -> tuple[LatentState, LatentState]:
-            return euler_denoising_loop(
-                sigmas=sigmas,
-                video_state=video_state,
-                audio_state=audio_state,
-                stepper=stepper,
-                denoise_fn=multi_modal_guider_denoising_func(
-                    video_guider=MultiModalGuider(
-                        params=video_guider_params,
-                        negative_context=v_context_n,
-                    ),
-                    audio_guider=MultiModalGuider(
-                        params=MultiModalGuiderParams(),
-                    ),
-                    v_context=v_context_p,
-                    a_context=a_context_p,
-                    transformer=transformer,  # noqa: F821
+        video_state, _ = self.stage_1(
+            denoiser=GuidedDenoiser(
+                v_context=v_context_p,
+                a_context=a_context_p,
+                video_guider=MultiModalGuider(
+                    params=video_guider_params,
+                    negative_context=v_context_n,
                 ),
-            )
-
-        video_state = denoise_video_only(
-            output_shape=stage_1_output_shape,
-            conditionings=stage_1_conditionings,
-            noiser=noiser,
+                audio_guider=MultiModalGuider(
+                    params=MultiModalGuiderParams(),
+                ),
+            ),
             sigmas=sigmas,
-            stepper=stepper,
-            denoising_loop_fn=first_stage_denoising_loop,
-            components=self.pipeline_components,
-            dtype=dtype,
-            device=self.device,
-            initial_audio_latent=encoded_audio_latent,
+            noiser=noiser,
+            width=stage_1_output_shape.width,
+            height=stage_1_output_shape.height,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(
+                context=v_context_p,
+                conditionings=stage_1_conditionings,
+            ),
+            audio=ModalitySpec(
+                context=a_context_p,
+                frozen=True,
+                noise_scale=0.0,
+                initial_latent=encoded_audio_latent,
+            ),
+            streaming_prefetch_count=streaming_prefetch_count,
+            max_batch_size=max_batch_size,
         )
 
-        torch.cuda.synchronize()
-        del transformer
-        cleanup_memory()
-
         # Stage 2: Upsample and refine the video at higher resolution with distilled LoRA.
-        video_encoder = self.stage_1_model_ledger.video_encoder()
-        upscaled_video_latent = upsample_video(
-            latent=video_state.latent[:1],
-            video_encoder=video_encoder,
-            upsampler=self.stage_2_model_ledger.spatial_upsampler(),
-        )
+        upscaled_video_latent = self.upsampler(video_state.latent[:1])
 
-        stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
-        stage_2_conditionings = combined_image_conditionings(
-            images=images,
-            height=stage_2_output_shape.height,
-            width=stage_2_output_shape.width,
-            video_encoder=video_encoder,
-            dtype=dtype,
-            device=self.device,
-        )
-        del video_encoder
-        torch.cuda.synchronize()
-        cleanup_memory()
-
-        transformer = self.stage_2_model_ledger.transformer()
         distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
-
-        def second_stage_denoising_loop(
-            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
-        ) -> tuple[LatentState, LatentState]:
-            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
-                ),
+        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: combined_image_conditionings(
+                images=images,
+                height=stage_2_output_shape.height,
+                width=stage_2_output_shape.width,
+                video_encoder=enc,
+                dtype=dtype,
+                device=self.device,
             )
+        )
 
-        video_state = denoise_video_only(
-            output_shape=stage_2_output_shape,
-            conditionings=stage_2_conditionings,
-            noiser=noiser,
+        video_state, _ = self.stage_2(
+            denoiser=SimpleDenoiser(v_context_p, a_context_p),
             sigmas=distilled_sigmas,
-            stepper=stepper,
-            denoising_loop_fn=second_stage_denoising_loop,
-            components=self.pipeline_components,
-            dtype=dtype,
-            device=self.device,
-            noise_scale=distilled_sigmas[0],
-            initial_video_latent=upscaled_video_latent,
-            initial_audio_latent=encoded_audio_latent,
+            noiser=noiser,
+            width=width,
+            height=height,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(
+                context=v_context_p,
+                conditionings=stage_2_conditionings,
+                noise_scale=distilled_sigmas[0].item(),
+                initial_latent=upscaled_video_latent,
+            ),
+            audio=ModalitySpec(
+                context=a_context_p,
+                frozen=True,
+                noise_scale=0.0,
+                initial_latent=encoded_audio_latent,
+            ),
+            streaming_prefetch_count=streaming_prefetch_count,
         )
 
-        torch.cuda.synchronize()
-        del transformer
-        cleanup_memory()
-
-        decoded_video = vae_decode_video(
-            video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config, generator
-        )
+        decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
 
         # Return the original input audio instead of VAE-decoded audio to preserve fidelity.
         # decode_audio_from_file already returns normalised [-1, 1] float values.
@@ -280,6 +260,7 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
+        torch_compile=args.compile,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -308,6 +289,8 @@ def main() -> None:
         audio_max_duration=args.audio_max_duration
         if args.audio_max_duration is not None
         else args.num_frames / args.frame_rate,
+        streaming_prefetch_count=args.streaming_prefetch_count,
+        max_batch_size=args.max_batch_size,
     )
 
     encode_video(

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

@@ -3,40 +3,38 @@ from collections.abc import Iterator
 
 import torch
 
-from ltx_core.components.diffusion_steps import EulerDiffusionStep
 from ltx_core.components.noisers import GaussianNoiser
-from ltx_core.components.protocols import DiffusionStepProtocol
 from ltx_core.loader import LoraPathStrengthAndSDOps
-from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
-from ltx_core.model.upsampler import upsample_video
+from ltx_core.loader.registry import Registry
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
-from ltx_core.model.video_vae import decode_video as vae_decode_video
 from ltx_core.quantization import QuantizationPolicy
-from ltx_core.types import Audio, LatentState, VideoPixelShape
-from ltx_pipelines.utils import ModelLedger, euler_denoising_loop
+from ltx_core.types import Audio
 from ltx_pipelines.utils.args import (
     ImageConditioningInput,
     default_2_stage_distilled_arg_parser,
     detect_checkpoint_path,
 )
+from ltx_pipelines.utils.blocks import (
+    AudioDecoder,
+    DiffusionStage,
+    ImageConditioner,
+    PromptEncoder,
+    VideoDecoder,
+    VideoUpsampler,
+)
 from ltx_pipelines.utils.constants import (
     DISTILLED_SIGMA_VALUES,
     STAGE_2_DISTILLED_SIGMA_VALUES,
     detect_params,
 )
+from ltx_pipelines.utils.denoisers import SimpleDenoiser
 from ltx_pipelines.utils.helpers import (
     assert_resolution,
-    cleanup_memory,
     combined_image_conditionings,
-    denoise_audio_video,
-    encode_prompts,
     get_device,
-    simple_denoising_func,
 )
 from ltx_pipelines.utils.media_io import encode_video
-from ltx_pipelines.utils.types import PipelineComponents
-
-device = get_device()
+from ltx_pipelines.utils.types import ModalitySpec
 
 
 class DistilledPipeline:
@@ -52,26 +50,32 @@ class DistilledPipeline:
         gemma_root: str,
         spatial_upsampler_path: str,
         loras: list[LoraPathStrengthAndSDOps],
-        device: torch.device = device,
+        device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
+        registry: Registry | None = None,
+        torch_compile: bool = False,
     ):
-        self.device = device
+        self.device = device or get_device()
         self.dtype = torch.bfloat16
 
-        self.model_ledger = ModelLedger(
-            dtype=self.dtype,
-            device=device,
-            checkpoint_path=distilled_checkpoint_path,
-            spatial_upsampler_path=spatial_upsampler_path,
-            gemma_root_path=gemma_root,
-            loras=loras,
+        self.prompt_encoder = PromptEncoder(
+            distilled_checkpoint_path, gemma_root, self.dtype, self.device, registry=registry
+        )
+        self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
+        self.stage = DiffusionStage(
+            distilled_checkpoint_path,
+            self.dtype,
+            self.device,
+            loras=tuple(loras),
             quantization=quantization,
+            registry=registry,
+            torch_compile=torch_compile,
         )
-
-        self.pipeline_components = PipelineComponents(
-            dtype=self.dtype,
-            device=device,
+        self.upsampler = VideoUpsampler(
+            distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
         )
+        self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
+        self.audio_decoder = AudioDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
 
     def __call__(
         self,
@@ -84,114 +88,88 @@ class DistilledPipeline:
         images: list[ImageConditioningInput],
         tiling_config: TilingConfig | None = None,
         enhance_prompt: bool = False,
+        streaming_prefetch_count: int | None = None,
     ) -> tuple[Iterator[torch.Tensor], Audio]:
         assert_resolution(height=height, width=width, is_two_stage=True)
 
         generator = torch.Generator(device=self.device).manual_seed(seed)
         noiser = GaussianNoiser(generator=generator)
-        stepper = EulerDiffusionStep()
         dtype = torch.bfloat16
 
-        (ctx_p,) = encode_prompts(
+        (ctx_p,) = self.prompt_encoder(
             [prompt],
-            self.model_ledger,
             enhance_first_prompt=enhance_prompt,
             enhance_prompt_image=images[0][0] if len(images) > 0 else None,
+            streaming_prefetch_count=streaming_prefetch_count,
         )
         video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
 
         # Stage 1: Initial low resolution video generation.
-        video_encoder = self.model_ledger.video_encoder()
-        transformer = self.model_ledger.transformer()
         stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
-
-        def denoising_loop(
-            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
-        ) -> tuple[LatentState, LatentState]:
-            return euler_denoising_loop(
-                sigmas=sigmas,
-                video_state=video_state,
-                audio_state=audio_state,
-                stepper=stepper,
-                denoise_fn=simple_denoising_func(
-                    video_context=video_context,
-                    audio_context=audio_context,
-                    transformer=transformer,  # noqa: F821
-                ),
+        stage_1_w, stage_1_h = width // 2, height // 2
+        stage_1_conditionings = self.image_conditioner(
+            lambda enc: combined_image_conditionings(
+                images=images,
+                height=stage_1_h,
+                width=stage_1_w,
+                video_encoder=enc,
+                dtype=dtype,
+                device=self.device,
             )
-
-        stage_1_output_shape = VideoPixelShape(
-            batch=1,
-            frames=num_frames,
-            width=width // 2,
-            height=height // 2,
-            fps=frame_rate,
-        )
-        stage_1_conditionings = combined_image_conditionings(
-            images=images,
-            height=stage_1_output_shape.height,
-            width=stage_1_output_shape.width,
-            video_encoder=video_encoder,
-            dtype=dtype,
-            device=self.device,
         )
 
-        video_state, audio_state = denoise_audio_video(
-            output_shape=stage_1_output_shape,
-            conditionings=stage_1_conditionings,
-            noiser=noiser,
+        video_state, audio_state = self.stage(
+            denoiser=SimpleDenoiser(video_context, audio_context),
             sigmas=stage_1_sigmas,
-            stepper=stepper,
-            denoising_loop_fn=denoising_loop,
-            components=self.pipeline_components,
-            dtype=dtype,
-            device=self.device,
+            noiser=noiser,
+            width=stage_1_w,
+            height=stage_1_h,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(context=video_context, conditionings=stage_1_conditionings),
+            audio=ModalitySpec(context=audio_context),
+            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         # Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
-        upscaled_video_latent = upsample_video(
-            latent=video_state.latent[:1], video_encoder=video_encoder, upsampler=self.model_ledger.spatial_upsampler()
-        )
-
-        torch.cuda.synchronize()
-        cleanup_memory()
+        upscaled_video_latent = self.upsampler(video_state.latent[:1])
 
         stage_2_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
-        stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
-        stage_2_conditionings = combined_image_conditionings(
-            images=images,
-            height=stage_2_output_shape.height,
-            width=stage_2_output_shape.width,
-            video_encoder=video_encoder,
-            dtype=dtype,
-            device=self.device,
+        stage_2_conditionings = self.image_conditioner(
+            lambda enc: combined_image_conditionings(
+                images=images,
+                height=height,
+                width=width,
+                video_encoder=enc,
+                dtype=dtype,
+                device=self.device,
+            )
         )
-        video_state, audio_state = denoise_audio_video(
-            output_shape=stage_2_output_shape,
-            conditionings=stage_2_conditionings,
-            noiser=noiser,
+
+        video_state, audio_state = self.stage(
+            denoiser=SimpleDenoiser(video_context, audio_context),
             sigmas=stage_2_sigmas,
-            stepper=stepper,
-            denoising_loop_fn=denoising_loop,
-            components=self.pipeline_components,
-            dtype=dtype,
-            device=self.device,
-            noise_scale=stage_2_sigmas[0],
-            initial_video_latent=upscaled_video_latent,
-            initial_audio_latent=audio_state.latent,
+            noiser=noiser,
+            width=width,
+            height=height,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(
+                context=video_context,
+                conditionings=stage_2_conditionings,
+                noise_scale=stage_2_sigmas[0].item(),
+                initial_latent=upscaled_video_latent,
+            ),
+            audio=ModalitySpec(
+                context=audio_context,
+                noise_scale=stage_2_sigmas[0].item(),
+                initial_latent=audio_state.latent,
+            ),
+            streaming_prefetch_count=streaming_prefetch_count,
         )
 
-        torch.cuda.synchronize()
-        del transformer
-        del video_encoder
-        cleanup_memory()
-
-        decoded_video = vae_decode_video(
-            video_state.latent, self.model_ledger.video_decoder(), tiling_config, generator
-        )
-        decoded_audio = vae_decode_audio(
-            audio_state.latent, self.model_ledger.audio_decoder(), self.model_ledger.vocoder()
-        )
+        decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
+        decoded_audio = self.audio_decoder(audio_state.latent)
         return decoded_video, decoded_audio
 
 
@@ -208,6 +186,7 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
+        torch_compile=args.compile,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -221,6 +200,7 @@ def main() -> None:
         images=args.images,
         tiling_config=tiling_config,
         enhance_prompt=args.enhance_prompt,
+        streaming_prefetch_count=args.streaming_prefetch_count,
     )
 
     encode_video(

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

@@ -5,32 +5,17 @@ import torch
 from einops import rearrange
 from safetensors import safe_open
 
-from ltx_core.components.diffusion_steps import EulerDiffusionStep
 from ltx_core.components.noisers import GaussianNoiser
-from ltx_core.components.protocols import DiffusionStepProtocol
 from ltx_core.conditioning import (
     ConditioningItem,
     ConditioningItemAttentionStrengthWrapper,
     VideoConditionByReferenceLatent,
 )
 from ltx_core.loader import LoraPathStrengthAndSDOps
-from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
-from ltx_core.model.upsampler import upsample_video
+from ltx_core.loader.registry import Registry
 from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
-from ltx_core.model.video_vae import decode_video as vae_decode_video
 from ltx_core.quantization import QuantizationPolicy
-from ltx_core.types import Audio, LatentState, VideoLatentShape, VideoPixelShape
-from ltx_pipelines.utils import (
-    ModelLedger,
-    assert_resolution,
-    cleanup_memory,
-    combined_image_conditionings,
-    denoise_audio_video,
-    encode_prompts,
-    euler_denoising_loop,
-    get_device,
-    simple_denoising_func,
-)
+from ltx_core.types import Audio, VideoLatentShape, VideoPixelShape
 from ltx_pipelines.utils.args import (
     ImageConditioningInput,
     VideoConditioningAction,
@@ -38,15 +23,23 @@ from ltx_pipelines.utils.args import (
     default_2_stage_distilled_arg_parser,
     detect_checkpoint_path,
 )
+from ltx_pipelines.utils.blocks import (
+    AudioDecoder,
+    DiffusionStage,
+    ImageConditioner,
+    PromptEncoder,
+    VideoDecoder,
+    VideoUpsampler,
+)
 from ltx_pipelines.utils.constants import (
     DISTILLED_SIGMA_VALUES,
     STAGE_2_DISTILLED_SIGMA_VALUES,
     detect_params,
 )
-from ltx_pipelines.utils.media_io import encode_video, load_video_conditioning
-from ltx_pipelines.utils.types import PipelineComponents
-
-device = get_device()
+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
 
 
 class ICLoraPipeline:
@@ -66,33 +59,41 @@ class ICLoraPipeline:
         spatial_upsampler_path: str,
         gemma_root: str,
         loras: list[LoraPathStrengthAndSDOps],
-        device: torch.device = device,
+        device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
+        registry: Registry | None = None,
+        torch_compile: bool = False,
     ):
+        self.device = device or get_device()
         self.dtype = torch.bfloat16
-        self.stage_1_model_ledger = ModelLedger(
-            dtype=self.dtype,
-            device=device,
-            checkpoint_path=distilled_checkpoint_path,
-            spatial_upsampler_path=spatial_upsampler_path,
-            gemma_root_path=gemma_root,
-            loras=loras,
+
+        self.prompt_encoder = PromptEncoder(
+            distilled_checkpoint_path, gemma_root, self.dtype, self.device, registry=registry
+        )
+        self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
+        self.stage_1 = DiffusionStage(
+            distilled_checkpoint_path,
+            self.dtype,
+            self.device,
+            loras=tuple(loras),
             quantization=quantization,
+            registry=registry,
+            torch_compile=torch_compile,
         )
-        self.stage_2_model_ledger = ModelLedger(
-            dtype=self.dtype,
-            device=device,
-            checkpoint_path=distilled_checkpoint_path,
-            spatial_upsampler_path=spatial_upsampler_path,
-            gemma_root_path=gemma_root,
-            loras=[],
+        self.stage_2 = DiffusionStage(
+            distilled_checkpoint_path,
+            self.dtype,
+            self.device,
+            loras=(),
             quantization=quantization,
+            registry=registry,
+            torch_compile=torch_compile,
         )
-        self.pipeline_components = PipelineComponents(
-            dtype=self.dtype,
-            device=device,
+        self.upsampler = VideoUpsampler(
+            distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
         )
-        self.device = device
+        self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
+        self.audio_decoder = AudioDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
 
         # Read reference downscale factor from LoRA metadata.
         # IC-LoRAs trained with low-resolution reference videos store this factor
@@ -124,6 +125,7 @@ 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,
     ) -> tuple[Iterator[torch.Tensor], Audio]:
         """
         Generate video with IC-LoRA conditioning.
@@ -165,15 +167,13 @@ class ICLoraPipeline:
 
         generator = torch.Generator(device=self.device).manual_seed(seed)
         noiser = GaussianNoiser(generator=generator)
-        stepper = EulerDiffusionStep()
-        dtype = torch.bfloat16
 
-        (ctx_p,) = encode_prompts(
+        (ctx_p,) = self.prompt_encoder(
             [prompt],
-            self.stage_1_model_ledger,
             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
 
@@ -186,130 +186,87 @@ class ICLoraPipeline:
             fps=frame_rate,
         )
 
-        # Encode conditionings before loading transformer to reduce peak VRAM
-        video_encoder = self.stage_1_model_ledger.video_encoder()
-        stage_1_conditionings = self._create_conditionings(
-            images=images,
-            video_conditioning=video_conditioning,
-            height=stage_1_output_shape.height,
-            width=stage_1_output_shape.width,
-            video_encoder=video_encoder,
-            num_frames=num_frames,
-            conditioning_attention_strength=conditioning_attention_strength,
-            conditioning_attention_mask=conditioning_attention_mask,
+        # Encode conditionings using the video encoder block
+        stage_1_conditionings = self.image_conditioner(
+            lambda enc: self._create_conditionings(
+                images=images,
+                video_conditioning=video_conditioning,
+                height=stage_1_output_shape.height,
+                width=stage_1_output_shape.width,
+                video_encoder=enc,
+                num_frames=num_frames,
+                conditioning_attention_strength=conditioning_attention_strength,
+                conditioning_attention_mask=conditioning_attention_mask,
+            )
         )
 
-        transformer = self.stage_1_model_ledger.transformer()
         stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
 
-        def first_stage_denoising_loop(
-            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
-        ) -> tuple[LatentState, LatentState]:
-            return euler_denoising_loop(
-                sigmas=sigmas,
-                video_state=video_state,
-                audio_state=audio_state,
-                stepper=stepper,
-                denoise_fn=simple_denoising_func(
-                    video_context=video_context,
-                    audio_context=audio_context,
-                    transformer=transformer,  # noqa: F821
-                ),
-            )
-
-        video_state, audio_state = denoise_audio_video(
-            output_shape=stage_1_output_shape,
-            conditionings=stage_1_conditionings,
-            noiser=noiser,
+        video_state, audio_state = self.stage_1(
+            denoiser=SimpleDenoiser(video_context, audio_context),
             sigmas=stage_1_sigmas,
-            stepper=stepper,
-            denoising_loop_fn=first_stage_denoising_loop,
-            components=self.pipeline_components,
-            dtype=dtype,
-            device=self.device,
+            noiser=noiser,
+            width=stage_1_output_shape.width,
+            height=stage_1_output_shape.height,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(
+                context=video_context,
+                conditionings=stage_1_conditionings,
+            ),
+            audio=ModalitySpec(
+                context=audio_context,
+            ),
+            streaming_prefetch_count=streaming_prefetch_count,
         )
 
-        torch.cuda.synchronize()
-        del transformer
-        cleanup_memory()
-
         if skip_stage_2:
             # Skip Stage 2: Decode directly from Stage 1 output at half resolution
             logging.info("[IC-LoRA] Skipping Stage 2 (--skip-stage-2 enabled)")
-            decoded_video = vae_decode_video(
-                video_state.latent, self.stage_1_model_ledger.video_decoder(), tiling_config, generator
-            )
-            decoded_audio = vae_decode_audio(
-                audio_state.latent, self.stage_1_model_ledger.audio_decoder(), self.stage_1_model_ledger.vocoder()
-            )
-            del video_encoder
-            cleanup_memory()
+            decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
+            decoded_audio = self.audio_decoder(audio_state.latent)
             return decoded_video, decoded_audio
 
         # Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
-        upscaled_video_latent = upsample_video(
-            latent=video_state.latent[:1],
-            video_encoder=video_encoder,
-            upsampler=self.stage_2_model_ledger.spatial_upsampler(),
-        )
+        upscaled_video_latent = self.upsampler(video_state.latent[:1])
 
-        torch.cuda.synchronize()
-        cleanup_memory()
-
-        transformer = self.stage_2_model_ledger.transformer()
         distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
-
-        def second_stage_denoising_loop(
-            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
-        ) -> tuple[LatentState, LatentState]:
-            return euler_denoising_loop(
-                sigmas=sigmas,
-                video_state=video_state,
-                audio_state=audio_state,
-                stepper=stepper,
-                denoise_fn=simple_denoising_func(
-                    video_context=video_context,
-                    audio_context=audio_context,
-                    transformer=transformer,  # noqa: F821
-                ),
-            )
-
         stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
-        stage_2_conditionings = combined_image_conditionings(
-            images=images,
-            height=stage_2_output_shape.height,
-            width=stage_2_output_shape.width,
-            video_encoder=video_encoder,
-            dtype=self.dtype,
-            device=self.device,
+        stage_2_conditionings = self.image_conditioner(
+            lambda enc: combined_image_conditionings(
+                images=images,
+                height=stage_2_output_shape.height,
+                width=stage_2_output_shape.width,
+                video_encoder=enc,
+                dtype=self.dtype,
+                device=self.device,
+            )
         )
 
-        video_state, audio_state = denoise_audio_video(
-            output_shape=stage_2_output_shape,
-            conditionings=stage_2_conditionings,
-            noiser=noiser,
+        video_state, audio_state = self.stage_2(
+            denoiser=SimpleDenoiser(video_context, audio_context),
             sigmas=distilled_sigmas,
-            stepper=stepper,
-            denoising_loop_fn=second_stage_denoising_loop,
-            components=self.pipeline_components,
-            dtype=dtype,
-            device=self.device,
-            noise_scale=distilled_sigmas[0],
-            initial_video_latent=upscaled_video_latent,
-            initial_audio_latent=audio_state.latent,
+            noiser=noiser,
+            width=width,
+            height=height,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(
+                context=video_context,
+                conditionings=stage_2_conditionings,
+                noise_scale=distilled_sigmas[0].item(),
+                initial_latent=upscaled_video_latent,
+            ),
+            audio=ModalitySpec(
+                context=audio_context,
+                noise_scale=distilled_sigmas[0].item(),
+                initial_latent=audio_state.latent,
+            ),
+            streaming_prefetch_count=streaming_prefetch_count,
         )
 
-        torch.cuda.synchronize()
-        del transformer
-        del video_encoder
-        cleanup_memory()
-
-        decoded_video = vae_decode_video(
-            video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config, generator
-        )
-        decoded_audio = vae_decode_audio(
-            audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder()
-        )
+        decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
+        decoded_audio = self.audio_decoder(audio_state.latent)
         return decoded_video, decoded_audio
 
     def _create_conditionings(
@@ -358,14 +315,8 @@ class ICLoraPipeline:
 
         for video_path, strength in video_conditioning:
             # Load video at scaled-down resolution (if scale > 1)
-            video = load_video_conditioning(
-                video_path=video_path,
-                height=ref_height,
-                width=ref_width,
-                frame_cap=num_frames,
-                dtype=self.dtype,
-                device=self.device,
-            )
+            frame_gen = decode_video_by_frame(path=video_path, frame_cap=num_frames, device=self.device)
+            video = video_preprocess(frame_gen, ref_height, ref_width, self.dtype, self.device)
             encoded_video = video_encoder(video)
             reference_video_shape = VideoLatentShape.from_torch_shape(encoded_video.shape)
 
@@ -509,6 +460,7 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
+        torch_compile=args.compile,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -525,6 +477,7 @@ 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(
@@ -553,17 +506,12 @@ def _load_mask_video(
     Returns:
         Tensor of shape ``(1, 1, F, H, W)`` with values in ``[0, 1]``.
     """
-    mask_video = load_video_conditioning(
-        video_path=mask_path,
-        height=height,
-        width=width,
-        frame_cap=num_frames,
-        dtype=torch.bfloat16,
-        device=device,
-    )
+    device = get_device()
+    frame_gen = decode_video_by_frame(path=mask_path, frame_cap=num_frames, device=device)
+    mask_video = video_preprocess(frame_gen, height, width, torch.bfloat16, device)
     # mask_video shape: (1, C, F, H, W) — take mean over channels for grayscale
     mask = mask_video.mean(dim=1, keepdim=True)  # (1, 1, F, H, W)
-    # Normalise to [0, 1] — load_video_conditioning applies normalize_latent,
+    # Normalise to [0, 1] — video_preprocess applies normalize_latent,
     # so undo that: values are in [-1, 1], remap to [0, 1]
     mask = (mask + 1.0) / 2.0
     return mask.clamp(0.0, 1.0)

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

@@ -3,40 +3,39 @@ from collections.abc import Iterator
 
 import torch
 
-from ltx_core.components.diffusion_steps import EulerDiffusionStep
 from ltx_core.components.guiders import (
     MultiModalGuiderFactory,
     MultiModalGuiderParams,
     create_multimodal_guider_factory,
 )
 from ltx_core.components.noisers import GaussianNoiser
-from ltx_core.components.protocols import DiffusionStepProtocol
 from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.loader import LoraPathStrengthAndSDOps
-from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
-from ltx_core.model.upsampler import upsample_video
+from ltx_core.loader.registry import Registry
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
-from ltx_core.model.video_vae import decode_video as vae_decode_video
 from ltx_core.quantization import QuantizationPolicy
-from ltx_core.types import Audio, LatentState, VideoPixelShape
-from ltx_pipelines.utils import ModelLedger
+from ltx_core.types import Audio, VideoPixelShape
 from ltx_pipelines.utils.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path
-from ltx_pipelines.utils.constants import STAGE_2_DISTILLED_SIGMA_VALUES, detect_params
+from ltx_pipelines.utils.blocks import (
+    AudioDecoder,
+    DiffusionStage,
+    ImageConditioner,
+    PromptEncoder,
+    VideoDecoder,
+    VideoUpsampler,
+)
+from ltx_pipelines.utils.constants import (
+    STAGE_2_DISTILLED_SIGMA_VALUES,
+    detect_params,
+)
+from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser, SimpleDenoiser
 from ltx_pipelines.utils.helpers import (
     assert_resolution,
-    cleanup_memory,
-    denoise_audio_video,
-    encode_prompts,
     get_device,
     image_conditionings_by_adding_guiding_latent,
-    multi_modal_guider_factory_denoising_func,
-    simple_denoising_func,
 )
 from ltx_pipelines.utils.media_io import encode_video
-from ltx_pipelines.utils.samplers import euler_denoising_loop
-from ltx_pipelines.utils.types import PipelineComponents
-
-device = get_device()
+from ltx_pipelines.utils.types import ModalitySpec
 
 
 class KeyframeInterpolationPipeline:
@@ -56,27 +55,40 @@ class KeyframeInterpolationPipeline:
         spatial_upsampler_path: str,
         gemma_root: str,
         loras: list[LoraPathStrengthAndSDOps],
-        device: torch.device = device,
+        device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
+        registry: Registry | None = None,
+        torch_compile: bool = False,
     ):
-        self.device = device
+        self.device = device or get_device()
         self.dtype = torch.bfloat16
-        self.stage_1_model_ledger = ModelLedger(
-            dtype=self.dtype,
-            device=device,
-            checkpoint_path=checkpoint_path,
-            spatial_upsampler_path=spatial_upsampler_path,
-            gemma_root_path=gemma_root,
-            loras=loras,
+
+        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)
+        self.stage_1 = DiffusionStage(
+            checkpoint_path,
+            self.dtype,
+            self.device,
+            loras=tuple(loras),
             quantization=quantization,
+            registry=registry,
+            torch_compile=torch_compile,
         )
-        self.stage_2_model_ledger = self.stage_1_model_ledger.with_additional_loras(
-            loras=distilled_lora,
+        stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
+        self.stage_2 = DiffusionStage(
+            checkpoint_path,
+            self.dtype,
+            self.device,
+            loras=stage_2_loras,
+            quantization=quantization,
+            registry=registry,
+            torch_compile=torch_compile,
         )
-        self.pipeline_components = PipelineComponents(
-            dtype=self.dtype,
-            device=device,
+        self.upsampler = VideoUpsampler(
+            checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
         )
+        self.video_decoder = VideoDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
+        self.audio_decoder = AudioDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
 
     def __call__(  # noqa: PLR0913
         self,
@@ -93,52 +105,28 @@ class KeyframeInterpolationPipeline:
         images: list[ImageConditioningInput],
         tiling_config: TilingConfig | None = None,
         enhance_prompt: bool = False,
+        streaming_prefetch_count: int | None = None,
+        max_batch_size: int = 1,
     ) -> tuple[Iterator[torch.Tensor], Audio]:
         assert_resolution(height=height, width=width, is_two_stage=True)
 
         generator = torch.Generator(device=self.device).manual_seed(seed)
         noiser = GaussianNoiser(generator=generator)
-        stepper = EulerDiffusionStep()
         dtype = torch.bfloat16
 
-        ctx_p, ctx_n = encode_prompts(
+        ctx_p, ctx_n = self.prompt_encoder(
             [prompt, negative_prompt],
-            self.stage_1_model_ledger,
             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,
         )
         v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
         v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
 
         # Stage 1: Initial low resolution video generation.
-        video_encoder = self.stage_1_model_ledger.video_encoder()
-        transformer = self.stage_1_model_ledger.transformer()
         sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
 
-        def first_stage_denoising_loop(
-            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
-        ) -> tuple[LatentState, LatentState]:
-            return euler_denoising_loop(
-                sigmas=sigmas,
-                video_state=video_state,
-                audio_state=audio_state,
-                stepper=stepper,
-                denoise_fn=multi_modal_guider_factory_denoising_func(
-                    video_guider_factory=create_multimodal_guider_factory(
-                        params=video_guider_params,
-                        negative_context=v_context_n,
-                    ),
-                    audio_guider_factory=create_multimodal_guider_factory(
-                        params=audio_guider_params,
-                        negative_context=a_context_n,
-                    ),
-                    v_context=v_context_p,
-                    a_context=a_context_p,
-                    transformer=transformer,  # noqa: F821
-                ),
-            )
-
         stage_1_output_shape = VideoPixelShape(
             batch=1,
             frames=num_frames,
@@ -146,93 +134,90 @@ class KeyframeInterpolationPipeline:
             height=height // 2,
             fps=frame_rate,
         )
-        stage_1_conditionings = image_conditionings_by_adding_guiding_latent(
-            images=images,
-            height=stage_1_output_shape.height,
-            width=stage_1_output_shape.width,
-            video_encoder=video_encoder,
-            dtype=dtype,
-            device=self.device,
-        )
-        video_state, audio_state = denoise_audio_video(
-            output_shape=stage_1_output_shape,
-            conditionings=stage_1_conditionings,
-            noiser=noiser,
-            sigmas=sigmas,
-            stepper=stepper,
-            denoising_loop_fn=first_stage_denoising_loop,
-            components=self.pipeline_components,
-            dtype=dtype,
-            device=self.device,
+        stage_1_conditionings = self.image_conditioner(
+            lambda enc: image_conditionings_by_adding_guiding_latent(
+                images=images,
+                height=stage_1_output_shape.height,
+                width=stage_1_output_shape.width,
+                video_encoder=enc,
+                dtype=dtype,
+                device=self.device,
+            )
         )
 
-        torch.cuda.synchronize()
-        del transformer
-        cleanup_memory()
+        video_guider_factory = create_multimodal_guider_factory(
+            params=video_guider_params,
+            negative_context=v_context_n,
+        )
+        audio_guider_factory = create_multimodal_guider_factory(
+            params=audio_guider_params,
+            negative_context=a_context_n,
+        )
+
+        video_state, audio_state = self.stage_1(
+            denoiser=FactoryGuidedDenoiser(
+                v_context=v_context_p,
+                a_context=a_context_p,
+                video_guider_factory=video_guider_factory,
+                audio_guider_factory=audio_guider_factory,
+            ),
+            sigmas=sigmas,
+            noiser=noiser,
+            width=stage_1_output_shape.width,
+            height=stage_1_output_shape.height,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(
+                context=v_context_p,
+                conditionings=stage_1_conditionings,
+            ),
+            audio=ModalitySpec(
+                context=a_context_p,
+            ),
+            streaming_prefetch_count=streaming_prefetch_count,
+            max_batch_size=max_batch_size,
+        )
 
         # Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
-        upscaled_video_latent = upsample_video(
-            latent=video_state.latent[:1],
-            video_encoder=video_encoder,
-            upsampler=self.stage_2_model_ledger.spatial_upsampler(),
-        )
+        upscaled_video_latent = self.upsampler(video_state.latent[:1])
 
-        torch.cuda.synchronize()
-        cleanup_memory()
-
-        transformer = self.stage_2_model_ledger.transformer()
         distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
-
-        def second_stage_denoising_loop(
-            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
-        ) -> tuple[LatentState, LatentState]:
-            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
-                ),
-            )
-
         stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
-        stage_2_conditionings = image_conditionings_by_adding_guiding_latent(
-            images=images,
-            height=stage_2_output_shape.height,
-            width=stage_2_output_shape.width,
-            video_encoder=video_encoder,
-            dtype=dtype,
-            device=self.device,
+        stage_2_conditionings = self.image_conditioner(
+            lambda enc: image_conditionings_by_adding_guiding_latent(
+                images=images,
+                height=stage_2_output_shape.height,
+                width=stage_2_output_shape.width,
+                video_encoder=enc,
+                dtype=dtype,
+                device=self.device,
+            )
         )
-        video_state, audio_state = denoise_audio_video(
-            output_shape=stage_2_output_shape,
-            conditionings=stage_2_conditionings,
-            noiser=noiser,
+
+        video_state, audio_state = self.stage_2(
+            denoiser=SimpleDenoiser(v_context_p, a_context_p),
             sigmas=distilled_sigmas,
-            stepper=stepper,
-            denoising_loop_fn=second_stage_denoising_loop,
-            components=self.pipeline_components,
-            dtype=dtype,
-            device=self.device,
-            noise_scale=distilled_sigmas[0],
-            initial_video_latent=upscaled_video_latent,
-            initial_audio_latent=audio_state.latent,
+            noiser=noiser,
+            width=width,
+            height=height,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(
+                context=v_context_p,
+                conditionings=stage_2_conditionings,
+                noise_scale=distilled_sigmas[0].item(),
+                initial_latent=upscaled_video_latent,
+            ),
+            audio=ModalitySpec(
+                context=a_context_p,
+                noise_scale=distilled_sigmas[0].item(),
+                initial_latent=audio_state.latent,
+            ),
+            streaming_prefetch_count=streaming_prefetch_count,
         )
 
-        torch.cuda.synchronize()
-        del transformer
-        del video_encoder
-        cleanup_memory()
-
-        decoded_video = vae_decode_video(
-            video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config, generator
-        )
-        decoded_audio = vae_decode_audio(
-            audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder()
-        )
+        decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
+        decoded_audio = self.audio_decoder(audio_state.latent)
         return decoded_video, decoded_audio
 
 
@@ -250,6 +235,7 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
+        torch_compile=args.compile,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -280,6 +266,8 @@ def main() -> None:
         ),
         images=args.images,
         tiling_config=tiling_config,
+        streaming_prefetch_count=args.streaming_prefetch_count,
+        max_batch_size=args.max_batch_size,
     )
 
     encode_video(

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

@@ -1,153 +1,42 @@
 from __future__ import annotations
 
-import argparse
 import logging
 from collections.abc import Iterator
-from dataclasses import dataclass
 
 import torch
 
-from ltx_core.components.diffusion_steps import EulerDiffusionStep
 from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
 from ltx_core.components.noisers import GaussianNoiser
-from ltx_core.components.patchifiers import get_pixel_coords
-from ltx_core.components.protocols import DiffusionStepProtocol
 from ltx_core.components.schedulers import LTX2Scheduler
-from ltx_core.conditioning import ConditioningItem
+from ltx_core.conditioning.types.noise_mask_cond import TemporalRegionMask
 from ltx_core.loader import LoraPathStrengthAndSDOps
-from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
-from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
+from ltx_core.loader.registry import Registry
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
-from ltx_core.model.video_vae import decode_video as vae_decode_video
 from ltx_core.quantization import QuantizationPolicy
-from ltx_core.tools import LatentTools
 from ltx_core.types import (
-    Audio,
-    AudioLatentShape,
-    LatentState,
     SpatioTemporalScaleFactors,
-    VideoPixelShape,
 )
-from ltx_pipelines.utils import ModelLedger
-from ltx_pipelines.utils.args import QuantizationAction
+from ltx_pipelines.utils.args import video_editing_arg_parser
+from ltx_pipelines.utils.blocks import (
+    AudioConditioner,
+    AudioDecoder,
+    DiffusionStage,
+    ImageConditioner,
+    PromptEncoder,
+    VideoDecoder,
+)
 from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES, detect_params
+from ltx_pipelines.utils.denoisers import GuidedDenoiser, SimpleDenoiser
 from ltx_pipelines.utils.helpers import (
-    cleanup_memory,
-    encode_prompts,
+    audio_latent_from_file,
     get_device,
-    multi_modal_guider_denoising_func,
-    noise_audio_state,
-    noise_video_state,
-    simple_denoising_func,
+    video_latent_from_file,
 )
 from ltx_pipelines.utils.media_io import (
-    decode_audio_from_file,
     encode_video,
     get_videostream_metadata,
-    load_video_conditioning,
 )
-from ltx_pipelines.utils.samplers import euler_denoising_loop
-from ltx_pipelines.utils.types import PipelineComponents
-
-device = get_device()
-
-
-def _encode_video_for_retake(
-    video_encoder: torch.nn.Module,
-    video_path: str,
-    output_shape: VideoPixelShape,
-    dtype: torch.dtype,
-    device: torch.device,
-) -> torch.Tensor:
-    """Load video and encode to latents."""
-    pixel_video = load_video_conditioning(
-        video_path=video_path,
-        height=output_shape.height,
-        width=output_shape.width,
-        frame_cap=output_shape.frames,
-        dtype=dtype,
-        device=device,
-    )  # (1, C, F, H, W)
-    return video_encoder(pixel_video)
-
-
-def _encode_audio_for_retake(
-    audio_encoder: torch.nn.Module,
-    waveform: torch.Tensor,
-    waveform_sr: int,
-    output_shape: VideoPixelShape,
-    dtype: torch.dtype,
-) -> torch.Tensor:
-    """Encode audio to latents and trim/pad to match output_shape."""
-    waveform_batch = waveform.unsqueeze(0) if waveform.dim() == 2 else waveform
-    initial_audio_latent = vae_encode_audio(
-        Audio(waveform=waveform_batch.to(dtype), sampling_rate=waveform_sr), audio_encoder, None
-    )
-    expected_audio_shape = AudioLatentShape.from_video_pixel_shape(output_shape)
-    expected_frames = expected_audio_shape.frames
-    actual_frames = initial_audio_latent.shape[2]
-    if actual_frames > expected_frames:
-        initial_audio_latent = initial_audio_latent[:, :, :expected_frames, :]
-    elif actual_frames < expected_frames:
-        pad = torch.zeros(
-            initial_audio_latent.shape[0],
-            initial_audio_latent.shape[1],
-            expected_frames - actual_frames,
-            initial_audio_latent.shape[3],
-            device=initial_audio_latent.device,
-            dtype=initial_audio_latent.dtype,
-        )
-        initial_audio_latent = torch.cat([initial_audio_latent, pad], dim=2)
-    return initial_audio_latent
-
-
-# ---------------------------------------------------------------------------
-# Custom conditioning item: temporal region mask
-# ---------------------------------------------------------------------------
-
-
-@dataclass(frozen=True)
-class TemporalRegionMask:
-    """Conditioning item that sets ``denoise_mask = 0`` outside a time range
-    and ``1`` inside, so only the specified temporal region is regenerated.
-    Uses ``start_time`` and ``end_time`` in seconds. Works in *patchified*
-    (token) space using the patchifier's ``get_patch_grid_bounds``: for video
-    coords are latent frame indices (converted from seconds via ``fps``), for
-    audio coords are already in seconds.
-    """
-
-    start_time: float  # seconds, inclusive
-    end_time: float  # seconds, exclusive
-    fps: float
-
-    def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
-        coords = latent_tools.patchifier.get_patch_grid_bounds(
-            latent_tools.target_shape, device=latent_state.denoise_mask.device
-        )
-        # coords: [B, 3, N, 2] (video) or [B, 1, N, 2] (audio); temporal dim is index 0
-        if coords.shape[1] == 1:
-            # Audio: patchifier returns seconds
-            t_start = coords[:, 0, :, 0]  # [B, N]
-            t_end = coords[:, 0, :, 1]  # [B, N]
-            in_region = (t_end > self.start_time) & (t_start < self.end_time)
-        else:
-            # Video: get pixel bounds per patch, find patches for start/end frame, read latent from coords.
-            scale_factors = getattr(latent_tools, "scale_factors", SpatioTemporalScaleFactors.default())
-            pixel_bounds = get_pixel_coords(coords, scale_factors, causal_fix=getattr(latent_tools, "causal_fix", True))
-            timestamp_bounds = pixel_bounds[0, 0] / self.fps
-            t_start, t_end = timestamp_bounds.unbind(dim=-1)
-            in_region = (t_end > self.start_time) & (t_start < self.end_time)
-        state = latent_state.clone()
-        mask_val = in_region.to(state.denoise_mask.dtype)
-        if state.denoise_mask.dim() == 3:
-            mask_val = mask_val.unsqueeze(-1)
-        state.denoise_mask.copy_(mask_val)
-        return state
-
-
-# ---------------------------------------------------------------------------
-# Pipeline
-# ---------------------------------------------------------------------------
+from ltx_pipelines.utils.types import ModalitySpec
 
 
 class RetakePipeline:
@@ -168,6 +57,11 @@ class RetakePipeline:
         Target device (default: CUDA if available).
     quantization : QuantizationPolicy | None
         Optional quantization policy for the transformer.
+    distilled : bool
+        Set to ``True`` if using distilled model or passing distillation
+        lora with full model. If set to ``True``, distilled sigma schedule
+        (``DISTILLED_SIGMA_VALUES``) and a simple (non-guided) denoising
+        function will be used during ``__call__``.
     """
 
     def __init__(
@@ -175,29 +69,61 @@ class RetakePipeline:
         checkpoint_path: str,
         gemma_root: str,
         loras: list[LoraPathStrengthAndSDOps],
-        device: torch.device = device,
+        device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
+        registry: Registry | None = None,
+        distilled: bool = True,
+        torch_compile: bool = False,
     ):
-        self.device = device
+        self.device = device or get_device()
         self.dtype = torch.bfloat16
-        self.model_ledger = ModelLedger(
-            dtype=self.dtype,
-            device=device,
+        self.distilled = distilled
+        self.prompt_encoder = PromptEncoder(
             checkpoint_path=checkpoint_path,
-            gemma_root_path=gemma_root,
-            loras=loras,
-            quantization=quantization,
-        )
-        self.pipeline_components = PipelineComponents(
+            gemma_root=gemma_root,
             dtype=self.dtype,
-            device=device,
+            device=self.device,
+            registry=registry,
+        )
+        self.image_conditioner = ImageConditioner(
+            checkpoint_path=checkpoint_path,
+            dtype=self.dtype,
+            device=self.device,
+            registry=registry,
+        )
+        self.audio_conditioner = AudioConditioner(
+            checkpoint_path=checkpoint_path,
+            dtype=self.dtype,
+            device=self.device,
+            registry=registry,
+        )
+        self.stage = DiffusionStage(
+            checkpoint_path=checkpoint_path,
+            dtype=self.dtype,
+            device=self.device,
+            loras=tuple(loras),
+            quantization=quantization,
+            registry=registry,
+            torch_compile=torch_compile,
+        )
+        self.video_decoder = VideoDecoder(
+            checkpoint_path=checkpoint_path,
+            dtype=self.dtype,
+            device=self.device,
+            registry=registry,
+        )
+        self.audio_decoder = AudioDecoder(
+            checkpoint_path=checkpoint_path,
+            dtype=self.dtype,
+            device=self.device,
+            registry=registry,
         )
 
     # --------------------------------------------------------------------- #
     #  Public entry point                                                     #
     # --------------------------------------------------------------------- #
 
-    def __call__(  # noqa: PLR0913, PLR0915
+    def __call__(  # noqa: PLR0913
         self,
         video_path: str,
         prompt: str,
@@ -212,8 +138,9 @@ class RetakePipeline:
         regenerate_video: bool = True,
         regenerate_audio: bool = True,
         enhance_prompt: bool = False,
-        distilled: bool = False,
         tiling_config: TilingConfig | None = None,
+        streaming_prefetch_count: int | None = None,
+        max_batch_size: int = 1,
     ) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
         """Regenerate ``[start_time, end_time]`` of the source video (retake).
         Parameters
@@ -235,19 +162,13 @@ class RetakePipeline:
             Guidance parameters for video and audio modalities.  Ignored in
             distilled mode.
         regenerate_video : bool
-            If ``True`` (default), preserve video outside ``[start_time, end_time]``
-            and only regenerate the masked region.  If ``False``, fully regenerate
-            all video frames (the encoded video is still used as the initial latent
-            but with ``denoise_mask = 1`` everywhere).
+            If ``True`` (default), regenerate video inside ``[start_time, end_time]``.
+            If ``False``, video is preserved as-is (no regeneration).
         regenerate_audio : bool
             If True, regenerate audio in the [start_time, end_time] window; if False,
             audio is preserved as-is (no regeneration).
         enhance_prompt : bool
             Whether to enhance the prompt via the text encoder.
-        distilled : bool
-            If ``True``, use the distilled sigma schedule
-            (``DISTILLED_SIGMA_VALUES``) and a simple (non-guided) denoising
-            function.  The model checkpoint must be the distilled variant.
         Returns
         -------
         tuple[Iterator[torch.Tensor], torch.Tensor]
@@ -256,95 +177,66 @@ class RetakePipeline:
         if start_time >= end_time:
             raise ValueError(f"start_time ({start_time}) must be less than end_time ({end_time})")
 
-        effective_seed = torch.randint(0, 2**31, (1,), device=self.device).item() if seed < 0 else seed
-        generator = torch.Generator(device=self.device).manual_seed(effective_seed)
+        generator = torch.Generator(device=self.device).manual_seed(seed)
         noiser = GaussianNoiser(generator=generator)
-        stepper = EulerDiffusionStep()
         dtype = self.dtype
 
-        video_encoder = self.model_ledger.video_encoder()
-
-        # Use av to get metadata
-        fps, num_pixel_frames, src_width, src_height = get_videostream_metadata(video_path)
-
-        output_shape = VideoPixelShape(
-            batch=1,
-            frames=num_pixel_frames,
-            width=src_width,
-            height=src_height,
-            fps=fps,
-        )
-        initial_video_latent = _encode_video_for_retake(
-            video_encoder=video_encoder,
-            video_path=video_path,
-            output_shape=output_shape,
-            dtype=dtype,
-            device=self.device,
-        )
-        video_conditionings: list[ConditioningItem] = [
-            TemporalRegionMask(
-                start_time=start_time if regenerate_video else 0.0,
-                end_time=end_time if regenerate_video else 0.0,
-                fps=fps,
-            )
-        ]
-        del video_encoder
-        cleanup_memory()
-
-        initial_audio_latent: torch.Tensor | None = None
-        audio_conditionings: list[ConditioningItem] = []
-
-        audio_in = decode_audio_from_file(video_path, self.device)
-        audio_encoder = self.model_ledger.audio_encoder()
-
-        if audio_in is not None:
-            waveform = audio_in.waveform.squeeze(0)
-            waveform_sr = audio_in.sampling_rate
-        else:
-            waveform, waveform_sr = None, None
-        if waveform is not None:
-            initial_audio_latent = _encode_audio_for_retake(
-                audio_encoder=audio_encoder,
-                waveform=waveform,
-                waveform_sr=waveform_sr,
+        output_shape = get_videostream_metadata(video_path)
+        initial_video_latent = self.image_conditioner(
+            lambda enc: video_latent_from_file(
+                video_encoder=enc,
+                file_path=video_path,
                 output_shape=output_shape,
                 dtype=dtype,
+                device=self.device,
             )
-            audio_conditionings = [
-                TemporalRegionMask(
-                    start_time=start_time if regenerate_audio else 0.0,
-                    end_time=end_time if regenerate_audio else 0.0,
-                    fps=fps,
-                )
-            ]
+        )
 
-        del audio_encoder
-        cleanup_memory()
+        initial_audio_latent = self.audio_conditioner(
+            lambda enc: audio_latent_from_file(
+                audio_encoder=enc,
+                file_path=video_path,
+                output_shape=output_shape,
+                dtype=dtype,
+                device=self.device,
+            )
+        )
 
-        prompts_to_encode = [prompt] if distilled else [prompt, negative_prompt]
-        contexts = encode_prompts(
+        prompts_to_encode = [prompt] if self.distilled else [prompt, negative_prompt]
+        contexts = self.prompt_encoder(
             prompts_to_encode,
-            self.model_ledger,
             enhance_first_prompt=enhance_prompt,
-            enhance_prompt_seed=effective_seed,
+            enhance_prompt_seed=seed,
+            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         v_context_p, a_context_p = contexts[0].video_encoding, contexts[0].audio_encoding
-        if not distilled:
-            v_context_n, a_context_n = contexts[1].video_encoding, contexts[1].audio_encoding
-
-        transformer = self.model_ledger.transformer()
-
-        sigmas = (
-            torch.tensor(DISTILLED_SIGMA_VALUES) if distilled else LTX2Scheduler().execute(steps=num_inference_steps)
-        ).to(dtype=torch.float32, device=self.device)
-        if distilled:
-            denoise_fn = simple_denoising_func(
-                video_context=v_context_p,
-                audio_context=a_context_p,
-                transformer=transformer,
+        video_modality_spec = ModalitySpec(
+            context=v_context_p,
+            conditionings=[TemporalRegionMask(start_time=start_time, end_time=end_time, fps=output_shape.fps)]
+            if regenerate_video
+            else [],
+            initial_latent=initial_video_latent,
+            frozen=not regenerate_video,
+        )
+        audio_modality_spec = ModalitySpec(
+            context=a_context_p,
+            conditionings=[TemporalRegionMask(start_time=start_time, end_time=end_time, fps=output_shape.fps)]
+            if (initial_audio_latent is not None and regenerate_audio)
+            else [],
+            initial_latent=initial_audio_latent,
+            frozen=initial_audio_latent is not None and not regenerate_audio,
+        )
+        # Build denoiser
+        if self.distilled:
+            sigmas = torch.tensor(DISTILLED_SIGMA_VALUES).to(dtype=torch.float32, device=self.device)
+            denoiser = SimpleDenoiser(
+                v_context=v_context_p,
+                a_context=a_context_p,
             )
         else:
+            sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
+            v_context_n, a_context_n = contexts[1].video_encoding, contexts[1].audio_encoding
             video_guider = MultiModalGuider(
                 params=video_guider_params,
                 negative_context=v_context_n,
@@ -353,66 +245,31 @@ class RetakePipeline:
                 params=audio_guider_params,
                 negative_context=a_context_n,
             )
-            denoise_fn = multi_modal_guider_denoising_func(
-                video_guider,
-                audio_guider,
+            denoiser = GuidedDenoiser(
                 v_context=v_context_p,
                 a_context=a_context_p,
-                transformer=transformer,
+                video_guider=video_guider,
+                audio_guider=audio_guider,
             )
 
-        def denoising_loop(
-            sigmas: torch.Tensor,
-            video_state: LatentState,
-            audio_state: LatentState,
-            stepper: DiffusionStepProtocol,
-        ) -> tuple[LatentState, LatentState]:
-            return euler_denoising_loop(
-                sigmas=sigmas,
-                video_state=video_state,
-                audio_state=audio_state,
-                stepper=stepper,
-                denoise_fn=denoise_fn,
-            )
-
-        # Build noised states with the encoded latents as initial values and
-        # the temporal masks applied via conditionings.
-        video_state, video_tools = noise_video_state(
-            output_shape=output_shape,
+        # Run diffusion stage
+        video_state, audio_state = self.stage(
+            denoiser=denoiser,
+            sigmas=sigmas,
             noiser=noiser,
-            conditionings=video_conditionings,
-            components=self.pipeline_components,
-            dtype=dtype,
-            device=self.device,
-            initial_latent=initial_video_latent,
-        )
-        audio_state, audio_tools = noise_audio_state(
-            output_shape=output_shape,
-            noiser=noiser,
-            conditionings=audio_conditionings,
-            components=self.pipeline_components,
-            dtype=dtype,
-            device=self.device,
-            initial_latent=initial_audio_latent,
+            width=output_shape.width,
+            height=output_shape.height,
+            frames=output_shape.frames,
+            fps=output_shape.fps,
+            video=video_modality_spec,
+            audio=audio_modality_spec,
+            streaming_prefetch_count=streaming_prefetch_count,
+            max_batch_size=max_batch_size,
         )
 
-        video_state, audio_state = denoising_loop(sigmas, video_state, audio_state, stepper)
-
-        video_state = video_tools.clear_conditioning(video_state)
-        video_state = video_tools.unpatchify(video_state)
-        audio_state = audio_tools.clear_conditioning(audio_state)
-        audio_state = audio_tools.unpatchify(audio_state)
-
-        torch.cuda.synchronize()
-        del transformer
-        cleanup_memory()
-
-        decoded_video = vae_decode_video(
-            video_state.latent, self.model_ledger.video_decoder(), tiling_config, generator
-        )
-        decoded_audio = vae_decode_audio(
-            audio_state.latent, self.model_ledger.audio_decoder(), self.model_ledger.vocoder()
-        )
+        # Decode
+        decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
+        decoded_audio = self.audio_decoder(audio_state.latent)
 
         return decoded_video, decoded_audio
 
@@ -421,25 +278,8 @@ class RetakePipeline:
 def main() -> None:
     """CLI entry point for retake (regenerate a time region)."""
     logging.getLogger().setLevel(logging.INFO)
-    parser = argparse.ArgumentParser(description="Retake: regenerate a time region of a video with LTX-2.")
-    parser.add_argument("--video-path", type=str, required=True, help="Path to the source video.")
-    parser.add_argument("--prompt", type=str, required=True, help="Text prompt for the regenerated region.")
-    parser.add_argument("--start-time", type=float, required=True, help="Start time of the region to regenerate (s).")
-    parser.add_argument("--end-time", type=float, required=True, help="End time of the region to regenerate (s).")
-    parser.add_argument("--output-path", type=str, required=True, help="Path for the output video.")
-    parser.add_argument("--checkpoint-path", type=str, required=True, help="Path to the LTX-2 checkpoint.")
-    parser.add_argument("--gemma-root", type=str, required=True, help="Path to Gemma text encoder weights.")
-    parser.add_argument("--seed", type=int, default=42, help="Random seed. Use -1 for a random seed.")
-    parser.add_argument("--loras", nargs="*", default=[], help="LoRA paths (optional).")
-    parser.add_argument(
-        "--quantization",
-        dest="quantization",
-        action=QuantizationAction,
-        nargs="+",
-        metavar=("POLICY", "AMAX_PATH"),
-        default=None,
-        help="Quantization policy: fp8-cast or fp8-scaled-mm [AMAX_PATH].",
-    )
+    parser = video_editing_arg_parser(distilled=True)
+    parser.description = "Retake: regenerate a time region of a video with LTX-2."
     args = parser.parse_args()
 
     if args.start_time >= args.end_time:
@@ -447,22 +287,24 @@ def main() -> None:
 
     # Validate frame count (8k+1) and resolution (multiples of 32) at CLI stage
     video_scale = SpatioTemporalScaleFactors.default()
-    fps, num_frames, width, height = get_videostream_metadata(args.video_path)
-    if (num_frames - 1) % video_scale.time != 0:
-        snapped = ((num_frames - 1) // video_scale.time) * video_scale.time + 1
+    src = get_videostream_metadata(args.video_path)
+    if (src.frames - 1) % video_scale.time != 0:
+        snapped = ((src.frames - 1) // video_scale.time) * video_scale.time + 1
         raise ValueError(
-            f"Video frame count must satisfy 8k+1 (e.g. 97, 193). Got {num_frames}; use a video with {snapped} frames."
+            f"Video frame count must satisfy 8k+1 (e.g. 97, 193). Got {src.frames}; use a video with {snapped} frames."
         )
-    if width % 32 != 0 or height % 32 != 0:
-        raise ValueError(f"Video width and height must be multiples of 32. Got {width}x{height}.")
+    if src.width % 32 != 0 or src.height % 32 != 0:
+        raise ValueError(f"Video width and height must be multiples of 32. Got {src.width}x{src.height}.")
 
     pipeline = RetakePipeline(
-        checkpoint_path=args.checkpoint_path,
+        checkpoint_path=args.distilled_checkpoint_path,
         gemma_root=args.gemma_root,
-        loras=tuple(args.loras) if args.loras else (),
+        loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
+        distilled=args.distilled,
+        torch_compile=args.compile,
     )
-    params = detect_params(args.checkpoint_path)
+    params = detect_params(args.distilled_checkpoint_path)
     tiling_config = TilingConfig.default()
     video_iter, audio = pipeline(
         video_path=args.video_path,
@@ -473,11 +315,13 @@ def main() -> None:
         video_guider_params=params.video_guider_params,
         audio_guider_params=params.audio_guider_params,
         tiling_config=tiling_config,
+        streaming_prefetch_count=args.streaming_prefetch_count,
+        max_batch_size=args.max_batch_size,
     )
-    video_chunks_number = get_video_chunks_number(num_frames, tiling_config)
+    video_chunks_number = get_video_chunks_number(src.frames, tiling_config)
     encode_video(
         video=video_iter,
-        fps=int(fps),
+        fps=int(src.fps),
         audio=audio,
         output_path=args.output_path,
         video_chunks_number=video_chunks_number,

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

@@ -3,37 +3,35 @@ from collections.abc import Iterator
 
 import torch
 
-from ltx_core.components.diffusion_steps import EulerDiffusionStep
 from ltx_core.components.guiders import (
     MultiModalGuiderFactory,
     MultiModalGuiderParams,
     create_multimodal_guider_factory,
 )
 from ltx_core.components.noisers import GaussianNoiser
-from ltx_core.components.protocols import DiffusionStepProtocol
 from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.loader import LoraPathStrengthAndSDOps
-from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
-from ltx_core.model.video_vae import decode_video as vae_decode_video
+from ltx_core.loader.registry import Registry
+from ltx_core.model.video_vae.tiling import TilingConfig
 from ltx_core.quantization import QuantizationPolicy
-from ltx_core.types import Audio, LatentState, VideoPixelShape
+from ltx_core.types import Audio
 from ltx_pipelines.utils import (
-    ModelLedger,
     assert_resolution,
-    cleanup_memory,
     combined_image_conditionings,
-    denoise_audio_video,
-    encode_prompts,
-    euler_denoising_loop,
     get_device,
-    multi_modal_guider_factory_denoising_func,
 )
 from ltx_pipelines.utils.args import ImageConditioningInput, default_1_stage_arg_parser, detect_checkpoint_path
+from ltx_pipelines.utils.blocks import (
+    AudioDecoder,
+    DiffusionStage,
+    ImageConditioner,
+    PromptEncoder,
+    VideoDecoder,
+)
 from ltx_pipelines.utils.constants import detect_params
+from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser
 from ltx_pipelines.utils.media_io import encode_video
-from ltx_pipelines.utils.types import PipelineComponents
-
-device = get_device()
+from ltx_pipelines.utils.types import ModalitySpec
 
 
 class TI2VidOneStagePipeline:
@@ -50,22 +48,46 @@ class TI2VidOneStagePipeline:
         checkpoint_path: str,
         gemma_root: str,
         loras: list[LoraPathStrengthAndSDOps],
-        device: torch.device = device,
+        device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
+        registry: Registry | None = None,
+        torch_compile: bool = False,
     ):
         self.dtype = torch.bfloat16
-        self.device = device
-        self.model_ledger = ModelLedger(
-            dtype=self.dtype,
-            device=device,
+        self.device = device or get_device()
+        self.prompt_encoder = PromptEncoder(
             checkpoint_path=checkpoint_path,
-            gemma_root_path=gemma_root,
-            loras=loras,
-            quantization=quantization,
-        )
-        self.pipeline_components = PipelineComponents(
+            gemma_root=gemma_root,
             dtype=self.dtype,
-            device=device,
+            device=self.device,
+            registry=registry,
+        )
+        self.image_conditioner = ImageConditioner(
+            checkpoint_path=checkpoint_path,
+            dtype=self.dtype,
+            device=self.device,
+            registry=registry,
+        )
+        self.stage = DiffusionStage(
+            checkpoint_path=checkpoint_path,
+            dtype=self.dtype,
+            device=self.device,
+            loras=tuple(loras),
+            quantization=quantization,
+            registry=registry,
+            torch_compile=torch_compile,
+        )
+        self.video_decoder = VideoDecoder(
+            checkpoint_path=checkpoint_path,
+            dtype=self.dtype,
+            device=self.device,
+            registry=registry,
+        )
+        self.audio_decoder = AudioDecoder(
+            checkpoint_path=checkpoint_path,
+            dtype=self.dtype,
+            device=self.device,
+            registry=registry,
         )
 
     def __call__(  # noqa: PLR0913
@@ -82,41 +104,37 @@ class TI2VidOneStagePipeline:
         audio_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
         images: list[ImageConditioningInput],
         enhance_prompt: bool = False,
+        streaming_prefetch_count: int | None = None,
+        tiling_config: TilingConfig | None = None,
+        max_batch_size: int = 1,
     ) -> tuple[Iterator[torch.Tensor], Audio]:
         assert_resolution(height=height, width=width, is_two_stage=False)
 
         generator = torch.Generator(device=self.device).manual_seed(seed)
         noiser = GaussianNoiser(generator=generator)
-        stepper = EulerDiffusionStep()
         dtype = torch.bfloat16
 
-        ctx_p, ctx_n = encode_prompts(
+        ctx_p, ctx_n = self.prompt_encoder(
             [prompt, negative_prompt],
-            self.model_ledger,
             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,
         )
         v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
         v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
 
-        # Encode image conditionings with the VAE encoder, then free it
-        # before loading the transformer to reduce peak VRAM.
-        stage_1_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
-        video_encoder = self.model_ledger.video_encoder()
-        stage_1_conditionings = combined_image_conditionings(
-            images=images,
-            height=stage_1_output_shape.height,
-            width=stage_1_output_shape.width,
-            video_encoder=video_encoder,
-            dtype=dtype,
-            device=self.device,
+        stage_1_conditionings = self.image_conditioner(
+            lambda enc: combined_image_conditionings(
+                images=images,
+                height=height,
+                width=width,
+                video_encoder=enc,
+                dtype=dtype,
+                device=self.device,
+            )
         )
-        torch.cuda.synchronize()
-        del video_encoder
-        cleanup_memory()
 
-        transformer = self.model_ledger.transformer()
         sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
 
         video_guider_factory = create_multimodal_guider_factory(
@@ -128,43 +146,32 @@ class TI2VidOneStagePipeline:
             negative_context=a_context_n,
         )
 
-        def first_stage_denoising_loop(
-            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
-        ) -> tuple[LatentState, LatentState]:
-            return euler_denoising_loop(
-                sigmas=sigmas,
-                video_state=video_state,
-                audio_state=audio_state,
-                stepper=stepper,
-                denoise_fn=multi_modal_guider_factory_denoising_func(
-                    video_guider_factory=video_guider_factory,
-                    audio_guider_factory=audio_guider_factory,
-                    v_context=v_context_p,
-                    a_context=a_context_p,
-                    transformer=transformer,  # noqa: F821
-                ),
-            )
-
-        video_state, audio_state = denoise_audio_video(
-            output_shape=stage_1_output_shape,
-            conditionings=stage_1_conditionings,
-            noiser=noiser,
+        video_state, audio_state = self.stage(
+            denoiser=FactoryGuidedDenoiser(
+                v_context=v_context_p,
+                a_context=a_context_p,
+                video_guider_factory=video_guider_factory,
+                audio_guider_factory=audio_guider_factory,
+            ),
             sigmas=sigmas,
-            stepper=stepper,
-            denoising_loop_fn=first_stage_denoising_loop,
-            components=self.pipeline_components,
-            dtype=dtype,
-            device=self.device,
+            noiser=noiser,
+            width=width,
+            height=height,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(
+                context=v_context_p,
+                conditionings=stage_1_conditionings,
+            ),
+            audio=ModalitySpec(
+                context=a_context_p,
+            ),
+            streaming_prefetch_count=streaming_prefetch_count,
+            max_batch_size=max_batch_size,
         )
 
-        torch.cuda.synchronize()
-        del transformer
-        cleanup_memory()
-
-        decoded_video = vae_decode_video(video_state.latent, self.model_ledger.video_decoder(), generator=generator)
-        decoded_audio = vae_decode_audio(
-            audio_state.latent, self.model_ledger.audio_decoder(), self.model_ledger.vocoder()
-        )
+        decoded_video = self.video_decoder(video_state.latent, tiling_config, generator=generator)
+        decoded_audio = self.audio_decoder(audio_state.latent)
         return decoded_video, decoded_audio
 
 
@@ -180,6 +187,7 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
+        torch_compile=args.compile,
     )
     video, audio = pipeline(
         prompt=args.prompt,
@@ -207,6 +215,8 @@ def main() -> None:
             stg_blocks=args.audio_stg_blocks,
         ),
         images=args.images,
+        streaming_prefetch_count=args.streaming_prefetch_count,
+        max_batch_size=args.max_batch_size,
     )
 
     encode_video(

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

@@ -3,40 +3,39 @@ from collections.abc import Iterator
 
 import torch
 
-from ltx_core.components.diffusion_steps import EulerDiffusionStep
 from ltx_core.components.guiders import (
     MultiModalGuiderFactory,
     MultiModalGuiderParams,
     create_multimodal_guider_factory,
 )
 from ltx_core.components.noisers import GaussianNoiser
-from ltx_core.components.protocols import DiffusionStepProtocol
 from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.loader import LoraPathStrengthAndSDOps
-from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
-from ltx_core.model.upsampler import upsample_video
+from ltx_core.loader.registry import Registry
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
-from ltx_core.model.video_vae import decode_video as vae_decode_video
 from ltx_core.quantization import QuantizationPolicy
-from ltx_core.types import Audio, LatentState, VideoPixelShape
-from ltx_pipelines.utils import (
-    ModelLedger,
-    assert_resolution,
-    cleanup_memory,
-    combined_image_conditionings,
-    denoise_audio_video,
-    encode_prompts,
-    euler_denoising_loop,
-    get_device,
-    multi_modal_guider_factory_denoising_func,
-    simple_denoising_func,
-)
+from ltx_core.types import Audio, VideoPixelShape
 from ltx_pipelines.utils.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path
-from ltx_pipelines.utils.constants import STAGE_2_DISTILLED_SIGMA_VALUES, detect_params
+from ltx_pipelines.utils.blocks import (
+    AudioDecoder,
+    DiffusionStage,
+    ImageConditioner,
+    PromptEncoder,
+    VideoDecoder,
+    VideoUpsampler,
+)
+from ltx_pipelines.utils.constants import (
+    STAGE_2_DISTILLED_SIGMA_VALUES,
+    detect_params,
+)
+from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser, SimpleDenoiser
+from ltx_pipelines.utils.helpers import (
+    assert_resolution,
+    combined_image_conditionings,
+    get_device,
+)
 from ltx_pipelines.utils.media_io import encode_video
-from ltx_pipelines.utils.types import PipelineComponents
-
-device = get_device()
+from ltx_pipelines.utils.types import ModalitySpec
 
 
 class TI2VidTwoStagesPipeline:
@@ -55,28 +54,39 @@ class TI2VidTwoStagesPipeline:
         spatial_upsampler_path: str,
         gemma_root: str,
         loras: list[LoraPathStrengthAndSDOps],
-        device: torch.device = device,
+        device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
+        registry: Registry | None = None,
+        torch_compile: bool = False,
     ):
-        self.device = device
+        self.device = device or get_device()
         self.dtype = torch.bfloat16
-        self.stage_1_model_ledger = ModelLedger(
-            dtype=self.dtype,
-            device=device,
-            checkpoint_path=checkpoint_path,
-            gemma_root_path=gemma_root,
-            spatial_upsampler_path=spatial_upsampler_path,
-            loras=loras,
+
+        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)
+        self.upsampler = VideoUpsampler(
+            checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
+        )
+        self.video_decoder = VideoDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
+        self.audio_decoder = AudioDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
+
+        self.stage_1 = DiffusionStage(
+            checkpoint_path,
+            self.dtype,
+            self.device,
+            loras=tuple(loras),
             quantization=quantization,
+            registry=registry,
+            torch_compile=torch_compile,
         )
-
-        self.stage_2_model_ledger = self.stage_1_model_ledger.with_additional_loras(
-            loras=distilled_lora,
-        )
-
-        self.pipeline_components = PipelineComponents(
-            dtype=self.dtype,
-            device=device,
+        self.stage_2 = DiffusionStage(
+            checkpoint_path,
+            self.dtype,
+            self.device,
+            loras=(*tuple(loras), *distilled_lora),
+            quantization=quantization,
+            registry=registry,
+            torch_compile=torch_compile,
         )
 
     def __call__(  # noqa: PLR0913
@@ -94,26 +104,26 @@ class TI2VidTwoStagesPipeline:
         images: list[ImageConditioningInput],
         tiling_config: TilingConfig | None = None,
         enhance_prompt: bool = False,
+        streaming_prefetch_count: int | None = None,
+        max_batch_size: int = 1,
     ) -> tuple[Iterator[torch.Tensor], Audio]:
         assert_resolution(height=height, width=width, is_two_stage=True)
 
         generator = torch.Generator(device=self.device).manual_seed(seed)
         noiser = GaussianNoiser(generator=generator)
-        stepper = EulerDiffusionStep()
         dtype = torch.bfloat16
 
-        ctx_p, ctx_n = encode_prompts(
+        ctx_p, ctx_n = self.prompt_encoder(
             [prompt, negative_prompt],
-            self.stage_1_model_ledger,
             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,
         )
         v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
         v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
 
-        # Stage 1: encode image conditionings with the VAE encoder, then free it
-        # before loading the transformer to reduce peak VRAM.
+        # Stage 1: Generate video at half resolution with CFG guidance.
         stage_1_output_shape = VideoPixelShape(
             batch=1,
             frames=num_frames,
@@ -121,125 +131,83 @@ class TI2VidTwoStagesPipeline:
             height=height // 2,
             fps=frame_rate,
         )
-        video_encoder = self.stage_1_model_ledger.video_encoder()
-        stage_1_conditionings = combined_image_conditionings(
-            images=images,
-            height=stage_1_output_shape.height,
-            width=stage_1_output_shape.width,
-            video_encoder=video_encoder,
-            dtype=dtype,
-            device=self.device,
+        stage_1_conditionings = self.image_conditioner(
+            lambda enc: combined_image_conditionings(
+                images=images,
+                height=stage_1_output_shape.height,
+                width=stage_1_output_shape.width,
+                video_encoder=enc,
+                dtype=dtype,
+                device=self.device,
+            )
         )
-        torch.cuda.synchronize()
-        del video_encoder
-        cleanup_memory()
 
-        transformer = self.stage_1_model_ledger.transformer()
         sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
 
-        def first_stage_denoising_loop(
-            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
-        ) -> tuple[LatentState, LatentState]:
-            return euler_denoising_loop(
-                sigmas=sigmas,
-                video_state=video_state,
-                audio_state=audio_state,
-                stepper=stepper,
-                denoise_fn=multi_modal_guider_factory_denoising_func(
-                    video_guider_factory=create_multimodal_guider_factory(
-                        params=video_guider_params,
-                        negative_context=v_context_n,
-                    ),
-                    audio_guider_factory=create_multimodal_guider_factory(
-                        params=audio_guider_params,
-                        negative_context=a_context_n,
-                    ),
-                    v_context=v_context_p,
-                    a_context=a_context_p,
-                    transformer=transformer,  # noqa: F821
+        video_state, audio_state = self.stage_1(
+            denoiser=FactoryGuidedDenoiser(
+                v_context=v_context_p,
+                a_context=a_context_p,
+                video_guider_factory=create_multimodal_guider_factory(
+                    params=video_guider_params,
+                    negative_context=v_context_n,
                 ),
-            )
-
-        video_state, audio_state = denoise_audio_video(
-            output_shape=stage_1_output_shape,
-            conditionings=stage_1_conditionings,
-            noiser=noiser,
+                audio_guider_factory=create_multimodal_guider_factory(
+                    params=audio_guider_params,
+                    negative_context=a_context_n,
+                ),
+            ),
             sigmas=sigmas,
-            stepper=stepper,
-            denoising_loop_fn=first_stage_denoising_loop,
-            components=self.pipeline_components,
-            dtype=dtype,
-            device=self.device,
-        )
-
-        torch.cuda.synchronize()
-        del transformer
-        cleanup_memory()
-
-        # Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
-        video_encoder = self.stage_1_model_ledger.video_encoder()
-        upscaled_video_latent = upsample_video(
-            latent=video_state.latent[:1],
-            video_encoder=video_encoder,
-            upsampler=self.stage_2_model_ledger.spatial_upsampler(),
-        )
-
-        stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
-        stage_2_conditionings = combined_image_conditionings(
-            images=images,
-            height=stage_2_output_shape.height,
-            width=stage_2_output_shape.width,
-            video_encoder=video_encoder,
-            dtype=dtype,
-            device=self.device,
-        )
-        del video_encoder
-        torch.cuda.synchronize()
-        cleanup_memory()
-
-        transformer = self.stage_2_model_ledger.transformer()
-        distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
-
-        def second_stage_denoising_loop(
-            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
-        ) -> tuple[LatentState, LatentState]:
-            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
-                ),
-            )
-
-        video_state, audio_state = denoise_audio_video(
-            output_shape=stage_2_output_shape,
-            conditionings=stage_2_conditionings,
             noiser=noiser,
+            width=stage_1_output_shape.width,
+            height=stage_1_output_shape.height,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(context=v_context_p, conditionings=stage_1_conditionings),
+            audio=ModalitySpec(context=a_context_p),
+            streaming_prefetch_count=streaming_prefetch_count,
+            max_batch_size=max_batch_size,
+        )
+
+        # 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_conditionings = self.image_conditioner(
+            lambda enc: combined_image_conditionings(
+                images=images,
+                height=height,
+                width=width,
+                video_encoder=enc,
+                dtype=dtype,
+                device=self.device,
+            )
+        )
+
+        video_state, audio_state = self.stage_2(
+            denoiser=SimpleDenoiser(v_context=v_context_p, a_context=a_context_p),
             sigmas=distilled_sigmas,
-            stepper=stepper,
-            denoising_loop_fn=second_stage_denoising_loop,
-            components=self.pipeline_components,
-            dtype=dtype,
-            device=self.device,
-            noise_scale=distilled_sigmas[0],
-            initial_video_latent=upscaled_video_latent,
-            initial_audio_latent=audio_state.latent,
+            noiser=noiser,
+            width=width,
+            height=height,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(
+                context=v_context_p,
+                conditionings=stage_2_conditionings,
+                noise_scale=distilled_sigmas[0].item(),
+                initial_latent=upscaled_video_latent,
+            ),
+            audio=ModalitySpec(
+                context=a_context_p,
+                noise_scale=distilled_sigmas[0].item(),
+                initial_latent=audio_state.latent,
+            ),
+            streaming_prefetch_count=streaming_prefetch_count,
         )
 
-        torch.cuda.synchronize()
-        del transformer
-        cleanup_memory()
-
-        decoded_video = vae_decode_video(
-            video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config, generator
-        )
-        decoded_audio = vae_decode_audio(
-            audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder()
-        )
+        decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
+        decoded_audio = self.audio_decoder(audio_state.latent)
         return decoded_video, decoded_audio
 
 
@@ -257,6 +225,7 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
+        torch_compile=args.compile,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -287,6 +256,8 @@ def main() -> None:
         ),
         images=args.images,
         tiling_config=tiling_config,
+        streaming_prefetch_count=args.streaming_prefetch_count,
+        max_batch_size=args.max_batch_size,
     )
 
     encode_video(

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

@@ -6,34 +6,34 @@ import torch
 from ltx_core.components.diffusion_steps import Res2sDiffusionStep
 from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
 from ltx_core.components.noisers import GaussianNoiser
-from ltx_core.components.protocols import DiffusionStepProtocol
 from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.loader import LoraPathStrengthAndSDOps
-from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
-from ltx_core.model.upsampler import upsample_video
+from ltx_core.loader.registry import Registry
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
-from ltx_core.model.video_vae import decode_video as vae_decode_video
 from ltx_core.quantization import QuantizationPolicy
-from ltx_core.tools import VideoLatentShape
-from ltx_core.types import Audio, LatentState, VideoPixelShape
-from ltx_pipelines.utils import (
-    ModelLedger,
-    assert_resolution,
-    cleanup_memory,
-    combined_image_conditionings,
-    denoise_audio_video,
-    encode_prompts,
-    get_device,
-    multi_modal_guider_denoising_func,
-    res2s_audio_video_denoising_loop,
-    simple_denoising_func,
-)
+from ltx_core.types import Audio, VideoLatentShape, VideoPixelShape
 from ltx_pipelines.utils.args import ImageConditioningInput, hq_2_stage_arg_parser
-from ltx_pipelines.utils.constants import LTX_2_3_HQ_PARAMS, STAGE_2_DISTILLED_SIGMA_VALUES
+from ltx_pipelines.utils.blocks import (
+    AudioDecoder,
+    DiffusionStage,
+    ImageConditioner,
+    PromptEncoder,
+    VideoDecoder,
+    VideoUpsampler,
+)
+from ltx_pipelines.utils.constants import (
+    LTX_2_3_HQ_PARAMS,
+    STAGE_2_DISTILLED_SIGMA_VALUES,
+)
+from ltx_pipelines.utils.denoisers import GuidedDenoiser, SimpleDenoiser
+from ltx_pipelines.utils.helpers import (
+    assert_resolution,
+    combined_image_conditionings,
+    get_device,
+)
 from ltx_pipelines.utils.media_io import encode_video
-from ltx_pipelines.utils.types import PipelineComponents
-
-device = get_device()
+from ltx_pipelines.utils.samplers import res2s_audio_video_denoising_loop
+from ltx_pipelines.utils.types import ModalitySpec
 
 
 class TI2VidTwoStagesHQPipeline:
@@ -48,7 +48,7 @@ class TI2VidTwoStagesHQPipeline:
     the images parameter.
     """
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         checkpoint_path: str,
         distilled_lora: list[LoraPathStrengthAndSDOps],
@@ -57,11 +57,14 @@ class TI2VidTwoStagesHQPipeline:
         spatial_upsampler_path: str,
         gemma_root: str,
         loras: tuple[LoraPathStrengthAndSDOps, ...],
-        device: str = device,
+        device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
+        registry: Registry | None = None,
+        torch_compile: bool = False,
     ):
-        self.device = device
+        self.device = device or get_device()
         self.dtype = torch.bfloat16
+
         distilled_lora_stage_1 = LoraPathStrengthAndSDOps(
             path=distilled_lora[0].path,
             strength=distilled_lora_strength_stage_1,
@@ -72,23 +75,32 @@ class TI2VidTwoStagesHQPipeline:
             strength=distilled_lora_strength_stage_2,
             sd_ops=distilled_lora[0].sd_ops,
         )
-        self.stage_1_model_ledger = ModelLedger(
-            dtype=self.dtype,
-            device=device,
-            checkpoint_path=checkpoint_path,
-            gemma_root_path=gemma_root,
-            spatial_upsampler_path=spatial_upsampler_path,
+
+        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)
+        self.upsampler = VideoUpsampler(
+            checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
+        )
+        self.video_decoder = VideoDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
+        self.audio_decoder = AudioDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
+
+        self.stage_1 = DiffusionStage(
+            checkpoint_path,
+            self.dtype,
+            self.device,
             loras=(*loras, distilled_lora_stage_1),
             quantization=quantization,
+            registry=registry,
+            torch_compile=torch_compile,
         )
-
-        self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
+        self.stage_2 = DiffusionStage(
+            checkpoint_path,
+            self.dtype,
+            self.device,
             loras=(*loras, distilled_lora_stage_2),
-        )
-
-        self.pipeline_components = PipelineComponents(
-            dtype=self.dtype,
-            device=device,
+            quantization=quantization,
+            registry=registry,
+            torch_compile=torch_compile,
         )
 
     @torch.inference_mode()
@@ -107,6 +119,8 @@ class TI2VidTwoStagesHQPipeline:
         images: list[ImageConditioningInput],
         tiling_config: TilingConfig | None = None,
         enhance_prompt: bool = False,
+        streaming_prefetch_count: int | None = None,
+        max_batch_size: int = 1,
     ) -> tuple[Iterator[torch.Tensor], Audio]:
         assert_resolution(height=height, width=width, is_two_stage=True)
 
@@ -114,18 +128,17 @@ class TI2VidTwoStagesHQPipeline:
         noiser = GaussianNoiser(generator=generator)
         dtype = torch.bfloat16
 
-        ctx_p, ctx_n = encode_prompts(
+        ctx_p, ctx_n = self.prompt_encoder(
             [prompt, negative_prompt],
-            self.stage_1_model_ledger,
             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,
         )
         v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
         v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
 
-        # Stage 1: encode image conditionings with the VAE encoder, then free it
-        # before loading the transformer to reduce peak VRAM.
+        # Stage 1: Generate video at half resolution with CFG guidance using res2s sampler.
         stage_1_output_shape = VideoPixelShape(
             batch=1,
             frames=num_frames,
@@ -133,20 +146,16 @@ class TI2VidTwoStagesHQPipeline:
             height=height // 2,
             fps=frame_rate,
         )
-        video_encoder = self.stage_1_model_ledger.video_encoder()
-        stage_1_conditionings = combined_image_conditionings(
-            images=images,
-            height=stage_1_output_shape.height,
-            width=stage_1_output_shape.width,
-            video_encoder=video_encoder,
-            dtype=dtype,
-            device=self.device,
+        stage_1_conditionings = self.image_conditioner(
+            lambda enc: combined_image_conditionings(
+                images=images,
+                height=stage_1_output_shape.height,
+                width=stage_1_output_shape.width,
+                video_encoder=enc,
+                dtype=dtype,
+                device=self.device,
+            )
         )
-        torch.cuda.synchronize()
-        del video_encoder
-        cleanup_memory()
-
-        transformer = self.stage_1_model_ledger.transformer()
 
         empty_latent = torch.empty(VideoLatentShape.from_pixel_shape(stage_1_output_shape).to_torch_shape())
         stepper = Res2sDiffusionStep()
@@ -156,109 +165,75 @@ class TI2VidTwoStagesHQPipeline:
             .to(dtype=torch.float32, device=self.device)
         )
 
-        def first_stage_denoising_loop(
-            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
-        ) -> tuple[LatentState, LatentState]:
-            return res2s_audio_video_denoising_loop(
-                sigmas=sigmas,
-                video_state=video_state,
-                audio_state=audio_state,
-                stepper=stepper,
-                denoise_fn=multi_modal_guider_denoising_func(
-                    video_guider=MultiModalGuider(
-                        params=video_guider_params,
-                        negative_context=v_context_n,
-                    ),
-                    audio_guider=MultiModalGuider(
-                        params=audio_guider_params,
-                        negative_context=a_context_n,
-                    ),
-                    v_context=v_context_p,
-                    a_context=a_context_p,
-                    transformer=transformer,  # noqa: F821
+        video_state, audio_state = self.stage_1(
+            denoiser=GuidedDenoiser(
+                v_context=v_context_p,
+                a_context=a_context_p,
+                video_guider=MultiModalGuider(
+                    params=video_guider_params,
+                    negative_context=v_context_n,
                 ),
-            )
-
-        video_state, audio_state = denoise_audio_video(
-            output_shape=stage_1_output_shape,
-            conditionings=stage_1_conditionings,
-            noiser=noiser,
+                audio_guider=MultiModalGuider(
+                    params=audio_guider_params,
+                    negative_context=a_context_n,
+                ),
+            ),
             sigmas=sigmas,
-            stepper=stepper,
-            denoising_loop_fn=first_stage_denoising_loop,
-            components=self.pipeline_components,
-            dtype=dtype,
-            device=self.device,
-        )
-
-        torch.cuda.synchronize()
-        del transformer
-        cleanup_memory()
-
-        # Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
-        video_encoder = self.stage_1_model_ledger.video_encoder()
-        upscaled_video_latent = upsample_video(
-            latent=video_state.latent[:1],
-            video_encoder=video_encoder,
-            upsampler=self.stage_2_model_ledger.spatial_upsampler(),
-        )
-
-        stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
-        stage_2_conditionings = combined_image_conditionings(
-            images=images,
-            height=stage_2_output_shape.height,
-            width=stage_2_output_shape.width,
-            video_encoder=video_encoder,
-            dtype=dtype,
-            device=self.device,
-        )
-        torch.cuda.synchronize()
-        del video_encoder
-        cleanup_memory()
-
-        transformer = self.stage_2_model_ledger.transformer()
-        distilled_sigmas = torch.tensor(STAGE_2_DISTILLED_SIGMA_VALUES, device=self.device)
-
-        def second_stage_denoising_loop(
-            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
-        ) -> tuple[LatentState, LatentState]:
-            return res2s_audio_video_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
-                ),
-            )
-
-        video_state, audio_state = denoise_audio_video(
-            output_shape=stage_2_output_shape,
-            conditionings=stage_2_conditionings,
             noiser=noiser,
-            sigmas=distilled_sigmas,
             stepper=stepper,
-            denoising_loop_fn=second_stage_denoising_loop,
-            components=self.pipeline_components,
-            dtype=dtype,
-            device=self.device,
-            noise_scale=distilled_sigmas[0],
-            initial_video_latent=upscaled_video_latent,
-            initial_audio_latent=audio_state.latent,
+            width=stage_1_output_shape.width,
+            height=stage_1_output_shape.height,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(context=v_context_p, conditionings=stage_1_conditionings),
+            audio=ModalitySpec(context=a_context_p),
+            loop=res2s_audio_video_denoising_loop,
+            streaming_prefetch_count=streaming_prefetch_count,
+            max_batch_size=max_batch_size,
         )
 
-        torch.cuda.synchronize()
-        del transformer
-        cleanup_memory()
+        # Stage 2: Upsample and refine the video at higher resolution with distilled LoRA.
+        upscaled_video_latent = self.upsampler(video_state.latent[:1])
 
-        decoded_video = vae_decode_video(
-            video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config, generator
+        distilled_sigmas = torch.tensor(STAGE_2_DISTILLED_SIGMA_VALUES, 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: combined_image_conditionings(
+                images=images,
+                height=stage_2_output_shape.height,
+                width=stage_2_output_shape.width,
+                video_encoder=enc,
+                dtype=dtype,
+                device=self.device,
+            )
         )
-        decoded_audio = vae_decode_audio(
-            audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder()
+
+        video_state, audio_state = self.stage_2(
+            denoiser=SimpleDenoiser(v_context=v_context_p, a_context=a_context_p),
+            sigmas=distilled_sigmas,
+            noiser=noiser,
+            stepper=stepper,
+            width=width,
+            height=height,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(
+                context=v_context_p,
+                conditionings=stage_2_conditionings,
+                noise_scale=distilled_sigmas[0].item(),
+                initial_latent=upscaled_video_latent,
+            ),
+            audio=ModalitySpec(
+                context=a_context_p,
+                noise_scale=distilled_sigmas[0].item(),
+                initial_latent=audio_state.latent,
+            ),
+            loop=res2s_audio_video_denoising_loop,
+            streaming_prefetch_count=streaming_prefetch_count,
         )
+
+        decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
+        decoded_audio = self.audio_decoder(audio_state.latent)
         return decoded_video, decoded_audio
 
 
@@ -276,6 +251,7 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
+        torch_compile=args.compile,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -306,6 +282,8 @@ def main() -> None:
         ),
         images=args.images,
         tiling_config=tiling_config,
+        streaming_prefetch_count=args.streaming_prefetch_count,
+        max_batch_size=args.max_batch_size,
     )
 
     encode_video(

packages/ltx-pipelines/src/ltx_pipelines/utils/__init__.py

@@ -1,35 +1,46 @@
+from ltx_pipelines.utils.blocks import (
+    AudioConditioner,
+    AudioDecoder,
+    DiffusionStage,
+    ImageConditioner,
+    PromptEncoder,
+    VideoDecoder,
+    VideoUpsampler,
+)
+from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser, GuidedDenoiser, SimpleDenoiser
 from ltx_pipelines.utils.helpers import (
     assert_resolution,
     cleanup_memory,
     combined_image_conditionings,
-    denoise_audio_video,
-    encode_prompts,
-    generate_enhanced_prompt,
     get_device,
-    multi_modal_guider_denoising_func,
-    multi_modal_guider_factory_denoising_func,
-    simple_denoising_func,
+    image_conditionings_by_adding_guiding_latent,
 )
-from ltx_pipelines.utils.model_ledger import ModelLedger
 from ltx_pipelines.utils.samplers import (
     euler_denoising_loop,
     gradient_estimating_euler_denoising_loop,
     res2s_audio_video_denoising_loop,
 )
+from ltx_pipelines.utils.types import Denoiser, ModalitySpec
 
 __all__ = [
-    "ModelLedger",
+    "AudioConditioner",
+    "AudioDecoder",
+    "Denoiser",
+    "DiffusionStage",
+    "FactoryGuidedDenoiser",
+    "GuidedDenoiser",
+    "ImageConditioner",
+    "ModalitySpec",
+    "PromptEncoder",
+    "SimpleDenoiser",
+    "VideoDecoder",
+    "VideoUpsampler",
     "assert_resolution",
     "cleanup_memory",
     "combined_image_conditionings",
-    "denoise_audio_video",
-    "encode_prompts",
     "euler_denoising_loop",
-    "generate_enhanced_prompt",
     "get_device",
     "gradient_estimating_euler_denoising_loop",
-    "multi_modal_guider_denoising_func",
-    "multi_modal_guider_factory_denoising_func",
+    "image_conditionings_by_adding_guiding_latent",
     "res2s_audio_video_denoising_loop",
-    "simple_denoising_func",
 ]

packages/ltx-pipelines/src/ltx_pipelines/utils/args.py

@@ -173,6 +173,15 @@ def basic_arg_parser(
             required=True,
             help="Path to LTX-2 model checkpoint (.safetensors file).",
         )
+        parser.add_argument(
+            "--num-inference-steps",
+            type=int,
+            default=params.num_inference_steps,
+            help=(
+                f"Number of denoising steps in the diffusion sampling process. "
+                f"Higher values improve quality but increase generation time (default: {params.num_inference_steps})."
+            ),
+        )
     parser.add_argument(
         "--gemma-root",
         type=resolve_path,
@@ -197,6 +206,85 @@ def basic_arg_parser(
         default=params.seed,
         help=f"Random seed for reproducible generation (default: {params.seed}).",
     )
+    parser.add_argument(
+        "--lora",
+        dest="lora",
+        action=LoraAction,
+        nargs="+",  # Accept 1-2 arguments per use (path and optional strength); validation is handled in LoraAction
+        metavar=("PATH", "STRENGTH"),
+        default=[],
+        help=(
+            "LoRA (Low-Rank Adaptation) model: path to model file and optional strength "
+            f"(default strength: {DEFAULT_LORA_STRENGTH}). Can be specified multiple times. "
+            "Example: --lora path/to/lora1.safetensors 0.8 --lora path/to/lora2.safetensors"
+        ),
+    )
+
+    parser.add_argument("--enhance-prompt", action="store_true")
+
+    def _positive_int(value: str) -> int:
+        try:
+            int_value = int(value)
+            if int_value < 1:
+                raise argparse.ArgumentTypeError("must be >= 1")
+            return int_value
+        except ValueError as e:
+            raise argparse.ArgumentTypeError(f"must be an integer, got {value}") from e
+
+    # Layer streaming
+    parser.add_argument(
+        "--streaming-prefetch-count",
+        type=_positive_int,
+        default=None,
+        metavar="N",
+        help=(
+            "Enable layer streaming prefetching N layers ahead. "
+            "At most 1 + N layers reside on GPU at once. "
+            "Must be >= 1. Example: --streaming-prefetch-count 2"
+        ),
+    )
+
+    parser.add_argument(
+        "--max-batch-size",
+        type=_positive_int,
+        default=1,
+        metavar="N",
+        help=(
+            "Maximum batch size per transformer forward pass. "
+            "Guided denoisers batch up to 4 guidance passes into a single call. "
+            "Default 1 runs passes sequentially. Set to 4 to batch all passes "
+            "together, which reduces layer-streaming PCIe transfers. "
+            "Example: --max-batch-size 4"
+        ),
+    )
+
+    parser.add_argument(
+        "--quantization",
+        dest="quantization",
+        action=QuantizationAction,
+        nargs="+",
+        metavar=("POLICY", "AMAX_PATH"),
+        default=None,
+        help=(
+            f"Quantization policy: {', '.join(QUANTIZATION_POLICIES)}. "
+            "fp8-cast uses FP8 casting with upcasting during inference. "
+            "fp8-scaled-mm uses FP8 scaled matrix multiplication (optionally provide amax calibration file path). "
+            "Example: --quantization fp8-cast or --quantization fp8-scaled-mm /path/to/amax.json"
+        ),
+    )
+    parser.add_argument(
+        "--compile",
+        action="store_true",
+        help="Enable torch.compile for transformer blocks to optimize performance.",
+    )
+    return parser
+
+
+def new_video_gen_arg_parser(
+    params: PipelineParams = LTX_2_3_PARAMS,
+    distilled: bool = False,
+) -> argparse.ArgumentParser:
+    parser = basic_arg_parser(params=params, distilled=distilled)
     parser.add_argument(
         "--height",
         type=int,
@@ -222,15 +310,6 @@ def basic_arg_parser(
         default=params.frame_rate,
         help=f"Frame rate of the generated video (fps) (default: {params.frame_rate}).",
     )
-    parser.add_argument(
-        "--num-inference-steps",
-        type=int,
-        default=params.num_inference_steps,
-        help=(
-            f"Number of denoising steps in the diffusion sampling process. "
-            f"Higher values improve quality but increase generation time (default: {params.num_inference_steps})."
-        ),
-    )
     parser.add_argument(
         "--image",
         dest="images",
@@ -247,42 +326,28 @@ def basic_arg_parser(
             "--image path/to/image2.jpg 160 0.9 0"
         ),
     )
-    parser.add_argument(
-        "--lora",
-        dest="lora",
-        action=LoraAction,
-        nargs="+",  # Accept 1-2 arguments per use (path and optional strength); validation is handled in LoraAction
-        metavar=("PATH", "STRENGTH"),
-        default=[],
-        help=(
-            "LoRA (Low-Rank Adaptation) model: path to model file and optional strength "
-            f"(default strength: {DEFAULT_LORA_STRENGTH}). Can be specified multiple times. "
-            "Example: --lora path/to/lora1.safetensors 0.8 --lora path/to/lora2.safetensors"
-        ),
-    )
 
-    parser.add_argument("--enhance-prompt", action="store_true")
-    parser.add_argument(
-        "--quantization",
-        dest="quantization",
-        action=QuantizationAction,
-        nargs="+",
-        metavar=("POLICY", "AMAX_PATH"),
-        default=None,
-        help=(
-            f"Quantization policy: {', '.join(QUANTIZATION_POLICIES)}. "
-            "fp8-cast uses FP8 casting with upcasting during inference. "
-            "fp8-scaled-mm uses FP8 scaled matrix multiplication (optionally provide amax calibration file path). "
-            "Example: --quantization fp8-cast or --quantization fp8-scaled-mm /path/to/amax.json"
-        ),
-    )
+    return parser
+
+
+def video_editing_arg_parser(
+    distilled: bool = True,
+) -> argparse.ArgumentParser:
+    """Base argument parser for video-editing pipelines (retake, extension, inpainting, sticker movement).
+    Uses the same actions and conventions as basic_arg_parser but only the args needed for editing
+    (no height/width/num-frames; resolution comes from input video). Default is distilled checkpoint only.
+    """
+    parser = basic_arg_parser(distilled=distilled)
+    parser.add_argument("--video-path", type=resolve_path, required=True, help="Path to the source video.")
+    parser.add_argument("--start-time", type=float, required=True, help="Start time of the region to regenerate (s).")
+    parser.add_argument("--end-time", type=float, required=True, help="End time of the region to regenerate (s).")
     return parser
 
 
 def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
     video_guider = params.video_guider_params
     audio_guider = params.audio_guider_params
-    parser = basic_arg_parser(params=params)
+    parser = new_video_gen_arg_parser(params=params)
     parser.add_argument(
         "--negative-prompt",
         type=str,
@@ -476,7 +541,7 @@ def hq_2_stage_arg_parser(params: PipelineParams = LTX_2_3_HQ_PARAMS) -> argpars
 
 
 def default_2_stage_distilled_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
-    parser = basic_arg_parser(params=params, distilled=True)
+    parser = new_video_gen_arg_parser(params=params, distilled=True)
     parser.set_defaults(height=params.stage_2_height, width=params.stage_2_width)
     # Update help text to reflect 2-stage defaults
     for action in parser._actions:

packages/ltx-pipelines/src/ltx_pipelines/utils/blocks.py

@@ -0,0 +1,574 @@
+"""Pipeline blocks — each block owns its model lifecycle.
+Blocks build a model on each ``__call__``, use it, then free GPU memory.
+This eliminates manual ``del model; cleanup_memory()`` in pipelines and
+removes the need for :class:`ModelLedger`.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator
+from contextlib import AbstractContextManager, contextmanager
+from dataclasses import replace
+from typing import Callable, TypeVar
+
+import torch
+
+from ltx_core.batch_split import BatchSplitAdapter
+from ltx_core.components.diffusion_steps import EulerDiffusionStep
+from ltx_core.components.noisers import Noiser
+from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier
+from ltx_core.components.protocols import DiffusionStepProtocol
+from ltx_core.layer_streaming import LayerStreamingWrapper
+from ltx_core.loader import SDOps
+from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
+from ltx_core.loader.registry import DummyRegistry, Registry
+from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
+from ltx_core.model.audio_vae import (
+    AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
+    AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
+    VOCODER_COMFY_KEYS_FILTER,
+    AudioDecoderConfigurator,
+    AudioEncoderConfigurator,
+    VocoderConfigurator,
+)
+from ltx_core.model.audio_vae import (
+    decode_audio as vae_decode_audio,
+)
+from ltx_core.model.transformer import (
+    LTXV_MODEL_COMFY_RENAMING_MAP,
+    LTXModelConfigurator,
+    X0Model,
+)
+from ltx_core.model.transformer.compiling import COMPILE_TRANSFORMER, modify_sd_ops_for_compilation
+from ltx_core.model.upsampler import LatentUpsamplerConfigurator, upsample_video
+from ltx_core.model.video_vae import (
+    VAE_DECODER_COMFY_KEYS_FILTER,
+    VAE_ENCODER_COMFY_KEYS_FILTER,
+    TilingConfig,
+    VideoDecoderConfigurator,
+    VideoEncoder,
+    VideoEncoderConfigurator,
+)
+from ltx_core.quantization import QuantizationPolicy
+from ltx_core.text_encoders.gemma import (
+    EMBEDDINGS_PROCESSOR_KEY_OPS,
+    GEMMA_LLM_KEY_OPS,
+    GEMMA_MODEL_OPS,
+    EmbeddingsProcessorConfigurator,
+    GemmaTextEncoderConfigurator,
+    module_ops_from_gemma_root,
+)
+from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessorOutput
+from ltx_core.tools import AudioLatentTools, LatentTools, VideoLatentTools
+from ltx_core.types import Audio, AudioLatentShape, LatentState, VideoLatentShape, VideoPixelShape
+from ltx_core.utils import find_matching_file
+from ltx_pipelines.utils.gpu_model import gpu_model
+from ltx_pipelines.utils.helpers import (
+    cleanup_memory,
+    create_noised_state,
+    generate_enhanced_prompt,
+)
+from ltx_pipelines.utils.samplers import euler_denoising_loop
+from ltx_pipelines.utils.types import Denoiser, ModalitySpec
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+_M = TypeVar("_M", bound=torch.nn.Module)
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+
+@contextmanager
+def _streaming_model(
+    model: _M,
+    layers_attr: str,
+    target_device: torch.device,
+    prefetch_count: int,
+) -> Iterator[_M]:
+    """Wrap *model* with :class:`LayerStreamingWrapper`, yield it, then tear down."""
+    wrapped = LayerStreamingWrapper(
+        model,
+        layers_attr=layers_attr,
+        target_device=target_device,
+        prefetch_count=prefetch_count,
+    )
+    try:
+        yield wrapped  # type: ignore[misc]
+    finally:
+        wrapped.teardown()
+        wrapped.to("meta")
+        cleanup_memory()
+        # Flush the host (pinned) memory cache so that freed pinned pages are
+        # returned to the OS.  Without this, sequential streaming models
+        # (e.g. text encoder then transformer) exhaust host memory because the
+        # CachingHostAllocator keeps freed blocks cached indefinitely.
+        torch.cuda.synchronize(device=target_device)
+        try:
+            if hasattr(torch._C, "_host_emptyCache"):
+                torch._C._host_emptyCache()
+        except Exception:
+            logger.warning("Host empty cache cleanup failed; ignoring.", exc_info=True)
+
+
+def _build_state(
+    spec: ModalitySpec,
+    tools: LatentTools,
+    noiser: Noiser,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> LatentState:
+    """Create a noised latent state from a modality spec and tools."""
+    state = create_noised_state(
+        tools=tools,
+        conditionings=spec.conditionings,
+        noiser=noiser,
+        dtype=dtype,
+        device=device,
+        noise_scale=spec.noise_scale,
+        initial_latent=spec.initial_latent,
+    )
+    if spec.frozen:
+        state = replace(state, denoise_mask=torch.zeros_like(state.denoise_mask))
+    return state
+
+
+def _cleanup_iter(it: Iterator[torch.Tensor], model: torch.nn.Module) -> Iterator[torch.Tensor]:
+    """Wrap an iterator to clean up *model* memory once it is exhausted or abandoned."""
+    with gpu_model(model):
+        yield from it
+
+
+# ---------------------------------------------------------------------------
+# DiffusionStage
+# ---------------------------------------------------------------------------
+
+
+class DiffusionStage:
+    """Owns transformer lifecycle. Builds on each call, frees on exit.
+    Replaces the manual ``model_ledger.transformer()`` / ``del transformer``
+    pattern in every pipeline.
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        dtype: torch.dtype,
+        device: torch.device,
+        loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
+        quantization: QuantizationPolicy | None = None,
+        registry: Registry | None = None,
+        torch_compile: bool = False,
+    ) -> None:
+        self._dtype = dtype
+        self._device = device
+        self._quantization = quantization
+        self._torch_compile = torch_compile
+        self._transformer_builder = Builder(
+            model_path=checkpoint_path,
+            model_class_configurator=LTXModelConfigurator,
+            model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
+            loras=tuple(loras),
+            registry=registry or DummyRegistry(),
+        )
+
+    def _build_transformer(self, *, device: torch.device | None = None, **kwargs: object) -> X0Model:
+        target = device or self._device
+        sd_ops = self._transformer_builder.model_sd_ops
+        module_ops = self._transformer_builder.module_ops
+        loras = self._transformer_builder.loras
+        if self._torch_compile:
+            module_ops = (*module_ops, COMPILE_TRANSFORMER)
+            number_of_layers = self._transformer_builder.model_config()["transformer"]["num_layers"]
+            sd_ops = modify_sd_ops_for_compilation(sd_ops, number_of_layers)
+            loras = tuple(
+                LoraPathStrengthAndSDOps(
+                    lora.path,
+                    lora.strength,
+                    modify_sd_ops_for_compilation(
+                        lora.sd_ops if lora.sd_ops is not None else SDOps(name="identity"), number_of_layers
+                    ),
+                )
+                for lora in loras
+            )
+        if self._quantization is not None:
+            module_ops = (*module_ops, *self._quantization.module_ops)
+            sd_ops = SDOps(
+                name=f"sd_ops_chain_{sd_ops.name}+{self._quantization.sd_ops.name}",
+                mapping=(*sd_ops.mapping, *self._quantization.sd_ops.mapping),
+            )
+
+        builder = self._transformer_builder.with_module_ops(module_ops).with_sd_ops(sd_ops).with_loras(loras)
+        return X0Model(builder.build(device=target, **kwargs)).to(target).eval()
+
+    def _transformer_ctx(
+        self,
+        streaming_prefetch_count: int | None,
+        **kwargs: object,
+    ) -> AbstractContextManager:
+        if streaming_prefetch_count is not None:
+            return _streaming_model(
+                self._build_transformer(device=torch.device("cpu"), **kwargs),
+                layers_attr="velocity_model.transformer_blocks",
+                target_device=self._device,
+                prefetch_count=streaming_prefetch_count,
+            )
+        return gpu_model(self._build_transformer(**kwargs))
+
+    def __call__(  # noqa: PLR0913
+        self,
+        denoiser: Denoiser,
+        sigmas: torch.Tensor,
+        noiser: Noiser,
+        width: int,
+        height: int,
+        frames: int,
+        fps: float,
+        video: ModalitySpec | None = None,
+        audio: ModalitySpec | None = None,
+        stepper: DiffusionStepProtocol | None = None,
+        loop: Callable[..., tuple[LatentState | None, LatentState | None]] | None = None,
+        streaming_prefetch_count: int | None = None,
+        max_batch_size: int = 1,
+    ) -> tuple[LatentState | None, LatentState | None]:
+        """Build transformer → run denoising loop → free transformer.
+        Args:
+            width: Output width in pixels.
+            height: Output height in pixels.
+            frames: Number of output frames.
+            fps: Frame rate.
+            loop: Denoising loop function. Must accept
+                ``(sigmas, video_state, audio_state, stepper, transformer, denoiser)``
+                as the first six positional arguments. When ``None``, resolves to
+                :func:`euler_denoising_loop` at call time.
+            streaming_prefetch_count: When set, build the transformer on CPU and
+                wrap with :class:`LayerStreamingWrapper` for memory-efficient
+                inference, prefetching this many layers ahead.
+            max_batch_size: Maximum batch size per transformer forward pass.
+                Guided denoisers make up to 4 transformer calls per step.
+                When set to a value > 1, the transformer batches multiple
+                calls together, reducing layer-streaming PCIe transfers.
+                Default ``1`` preserves sequential behavior.
+        Returns ``(video_state | None, audio_state | None)`` with cleared
+        conditionings and unpatchified latents for present modalities.
+        """
+        if video is None and audio is None:
+            raise ValueError("At least one of `video` or `audio` must be provided")
+
+        if loop is None:
+            loop = euler_denoising_loop
+
+        if stepper is None:
+            stepper = EulerDiffusionStep()
+
+        pixel_shape = VideoPixelShape(batch=1, frames=frames, height=height, width=width, fps=fps)
+
+        video_state: LatentState | None = None
+        video_tools: LatentTools | None = None
+        if video is not None:
+            v_shape = VideoLatentShape.from_pixel_shape(pixel_shape)
+            video_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape, fps)
+            video_state = _build_state(video, video_tools, noiser, self._dtype, self._device)
+
+        audio_state: LatentState | None = None
+        audio_tools: LatentTools | None = None
+        if audio is not None:
+            a_shape = AudioLatentShape.from_video_pixel_shape(pixel_shape)
+            audio_tools = AudioLatentTools(AudioPatchifier(patch_size=1), a_shape)
+            audio_state = _build_state(audio, audio_tools, noiser, self._dtype, self._device)
+
+        with self._transformer_ctx(streaming_prefetch_count, video_tools=video_tools) as base_transformer:
+            transformer = BatchSplitAdapter(base_transformer, max_batch_size=max_batch_size)
+            video_state, audio_state = loop(
+                sigmas=sigmas,
+                video_state=video_state,
+                audio_state=audio_state,
+                stepper=stepper,
+                transformer=transformer,
+                denoiser=denoiser,
+            )
+
+        # Post-process: clear conditionings and unpatchify
+        if video_state is not None and video_tools is not None:
+            video_state = video_tools.clear_conditioning(video_state)
+            video_state = video_tools.unpatchify(video_state)
+
+        if audio_state is not None and audio_tools is not None:
+            audio_state = audio_tools.clear_conditioning(audio_state)
+            audio_state = audio_tools.unpatchify(audio_state)
+
+        return video_state, audio_state
+
+
+# ---------------------------------------------------------------------------
+# PromptEncoder
+# ---------------------------------------------------------------------------
+
+
+class PromptEncoder:
+    """Owns text encoder + embeddings processor lifecycle.
+    Loads Gemma, encodes prompts, frees Gemma, then loads the embeddings
+    processor to produce final outputs.
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        gemma_root: str,
+        dtype: torch.dtype,
+        device: torch.device,
+        registry: Registry | None = None,
+    ) -> None:
+        self._dtype = dtype
+        self._device = device
+
+        module_ops = module_ops_from_gemma_root(gemma_root)
+        model_folder = find_matching_file(gemma_root, "model*.safetensors").parent
+        weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")]
+
+        self._text_encoder_builder = Builder(
+            model_path=tuple(weight_paths),
+            model_class_configurator=GemmaTextEncoderConfigurator,
+            model_sd_ops=GEMMA_LLM_KEY_OPS,
+            module_ops=(GEMMA_MODEL_OPS, *module_ops),
+            registry=registry or DummyRegistry(),
+        )
+        self._embeddings_processor_builder = Builder(
+            model_path=checkpoint_path,
+            model_class_configurator=EmbeddingsProcessorConfigurator,
+            model_sd_ops=EMBEDDINGS_PROCESSOR_KEY_OPS,
+            registry=registry or DummyRegistry(),
+        )
+
+    def _text_encoder_ctx(
+        self,
+        streaming_prefetch_count: int | None,
+    ) -> AbstractContextManager:
+        if streaming_prefetch_count is not None:
+            return _streaming_model(
+                self._text_encoder_builder.build(device=torch.device("cpu"), dtype=self._dtype).eval(),
+                layers_attr="model.model.language_model.layers",
+                target_device=self._device,
+                prefetch_count=streaming_prefetch_count,
+            )
+        return gpu_model(self._text_encoder_builder.build(device=self._device, dtype=self._dtype).eval())
+
+    def __call__(
+        self,
+        prompts: list[str],
+        *,
+        enhance_first_prompt: bool = False,
+        enhance_prompt_image: str | None = None,
+        enhance_prompt_seed: int = 42,
+        streaming_prefetch_count: int | None = None,
+    ) -> list[EmbeddingsProcessorOutput]:
+        """Encode *prompts* through Gemma → embeddings processor, freeing each model after use."""
+        with self._text_encoder_ctx(streaming_prefetch_count) as text_encoder:
+            if enhance_first_prompt:
+                prompts = list(prompts)
+                prompts[0] = generate_enhanced_prompt(
+                    text_encoder, prompts[0], enhance_prompt_image, seed=enhance_prompt_seed
+                )
+            raw_outputs = [text_encoder.encode(p) for p in prompts]
+
+        with gpu_model(
+            self._embeddings_processor_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
+        ) as embeddings_processor:
+            return [embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs]
+
+
+# ---------------------------------------------------------------------------
+# ImageConditioner
+# ---------------------------------------------------------------------------
+
+
+class ImageConditioner:
+    """Owns video encoder lifecycle.
+    Builds the encoder, passes it to the user-supplied callable, then frees it.
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        dtype: torch.dtype,
+        device: torch.device,
+        registry: Registry | None = None,
+    ) -> None:
+        self._dtype = dtype
+        self._device = device
+        self._encoder_builder = Builder(
+            model_path=checkpoint_path,
+            model_class_configurator=VideoEncoderConfigurator,
+            model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
+            registry=registry or DummyRegistry(),
+        )
+
+    def _build_encoder(self) -> VideoEncoder:
+        return self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
+
+    def __call__(self, fn: Callable[[VideoEncoder], T]) -> T:
+        """Build video encoder → call *fn(encoder)* → free encoder."""
+        with gpu_model(self._build_encoder()) as encoder:
+            return fn(encoder)
+
+
+# ---------------------------------------------------------------------------
+# VideoUpsampler
+# ---------------------------------------------------------------------------
+
+
+class VideoUpsampler:
+    """Owns video encoder + spatial upsampler lifecycle."""
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        upsampler_path: str,
+        dtype: torch.dtype,
+        device: torch.device,
+        registry: Registry | None = None,
+    ) -> None:
+        self._dtype = dtype
+        self._device = device
+        self._encoder_builder = Builder(
+            model_path=checkpoint_path,
+            model_class_configurator=VideoEncoderConfigurator,
+            model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
+            registry=registry or DummyRegistry(),
+        )
+        self._upsampler_builder = Builder(
+            model_path=upsampler_path,
+            model_class_configurator=LatentUpsamplerConfigurator,
+            registry=registry or DummyRegistry(),
+        )
+
+    def __call__(self, latent: torch.Tensor) -> torch.Tensor:
+        """Upsample *latent* using video encoder + spatial upsampler, then free both."""
+        with (
+            gpu_model(
+                self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
+            ) as encoder,
+            gpu_model(
+                self._upsampler_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
+            ) as upsampler,
+        ):
+            return upsample_video(latent=latent, video_encoder=encoder, upsampler=upsampler)
+
+
+# ---------------------------------------------------------------------------
+# VideoDecoder
+# ---------------------------------------------------------------------------
+
+
+class VideoDecoder:
+    """Owns video decoder lifecycle.
+    Returns an iterator that cleans up the decoder after all chunks are consumed.
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        dtype: torch.dtype,
+        device: torch.device,
+        registry: Registry | None = None,
+    ) -> None:
+        self._dtype = dtype
+        self._device = device
+        self._decoder_builder = Builder(
+            model_path=checkpoint_path,
+            model_class_configurator=VideoDecoderConfigurator,
+            model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
+            registry=registry or DummyRegistry(),
+        )
+
+    def __call__(
+        self,
+        latent: torch.Tensor,
+        tiling_config: TilingConfig | None = None,
+        generator: torch.Generator | None = None,
+    ) -> Iterator[torch.Tensor]:
+        """Decode *latent* to pixel-space video chunks. Decoder freed after exhaustion."""
+        decoder = self._decoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
+        return _cleanup_iter(decoder.decode_video(latent, tiling_config, generator), decoder)
+
+
+# ---------------------------------------------------------------------------
+# AudioDecoder
+# ---------------------------------------------------------------------------
+
+
+class AudioDecoder:
+    """Owns audio decoder + vocoder lifecycle."""
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        dtype: torch.dtype,
+        device: torch.device,
+        registry: Registry | None = None,
+    ) -> None:
+        self._dtype = dtype
+        self._device = device
+        self._decoder_builder = Builder(
+            model_path=checkpoint_path,
+            model_class_configurator=AudioDecoderConfigurator,
+            model_sd_ops=AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
+            registry=registry or DummyRegistry(),
+        )
+        self._vocoder_builder = Builder(
+            model_path=checkpoint_path,
+            model_class_configurator=VocoderConfigurator,
+            model_sd_ops=VOCODER_COMFY_KEYS_FILTER,
+            registry=registry or DummyRegistry(),
+        )
+
+    def __call__(self, latent: torch.Tensor) -> Audio:
+        """Decode audio *latent* through VAE decoder + vocoder, then free both."""
+        with (
+            gpu_model(
+                self._decoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
+            ) as decoder,
+            gpu_model(
+                self._vocoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
+            ) as vocoder,
+        ):
+            return vae_decode_audio(latent, decoder, vocoder)
+
+
+# ---------------------------------------------------------------------------
+# AudioEncoder
+# ---------------------------------------------------------------------------
+
+
+class AudioConditioner:
+    """Owns audio encoder lifecycle.
+    Builds the encoder, passes it to the user-supplied callable, then frees it.
+    Mirrors :class:`ImageConditioner` for the audio modality.
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        dtype: torch.dtype,
+        device: torch.device,
+        registry: Registry | None = None,
+    ) -> None:
+        self._dtype = dtype
+        self._device = device
+        self._encoder_builder = Builder(
+            model_path=checkpoint_path,
+            model_class_configurator=AudioEncoderConfigurator,
+            model_sd_ops=AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
+            registry=registry or DummyRegistry(),
+        )
+
+    def __call__(self, fn: Callable[[torch.nn.Module], T]) -> T:
+        """Build audio encoder → call *fn(encoder)* → free encoder."""
+        with gpu_model(
+            self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
+        ) as encoder:
+            return fn(encoder)

packages/ltx-pipelines/src/ltx_pipelines/utils/denoisers.py

@@ -0,0 +1,305 @@
+"""Flat denoiser classes — transformer received at call time, not stored.
+Three implementations of the :class:`~ltx_pipelines.utils.types.Denoiser` protocol:
+* :class:`SimpleDenoiser` — single transformer call, no guidance.
+* :class:`GuidedDenoiser` — static guiders, handles CFG + STG + isolated modality.
+* :class:`FactoryGuidedDenoiser` — resolves guiders per-step from sigma.
+``GuidedDenoiser`` and ``FactoryGuidedDenoiser`` share the core multi-pass
+logic via the module-level :func:`_guided_denoise` function, which batches
+all guidance passes into a single transformer call.
+"""
+
+import torch
+
+from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderFactory, MultiModalGuiderParams
+from ltx_core.guidance.perturbations import (
+    BatchedPerturbationConfig,
+    Perturbation,
+    PerturbationConfig,
+    PerturbationType,
+)
+from ltx_core.model.transformer import X0Model
+from ltx_core.types import LatentState
+from ltx_pipelines.utils.helpers import modality_from_latent_state
+
+_POSITIVE_ONLY_GUIDER = MultiModalGuider(
+    params=MultiModalGuiderParams(cfg_scale=1.0, stg_scale=0.0, modality_scale=1.0),
+)
+"""Guider that only runs the conditioned pass and returns cond unchanged."""
+
+
+def _ensure_guider(guider: MultiModalGuider | None) -> MultiModalGuider:
+    """Return the guider as-is, or a positive-only guider for absent modalities."""
+    return guider if guider is not None else _POSITIVE_ONLY_GUIDER
+
+
+def _repeat_state(state: LatentState, n: int) -> LatentState:
+    """Repeat a ``LatentState`` *n* times along the batch dimension.
+    ``(B, ...) → (n*B, ...)`` by tiling the whole tensor n times, so the
+    ordering is ``[item0, item1, ..., item0, item1, ...]`` — matching
+    ``torch.cat`` of n per-pass contexts.
+    """
+
+    def _repeat(t: torch.Tensor) -> torch.Tensor:
+        repeats = [1] * t.dim()
+        repeats[0] = n
+        return t.repeat(repeats)
+
+    return LatentState(
+        latent=_repeat(state.latent),
+        denoise_mask=_repeat(state.denoise_mask),
+        positions=_repeat(state.positions),
+        clean_latent=_repeat(state.clean_latent),
+        attention_mask=_repeat(state.attention_mask) if state.attention_mask is not None else None,
+    )
+
+
+def _guided_denoise(  # noqa: PLR0913
+    transformer: X0Model,
+    video_state: LatentState | None,
+    audio_state: LatentState | None,
+    sigma: torch.Tensor,
+    video_guider: MultiModalGuider,
+    audio_guider: MultiModalGuider,
+    v_context: torch.Tensor | None,
+    a_context: torch.Tensor | None,
+    *,
+    last_denoised_video: torch.Tensor | None,
+    last_denoised_audio: torch.Tensor | None,
+    step_index: int,
+) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    """Core guided denoising — batches all guidance passes into one transformer call.
+    Collects per-pass contexts first, then builds a single batched Modality
+    per present modality via :func:`modality_from_latent_state`.  When wrapped
+    with :class:`~ltx_core.batch_split.BatchSplitAdapter`, the transformer may
+    split this batch into sequential chunks internally.
+    Guiders must not be ``None``. For absent modalities, callers should pass
+    :data:`_POSITIVE_ONLY_GUIDER` (via :func:`_ensure_guider`) so that only
+    the conditioned pass runs and ``calculate()`` returns cond unchanged.
+    """
+    v_skip = video_guider.should_skip_step(step_index)
+    a_skip = audio_guider.should_skip_step(step_index)
+
+    if v_skip and a_skip:
+        return last_denoised_video, last_denoised_audio
+
+    if video_state is not None and v_context is None:
+        raise ValueError("v_context is required when video_state is provided")
+    if audio_state is not None and a_context is None:
+        raise ValueError("a_context is required when audio_state is provided")
+    # Define passes: (name, video_context, audio_context, perturbation_config).
+    # Context is None for absent modalities — filtered out during collection.
+    _pass = tuple[str, torch.Tensor | None, torch.Tensor | None, PerturbationConfig]
+    passes: list[_pass] = [("cond", v_context, a_context, PerturbationConfig.empty())]
+
+    if video_guider.do_unconditional_generation() or audio_guider.do_unconditional_generation():
+        if video_guider.do_unconditional_generation() and video_guider.negative_context is None:
+            raise ValueError("Negative context is required for unconditioned denoising")
+        if audio_guider.do_unconditional_generation() and audio_guider.negative_context is None:
+            raise ValueError("Negative context is required for unconditioned denoising")
+        v_neg = video_guider.negative_context if video_guider.negative_context is not None else v_context
+        a_neg = audio_guider.negative_context if audio_guider.negative_context is not None else a_context
+        passes.append(("uncond", v_neg, a_neg, PerturbationConfig.empty()))
+
+    stg_perturbations: list[Perturbation] = []
+    if video_guider.do_perturbed_generation():
+        stg_perturbations.append(
+            Perturbation(type=PerturbationType.SKIP_VIDEO_SELF_ATTN, blocks=video_guider.params.stg_blocks)
+        )
+    if audio_guider.do_perturbed_generation():
+        stg_perturbations.append(
+            Perturbation(type=PerturbationType.SKIP_AUDIO_SELF_ATTN, blocks=audio_guider.params.stg_blocks)
+        )
+    if stg_perturbations:
+        passes.append(("ptb", v_context, a_context, PerturbationConfig(stg_perturbations)))
+
+    if video_guider.do_isolated_modality_generation() or audio_guider.do_isolated_modality_generation():
+        passes.append(
+            (
+                "mod",
+                v_context,
+                a_context,
+                PerturbationConfig(
+                    [
+                        Perturbation(type=PerturbationType.SKIP_A2V_CROSS_ATTN, blocks=None),
+                        Perturbation(type=PerturbationType.SKIP_V2A_CROSS_ATTN, blocks=None),
+                    ]
+                ),
+            )
+        )
+
+    # Collect contexts, repeat states, and build batched modalities.
+    pass_names = [name for name, _, _, _ in passes]
+    ptb_configs = [ptb for _, _, _, ptb in passes]
+    n = len(passes)
+
+    def _batched_sigma(state: LatentState) -> torch.Tensor:
+        """Expand scalar sigma to (n * B,) matching the repeated state."""
+        return sigma.expand(state.latent.shape[0] * n)
+
+    batched_video = None
+    if video_state is not None:
+        v_context = torch.cat([vc for _, vc, _, _ in passes], dim=0)
+        batched_video = modality_from_latent_state(
+            _repeat_state(video_state, n),
+            v_context,
+            _batched_sigma(video_state),
+            enabled=not v_skip,
+        )
+
+    batched_audio = None
+    if audio_state is not None:
+        a_context = torch.cat([ac for _, _, ac, _ in passes], dim=0)
+        batched_audio = modality_from_latent_state(
+            _repeat_state(audio_state, n),
+            a_context,
+            _batched_sigma(audio_state),
+            enabled=not a_skip,
+        )
+
+    all_v, all_a = transformer(
+        video=batched_video, audio=batched_audio, perturbations=BatchedPerturbationConfig(ptb_configs)
+    )
+
+    # Split results back and combine via guiders.
+    splits_v = list(all_v.chunk(n)) if all_v is not None else [0.0] * n
+    splits_a = list(all_a.chunk(n)) if all_a is not None else [0.0] * n
+    r = dict(zip(pass_names, zip(splits_v, splits_a, strict=True), strict=True))
+
+    cond_v, cond_a = r["cond"]
+    uncond_v, uncond_a = r.get("uncond", (0.0, 0.0))
+    ptb_v, ptb_a = r.get("ptb", (0.0, 0.0))
+    mod_v, mod_a = r.get("mod", (0.0, 0.0))
+
+    denoised_video = last_denoised_video if v_skip else video_guider.calculate(cond_v, uncond_v, ptb_v, mod_v)
+    denoised_audio = last_denoised_audio if a_skip else audio_guider.calculate(cond_a, uncond_a, ptb_a, mod_a)
+    return denoised_video, denoised_audio
+
+
+class SimpleDenoiser:
+    """Single transformer call, no guidance.
+    Passes ``None`` Modality for absent modalities.
+    """
+
+    def __init__(
+        self,
+        v_context: torch.Tensor | None,
+        a_context: torch.Tensor | None,
+    ) -> None:
+        self.v_context = v_context
+        self.a_context = a_context
+
+    def __call__(
+        self,
+        transformer: X0Model,
+        video_state: LatentState | None,
+        audio_state: LatentState | None,
+        sigmas: torch.Tensor,
+        step_index: int,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        sigma = sigmas[step_index]
+        pos_video = modality_from_latent_state(video_state, self.v_context, sigma) if video_state is not None else None
+        pos_audio = modality_from_latent_state(audio_state, self.a_context, sigma) if audio_state is not None else None
+        return transformer(video=pos_video, audio=pos_audio, perturbations=None)
+
+
+class GuidedDenoiser:
+    """Static guiders — handles CFG + STG + isolated modality.
+    Context/guider can be ``None`` for absent modalities (a positive-only
+    guider is substituted at call time).
+    """
+
+    def __init__(
+        self,
+        v_context: torch.Tensor | None,
+        a_context: torch.Tensor | None,
+        video_guider: MultiModalGuider | None = None,
+        audio_guider: MultiModalGuider | None = None,
+    ) -> None:
+        self.v_context = v_context
+        self.a_context = a_context
+        self.video_guider = video_guider
+        self.audio_guider = audio_guider
+        self._last_denoised_video: torch.Tensor | None = None
+        self._last_denoised_audio: torch.Tensor | None = None
+
+    def __call__(
+        self,
+        transformer: X0Model,
+        video_state: LatentState | None,
+        audio_state: LatentState | None,
+        sigmas: torch.Tensor,
+        step_index: int,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        denoised_video, denoised_audio = _guided_denoise(
+            transformer=transformer,
+            video_state=video_state,
+            audio_state=audio_state,
+            sigma=sigmas[step_index],
+            video_guider=_ensure_guider(self.video_guider),
+            audio_guider=_ensure_guider(self.audio_guider),
+            v_context=self.v_context,
+            a_context=self.a_context,
+            last_denoised_video=self._last_denoised_video,
+            last_denoised_audio=self._last_denoised_audio,
+            step_index=step_index,
+        )
+        self._last_denoised_video = denoised_video
+        self._last_denoised_audio = denoised_audio
+        return denoised_video, denoised_audio
+
+
+class FactoryGuidedDenoiser:
+    """Resolves guiders per-step from sigma, then delegates to shared guided logic."""
+
+    def __init__(
+        self,
+        v_context: torch.Tensor | None,
+        a_context: torch.Tensor | None,
+        video_guider_factory: MultiModalGuiderFactory | None = None,
+        audio_guider_factory: MultiModalGuiderFactory | None = None,
+    ) -> None:
+        self.v_context = v_context
+        self.a_context = a_context
+        self.video_guider_factory = video_guider_factory
+        self.audio_guider_factory = audio_guider_factory
+        self._last_denoised_video: torch.Tensor | None = None
+        self._last_denoised_audio: torch.Tensor | None = None
+        self._sigma_vals_cached: list[float] | None = None
+
+    def __call__(
+        self,
+        transformer: X0Model,
+        video_state: LatentState | None,
+        audio_state: LatentState | None,
+        sigmas: torch.Tensor,
+        step_index: int,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        if self._sigma_vals_cached is None:
+            self._sigma_vals_cached = sigmas.detach().cpu().tolist()
+        sigma_val = self._sigma_vals_cached[step_index]
+
+        video_guider = _ensure_guider(
+            self.video_guider_factory.build_from_sigma(sigma_val) if self.video_guider_factory else None
+        )
+        audio_guider = _ensure_guider(
+            (self.audio_guider_factory or self.video_guider_factory).build_from_sigma(sigma_val)
+            if self.video_guider_factory or self.audio_guider_factory
+            else None
+        )
+
+        denoised_video, denoised_audio = _guided_denoise(
+            transformer=transformer,
+            video_state=video_state,
+            audio_state=audio_state,
+            sigma=sigmas[step_index],
+            video_guider=video_guider,
+            audio_guider=audio_guider,
+            v_context=self.v_context,
+            a_context=self.a_context,
+            last_denoised_video=self._last_denoised_video,
+            last_denoised_audio=self._last_denoised_audio,
+            step_index=step_index,
+        )
+        self._last_denoised_video = denoised_video
+        self._last_denoised_audio = denoised_audio
+        return denoised_video, denoised_audio

packages/ltx-pipelines/src/ltx_pipelines/utils/gpu_model.py

@@ -0,0 +1,30 @@
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import TypeVar
+
+import torch
+
+from ltx_pipelines.utils.helpers import cleanup_memory
+
+_M = TypeVar("_M", bound=torch.nn.Module)
+
+
+@contextmanager
+def gpu_model(model: _M) -> Iterator[_M]:
+    """Context manager that yields a model and releases its memory on exit.
+    Moves all parameters and buffers to ``meta`` device on exit, which
+    immediately releases the underlying storage on **both** GPU and CPU,
+    then runs ``cleanup_memory()`` to reclaim fragmented CUDA memory.
+    Usage::
+        with gpu_model(build_encoder()) as encoder:
+            ...  # use encoder — typed as the concrete class
+        # GPU + CPU memory freed automatically
+    """
+    try:
+        yield model
+    finally:
+        torch.cuda.synchronize()
+        # .to("meta") releases storage for all parameters/buffers regardless
+        # of their original device (CUDA or CPU).
+        model.to("meta")
+        cleanup_memory()

packages/ltx-pipelines/src/ltx_pipelines/utils/helpers.py

@@ -1,41 +1,35 @@
 import gc
 import logging
-from dataclasses import replace
 
 import torch
 
-from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderFactory
 from ltx_core.components.noisers import Noiser
-from ltx_core.components.protocols import DiffusionStepProtocol, GuiderProtocol
 from ltx_core.conditioning import (
     ConditioningItem,
     VideoConditionByKeyframeIndex,
     VideoConditionByLatentIndex,
 )
-from ltx_core.guidance.perturbations import (
-    BatchedPerturbationConfig,
-    Perturbation,
-    PerturbationConfig,
-    PerturbationType,
-)
-from ltx_core.model.transformer import Modality, X0Model
-from ltx_core.model.video_vae import VideoEncoder
+from ltx_core.model.audio_vae import encode_audio
+from ltx_core.model.transformer import Modality
+from ltx_core.model.video_vae import TilingConfig, VideoEncoder
 from ltx_core.text_encoders.gemma import GemmaTextEncoder
-from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessorOutput
-from ltx_core.tools import AudioLatentTools, LatentTools, VideoLatentTools
+from ltx_core.tools import LatentTools
 from ltx_core.types import AudioLatentShape, LatentState, VideoLatentShape, VideoPixelShape
 from ltx_pipelines.utils.args import ImageConditioningInput
-from ltx_pipelines.utils.media_io import decode_image, load_image_conditioning, resize_aspect_ratio_preserving
-from ltx_pipelines.utils.types import (
-    DenoisingFunc,
-    DenoisingLoopFunc,
-    PipelineComponents,
+from ltx_pipelines.utils.media_io import (
+    decode_audio_from_file,
+    decode_image,
+    decode_video_from_file,
+    get_videostream_fps,
+    load_image_and_preprocess,
+    resize_aspect_ratio_preserving,
+    video_preprocess,
 )
 
 
 def get_device() -> torch.device:
     if torch.cuda.is_available():
-        return torch.device("cuda")
+        return torch.device("cuda", torch.cuda.current_device())
     return torch.device("cpu")
 
 
@@ -45,45 +39,89 @@ def cleanup_memory() -> None:
     torch.cuda.synchronize()
 
 
-def encode_prompts(
-    prompts: list[str],
-    model_ledger: object,
-    *,
-    enhance_prompt_image: str | None = None,
-    enhance_prompt_seed: int = 42,
-    enhance_first_prompt: bool = False,
-) -> list[EmbeddingsProcessorOutput]:
-    """Encode prompts through Gemma → embeddings processor, freeing each after use.
-    Loads the text encoder from *model_ledger*, optionally enhances the first
-    prompt, encodes all *prompts*, frees the text encoder, then loads the
-    embeddings processor to produce the final outputs.  Because the text encoder
-    is loaded and freed entirely within this function, there are no lingering
-    references that could prevent GPU memory reclamation.
-    Args:
-        prompts: Text prompts to encode.
-        model_ledger: ModelLedger instance (used to load text encoder and embeddings processor).
-        enhance_prompt_image: Optional image path for prompt enhancement.
-        enhance_prompt_seed: Seed for prompt enhancement (default 42).
-        enhance_first_prompt: If True, enhance ``prompts[0]`` before encoding.
-    Returns:
-        List of EmbeddingsProcessorOutput, one per prompt.
-    """
-    text_encoder = model_ledger.text_encoder()
-    if enhance_first_prompt:
-        prompts = list(prompts)
-        prompts[0] = generate_enhanced_prompt(text_encoder, prompts[0], enhance_prompt_image, seed=enhance_prompt_seed)
-    raw_outputs = [text_encoder.encode(p) for p in prompts]
-    torch.cuda.synchronize()
-    del text_encoder
-    cleanup_memory()
+def _conform_latent_length(latent: torch.Tensor, expected_frames_count: int) -> torch.Tensor:
+    actual_frames = latent.shape[2]
+    if actual_frames > expected_frames_count:
+        latent = latent[:, :, :expected_frames_count]
+    elif actual_frames < expected_frames_count:
+        shape_as_list = list(latent.shape)
+        shape_as_list[2] = expected_frames_count - actual_frames
+        pad = torch.zeros(
+            shape_as_list,
+            device=latent.device,
+            dtype=latent.dtype,
+        )
+        latent = torch.cat([latent, pad], dim=2)
+    return latent
 
-    embeddings_processor = model_ledger.gemma_embeddings_processor()
-    results: list[EmbeddingsProcessorOutput] = [
-        embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs
-    ]
-    del embeddings_processor
-    cleanup_memory()
-    return results
+
+def video_latent_from_file(
+    video_encoder: VideoEncoder,
+    file_path: str,
+    output_shape: VideoPixelShape,
+    device: torch.device,
+    dtype: torch.dtype,
+    start_time: float = 0.0,
+    max_duration: float | None = None,
+    tiling_config: TilingConfig | None = None,
+) -> torch.Tensor | None:
+    """Load video from a file, and construct the video latent conforming to video output shape.
+    Args:
+        video_encoder: Model used to encode pixel frames to latent space.
+        file_path: Path to the video file.
+        output_shape: Target pixel shape (height, width, frames, fps) for the conditioning.
+        device: Device to run the encoder and hold tensors on.
+        dtype: Dtype for the output latents.
+        start_time: Start time in seconds to begin reading the video (default 0.0).
+        max_duration: Maximum duration in seconds. If None, uses output_shape.frames at
+            output_shape.fps (default None).
+        tiling_config: Tiling configuration for the encoder. Defaults to TilingConfig.default().
+    Returns:
+        Encoded video latents of shape (1, C, T, H, W) with T = required_latent_frames, or
+        None (currently this function always returns a tensor).
+    """
+    fps = get_videostream_fps(file_path)
+    if fps != output_shape.fps:
+        raise ValueError(f"Input video FPS {fps} does not match output FPS {output_shape.fps}, not supported")
+    max_duration = max_duration or output_shape.frames / fps
+    frame_gen = decode_video_from_file(path=file_path, device=device, start_time=start_time, max_duration=max_duration)
+    frames = video_preprocess(frame_gen, output_shape.height, output_shape.width, dtype, device)
+    latents = video_encoder.tiled_encode(frames, tiling_config or TilingConfig.default())
+    required_latent_frames = VideoLatentShape.from_pixel_shape(output_shape).frames
+    return _conform_latent_length(latents, required_latent_frames)
+
+
+def audio_latent_from_file(
+    audio_encoder: torch.nn.Module,
+    file_path: str,
+    output_shape: VideoPixelShape,
+    device: torch.device,
+    dtype: torch.dtype,
+    start_time: float = 0.0,
+    max_duration: float | None = None,
+) -> torch.Tensor | None:
+    """Load audio from a file, and construct the audio latent conforming to video output shape.
+    Args:
+        audio_encoder: Model used to encode audio to latent space.
+        file_path: Path to the audio or video file containing an audio stream.
+        output_shape: Target video pixel shape; used to derive required latent frames
+            and, when max_duration is None, the audio duration (output_shape.frames / fps).
+        device: Device to run the encoder and hold tensors on.
+        dtype: Dtype for the output latents.
+        start_time: Start time in seconds to begin reading the audio (default 0.0).
+        max_duration: Maximum duration in seconds. If None, uses the full span implied
+            by output_shape (default None).
+    Returns:
+        Encoded audio latents of shape (1, C, T, ...) with T = required_latent_frames, or
+        None if the file has no audio stream.
+    """
+    max_duration = max_duration or output_shape.frames / output_shape.fps
+    audio_in = decode_audio_from_file(file_path, device, start_time, max_duration)
+    if audio_in is None:
+        return None
+    latents = encode_audio(audio_in, audio_encoder, None).to(device, dtype)
+    required_latent_frames = AudioLatentShape.from_video_pixel_shape(output_shape).frames
+    return _conform_latent_length(latents, required_latent_frames)
 
 
 def combined_image_conditionings(
@@ -98,7 +136,7 @@ def combined_image_conditionings(
     and using other encoded images as the keyframe conditionings."""
     conditionings = []
     for img in images:
-        image = load_image_conditioning(
+        image = load_image_and_preprocess(
             image_path=img.path,
             height=height,
             width=width,
@@ -133,7 +171,7 @@ def image_conditionings_by_replacing_latent(
 ) -> list[ConditioningItem]:
     conditionings = []
     for img in images:
-        image = load_image_conditioning(
+        image = load_image_and_preprocess(
             image_path=img.path,
             height=height,
             width=width,
@@ -163,7 +201,7 @@ def image_conditionings_by_adding_guiding_latent(
 ) -> list[ConditioningItem]:
     conditionings = []
     for img in images:
-        image = load_image_conditioning(
+        image = load_image_and_preprocess(
             image_path=img.path,
             height=height,
             width=width,
@@ -178,72 +216,6 @@ def image_conditionings_by_adding_guiding_latent(
     return conditionings
 
 
-def noise_video_state(
-    output_shape: VideoPixelShape,
-    noiser: Noiser,
-    conditionings: list[ConditioningItem],
-    components: PipelineComponents,
-    dtype: torch.dtype,
-    device: torch.device,
-    noise_scale: float = 1.0,
-    initial_latent: torch.Tensor | None = None,
-) -> tuple[LatentState, VideoLatentTools]:
-    """Initialize and noise a video latent state for the diffusion pipeline.
-    Creates a video latent state from the output shape, applies conditionings,
-    and adds noise using the provided noiser. Returns the noised state and
-    video latent tools for further processing. If initial_latent is provided, it will be used to create the initial
-    state, otherwise an empty initial state will be created.
-    """
-    video_latent_shape = VideoLatentShape.from_pixel_shape(
-        shape=output_shape,
-        latent_channels=components.video_latent_channels,
-        scale_factors=components.video_scale_factors,
-    )
-    video_tools = VideoLatentTools(components.video_patchifier, video_latent_shape, output_shape.fps)
-    video_state = create_noised_state(
-        tools=video_tools,
-        conditionings=conditionings,
-        noiser=noiser,
-        dtype=dtype,
-        device=device,
-        noise_scale=noise_scale,
-        initial_latent=initial_latent,
-    )
-
-    return video_state, video_tools
-
-
-def noise_audio_state(
-    output_shape: VideoPixelShape,
-    noiser: Noiser,
-    conditionings: list[ConditioningItem],
-    components: PipelineComponents,
-    dtype: torch.dtype,
-    device: torch.device,
-    noise_scale: float = 1.0,
-    initial_latent: torch.Tensor | None = None,
-) -> tuple[LatentState, AudioLatentTools]:
-    """Initialize and noise an audio latent state for the diffusion pipeline.
-    Creates an audio latent state from the output shape, applies conditionings,
-    and adds noise using the provided noiser. Returns the noised state and
-    audio latent tools for further processing. If initial_latent is provided, it will be used to create the initial
-    state, otherwise an empty initial state will be created.
-    """
-    audio_latent_shape = AudioLatentShape.from_video_pixel_shape(output_shape)
-    audio_tools = AudioLatentTools(components.audio_patchifier, audio_latent_shape)
-    audio_state = create_noised_state(
-        tools=audio_tools,
-        conditionings=conditionings,
-        noiser=noiser,
-        dtype=dtype,
-        device=device,
-        noise_scale=noise_scale,
-        initial_latent=initial_latent,
-    )
-
-    return audio_state, audio_tools
-
-
 def create_noised_state(
     tools: LatentTools,
     conditionings: list[ConditioningItem],
@@ -308,301 +280,14 @@ def timesteps_from_mask(denoise_mask: torch.Tensor, sigma: float | torch.Tensor)
     """Compute timesteps from a denoise mask and sigma value.
     Multiplies the denoise mask by sigma to produce timesteps for each position
     in the latent state. Areas where the mask is 0 will have zero timesteps.
+    When sigma is ``(B,)`` it is reshaped to ``(B, 1, ...)`` so the batch
+    dimension aligns correctly with ``denoise_mask``.
     """
+    if isinstance(sigma, torch.Tensor) and sigma.dim() == 1:
+        sigma = sigma.view(-1, *([1] * (denoise_mask.dim() - 1)))
     return denoise_mask * sigma
 
 
-def simple_denoising_func(
-    video_context: torch.Tensor, audio_context: torch.Tensor, transformer: X0Model
-) -> DenoisingFunc:
-    def simple_denoising_step(
-        video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        sigma = sigmas[step_index]
-        pos_video = modality_from_latent_state(video_state, video_context, sigma)
-        pos_audio = modality_from_latent_state(audio_state, audio_context, sigma)
-
-        denoised_video, denoised_audio = transformer(video=pos_video, audio=pos_audio, perturbations=None)
-        return denoised_video, denoised_audio
-
-    return simple_denoising_step
-
-
-def guider_denoising_func(
-    guider: GuiderProtocol,
-    v_context_p: torch.Tensor,
-    v_context_n: torch.Tensor,
-    a_context_p: torch.Tensor,
-    a_context_n: torch.Tensor,
-    transformer: X0Model,
-) -> DenoisingFunc:
-    def guider_denoising_step(
-        video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        sigma = sigmas[step_index]
-        pos_video = modality_from_latent_state(video_state, v_context_p, sigma)
-        pos_audio = modality_from_latent_state(audio_state, a_context_p, sigma)
-
-        denoised_video, denoised_audio = transformer(video=pos_video, audio=pos_audio, perturbations=None)
-        if guider.enabled():
-            neg_video = modality_from_latent_state(video_state, v_context_n, sigma)
-            neg_audio = modality_from_latent_state(audio_state, a_context_n, sigma)
-
-            neg_denoised_video, neg_denoised_audio = transformer(video=neg_video, audio=neg_audio, perturbations=None)
-
-            denoised_video = denoised_video + guider.delta(denoised_video, neg_denoised_video)
-            denoised_audio = denoised_audio + guider.delta(denoised_audio, neg_denoised_audio)
-
-        return denoised_video, denoised_audio
-
-    return guider_denoising_step
-
-
-def multi_modal_guider_denoising_func(
-    video_guider: MultiModalGuider,
-    audio_guider: MultiModalGuider,
-    v_context: torch.Tensor,
-    a_context: torch.Tensor,
-    transformer: X0Model,
-    *,
-    last_denoised_video: torch.Tensor | None = None,
-    last_denoised_audio: torch.Tensor | None = None,
-) -> DenoisingFunc:
-    def guider_denoising_step(
-        video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        nonlocal last_denoised_video, last_denoised_audio
-
-        if video_guider.should_skip_step(step_index) and audio_guider.should_skip_step(step_index):
-            return last_denoised_video, last_denoised_audio
-
-        sigma = sigmas[step_index]
-        pos_video_modality = modality_from_latent_state(
-            video_state, v_context, sigma, enabled=not video_guider.should_skip_step(step_index)
-        )
-        pos_audio_modality = modality_from_latent_state(
-            audio_state, a_context, sigma, enabled=not audio_guider.should_skip_step(step_index)
-        )
-
-        denoised_video, denoised_audio = transformer(
-            video=pos_video_modality, audio=pos_audio_modality, perturbations=None
-        )
-        neg_denoised_video, neg_denoised_audio = 0.0, 0.0
-        if video_guider.do_unconditional_generation() or audio_guider.do_unconditional_generation():
-            if video_guider.do_unconditional_generation() and video_guider.negative_context is None:
-                raise ValueError("Negative context is required for unconditioned denoising")
-            if audio_guider.do_unconditional_generation() and audio_guider.negative_context is None:
-                raise ValueError("Negative context is required for unconditioned denoising")
-            neg_video_modality = modality_from_latent_state(
-                video_state,
-                video_guider.negative_context
-                if video_guider.negative_context is not None
-                else pos_video_modality.context,
-                sigma,
-            )
-            neg_audio_modality = modality_from_latent_state(
-                audio_state,
-                audio_guider.negative_context
-                if audio_guider.negative_context is not None
-                else pos_audio_modality.context,
-                sigma,
-            )
-
-            neg_denoised_video, neg_denoised_audio = transformer(
-                video=neg_video_modality, audio=neg_audio_modality, perturbations=None
-            )
-
-        ptb_denoised_video, ptb_denoised_audio = 0.0, 0.0
-        if video_guider.do_perturbed_generation() or audio_guider.do_perturbed_generation():
-            perturbations = []
-            if video_guider.do_perturbed_generation():
-                perturbations.append(
-                    Perturbation(type=PerturbationType.SKIP_VIDEO_SELF_ATTN, blocks=video_guider.params.stg_blocks)
-                )
-            if audio_guider.do_perturbed_generation():
-                perturbations.append(
-                    Perturbation(type=PerturbationType.SKIP_AUDIO_SELF_ATTN, blocks=audio_guider.params.stg_blocks)
-                )
-            perturbation_config = PerturbationConfig(perturbations=perturbations)
-            ptb_denoised_video, ptb_denoised_audio = transformer(
-                video=pos_video_modality,
-                audio=pos_audio_modality,
-                perturbations=BatchedPerturbationConfig(perturbations=[perturbation_config]),
-            )
-
-        mod_denoised_video, mod_denoised_audio = 0.0, 0.0
-        if video_guider.do_isolated_modality_generation() or audio_guider.do_isolated_modality_generation():
-            perturbations = [
-                Perturbation(type=PerturbationType.SKIP_A2V_CROSS_ATTN, blocks=None),
-                Perturbation(type=PerturbationType.SKIP_V2A_CROSS_ATTN, blocks=None),
-            ]
-            perturbation_config = PerturbationConfig(perturbations=perturbations)
-            mod_denoised_video, mod_denoised_audio = transformer(
-                video=pos_video_modality,
-                audio=pos_audio_modality,
-                perturbations=BatchedPerturbationConfig(perturbations=[perturbation_config]),
-            )
-
-        if video_guider.should_skip_step(step_index):
-            denoised_video = last_denoised_video
-        else:
-            denoised_video = video_guider.calculate(
-                denoised_video, neg_denoised_video, ptb_denoised_video, mod_denoised_video
-            )
-
-        if audio_guider.should_skip_step(step_index):
-            denoised_audio = last_denoised_audio
-        else:
-            denoised_audio = audio_guider.calculate(
-                denoised_audio, neg_denoised_audio, ptb_denoised_audio, mod_denoised_audio
-            )
-
-        last_denoised_video = denoised_video
-        last_denoised_audio = denoised_audio
-
-        return denoised_video, denoised_audio
-
-    return guider_denoising_step
-
-
-def multi_modal_guider_factory_denoising_func(
-    video_guider_factory: MultiModalGuiderFactory,
-    audio_guider_factory: MultiModalGuiderFactory | None,
-    v_context: torch.Tensor,
-    a_context: torch.Tensor,
-    transformer: X0Model,
-) -> DenoisingFunc:
-    """Resolve guiders per step via factory.build_from_sigma, then multi_modal_guider_denoising_func."""
-    last_denoised_video: torch.Tensor | None = None
-    last_denoised_audio: torch.Tensor | None = None
-    sigma_vals_cached: list[float] | None = None
-
-    def guider_denoising_step(
-        video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        nonlocal last_denoised_video, last_denoised_audio, sigma_vals_cached
-        if sigma_vals_cached is None:
-            sigma_vals_cached = sigmas.detach().cpu().tolist()
-        sigma_val = sigma_vals_cached[step_index]
-        video_guider = video_guider_factory.build_from_sigma(sigma_val)
-        audio_guider = (audio_guider_factory or video_guider_factory).build_from_sigma(sigma_val)
-        denoise_fn = multi_modal_guider_denoising_func(
-            video_guider,
-            audio_guider,
-            v_context,
-            a_context,
-            transformer,
-            last_denoised_video=last_denoised_video,
-            last_denoised_audio=last_denoised_audio,
-        )
-        denoised_video, denoised_audio = denoise_fn(video_state, audio_state, sigmas, step_index)
-        last_denoised_video, last_denoised_audio = denoised_video, denoised_audio
-        return denoised_video, denoised_audio
-
-    return guider_denoising_step
-
-
-def denoise_audio_video(  # noqa: PLR0913
-    output_shape: VideoPixelShape,
-    conditionings: list[ConditioningItem],
-    noiser: Noiser,
-    sigmas: torch.Tensor,
-    stepper: DiffusionStepProtocol,
-    denoising_loop_fn: DenoisingLoopFunc,
-    components: PipelineComponents,
-    dtype: torch.dtype,
-    device: torch.device,
-    noise_scale: float = 1.0,
-    initial_video_latent: torch.Tensor | None = None,
-    initial_audio_latent: torch.Tensor | None = None,
-) -> tuple[LatentState, LatentState]:
-    video_state, video_tools = noise_video_state(
-        output_shape=output_shape,
-        noiser=noiser,
-        conditionings=conditionings,
-        components=components,
-        dtype=dtype,
-        device=device,
-        noise_scale=noise_scale,
-        initial_latent=initial_video_latent,
-    )
-    audio_state, audio_tools = noise_audio_state(
-        output_shape=output_shape,
-        noiser=noiser,
-        conditionings=[],
-        components=components,
-        dtype=dtype,
-        device=device,
-        noise_scale=noise_scale,
-        initial_latent=initial_audio_latent,
-    )
-
-    video_state, audio_state = denoising_loop_fn(
-        sigmas,
-        video_state,
-        audio_state,
-        stepper,
-    )
-
-    video_state = video_tools.clear_conditioning(video_state)
-    video_state = video_tools.unpatchify(video_state)
-    audio_state = audio_tools.clear_conditioning(audio_state)
-    audio_state = audio_tools.unpatchify(audio_state)
-
-    return video_state, audio_state
-
-
-def denoise_video_only(  # noqa: PLR0913
-    output_shape: VideoPixelShape,
-    conditionings: list[ConditioningItem],
-    noiser: Noiser,
-    sigmas: torch.Tensor,
-    stepper: DiffusionStepProtocol,
-    denoising_loop_fn: DenoisingLoopFunc,
-    components: PipelineComponents,
-    dtype: torch.dtype,
-    device: torch.device,
-    noise_scale: float = 1.0,
-    initial_video_latent: torch.Tensor | None = None,
-    initial_audio_latent: torch.Tensor | None = None,
-) -> LatentState:
-    video_state, video_tools = noise_video_state(
-        output_shape=output_shape,
-        noiser=noiser,
-        conditionings=conditionings,
-        components=components,
-        dtype=dtype,
-        device=device,
-        noise_scale=noise_scale,
-        initial_latent=initial_video_latent,
-    )
-
-    audio_state, _ = noise_audio_state(
-        output_shape=output_shape,
-        noiser=noiser,
-        conditionings=[],
-        components=components,
-        dtype=dtype,
-        device=device,
-        noise_scale=0.0,
-        initial_latent=initial_audio_latent,
-    )
-
-    audio_state = replace(audio_state, denoise_mask=torch.zeros_like(audio_state.denoise_mask))
-
-    video_state, audio_state = denoising_loop_fn(
-        sigmas,
-        video_state,
-        audio_state,
-        stepper,
-    )
-
-    video_state = video_tools.clear_conditioning(video_state)
-    video_state = video_tools.unpatchify(video_state)
-
-    return video_state
-
-
 _UNICODE_REPLACEMENTS = str.maketrans("\u2018\u2019\u201c\u201d\u2014\u2013\u00a0\u2032\u2212", "''\"\"-- '-")
 
 

packages/ltx-pipelines/src/ltx_pipelines/utils/media_io.py

@@ -12,7 +12,7 @@ from PIL import Image
 from torch._prims_common import DeviceLikeType
 from tqdm import tqdm
 
-from ltx_core.types import Audio
+from ltx_core.types import Audio, VideoPixelShape
 from ltx_pipelines.utils.constants import DEFAULT_IMAGE_CRF
 
 logger = logging.getLogger(__name__)
@@ -79,7 +79,7 @@ def normalize_latent(latent: torch.Tensor, device: torch.device, dtype: torch.dt
     return (latent / 127.5 - 1.0).to(device=device, dtype=dtype)
 
 
-def load_image_conditioning(
+def load_image_and_preprocess(
     image_path: str,
     height: int,
     width: int,
@@ -99,14 +99,23 @@ def load_image_conditioning(
     return image
 
 
-def load_video_conditioning(
-    video_path: str, height: int, width: int, frame_cap: int, dtype: torch.dtype, device: torch.device
+def video_preprocess(
+    frames: Generator[torch.Tensor],
+    height: int,
+    width: int,
+    dtype: torch.dtype,
+    device: torch.device,
 ) -> torch.Tensor:
+    """Preprocesses a video frame generator for conditioning.
+    Args:
+        frames: Generator of video frames as tensors of shape (1, H, W, C), dtype uint8.
+        height: Target height in pixels.
+        width: Target width in pixels.
+        dtype: Target dtype for the output tensor.
+        device: Target device for the output tensor.
+    Returns:
+        Tensor of shape (1, C, F, height, width) with values in [-1, 1].
     """
-    Loads a video from a path and preprocesses it for conditioning.
-    Note: The video is resized to the nearest multiple of 2 for compatibility with video codecs.
-    """
-    frames = decode_video_from_file(path=video_path, frame_cap=frame_cap, device=device)
     result = None
     for f in frames:
         frame = resize_and_center_crop(f.to(torch.float32), height, width)
@@ -257,9 +266,23 @@ def _audio_frame_to_float(frame: av.AudioFrame) -> np.ndarray:
     return arr
 
 
-def get_videostream_metadata(path: str) -> tuple[float, int, int, int]:
-    """Read video stream metadata: (fps, num_frames, width, height).
+def get_videostream_fps(path: str) -> float:
+    """Read video stream FPS."""
+    container = av.open(path)
+    try:
+        video_stream = next(s for s in container.streams if s.type == "video")
+        return float(video_stream.average_rate)
+    finally:
+        container.close()
+
+
+def get_videostream_metadata(path: str) -> VideoPixelShape:
+    """Read video stream metadata as a VideoPixelShape with batch=1.
     If frame count is missing in the container, decodes the stream to count frames.
+    Args:
+        path: Path to the video file.
+    Returns:
+        VideoPixelShape with batch=1, frames, height, width, and fps populated from the stream.
     """
     container = av.open(path)
     try:
@@ -270,7 +293,7 @@ def get_videostream_metadata(path: str) -> tuple[float, int, int, int]:
             num_frames = sum(1 for _ in container.decode(video_stream))
         width = video_stream.codec_context.width
         height = video_stream.codec_context.height
-        return fps, num_frames, width, height
+        return VideoPixelShape(batch=1, frames=num_frames, height=height, width=width, fps=fps)
     finally:
         container.close()
 
@@ -338,16 +361,85 @@ def decode_audio_from_file(
     return Audio(waveform=waveform, sampling_rate=sample_rate)
 
 
-def decode_video_from_file(path: str, frame_cap: int, device: DeviceLikeType) -> Generator[torch.Tensor]:
+def decode_video_by_frame(
+    path: str,
+    device: DeviceLikeType,
+    starting_frame: int = 0,
+    frame_cap: int | None = None,
+) -> Generator[torch.Tensor]:
+    """Decodes video from a file by sequential frame index, without relying on pts.
+    Args:
+        path: Path to the video file.
+        device: Device to place the resulting tensors on.
+        starting_frame: Number of leading frames to skip (default 0).
+        frame_cap: Maximum number of frames to yield. If None, no frame limit (default None).
+    Yields:
+        Frames as tensors of shape (1, H, W, C), dtype uint8.
+    """
     container = av.open(path)
     try:
         video_stream = next(s for s in container.streams if s.type == "video")
-        for frame in container.decode(video_stream):
+        for index, frame in enumerate(container.decode(video_stream)):
+            if index < starting_frame:
+                continue
             tensor = torch.tensor(frame.to_rgb().to_ndarray(), dtype=torch.uint8, device=device).unsqueeze(0)
             yield tensor
-            frame_cap = frame_cap - 1
-            if frame_cap == 0:
+            if frame_cap is not None:
+                frame_cap -= 1
+                if frame_cap == 0:
+                    break
+    finally:
+        container.close()
+
+
+def decode_video_from_file(
+    path: str,
+    device: DeviceLikeType,
+    start_time: float = 0.0,
+    max_duration: float | None = None,
+) -> Generator[torch.Tensor]:
+    """Decodes video from a file using presentation timestamps for time-based trimming.
+    If a frame with no pts is encountered, falls back to :func:`decode_video_by_frame`
+    using FPS-derived frame indices.
+    Args:
+        path: Path to the video file.
+        device: Device to place the resulting tensors on.
+        start_time: Start time in seconds (default 0.0).
+        max_duration: Maximum duration in seconds to decode. If None, reads to end of
+            stream (default None).
+    Yields:
+        Frames as tensors of shape (1, H, W, C), dtype uint8.
+    """
+    container = av.open(path)
+    try:
+        video_stream = next(s for s in container.streams if s.type == "video")
+        time_base = float(video_stream.time_base)
+
+        if start_time > 0:
+            container.seek(int(start_time / time_base), stream=video_stream)
+
+        end_time = start_time + max_duration if max_duration is not None else None
+
+        for frame in container.decode(video_stream):
+            # PyAV may leave pts unset when the demuxer does not expose per-frame
+            # timestamps (e.g. some raw/elementary streams, stripped or missing
+            # metadata, or certain remux paths). Without pts we cannot map frames to
+            # wall-clock time, so we fall back to sequential frame indices using the
+            # stream's average frame rate.
+            if frame.pts is None:
+                fps = float(video_stream.average_rate)
+                starting_frame = round(start_time * fps)
+                frame_cap = round(max_duration * fps) if max_duration is not None else None
+                yield from decode_video_by_frame(
+                    path=path, device=device, starting_frame=starting_frame, frame_cap=frame_cap
+                )
+                return
+            frame_time = frame.pts * time_base
+            if frame_time < start_time:
+                continue
+            if end_time is not None and frame_time >= end_time:
                 break
+            yield torch.tensor(frame.to_rgb().to_ndarray(), dtype=torch.uint8, device=device).unsqueeze(0)
     finally:
         container.close()
 

packages/ltx-pipelines/src/ltx_pipelines/utils/model_ledger.py

@@ -1,304 +0,0 @@
-from dataclasses import replace
-
-import torch
-
-from ltx_core.loader import SDOps
-from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
-from ltx_core.loader.registry import DummyRegistry, Registry
-from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
-from ltx_core.model.audio_vae import (
-    AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
-    AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
-    VOCODER_COMFY_KEYS_FILTER,
-    AudioDecoder,
-    AudioDecoderConfigurator,
-    AudioEncoder,
-    AudioEncoderConfigurator,
-    Vocoder,
-    VocoderConfigurator,
-)
-from ltx_core.model.transformer import (
-    LTXV_MODEL_COMFY_RENAMING_MAP,
-    LTXModelConfigurator,
-    X0Model,
-)
-from ltx_core.model.upsampler import LatentUpsampler, LatentUpsamplerConfigurator
-from ltx_core.model.video_vae import (
-    VAE_DECODER_COMFY_KEYS_FILTER,
-    VAE_ENCODER_COMFY_KEYS_FILTER,
-    VideoDecoder,
-    VideoDecoderConfigurator,
-    VideoEncoder,
-    VideoEncoderConfigurator,
-)
-from ltx_core.quantization import QuantizationPolicy
-from ltx_core.text_encoders.gemma import (
-    EMBEDDINGS_PROCESSOR_KEY_OPS,
-    GEMMA_LLM_KEY_OPS,
-    GEMMA_MODEL_OPS,
-    EmbeddingsProcessor,
-    EmbeddingsProcessorConfigurator,
-    GemmaTextEncoder,
-    GemmaTextEncoderConfigurator,
-    module_ops_from_gemma_root,
-)
-from ltx_core.utils import find_matching_file
-
-
-class ModelLedger:
-    """
-    Central coordinator for loading and building models used in an LTX pipeline.
-    The ledger wires together multiple model builders (transformer, video VAE encoder/decoder,
-    audio VAE decoder, vocoder, text encoder, and optional latent upsampler) and exposes
-    factory methods for constructing model instances.
-    ### Model Building
-    Each model method (e.g. :meth:`transformer`, :meth:`video_decoder`, :meth:`text_encoder`)
-    constructs a new model instance on each call. The builder uses the
-    :class:`~ltx_core.loader.registry.Registry` to load weights from the checkpoint,
-    instantiates the model with the configured ``dtype``, and moves it to ``self.device``.
-    .. note::
-        Models are **not cached**. Each call to a model method creates a new instance.
-        Callers are responsible for storing references to models they wish to reuse
-        and for freeing GPU memory (e.g. by deleting references and calling
-        ``torch.cuda.empty_cache()``).
-    ### Constructor parameters
-    dtype:
-        Torch dtype used when constructing all models (e.g. ``torch.bfloat16``).
-    device:
-        Target device to which models are moved after construction (e.g. ``torch.device("cuda")``).
-    checkpoint_path:
-        Path to a checkpoint directory or file containing the core model weights
-        (transformer, video VAE, audio VAE, text encoder, vocoder). If ``None``, the
-        corresponding builders are not created and calling those methods will raise
-        a :class:`ValueError`.
-    gemma_root_path:
-        Base path to Gemma-compatible CLIP/text encoder weights. Required to
-        initialize the text encoder builder; if omitted, :meth:`text_encoder` cannot be used.
-    spatial_upsampler_path:
-        Optional path to a latent upsampler checkpoint. If provided, the
-        :meth:`spatial_upsampler` method becomes available; otherwise calling it raises
-        a :class:`ValueError`.
-    loras:
-        Tuple of LoRA configurations (path, strength, sd_ops) applied on top of the base
-        transformer weights. Use ``()`` for none.
-    registry:
-        Optional :class:`Registry` instance for weight caching across builders.
-        Defaults to :class:`DummyRegistry` which performs no cross-builder caching.
-    quantization:
-        Optional :class:`QuantizationPolicy` controlling how transformer weights
-        are stored and how matmul is executed. Defaults to None, which means no quantization.
-    ### Creating Variants
-    Use :meth:`with_additional_loras` to create a new ``ModelLedger`` instance that
-    includes additional LoRA configurations or :meth:`with_loras` to replace existing
-    lora configurations while sharing the same registry for weight caching.
-    """
-
-    def __init__(
-        self,
-        dtype: torch.dtype,
-        device: torch.device,
-        checkpoint_path: str | None = None,
-        gemma_root_path: str | None = None,
-        spatial_upsampler_path: str | None = None,
-        loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
-        registry: Registry | None = None,
-        quantization: QuantizationPolicy | None = None,
-    ):
-        self.dtype = dtype
-        self.device = device
-        self.checkpoint_path = checkpoint_path
-        self.gemma_root_path = gemma_root_path
-        self.spatial_upsampler_path = spatial_upsampler_path
-        self.loras = loras
-        self.registry = registry or DummyRegistry()
-        self.quantization = quantization
-        self.build_model_builders()
-
-    def build_model_builders(self) -> None:
-        if self.checkpoint_path is not None:
-            self.transformer_builder = Builder(
-                model_path=self.checkpoint_path,
-                model_class_configurator=LTXModelConfigurator,
-                model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
-                loras=tuple(self.loras),
-                registry=self.registry,
-            )
-
-            self.vae_decoder_builder = Builder(
-                model_path=self.checkpoint_path,
-                model_class_configurator=VideoDecoderConfigurator,
-                model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
-                registry=self.registry,
-            )
-
-            self.vae_encoder_builder = Builder(
-                model_path=self.checkpoint_path,
-                model_class_configurator=VideoEncoderConfigurator,
-                model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
-                registry=self.registry,
-            )
-
-            self.audio_encoder_builder = Builder[AudioEncoder](
-                model_path=self.checkpoint_path,
-                model_class_configurator=AudioEncoderConfigurator,
-                model_sd_ops=AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
-                registry=self.registry,
-            )
-
-            self.audio_decoder_builder = Builder(
-                model_path=self.checkpoint_path,
-                model_class_configurator=AudioDecoderConfigurator,
-                model_sd_ops=AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
-                registry=self.registry,
-            )
-
-            self.vocoder_builder = Builder(
-                model_path=self.checkpoint_path,
-                model_class_configurator=VocoderConfigurator,
-                model_sd_ops=VOCODER_COMFY_KEYS_FILTER,
-                registry=self.registry,
-            )
-
-            # Embeddings processor only needs the LTX checkpoint (no Gemma weights)
-            self.embeddings_processor_builder = Builder(
-                model_path=self.checkpoint_path,
-                model_class_configurator=EmbeddingsProcessorConfigurator,
-                model_sd_ops=EMBEDDINGS_PROCESSOR_KEY_OPS,
-                registry=self.registry,
-            )
-
-            if self.gemma_root_path is not None:
-                module_ops = module_ops_from_gemma_root(self.gemma_root_path)
-                model_folder = find_matching_file(self.gemma_root_path, "model*.safetensors").parent
-                weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")]
-
-                self.text_encoder_builder = Builder(
-                    model_path=tuple(weight_paths),
-                    model_class_configurator=GemmaTextEncoderConfigurator,
-                    model_sd_ops=GEMMA_LLM_KEY_OPS,
-                    registry=self.registry,
-                    module_ops=(GEMMA_MODEL_OPS, *module_ops),
-                )
-
-        if self.spatial_upsampler_path is not None:
-            self.upsampler_builder = Builder(
-                model_path=self.spatial_upsampler_path,
-                model_class_configurator=LatentUpsamplerConfigurator,
-                registry=self.registry,
-            )
-
-    def _target_device(self) -> torch.device:
-        if isinstance(self.registry, DummyRegistry) or self.registry is None:
-            return self.device
-        else:
-            return torch.device("cpu")
-
-    def with_additional_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "ModelLedger":
-        """Add new lora configurations to the existing ones."""
-        return self.with_loras((*self.loras, *loras))
-
-    def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "ModelLedger":
-        """Replace existing lora configurations with new ones."""
-        return ModelLedger(
-            dtype=self.dtype,
-            device=self.device,
-            checkpoint_path=self.checkpoint_path,
-            gemma_root_path=self.gemma_root_path,
-            spatial_upsampler_path=self.spatial_upsampler_path,
-            loras=loras,
-            registry=self.registry,
-            quantization=self.quantization,
-        )
-
-    def transformer(self) -> X0Model:
-        if not hasattr(self, "transformer_builder"):
-            raise ValueError(
-                "Transformer not initialized. Please provide a checkpoint path to the ModelLedger constructor."
-            )
-
-        if self.quantization is None:
-            return (
-                X0Model(self.transformer_builder.build(device=self._target_device(), dtype=self.dtype))
-                .to(self.device)
-                .eval()
-            )
-        else:
-            sd_ops = self.transformer_builder.model_sd_ops
-            if self.quantization.sd_ops is not None:
-                sd_ops = SDOps(
-                    name=f"sd_ops_chain_{sd_ops.name}+{self.quantization.sd_ops.name}",
-                    mapping=(*sd_ops.mapping, *self.quantization.sd_ops.mapping),
-                )
-            builder = replace(
-                self.transformer_builder,
-                module_ops=(*self.transformer_builder.module_ops, *self.quantization.module_ops),
-                model_sd_ops=sd_ops,
-            )
-            return X0Model(builder.build(device=self._target_device())).to(self.device).eval()
-
-    def video_decoder(self) -> VideoDecoder:
-        if not hasattr(self, "vae_decoder_builder"):
-            raise ValueError(
-                "Video decoder not initialized. Please provide a checkpoint path to the ModelLedger constructor."
-            )
-
-        return self.vae_decoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
-
-    def video_encoder(self) -> VideoEncoder:
-        if not hasattr(self, "vae_encoder_builder"):
-            raise ValueError(
-                "Video encoder not initialized. Please provide a checkpoint path to the ModelLedger constructor."
-            )
-
-        return self.vae_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
-
-    def text_encoder(self) -> GemmaTextEncoder:
-        if not hasattr(self, "text_encoder_builder"):
-            raise ValueError(
-                "Text encoder not initialized. Please provide a checkpoint path and gemma root path to the "
-                "ModelLedger constructor."
-            )
-
-        return self.text_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
-
-    def gemma_embeddings_processor(self) -> EmbeddingsProcessor:
-        if not hasattr(self, "embeddings_processor_builder"):
-            raise ValueError(
-                "Embeddings processor not initialized. Please provide a checkpoint path to the ModelLedger constructor."
-            )
-
-        return (
-            self.embeddings_processor_builder.build(device=self._target_device(), dtype=self.dtype)
-            .to(self.device)
-            .eval()
-        )
-
-    def audio_encoder(self) -> AudioEncoder:
-        if not hasattr(self, "audio_encoder_builder"):
-            raise ValueError(
-                "Audio encoder not initialized. Please provide a checkpoint path to the ModelLedger constructor."
-            )
-
-        return self.audio_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
-
-    def audio_decoder(self) -> AudioDecoder:
-        if not hasattr(self, "audio_decoder_builder"):
-            raise ValueError(
-                "Audio decoder not initialized. Please provide a checkpoint path to the ModelLedger constructor."
-            )
-
-        return self.audio_decoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
-
-    def vocoder(self) -> Vocoder:
-        if not hasattr(self, "vocoder_builder"):
-            raise ValueError(
-                "Vocoder not initialized. Please provide a checkpoint path to the ModelLedger constructor."
-            )
-
-        return self.vocoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
-
-    def spatial_upsampler(self) -> LatentUpsampler:
-        if not hasattr(self, "upsampler_builder"):
-            raise ValueError("Upsampler not initialized. Please provide upsampler path to the ModelLedger constructor.")
-
-        return self.upsampler_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()

packages/ltx-pipelines/src/ltx_pipelines/utils/samplers.py

@@ -8,86 +8,91 @@ from tqdm import tqdm
 
 from ltx_core.components.diffusion_steps import Res2sDiffusionStep
 from ltx_core.components.protocols import DiffusionStepProtocol
+from ltx_core.model.transformer import X0Model
 from ltx_core.utils import to_denoised, to_velocity
 from ltx_pipelines.utils.helpers import post_process_latent, timesteps_from_mask
 from ltx_pipelines.utils.res2s import get_res2s_coefficients
-from ltx_pipelines.utils.types import DenoisingFunc, LatentState
+from ltx_pipelines.utils.types import Denoiser, LatentState
 
 logger = logging.getLogger(__name__)
 
 
+def _step_state(
+    state: LatentState | None,
+    denoised: torch.Tensor | None,
+    stepper: DiffusionStepProtocol,
+    sigmas: torch.Tensor,
+    step_idx: int,
+) -> LatentState | None:
+    """Advance one diffusion step for a single modality, or return ``None`` if absent."""
+    if state is None or denoised is None:
+        return state
+    denoised = post_process_latent(denoised, state.denoise_mask, state.clean_latent)
+    return replace(state, latent=stepper.step(state.latent, denoised, sigmas, step_idx))
+
+
 def euler_denoising_loop(
     sigmas: torch.Tensor,
-    video_state: LatentState,
-    audio_state: LatentState,
+    video_state: LatentState | None,
+    audio_state: LatentState | None,
     stepper: DiffusionStepProtocol,
-    denoise_fn: DenoisingFunc,
-) -> tuple[LatentState, LatentState]:
+    transformer: X0Model,
+    denoiser: Denoiser,
+) -> tuple[LatentState | None, LatentState | None]:
     """
     Perform the joint audio-video denoising loop over a diffusion schedule.
-    This function iterates over all but the final value in ``sigmas`` and, at
-    each diffusion step, calls ``denoise_fn`` to obtain denoised video and
-    audio latents. The denoised latents are post-processed with their
-    respective denoise masks and clean latents, then passed to ``stepper`` to
-    advance the noisy latents one step along the diffusion schedule.
+    Either ``video_state`` or ``audio_state`` may be ``None`` for absent
+    modalities; the absent modality is passed through unchanged.
     ### Parameters
     sigmas:
         A 1D tensor of noise levels (diffusion sigmas) defining the sampling
         schedule. All steps except the last element are iterated over.
     video_state:
-        The current video :class:`LatentState`, containing the noisy latent,
-        its clean reference latent, and the denoising mask.
+        The current video :class:`LatentState`, or ``None`` if video is absent.
     audio_state:
-        The current audio :class:`LatentState`, analogous to ``video_state``
-        but for the audio modality.
+        The current audio :class:`LatentState`, or ``None`` if audio is absent.
     stepper:
         An implementation of :class:`DiffusionStepProtocol` that updates a
         latent given the current latent, its denoised estimate, the full
         ``sigmas`` schedule, and the current step index.
-    denoise_fn:
-        A callable implementing :class:`DenoisingFunc`. It is invoked as
-        ``denoise_fn(video_state, audio_state, sigmas, step_index)`` and must
-        return a tuple ``(denoised_video, denoised_audio)``, where each element
-        is a tensor with the same shape as the corresponding latent.
+    transformer:
+        The diffusion model passed to the denoiser at each step.
+    denoiser:
+        A callable implementing :class:`Denoiser`. It is invoked as
+        ``denoiser(transformer, video_state, audio_state, sigmas, step_index)``
+        and must return ``(denoised_video, denoised_audio)``.
     ### Returns
-    tuple[LatentState, LatentState]
-        A pair ``(video_state, audio_state)`` containing the final video and
-        audio latent states after completing the denoising loop.
+    tuple[LatentState | None, LatentState | None]
+        Final ``(video_state, audio_state)`` after the denoising loop.
     """
     for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
-        denoised_video, denoised_audio = denoise_fn(video_state, audio_state, sigmas, step_idx)
+        denoised_video, denoised_audio = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
 
-        denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
-        denoised_audio = post_process_latent(denoised_audio, audio_state.denoise_mask, audio_state.clean_latent)
-
-        video_state = replace(video_state, latent=stepper.step(video_state.latent, denoised_video, sigmas, step_idx))
-        audio_state = replace(audio_state, latent=stepper.step(audio_state.latent, denoised_audio, sigmas, step_idx))
+        video_state = _step_state(video_state, denoised_video, stepper, sigmas, step_idx)
+        audio_state = _step_state(audio_state, denoised_audio, stepper, sigmas, step_idx)
 
     return (video_state, audio_state)
 
 
 def gradient_estimating_euler_denoising_loop(
     sigmas: torch.Tensor,
-    video_state: LatentState,
-    audio_state: LatentState,
+    video_state: LatentState | None,
+    audio_state: LatentState | None,
     stepper: DiffusionStepProtocol,
-    denoise_fn: DenoisingFunc,
+    transformer: X0Model,
+    denoiser: Denoiser,
     ge_gamma: float = 2.0,
-) -> tuple[LatentState, LatentState]:
+) -> tuple[LatentState | None, LatentState | None]:
     """
     Perform the joint audio-video denoising loop using gradient-estimation sampling.
-    This function is similar to :func:`euler_denoising_loop`, but applies
-    gradient estimation to improve the denoised estimates by tracking velocity
-    changes across steps. See the referenced function for detailed parameter
-    documentation.
+    Same interface as :func:`euler_denoising_loop` with an additional
+    ``ge_gamma`` parameter for velocity correction.
     ### Parameters
     ge_gamma:
         Gradient estimation coefficient controlling the velocity correction term.
         Default is 2.0. Paper: https://openreview.net/pdf?id=o2ND9v0CeK
-    sigmas, video_state, audio_state, stepper, denoise_fn:
-        See :func:`euler_denoising_loop` for parameter descriptions.
     ### Returns
-    tuple[LatentState, LatentState]
+    tuple[LatentState | None, LatentState | None]
         See :func:`euler_denoising_loop` for return value description.
     """
 
@@ -105,23 +110,35 @@ def gradient_estimating_euler_denoising_loop(
         return current_velocity, denoised_sample
 
     for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
-        denoised_video, denoised_audio = denoise_fn(video_state, audio_state, sigmas, step_idx)
+        denoised_video, denoised_audio = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
 
-        denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
-        denoised_audio = post_process_latent(denoised_audio, audio_state.denoise_mask, audio_state.clean_latent)
+        if video_state is not None and denoised_video is not None:
+            denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
+        if audio_state is not None and denoised_audio is not None:
+            denoised_audio = post_process_latent(denoised_audio, audio_state.denoise_mask, audio_state.clean_latent)
 
         if sigmas[step_idx + 1] == 0:
-            return replace(video_state, latent=denoised_video), replace(audio_state, latent=denoised_audio)
+            if video_state is not None and denoised_video is not None:
+                video_state = replace(video_state, latent=denoised_video)
+            if audio_state is not None and denoised_audio is not None:
+                audio_state = replace(audio_state, latent=denoised_audio)
+            return video_state, audio_state
 
-        previous_video_velocity, denoised_video = update_velocity_and_sample(
-            video_state.latent, denoised_video, sigmas[step_idx], previous_video_velocity
-        )
-        previous_audio_velocity, denoised_audio = update_velocity_and_sample(
-            audio_state.latent, denoised_audio, sigmas[step_idx], previous_audio_velocity
-        )
+        if video_state is not None and denoised_video is not None:
+            previous_video_velocity, denoised_video = update_velocity_and_sample(
+                video_state.latent, denoised_video, sigmas[step_idx], previous_video_velocity
+            )
+            video_state = replace(
+                video_state, latent=stepper.step(video_state.latent, denoised_video, sigmas, step_idx)
+            )
 
-        video_state = replace(video_state, latent=stepper.step(video_state.latent, denoised_video, sigmas, step_idx))
-        audio_state = replace(audio_state, latent=stepper.step(audio_state.latent, denoised_audio, sigmas, step_idx))
+        if audio_state is not None and denoised_audio is not None:
+            previous_audio_velocity, denoised_audio = update_velocity_and_sample(
+                audio_state.latent, denoised_audio, sigmas[step_idx], previous_audio_velocity
+            )
+            audio_state = replace(
+                audio_state, latent=stepper.step(audio_state.latent, denoised_audio, sigmas, step_idx)
+            )
 
     return (video_state, audio_state)
 
@@ -146,6 +163,7 @@ def _inject_sde_noise(
     sigmas: torch.Tensor,
     step_idx: int,
     legacy_mode: bool = False,
+    eta: float = 0.5,
 ) -> torch.Tensor:
     sigmas_copy = sigmas.clone()
     new_noise = new_noise_fn(state.latent, step_noise_generator)
@@ -160,6 +178,7 @@ def _inject_sde_noise(
         sigmas=sigmas,
         step_index=step_idx,
         noise=new_noise,
+        eta=eta,
     )
 
     if legacy_mode:
@@ -168,20 +187,22 @@ def _inject_sde_noise(
     return x_next
 
 
-def res2s_audio_video_denoising_loop(  # noqa: PLR0913,PLR0915
+def res2s_audio_video_denoising_loop(  # noqa: PLR0913,PLR0915,PLR0912
     sigmas: torch.Tensor,
-    video_state: LatentState,
-    audio_state: LatentState,
+    video_state: LatentState | None,
+    audio_state: LatentState | None,
     stepper: DiffusionStepProtocol,
-    denoise_fn: DenoisingFunc,
+    transformer: X0Model,
+    denoiser: Denoiser,
     noise_seed: int = -1,
     noise_seed_substep: int | None = None,
+    eta: float = 0.5,
     bongmath: bool = True,
     bongmath_max_iter: int = 100,
     new_noise_fn: Callable[[torch.Tensor, torch.Generator], torch.Tensor] = _get_new_noise,
     model_dtype: torch.dtype = torch.bfloat16,
     legacy_mode: bool = True,
-) -> tuple[LatentState, LatentState]:
+) -> tuple[LatentState | None, LatentState | None]:
     """
     Joint audio-video denoising loop using the res_2s second-order sampler.
     Iterates over the diffusion schedule with a two-stage Runge-Kutta step:
@@ -189,46 +210,48 @@ def res2s_audio_video_denoising_loop(  # noqa: PLR0913,PLR0915
     noise), then combines both with RK coefficients. Supports anchor-point
     refinement (bong iteration) and optional SDE noise injection. Requires
     :class:`Res2sDiffusionStep` as ``stepper``.
+    Either modality may be ``None`` (absent).
     ### Parameters
-    sigmas:
-        A 1D tensor of noise levels defining the sampling schedule.
-    video_state:
-        Current video :class:`LatentState` (noisy latent, clean reference, mask).
-    audio_state:
-        Current audio :class:`LatentState`, same structure as ``video_state``.
-    stepper:
-        Must be an instance of :class:`Res2sDiffusionStep`; performs SDE step
-        with noise injection.
-    denoise_fn:
-        Callable ``(video_state, audio_state, sigmas, step_index)`` returning
-        ``(denoised_video, denoised_audio)``.
+    transformer:
+        The diffusion model passed to the denoiser at each step.
+    denoiser:
+        Callable implementing :class:`Denoiser`.
     noise_seed:
         Seed for step-level SDE noise; substep seed defaults to ``noise_seed + 10000``.
     noise_seed_substep:
         Optional seed for substep SDE noise; if None, derived from ``noise_seed``.
+    eta:
+        Controls stochastic noise injection strength (0=deterministic, 1=maximum).
+        Applies to main diffusion steps; substeps always use 0.5. Default 0.5.
     bongmath:
         Whether to run iterative anchor refinement (bong iteration) when step size is small.
     bongmath_max_iter:
         Max iterations for bong refinement when enabled.
     new_noise_fn:
-        Callable ``(latent, generator) -> noise`` for SDE injection; default
-        uses normalized channel-wise Gaussian noise.
+        Callable ``(latent, generator) -> noise`` for SDE injection.
     model_dtype:
         Dtype for latent state updates (e.g. bfloat16).
     ### Returns
-    tuple[LatentState, LatentState]
+    tuple[LatentState | None, LatentState | None]
         Final ``(video_state, audio_state)`` after the denoising loop.
     """
+    # Determine device from whichever state is present
+    present_state = video_state or audio_state
+    if present_state is None:
+        raise ValueError("At least one of video_state or audio_state must be provided")
+    state_device = present_state.latent.device
+
     # Initialize noise generators with different seeds
     if noise_seed_substep is None:
         noise_seed_substep = noise_seed + 10000  # Offset to ensure different seeds
-    step_noise_generator = torch.Generator(device=video_state.latent.device).manual_seed(noise_seed)
-    substep_noise_generator = torch.Generator(device=video_state.latent.device).manual_seed(noise_seed_substep)
+    step_noise_generator = torch.Generator(device=state_device).manual_seed(noise_seed)
+    substep_noise_generator = torch.Generator(device=state_device).manual_seed(noise_seed_substep)
     sde_noise_injecting_fn = partial(
         _inject_sde_noise, stepper=stepper, new_noise_fn=new_noise_fn, legacy_mode=legacy_mode
     )
-    step_noise_injecting_fn = partial(sde_noise_injecting_fn, step_noise_generator=step_noise_generator)
-    substep_noise_injecting_fn = partial(sde_noise_injecting_fn, step_noise_generator=substep_noise_generator)
+    step_noise_injecting_fn = partial(sde_noise_injecting_fn, step_noise_generator=step_noise_generator, eta=eta)
+    # substep eta is always default 0.5 for compatibility with original implementation.
+    substep_noise_injecting_fn = partial(sde_noise_injecting_fn, step_noise_generator=substep_noise_generator, eta=0.5)
 
     if not isinstance(stepper, Res2sDiffusionStep):
         raise ValueError("stepper must be an instance of Res2sDiffusionStep")
@@ -241,26 +264,25 @@ def res2s_audio_video_denoising_loop(  # noqa: PLR0913,PLR0915
     hs = -torch.log(sigmas[1:].double().cpu() / (sigmas[:-1].double().cpu()))
 
     # Initialize phi cache for reuse across loop iterations
-    # Cache key: (j, neg_h) where j is phi order and neg_h is negative step value
     phi_cache = {}
     c2 = 0.5  # Midpoint for res_2s
 
-    # Progress bar shows only full two-stage steps; final (sigma_next==0) step is done silently
-
     for step_idx in tqdm(range(n_full_steps)):
         sigma = sigmas[step_idx].double()
         sigma_next = sigmas[step_idx + 1].double()
 
         # Initialize anchor point
-        x_anchor_video = video_state.latent.clone().double()
-        x_anchor_audio = audio_state.latent.clone().double()
+        x_anchor_video = video_state.latent.clone().double() if video_state is not None else None
+        x_anchor_audio = audio_state.latent.clone().double() if audio_state is not None else None
 
         # ====================================================================
         # STAGE 1: Evaluate at current point
         # ====================================================================
-        denoised_video_1, denoised_audio_1 = denoise_fn(video_state, audio_state, sigmas, step_idx)
-        denoised_video_1 = post_process_latent(denoised_video_1, video_state.denoise_mask, video_state.clean_latent)
-        denoised_audio_1 = post_process_latent(denoised_audio_1, audio_state.denoise_mask, audio_state.clean_latent)
+        denoised_video_1, denoised_audio_1 = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
+        if video_state is not None and denoised_video_1 is not None:
+            denoised_video_1 = post_process_latent(denoised_video_1, video_state.denoise_mask, video_state.clean_latent)
+        if audio_state is not None and denoised_audio_1 is not None:
+            denoised_audio_1 = post_process_latent(denoised_audio_1, audio_state.denoise_mask, audio_state.clean_latent)
 
         h = hs[step_idx].item()
 
@@ -273,91 +295,127 @@ def res2s_audio_video_denoising_loop(  # noqa: PLR0913,PLR0915
         # ====================================================================
         # Compute substep x using RK coefficient a21
         # ====================================================================
-        eps_1_video = denoised_video_1.double() - x_anchor_video
-        eps_1_audio = denoised_audio_1.double() - x_anchor_audio
+        if x_anchor_video is not None and denoised_video_1 is not None:
+            eps_1_video = denoised_video_1.double() - x_anchor_video
+            x_mid_video = x_anchor_video.double() + h * a21 * eps_1_video
+        else:
+            eps_1_video = None
+            x_mid_video = None
 
-        x_mid_video = x_anchor_video.double() + h * a21 * eps_1_video
-        x_mid_audio = x_anchor_audio.double() + h * a21 * eps_1_audio
+        if x_anchor_audio is not None and denoised_audio_1 is not None:
+            eps_1_audio = denoised_audio_1.double() - x_anchor_audio
+            x_mid_audio = x_anchor_audio.double() + h * a21 * eps_1_audio
+        else:
+            eps_1_audio = None
+            x_mid_audio = None
 
         # ====================================================================
         # SDE noise injection at substep
         # ====================================================================
-        x_mid_video = substep_noise_injecting_fn(
-            state=video_state,
-            sample=x_anchor_video,
-            denoised_sample=x_mid_video,
-            sigmas=torch.stack([sigma, sub_sigma]),
-            step_idx=0,
-        )
-        x_mid_audio = substep_noise_injecting_fn(
-            state=audio_state,
-            sample=x_anchor_audio,
-            denoised_sample=x_mid_audio,
-            sigmas=torch.stack([sigma, sub_sigma]),
-            step_idx=0,
-        )
+        if x_mid_video is not None and video_state is not None:
+            x_mid_video = substep_noise_injecting_fn(
+                state=video_state,
+                sample=x_anchor_video,
+                denoised_sample=x_mid_video,
+                sigmas=torch.stack([sigma, sub_sigma]),
+                step_idx=0,
+            )
+        if x_mid_audio is not None and audio_state is not None:
+            x_mid_audio = substep_noise_injecting_fn(
+                state=audio_state,
+                sample=x_anchor_audio,
+                denoised_sample=x_mid_audio,
+                sigmas=torch.stack([sigma, sub_sigma]),
+                step_idx=0,
+            )
+
         # ====================================================================
-        # ITERATIVE REFINEMENT (Bong Iteration) - Stabilize anchor point
+        # ITERATIVE REFINEMENT (Bong Iteration)
         # ====================================================================
         if bongmath and h < 0.5 and sigma > 0.03:
             for _ in range(bongmath_max_iter):
-                x_anchor_video = x_mid_video - h * a21 * eps_1_video
-                eps_1_video = denoised_video_1.double() - x_anchor_video
-                x_anchor_audio = x_mid_audio - h * a21 * eps_1_audio
-                eps_1_audio = denoised_audio_1.double() - x_anchor_audio
+                if x_mid_video is not None and eps_1_video is not None:
+                    x_anchor_video = x_mid_video - h * a21 * eps_1_video
+                    eps_1_video = denoised_video_1.double() - x_anchor_video
+                if x_mid_audio is not None and eps_1_audio is not None:
+                    x_anchor_audio = x_mid_audio - h * a21 * eps_1_audio
+                    eps_1_audio = denoised_audio_1.double() - x_anchor_audio
 
         # ====================================================================
         # STAGE 2: Evaluate at substep point (WITH NOISE)
         # ====================================================================
-        mid_video_state = replace(video_state, latent=x_mid_video.to(model_dtype))
-        mid_audio_state = replace(audio_state, latent=x_mid_audio.to(model_dtype))
+        mid_video_state = (
+            replace(video_state, latent=x_mid_video.to(model_dtype))
+            if video_state is not None and x_mid_video is not None
+            else None
+        )
+        mid_audio_state = (
+            replace(audio_state, latent=x_mid_audio.to(model_dtype))
+            if audio_state is not None and x_mid_audio is not None
+            else None
+        )
 
-        denoised_video_2, denoised_audio_2 = denoise_fn(
+        denoised_video_2, denoised_audio_2 = denoiser(
+            transformer,
             video_state=mid_video_state,
             audio_state=mid_audio_state,
             sigmas=torch.stack([sub_sigma]).to(sigmas.device),
             step_index=0,
         )
-        denoised_video_2 = post_process_latent(denoised_video_2, video_state.denoise_mask, video_state.clean_latent)
-        denoised_audio_2 = post_process_latent(denoised_audio_2, audio_state.denoise_mask, audio_state.clean_latent)
+        if video_state is not None and denoised_video_2 is not None:
+            denoised_video_2 = post_process_latent(denoised_video_2, video_state.denoise_mask, video_state.clean_latent)
+        if audio_state is not None and denoised_audio_2 is not None:
+            denoised_audio_2 = post_process_latent(denoised_audio_2, audio_state.denoise_mask, audio_state.clean_latent)
 
         # ====================================================================
         # FINAL COMBINATION: Compute x_next using RK coefficients
         # ====================================================================
-        eps_2_video = denoised_video_2.double() - x_anchor_video
-        eps_2_audio = denoised_audio_2.double() - x_anchor_audio
+        if x_anchor_video is not None and eps_1_video is not None and denoised_video_2 is not None:
+            eps_2_video = denoised_video_2.double() - x_anchor_video
+            x_next_video = x_anchor_video + h * (b1 * eps_1_video + b2 * eps_2_video)
+        else:
+            x_next_video = None
 
-        x_next_video = x_anchor_video + h * (b1 * eps_1_video + b2 * eps_2_video)
-        x_next_audio = x_anchor_audio + h * (b1 * eps_1_audio + b2 * eps_2_audio)
+        if x_anchor_audio is not None and eps_1_audio is not None and denoised_audio_2 is not None:
+            eps_2_audio = denoised_audio_2.double() - x_anchor_audio
+            x_next_audio = x_anchor_audio + h * (b1 * eps_1_audio + b2 * eps_2_audio)
+        else:
+            x_next_audio = None
 
         # ====================================================================
         # SDE NOISE INJECTION AT STEP LEVEL
         # ====================================================================
-        x_next_video = step_noise_injecting_fn(
-            state=video_state,
-            sample=x_anchor_video,
-            denoised_sample=x_next_video,
-            sigmas=sigmas,
-            step_idx=step_idx,
-        )
-        x_next_audio = step_noise_injecting_fn(
-            state=audio_state,
-            sample=x_anchor_audio,
-            denoised_sample=x_next_audio,
-            sigmas=sigmas,
-            step_idx=step_idx,
-        )
+        if x_next_video is not None and video_state is not None:
+            x_next_video = step_noise_injecting_fn(
+                state=video_state,
+                sample=x_anchor_video,
+                denoised_sample=x_next_video,
+                sigmas=sigmas,
+                step_idx=step_idx,
+            )
+        if x_next_audio is not None and audio_state is not None:
+            x_next_audio = step_noise_injecting_fn(
+                state=audio_state,
+                sample=x_anchor_audio,
+                denoised_sample=x_next_audio,
+                sigmas=sigmas,
+                step_idx=step_idx,
+            )
 
         # Update states
-        video_state = replace(video_state, latent=x_next_video.to(model_dtype))
-        audio_state = replace(audio_state, latent=x_next_audio.to(model_dtype))
+        if video_state is not None and x_next_video is not None:
+            video_state = replace(video_state, latent=x_next_video.to(model_dtype))
+        if audio_state is not None and x_next_audio is not None:
+            audio_state = replace(audio_state, latent=x_next_audio.to(model_dtype))
 
     # Final step if we need to fully remove the noise
     if sigmas[-1] == 0:
-        denoised_video_1, denoised_audio_1 = denoise_fn(video_state, audio_state, sigmas, n_full_steps)
-        denoised_video_1 = post_process_latent(denoised_video_1, video_state.denoise_mask, video_state.clean_latent)
-        denoised_audio_1 = post_process_latent(denoised_audio_1, audio_state.denoise_mask, audio_state.clean_latent)
-        video_state = replace(video_state, latent=denoised_video_1.to(model_dtype))
-        audio_state = replace(audio_state, latent=denoised_audio_1.to(model_dtype))
+        denoised_video_1, denoised_audio_1 = denoiser(transformer, video_state, audio_state, sigmas, n_full_steps)
+        if video_state is not None and denoised_video_1 is not None:
+            denoised_video_1 = post_process_latent(denoised_video_1, video_state.denoise_mask, video_state.clean_latent)
+            video_state = replace(video_state, latent=denoised_video_1.to(model_dtype))
+        if audio_state is not None and denoised_audio_1 is not None:
+            denoised_audio_1 = post_process_latent(denoised_audio_1, audio_state.denoise_mask, audio_state.clean_latent)
+            audio_state = replace(audio_state, latent=denoised_audio_1.to(model_dtype))
 
     return video_state, audio_state

packages/ltx-pipelines/src/ltx_pipelines/utils/types.py

@@ -1,9 +1,11 @@
+from dataclasses import dataclass, field
 from typing import Protocol
 
 import torch
 
 from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier
-from ltx_core.components.protocols import DiffusionStepProtocol
+from ltx_core.conditioning import ConditioningItem
+from ltx_core.model.transformer import X0Model
 from ltx_core.types import LatentState
 from ltx_pipelines.utils.constants import VIDEO_LATENT_CHANNELS, VIDEO_SCALE_FACTORS
 
@@ -35,39 +37,40 @@ class PipelineComponents:
         self.audio_patchifier = AudioPatchifier(patch_size=1)
 
 
-class DenoisingFunc(Protocol):
-    """
-    Protocol for a denoising function used in the LTX pipeline.
+class Denoiser(Protocol):
+    """Protocol for a denoiser that receives the transformer at call time.
+    The transformer is not stored — it is passed as the first argument so the
+    caller (a denoising loop or a pipeline block) controls its lifecycle.
     Args:
-        video_state (LatentState): The current latent state for video.
-        audio_state (LatentState): The current latent state for audio.
-        sigmas (torch.Tensor): A 1D tensor of sigma values for each diffusion step.
-        step_index (int): Index of the current denoising step.
+        transformer: The diffusion model.
+        video_state: Current video latent state, or ``None`` if absent.
+        audio_state: Current audio latent state, or ``None`` if absent.
+        sigmas: 1-D tensor of sigma values for each diffusion step.
+        step_index: Index of the current denoising step.
     Returns:
-        tuple[torch.Tensor, torch.Tensor]: The denoised video and audio tensors.
-    """
-
-    def __call__(
-        self, video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
-    ) -> tuple[torch.Tensor, torch.Tensor]: ...
-
-
-class DenoisingLoopFunc(Protocol):
-    """
-    Protocol for a denoising loop function used in the LTX pipeline.
-    Args:
-        sigmas (torch.Tensor): A 1D tensor of sigma values for each diffusion step.
-        video_state (LatentState): The current latent state for video.
-        audio_state (LatentState): The current latent state for audio.
-        stepper (DiffusionStepProtocol): The diffusion step protocol to use.
-    Returns:
-        tuple[LatentState, LatentState]: The denoised video and audio latent states.
+        ``(denoised_video, denoised_audio)`` tensors (either may be ``None``).
     """
 
     def __call__(
         self,
+        transformer: X0Model,
+        video_state: LatentState | None,
+        audio_state: LatentState | None,
         sigmas: torch.Tensor,
-        video_state: LatentState,
-        audio_state: LatentState,
-        stepper: DiffusionStepProtocol,
-    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+        step_index: int,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]: ...
+
+
+@dataclass(frozen=True)
+class ModalitySpec:
+    """Specification for one modality passed to a diffusion stage.
+    Carries everything needed to build the initial noised latent state
+    and run the denoising loop for a single modality (video or audio).
+    Tools are created by ``DiffusionStage`` from pixel-space dimensions.
+    """
+
+    context: torch.Tensor
+    conditionings: list[ConditioningItem] = field(default_factory=list)
+    noise_scale: float = 1.0
+    frozen: bool = False
+    initial_latent: torch.Tensor | None = None