Automated PR - 2026-05-11

Public sync - 2026-04-13

.gitattributes

@@ -7,3 +7,4 @@
 *.jpeg filter=lfs diff=lfs merge=lfs -text
 *.jpg filter=lfs diff=lfs merge=lfs -text
 *.webp filter=lfs diff=lfs merge=lfs -text
+*.exr filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -27,6 +27,7 @@ tmp
 *.sft
 
 # Media files
+*.exr
 *.gif
 *.heic
 *.heif
@@ -40,5 +41,9 @@ tmp
 *.wav
 *.webp
 
+# HDR IC-LoRA e2e test baseline (checked in via Git LFS)
+!packages/ltx-pipelines/tests/assets/expected_hdr_ic_lora_exr/frame_*.exr
+!packages/ltx-pipelines/tests/assets/hdr_ic_lora_test_input.mp4
+
 # Binary files
 *.so

README.md

@@ -39,7 +39,7 @@ Download the following models from the [LTX-2.3 HuggingFace repository](https://
 **Temporal Upscaler** - Supported by the model and will be required for future pipeline implementations
   * [`ltx-2.3-temporal-upscaler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-temporal-upscaler-x2-1.0.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-temporal-upscaler-x2-1.0.safetensors)
 
-**Distilled LoRA** - Required for current two-stage pipeline implementations in this repository (except DistilledPipeline and ICLoraPipeline)
+**Distilled LoRA** - Required for current two-stage pipeline implementations in this repository (except DistilledPipeline, ICLoraPipeline, and LipDubPipeline)
   * [`ltx-2.3-22b-distilled-lora-384-1.1.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-distilled-lora-384-1.1.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-22b-distilled-lora-384-1.1.safetensors)
 
 **Gemma Text Encoder** (download all assets from the repository)
@@ -57,6 +57,8 @@ Download the following models from the [LTX-2.3 HuggingFace repository](https://
   * [`LTX-2-19b-LoRA-Camera-Control-Jib-Down`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down/resolve/main/ltx-2-19b-lora-camera-control-jib-down.safetensors)
   * [`LTX-2-19b-LoRA-Camera-Control-Jib-Up`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up/resolve/main/ltx-2-19b-lora-camera-control-jib-up.safetensors)
   * [`LTX-2-19b-LoRA-Camera-Control-Static`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Static) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Static/resolve/main/ltx-2-19b-lora-camera-control-static.safetensors)
+  * [`LTX-2.3-22b-IC-LoRA-HDR`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-HDR) - HDR IC-LoRA and pre-computed text embeddings for `HDRICLoraPipeline`
+  * [`LTX-2.3-22b-IC-LoRA-LipDub`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-LipDub) - [Download](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-LipDub/resolve/main/ltx-2.3-22b-ic-lora-lipdub-0.9.safetensors)
 
 ### Available Pipelines
 
@@ -68,6 +70,8 @@ Download the following models from the [LTX-2.3 HuggingFace repository](https://
 * **[KeyframeInterpolationPipeline](packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py)** - Interpolate between keyframe images
 * **[A2VidPipelineTwoStage](packages/ltx-pipelines/src/ltx_pipelines/a2vid_two_stage.py)** - Audio-to-video generation conditioned on an input audio file
 * **[RetakePipeline](packages/ltx-pipelines/src/ltx_pipelines/retake.py)** - Regenerate a specific time region of an existing video
+* **[HDRICLoraPipeline](packages/ltx-pipelines/src/ltx_pipelines/hdr_ic_lora.py)** - Video-to-video with HDR output (linear float frames via LogC3 inverse decode, suitable for EXR export and tonemapping)
+* **[LipDubPipeline](packages/ltx-pipelines/src/ltx_pipelines/lipdub.py)** - Lip dubbing, rephrasing, matching speaker identity (distilled model, single IC-LoRA, Two stages).
 
 ### ⚡ Optimization Tips
 

packages/ltx-core/README.md

@@ -77,13 +77,17 @@ model = builder.build(device=torch.device("cuda"))
 Use the `.lora()` method to attach one or more LoRA adapters before calling `.build()`:
 
 ```python
+from ltx_core.loader import SDOps
+
+lora_sd_ops = SDOps(name="identity").with_matching()  # or a model-specific key-renaming SDOps
+
 builder = (
     SingleGPUModelBuilder(
         model_class_configurator=MyModelConfigurator,
         model_path="/path/to/model.safetensors",
     )
-    .lora("/path/to/lora_a.safetensors", strength=0.8)
+    .lora("/path/to/lora_a.safetensors", 0.8, lora_sd_ops)
-    .lora("/path/to/lora_b.safetensors", strength=0.5)
+    .lora("/path/to/lora_b.safetensors", 0.5, lora_sd_ops)
 )
 model = builder.build(device=torch.device("cuda"))
 ```
@@ -103,7 +107,7 @@ builder = SingleGPUModelBuilder(
     model_class_configurator=MyModelConfigurator,
     model_path="/path/to/model.safetensors",
     lora_load_device=torch.device("cuda"),
-).lora("/path/to/lora.safetensors", strength=1.0)
+).lora("/path/to/lora.safetensors", 1.0, lora_sd_ops)
 
 model = builder.build(device=torch.device("cuda"))
 ```

packages/ltx-core/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ltx-core"
-version = "1.1.1"
+version = "1.1.3"
 description = "Core implementation of Lightricks' LTX-2 model"
 readme = "README.md"
 requires-python = ">=3.10"

packages/ltx-core/src/ltx_core/batch_split.py

@@ -1,5 +1,5 @@
 """Batch-splitting adapter for the transformer.
-Wraps an ``X0Model`` (or ``LayerStreamingWrapper``) and splits batched inputs
+Wraps an ``X0Model`` (or ``BlockStreamingWrapper``) and splits batched inputs
 into smaller chunks before forwarding, then concatenates the results. This
 controls peak activation memory at the cost of more forward passes.
 The adapter is transparent — it has the same ``forward`` signature as
@@ -42,7 +42,7 @@ class BatchSplitAdapter(nn.Module):
     Has the same ``forward`` signature as ``X0Model``:
     ``(video, audio, perturbations) -> (denoised_video, denoised_audio)``.
     Args:
-        model: The model to wrap (``X0Model``, ``LayerStreamingWrapper``, etc.).
+        model: The model to wrap (``X0Model``, ``BlockStreamingWrapper``, etc.).
         max_batch_size: Maximum batch size per forward pass. Input batches
             larger than this are split into sequential chunks.
     """

packages/ltx-core/src/ltx_core/block_streaming/__init__.py

@@ -0,0 +1,19 @@
+"""Block streaming: memory-efficient sequential-block inference.
+Streams transformer blocks from safetensors to GPU one at a time.
+Block weights are provided by a :class:`WeightsProvider` which handles
+CPU-to-GPU copies, caching, and stream synchronization.  Two weight
+source strategies are available:
+- **RAM streaming** (default): all blocks pre-loaded into pinned CPU
+  buffers with LoRA fusion at build time.  Fast, higher CPU memory.
+- **Disk streaming** (``cpu_slots < num_blocks``): blocks read from
+  disk on demand with FIFO eviction.  Slower, lower CPU memory.
+"""
+
+from ltx_core.block_streaming.builder import DISK_CPU_SLOTS, StreamingModelBuilder
+from ltx_core.block_streaming.wrapper import BlockStreamingWrapper
+
+__all__ = [
+    "DISK_CPU_SLOTS",
+    "BlockStreamingWrapper",
+    "StreamingModelBuilder",
+]

packages/ltx-core/src/ltx_core/block_streaming/builder.py

@@ -0,0 +1,351 @@
+"""Builder that constructs a BlockStreamingWrapper from safetensors checkpoints."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass, field, replace
+from typing import Generic
+
+import safetensors
+import torch
+from torch import nn
+
+from ltx_core.block_streaming.disk import DiskBlockReader, DiskTensorReader, LoraSource
+from ltx_core.block_streaming.pool import WeightPool
+from ltx_core.block_streaming.provider import WeightsProvider
+from ltx_core.block_streaming.source import DiskWeightSource, PinnedWeightSource, WeightSource
+from ltx_core.block_streaming.utils import allocate_layout_views, derive_layout, make_block_key, resolve_attr
+from ltx_core.block_streaming.wrapper import BlockStreamingWrapper
+from ltx_core.loader.fuse_loras import aggregate_lora_products, fuse_lora_weights
+from ltx_core.loader.helpers import create_meta_model, load_state_dict, read_model_config
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.primitives import (
+    LoraPathStrengthAndSDOps,
+    LoraStateDictWithStrength,
+    ModelBuilderProtocol,
+    StateDictLoader,
+)
+from ltx_core.loader.registry import DummyRegistry, Registry
+from ltx_core.loader.sd_ops import SDOps
+from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
+from ltx_core.model.model_protocol import ModelConfigurator, ModelType
+
+logger = logging.getLogger(__name__)
+
+DISK_CPU_SLOTS = 2
+_DEFAULT_GPU_SLOTS = 2
+
+
+@dataclass(frozen=True)
+class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType]):
+    """Immutable builder for :class:`BlockStreamingWrapper`.
+    Reads block weights from safetensors on demand.  ``cpu_slots`` and
+    ``gpu_slots`` control the memory/speed trade-off (see :meth:`build`).
+    Args:
+        model_class_configurator: Creates the model from a config dict.
+        model_path: One or more ``.safetensors`` checkpoint paths.
+        model_sd_ops: Key remapping applied to safetensors keys.
+        module_ops: Module-level mutations for the meta model.
+        loras: LoRA adapters fused into weights at load time.
+        model_loader: Strategy for reading checkpoint metadata.
+        registry: Shared cache for loaded state dicts.
+        blocks_attr: Dotted path to the ``nn.ModuleList`` (e.g.
+            ``"velocity_model.transformer_blocks"``).
+        blocks_prefix: State-dict key prefix for block weights
+            (e.g. ``"transformer_blocks"``).
+        state_dict_prefix: Wrapper offset prepended to keys when loading into
+            the meta model (e.g. ``"velocity_model."`` when wrapped by ``X0Model``).
+        model_wrapper: Optional callable wrapping the model
+            (e.g. ``X0Model``).
+    """
+
+    model_class_configurator: type[ModelConfigurator[ModelType]]
+    model_path: str | tuple[str, ...]
+    model_sd_ops: SDOps | None = None
+    module_ops: tuple[ModuleOps, ...] = field(default_factory=tuple)
+    loras: tuple[LoraPathStrengthAndSDOps, ...] = field(default_factory=tuple)
+    model_loader: StateDictLoader = field(default_factory=SafetensorsModelStateDictLoader)
+    registry: Registry = field(default_factory=DummyRegistry)
+
+    # Streaming-specific
+    blocks_attr: str = ""
+    blocks_prefix: str = ""
+    state_dict_prefix: str = ""
+    model_wrapper: Callable[[ModelType], nn.Module] | None = None
+
+    def with_sd_ops(self, sd_ops: SDOps | None) -> StreamingModelBuilder:
+        return replace(self, model_sd_ops=sd_ops)
+
+    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> StreamingModelBuilder:
+        return replace(self, module_ops=module_ops)
+
+    def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> StreamingModelBuilder:
+        return replace(self, loras=loras)
+
+    def model_config(self) -> dict:
+        """Read model configuration from the checkpoint metadata."""
+        return read_model_config(self.model_path, self.model_loader)
+
+    def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType:
+        """Create a model on the meta device and apply module operations."""
+        return create_meta_model(self.model_class_configurator, config, module_ops)
+
+    def build(
+        self,
+        target_device: torch.device,
+        dtype: torch.dtype,
+        cpu_slots_count: int | None = None,
+        gpu_slots_count: int | None = None,
+        **_kwargs: object,
+    ) -> BlockStreamingWrapper:
+        """Build and return a ready-to-use :class:`BlockStreamingWrapper`.
+        Args:
+            target_device: GPU device for compute.
+            dtype: Weight dtype (e.g. ``torch.bfloat16``).
+            cpu_slots_count: Number of pinned CPU buffer slots.
+                ``None`` = RAM streaming (all blocks pre-loaded with LoRA fusion).
+            gpu_slots_count: Number of GPU buffer slots.
+                ``None`` = ``_DEFAULT_GPU_SLOTS`` (2).
+        """
+        if not self.blocks_prefix:
+            raise ValueError("blocks_prefix must be non-empty for streaming")
+
+        config = read_model_config(self.model_path, self.model_loader)
+        meta_model: nn.Module = create_meta_model(self.model_class_configurator, config, self.module_ops)
+        if self.model_wrapper is not None:
+            meta_model = self.model_wrapper(meta_model)
+        meta_model.eval()
+
+        blocks = resolve_attr(meta_model, self.blocks_attr)
+
+        checkpoint_paths = list(self.model_path) if isinstance(self.model_path, tuple) else [self.model_path]
+        block_key_map, non_block_keys = _scan_checkpoint_keys(checkpoint_paths, self.model_sd_ops, self.blocks_prefix)
+
+        cpu_slots_count = cpu_slots_count if cpu_slots_count is not None else len(blocks)
+        gpu_slots_count = gpu_slots_count if gpu_slots_count is not None else _DEFAULT_GPU_SLOTS
+
+        if cpu_slots_count >= len(blocks):
+            source, lora_sources = self._build_pinned_source(
+                meta_model, target_device, dtype, cpu_slots_count, block_key_map, non_block_keys
+            )
+        else:
+            reader = DiskTensorReader(checkpoint_paths)
+            source, lora_sources = self._build_disk_source(
+                meta_model, target_device, dtype, cpu_slots_count, reader, block_key_map, non_block_keys
+            )
+
+        copy_stream = torch.cuda.Stream(device=target_device)
+        gpu_pool = WeightPool(
+            source.block_layout,
+            gpu_slots_count,
+            target_device,
+            reuse_barrier=lambda event: copy_stream.wait_event(event),
+        )
+        provider = WeightsProvider(gpu_pool, copy_stream, target_device, source, lora_sources, self.blocks_prefix)
+        return BlockStreamingWrapper(
+            model=meta_model,
+            blocks=blocks,
+            provider=provider,
+            target_device=target_device,
+        )
+
+    def _build_pinned_source(
+        self,
+        meta_model: nn.Module,
+        target_device: torch.device,
+        dtype: torch.dtype,
+        cpu_slots_count: int,
+        block_key_map: dict[int, list[tuple[str, str]]],
+        non_block_keys: list[tuple[str, str]],
+    ) -> tuple[WeightSource, list[LoraSource]]:
+        """Pre-load all blocks into pinned CPU buffers with LoRA fusion."""
+        model_sd = load_state_dict(
+            self.model_path, self.model_loader, self.registry, torch.device("cpu"), self.model_sd_ops
+        )
+
+        lora_sd_and_strengths = [
+            LoraStateDictWithStrength(
+                load_state_dict([lora.path], self.model_loader, self.registry, torch.device("cpu"), lora.sd_ops),
+                lora.strength,
+            )
+            for lora in self.loras
+        ]
+
+        for block_idx in block_key_map:
+            if block_idx >= cpu_slots_count:
+                raise ValueError(
+                    f"Pinned source requires one CPU slot per block; "
+                    f"got block index {block_idx} with only {cpu_slots_count} slots."
+                )
+
+        blocks = resolve_attr(meta_model, self.blocks_attr)
+        block_tensors: dict[str, torch.Tensor] = {}
+        for block_idx, entries in block_key_map.items():
+            block_params = dict(blocks[block_idx].named_parameters())
+            for _sft_key, param_name in entries:
+                key = make_block_key(self.blocks_prefix, block_idx, param_name)
+                block_tensors[key] = block_params[param_name]
+        blocks_layout = derive_layout(block_tensors, dtype)
+        pinned_blocks = allocate_layout_views(blocks_layout, pin_memory=True)
+
+        should_sync = False
+        for key, fused in fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype=None, preserve_input_device=False):
+            if key in pinned_blocks:
+                pinned_blocks[key].copy_(fused, non_blocking=True)
+                model_sd.sd[key] = None
+                should_sync = True
+            else:
+                model_sd.sd[key] = fused
+        if should_sync:
+            torch.cuda.synchronize()
+
+        # Fill remaining pinned keys from the source state dict.
+        for key in blocks_layout:
+            if model_sd.sd[key] is None:
+                continue
+            pinned_blocks[key].copy_(model_sd.sd[key])
+            model_sd.sd[key] = None
+
+        pinned: dict[int, dict[str, torch.Tensor]] = {
+            block_idx: {
+                param_name: pinned_blocks[make_block_key(self.blocks_prefix, block_idx, param_name)]
+                for _sft_key, param_name in entries
+            }
+            for block_idx, entries in block_key_map.items()
+        }
+
+        non_block_sd: dict[str, torch.Tensor] = {
+            self.state_dict_prefix + model_key: model_sd.sd[model_key].to(device=target_device, dtype=dtype)
+            for _sft_key, model_key in non_block_keys
+        }
+
+        meta_model.load_state_dict(non_block_sd, strict=False, assign=True)
+
+        return PinnedWeightSource(pinned), []
+
+    def _build_disk_source(
+        self,
+        meta_model: nn.Module,
+        target_device: torch.device,
+        dtype: torch.dtype,
+        cpu_slots_count: int,
+        reader: DiskTensorReader,
+        block_key_map: dict[int, list[tuple[str, str]]],
+        non_block_keys: list[tuple[str, str]],
+    ) -> tuple[WeightSource, list[LoraSource]]:
+        """Create a DiskWeightSource backed by a DiskBlockReader for lazy loading.
+        Derives the shared pool layout from the meta model's block 0 — this
+        relies on module_ops (e.g. fp8_cast) leaving the meta param dtype in
+        sync with the post-sd_ops checkpoint dtype.
+        """
+        lora_sources = [LoraSource(lora.path, lora.sd_ops, lora.strength) for lora in self.loras]
+
+        self._load_non_block_weights(
+            reader,
+            non_block_keys,
+            meta_model,
+            target_device,
+            dtype,
+            sd_ops=self.model_sd_ops,
+            key_prefix=self.state_dict_prefix,
+            lora_sources=lora_sources,
+        )
+
+        blocks = resolve_attr(meta_model, self.blocks_attr)
+        layout = derive_layout(dict(blocks[0].named_parameters()), dtype)
+
+        cpu_pool = WeightPool(
+            layout,
+            cpu_slots_count,
+            torch.device("cpu"),
+            reuse_barrier=lambda event: event.synchronize(),
+            pin_memory=True,
+        )
+        block_reader = DiskBlockReader(
+            reader=reader,
+            block_key_map=block_key_map,
+            sd_ops=self.model_sd_ops,
+            blocks_prefix=self.blocks_prefix,
+        )
+        source = DiskWeightSource(cpu_pool, block_reader)
+        return source, lora_sources
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _fuse_lora_delta(
+        model_key: str,
+        tensor: torch.Tensor,
+        lora_sources: list[LoraSource],
+    ) -> torch.Tensor:
+        """Add all matching LoRA deltas to *tensor* in-place via ``addmm_``."""
+        if not lora_sources or not model_key.endswith(".weight"):
+            return tensor
+        prefix = model_key[: -len(".weight")]
+        products = (
+            ab
+            for ab in (s.get_ab(prefix, device=tensor.device, dtype=tensor.dtype) for s in lora_sources)
+            if ab is not None
+        )
+        aggregate_lora_products(products, out=tensor)
+        return tensor
+
+    @staticmethod
+    @torch.inference_mode()
+    def _load_non_block_weights(
+        reader: DiskTensorReader,
+        non_block_keys: list[tuple[str, str]],
+        model: nn.Module,
+        device: torch.device,
+        dtype: torch.dtype,
+        sd_ops: SDOps | None = None,
+        key_prefix: str = "",
+        lora_sources: list[LoraSource] | None = None,
+    ) -> None:
+        """Load non-block weights into *model* on *device*."""
+        state_dict: dict[str, torch.Tensor] = {}
+        sources = lora_sources or []
+        for sft_key, model_key in non_block_keys:
+            tensor = reader.get_tensor(sft_key).to(device=device, dtype=dtype)
+            tensor = StreamingModelBuilder._fuse_lora_delta(model_key, tensor, sources)
+            if sd_ops is not None:
+                for kv in sd_ops.apply_to_key_value(model_key, tensor):
+                    state_dict[key_prefix + kv.new_key] = kv.new_value
+                continue
+            state_dict[key_prefix + model_key] = tensor
+        model.load_state_dict(state_dict, strict=False, assign=True)
+
+
+def _scan_checkpoint_keys(
+    checkpoint_paths: list[str],
+    sd_ops: SDOps | None,
+    blocks_prefix: str,
+) -> tuple[dict[int, list[tuple[str, str]]], list[tuple[str, str]]]:
+    """Partition checkpoint keys into per-block and non-block lists.
+    Opens the safetensors files for header-only key enumeration; no tensor data
+    is read.
+    """
+    block_key_map: dict[int, list[tuple[str, str]]] = {}
+    non_block_keys: list[tuple[str, str]] = []
+    prefix_dot = blocks_prefix + "."
+    for path in checkpoint_paths:
+        with safetensors.safe_open(path, framework="pt", device="cpu") as handle:
+            for sft_key in handle.keys():  # noqa: SIM118
+                model_key = sd_ops.apply_to_key(sft_key) if sd_ops else sft_key
+                if model_key is None:
+                    continue
+                if model_key.startswith(prefix_dot):
+                    rest = model_key[len(prefix_dot) :]
+                    idx_str, _, param_name = rest.partition(".")
+                    try:
+                        block_idx = int(idx_str)
+                    except ValueError:
+                        non_block_keys.append((sft_key, model_key))
+                        continue
+                    block_key_map.setdefault(block_idx, []).append((sft_key, param_name))
+                else:
+                    non_block_keys.append((sft_key, model_key))
+    return block_key_map, non_block_keys

packages/ltx-core/src/ltx_core/block_streaming/disk.py

@@ -0,0 +1,148 @@
+"""Safetensors I/O and LoRA fusion for block streaming."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+
+import safetensors
+import torch
+
+from ltx_core.block_streaming.utils import allocate_layout_views, make_block_key
+from ltx_core.loader.fuse_loras import LoraProduct
+from ltx_core.loader.sd_ops import SDOps
+
+_SAFETENSORS_DTYPE_TO_TORCH: dict[str, torch.dtype] = {
+    "F64": torch.float64,
+    "F32": torch.float32,
+    "F16": torch.float16,
+    "BF16": torch.bfloat16,
+}
+
+
+class DiskTensorReader:
+    """Key-based tensor accessor over one or more safetensors files."""
+
+    def __init__(self, paths: list[str]) -> None:
+        self._handles: list[safetensors.safe_open] = []
+        self._key_to_handle_idx: dict[str, int] = {}
+        for path in paths:
+            handle = safetensors.safe_open(path, framework="pt", device="cpu")
+            handle_idx = len(self._handles)
+            self._handles.append(handle)
+            for sft_key in handle.keys():  # noqa: SIM118
+                self._key_to_handle_idx[sft_key] = handle_idx
+
+    def get_tensor(self, key: str) -> torch.Tensor:
+        return self._handles[self._key_to_handle_idx[key]].get_tensor(key)
+
+    def close(self) -> None:
+        self._handles.clear()
+        self._key_to_handle_idx.clear()
+
+    def __contains__(self, key: str) -> bool:
+        return key in self._key_to_handle_idx
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._key_to_handle_idx)
+
+
+class DiskBlockReader:
+    """Reads one block at a time from safetensors into provided buffers."""
+
+    def __init__(
+        self,
+        reader: DiskTensorReader,
+        block_key_map: dict[int, list[tuple[str, str]]],
+        sd_ops: SDOps | None = None,
+        blocks_prefix: str = "",
+    ) -> None:
+        self._reader = reader
+        self._block_key_map = block_key_map
+        self._sd_ops = sd_ops
+        self._blocks_prefix = blocks_prefix
+
+    def read_into(self, target: dict[str, torch.Tensor], block_idx: int) -> None:
+        block_prefix = make_block_key(self._blocks_prefix, block_idx, "")
+        for sft_key, param_name in self._block_key_map[block_idx]:
+            tensor = self._reader.get_tensor(sft_key)
+            if self._sd_ops is None:
+                target[param_name].copy_(tensor)
+                continue
+            full_key = make_block_key(self._blocks_prefix, block_idx, param_name)
+            for result in self._sd_ops.apply_to_key_value(full_key, tensor):
+                if not result.new_key.startswith(block_prefix):
+                    raise ValueError(
+                        f"SDOps output key '{result.new_key}' is outside block {block_idx} "
+                        f"(expected prefix '{block_prefix}'); cannot route to a per-block buffer."
+                    )
+                target[result.new_key[len(block_prefix) :]].copy_(result.new_value)
+
+    def cleanup(self) -> None:
+        self._reader.close()
+
+
+class LoraSource:
+    """Pinned-memory cache of matched LoRA A/B factors backed by a single buffer."""
+
+    def __init__(self, path: str, sd_ops: SDOps | None, strength: float) -> None:
+        self.strength = strength
+        self._pinned_ab: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
+
+        a_keys: dict[str, str] = {}
+        b_keys: dict[str, str] = {}
+        with safetensors.safe_open(path, framework="pt", device="cpu") as handle:
+            for sft_key in handle.keys():  # noqa: SIM118
+                model_key = sd_ops.apply_to_key(sft_key) if sd_ops is not None else sft_key
+                if model_key is None:
+                    continue
+                if model_key.endswith(".lora_A.weight"):
+                    a_keys[model_key[: -len(".lora_A.weight")]] = sft_key
+                elif model_key.endswith(".lora_B.weight"):
+                    b_keys[model_key[: -len(".lora_B.weight")]] = sft_key
+
+            matched_prefixes = list(a_keys.keys() & b_keys.keys())
+
+            # Build the layout from safetensors header metadata only — no tensor data is read.
+            layout: dict[str, tuple[torch.Size, torch.dtype]] = {}
+            for prefix in matched_prefixes:
+                a_slice_view = handle.get_slice(a_keys[prefix])
+                b_slice_view = handle.get_slice(b_keys[prefix])
+                layout[f"{prefix}.A"] = (
+                    torch.Size(a_slice_view.get_shape()),
+                    _SAFETENSORS_DTYPE_TO_TORCH[a_slice_view.get_dtype()],
+                )
+                layout[f"{prefix}.B"] = (
+                    torch.Size(b_slice_view.get_shape()),
+                    _SAFETENSORS_DTYPE_TO_TORCH[b_slice_view.get_dtype()],
+                )
+
+            all_views = allocate_layout_views(layout, pin_memory=True)
+
+            for prefix in matched_prefixes:
+                a_view = all_views[f"{prefix}.A"]
+                b_view = all_views[f"{prefix}.B"]
+                a_view.copy_(handle.get_tensor(a_keys[prefix]))
+                b_view.copy_(handle.get_tensor(b_keys[prefix]))
+                self._pinned_ab[prefix] = (a_view, b_view)
+
+    def get_ab(
+        self,
+        param_prefix: str,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> LoraProduct | None:
+        """Return the :class:`LoraProduct` for *param_prefix*, or ``None``."""
+        pair = self._pinned_ab.get(param_prefix)
+        if pair is None:
+            return None
+        a, b = pair
+        if device is not None and device.type == "cuda":
+            a = a.to(device=device, non_blocking=True)
+            b = b.to(device=device, non_blocking=True)
+        if dtype is not None:
+            a = a.to(dtype=dtype)
+            b = b.to(dtype=dtype)
+        return LoraProduct(a, b, self.strength)
+
+    def cleanup(self) -> None:
+        self._pinned_ab.clear()

packages/ltx-core/src/ltx_core/block_streaming/pool.py

@@ -0,0 +1,75 @@
+"""Weight buffer pool for block streaming."""
+
+from __future__ import annotations
+
+from collections import deque
+from typing import Callable
+
+import torch
+
+from ltx_core.block_streaming.utils import allocate_layout_views
+from ltx_core.loader.primitives import TensorLayout
+
+
+class WeightPool:
+    """Fixed pool of pre-allocated weight buffers with event-based reuse.
+    All slots share a single buffer (CPU or GPU); each slot is a
+    contiguous slice carved out of it via :func:`allocate_layout_views`.
+    Args:
+        buffer_layout: ``{name: (shape, dtype)}`` for each buffer.
+        capacity: Number of buffers to pre-allocate.
+        device: Device for allocation.
+        reuse_barrier: Called with the pending event before a buffer is reused.
+        pin_memory: Pin buffers (for async H2D copies from CPU).
+    """
+
+    def __init__(
+        self,
+        buffer_layout: TensorLayout,
+        capacity: int,
+        device: torch.device,
+        reuse_barrier: Callable[[torch.cuda.Event], None],
+        pin_memory: bool = False,
+    ) -> None:
+        self._buffer_layout = buffer_layout
+        self._capacity = capacity
+        self._free: deque[dict[str, torch.Tensor]] = deque()
+        self._events: dict[int, torch.cuda.Event] = {}
+        self._reuse_barrier = reuse_barrier
+        memory_layout = {
+            _make_key(slot, name): (shape, dtype)
+            for slot in range(capacity)
+            for name, (shape, dtype) in buffer_layout.items()
+        }
+        all_views = allocate_layout_views(memory_layout, device=device, pin_memory=pin_memory)
+        for slot in range(capacity):
+            self._free.append({name: all_views[_make_key(slot, name)] for name in buffer_layout})
+
+    @property
+    def capacity(self) -> int:
+        return self._capacity
+
+    @property
+    def buffer_layout(self) -> TensorLayout:
+        return self._buffer_layout
+
+    def acquire(self) -> dict[str, torch.Tensor]:
+        """Take a free buffer, waiting any pending event before returning."""
+        weights = self._free.popleft()
+        event = self._events.pop(id(weights), None)
+        if event is not None:
+            self._reuse_barrier(event)
+        return weights
+
+    def release(self, weights: dict[str, torch.Tensor], event: torch.cuda.Event | None = None) -> None:
+        """Return a buffer to the free list.
+        If *event* is given it is waited on the next :meth:`acquire`
+        of this buffer, ensuring the prior operation has completed.
+        """
+        if event is not None:
+            self._events[id(weights)] = event
+        self._free.append(weights)
+
+
+def _make_key(slot: int, name: str) -> str:
+    return f"{slot}/{name}"

packages/ltx-core/src/ltx_core/block_streaming/provider.py

@@ -0,0 +1,147 @@
+"""GPU weights provider for block streaming."""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+
+import torch
+
+from ltx_core.block_streaming.disk import LoraSource
+from ltx_core.block_streaming.pool import WeightPool
+from ltx_core.block_streaming.source import WeightSource
+from ltx_core.block_streaming.utils import FP8_DTYPES
+from ltx_core.loader.fuse_loras import aggregate_lora_products, fuse_cast_fp8_weight
+
+
+def _contiguous_byte_view(weights: dict[str, torch.Tensor]) -> torch.Tensor | None:
+    """Return a ``uint8`` view spanning every tensor in *weights*, or ``None`` if
+    they don't share one contiguous storage region."""
+    tensors = list(weights.values())
+    if not tensors:
+        return None
+    storage = tensors[0].untyped_storage()
+    storage_ptr = storage.data_ptr()
+    start = end = tensors[0].storage_offset() * tensors[0].element_size()
+    for t in tensors:
+        if t.untyped_storage().data_ptr() != storage_ptr or not t.is_contiguous():
+            return None
+        offset = t.storage_offset() * t.element_size()
+        nbytes = t.numel() * t.element_size()
+        start = min(start, offset)
+        end = max(end, offset + nbytes)
+    view = torch.empty(0, dtype=torch.uint8, device=tensors[0].device)
+    view.set_(storage, start, (end - start,), (1,))
+    return view
+
+
+class WeightsProvider:
+    """Provides GPU-ready block weights via H2D copy from a pinned CPU weight source.
+    Args:
+        pool: Pre-allocated GPU weight buffer pool.
+        copy_stream: Dedicated CUDA stream for async H2D copies.
+        target_device: GPU device for compute.
+        source: Pinned CPU weight source.
+        lora_sources: LoRA adapters fused on H2D copy.
+        blocks_prefix: State-dict prefix for LoRA key matching.
+    """
+
+    def __init__(
+        self,
+        pool: WeightPool,
+        copy_stream: torch.cuda.Stream,
+        target_device: torch.device,
+        source: WeightSource,
+        lora_sources: list[LoraSource] | None = None,
+        blocks_prefix: str = "",
+    ) -> None:
+        self._copy_stream = copy_stream
+        self._pool = pool
+        self._cache: OrderedDict[int, dict[str, torch.Tensor]] = OrderedDict()
+        self._events: dict[int, torch.cuda.Event] = {}
+        self._target_device = target_device
+        self._source = source
+        self._lora_sources = lora_sources or []
+        self._blocks_prefix = blocks_prefix
+
+    def get(self, idx: int) -> dict[str, torch.Tensor]:
+        """Return GPU weights for block *idx*. Does H2D copy on miss."""
+        if idx in self._cache:
+            return self._cache[idx]
+
+        # Evict oldest GPU buffer if at capacity.
+        if len(self._cache) >= self._pool.capacity:
+            evicted_idx, evicted_weights = self._cache.popitem(last=False)
+            self._pool.release(evicted_weights, event=self._events.pop(evicted_idx, None))
+
+        gpu_weights = self._pool.acquire()
+        cpu_weights = self._source.get(idx)
+
+        h2d_event = self._copy_to_gpu(idx, gpu_weights, cpu_weights)
+        self._source.release(idx, event=h2d_event)
+
+        self._cache[idx] = gpu_weights
+        return gpu_weights
+
+    def _copy_to_gpu(
+        self,
+        idx: int,
+        gpu_weights: dict[str, torch.Tensor],
+        cpu_weights: dict[str, torch.Tensor],
+    ) -> torch.cuda.Event:
+        """Enqueue H2D copy + LoRA fusion on the copy stream and wait on compute.
+        The wait is intentionally inside this method so callers -- and
+        instrumentation regions wrapping it -- observe the full transfer time.
+        """
+        with torch.cuda.stream(self._copy_stream):
+            gpu_view = _contiguous_byte_view(gpu_weights)
+            cpu_view = _contiguous_byte_view(cpu_weights)
+            if gpu_view is not None and cpu_view is not None and gpu_view.numel() == cpu_view.numel():
+                gpu_view.copy_(cpu_view, non_blocking=True)
+            else:
+                for name, gpu_tensor in gpu_weights.items():
+                    gpu_tensor.copy_(cpu_weights[name], non_blocking=True)
+            if self._lora_sources:
+                self._fuse_block_loras(idx, gpu_weights)
+            h2d_event = torch.cuda.Event()
+            h2d_event.record(self._copy_stream)
+
+        torch.cuda.current_stream(self._target_device).wait_event(h2d_event)
+        return h2d_event
+
+    def release(self, idx: int, event: torch.cuda.Event) -> None:
+        """Attach a compute-done event -- waited before this buffer is recycled."""
+        self._events[idx] = event
+
+    def cleanup(self) -> None:
+        """Synchronize streams and release all resources."""
+        self._copy_stream.synchronize()
+        torch.cuda.current_stream(self._target_device).synchronize()
+        self._cache.clear()
+        self._events.clear()
+        self._source.cleanup()
+        for lora in self._lora_sources:
+            lora.cleanup()
+
+    def __len__(self) -> int:
+        return len(self._cache)
+
+    def _fuse_block_loras(self, idx: int, weights: dict[str, torch.Tensor]) -> None:
+        """Fuse LoRA deltas directly into GPU block weights."""
+        for name, tensor in weights.items():
+            if not name.endswith(".weight"):
+                continue
+            prefix = f"{self._blocks_prefix}.{idx}.{name}".removesuffix(".weight")
+            is_fp8 = tensor.dtype in FP8_DTYPES
+            agg_dtype = torch.bfloat16 if is_fp8 else tensor.dtype
+            products = (
+                ab
+                for ab in (s.get_ab(prefix, device=self._target_device, dtype=agg_dtype) for s in self._lora_sources)
+                if ab is not None
+            )
+            aggregated = aggregate_lora_products(products, agg_dtype)
+            if aggregated is None:
+                continue
+            if is_fp8:
+                tensor.copy_(fuse_cast_fp8_weight(aggregated, tensor, tensor.dtype))
+            else:
+                tensor.add_(aggregated)

packages/ltx-core/src/ltx_core/block_streaming/source.py

@@ -0,0 +1,102 @@
+"""Weight sources for block streaming: protocol and implementations."""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from typing import Protocol
+
+import torch
+
+from ltx_core.block_streaming.disk import DiskBlockReader
+from ltx_core.block_streaming.pool import WeightPool
+from ltx_core.loader.primitives import TensorLayout
+
+
+class WeightSource(Protocol):
+    """Provides pinned CPU weights for a given block index.
+    Assumes all buffers share an identical layout across all block indices.
+    """
+
+    @property
+    def block_layout(self) -> TensorLayout:
+        """Shared per-block buffer layout (shape + dtype for each param)."""
+        ...
+
+    def get(self, idx: int) -> dict[str, torch.Tensor]:
+        """Return CPU weights for block *idx*."""
+        ...
+
+    def release(self, idx: int, event: torch.cuda.Event) -> None:
+        """Signal that an async operation using these weights is guarded by *event*."""
+        ...
+
+    def cleanup(self) -> None:
+        """Release all resources (buffers, readers, events)."""
+        ...
+
+
+class DiskWeightSource(WeightSource):
+    """Reads block weights from disk into pinned CPU buffers on demand."""
+
+    def __init__(self, pool: WeightPool, reader: DiskBlockReader) -> None:
+        self._pool = pool
+        self._cache: OrderedDict[int, dict[str, torch.Tensor]] = OrderedDict()
+        self._events: dict[int, torch.cuda.Event] = {}
+        self._reader = reader
+
+    @property
+    def block_layout(self) -> TensorLayout:
+        return self._pool.buffer_layout
+
+    def get(self, idx: int) -> dict[str, torch.Tensor]:
+        """Return CPU weights for block *idx*. Reads from disk on miss."""
+        if idx in self._cache:
+            return self._cache[idx]
+
+        if len(self._cache) >= self._pool.capacity:
+            evicted_idx, evicted_weights = self._cache.popitem(last=False)
+            self._pool.release(evicted_weights, event=self._events.pop(evicted_idx, None))
+
+        weights = self._pool.acquire()
+        self._reader.read_into(weights, idx)
+        self._cache[idx] = weights
+        return weights
+
+    def release(self, idx: int, event: torch.cuda.Event) -> None:
+        """Attach an H2D event -- waited before this buffer is recycled."""
+        self._events[idx] = event
+
+    def cleanup(self) -> None:
+        """Clear cache and close the disk reader."""
+        self._cache.clear()
+        self._events.clear()
+        self._reader.cleanup()
+
+    def __len__(self) -> int:
+        return len(self._cache)
+
+
+class PinnedWeightSource(WeightSource):
+    """Pre-loaded pinned CPU weights."""
+
+    def __init__(self, weights: dict[int, dict[str, torch.Tensor]]) -> None:
+        if not weights:
+            raise ValueError("PinnedWeightSource requires at least one block")
+        self._weights = weights
+
+    @property
+    def block_layout(self) -> TensorLayout:
+        first_block = self._weights[min(self._weights)]
+        return {name: (t.shape, t.dtype) for name, t in first_block.items()}
+
+    def get(self, idx: int) -> dict[str, torch.Tensor]:
+        return self._weights[idx]
+
+    def release(self, idx: int, event: torch.cuda.Event) -> None:
+        pass
+
+    def cleanup(self) -> None:
+        self._weights.clear()
+
+    def __len__(self) -> int:
+        return len(self._weights)

packages/ltx-core/src/ltx_core/block_streaming/utils.py

@@ -0,0 +1,134 @@
+"""Shared utilities for the block_streaming package."""
+
+from __future__ import annotations
+
+import math
+import weakref
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import nn
+
+from ltx_core.loader.primitives import TensorLayout
+
+FP8_DTYPES = frozenset({torch.float8_e4m3fn, torch.float8_e5m2})
+
+_BUFFER_ALIGN = 16
+
+
+def make_block_key(blocks_prefix: str, block_idx: int, param_name: str) -> str:
+    """Return the state-dict key for *param_name* under block *block_idx*."""
+    return f"{blocks_prefix}.{block_idx}.{param_name}"
+
+
+def resolve_attr(module: nn.Module, dotted_path: str) -> nn.ModuleList:
+    """Resolve a dotted attribute path like ``'model.language_model.layers'``."""
+    obj: Any = module
+    for part in dotted_path.split("."):
+        obj = getattr(obj, part)
+    if not isinstance(obj, nn.ModuleList):
+        raise TypeError(f"Expected nn.ModuleList at '{dotted_path}', got {type(obj).__name__}")
+    return obj
+
+
+def assign_tensor_to_module(root: nn.Module, dotted_name: str, tensor: torch.Tensor) -> None:
+    """Assign *tensor* to the parameter/buffer at *dotted_name* inside *root*.
+    Unlike ``param.data = tensor``, this works even when the existing parameter
+    lives on the ``meta`` device (which has an incompatible storage type).
+    """
+    parts = dotted_name.split(".")
+    parent = root
+    for part in parts[:-1]:
+        parent = getattr(parent, part)
+    leaf = parts[-1]
+    if leaf in parent._parameters:
+        parent._parameters[leaf] = nn.Parameter(tensor, requires_grad=False)
+    elif leaf in parent._buffers:
+        parent._buffers[leaf] = tensor
+    else:
+        raise AttributeError(f"{leaf} is not a parameter or buffer of {type(parent).__name__}")
+
+
+def derive_layout(tensors: dict[str, torch.Tensor], dtype: torch.dtype | None = None) -> TensorLayout:
+    """Derive a layout from a ``{name: tensor}`` dict.
+    If ``dtype`` is given, non-FP8 dtypes are coerced to it (FP8 preserved). If
+    ``None``, the source dtype is preserved as-is.
+    """
+    return {
+        name: (t.shape, t.dtype if dtype is None or t.dtype in FP8_DTYPES else dtype) for name, t in tensors.items()
+    }
+
+
+def _align_up(offset: int, alignment: int) -> int:
+    return (offset + alignment - 1) & ~(alignment - 1)
+
+
+def _alloc_pinned_exact(nbytes: int) -> torch.Tensor | None:
+    """Allocate exactly ``nbytes`` of pinned host memory via ``cudaHostRegister``.
+    Bypasses PyTorch's ``CachingHostAllocator``, which rounds every
+    ``pin_memory=True`` request up to ``PowerOf2Ceil(N)`` (see
+    ``aten/src/ATen/core/CachingHostAllocator.h``). Returns ``None`` if
+    registration fails. The unregister hook is bound to the storage (not the
+    tensor) so views of the buffer keep the registration alive until the
+    memory is actually freed. Caller is responsible for ensuring CUDA is
+    available.
+    """
+    cudart = torch.cuda.cudart()
+    buf = torch.empty(nbytes, dtype=torch.uint8)
+    ptr = buf.data_ptr()
+    err = int(cudart.cudaHostRegister(ptr, nbytes, 0))
+    if err != 0:
+        return None
+    weakref.finalize(buf.untyped_storage(), lambda p=ptr: cudart.cudaHostUnregister(p))
+    return buf
+
+
+def _alloc_buffer(nbytes: int, device: torch.device | None, pin_memory: bool) -> torch.Tensor:
+    """Allocate one ``uint8`` buffer for :func:`allocate_layout_views`.
+    For pinned host buffers, prefer ``cudaHostRegister`` to dodge the caching
+    allocator's power-of-2 rounding. Falls back to the caching allocator if
+    registration fails. Raises if pinning is requested without a CUDA runtime,
+    since pinning is fundamentally a CUDA driver operation.
+    """
+    if pin_memory and (device is None or torch.device(device).type == "cpu"):
+        if not torch.cuda.is_available():
+            raise RuntimeError("pin_memory=True requires CUDA, which is not available")
+        buf = _alloc_pinned_exact(nbytes)
+        if buf is not None:
+            return buf
+    return torch.empty(nbytes, dtype=torch.uint8, device=device, pin_memory=pin_memory)
+
+
+@dataclass(frozen=True)
+class _TensorSlice:
+    """Location of a single tensor view within the buffer."""
+
+    offset: int
+    shape: torch.Size
+    dtype: torch.dtype
+
+    def size(self) -> int:
+        return math.prod(self.shape) * self.dtype.itemsize
+
+
+def allocate_layout_views(
+    layout: TensorLayout,
+    device: torch.device | None = None,
+    pin_memory: bool = False,
+) -> dict[str, torch.Tensor]:
+    """Allocate a single ``uint8`` buffer and return per-key tensor views into it.
+    All keys in *layout* live in one contiguous allocation; each returned
+    tensor is a non-overlapping slice of that buffer reinterpreted at the
+    requested shape and dtype. The views keep the underlying storage alive
+    via PyTorch refcounting — drop them all to release the memory.
+    """
+    slices: dict[str, _TensorSlice] = {}
+    cursor = 0
+    for key, (shape, dtype) in layout.items():
+        cursor = _align_up(cursor, _BUFFER_ALIGN)
+        slices[key] = _TensorSlice(offset=cursor, shape=shape, dtype=dtype)
+        cursor += slices[key].size()
+    # Allocate at least one byte so empty layouts still produce a valid buffer.
+    buffer = _alloc_buffer(max(_align_up(cursor, _BUFFER_ALIGN), 1), device, pin_memory)
+    return {key: buffer[s.offset : s.offset + s.size()].view(s.dtype).view(s.shape) for key, s in slices.items()}

packages/ltx-core/src/ltx_core/block_streaming/wrapper.py

@@ -0,0 +1,96 @@
+"""Block streaming wrapper: streams transformer blocks through a WeightsProvider."""
+
+from __future__ import annotations
+
+import itertools
+from typing import Any
+
+import torch
+from torch import nn
+
+from ltx_core.block_streaming.provider import WeightsProvider
+from ltx_core.block_streaming.utils import assign_tensor_to_module
+
+
+class BlockStreamingWrapper(nn.Module):
+    """Streams sequential model blocks through GPU buffer caches.
+    The wrapper delegates all weight management to a :class:`WeightsProvider`
+    which handles CPU-to-GPU copies, caching, LoRA fusion, and stream
+    synchronization internally.
+    Use :class:`StreamingModelBuilder` to construct this wrapper -- it
+    handles checkpoint parsing, source selection, and provider creation.
+    Args:
+        model: The wrapped model (non-block params already on GPU).
+        blocks: Sequential blocks to stream (``nn.ModuleList``).
+        provider: Provides GPU-ready weights on demand.
+        target_device: GPU device for compute.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        blocks: nn.ModuleList,
+        provider: WeightsProvider,
+        target_device: torch.device,
+    ) -> None:
+        super().__init__()
+        self._model = model
+        self._blocks = blocks
+        self._target_device = target_device
+        self._provider = provider
+
+        self._hooks: list[torch.utils.hooks.RemovableHandle] = []
+        self._register_hooks()
+
+    # ------------------------------------------------------------------
+    # Hook registration
+    # ------------------------------------------------------------------
+
+    def _pre_hook(self, block_idx: int) -> None:
+        """Load GPU weights for a block and inject them into its parameters."""
+        gpu_weights = self._provider.get(block_idx)
+
+        block = self._blocks[block_idx]
+        for name, _param in itertools.chain(block.named_parameters(), block.named_buffers()):
+            assign_tensor_to_module(block, name, gpu_weights[name])
+
+    def _post_hook(self, block_idx: int) -> None:
+        """Record a compute-done event and release the block weights."""
+        compute_done = torch.cuda.Event()
+        compute_done.record(torch.cuda.current_stream(self._target_device))
+        self._provider.release(block_idx, event=compute_done)
+
+    def _register_hooks(self) -> None:
+        for idx, block in enumerate(self._blocks):
+            pre = block.register_forward_pre_hook(
+                lambda _mod, _args, *, idx=idx: self._pre_hook(idx),
+            )
+            post = block.register_forward_hook(
+                lambda _mod, _args, _out, *, idx=idx: self._post_hook(idx),
+            )
+            self._hooks.extend([pre, post])
+
+    # ------------------------------------------------------------------
+    # Teardown
+    # ------------------------------------------------------------------
+
+    def teardown(self) -> None:
+        """Remove hooks and release all resources."""
+        for h in self._hooks:
+            h.remove()
+        self._hooks.clear()
+        self._provider.cleanup()
+
+    # ------------------------------------------------------------------
+    # Forward and attribute delegation
+    # ------------------------------------------------------------------
+
+    def forward(self, *args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+        return self._model(*args, **kwargs)
+
+    def __getattr__(self, name: str) -> Any:  # noqa: ANN401
+        """Proxy attribute access to the wrapped model."""
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self._model, name)

packages/ltx-core/src/ltx_core/components/diffusion_steps.py

@@ -4,6 +4,24 @@ from ltx_core.components.protocols import DiffusionStepProtocol
 from ltx_core.utils import to_velocity
 
 
+def _get_ancestral_step(
+    sigma_from: torch.Tensor,
+    sigma_to: torch.Tensor,
+    eta: float = 1.0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute ``(sigma_down, sigma_up)`` for one DDIM ancestral sampling step.
+    Both inputs are in the rescaled parameterization ``sigma / alpha``.
+    Returns ``sigma_down`` (deterministic component) and ``sigma_up``
+    (stochastic component) in the same rescaled space.
+    """
+    if not eta:
+        return sigma_to, torch.zeros_like(sigma_to)
+    variance = sigma_to**2 * (sigma_from**2 - sigma_to**2).clamp(min=0) / sigma_from**2
+    sigma_up = (eta * variance**0.5).clamp(max=sigma_to)
+    sigma_down = (sigma_to**2 - sigma_up**2).clamp(min=0) ** 0.5
+    return sigma_down, sigma_up
+
+
 class EulerDiffusionStep(DiffusionStepProtocol):
     """
     First-order Euler method for diffusion sampling.
@@ -104,3 +122,65 @@ class Res2sDiffusionStep(DiffusionStepProtocol):
         # Mix deterministic and stochastic components
         x_noised = alpha_ratio * (denoised_next + sigma_down * eps_next) + sigma_up * noise
         return x_noised.to(output_dtype)
+
+
+class EulerCfgPpDiffusionStep(DiffusionStepProtocol):
+    """Euler step using the CFG++ correction for the ODE derivative.
+    Instead of the standard velocity formula, the ODE derivative is computed
+    from the unconditioned prediction, keeping the conditioned prediction as
+    the target denoised state.  Ancestral (DDIM) noise injection is applied
+    in the rescaled sigma parameterization (sigma / alpha).
+    All diffusion quantities (alpha, ODE derivative, ancestral coefficients)
+    are computed internally from ``sigmas`` and ``uncond_denoised``.
+    Reference: CFG++ (https://arxiv.org/abs/2406.08070).
+    """
+
+    def __init__(self, eta: float = 1.0, s_noise: float = 1.0) -> None:
+        self.eta = eta
+        self.s_noise = s_noise
+
+    def step(
+        self,
+        sample: torch.Tensor,
+        denoised_sample: torch.Tensor,
+        sigmas: torch.Tensor,
+        step_index: int,
+        uncond_denoised: torch.Tensor,
+        noise: torch.Tensor | None = None,
+        **_kwargs,
+    ) -> torch.Tensor:
+        """Advance one CFG++ Euler step.
+        Args:
+            sample: Current noisy latent x_t.
+            denoised_sample: Conditioned denoised prediction x_0^cond.
+            sigmas: Full sigma schedule tensor.
+            step_index: Current step index.
+            uncond_denoised: Unconditioned denoised prediction x_0^uncond,
+                used to compute the ODE derivative direction.
+            noise: Noise tensor for stochastic injection; ignored when
+                ``eta=0`` or ``s_noise=0``.
+        Returns:
+            Updated latent x_{t-1}.
+        """
+        sigma_s = sigmas[step_index].to(torch.float32)
+        sigma_t = sigmas[step_index + 1].to(torch.float32)
+        _eps = torch.finfo(torch.float32).eps
+        # Clamp to avoid division by zero when sigma == 1.0 exactly.
+        alpha_s = (1.0 - sigma_s).clamp(min=_eps)
+        alpha_t = (1.0 - sigma_t).clamp(min=_eps)
+
+        x = sample.to(torch.float32)
+        denoised = denoised_sample.to(torch.float32)
+        uncond = uncond_denoised.to(torch.float32)
+
+        # ODE derivative: direction toward noise using uncond prediction (CFG++ correction)
+        d = (x - alpha_s * uncond) / sigma_s
+
+        # Ancestral step in rescaled sigma space (sigma / alpha)
+        sigma_down, sigma_up = _get_ancestral_step(sigma_s / alpha_s, sigma_t / alpha_t, eta=self.eta)
+        sigma_down = alpha_t * sigma_down
+
+        x_next = alpha_t * denoised + sigma_down * d
+        if noise is not None and self.eta > 0 and self.s_noise > 0:
+            x_next = x_next + alpha_t * noise.to(torch.float32) * self.s_noise * sigma_up
+        return x_next.to(sample.dtype)

packages/ltx-core/src/ltx_core/components/patchifiers.py

@@ -151,17 +151,22 @@ def get_pixel_coords(
             that treat frame zero differently still yield non-negative timestamps.
     """
     # Broadcast the VAE scale factors so they align with the `(batch, axis, patch, bound)` layout.
+    # Axis 1 of `latent_coords` is ordered (frame/time, height, width) — match that explicitly by
+    # pulling fields from the NamedTuple rather than relying on tuple iteration order.
     broadcast_shape = [1] * latent_coords.ndim
     broadcast_shape[1] = -1  # axis dimension corresponds to (frame/time, height, width)
-    scale_tensor = torch.tensor(scale_factors, device=latent_coords.device).view(*broadcast_shape)
+    scale_tensor = torch.tensor(
+        [scale_factors.time, scale_factors.height, scale_factors.width],
+        device=latent_coords.device,
+    ).view(*broadcast_shape)
 
     # Apply per-axis scaling to convert latent bounds into pixel-space coordinates.
     pixel_coords = latent_coords * scale_tensor
 
     if causal_fix:
-        # VAE temporal stride for the very first frame is 1 instead of `scale_factors[0]`.
+        # VAE temporal stride for the very first frame is 1 instead of `scale_factors.time`.
         # Shift and clamp to keep the first-frame timestamps causal and non-negative.
-        pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + 1 - scale_factors[0]).clamp(min=0)
+        pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + 1 - scale_factors.time).clamp(min=0)
 
     return pixel_coords
 

packages/ltx-core/src/ltx_core/conditioning/__init__.py

@@ -3,6 +3,7 @@
 from ltx_core.conditioning.exceptions import ConditioningError
 from ltx_core.conditioning.item import ConditioningItem
 from ltx_core.conditioning.types import (
+    AudioConditionByReferenceLatent,
     ConditioningItemAttentionStrengthWrapper,
     VideoConditionByKeyframeIndex,
     VideoConditionByLatentIndex,
@@ -10,6 +11,7 @@ from ltx_core.conditioning.types import (
 )
 
 __all__ = [
+    "AudioConditionByReferenceLatent",
     "ConditioningError",
     "ConditioningItem",
     "ConditioningItemAttentionStrengthWrapper",

packages/ltx-core/src/ltx_core/conditioning/types/__init__.py

@@ -3,9 +3,11 @@
 from ltx_core.conditioning.types.attention_strength_wrapper import ConditioningItemAttentionStrengthWrapper
 from ltx_core.conditioning.types.keyframe_cond import VideoConditionByKeyframeIndex
 from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex
+from ltx_core.conditioning.types.reference_audio_cond import AudioConditionByReferenceLatent
 from ltx_core.conditioning.types.reference_video_cond import VideoConditionByReferenceLatent
 
 __all__ = [
+    "AudioConditionByReferenceLatent",
     "ConditioningItemAttentionStrengthWrapper",
     "VideoConditionByKeyframeIndex",
     "VideoConditionByLatentIndex",

packages/ltx-core/src/ltx_core/conditioning/types/keyframe_cond.py

@@ -17,12 +17,20 @@ class VideoConditionByKeyframeIndex(ConditioningItem):
         keyframes: Keyframe latents [B, C, F, H, W].
         frame_idx: Frame index offset for positional encoding.
         strength: Conditioning strength (1.0 = clean, 0.0 = fully denoised).
+        num_pixel_frames: Number of pixel frames the keyframe latent originally encodes.
     """
 
-    def __init__(self, keyframes: torch.Tensor, frame_idx: int, strength: float):
+    def __init__(
+        self,
+        keyframes: torch.Tensor,
+        frame_idx: int,
+        strength: float,
+        num_pixel_frames: int = 1,
+    ):
         self.keyframes = keyframes
         self.frame_idx = frame_idx
         self.strength = strength
+        self.num_pixel_frames = num_pixel_frames
 
     def apply_to(
         self,
@@ -41,6 +49,11 @@ class VideoConditionByKeyframeIndex(ConditioningItem):
         )
 
         positions[:, 0, ...] += self.frame_idx
+        # If the keyframe latent encodes a single pixel frame,
+        # narrow the temporal end to [start, start + 1) instead of the
+        # VAE-scaled range.
+        if self.num_pixel_frames == 1:
+            positions[:, 0, ..., 1:] = positions[:, 0, ..., :1] + 1
         positions = positions.to(dtype=torch.float32)
         positions[:, 0, ...] /= latent_tools.fps
 

packages/ltx-core/src/ltx_core/conditioning/types/reference_audio_cond.py

@@ -0,0 +1,59 @@
+"""Audio reference conditioning items."""
+
+from __future__ import annotations
+
+import torch
+
+from ltx_core.conditioning.mask_utils import update_attention_mask
+from ltx_core.tools import LatentTools
+from ltx_core.types import LatentState
+
+
+class AudioConditionByReferenceLatent:
+    """Append patchified reference audio tokens after the target audio sequence.
+    Mirrors :class:`ltx_core.conditioning.types.reference_video_cond.VideoConditionByReferenceLatent`
+    but for audio. The reference tokens are appended so the target audio tokens stay
+    in the first ``num_noisy_tokens`` positions and can be kept by
+    :meth:`ltx_core.tools.LatentTools.clear_conditioning`.
+    Args:
+        patchified: Patchified reference latent ``[B, T_ref, C]``.
+        positions: RoPE positions for reference tokens, ``[B, 1, T_ref, 2]``.
+        strength: 1.0 keeps reference clean; 0.0 would fully denoise it.
+    """
+
+    def __init__(
+        self,
+        patchified: torch.Tensor,
+        positions: torch.Tensor,
+        strength: float = 1.0,
+    ) -> None:
+        self.patchified = patchified
+        self.positions = positions.to(dtype=torch.float32)
+        self.strength = strength
+
+    def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
+        tokens = self.patchified
+        denoise_mask = torch.full(
+            size=(*tokens.shape[:2], 1),
+            fill_value=1.0 - self.strength,
+            device=tokens.device,
+            dtype=tokens.dtype,
+        )
+
+        new_attention_mask = update_attention_mask(
+            latent_state=latent_state,
+            attention_mask=None,
+            num_noisy_tokens=latent_tools.patchifier.get_token_count(latent_tools.target_shape),
+            num_new_tokens=tokens.shape[1],
+            batch_size=tokens.shape[0],
+            device=tokens.device,
+            dtype=tokens.dtype,
+        )
+
+        return LatentState(
+            latent=torch.cat([latent_state.latent, tokens], dim=1),
+            denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
+            positions=torch.cat([latent_state.positions, self.positions], dim=2),
+            clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
+            attention_mask=new_attention_mask,
+        )

packages/ltx-core/src/ltx_core/hdr.py

@@ -0,0 +1,71 @@
+"""HDR utilities: LogC3 compression for HDR IC-LoRA training and inference.
+Provides compress/decompress and postprocess helpers for HDR video generation.
+Used by ltx-pipelines for HDR IC-LoRA and by ltx-trainer for HDR validation.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+from torch import Tensor
+
+
+class LogC3:
+    """ARRI LogC3 (EI 800) HDR compression.
+    Maps linear [0, ∞) <-> LogC3 [0, 1] via the camera log curve. The log
+    curve allocates more precision to shadows/midtones and compresses
+    highlights smoothly. Callers are responsible for mapping the [0, 1]
+    output to the VAE's [-1, 1] input range.
+    """
+
+    name = "LogC3"
+    A = 5.555556
+    B = 0.052272
+    C = 0.247190
+    D = 0.385537
+    E = 5.367655
+    F = 0.092809
+    CUT = 0.010591
+
+    def compress(self, hdr: Tensor) -> Tensor:
+        """Compress linear HDR [0, ∞) → LogC3 [0, 1]."""
+        x = torch.clamp(hdr, min=0.0)
+        log_part = self.C * torch.log10(self.A * x + self.B) + self.D
+        lin_part = self.E * x + self.F
+        logc = torch.where(x >= self.CUT, log_part, lin_part)
+        return torch.clamp(logc, 0.0, 1.0)
+
+    def compress_ldr(self, ldr: Tensor) -> Tensor:
+        """Compress LDR [0, 1] → [0, 1] (no log curve, just clamp)."""
+        return torch.clamp(ldr, 0.0, 1.0)
+
+    def decompress(self, logc: Tensor) -> Tensor:
+        """Decompress LogC3 [0, 1] → linear HDR [0, ∞)."""
+        logc = torch.clamp(logc, 0.0, 1.0)
+        cut_log = self.E * self.CUT + self.F
+        lin_from_log = (torch.pow(10.0, (logc - self.D) / self.C) - self.B) / self.A
+        lin_from_lin = (logc - self.F) / self.E
+        return torch.where(logc >= cut_log, lin_from_log, lin_from_lin)
+
+    def decompress_ldr(self, logc: Tensor) -> Tensor:
+        """Decompress [0, 1] → LDR [0, 1] (identity clamp)."""
+        return torch.clamp(logc, 0.0, 1.0)
+
+
+def apply_hdr_decode_postprocess(
+    decoded_video: Tensor,
+    transform: Literal["logc3"] = "logc3",
+) -> Tensor:
+    """Apply HDR decompress to VAE decode output for HDR recovery.
+    Args:
+        decoded_video: Tensor from VAE decode in [0, 1], shape [B, C, F, H, W].
+            Must be float32 for sufficient color resolution.
+        transform: "logc3".
+    Returns:
+        HDR video tensor float32.
+    """
+    decoded_video = decoded_video.float()
+    if transform == "logc3":
+        return LogC3().decompress(decoded_video)
+    raise ValueError(f"Unsupported HDR transform: {transform}")

packages/ltx-core/src/ltx_core/layer_streaming.py

@@ -1,306 +0,0 @@
-"""Layer streaming wrapper for memory-efficient inference.
-Keeps most transformer/decoder layers on CPU pinned memory and streams them
-to GPU on demand, using a secondary CUDA stream to prefetch upcoming layers
-so that data transfer overlaps with compute.
-General-purpose: works with any ``nn.Module`` whose forward iterates over a
-``nn.ModuleList`` attribute (e.g. ``transformer_blocks``, ``layers``).
-Each layer is evicted back to CPU immediately after its forward completes,
-and prefetch uses modular indexing so the last layer's prefetch wraps around
-to prepare early layers for the next forward pass.
-Example
--------
->>> model = build_my_model(device=torch.device("cpu"))
->>> model = LayerStreamingWrapper(
-...     model,
-...     layers_attr="transformer_blocks",
-...     target_device=torch.device("cuda:0"),
-...     prefetch_count=2,
-... )
->>> out = model(inputs)            # hooks handle layer streaming
->>> model.teardown()               # move everything back to CPU
-"""
-
-from __future__ import annotations
-
-import functools
-import itertools
-import logging
-from typing import Any
-
-import torch
-from torch import nn
-
-logger = logging.getLogger(__name__)
-
-
-def _resolve_attr(module: nn.Module, dotted_path: str) -> nn.ModuleList:
-    """Resolve a dotted attribute path like ``'model.language_model.layers'``."""
-    obj: Any = module
-    for part in dotted_path.split("."):
-        obj = getattr(obj, part)
-    if not isinstance(obj, nn.ModuleList):
-        raise TypeError(f"Expected nn.ModuleList at '{dotted_path}', got {type(obj).__name__}")
-    return obj
-
-
-class _LayerStore:
-    """Manages CPU-pinned copies of layer parameters/buffers.
-    Tracks which layers currently reside on GPU so the prefetcher and evictor
-    can make correct decisions.
-    """
-
-    def __init__(self, layers: nn.ModuleList, target_device: torch.device) -> None:
-        self.target_device = target_device
-        self.num_layers = len(layers)
-
-        # CPU-pinned copies keyed by (layer_idx, param_name)
-        self._pinned: list[dict[str, torch.Tensor]] = []
-        self._on_gpu: set[int] = set()
-
-        for layer in layers:
-            pinned: dict[str, torch.Tensor] = {}
-            for name, tensor in itertools.chain(layer.named_parameters(), layer.named_buffers()):
-                pinned_tensor = tensor.data.pin_memory()
-                tensor.data = pinned_tensor
-                pinned[name] = pinned_tensor
-            self._pinned.append(pinned)
-
-    def _check_idx(self, idx: int) -> None:
-        if idx < 0 or idx >= self.num_layers:
-            raise IndexError(f"Layer index {idx} out of range [0, {self.num_layers})")
-
-    def is_on_gpu(self, idx: int) -> bool:
-        return idx in self._on_gpu
-
-    def move_to_gpu(self, idx: int, layer: nn.Module, *, non_blocking: bool = False) -> None:
-        """Move layer *idx* parameters from pinned CPU to *target_device*."""
-        self._check_idx(idx)
-        if idx in self._on_gpu:
-            return
-        pinned = self._pinned[idx]
-        for name, param in itertools.chain(layer.named_parameters(), layer.named_buffers()):
-            param.data = pinned[name].to(self.target_device, non_blocking=non_blocking)
-        self._on_gpu.add(idx)
-
-    def evict_to_cpu(self, idx: int, layer: nn.Module) -> None:
-        """Swap layer *idx* parameters back to their pinned CPU copies."""
-        self._check_idx(idx)
-        if idx not in self._on_gpu:
-            return
-        pinned = self._pinned[idx]
-        for name, param in itertools.chain(layer.named_parameters(), layer.named_buffers()):
-            param.data = pinned[name]
-        self._on_gpu.discard(idx)
-
-    def cleanup(self) -> None:
-        """Release all pinned memory references.
-        After this call, the pinned tensors can be garbage-collected once
-        the layer parameters (which still reference them via ``.data``) are
-        also released (e.g. via ``.to("meta")``).
-        """
-        for pinned_dict in self._pinned:
-            pinned_dict.clear()
-        self._pinned.clear()
-
-
-class _AsyncPrefetcher:
-    """Issues H2D transfers on a dedicated CUDA stream.
-    Uses per-layer CUDA events so that the compute stream only waits for the
-    specific layer it needs, not all pending transfers.
-    """
-
-    def __init__(self, store: _LayerStore, layers: nn.ModuleList) -> None:
-        self._store = store
-        self._layers = layers
-        self._stream = torch.cuda.Stream(device=store.target_device)
-        self._events: dict[int, torch.cuda.Event] = {}
-
-    def prefetch(self, idx: int) -> None:
-        """Begin async transfer of layer *idx* to GPU (no-op if already there)."""
-        if self._store.is_on_gpu(idx) or idx in self._events:
-            return
-        with torch.cuda.stream(self._stream):
-            self._store.move_to_gpu(idx, self._layers[idx], non_blocking=True)
-            event = torch.cuda.Event()
-            event.record(self._stream)
-            self._events[idx] = event
-
-    def wait(self, idx: int) -> None:
-        """Block the compute stream until layer *idx* transfer is complete."""
-        event = self._events.pop(idx, None)
-        if event is not None:
-            torch.cuda.current_stream(self._store.target_device).wait_event(event)
-
-    def cleanup(self) -> None:
-        """Drain pending work and release CUDA stream/event resources."""
-        self._events.clear()
-        self._stream = None
-        self._layers = None
-        self._store = None
-
-
-class LayerStreamingWrapper(nn.Module):
-    """Wraps a model to stream its sequential layers between CPU and GPU.
-    Each layer is evicted immediately after its forward completes, and
-    prefetch wraps around using modular indexing so the end of one forward
-    pass prepares early layers for the next.
-    Parameters
-    ----------
-    model:
-        The model to wrap, with all parameters on **CPU**.
-    layers_attr:
-        Dotted attribute path to the ``nn.ModuleList`` of sequential layers
-        (e.g. ``"transformer_blocks"`` or ``"model.language_model.layers"``).
-    target_device:
-        The GPU device to use for compute.
-    prefetch_count:
-        How many layers ahead to prefetch.  The maximum number of layers on
-        GPU at once is ``1 + prefetch_count``.  Must be >= 1.
-    """
-
-    def __init__(
-        self,
-        model: nn.Module,
-        layers_attr: str,
-        target_device: torch.device,
-        prefetch_count: int = 2,
-    ) -> None:
-        if prefetch_count < 1:
-            raise ValueError("prefetch_count must be >= 1")
-        super().__init__()
-        # Store the wrapped model as a submodule so parameters are discoverable.
-        self._model = model
-        self._layers = _resolve_attr(model, layers_attr)
-        self._target_device = target_device
-        # Clamp: no point prefetching more than num_layers - 1 (the rest are evicted).
-        self._prefetch_count = min(prefetch_count, len(self._layers) - 1)
-        self._hooks: list[torch.utils.hooks.RemovableHandle] = []
-
-        self._setup()
-
-    # ------------------------------------------------------------------
-    # Setup / teardown
-    # ------------------------------------------------------------------
-
-    def _setup(self) -> None:
-        # 1. Build the pinned CPU store (copies all layer tensors to pinned memory).
-        self._store = _LayerStore(self._layers, self._target_device)
-
-        # 2. Move all NON-layer params/buffers to GPU.
-        layer_tensor_ids: set[int] = set()
-        for layer in self._layers:
-            for t in itertools.chain(layer.parameters(), layer.buffers()):
-                layer_tensor_ids.add(id(t))
-
-        for p in self._model.parameters():
-            if id(p) not in layer_tensor_ids:
-                p.data = p.data.to(self._target_device)
-        for b in self._model.buffers():
-            if id(b) not in layer_tensor_ids:
-                b.data = b.data.to(self._target_device)
-
-        # 3. Pre-load the first (1 + prefetch_count) layers synchronously.
-        for idx in range(min(self._prefetch_count + 1, len(self._layers))):
-            self._store.move_to_gpu(idx, self._layers[idx])
-
-        # 4. Create the async prefetcher and register hooks.
-        self._prefetcher = _AsyncPrefetcher(self._store, self._layers)
-        self._register_hooks()
-
-    def _register_hooks(self) -> None:
-        idx_map: dict[int, int] = {id(layer): idx for idx, layer in enumerate(self._layers)}
-        num_layers = len(self._layers)
-
-        def _pre_hook(
-            module: nn.Module,
-            _args: Any,  # noqa: ANN401
-            *,
-            idx: int,
-        ) -> None:
-            # Wait only for THIS layer's H2D transfer (not all pending ones).
-            self._prefetcher.wait(idx)
-            if not self._store.is_on_gpu(idx):
-                self._store.move_to_gpu(idx, module)
-
-            # Record that the compute stream will read these weight tensors.
-            # They were allocated on the prefetch stream, so without this the
-            # caching allocator would allow the prefetch stream to reuse their
-            # memory immediately after eviction — even if the compute kernel
-            # that reads them hasn't finished yet.
-            compute_stream = torch.cuda.current_stream(self._target_device)
-            for param in itertools.chain(module.parameters(), module.buffers()):
-                param.data.record_stream(compute_stream)
-
-            # Kick off prefetch for upcoming layers (wraps around for next pass).
-            for offset in range(1, self._prefetch_count + 1):
-                self._prefetcher.prefetch((idx + offset) % num_layers)
-
-        def _post_hook(
-            module: nn.Module,
-            _args: Any,  # noqa: ANN401
-            _output: Any,  # noqa: ANN401
-            *,
-            idx: int,
-        ) -> None:
-            # Evict this layer immediately — its computation is done.
-            self._store.evict_to_cpu(idx, module)
-
-        for layer in self._layers:
-            idx = idx_map[id(layer)]
-            h1 = layer.register_forward_pre_hook(functools.partial(_pre_hook, idx=idx))
-            h2 = layer.register_forward_hook(functools.partial(_post_hook, idx=idx))
-            self._hooks.extend([h1, h2])
-
-    def teardown(self) -> None:
-        """Remove hooks, release pinned memory, and move parameters back to CPU.
-        After this call the wrapper is inert: hooks are removed, the prefetch
-        stream is drained and destroyed, all parameters reside on regular
-        (non-pinned) CPU memory, and the ``_LayerStore`` pinned-tensor cache is
-        cleared.  Callers should still follow up with ``.to("meta")`` to release
-        the CPU copies if the model is no longer needed.
-        """
-        for h in self._hooks:
-            h.remove()
-        self._hooks.clear()
-
-        # Drain all in-flight async H2D copies, then release stream resources.
-        # Without the synchronize, clearing the stream/events can trigger
-        # use-after-free at the CUDA driver level.
-        torch.cuda.synchronize(device=self._target_device)
-        if self._prefetcher is not None:
-            self._prefetcher.cleanup()
-            self._prefetcher = None
-
-        # Move everything to CPU.
-        for idx, layer in enumerate(self._layers):
-            self._store.evict_to_cpu(idx, layer)
-
-        for p in self._model.parameters():
-            p.data = p.data.to("cpu")
-        for b in self._model.buffers():
-            b.data = b.data.to("cpu")
-
-        # Release pinned memory.  After evict_to_cpu() the layer parameters
-        # still reference the pinned tensors (since .to("cpu") on a pinned
-        # tensor is a no-op).  The caller is expected to follow up with
-        # .to("meta") to drop the param refs; cleanup() drops the store's refs.
-        self._store.cleanup()
-
-    # ------------------------------------------------------------------
-    # Forward and attribute delegation
-    # ------------------------------------------------------------------
-
-    def forward(self, *args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
-        return self._model(*args, **kwargs)
-
-    def __getattr__(self, name: str) -> Any:  # noqa: ANN401
-        """Proxy attribute access to the wrapped model.
-        This allows calling methods like ``encode()`` on a wrapped
-        GemmaTextEncoder without the caller needing to know about the wrapper.
-        ``nn.Module.__getattr__`` is only called when normal attribute lookup
-        fails, so ``_model``, ``_store``, etc. are found first via ``__dict__``.
-        """
-        try:
-            return super().__getattr__(name)
-        except AttributeError:
-            return getattr(self._model, name)

packages/ltx-core/src/ltx_core/loader/__init__.py

@@ -1,6 +1,11 @@
 """Loader utilities for model weights, LoRAs, and safetensor operations."""
 
 from ltx_core.loader.fuse_loras import apply_loras
+from ltx_core.loader.helpers import (
+    create_meta_model,
+    load_state_dict,
+    read_model_config,
+)
 from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.primitives import (
     LoRAAdaptableProtocol,
@@ -45,4 +50,7 @@ __all__ = [
     "StateDictLoader",
     "StateDictRegistry",
     "apply_loras",
+    "create_meta_model",
+    "load_state_dict",
+    "read_model_config",
 ]

packages/ltx-core/src/ltx_core/loader/fuse_loras.py

@@ -1,26 +1,81 @@
-from collections.abc import Iterator
+from collections.abc import Iterable, Iterator
+from typing import NamedTuple
 
 import torch
 
+from ltx_core.loader.kernels import TRITON_AVAILABLE
 from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict
-from ltx_core.quantization.fp8_cast import _fused_add_round_launch
+from ltx_core.quantization.fp8_cast import fused_add_round_launch
 from ltx_core.quantization.fp8_scaled_mm import quantize_weight_to_fp8_per_tensor
 
 
+class LoraProduct(NamedTuple):
+    """A LoRA's ``A``, ``B`` factors and its strength scalar."""
+
+    a: torch.Tensor
+    b: torch.Tensor
+    strength: float
+
+
 def _get_device() -> torch.device:
     if torch.cuda.is_available():
         return torch.device("cuda", torch.cuda.current_device())
     return torch.device("cpu")
 
 
+def aggregate_lora_products(
+    products: Iterable[LoraProduct],
+    dtype: torch.dtype | None = None,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor | None:
+    """Accumulate ``sum((B * strength) @ A)`` across :class:`LoraProduct` items.
+    If ``out`` is provided, ``addmm_`` accumulates directly into it — caller
+    ensures A/B dtypes and devices match ``out``. Otherwise the first product
+    materializes the ``(out, in)``-shape aggregator at ``dtype``; subsequent
+    products use ``addmm_`` to avoid allocating the full intermediate delta.
+    Returns ``out`` (or the new aggregator), or ``None`` if ``products`` was empty
+    and ``out`` was not given.
+    """
+    aggregated = out
+    for product in products:
+        if aggregated is None:
+            aggregated = torch.matmul(product.b * product.strength, product.a).to(dtype=dtype)
+        else:
+            aggregated.addmm_(product.b, product.a, alpha=product.strength)
+    return aggregated
+
+
+def fuse_cast_fp8_weight(
+    delta_bf16: torch.Tensor,
+    weight_fp8: torch.Tensor,
+    target_dtype: torch.dtype,
+) -> torch.Tensor:
+    """Return ``(delta_bf16 + dequantize(weight_fp8)).to(target_dtype)``.
+    CUDA with Triton uses stochastic rounding; otherwise uses a deterministic bf16 add.
+    ``delta_bf16`` is the bf16 accumulator and is mutated in place.
+    """
+    if delta_bf16.dtype != torch.bfloat16:
+        raise ValueError(f"delta_bf16 must be bfloat16, got {delta_bf16.dtype}")
+    if str(weight_fp8.device).startswith("cuda") and TRITON_AVAILABLE:
+        fused_add_round_launch(delta_bf16, weight_fp8, seed=0)
+    else:
+        delta_bf16.add_(weight_fp8.to(dtype=torch.bfloat16))
+    return delta_bf16.to(dtype=target_dtype)
+
+
 def fuse_lora_weights(
     model_sd: StateDict,
     lora_sd_and_strengths: list[LoraStateDictWithStrength],
     dtype: torch.dtype | None = None,
+    preserve_input_device: bool = True,
 ) -> Iterator[tuple[str, torch.Tensor]]:
     """Yield ``(key, fused_tensor)`` for each weight modified by at least one LoRA.
     For scaled-FP8 weights, this includes both the updated ``.weight`` tensor
     and its corresponding ``.weight_scale`` tensor.
+    When ``preserve_input_device`` is False, fused tensors are yielded on the device
+    used for fusion; caller is responsible for moving them to their final
+    destination.
     """
     for key, original_weight in model_sd.sd.items():
         if original_weight is None or key.endswith(".weight_scale"):
@@ -30,7 +85,7 @@ def fuse_lora_weights(
         target_dtype = dtype if dtype is not None else weight.dtype
         deltas_dtype = target_dtype if target_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2] else torch.bfloat16
 
-        deltas = _prepare_deltas(lora_sd_and_strengths, key, deltas_dtype, weight.device)
+        deltas = _aggregate_deltas(lora_sd_and_strengths, key, deltas_dtype, weight.device)
         if deltas is None:
             continue
 
@@ -41,14 +96,15 @@ def fuse_lora_weights(
             if is_scaled_fp8:
                 fused = _fuse_delta_with_scaled_fp8(deltas, weight, key, scale_key, model_sd)
             else:
-                fused = _fuse_delta_with_cast_fp8(deltas, weight, key, target_dtype)
+                fused = {key: fuse_cast_fp8_weight(deltas, weight, target_dtype)}
         elif weight.dtype == torch.bfloat16:
-            fused = _fuse_delta_with_bfloat16(deltas, weight, key, target_dtype)
+            deltas.add_(weight)
+            fused = {key: deltas.to(dtype=target_dtype)}
         else:
             raise ValueError(f"Unsupported dtype: {weight.dtype}")
 
         for k, v in fused.items():
-            yield k, v.to(device=original_device)
+            yield k, v.to(device=original_device) if preserve_input_device else v
 
 
 def apply_loras(
@@ -57,10 +113,12 @@ def apply_loras(
     dtype: torch.dtype | None = None,
     destination_sd: StateDict | None = None,
 ) -> StateDict:
+    """Fuse LoRAs into ``model_sd`` and place the results in ``destination_sd``.
+    When ``destination_sd`` is provided, the fused tensors are placed directly into it.
+    """
     if destination_sd is not None:
-        sd = destination_sd.sd
+        for key, fused in fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype):
-        for key, tensor in fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype):
+            destination_sd.sd[key] = fused
-            sd[key] = tensor
         return destination_sd
 
     fused = dict(fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype))
@@ -68,26 +126,22 @@ def apply_loras(
     return StateDict(sd, model_sd.device, model_sd.size, model_sd.dtype)
 
 
-def _prepare_deltas(
+def _aggregate_deltas(
     lora_sd_and_strengths: list[LoraStateDictWithStrength], key: str, dtype: torch.dtype, device: torch.device
 ) -> torch.Tensor | None:
-    deltas = []
     prefix = key[: -len(".weight")]
     key_a = f"{prefix}.lora_A.weight"
     key_b = f"{prefix}.lora_B.weight"
-    for lsd, coef in lora_sd_and_strengths:
+
-        if key_a not in lsd.sd or key_b not in lsd.sd:
+    def _ab_products() -> Iterator[LoraProduct]:
-            continue
+        for lsd, coef in lora_sd_and_strengths:
-        a = lsd.sd[key_a].to(device=device)
+            if key_a not in lsd.sd or key_b not in lsd.sd:
-        b = lsd.sd[key_b].to(device=device)
+                continue
-        product = torch.matmul(b * coef, a)
+            a = lsd.sd[key_a].to(device=device, dtype=dtype, non_blocking=True)
-        del a, b
+            b = lsd.sd[key_b].to(device=device, dtype=dtype, non_blocking=True)
-        deltas.append(product.to(dtype=dtype))
+            yield LoraProduct(a, b, coef)
-    if len(deltas) == 0:
+
-        return None
+    return aggregate_lora_products(_ab_products(), dtype)
-    elif len(deltas) == 1:
-        return deltas[0]
-    return torch.sum(torch.stack(deltas, dim=0), dim=0)
 
 
 def _fuse_delta_with_scaled_fp8(
@@ -100,34 +154,9 @@ def _fuse_delta_with_scaled_fp8(
     """Dequantize scaled FP8 weight, add LoRA delta, and re-quantize."""
     weight_scale = model_sd.sd[scale_key]
 
-    original_weight = weight.t().to(torch.float32) * weight_scale
+    original_weight = weight.to(torch.float32) * weight_scale
 
     new_weight = original_weight + deltas.to(torch.float32)
 
     new_fp8_weight, new_weight_scale = quantize_weight_to_fp8_per_tensor(new_weight)
     return {key: new_fp8_weight, scale_key: new_weight_scale}
-
-
-def _fuse_delta_with_cast_fp8(
-    deltas: torch.Tensor,
-    weight: torch.Tensor,
-    key: str,
-    target_dtype: torch.dtype,
-) -> dict[str, torch.Tensor]:
-    """Fuse LoRA delta with cast-only FP8 weight (no scale factor)."""
-    if str(weight.device).startswith("cuda"):
-        _fused_add_round_launch(deltas, weight, seed=0)
-    else:
-        deltas.add_(weight.to(dtype=deltas.dtype))
-    return {key: deltas.to(dtype=target_dtype)}
-
-
-def _fuse_delta_with_bfloat16(
-    deltas: torch.Tensor,
-    weight: torch.Tensor,
-    key: str,
-    target_dtype: torch.dtype,
-) -> dict[str, torch.Tensor]:
-    """Fuse LoRA delta with bfloat16 weight."""
-    deltas.add_(weight)
-    return {key: deltas.to(dtype=target_dtype)}

packages/ltx-core/src/ltx_core/loader/helpers.py

@@ -0,0 +1,61 @@
+"""Shared model-construction helpers used by both SingleGPUModelBuilder and StreamingModelBuilder."""
+
+from __future__ import annotations
+
+from typing import TypeVar
+
+import torch
+from torch import nn
+
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.primitives import StateDict, StateDictLoader
+from ltx_core.loader.registry import Registry
+from ltx_core.loader.sd_ops import SDOps
+from ltx_core.model.model_protocol import ModelConfigurator
+
+_M = TypeVar("_M", bound=nn.Module)
+
+
+def load_state_dict(
+    paths: str | tuple[str, ...] | list[str],
+    loader: StateDictLoader,
+    registry: Registry,
+    device: torch.device | None,
+    sd_ops: SDOps | None = None,
+) -> StateDict:
+    """Load a state dict from disk, using registry caching."""
+    if isinstance(paths, str):
+        path_list = [paths]
+    elif isinstance(paths, tuple):
+        path_list = list(paths)
+    else:
+        path_list = paths
+    cached = registry.get(path_list, sd_ops)
+    if cached is not None:
+        return cached
+    result = loader.load(path_list, sd_ops=sd_ops, device=device)
+    registry.add(path_list, sd_ops=sd_ops, state_dict=result)
+    return result
+
+
+def read_model_config(
+    model_path: str | tuple[str, ...],
+    loader: StateDictLoader,
+) -> dict:
+    """Read metadata from the first shard of a checkpoint."""
+    first = model_path[0] if isinstance(model_path, tuple) else model_path
+    return loader.metadata(first)
+
+
+def create_meta_model(
+    configurator: type[ModelConfigurator[_M]],
+    config: dict,
+    module_ops: tuple[ModuleOps, ...] = (),
+) -> _M:
+    """Create a model on the meta device and apply module operations."""
+    with torch.device("meta"):
+        model = configurator.from_config(config)
+    for op in module_ops:
+        if op.matcher(model):
+            model = op.mutator(model)
+    return model

packages/ltx-core/src/ltx_core/loader/kernels.py

@@ -1,72 +1,79 @@
 # ruff: noqa: ANN001, ANN201, ERA001, N803, N806
-import triton
+try:
-import triton.language as tl
+    import triton
+    import triton.language as tl
+
+    TRITON_AVAILABLE = True
+except (ImportError, OSError):
+    TRITON_AVAILABLE = False
 
 
-@triton.jit
+if TRITON_AVAILABLE:
-def fused_add_round_kernel(
-    x_ptr,
-    output_ptr,  # contents will be added to the output
-    seed,
-    n_elements,
-    EXPONENT_BIAS,
-    MANTISSA_BITS,
-    BLOCK_SIZE: tl.constexpr,
-):
-    """
-    A kernel to upcast 8bit quantized weights to bfloat16 with stochastic rounding
-    and add them to bfloat16 output weights. Might be used to upcast original model weights
-    and to further add them to precalculated deltas coming from LoRAs.
-    """
-    # Get program ID and compute offsets
-    pid = tl.program_id(axis=0)
-    block_start = pid * BLOCK_SIZE
-    offsets = block_start + tl.arange(0, BLOCK_SIZE)
-    mask = offsets < n_elements
 
-    # Load data
+    @triton.jit
-    x = tl.load(x_ptr + offsets, mask=mask)
+    def fused_add_round_kernel(
-    rand_vals = tl.rand(seed, offsets) - 0.5
+        x_ptr,
+        output_ptr,  # contents will be added to the output
+        seed,
+        n_elements,
+        EXPONENT_BIAS,
+        MANTISSA_BITS,
+        BLOCK_SIZE: tl.constexpr,
+    ):
+        """
+        A kernel to upcast 8bit quantized weights to bfloat16 with stochastic rounding
+        and add them to bfloat16 output weights. Might be used to upcast original model weights
+        and to further add them to precalculated deltas coming from LoRAs.
+        """
+        # Get program ID and compute offsets
+        pid = tl.program_id(axis=0)
+        block_start = pid * BLOCK_SIZE
+        offsets = block_start + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_elements
 
-    x = tl.cast(x, tl.float16)
+        # Load data
-    delta = tl.load(output_ptr + offsets, mask=mask)
+        x = tl.load(x_ptr + offsets, mask=mask)
-    delta = tl.cast(delta, tl.float16)
+        rand_vals = tl.rand(seed, offsets) - 0.5
-    x = x + delta
 
-    x_bits = tl.cast(x, tl.int16, bitcast=True)
+        x = tl.cast(x, tl.float16)
+        delta = tl.load(output_ptr + offsets, mask=mask)
+        delta = tl.cast(delta, tl.float16)
+        x = x + delta
 
-    # Calculate the exponent. Unbiased fp16 exponent is ((x_bits & 0x7C00) >> 10) - 15 for
+        x_bits = tl.cast(x, tl.int16, bitcast=True)
-    # normal numbers and -14 for subnormals.
-    fp16_exponent_bits = (x_bits & 0x7C00) >> 10
-    fp16_normals = fp16_exponent_bits > 0
-    fp16_exponent = tl.where(fp16_normals, fp16_exponent_bits - 15, -14)
 
-    # Add the target dtype's exponent bias and clamp to the target dtype's exponent range.
+        # Calculate the exponent. Unbiased fp16 exponent is ((x_bits & 0x7C00) >> 10) - 15 for
-    exponent = fp16_exponent + EXPONENT_BIAS
+        # normal numbers and -14 for subnormals.
-    MAX_EXPONENT = 2 * EXPONENT_BIAS + 1
+        fp16_exponent_bits = (x_bits & 0x7C00) >> 10
-    exponent = tl.where(exponent > MAX_EXPONENT, MAX_EXPONENT, exponent)
+        fp16_normals = fp16_exponent_bits > 0
-    exponent = tl.where(exponent < 0, 0, exponent)
+        fp16_exponent = tl.where(fp16_normals, fp16_exponent_bits - 15, -14)
 
-    # Normal ULP exponent, expressed as an fp16 exponent field:
+        # Add the target dtype's exponent bias and clamp to the target dtype's exponent range.
-    # (exponent - EXPONENT_BIAS - MANTISSA_BITS) + 15
+        exponent = fp16_exponent + EXPONENT_BIAS
-    # Simplifies to: fp16_exponent - MANTISSA_BITS + 15
+        MAX_EXPONENT = 2 * EXPONENT_BIAS + 1
-    # See https://en.wikipedia.org/wiki/Unit_in_the_last_place
+        exponent = tl.where(exponent > MAX_EXPONENT, MAX_EXPONENT, exponent)
-    eps_exp = tl.maximum(0, tl.minimum(31, exponent - EXPONENT_BIAS - MANTISSA_BITS + 15))
+        exponent = tl.where(exponent < 0, 0, exponent)
 
-    # Calculate epsilon in the target dtype
+        # Normal ULP exponent, expressed as an fp16 exponent field:
-    eps_normal = tl.cast(tl.cast(eps_exp << 10, tl.int16), tl.float16, bitcast=True)
+        # (exponent - EXPONENT_BIAS - MANTISSA_BITS) + 15
+        # Simplifies to: fp16_exponent - MANTISSA_BITS + 15
+        # See https://en.wikipedia.org/wiki/Unit_in_the_last_place
+        eps_exp = tl.maximum(0, tl.minimum(31, exponent - EXPONENT_BIAS - MANTISSA_BITS + 15))
 
-    # Subnormal ULP: 2^(1 - EXPONENT_BIAS - MANTISSA_BITS) ->
+        # Calculate epsilon in the target dtype
-    # fp16 exponent bits: (1 - EXPONENT_BIAS - MANTISSA_BITS) + 15 =
+        eps_normal = tl.cast(tl.cast(eps_exp << 10, tl.int16), tl.float16, bitcast=True)
-    # 16 - EXPONENT_BIAS - MANTISSA_BITS
-    eps_subnormal = tl.cast(tl.cast((16 - EXPONENT_BIAS - MANTISSA_BITS) << 10, tl.int16), tl.float16, bitcast=True)
-    eps = tl.where(exponent > 0, eps_normal, eps_subnormal)
 
-    # Apply zero mask to epsilon
+        # Subnormal ULP: 2^(1 - EXPONENT_BIAS - MANTISSA_BITS) ->
-    eps = tl.where(x == 0, 0.0, eps)
+        # fp16 exponent bits: (1 - EXPONENT_BIAS - MANTISSA_BITS) + 15 =
+        # 16 - EXPONENT_BIAS - MANTISSA_BITS
+        eps_subnormal = tl.cast(tl.cast((16 - EXPONENT_BIAS - MANTISSA_BITS) << 10, tl.int16), tl.float16, bitcast=True)
+        eps = tl.where(exponent > 0, eps_normal, eps_subnormal)
 
-    # Apply stochastic rounding
+        # Apply zero mask to epsilon
-    output = tl.cast(x + rand_vals * eps, tl.bfloat16)
+        eps = tl.where(x == 0, 0.0, eps)
 
-    # Store the result
+        # Apply stochastic rounding
-    tl.store(output_ptr + offsets, output, mask=mask)
+        output = tl.cast(x + rand_vals * eps, tl.bfloat16)
+
+        # Store the result
+        tl.store(output_ptr + offsets, output, mask=mask)

packages/ltx-core/src/ltx_core/loader/primitives.py

@@ -13,6 +13,10 @@ if TYPE_CHECKING:
     from ltx_core.loader.registry import Registry
 
 
+# Per-key shape and dtype description for a flat collection of tensors.
+TensorLayout = dict[str, tuple[torch.Size, torch.dtype]]
+
+
 @dataclass(frozen=True)
 class StateDict:
     """
@@ -52,7 +56,15 @@ class StateDictLoader(Protocol):
         """
 
 
-class ModelBuilderProtocol(Protocol[ModelType]):
+class BuilderProtocol(Protocol[ModelType]):
+    """Protocol for model builders that produce a model via ``build()``."""
+
+    def build(
+        self, device: torch.device | None = None, dtype: torch.dtype | None = None, **kwargs: object
+    ) -> ModelType: ...
+
+
+class ModelBuilderProtocol(BuilderProtocol[ModelType], Protocol[ModelType]):
     """
     Protocol for building PyTorch models from configuration dictionaries.
     Implementations must provide:

packages/ltx-core/src/ltx_core/loader/single_gpu_model_builder.py

@@ -3,8 +3,10 @@ from dataclasses import dataclass, field, replace
 from typing import Generic
 
 import torch
+from torch import nn
 
 from ltx_core.loader.fuse_loras import apply_loras
+from ltx_core.loader.helpers import create_meta_model, load_state_dict, read_model_config
 from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.primitives import (
     LoRAAdaptableProtocol,
@@ -22,6 +24,56 @@ from ltx_core.model.model_protocol import ModelConfigurator, ModelType
 logger: logging.Logger = logging.getLogger(__name__)
 
 
+def _check_uninitialized(model: nn.Module) -> list[str]:
+    """Return names of any parameters/buffers still on meta device."""
+    names = []
+    for name, param in model.named_parameters():
+        if str(param.device) == "meta":
+            names.append(name)
+    for name, buf in model.named_buffers():
+        if str(buf.device) == "meta":
+            names.append(name)
+    return names
+
+
+def _load_model_weights(
+    meta_model: nn.Module,
+    model_path: str | tuple[str, ...],
+    loras: tuple[LoraPathStrengthAndSDOps, ...],
+    loader: StateDictLoader,
+    registry: Registry,
+    device: torch.device,
+    dtype: torch.dtype | None,
+    model_sd_ops: SDOps | None = None,
+    lora_load_device: torch.device | None = None,
+) -> None:
+    """Load base weights and fuse LoRAs into *meta_model* in-place."""
+    if lora_load_device is None:
+        lora_load_device = device
+
+    model_sd = load_state_dict(model_path, loader, registry, device, model_sd_ops)
+
+    lora_strengths = [lora.strength for lora in loras]
+    if not lora_strengths or (min(lora_strengths) == 0 and max(lora_strengths) == 0):
+        sd = model_sd.sd
+        if dtype is not None:
+            sd = {key: value.to(dtype=dtype) for key, value in model_sd.sd.items()}
+        meta_model.load_state_dict(sd, strict=False, assign=True)
+        return
+
+    lora_state_dicts = [load_state_dict([lora.path], loader, registry, lora_load_device, lora.sd_ops) for lora in loras]
+    lora_sd_and_strengths = [
+        LoraStateDictWithStrength(sd, strength) for sd, strength in zip(lora_state_dicts, lora_strengths, strict=True)
+    ]
+    final_sd = apply_loras(
+        model_sd=model_sd,
+        lora_sd_and_strengths=lora_sd_and_strengths,
+        dtype=dtype,
+        destination_sd=model_sd if isinstance(registry, DummyRegistry) else None,
+    )
+    meta_model.load_state_dict(final_sd.sd, strict=False, assign=True)
+
+
 @dataclass(frozen=True)
 class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType], LoRAAdaptableProtocol):
     """
@@ -50,7 +102,7 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
     registry: Registry = field(default_factory=DummyRegistry)
     lora_load_device: torch.device = field(default_factory=lambda: torch.device("cpu"))
 
-    def lora(self, lora_path: str, strength: float = 1.0, sd_ops: SDOps | None = None) -> "SingleGPUModelBuilder":
+    def lora(self, lora_path: str, strength: float, sd_ops: SDOps) -> "SingleGPUModelBuilder":
         return replace(self, loras=(*self.loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops)))
 
     def with_sd_ops(self, sd_ops: SDOps | None) -> "SingleGPUModelBuilder":
@@ -69,34 +121,22 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
         return replace(self, lora_load_device=device)
 
     def model_config(self) -> dict:
-        first_shard_path = self.model_path[0] if isinstance(self.model_path, tuple) else self.model_path
+        return read_model_config(self.model_path, self.model_loader)
-        return self.model_loader.metadata(first_shard_path)
 
     def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType:
-        with torch.device("meta"):
+        return create_meta_model(self.model_class_configurator, config, module_ops)
-            model = self.model_class_configurator.from_config(config)
-        for module_op in module_ops:
-            if module_op.matcher(model):
-                model = module_op.mutator(model)
-        return model
 
     def load_sd(
         self, paths: list[str], registry: Registry, device: torch.device | None, sd_ops: SDOps | None = None
     ) -> StateDict:
-        state_dict = registry.get(paths, sd_ops)
+        return load_state_dict(paths, self.model_loader, registry, device, sd_ops)
-        if state_dict is None:
-            state_dict = self.model_loader.load(paths, sd_ops=sd_ops, device=device)
-            registry.add(paths, sd_ops=sd_ops, state_dict=state_dict)
-        return state_dict
 
     def _return_model(self, meta_model: ModelType, device: torch.device) -> ModelType:
-        uninitialized_params = [name for name, param in meta_model.named_parameters() if str(param.device) == "meta"]
+        uninitialized = _check_uninitialized(meta_model)
-        uninitialized_buffers = [name for name, buffer in meta_model.named_buffers() if str(buffer.device) == "meta"]
+        if uninitialized:
-        if uninitialized_params or uninitialized_buffers:
+            logger.warning(f"Uninitialized parameters or buffers: {uninitialized}")
-            logger.warning(f"Uninitialized parameters or buffers: {uninitialized_params + uninitialized_buffers}")
             return meta_model
-        retval = meta_model.to(device)
+        return meta_model.to(device)
-        return retval
 
     def build(
         self,
@@ -107,30 +147,16 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
         device = torch.device("cuda") if device is None else device
         config = self.model_config()
         meta_model = self.meta_model(config, self.module_ops)
-        model_paths = list(self.model_path) if isinstance(self.model_path, tuple) else [self.model_path]
-        model_state_dict = self.load_sd(model_paths, sd_ops=self.model_sd_ops, registry=self.registry, device=device)
 
-        lora_strengths = [lora.strength for lora in self.loras]
+        _load_model_weights(
-        if not lora_strengths or (min(lora_strengths) == 0 and max(lora_strengths) == 0):
+            meta_model=meta_model,
-            sd = model_state_dict.sd
+            model_path=self.model_path,
-            if dtype is not None:
+            loras=self.loras,
-                sd = {key: value.to(dtype=dtype) for key, value in model_state_dict.sd.items()}
+            loader=self.model_loader,
-            meta_model.load_state_dict(sd, strict=False, assign=True)
+            registry=self.registry,
-            return self._return_model(meta_model, device)
+            device=device,
-
-        lora_state_dicts = [
-            self.load_sd([lora.path], sd_ops=lora.sd_ops, registry=self.registry, device=self.lora_load_device)
-            for lora in self.loras
-        ]
-        lora_sd_and_strengths = [
-            LoraStateDictWithStrength(sd, strength)
-            for sd, strength in zip(lora_state_dicts, lora_strengths, strict=True)
-        ]
-        final_sd = apply_loras(
-            model_sd=model_state_dict,
-            lora_sd_and_strengths=lora_sd_and_strengths,
             dtype=dtype,
-            destination_sd=model_state_dict if isinstance(self.registry, DummyRegistry) else None,
+            model_sd_ops=self.model_sd_ops,
+            lora_load_device=self.lora_load_device,
         )
-        meta_model.load_state_dict(final_sd.sd, strict=False, assign=True)
         return self._return_model(meta_model, device)

packages/ltx-core/src/ltx_core/model/transformer/attention.py

@@ -144,7 +144,7 @@ class Attention(torch.nn.Module):
         heads: int = 8,
         dim_head: int = 64,
         norm_eps: float = 1e-6,
-        rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+        rope_type: LTXRopeType = LTXRopeType.SPLIT,
         attention_function: AttentionCallable | AttentionFunction = AttentionFunction.DEFAULT,
         apply_gated_attention: bool = False,
     ) -> None:

packages/ltx-core/src/ltx_core/model/transformer/model.py

@@ -57,7 +57,7 @@ class LTXModel(torch.nn.Module):
         audio_cross_attention_dim: int = 2048,
         audio_positional_embedding_max_pos: list[int] | None = None,
         av_ca_timestep_scale_multiplier: int = 1,
-        rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+        rope_type: LTXRopeType = LTXRopeType.SPLIT,
         double_precision_rope: bool = False,
         apply_gated_attention: bool = False,
         caption_projection: torch.nn.Module | None = None,

packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py

@@ -62,7 +62,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
             audio_cross_attention_dim=config.get("audio_cross_attention_dim", 2048),
             audio_positional_embedding_max_pos=config.get("audio_positional_embedding_max_pos", [20]),
             av_ca_timestep_scale_multiplier=config.get("av_ca_timestep_scale_multiplier", 1),
-            rope_type=LTXRopeType(config.get("rope_type", "interleaved")),
+            rope_type=LTXRopeType(config.get("rope_type", "split")),
             double_precision_rope=config.get("frequencies_precision", False) == "float64",
             apply_gated_attention=config.get("apply_gated_attention", False),
             caption_projection=caption_projection,
@@ -114,7 +114,7 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
             positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]),
             timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
             use_middle_indices_grid=config.get("use_middle_indices_grid", True),
-            rope_type=LTXRopeType(config.get("rope_type", "interleaved")),
+            rope_type=LTXRopeType(config.get("rope_type", "split")),
             double_precision_rope=config.get("frequencies_precision", False) == "float64",
             apply_gated_attention=config.get("apply_gated_attention", False),
             caption_projection=caption_projection,

packages/ltx-core/src/ltx_core/model/transformer/rope.py

@@ -16,9 +16,10 @@ class LTXRopeType(Enum):
 def apply_rotary_emb(
     input_tensor: torch.Tensor,
     freqs_cis: Tuple[torch.Tensor, torch.Tensor],
-    rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+    rope_type: LTXRopeType = LTXRopeType.SPLIT,
 ) -> torch.Tensor:
     if rope_type == LTXRopeType.INTERLEAVED:
+        # Note: INTERLEAVED rope is a legacy mode. Prefer SPLIT instead.
         return apply_interleaved_rotary_emb(input_tensor, *freqs_cis)
     elif rope_type == LTXRopeType.SPLIT:
         return apply_split_rotary_emb(input_tensor, *freqs_cis)
@@ -45,6 +46,11 @@ def apply_split_rotary_emb(
     needs_reshape = False
     if input_tensor.ndim != 4 and cos_freqs.ndim == 4:
         b, h, t, _ = cos_freqs.shape
+        if input_tensor.shape[0] != b:
+            raise ValueError(
+                f"apply_split_rotary_emb: input_tensor batch ({input_tensor.shape[0]}) "
+                f"must equal cos_freqs batch ({b})."
+            )
         input_tensor = input_tensor.reshape(b, t, h, -1).swapaxes(1, 2)
         needs_reshape = True
 
@@ -183,7 +189,7 @@ def precompute_freqs_cis(
     max_pos: list[int] | None = None,
     use_middle_indices_grid: bool = False,
     num_attention_heads: int = 32,
-    rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+    rope_type: LTXRopeType = LTXRopeType.SPLIT,
     freq_grid_generator: Callable[[float, int, int, torch.device], torch.Tensor] = generate_freq_grid_pytorch,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     if max_pos is None:

packages/ltx-core/src/ltx_core/model/transformer/transformer.py

@@ -27,7 +27,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
         idx: int,
         video: TransformerConfig | None = None,
         audio: TransformerConfig | None = None,
-        rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+        rope_type: LTXRopeType = LTXRopeType.SPLIT,
         norm_eps: float = 1e-6,
         attention_function: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT,
     ):

packages/ltx-core/src/ltx_core/model/video_vae/__init__.py

@@ -1,5 +1,6 @@
 """Video VAE package."""
 
+from ltx_core.model.video_vae.memory_efficient_decode import MEMORY_EFFICIENT_DECODE
 from ltx_core.model.video_vae.model_configurator import (
     VAE_DECODER_COMFY_KEYS_FILTER,
     VAE_ENCODER_COMFY_KEYS_FILTER,
@@ -10,6 +11,7 @@ from ltx_core.model.video_vae.tiling import SpatialTilingConfig, TemporalTilingC
 from ltx_core.model.video_vae.video_vae import VideoDecoder, VideoEncoder, get_video_chunks_number
 
 __all__ = [
+    "MEMORY_EFFICIENT_DECODE",
     "VAE_DECODER_COMFY_KEYS_FILTER",
     "VAE_ENCODER_COMFY_KEYS_FILTER",
     "SpatialTilingConfig",

packages/ltx-core/src/ltx_core/model/video_vae/memory_efficient_decode.py

@@ -0,0 +1,612 @@
+"""Memory-efficient VAE decoder operations.
+Reduces peak VRAM usage during video decoding through in-place operations
+and workspace buffer reuse.  The main optimizations are:
+1. **Workspace buffers** -- Pre-allocated tensors with temporal padding replace
+   dynamic padding (``F.pad`` / ``concatenate``) in ``CausalConv3d``.  A
+   workspace of shape ``[B, C, T+2, H, W]`` holds the data in positions
+   ``[1:-1]`` with replicate padding at ``[0]`` and ``[-1]``.
+2. **In-place temporal-chunked Conv3d** *(non-causal only)* -- The convolution
+   output is written back into the workspace buffer, avoiding a separate
+   output allocation.  Temporal chunking with boundary save/restore ensures
+   correct reads despite in-place writes.
+3. **In-place normalization and affine transforms** -- PixelNorm, scale/shift,
+   and SiLU are applied in-place on workspace views.
+4. **Free-before-conv** -- For ``DepthToSpaceUpsample`` blocks the input
+   tensor is freed before the convolution runs so that peak VRAM never holds
+   input *and* output simultaneously.
+Both causal and non-causal modes are supported.  Non-causal mode benefits
+from all four optimizations.  Causal mode benefits from optimizations 1, 3,
+and 4; in-place conv (2) is skipped because the asymmetric causal padding
+layout prevents clean in-place overwrites.
+Usage via the ``ModuleOps`` pattern (preferred)::
+    from ltx_core.model.video_vae import MEMORY_EFFICIENT_DECODE
+    builder = decoder_builder.with_module_ops(
+        (*decoder_builder.module_ops, MEMORY_EFFICIENT_DECODE)
+    )
+Or applied directly to an existing decoder::
+    from ltx_core.model.video_vae.memory_efficient_decode import (
+        enable_memory_efficient_decode,
+    )
+    enable_memory_efficient_decode(decoder)
+"""
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING
+
+import torch
+from einops import rearrange
+from torch import nn
+from torch.nn import functional as F
+
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.model.common.normalization import PixelNorm
+from ltx_core.model.video_vae.convolution import CausalConv3d
+from ltx_core.model.video_vae.ops import unpatchify
+from ltx_core.model.video_vae.resnet import ResnetBlock3D, UNetMidBlock3D
+from ltx_core.model.video_vae.sampling import DepthToSpaceUpsample
+
+if TYPE_CHECKING:
+    from ltx_core.model.video_vae.video_vae import VideoDecoder
+
+
+# ---------------------------------------------------------------------------
+# Low-level helpers
+# ---------------------------------------------------------------------------
+
+
+def _find_temporal_split_size(num_frames: int) -> int:
+    """Find chunk size for in-place temporal convolution.
+    The chunk size ensures the last chunk has at least 3 frames
+    (the temporal kernel size), avoiding degenerate chunks.
+    """
+    for s in range(16, 2, -1):
+        remainder = num_frames % s
+        if remainder == 0 or remainder >= 3:
+            return s
+
+    raise ValueError(
+        f"Unable to find a valid temporal split size for num_frames={num_frames}. "
+        "Expected a split size between 3 and 16 such that the final chunk is "
+        "either exact or has at least 3 frames."
+    )
+
+
+def _pad_workspace_temporal(workspace: torch.Tensor) -> None:
+    """Apply non-causal replicate padding to temporal boundaries.
+    Sets ``workspace[:, :, 0]`` to a copy of ``workspace[:, :, 1]`` and
+    ``workspace[:, :, -1]`` to a copy of ``workspace[:, :, -2]``.
+    """
+    workspace[:, :, 0, :, :].copy_(workspace[:, :, 1, :, :])
+    workspace[:, :, -1, :, :].copy_(workspace[:, :, -2, :, :])
+
+
+# ---------------------------------------------------------------------------
+# In-place Conv3d (non-causal only)
+# ---------------------------------------------------------------------------
+
+
+def inplace_conv3d_temporal_chunked(workspace: torch.Tensor, conv: nn.Conv3d) -> None:
+    """Run a 3x3x3 Conv3d in-place on a temporally-padded workspace.
+    The workspace has shape ``[B, C, T+2, H, W]`` where positions ``[1:-1]``
+    hold the real data and positions ``[0]`` and ``[-1]`` are padding slots.
+    The convolution must have ``kernel_size=(3,3,3)``, ``stride=(1,1,1)``,
+    ``padding=(0,1,1)`` -- no temporal padding, symmetric spatial padding.
+    The output (T frames) overwrites positions ``[1:-1]``.  Temporal chunking
+    with boundary save/restore ensures each chunk reads unmodified input even
+    though earlier chunks already wrote to the same buffer.
+    Only valid for **non-causal** mode (symmetric replicate padding).
+    Args:
+        workspace: Tensor ``[B, max(C_in, C_out), T+2, H, W]``.
+            Modified in-place; after the call ``workspace[:, :C_out, 1:-1]``
+            holds the convolution result.
+        conv: ``nn.Conv3d`` with the constraints above.
+    """
+    if conv.kernel_size != (3, 3, 3):
+        raise ValueError(f"Expected kernel_size=(3,3,3), got {conv.kernel_size}")
+    if conv.stride != (1, 1, 1):
+        raise ValueError(f"Expected stride=(1,1,1), got {conv.stride}")
+    if conv.padding != (0, 1, 1):
+        raise ValueError(f"Expected padding=(0,1,1), got {conv.padding}")
+
+    _pad_workspace_temporal(workspace)
+
+    total_frames = workspace.shape[2]
+    out_channels = conv.out_channels
+    in_channels = conv.in_channels
+
+    if total_frames > 16:
+        split_size = _find_temporal_split_size(total_frames)
+        num_splits = (total_frames + split_size - 1) // split_size
+    else:
+        split_size = total_frames - 1
+        num_splits = 1
+
+    # 1-frame buffers for saving / restoring boundary frames across chunks.
+    x_buf = torch.empty(
+        workspace.shape[0],
+        workspace.shape[1],
+        1,
+        workspace.shape[3],
+        workspace.shape[4],
+        device=workspace.device,
+        dtype=workspace.dtype,
+    )
+    o_buf = torch.empty_like(x_buf)
+
+    # Helper: extract a chunk and make it contiguous.  Workspace views can
+    # inherit strides > 2^31 from the full buffer, which makes Conv3d's
+    # reflect-padding path (F.pad) crash with "input tensor must fit into
+    # 32-bit index math".  A small .clone() per chunk avoids this.
+    needs_clone = workspace.untyped_storage().nbytes() > (2**31 - 1) * workspace.element_size()
+
+    def _chunk(t_start: int, t_end: int) -> torch.Tensor:
+        s = workspace[:, :in_channels, t_start:t_end]
+        return s.clone() if needs_clone else s
+
+    # --- First chunk ---
+    if num_splits > 1:
+        # Save the boundary now so the loop below can restore it.  Skipped
+        # when there is only one chunk: the loop never runs, and the save
+        # would be a wasted full HW slice copy.
+        x_buf[:, :, 0] = workspace[:, :, split_size - 1].clone()
+    workspace[:, :out_channels, 1:split_size] = conv(_chunk(0, split_size + 1))
+
+    # --- Remaining chunks ---
+    for i in range(1, num_splits):
+        start = i * split_size
+        end = min((i + 1) * split_size, total_frames - 1)
+
+        # Save the value at start-1 (now holds previous chunk's output).
+        o_buf[:, :, 0] = workspace[:, :, start - 1].clone()
+        # Restore the original input value needed by this chunk's conv.
+        workspace[:, :, start - 1] = x_buf[:, :, 0]
+        # Save the boundary for the *next* chunk before we overwrite it.
+        x_buf[:, :, 0] = workspace[:, :, end - 1].clone()
+
+        workspace[:, :out_channels, start:end] = conv(_chunk(start - 1, end + 1))
+
+        # Put back the previous chunk's output at the boundary.
+        workspace[:, :, start - 1] = o_buf[:, :, 0]
+
+
+# ---------------------------------------------------------------------------
+# Causal conv helper (free-before-conv)
+# ---------------------------------------------------------------------------
+
+
+def _causal_pad(x: torch.Tensor, pad_size: int) -> torch.Tensor:
+    """Build a causal-padded buffer of shape ``[B, C, T+pad_size, H, W]``.
+    Copies ``x`` into ``padded[:, :, pad_size:]`` and replicates the first
+    real frame into the leading ``pad_size`` slots.  The caller still owns
+    ``x`` after this returns.
+    """
+    padded = torch.empty(
+        x.shape[0],
+        x.shape[1],
+        x.shape[2] + pad_size,
+        x.shape[3],
+        x.shape[4],
+        device=x.device,
+        dtype=x.dtype,
+    )
+    padded[:, :, pad_size:].copy_(x)
+    for i in range(pad_size):
+        padded[:, :, i] = padded[:, :, pad_size]
+    return padded
+
+
+def _causal_pad_free_and_conv(x: torch.Tensor, causal_conv: CausalConv3d) -> torch.Tensor:
+    """Causal-pad *x*, free it, then run the raw ``nn.Conv3d``.
+    This avoids the peak where both the original and padded tensors are
+    live simultaneously (as happens inside ``CausalConv3d.forward``).
+    Args:
+        x: Input ``[B, C_in, T, H, W]``.  **Deleted** inside this function;
+            the caller must not use it afterwards.
+    Returns:
+        Convolution output ``[B, C_out, T, H, W]``.
+    """
+    padded = _causal_pad(x, causal_conv.time_kernel_size - 1)
+    del x
+    result = causal_conv.conv(padded)
+    del padded
+    return result
+
+
+# ---------------------------------------------------------------------------
+# In-place normalization
+# ---------------------------------------------------------------------------
+
+
+def _pixel_norm_inplace(x: torch.Tensor, eps: float = 1e-8) -> None:
+    """In-place RMS (pixel) normalization along the channel dimension."""
+    rms = torch.sqrt(torch.mean(x**2, dim=1, keepdim=True) + eps)
+    x.div_(rms)
+
+
+def _norm_inplace(norm: nn.Module, x: torch.Tensor) -> None:
+    """Apply *norm* in-place, using an optimised path for ``PixelNorm``."""
+    if isinstance(norm, PixelNorm):
+        _pixel_norm_inplace(x, eps=norm.eps)
+    else:
+        # GroupNorm or other -- fall back to allocating a temporary.
+        result = norm(x)
+        x.copy_(result)
+        del result
+
+
+# ---------------------------------------------------------------------------
+# Per-block efficient forwards
+# ---------------------------------------------------------------------------
+
+
+def _resnet_block_forward_inplace(
+    resnet: ResnetBlock3D,
+    workspace: torch.Tensor,
+    causal: bool,
+    timestep: torch.Tensor | None,
+    generator: torch.Generator | None,
+) -> None:
+    """Run a ``ResnetBlock3D`` in-place on a workspace buffer.
+    The workspace has shape ``[B, C, T+2, H, W]`` with real data in
+    ``[1:-1]``.  After this call ``workspace[:, :, 1:-1]`` holds the
+    residual-branch output ``F(x)`` (without the skip connection --
+    the caller adds it back to the hidden state).
+    Only valid when ``in_channels == out_channels`` (true for all
+    ``ResnetBlock3D`` instances inside a ``UNetMidBlock3D``).
+    """
+    if resnet.in_channels != resnet.out_channels:
+        raise ValueError(
+            "In-place resnet forward requires in_channels == out_channels, "
+            f"got {resnet.in_channels} != {resnet.out_channels}"
+        )
+
+    interior = workspace[:, :, 1:-1]
+
+    # --- norm1 + [ada scaling] + SiLU + conv1 ---
+    _norm_inplace(resnet.norm1, interior)
+
+    if resnet.timestep_conditioning and timestep is not None:
+        ada = resnet.scale_shift_table[None, ..., None, None, None].to(
+            device=interior.device, dtype=interior.dtype
+        ) + timestep.reshape(
+            interior.shape[0],
+            4,
+            -1,
+            timestep.shape[-3],
+            timestep.shape[-2],
+            timestep.shape[-1],
+        )
+        shift1, scale1, shift2, scale2 = ada.unbind(dim=1)
+        interior.mul_(1 + scale1).add_(shift1)
+
+    F.silu(interior, inplace=True)
+
+    if causal:
+        result = resnet.conv1(interior, causal=True)
+        interior.copy_(result)
+        del result
+    else:
+        inplace_conv3d_temporal_chunked(workspace, resnet.conv1.conv)
+
+    if resnet.inject_noise:
+        spatial_shape = interior.shape[-2:]
+        scale = resnet.per_channel_scale1.to(device=interior.device, dtype=interior.dtype)
+        noise = torch.randn(spatial_shape, device=interior.device, dtype=interior.dtype, generator=generator)
+        interior.add_((noise * scale)[None, :, None, ...])
+
+    # --- norm2 + [ada scaling] + SiLU + conv2 ---
+    _norm_inplace(resnet.norm2, interior)
+
+    if resnet.timestep_conditioning and timestep is not None:
+        interior.mul_(1 + scale2).add_(shift2)  # type: ignore[possibly-undefined]
+
+    F.silu(interior, inplace=True)
+    # dropout is always 0.0 during inference -- skip.
+
+    if causal:
+        result = resnet.conv2(interior, causal=True)
+        interior.copy_(result)
+        del result
+    else:
+        inplace_conv3d_temporal_chunked(workspace, resnet.conv2.conv)
+
+    if resnet.inject_noise:
+        spatial_shape = interior.shape[-2:]
+        scale = resnet.per_channel_scale2.to(device=interior.device, dtype=interior.dtype)
+        noise = torch.randn(spatial_shape, device=interior.device, dtype=interior.dtype, generator=generator)
+        interior.add_((noise * scale)[None, :, None, ...])
+
+
+def _midblock_forward_efficient(
+    block: UNetMidBlock3D,
+    hidden_states: torch.Tensor,
+    causal: bool,
+    timestep: torch.Tensor | None,
+    generator: torch.Generator | None,
+) -> torch.Tensor:
+    """Memory-efficient ``UNetMidBlock3D`` forward.
+    Allocates a single workspace buffer that is reused across all
+    ``ResnetBlock3D`` iterations.  For each block the workspace is
+    populated with the current hidden state, processed in-place, and
+    the result is added back (residual connection).
+    """
+    timestep_embed = None
+    if block.timestep_conditioning:
+        if timestep is None:
+            raise ValueError("'timestep' required when timestep_conditioning=True")
+        batch_size = hidden_states.shape[0]
+        timestep_embed = block.time_embedder(
+            timestep=timestep.flatten(),
+            hidden_dtype=hidden_states.dtype,
+        )
+        timestep_embed = timestep_embed.view(batch_size, timestep_embed.shape[-1], 1, 1, 1)
+
+    workspace = torch.empty(
+        hidden_states.shape[0],
+        hidden_states.shape[1],
+        hidden_states.shape[2] + 2,
+        hidden_states.shape[3],
+        hidden_states.shape[4],
+        device=hidden_states.device,
+        dtype=hidden_states.dtype,
+    )
+
+    for resnet in block.res_blocks:
+        workspace[:, :, 1:-1].copy_(hidden_states)
+        _resnet_block_forward_inplace(resnet, workspace, causal, timestep_embed, generator)
+        hidden_states.add_(workspace[:, :, 1:-1])
+
+    del workspace
+    return hidden_states
+
+
+def _upsample_forward_efficient(
+    block: DepthToSpaceUpsample,
+    x: torch.Tensor,
+    causal: bool,
+) -> torch.Tensor:
+    """Memory-efficient ``DepthToSpaceUpsample`` forward.
+    For non-causal mode the input is copied into a workspace and the
+    convolution runs in-place.  For causal mode the input is manually
+    padded and freed before the convolution runs.  Both paths avoid
+    the peak where input *and* output coexist.
+    """
+    if block.residual:
+        x_in = rearrange(
+            x,
+            "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
+            p1=block.stride[0],
+            p2=block.stride[1],
+            p3=block.stride[2],
+        )
+        num_repeat = math.prod(block.stride) // block.out_channels_reduction_factor
+        x_in = x_in.repeat(1, num_repeat, 1, 1, 1)
+        if block.stride[0] == 2:
+            x_in = x_in[:, :, 1:, :, :]
+
+    conv = block.conv.conv  # underlying nn.Conv3d inside CausalConv3d
+    in_channels = x.shape[1]
+    out_channels = conv.out_channels
+
+    if causal:
+        x = _causal_pad_free_and_conv(x, block.conv)
+    else:
+        workspace = torch.empty(
+            x.shape[0],
+            max(in_channels, out_channels),
+            x.shape[2] + 2,
+            x.shape[3],
+            x.shape[4],
+            device=x.device,
+            dtype=x.dtype,
+        )
+        workspace[:, :in_channels, 1:-1].copy_(x)
+        del x
+        inplace_conv3d_temporal_chunked(workspace, conv)
+        x = workspace[:, :out_channels, 1:-1].contiguous()
+        del workspace
+
+    x = rearrange(
+        x,
+        "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
+        p1=block.stride[0],
+        p2=block.stride[1],
+        p3=block.stride[2],
+    )
+    if block.stride[0] == 2:
+        x = x[:, :, 1:, :, :]
+    if block.residual:
+        x = x + x_in
+        del x_in
+    return x
+
+
+# ---------------------------------------------------------------------------
+# Final norm + conv_out
+# ---------------------------------------------------------------------------
+
+
+def _final_norm_and_conv_out(
+    decoder: VideoDecoder,
+    sample: torch.Tensor,
+    causal: bool,
+    scaled_timestep: torch.Tensor | None,
+    batch_size: int,
+) -> torch.Tensor:
+    """Workspace-based final norm + [ada] + SiLU + conv_out + unpatchify."""
+    conv_out_mod: CausalConv3d = decoder.conv_out  # type: ignore[assignment]
+    conv_out = conv_out_mod.conv
+    feature_channels = sample.shape[1]
+
+    workspace = torch.empty(
+        sample.shape[0],
+        max(feature_channels, conv_out.out_channels),
+        sample.shape[2] + 2,
+        sample.shape[3],
+        sample.shape[4],
+        device=sample.device,
+        dtype=sample.dtype,
+    )
+    workspace[:, :feature_channels, 1:-1].copy_(sample)
+    del sample
+
+    interior = workspace[:, :feature_channels, 1:-1]
+    _norm_inplace(decoder.conv_norm_out, interior)
+
+    if decoder.timestep_conditioning:
+        embedded_timestep = decoder.last_time_embedder(
+            timestep=scaled_timestep.flatten(),
+            hidden_dtype=interior.dtype,
+        )
+        embedded_timestep = embedded_timestep.view(batch_size, embedded_timestep.shape[-1], 1, 1, 1)
+        ada_values = decoder.last_scale_shift_table[None, ..., None, None, None].to(
+            device=interior.device, dtype=interior.dtype
+        ) + embedded_timestep.reshape(
+            batch_size,
+            2,
+            -1,
+            embedded_timestep.shape[-3],
+            embedded_timestep.shape[-2],
+            embedded_timestep.shape[-1],
+        )
+        shift, scale = ada_values.unbind(dim=1)
+        interior.mul_(1 + scale).add_(shift)
+
+    F.silu(interior, inplace=True)
+
+    if causal:
+        # Causal: build padded tensor directly from the interior view,
+        # then free the workspace before running the conv.
+        padded = _causal_pad(interior, conv_out_mod.time_kernel_size - 1)
+        del workspace, interior
+        result = conv_out(padded)
+        del padded
+    else:
+        inplace_conv3d_temporal_chunked(workspace, conv_out)
+        result = workspace[:, : conv_out.out_channels, 1:-1].contiguous()
+        del workspace, interior
+
+    return unpatchify(result, patch_size_hw=decoder.patch_size, patch_size_t=1)
+
+
+# ---------------------------------------------------------------------------
+# Top-level efficient decoder forward
+# ---------------------------------------------------------------------------
+
+
+def _memory_efficient_forward(
+    decoder: VideoDecoder,
+    sample: torch.Tensor,
+    timestep: torch.Tensor | None = None,
+    generator: torch.Generator | None = None,
+) -> torch.Tensor:
+    """Full memory-efficient ``VideoDecoder.forward`` replacement.
+    Orchestrates the entire decode through workspace-based operations:
+    ``UNetMidBlock3D`` and ``DepthToSpaceUpsample`` blocks use efficient
+    paths; standalone ``ResnetBlock3D`` blocks fall back to the standard
+    forward.  The final norm + ada + SiLU + conv_out is also workspace-based.
+    """
+    causal = decoder.causal
+    batch_size = sample.shape[0]
+    sample = sample.to(next(decoder.parameters()).dtype)
+
+    # --- Noise injection and de-normalisation (identical to standard path) ---
+    if decoder.timestep_conditioning:
+        noise = (
+            torch.randn(sample.size(), generator=generator, dtype=sample.dtype, device=sample.device)
+            * decoder.decode_noise_scale
+        )
+        sample = noise + (1.0 - decoder.decode_noise_scale) * sample
+
+    sample = decoder.per_channel_statistics.un_normalize(sample)
+
+    if timestep is None and decoder.timestep_conditioning:
+        timestep = torch.full((batch_size,), decoder.decode_timestep, device=sample.device, dtype=sample.dtype)
+
+    # --- conv_in (latent tensor is small -- standard path is fine) ---
+    sample = decoder.conv_in(sample, causal=causal)
+
+    upscale_dtype = next(iter(decoder.up_blocks.parameters())).dtype
+    sample = sample.to(upscale_dtype)
+
+    scaled_timestep = None
+    if decoder.timestep_conditioning:
+        if timestep is None:
+            raise ValueError("'timestep' required when timestep_conditioning=True")
+        scaled_timestep = timestep * decoder.timestep_scale_multiplier.to(sample)
+
+    # --- Up blocks (dispatch to efficient path per block type) ---
+    for up_block in decoder.up_blocks:
+        if isinstance(up_block, UNetMidBlock3D):
+            sample = _midblock_forward_efficient(
+                up_block,
+                sample,
+                causal=causal,
+                timestep=scaled_timestep if decoder.timestep_conditioning else None,
+                generator=generator,
+            )
+        elif isinstance(up_block, DepthToSpaceUpsample):
+            sample = _upsample_forward_efficient(up_block, sample, causal=causal)
+        elif isinstance(up_block, ResnetBlock3D):
+            sample = up_block(sample, causal=causal, generator=generator)
+        else:
+            sample = up_block(sample, causal=causal)
+
+    return _final_norm_and_conv_out(decoder, sample, causal, scaled_timestep, batch_size)
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def enable_memory_efficient_decode(decoder: nn.Module) -> nn.Module:
+    """Patch a ``VideoDecoder`` to use the memory-efficient forward path.
+    The original ``forward`` is saved as ``decoder._original_forward`` so
+    that it can be restored later with :func:`disable_memory_efficient_decode`.
+    """
+    # Import here to avoid circular dependency at module level.
+    from ltx_core.model.video_vae.video_vae import VideoDecoder  # noqa: PLC0415
+
+    if not isinstance(decoder, VideoDecoder):
+        raise TypeError(f"Expected VideoDecoder, got {type(decoder).__name__}")
+
+    if hasattr(decoder, "_original_forward"):
+        return decoder
+
+    original_forward = decoder.forward
+
+    def efficient_forward(
+        sample: torch.Tensor,
+        timestep: torch.Tensor | None = None,
+        generator: torch.Generator | None = None,
+    ) -> torch.Tensor:
+        return _memory_efficient_forward(decoder, sample, timestep, generator)
+
+    decoder._original_forward = original_forward  # type: ignore[attr-defined]
+    decoder.forward = efficient_forward  # type: ignore[assignment]
+    return decoder
+
+
+def disable_memory_efficient_decode(decoder: nn.Module) -> nn.Module:
+    """Restore the original ``forward`` method on a patched ``VideoDecoder``."""
+    if hasattr(decoder, "_original_forward"):
+        decoder.forward = decoder._original_forward  # type: ignore[attr-defined]
+        del decoder._original_forward  # type: ignore[attr-defined]
+    return decoder
+
+
+def _is_video_decoder(model: nn.Module) -> bool:
+    """Matcher for the ``MEMORY_EFFICIENT_DECODE`` module op."""
+    from ltx_core.model.video_vae.video_vae import VideoDecoder  # noqa: PLC0415
+
+    return isinstance(model, VideoDecoder)
+
+
+MEMORY_EFFICIENT_DECODE = ModuleOps(
+    name="memory_efficient_vae_decode",
+    matcher=_is_video_decoder,
+    mutator=enable_memory_efficient_decode,
+)

packages/ltx-core/src/ltx_core/model/video_vae/tiling.py

@@ -64,6 +64,6 @@ class TilingConfig:
     @classmethod
     def default(cls) -> "TilingConfig":
         return cls(
-            spatial_config=SpatialTilingConfig(tile_size_in_pixels=512, tile_overlap_in_pixels=64),
+            spatial_config=SpatialTilingConfig(tile_size_in_pixels=768, tile_overlap_in_pixels=64),
-            temporal_config=TemporalTilingConfig(tile_size_in_frames=64, tile_overlap_in_frames=24),
+            temporal_config=TemporalTilingConfig(tile_size_in_frames=80, tile_overlap_in_frames=24),
         )

packages/ltx-core/src/ltx_core/model/video_vae/video_vae.py

@@ -259,6 +259,7 @@ class VideoEncoder(nn.Module):
         Args:
             sample: Input video (B, C, F, H, W). F should be 1 + 8*k (e.g., 1, 9, 17, 25, 33...).
                 If not, the encoder crops the last frames to the nearest valid length.
+                Should be normalized to [-1, 1] range before encoding.
         Returns:
             Normalized latent means (B, 128, F', H', W') where F' = 1+(F-1)/8, H' = H/32, W' = W/32.
             Example: (B, 3, 33, 512, 512) -> (B, 128, 5, 16, 16).
@@ -605,8 +606,8 @@ class VideoDecoder(nn.Module):
         # many video frames and pixels correspond to a single latent cell.
         self.video_downscale_factors = SpatioTemporalScaleFactors(
             time=8,
-            width=32,
             height=32,
+            width=32,
         )
 
         self.patch_size = patch_size
@@ -698,6 +699,9 @@ class VideoDecoder(nn.Module):
             When causal=False, allows future frame dependencies in convolutions but maintains same output shape.
         """
         batch_size = sample.shape[0]
+        output_dtype = sample.dtype
+        weights_dtype = next(self.parameters()).dtype
+        sample = sample.to(weights_dtype)
 
         # Add noise if timestep conditioning is enabled
         if self.timestep_conditioning:
@@ -770,7 +774,7 @@ class VideoDecoder(nn.Module):
         # Example: (B, 48, F, 128, 128) -> (B, 3, F, 512, 512) with patch_size=4
         sample = unpatchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)
 
-        return sample
+        return sample.to(output_dtype)
 
     def _prepare_tiles(
         self,
@@ -903,22 +907,21 @@ class VideoDecoder(nn.Module):
         tiling_config: TilingConfig | None = None,
         generator: torch.Generator | None = None,
     ) -> Iterator[torch.Tensor]:
-        """Decode a video latent tensor, yielding uint8 chunks ``[f, h, w, c]``.
+        """Decode a video latent tensor, yielding float chunks ``[f, h, w, c]`` in ``[0, 1]``.
         Subclasses (e.g. ``DistributedVideoDecoder``) may override this to
         control eagerness or distribution across ranks.
         """
 
-        def convert_to_uint8(frames: torch.Tensor) -> torch.Tensor:
+        def to_rgb(frames: torch.Tensor) -> torch.Tensor:
-            frames = (((frames + 1.0) / 2.0).clamp(0.0, 1.0) * 255.0).to(torch.uint8)
+            video = rearrange(frames[0], "c f h w -> f h w c")
-            frames = rearrange(frames[0], "c f h w -> f h w c")
+            return video.add_(1.0).mul_(0.5).clamp_(0.0, 1.0)
-            return frames
 
         if tiling_config is not None:
             for frames in self.tiled_decode(latent, tiling_config, generator=generator):
-                yield convert_to_uint8(frames)
+                yield to_rgb(frames)
         else:
             decoded = self(latent, generator=generator)
-            yield convert_to_uint8(decoded)
+            yield to_rgb(decoded)
 
     def _group_tiles_by_temporal_slice(self, tiles: List[Tile]) -> List[List[Tile]]:
         """Group tiles by their temporal output slice."""

packages/ltx-core/src/ltx_core/quantization/__init__.py

@@ -3,12 +3,9 @@ from ltx_core.quantization.fp8_cast import (
     UPCAST_DURING_INFERENCE,
     UpcastWithStochasticRounding,
 )
-from ltx_core.quantization.fp8_scaled_mm import FP8_PREPARE_MODULE_OPS, FP8_TRANSPOSE_SD_OPS
 from ltx_core.quantization.policy import QuantizationPolicy
 
 __all__ = [
-    "FP8_PREPARE_MODULE_OPS",
-    "FP8_TRANSPOSE_SD_OPS",
     "TRANSFORMER_LINEAR_DOWNCAST_MAP",
     "UPCAST_DURING_INFERENCE",
     "QuantizationPolicy",

packages/ltx-core/src/ltx_core/quantization/fp8_cast.py

@@ -1,5 +1,6 @@
 import torch
 
+from ltx_core.loader.kernels import TRITON_AVAILABLE
 from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
 from ltx_core.model.transformer.model import LTXModel
@@ -7,8 +8,13 @@ from ltx_core.model.transformer.model import LTXModel
 BLOCK_SIZE = 1024
 
 
-def _fused_add_round_launch(target_weight: torch.Tensor, original_weight: torch.Tensor, seed: int) -> torch.Tensor:
+def fused_add_round_launch(target_weight: torch.Tensor, original_weight: torch.Tensor, seed: int) -> torch.Tensor:
-    # Lazy import triton - only available on CUDA platforms
+    if not TRITON_AVAILABLE:
+        raise RuntimeError(
+            "fused_add_round_launch requires Triton, which is not available on this platform. "
+            "Callers should gate on ltx_core.loader.kernels.TRITON_AVAILABLE and use a "
+            "deterministic-rounding fallback instead."
+        )
     import triton  # noqa: PLC0415
 
     from ltx_core.loader.kernels import fused_add_round_kernel  # noqa: PLC0415
@@ -53,10 +59,13 @@ def _upcast_and_round(
     """
     Upcast the weight to the given dtype and optionally apply stochastic rounding.
     Input weight needs to have float8_e4m3fn or float8_e5m2 dtype.
+    Stochastic rounding is implemented via a Triton kernel. When Triton is not
+    available (e.g., on Windows), this falls back to deterministic (nearest)
+    rounding via ``weight.to(dtype)``.
     """
-    if not with_stochastic_rounding:
+    if not with_stochastic_rounding or not TRITON_AVAILABLE or weight.device.type != "cuda":
         return weight.to(dtype)
-    return _fused_add_round_launch(torch.zeros_like(weight, dtype=dtype), weight, seed)
+    return fused_add_round_launch(torch.zeros_like(weight, dtype=dtype), weight, seed)
 
 
 class Fp8CastLinear(torch.nn.Linear):
@@ -82,14 +91,26 @@ class Fp8CastLinear(torch.nn.Linear):
 def _replace_fwd_with_upcast(layer: torch.nn.Linear, with_stochastic_rounding: bool = False, seed: int = 0) -> None:
     """
     Intended to be applied via __class__ reassignment to existing nn.Linear
-    instances so that their parameter and buffer tensors are preserved in-place,
+    instances. Forward remains defined at the class level, which is required for
-    avoiding re-instantiation. Forward remains defined at the class level, which
+    torch.compile compatibility — instance-level closure monkey-patches cause
-    is required for torch.compile compatibility — instance-level closure
+    graph breaks.
-    monkey-patches cause graph breaks.
+    Also retypes ``weight`` and ``bias`` to fp8 so the meta param dtype matches
+    the post-load tensor dtype (sd_ops downcasts checkpoint bf16 -> fp8 at load).
+    Block streaming relies on this to derive pool buffer layout from the meta
+    model without an eager checkpoint read.
     """
     layer.__class__ = Fp8CastLinear
     layer._with_stochastic_rounding = with_stochastic_rounding
     layer._seed = seed
+    layer.weight = torch.nn.Parameter(
+        torch.empty(layer.weight.shape, dtype=torch.float8_e4m3fn, device=layer.weight.device),
+        requires_grad=layer.weight.requires_grad,
+    )
+    if layer.bias is not None:
+        layer.bias = torch.nn.Parameter(
+            torch.empty(layer.bias.shape, dtype=torch.float8_e4m3fn, device=layer.bias.device),
+            requires_grad=layer.bias.requires_grad,
+        )
 
 
 def _amend_forward_with_upcast(

packages/ltx-core/src/ltx_core/quantization/fp8_scaled_mm.py

@@ -1,11 +1,21 @@
+import json
+import struct
 from typing import Callable
 
 import torch
 from torch import nn
 
 from ltx_core.loader.module_ops import ModuleOps
-from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
 from ltx_core.model.transformer import LTXModel
+from ltx_core.quantization.trtllm_scaled_usable import trtllm_scaled_mm_usable
+
+
+def _read_safetensors_dtypes(path: str) -> dict[str, str]:
+    """Return ``{tensor_name: dtype_string}`` from the safetensors header."""
+    with open(path, "rb") as f:
+        header_size = struct.unpack("<Q", f.read(8))[0]
+        header = json.loads(f.read(header_size).decode("utf-8"))
+    return {k: v["dtype"] for k, v in header.items() if k != "__metadata__"}
 
 
 class FP8Linear(nn.Module):
@@ -25,11 +35,8 @@ class FP8Linear(nn.Module):
         self.in_features = in_features
         self.out_features = out_features
 
-        fp8_shape = (in_features, out_features)
+        self.weight = nn.Parameter(torch.empty((out_features, in_features), dtype=torch.float8_e4m3fn, device=device))
-        self.weight = nn.Parameter(torch.empty(fp8_shape, dtype=torch.float8_e4m3fn, device=device))
-        # Weight scale for FP8 dequantization (shape matches checkpoint format)
         self.weight_scale = nn.Parameter(torch.empty((), dtype=torch.float32, device=device))
-        # Input scale for static quantization (pre-quantized checkpoints)
         self.input_scale = nn.Parameter(torch.empty((), dtype=torch.float32, device=device))
 
         if bias:
@@ -40,31 +47,38 @@ class FP8Linear(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         origin_shape = x.shape
 
-        # Static quantization: use pre-computed scale
+        if trtllm_scaled_mm_usable():
-        qinput, cur_input_scale = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x, self.input_scale)
+            qinput, cur_input_scale = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x, self.input_scale)
+            if qinput.dim() == 3:
+                qinput = qinput.reshape(-1, qinput.shape[-1])
+            output = torch.ops.trtllm.cublas_scaled_mm(
+                qinput,
+                self.weight.t(),
+                scale_a=cur_input_scale,
+                scale_b=self.weight_scale,
+                bias=None,
+                out_dtype=x.dtype,
+            )
+        else:
+            # Clamp before cast: out-of-range values cast to NaN/saturated FP8, which
+            # produces black-screen output on some checkpoints (e.g. ltx-2-19b-dev-fp8).
+            fp8_min = torch.finfo(torch.float8_e4m3fn).min
+            fp8_max = torch.finfo(torch.float8_e4m3fn).max
+            qinput = torch.clamp(x * self.input_scale.reciprocal(), fp8_min, fp8_max).to(torch.float8_e4m3fn)
+            if qinput.dim() == 3:
+                qinput = qinput.reshape(-1, qinput.shape[-1])
+            output = torch._scaled_mm(
+                qinput,
+                self.weight.t(),
+                scale_a=self.input_scale,
+                scale_b=self.weight_scale,
+                out_dtype=x.dtype,
+                use_fast_accum=True,
+            )
 
-        # Flatten to 2D for matmul
-        if qinput.dim() == 3:
-            qinput = qinput.reshape(-1, qinput.shape[-1])
-
-        # FP8 scaled matmul
-        output = torch.ops.trtllm.cublas_scaled_mm(
-            qinput,
-            self.weight,
-            scale_a=cur_input_scale,
-            scale_b=self.weight_scale,
-            bias=None,
-            out_dtype=x.dtype,
-        )
-
-        # Add bias
         if self.bias is not None:
-            bias = self.bias
+            output = output + self.bias.to(output.dtype)
-            if bias.dtype != output.dtype:
-                bias = bias.to(output.dtype)
-            output = output + bias
 
-        # Restore original shape
         if output.dim() != len(origin_shape):
             output_shape = list(origin_shape)
             output_shape[-1] = output.shape[-1]
@@ -74,15 +88,7 @@ class FP8Linear(nn.Module):
 
 
 def quantize_weight_to_fp8_per_tensor(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    """
+    """Quantize a weight tensor to ``float8_e4m3fn`` with a per-tensor scale."""
-    Quantize a weight tensor to FP8 (float8_e4m3fn) using per-tensor scaling.
-    Args:
-        weight: The weight tensor to quantize (any dtype, will be cast to float32)
-    Returns:
-        Tuple of (quantized_weight, weight_scale):
-        - quantized_weight: FP8 tensor, transposed for cublas_scaled_mm
-        - weight_scale: Per-tensor scale factor (reciprocal of quantization scale)
-    """
     weight_fp32 = weight.to(torch.float32)
 
     fp8_min = torch.finfo(torch.float8_e4m3fn).min
@@ -96,7 +102,6 @@ def quantize_weight_to_fp8_per_tensor(weight: torch.Tensor) -> tuple[torch.Tenso
         weight_fp32: torch.Tensor, scale: torch.Tensor, fp8_min: torch.Tensor, fp8_max: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor]:
         quantized_weight = torch.clamp(weight_fp32 * scale, min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
-        quantized_weight = quantized_weight.t()
         weight_scale = scale.reciprocal()
         return quantized_weight, weight_scale
 
@@ -104,36 +109,8 @@ def quantize_weight_to_fp8_per_tensor(weight: torch.Tensor) -> tuple[torch.Tenso
     return quantized_weight, weight_scale
 
 
-def _should_skip_layer(layer_name: str, excluded_layer_substrings: tuple[str, ...]) -> bool:
-    return any(substring in layer_name for substring in excluded_layer_substrings)
-
-
-EXCLUDED_LAYER_SUBSTRINGS = (
-    "patchify_proj",
-    "adaln_single",
-    "av_ca_video_scale_shift_adaln_single",
-    "av_ca_a2v_gate_adaln_single",
-    "caption_projection",
-    "proj_out",
-    "audio_patchify_proj",
-    "audio_adaln_single",
-    "av_ca_audio_scale_shift_adaln_single",
-    "av_ca_v2a_gate_adaln_single",
-    "audio_caption_projection",
-    "audio_proj_out",
-    "transformer_blocks.0.",
-    *[f"transformer_blocks.{i}." for i in range(43, 48)],
-)
-
-
 def _linear_to_fp8linear(layer: nn.Linear) -> FP8Linear:
-    """
+    """Create an ``FP8Linear`` matching the shape/bias of *layer*."""
-    Create an FP8Linear layer from an nn.Linear layer.
-    Args:
-        layer: The nn.Linear layer to convert (typically on meta device)
-    Returns:
-        A new FP8Linear with the same configuration
-    """
     return FP8Linear(
         in_features=layer.in_features,
         out_features=layer.out_features,
@@ -142,15 +119,14 @@ def _linear_to_fp8linear(layer: nn.Linear) -> FP8Linear:
     )
 
 
-def _apply_fp8_prepare_to_model(model: nn.Module, excluded_layer_substrings: tuple[str, ...]) -> nn.Module:
+def _swap_linears_to_fp8(model: nn.Module, should_swap: Callable[[str], bool]) -> nn.Module:
-    """Replace nn.Linear layers with FP8Linear in the module tree."""
+    """Replace nn.Linear layers with FP8Linear where ``should_swap(name)`` returns True."""
     replacements: list[tuple[nn.Module, str, nn.Linear]] = []
 
     for name, module in model.named_modules():
         if not isinstance(module, nn.Linear) or isinstance(module, FP8Linear):
             continue
-
+        if not should_swap(name):
-        if _should_skip_layer(name, excluded_layer_substrings):
             continue
 
         if "." in name:
@@ -168,40 +144,32 @@ def _apply_fp8_prepare_to_model(model: nn.Module, excluded_layer_substrings: tup
     return model
 
 
-def _create_transpose_kv_operation(
+def get_fp8_swap_module_ops(checkpoint_path: str) -> tuple[ModuleOps, ...]:
-    excluded_layer_substrings: tuple[str, ...],
+    """Return the FP8 swap ``ModuleOps`` for layers whose ``.weight`` is ``F8_E4M3``
-) -> Callable[[str, torch.Tensor], list[KeyValueOperationResult]]:
+    and which have a sibling ``.weight_scale`` tensor in the checkpoint.
-    def transpose_if_matches(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
+    Raises ``ValueError`` if no such layers are found — that combination is ambiguous
-        # Only process .weight keys
+    (a BF16 checkpoint with this policy would load as a no-op).
-        if not key.endswith(".weight"):
+    """
-            return [KeyValueOperationResult(key, value)]
+    dtypes = _read_safetensors_dtypes(checkpoint_path)
+    fp8_scale_paths = frozenset(
+        key.removesuffix(".weight_scale")
+        for key in dtypes
+        if key.endswith(".weight_scale") and dtypes.get(key.removesuffix(".weight_scale") + ".weight") == "F8_E4M3"
+    )
+    if not fp8_scale_paths:
+        raise ValueError(
+            f"fp8_scaled_mm requires a pre-quantized checkpoint with F8_E4M3 .weight + .weight_scale "
+            f"tensors, but {checkpoint_path!r} has none. Use QuantizationPolicy.fp8_cast() for BF16 checkpoints."
+        )
 
-        # Only transpose 2D FP8 tensors (Linear weights)
+    def _should_swap(name: str) -> bool:
-        if value.dim() != 2 or value.dtype != torch.float8_e4m3fn:
+        suffix = "." + name
-            return [KeyValueOperationResult(key, value)]
+        return any(p == name or p.endswith(suffix) for p in fp8_scale_paths)
 
-        # Check if the layer is excluded
+    return (
-        layer_name = key.rsplit(".weight", 1)[0]
+        ModuleOps(
-        if _should_skip_layer(layer_name, excluded_layer_substrings):
+            name="fp8_swap_linears",
-            return [KeyValueOperationResult(key, value)]
+            matcher=lambda model: isinstance(model, LTXModel),
-
+            mutator=lambda model: _swap_linears_to_fp8(model, _should_swap),
-        # Transpose to cuBLAS layout (in, out)
+        ),
-        transposed_weight = value.t()
+    )
-
-        return [KeyValueOperationResult(key, transposed_weight)]
-
-    return transpose_if_matches
-
-
-FP8_TRANSPOSE_SD_OPS = SDOps("fp8_transpose_weights").with_kv_operation(
-    _create_transpose_kv_operation(EXCLUDED_LAYER_SUBSTRINGS),
-    key_prefix="transformer_blocks.",
-    key_suffix=".weight",
-)
-
-
-FP8_PREPARE_MODULE_OPS = ModuleOps(
-    name="fp8_prepare_for_loading",
-    matcher=lambda model: isinstance(model, LTXModel),
-    mutator=lambda model: _apply_fp8_prepare_to_model(model, EXCLUDED_LAYER_SUBSTRINGS),
-)

packages/ltx-core/src/ltx_core/quantization/policy.py

@@ -1,39 +1,48 @@
 from dataclasses import dataclass
+from enum import Enum
 
 from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.sd_ops import SDOps
 from ltx_core.quantization.fp8_cast import TRANSFORMER_LINEAR_DOWNCAST_MAP, UPCAST_DURING_INFERENCE
-from ltx_core.quantization.fp8_scaled_mm import FP8_PREPARE_MODULE_OPS, FP8_TRANSPOSE_SD_OPS
+from ltx_core.quantization.fp8_scaled_mm import get_fp8_swap_module_ops
 
 
 @dataclass(frozen=True)
 class QuantizationPolicy:
     """Configuration for model quantization during loading.
     Attributes:
-        sd_ops: State dict operations for weight transformation.
+        kind: Discriminator for the policy variant.
-        module_ops: Post-load module transformations.
+        sd_ops: State-dict operations applied to each tensor during load.
+        module_ops: Post-load module transformations applied to the meta model.
     """
 
+    class Kind(str, Enum):
+        FP8_CAST = "fp8_cast"
+        FP8_SCALED_MM = "fp8_scaled_mm"
+
+    kind: Kind
     sd_ops: SDOps | None = None
     module_ops: tuple[ModuleOps, ...] = ()
 
     @classmethod
     def fp8_cast(cls) -> "QuantizationPolicy":
-        """Create policy using FP8 casting with upcasting during inference."""
+        """FP8 casting with upcasting during inference."""
         return cls(
+            kind=cls.Kind.FP8_CAST,
             sd_ops=TRANSFORMER_LINEAR_DOWNCAST_MAP,
             module_ops=(UPCAST_DURING_INFERENCE,),
         )
 
     @classmethod
-    def fp8_scaled_mm(cls) -> "QuantizationPolicy":
+    def fp8_scaled_mm(cls, checkpoint_path: str) -> "QuantizationPolicy":
-        """Create policy using FP8 scaled matrix multiplication."""
+        """FP8 scaled matmul for checkpoints pre-quantized with per-tensor scales.
-        try:
+        The set of layers to swap to ``FP8Linear`` is discovered from the
-            import tensorrt_llm  # noqa: F401, PLC0415
+        checkpoint's ``.weight_scale`` tensors via suffix-matching against the
-        except ImportError as e:
+        model's named modules. Requires a pre-quantized checkpoint; for BF16
-            raise ImportError("tensorrt_llm is not installed, skipping FP8 scaled MM quantization") from e
+        checkpoints, use :meth:`fp8_cast` instead.
-
+        """
         return cls(
-            sd_ops=FP8_TRANSPOSE_SD_OPS,
+            kind=cls.Kind.FP8_SCALED_MM,
-            module_ops=(FP8_PREPARE_MODULE_OPS,),
+            sd_ops=None,
+            module_ops=get_fp8_swap_module_ops(checkpoint_path),
         )

packages/ltx-core/src/ltx_core/quantization/trtllm_scaled_usable.py

@@ -0,0 +1,37 @@
+"""Runtime detection of TensorRT-LLM FP8 scaled-matmul availability.
+When the TRT-LLM ops are usable on the current host (Linux + Hopper-class CUDA
++ tensorrt_llm wheel installed) we use them since they outperform the PyTorch-native
+``torch._scaled_mm`` path. Otherwise we fall back to the native implementation,
+which is portable across platforms (Windows, macOS, AMD GPUs).
+The check runs once and is cached.
+"""
+
+from __future__ import annotations
+
+import platform
+from functools import cache
+
+import torch
+
+
+@cache
+def trtllm_scaled_mm_usable() -> bool:
+    if platform.system() != "Linux":
+        return False
+
+    if not torch.cuda.is_available():
+        return False
+
+    major, minor = torch.cuda.get_device_capability()
+    sm = major * 10 + minor
+
+    if sm < 90 or sm >= 120:
+        return False
+
+    # The import is load-bearing — registers the trtllm torch ops as a side effect.
+    try:
+        import tensorrt_llm  # noqa: F401, PLC0415
+    except Exception:
+        return False
+
+    return True

packages/ltx-core/src/ltx_core/text_encoders/gemma/embeddings_connector.py

@@ -18,7 +18,7 @@ class _BasicTransformerBlock1D(torch.nn.Module):
         dim: int,
         heads: int,
         dim_head: int,
-        rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+        rope_type: LTXRopeType = LTXRopeType.SPLIT,
         apply_gated_attention: bool = False,
     ):
         super().__init__()
@@ -39,7 +39,7 @@ class _BasicTransformerBlock1D(torch.nn.Module):
     def forward(
         self,
         hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
+        additive_attention_mask: torch.Tensor | None = None,
         pe: torch.Tensor | None = None,
     ) -> torch.Tensor:
         # Notice that normalization is always applied before the real computation in the following blocks.
@@ -49,8 +49,8 @@ class _BasicTransformerBlock1D(torch.nn.Module):
 
         norm_hidden_states = norm_hidden_states.squeeze(1)
 
-        # 2. Self-Attention
+        # 2. Self-Attention — `mask` is the kernel-boundary name for the additive mask.
-        attn_output = self.attn1(norm_hidden_states, mask=attention_mask, pe=pe)
+        attn_output = self.attn1(norm_hidden_states, mask=additive_attention_mask, pe=pe)
 
         hidden_states = attn_output + hidden_states
         if hidden_states.ndim == 4:
@@ -84,7 +84,7 @@ class Embeddings1DConnector(torch.nn.Module):
         causal_temporal_positioning (bool): If True, uses causal attention (default=False).
         num_learnable_registers (int | None): Number of learnable registers to replace padded tokens. If None, disables
             register replacement. (default=128)
-        rope_type (LTXRopeType): The RoPE variant to use (default=DEFAULT_ROPE_TYPE).
+        rope_type (LTXRopeType): The RoPE variant to use.
         double_precision_rope (bool): Use double precision rope calculation (default=False).
     """
 
@@ -99,7 +99,7 @@ class Embeddings1DConnector(torch.nn.Module):
         positional_embedding_max_pos: list[int] | None = None,
         causal_temporal_positioning: bool = False,
         num_learnable_registers: int | None = 128,
-        rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+        rope_type: LTXRopeType = LTXRopeType.SPLIT,
         double_precision_rope: bool = False,
         apply_gated_attention: bool = False,
     ):
@@ -133,51 +133,40 @@ class Embeddings1DConnector(torch.nn.Module):
             )
 
     def _replace_padded_with_learnable_registers(
-        self, hidden_states: torch.Tensor, attention_mask: torch.Tensor
+        self, hidden_states: torch.Tensor, additive_attention_mask: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        assert hidden_states.shape[1] % self.num_learnable_registers == 0, (
+        batch_size, seq_len, _ = hidden_states.shape
-            f"Hidden states sequence length {hidden_states.shape[1]} must be divisible by num_learnable_registers "
-            f"{self.num_learnable_registers}."
-        )
 
-        num_registers_duplications = hidden_states.shape[1] // self.num_learnable_registers
+        assert seq_len % self.num_learnable_registers == 0
-        learnable_registers = torch.tile(self.learnable_registers, (num_registers_duplications, 1))
-        attention_mask_binary = (attention_mask.squeeze(1).squeeze(1).unsqueeze(-1) >= -9000.0).int()
 
-        non_zero_hidden_states = hidden_states[:, attention_mask_binary.squeeze().bool(), :]
+        registers = self.learnable_registers.repeat(seq_len // self.num_learnable_registers, 1).to(hidden_states.dtype)
-        non_zero_nums = non_zero_hidden_states.shape[1]
+        registers = registers.unsqueeze(0).expand(batch_size, -1, -1)  # (B, seq_len, hidden_dim)
-        pad_length = hidden_states.shape[1] - non_zero_nums
+        binary_mask = additive_attention_mask[:, 0, 0, :].unsqueeze(-1) >= 0
-        adjusted_hidden_states = torch.nn.functional.pad(non_zero_hidden_states, pad=(0, 0, 0, pad_length), value=0)
+        binary_mask = binary_mask.to(hidden_states.dtype)
-        flipped_mask = torch.flip(attention_mask_binary, dims=[1])
+        hidden_states = binary_mask * hidden_states + (1 - binary_mask) * registers
-        hidden_states = flipped_mask * adjusted_hidden_states + (1 - flipped_mask) * learnable_registers
 
-        attention_mask = torch.full_like(
+        return hidden_states, torch.zeros_like(additive_attention_mask)
-            attention_mask,
-            0.0,
-            dtype=attention_mask.dtype,
-            device=attention_mask.device,
-        )
-
-        return hidden_states, attention_mask
 
     def forward(
         self,
         hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
+        additive_attention_mask: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
+        """Forward pass of Embeddings1DConnector.
-        Forward pass of Embeddings1DConnector.
         Args:
-            hidden_states (torch.Tensor): Input tensor of embeddings (shape [batch, seq_len, feature_dim]).
+            hidden_states: (B, S, D) input embeddings.
-            attention_mask (torch.Tensor|None): Optional mask for valid tokens (shape compatible with hidden_states).
+            additive_attention_mask: optional additive mask of shape (B, 1, 1, S), where
+                valid = 0.0 and padding = -torch.finfo(dtype).max.
         Returns:
-            tuple[torch.Tensor, torch.Tensor]: Processed features and the corresponding (possibly modified) mask.
+            (hidden_states, additive_attention_mask)
         """
         if self.num_learnable_registers:
-            hidden_states, attention_mask = self._replace_padded_with_learnable_registers(hidden_states, attention_mask)
+            hidden_states, additive_attention_mask = self._replace_padded_with_learnable_registers(
+                hidden_states, additive_attention_mask
+            )
 
         indices_grid = torch.arange(hidden_states.shape[1], dtype=torch.float32, device=hidden_states.device)
-        indices_grid = indices_grid[None, None, :]
+        indices_grid = indices_grid[None, None, :].expand(hidden_states.shape[0], -1, -1)
         freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch
         freqs_cis = precompute_freqs_cis(
             indices_grid=indices_grid,
@@ -191,11 +180,11 @@ class Embeddings1DConnector(torch.nn.Module):
         )
 
         for block in self.transformer_1d_blocks:
-            hidden_states = block(hidden_states, attention_mask=attention_mask, pe=freqs_cis)
+            hidden_states = block(hidden_states, additive_attention_mask=additive_attention_mask, pe=freqs_cis)
 
         hidden_states = rms_norm(hidden_states)
 
-        return hidden_states, attention_mask
+        return hidden_states, additive_attention_mask
 
 
 class Embeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConnector]):
@@ -204,7 +193,7 @@ class Embeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConnector]
     @classmethod
     def from_config(cls: type[Embeddings1DConnector], config: dict) -> Embeddings1DConnector:
         transformer_config = config.get("transformer", {})
-        rope_type = LTXRopeType(transformer_config.get("rope_type", "interleaved"))
+        rope_type = LTXRopeType(transformer_config.get("rope_type", "split"))
         double_precision_rope = transformer_config.get("frequencies_precision", False) == "float64"
         pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
 
@@ -231,7 +220,7 @@ class AudioEmbeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConne
     @classmethod
     def from_config(cls: type[Embeddings1DConnector], config: dict) -> Embeddings1DConnector:
         transformer_config = config.get("transformer", {})
-        rope_type = LTXRopeType(transformer_config.get("rope_type", "interleaved"))
+        rope_type = LTXRopeType(transformer_config.get("rope_type", "split"))
         double_precision_rope = transformer_config.get("frequencies_precision", False) == "float64"
         pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
 

packages/ltx-core/src/ltx_core/text_encoders/gemma/embeddings_processor.py

@@ -19,12 +19,32 @@ def convert_to_additive_mask(attention_mask: torch.Tensor, dtype: torch.dtype) -
     ) * torch.finfo(dtype).max
 
 
-def _to_binary_mask(encoded: torch.Tensor, encoded_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+def _compute_right_pad_order(additive_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    """Convert connector output mask to binary mask and apply to encoded tensor."""
+    """Compute the index permutation that places valid tokens before pads in each row.
-    binary_mask = (encoded_mask < 0.000001).to(torch.int64)
+    Stable sort: valid tokens keep their relative order. Idempotent for inputs already
-    binary_mask = binary_mask.reshape([encoded.shape[0], encoded.shape[1], 1])
+    right-padded. The sort and reordered mask depend only on the mask, so they can be
-    encoded = encoded * binary_mask
+    computed once and reused across multiple feature tensors that share the mask.
-    return encoded, binary_mask
+    Args:
+        additive_mask: (B, 1, 1, S) additive mask, ``0.0`` for valid, ``-finfo.max`` for pad.
+    Returns:
+        ``(sort_idx, reordered_additive_mask)``: ``sort_idx`` is (B, S); the reordered mask
+        has the same shape as the input.
+    """
+    binary = (additive_mask[:, 0, 0, :] >= 0).to(torch.int32)  # (B, S)
+    sort_idx = torch.argsort(binary, dim=-1, descending=True, stable=True)  # (B, S)
+    new_binary = torch.gather(binary, 1, sort_idx)
+    new_additive = (new_binary.to(additive_mask.dtype) - 1) * torch.finfo(additive_mask.dtype).max
+    return sort_idx, new_additive[:, None, None, :]
+
+
+def _apply_right_pad_order(features: torch.Tensor, sort_idx: torch.Tensor) -> torch.Tensor:
+    """Apply a precomputed right-pad permutation (from ``_compute_right_pad_order``) to features."""
+    return torch.gather(features, 1, sort_idx.unsqueeze(-1).expand_as(features))
+
+
+def _to_binary_mask(encoded_mask: torch.Tensor, lead_shape: tuple[int, int]) -> torch.Tensor:
+    """Convert connector output mask to a binary (0/1) mask shaped ``(B, S, 1)`` for broadcasting."""
+    return (encoded_mask < 0.000001).to(torch.int64).reshape([lead_shape[0], lead_shape[1], 1])
 
 
 class EmbeddingsProcessor(nn.Module):
@@ -57,12 +77,19 @@ class EmbeddingsProcessor(nn.Module):
         if self.audio_connector is None and audio_features is not None:
             raise ValueError("Audio features were provided but no audio connector is configured.")
 
-        video_encoded, video_mask = self.video_connector(video_features, additive_attention_mask)
+        # Connectors expect right-padded input ([valid, pad]). Normalize layout here so the
-        video_encoded, binary_mask = _to_binary_mask(video_encoded, video_mask)
+        # upstream tokenizer can keep using either side without coupling to the connector.
+        # The sort index depends only on the mask, so compute it once and reuse for audio.
+        sort_idx, mask_for_connector = _compute_right_pad_order(additive_attention_mask)
+        video_features = _apply_right_pad_order(video_features, sort_idx)
+        video_encoded, video_mask = self.video_connector(video_features, mask_for_connector)
+        binary_mask = _to_binary_mask(video_mask, video_encoded.shape[:2])
+        video_encoded = video_encoded * binary_mask
 
         audio_encoded = None
         if self.audio_connector is not None:
-            audio_encoded, _ = self.audio_connector(audio_features, additive_attention_mask)
+            audio_features = _apply_right_pad_order(audio_features, sort_idx)
+            audio_encoded, _ = self.audio_connector(audio_features, mask_for_connector)
 
         return video_encoded, audio_encoded, binary_mask.squeeze(-1)
 

packages/ltx-core/src/ltx_core/text_encoders/gemma/feature_extractor.py

@@ -11,37 +11,25 @@ from torch import nn
 
 def _norm_and_concat_padded_batch(
     encoded_text: torch.Tensor,
-    sequence_lengths: torch.Tensor,
+    attention_mask: torch.Tensor,
-    padding_side: str = "right",
 ) -> torch.Tensor:
     """Normalize and flatten multi-layer hidden states, respecting padding.
     Performs per-batch, per-layer normalization using masked mean and range,
-    then concatenates across the layer dimension.
+    then concatenates across the layer dimension. Padding-side agnostic: the
+    binary ``attention_mask`` already encodes which positions are valid.
     Args:
         encoded_text: Hidden states of shape [batch, seq_len, hidden_dim, num_layers].
-        sequence_lengths: Number of valid (non-padded) tokens per batch item.
+        attention_mask: Binary mask of shape [batch, seq_len], 1 for valid tokens, 0 for padding.
-        padding_side: Whether padding is on "left" or "right".
     Returns:
         Normalized tensor of shape [batch, seq_len, hidden_dim * num_layers],
         with padded positions zeroed out.
     """
-    b, t, d, l = encoded_text.shape  # noqa: E741
+    b, _, d, l = encoded_text.shape  # noqa: E741
-    device = encoded_text.device
-
-    token_indices = torch.arange(t, device=device)[None, :]
-
-    if padding_side == "right":
-        mask = token_indices < sequence_lengths[:, None]
-    elif padding_side == "left":
-        start_indices = t - sequence_lengths[:, None]
-        mask = token_indices >= start_indices
-    else:
-        raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}")
-
-    mask = rearrange(mask, "b t -> b t 1 1")
-
     eps = 1e-6
 
+    sequence_lengths = attention_mask.sum(dim=-1)
+    mask = rearrange(attention_mask.bool(), "b t -> b t 1 1")
+
     masked = encoded_text.masked_fill(~mask, 0.0)
     denom = (sequence_lengths * d).view(b, 1, 1, 1)
     mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + eps)
@@ -51,12 +39,10 @@ def _norm_and_concat_padded_batch(
     range_ = x_max - x_min
 
     normed = 8 * (encoded_text - mean) / (range_ + eps)
-    normed = normed.reshape(b, t, -1)
+    normed = normed.reshape(b, -1, d * l)
 
     mask_flattened = rearrange(mask, "b t 1 1 -> b t 1").expand(-1, -1, d * l)
-    normed = normed.masked_fill(~mask_flattened, 0.0)
+    return normed.masked_fill(~mask_flattened, 0.0)
-
-    return normed
 
 
 def norm_and_concat_per_token_rms(
@@ -97,12 +83,14 @@ class FeatureExtractorV1(nn.Module):
         self.is_av = is_av
 
     def forward(
-        self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, padding_side: str = "left"
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        padding_side: str = "left",  # noqa: ARG002 — kept for API stability; norm is layout-agnostic
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         encoded = torch.stack(hidden_states, dim=-1) if isinstance(hidden_states, (list, tuple)) else hidden_states
         dtype = encoded.dtype
-        sequence_lengths = attention_mask.sum(dim=-1)
+        normed = _norm_and_concat_padded_batch(encoded, attention_mask)
-        normed = _norm_and_concat_padded_batch(encoded, sequence_lengths, padding_side)
         features = self.aggregate_embed(normed.to(dtype))
         if self.is_av:
             return features, features

packages/ltx-core/src/ltx_core/text_encoders/gemma/tokenizer.py

@@ -1,6 +1,13 @@
+from enum import Enum
+
 from transformers import AutoTokenizer
 
 
+class PaddingSide(str, Enum):
+    LEFT = "left"
+    RIGHT = "right"
+
+
 class LTXVGemmaTokenizer:
     """
     Tokenizer wrapper for Gemma models compatible with LTXV processes.
@@ -8,18 +15,18 @@ class LTXVGemmaTokenizer:
     ensuring correct settings and output formatting for downstream consumption.
     """
 
-    def __init__(self, tokenizer_path: str, max_length: int = 256):
+    def __init__(self, tokenizer_path: str, max_length: int = 256, padding_side: PaddingSide = PaddingSide.LEFT):
         """
         Initialize the tokenizer.
         Args:
             tokenizer_path (str): Path to the pretrained tokenizer files or model directory.
             max_length (int, optional): Max sequence length for encoding. Defaults to 256.
+            padding_side (PaddingSide, optional): Side to pad on. Defaults to ``PaddingSide.LEFT``.
         """
         self.tokenizer = AutoTokenizer.from_pretrained(
             tokenizer_path, local_files_only=True, model_max_length=max_length
         )
-        # Gemma expects left padding for chat-style prompts; for plain text it doesn't matter much.
+        self.tokenizer.padding_side = padding_side.value
-        self.tokenizer.padding_side = "left"
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
 

packages/ltx-core/src/ltx_core/tools.py

@@ -138,7 +138,7 @@ class VideoLatentTools(LatentTools):
             LatentState(
                 latent=initial_latent,
                 denoise_mask=denoise_mask,
-                positions=positions.to(dtype),
+                positions=positions,
                 clean_latent=clean_latent,
             )
         )

packages/ltx-core/src/ltx_core/types.py

@@ -20,15 +20,17 @@ class SpatioTemporalScaleFactors(NamedTuple):
     """
     Describes the spatiotemporal downscaling between decoded video space and
     the corresponding VAE latent grid.
+    Field order matches the (frame/time, height, width) axis layout used by
+    latent tensors and meshgrid coordinates elsewhere in the codebase.
     """
 
     time: int
-    width: int
     height: int
+    width: int
 
     @classmethod
     def default(cls) -> "SpatioTemporalScaleFactors":
-        return cls(time=8, width=32, height=32)
+        return cls(time=8, height=32, width=32)
 
 
 VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
@@ -74,9 +76,9 @@ class VideoLatentShape(NamedTuple):
         latent_channels: int = 128,
         scale_factors: SpatioTemporalScaleFactors = VIDEO_SCALE_FACTORS,
     ) -> "VideoLatentShape":
-        frames = (shape.frames - 1) // scale_factors[0] + 1
+        frames = (shape.frames - 1) // scale_factors.time + 1
-        height = shape.height // scale_factors[1]
+        height = shape.height // scale_factors.height
-        width = shape.width // scale_factors[2]
+        width = shape.width // scale_factors.width
 
         return VideoLatentShape(
             batch=shape.batch,

packages/ltx-pipelines/CLAUDE.md

@@ -23,6 +23,7 @@ Inference pipelines for LTX-2 audio-video generation. Depends on `ltx-core` for
 | `KeyframeInterpolationPipeline` | `keyframe_interpolation.py` | 2 | Full + distilled LoRA | Euler | Keyframe interpolation |
 | `DistilledPipeline` | `distilled.py` | 2 | Distilled only | Euler | Fastest inference |
 | `ICLoraPipeline` | `ic_lora.py` | 2 | Distilled only | Euler | Video-to-video with IC-LoRA control |
+| `LipDubPipeline` | `lipdub.py` | 2 | Distilled only | Euler | Lip dubbing with IC-LoRA + audio ref conditioning |
 | `RetakePipeline` | `retake.py` | 1 | Full or distilled | Euler | Video region regeneration |
 
 ## Guidance
@@ -56,7 +57,7 @@ Inference pipelines for LTX-2 audio-video generation. Depends on `ltx-core` for
 ### Memory management
 
 - **Model lifecycle**: All blocks build their model on call and free it on exit. `gpu_model()` moves params to `"meta"` device on exit, immediately releasing storage. No model persists between calls.
-- **Layer streaming**: When `streaming_prefetch_count` is set, `DiffusionStage` wraps the transformer in `LayerStreamingWrapper`. Layers live on pinned CPU memory; only `1 + prefetch_count` layers are on GPU at a time, with async H2D prefetch on a separate CUDA stream.
+- **Block streaming**: When offloading is enabled, `DiffusionStage` wraps the transformer in `BlockStreamingWrapper`. Blocks live on pinned CPU memory; only 2 blocks are buffered on GPU at a time (one for compute, one for async H2D copy on a separate CUDA stream).
 - **Batch splitting**: `BatchSplitAdapter` wraps the transformer and splits inputs exceeding `max_batch_size` into sequential chunks. If guidance needs B=4 but `max_batch_size=1`, it runs 4 sequential B=1 passes. Higher `max_batch_size` reduces layer-streaming PCIe transfers at the cost of peak memory.
 
 ## Denoisers (`utils/denoisers.py`)
@@ -65,6 +66,10 @@ Inference pipelines for LTX-2 audio-video generation. Depends on `ltx-core` for
 - `GuidedDenoiser` -- CFG/STG with static `MultiModalGuider` instances (HQ, A2Vid, Retake non-distilled).
 - `FactoryGuidedDenoiser` -- per-step guider creation via factory (OneStageTI2Vid, TwoStagesTI2Vid, Keyframe).
 
+All denoisers return a `(video_result, audio_result)` tuple of `DenoisedLatentResult` (defined in `utils/types.py`), either element may be `None` for absent modalities. `DenoisedLatentResult.denoised` is the final blended tensor. Guided denoisers additionally populate per-pass fields (`.cond`, `.uncond`, `.ptb`, `.mod`) on each result; `SimpleDenoiser` leaves these `None`.
+
+`GuidedDenoiser` and `FactoryGuidedDenoiser` accept `force_uncond_pass=True` to run the uncond pass even when `cfg_scale=1.0` (required by CFG++ when the guidance scale is 1 but the uncond prediction is still needed for the ODE derivative). Requires `negative_context` to be set on the guider. When enabled, `DenoisedLatentResult.uncond` will be a tensor instead of `None`.
+
 Guided denoisers batch all guidance passes into a **single transformer call**: states are repeated along the batch dimension, contexts concatenated, and a `BatchedPerturbationConfig` controls which attention ops are skipped per sample. Pass count is dynamic: B=2 for CFG-only, up to B=4 with CFG+STG+modality isolation. Results are split back and blended by the guider.
 
 ## Per-pipeline unique features
@@ -72,6 +77,7 @@ Guided denoisers batch all guidance passes into a **single transformer call**: s
 - **HQ**: Res2s second-order sampler for **both** stages, latent-dependent sigma schedule, distilled LoRA on both stages with separate strengths.
 - **A2Vid**: Audio frozen in both stages (`frozen=True, noise_scale=0.0`). Returns original audio (not VAE-decoded); no `AudioDecoder`.
 - **IC-LoRA**: `VideoConditionByReferenceLatent`, `reference_downscale_factor` from LoRA metadata, `skip_stage_2`, attention mask downsampling. Stage 2 is LoRA-free and uses `combined_image_conditionings` (no IC-LoRA conditioning).
+- **LipDub**: Standalone pipeline; IC reference **video** helpers in `iclora_utils.py`, LipDub-only **audio** patchify/negative positions in `lipdub.py`. Appends frozen audio-reference tokens via `AudioConditionByReferenceLatent` (ltx-core), matching video token order (`[target | ref]`) while keeping reference RoPE positions negative (training-compatible). Single IC-LoRA on both stages; full IC-LoRA video conditioning at stage 1 and 2; stage-2 audio is frozen with S1 latent as initial state and uses S1-derived ref. Final audio decoded from stage 1 latent. The LipDub CLI does not expose `--conditioning-attention-mask`; use `ic_lora.py` if you need spatial IC attention masking.
 - **Keyframe**: Uses `image_conditionings_by_adding_guiding_latent` in both stages (all frames as keyframe guidance, no replacement) -- unlike TI2Vid which uses `combined_image_conditionings` (frame_idx=0 replaces, others guide).
 - **Retake**: `TemporalRegionMask` for selective time-window regeneration. `regenerate_video`/`regenerate_audio` flags. Conditional distilled/full behavior.
 - **Distilled**: Single `self.stage` reused for both stages (not `stage_1`/`stage_2`).

packages/ltx-pipelines/README.md

@@ -63,6 +63,8 @@ Available pipeline modules:
 - `ltx_pipelines.keyframe_interpolation` - Keyframe interpolation.
 - `ltx_pipelines.a2vid_two_stage` - Audio-to-video generation conditioned on an input audio.
 - `ltx_pipelines.retake` - Regenerate a time region of an existing video.
+- `ltx_pipelines.hdr_ic_lora` - Video-to-video with HDR output (linear float via LogC3 inverse decode).
+- `ltx_pipelines.lipdub` - Lip dubbing / re-voicing with IC-LoRA and audio reference conditioning.
 
 Use `--help` with any pipeline module to see all available options and parameters.
 
@@ -79,6 +81,9 @@ Do you have an existing video to modify?
 Do you have an audio file to drive generation?
 ├─ YES → Use A2VidPipelineTwoStage (audio-to-video)
 │
+Do you need HDR output (linear float frames for EXR / tonemapping)?
+├─ YES → Use HDRICLoraPipeline (video-to-video with LogC3 inverse decode)
+│
 Do you need to condition on existing images/videos?
 ├─ YES → Do you have reference videos for video-to-video?
 │  ├─ YES → Use ICLoraPipeline
@@ -108,6 +113,8 @@ Do you need to condition on existing images/videos?
 | **KeyframeInterpolationPipeline** | 2 | ✅ | ✅ | Keyframes | Animation, interpolation |
 | **A2VidPipelineTwoStage** | 2 | ✅ | ✅ | Audio + Image | Audio-driven video generation |
 | **RetakePipeline** | 1 | ✅ | ❌ | Source Video | Regenerating a time region of a video |
+| **HDRICLoraPipeline** | 2 | ❌ | ✅ | Video | HDR video-to-video (linear float output for EXR) |
+| **LipDubPipeline** | 2 | ✅ | ✅ | Video + Audio | Lip dubbing with audio ref conditioning |
 
 ---
 
@@ -219,6 +226,34 @@ Single-stage generation that encodes the source video and audio into latents, ap
 
 ---
 
+### 9. HDRICLoraPipeline
+
+**Best for:** Video-to-video generation with HDR output for EXR export and offline tonemapping.
+
+**Source**: [`src/ltx_pipelines/hdr_ic_lora.py`](src/ltx_pipelines/hdr_ic_lora.py)
+
+Two-stage video-to-video on the distilled model with an HDR IC-LoRA. Decoded latents pass through an HDR inverse transform (ARRI LogC3, auto-detected from LoRA metadata) to produce a **linear HDR float** tensor `[f, h, w, c]`. Video-only (audio skipped). Text embeddings are pre-computed externally and loaded from a `.safetensors` file. Tonemapping and EXR saving are the caller's responsibility. LoRA and embeddings: [`Lightricks/LTX-2.3-22b-IC-LoRA-HDR`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-HDR).
+
+**Extra CLI arguments:** `--input` (mp4 or directory, required), `--output-dir` (required), `--hdr-lora` (required), `--text-embeddings` (pre-computed `.safetensors`, required), `--num-frames`, `--spatial-tile` (tiled VAE decode tile size; reduce on lower-VRAM GPUs), `--skip-mp4` (EXR only, no H.264 preview), `--exr-half` (float16 EXR), `--high-quality` (generates 2x frames internally for smoother output, ~2x slower), `--offload {none,cpu,disk}` (weight offloading; disables FP8 quantization when not `none`).
+
+**Use when:** You need linear HDR float output for EXR export, color grading, or custom tonemapping workflows.
+
+---
+
+### 10. LipDubPipeline
+
+**Best for:** Lip dubbing, rephrasing while keeping the same speaker identity and matching lip movements to new audio.
+
+**Source**: [`src/ltx_pipelines/lipdub.py`](src/ltx_pipelines/lipdub.py)
+
+Uses IC-LoRA on a **distilled** checkpoint with a **single** lip-dub IC-LoRA applied in **both** stages. The reference clip provides video and audio reference tokens whose VAE latents are appended to the target audio sequence as frozen reference tokens. The frame count and frame rate are derived from the reference video (frame count is silently snapped to the nearest `8k+1`), so the CLI does not accept `--num-frames` or `--frame-rate`. Required: `--reference-video`. Optional: `--reference-strength`. LoRA: [`Lightricks/LTX-2.3-22b-IC-LoRA-LipDub`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-LipDub).
+
+**Note:** Requires a distilled model checkpoint and one lip-dub IC-LoRA (`--lora` exactly once).
+
+**Use when:** Dubbing, rephrasing with matched lips and speaker identity.
+
+---
+
 ## 🎨 Conditioning Types
 
 Pipelines use different conditioning methods from [`ltx-core`](../ltx-core/) for controlling generation. See the [ltx-core conditioning documentation](../ltx-core/README.md#conditioning--control) for details.

packages/ltx-pipelines/pyproject.toml

@@ -1,10 +1,10 @@
 [project]
 name = "ltx-pipelines"
-version = "1.1.1"
+version = "1.1.3"
 description = "Pipelines implementation for Lightricks' LTX-2 model"
 readme = "README.md"
 requires-python = ">=3.10"
-dependencies = ["ltx-core", "av", "tqdm", "pillow"]
+dependencies = ["ltx-core", "av", "tqdm", "pillow", "openimageio"]
 
 [build-system]
 requires = ["uv_build>=0.9.8,<0.10.0"]

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

@@ -5,6 +5,7 @@ This package provides ready-to-use pipelines for video generation:
 - TI2VidTwoStagesPipeline: Two-stage generation with upsampling
 - DistilledPipeline: Fast distilled two-stage generation
 - ICLoraPipeline: Image/video conditioning with distilled LoRA
+- LipDubPipeline: Lip dubbing with IC-LoRA and audio conditioning
 - KeyframeInterpolationPipeline: Keyframe-based video interpolation
 - RetakePipeline: Regenerate a time region (retake) of an existing video
 For more detailed components and utilities, import from specific submodules
@@ -15,6 +16,7 @@ from ltx_pipelines.a2vid_two_stage import A2VidPipelineTwoStage
 from ltx_pipelines.distilled import DistilledPipeline
 from ltx_pipelines.ic_lora import ICLoraPipeline
 from ltx_pipelines.keyframe_interpolation import KeyframeInterpolationPipeline
+from ltx_pipelines.lipdub import LipDubPipeline
 from ltx_pipelines.retake import RetakePipeline
 from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
 from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
@@ -24,6 +26,7 @@ __all__ = [
     "DistilledPipeline",
     "ICLoraPipeline",
     "KeyframeInterpolationPipeline",
+    "LipDubPipeline",
     "RetakePipeline",
     "TI2VidOneStagePipeline",
     "TI2VidTwoStagesPipeline",

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

@@ -31,7 +31,7 @@ from ltx_pipelines.utils.helpers import (
     get_device,
 )
 from ltx_pipelines.utils.media_io import decode_audio_from_file, encode_video
-from ltx_pipelines.utils.types import ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class A2VidPipelineTwoStage:
@@ -53,12 +53,15 @@ class A2VidPipelineTwoStage:
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
         torch_compile: bool = False,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
         self.dtype = torch.bfloat16
         self._scheduler = LTX2Scheduler()
 
-        self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
+        self.prompt_encoder = PromptEncoder(
+            checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
+        )
         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(
@@ -69,6 +72,7 @@ class A2VidPipelineTwoStage:
             quantization=quantization,
             registry=registry,
             torch_compile=torch_compile,
+            offload_mode=offload_mode,
         )
         stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
         self.stage_2 = DiffusionStage(
@@ -79,6 +83,7 @@ class A2VidPipelineTwoStage:
             quantization=quantization,
             registry=registry,
             torch_compile=torch_compile,
+            offload_mode=offload_mode,
         )
         self.upsampler = VideoUpsampler(
             checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -102,7 +107,6 @@ 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,
         stage_1_sigmas: torch.Tensor | None = None,
         stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
@@ -117,7 +121,6 @@ class A2VidPipelineTwoStage:
             [prompt, negative_prompt],
             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
@@ -183,7 +186,6 @@ class A2VidPipelineTwoStage:
                 noise_scale=0.0,
                 initial_latent=encoded_audio_latent,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
             max_batch_size=max_batch_size,
         )
 
@@ -223,7 +225,6 @@ class A2VidPipelineTwoStage:
                 noise_scale=0.0,
                 initial_latent=encoded_audio_latent,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -266,6 +267,7 @@ def main() -> None:
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
         torch_compile=args.compile,
+        offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -294,7 +296,6 @@ 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,
     )
 

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

@@ -34,7 +34,7 @@ from ltx_pipelines.utils.helpers import (
     get_device,
 )
 from ltx_pipelines.utils.media_io import encode_video
-from ltx_pipelines.utils.types import ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class DistilledPipeline:
@@ -54,12 +54,18 @@ class DistilledPipeline:
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
         torch_compile: bool = False,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
         self.dtype = torch.bfloat16
 
         self.prompt_encoder = PromptEncoder(
-            distilled_checkpoint_path, gemma_root, self.dtype, self.device, registry=registry
+            distilled_checkpoint_path,
+            gemma_root,
+            self.dtype,
+            self.device,
+            registry=registry,
+            offload_mode=offload_mode,
         )
         self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
         self.stage = DiffusionStage(
@@ -70,6 +76,7 @@ class DistilledPipeline:
             quantization=quantization,
             registry=registry,
             torch_compile=torch_compile,
+            offload_mode=offload_mode,
         )
         self.upsampler = VideoUpsampler(
             distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -88,7 +95,6 @@ class DistilledPipeline:
         images: list[ImageConditioningInput],
         tiling_config: TilingConfig | None = None,
         enhance_prompt: bool = False,
-        streaming_prefetch_count: int | None = None,
         stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS,
         stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
     ) -> tuple[Iterator[torch.Tensor], Audio]:
@@ -102,7 +108,6 @@ class DistilledPipeline:
             [prompt],
             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
 
@@ -130,7 +135,6 @@ class DistilledPipeline:
             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.
@@ -167,7 +171,6 @@ class DistilledPipeline:
                 noise_scale=stage_2_sigmas[0].item(),
                 initial_latent=audio_state.latent,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -189,6 +192,7 @@ def main() -> None:
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
         torch_compile=args.compile,
+        offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -202,7 +206,6 @@ 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/hdr_ic_lora.py

@@ -0,0 +1,901 @@
+"""HDR IC-LoRA pipeline: two-stage video generation with HDR output.
+Extends the standard IC-LoRA pipeline with HDR decode via LogC3 inverse
+transform.  ``__call__`` returns a **linear HDR float** tensor
+``[f, h, w, c]``; tonemapping and EXR saving are the caller's
+responsibility.
+Text embeddings must be pre-computed externally (e.g. using
+``PromptEncoder`` from ``ltx_pipelines.utils.blocks`` with a Gemma text
+encoder) and saved as a ``.safetensors`` file with ``video_context``
+and ``audio_context`` tensors (via ``safetensors.torch.save_file``).
+The path is passed via ``text_embeddings_path``.
+Run as a script for batch inference::
+    python -m ltx_pipelines.hdr_ic_lora \\
+        --input ./videos/ \\
+        --output-dir ./hdr-output \\
+        --distilled-checkpoint-path /models/ltx-2.3-22b-distilled.safetensors \\
+        --spatial-upsampler-path /models/ltx-2.3-spatial-upscaler-x2-1.0.safetensors \\
+        --hdr-lora /path/to/hdr_lora.safetensors \\
+        --text-embeddings /path/to/hdr_scene_emb.safetensors \\
+        --num-frames 161
+Supports resolutions up to 4K (3840x2160 @ 121 frames on 80 GB,
+49 frames on 48 GB).  The caller is responsible for choosing a resolution
+and frame count that fits in GPU memory.  See ``--help`` for a reference
+table, or use ``ltx_pipelines.utils.vram_budget.max_frames_for_resolution``
+to query your specific configuration.
+"""
+
+import dataclasses
+import logging
+from dataclasses import replace
+from pathlib import Path
+
+import torch
+from einops import rearrange
+from safetensors import safe_open
+
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.patchifiers import VideoLatentPatchifier
+from ltx_core.conditioning import (
+    ConditioningItem,
+    VideoConditionByReferenceLatent,
+)
+from ltx_core.hdr import apply_hdr_decode_postprocess
+from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.loader.registry import Registry
+from ltx_core.loader.sd_ops import LTXV_LORA_COMFY_RENAMING_MAP
+from ltx_core.modality_tiling import VideoModalityTilingHelper
+from ltx_core.model.video_vae import TilingConfig, VideoEncoder
+from ltx_core.quantization import QuantizationPolicy
+from ltx_core.tiling import DimensionTilingConfig, TileCountConfig
+from ltx_core.tools import VideoLatentTools
+from ltx_core.types import VideoLatentShape, VideoPixelShape
+from ltx_pipelines.utils.blocks import (
+    DiffusionStage,
+    ImageConditioner,
+    VideoDecoder,
+    VideoUpsampler,
+)
+from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
+from ltx_pipelines.utils.denoisers import SimpleDenoiser
+from ltx_pipelines.utils.helpers import get_device, modality_from_latent_state
+from ltx_pipelines.utils.media_io import ResizeMode, align_resolution, load_video_conditioning_hdr
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+
+DEFAULT_NUM_FRAMES = 161
+MIN_RESOLUTION = 64
+ALIGNMENT_DIVISOR = 64
+
+# Conditioning videos whose spatial resolution (H x W) exceeds this value are
+# encoded with the tiled encoder. The default (512 x 768) is suitable for
+# H100-80GB. On lower-VRAM GPUs pass tiled_vae_encode_pixel_threshold=256*256
+# to the pipeline constructor.
+TILED_VAE_ENCODE_PIXEL_THRESHOLD = 512 * 768
+
+_DEFAULT_QUANTIZATION = QuantizationPolicy.fp8_cast()
+
+# Default stage-2 configuration: one refinement phase with modest 2-way tiling
+# in every dimension and a short 2-step distilled sigma schedule.
+_S2 = STAGE_2_DISTILLED_SIGMA_VALUES
+
+_TILED_2F2H2W_OV8_6 = TileCountConfig(
+    frames=DimensionTilingConfig(2, 8),
+    height=DimensionTilingConfig(2, 6),
+    width=DimensionTilingConfig(2, 6),
+)
+
+STAGE2_TILINGS = [_TILED_2F2H2W_OV8_6]
+STAGE2_SIGMAS = [[_S2[0], _S2[1], 0.0]]
+STAGE2_USE_IC_LORA = [True]
+
+
+def _clamp_dim_tiling(cfg: DimensionTilingConfig, dim_size: int, axis: str) -> DimensionTilingConfig:
+    """Clamp a single dim's tile count and overlap to the latent's extent.
+    ``split_by_count`` requires ``overlap < tile_size``; with
+    ``tile_size = (dim_size + overlap*(n-1)) // n`` this reduces to
+    ``overlap <= dim_size - n``. When the configured overlap exceeds this
+    bound it is clamped; if the latent is too small to hold ``n`` tiles
+    at all, tiling falls back to a single tile on this axis.
+    """
+    n = cfg.num_tiles
+    if n <= 1:
+        return cfg
+    if dim_size < n:
+        logger.warning(
+            "%s tiling: dim_size=%d < num_tiles=%d; falling back to 1 tile on this axis.",
+            axis,
+            dim_size,
+            n,
+        )
+        return DimensionTilingConfig(1, 0)
+    max_overlap = dim_size - n
+    if cfg.overlap <= max_overlap:
+        return cfg
+    logger.warning(
+        "%s tiling: overlap=%d exceeds latent bound (%d); clamping to %d.",
+        axis,
+        cfg.overlap,
+        max_overlap,
+        max_overlap,
+    )
+    return DimensionTilingConfig(n, max_overlap)
+
+
+def _clamp_tile_to_latent(tiling: TileCountConfig, latent_shape: tuple[int, int, int]) -> TileCountConfig:
+    """Clamp frame, height, and width tilings to the latent's extents.
+    ``latent_shape`` is ``(F, H, W)`` in latent units.
+    """
+    f, h, w = latent_shape
+    return replace(
+        tiling,
+        frames=_clamp_dim_tiling(tiling.frames, f, "Frame"),
+        height=_clamp_dim_tiling(tiling.height, h, "Height"),
+        width=_clamp_dim_tiling(tiling.width, w, "Width"),
+    )
+
+
+# Default tiling config (spatial tile 1280 px, overlap 256 px; temporal 32
+# frames, overlap 16). On GPUs with < 80 GB VRAM you may need to shrink
+# the spatial tile size (e.g. 768) to avoid OOM during VAE decode.
+DEFAULT_SPATIAL_TILE = 1280
+DEFAULT_SPATIAL_OVERLAP = 256
+DEFAULT_TEMPORAL_TILE = 32
+DEFAULT_TEMPORAL_OVERLAP = 16
+
+
+# ---------------------------------------------------------------------------
+# HDR LoRA config
+# ---------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class HdrLoraConfig:
+    """Explicit HDR LoRA parameters.
+    Read from LoRA safetensors metadata by :func:`read_hdr_lora_config`, or
+    constructed manually for testing.
+    """
+
+    hdr_transform: str = "logc3"
+    reference_downscale_factor: int = 1
+
+
+def read_hdr_lora_config(lora_path: str) -> HdrLoraConfig | None:
+    """Read HDR config from LoRA safetensors metadata.
+    Returns ``None`` when the LoRA has no HDR metadata.
+    """
+    try:
+        with safe_open(lora_path, framework="pt") as f:
+            metadata = f.metadata() or {}
+    except (OSError, ValueError) as e:
+        logger.warning("Failed to read metadata from LoRA file '%s': %s", lora_path, e)
+        return None
+
+    hdr_transform = metadata.get("hdr_transform", "")
+    has_hdr = bool(hdr_transform or metadata.get("use_hdr_transform"))
+    if not has_hdr:
+        return None
+
+    transform = hdr_transform if hdr_transform and hdr_transform != "true" else "logc3"
+    scale = int(metadata.get("reference_downscale_factor", 1))
+    return HdrLoraConfig(hdr_transform=transform, reference_downscale_factor=scale)
+
+
+# ---------------------------------------------------------------------------
+# Pipeline
+# ---------------------------------------------------------------------------
+
+
+class HDRICLoraPipeline:
+    """Two-stage IC-LoRA pipeline with HDR support.
+    Same two-stage architecture as ICLoraPipeline (half-res generation + 2x
+    upscale refinement), with HDR decode via LogC3 inverse.
+    ``__call__`` returns a **linear HDR float** tensor ``[f, h, w, c]``.
+    Tonemapping and EXR saving are the caller's responsibility.
+    """
+
+    def __init__(
+        self,
+        distilled_checkpoint_path: str,
+        spatial_upsampler_path: str,
+        hdr_lora: str | Path,
+        text_embeddings_path: str | Path,
+        device: torch.device | None = None,
+        quantization: QuantizationPolicy = _DEFAULT_QUANTIZATION,
+        registry: Registry | None = None,
+        hdr_lora_config: HdrLoraConfig | None = None,
+        tiled_vae_encode_pixel_threshold: int = TILED_VAE_ENCODE_PIXEL_THRESHOLD,
+        offload_mode: OffloadMode = OffloadMode.NONE,
+    ):
+        """
+        Args:
+            distilled_checkpoint_path: Path to the distilled model checkpoint.
+            spatial_upsampler_path: Path to the spatial upsampler checkpoint.
+            hdr_lora: Path to the HDR IC-LoRA ``.safetensors`` file.
+            text_embeddings_path: Path to pre-computed text embeddings
+                (``.safetensors`` file with ``video_context`` and
+                ``audio_context`` tensors).
+            device: Target device. Auto-detected when ``None``.
+            quantization: Quantization policy. Defaults to ``fp8_cast``.
+            registry: Optional model registry for caching loaded components.
+            hdr_lora_config: Explicit HDR LoRA config override. When ``None``,
+                auto-detected from LoRA safetensors metadata.
+            tiled_vae_encode_pixel_threshold: Conditioning videos whose spatial
+                area (H x W) exceeds this value are encoded with the tiled
+                encoder. Default ``512 * 768`` is suitable for 80 GB GPUs.
+                Use ``256 * 256`` on GPUs with less VRAM.
+            offload_mode: Weight offloading strategy for diffusion stages.
+        """
+        self.device = device or get_device()
+        self._tiled_vae_encode_threshold = tiled_vae_encode_pixel_threshold
+        if offload_mode != OffloadMode.NONE and quantization is not None:
+            logger.info("Offload mode enabled — disabling quantization (not supported with layer streaming).")
+            quantization = None
+        self.dtype = torch.bfloat16
+
+        lora_path = str(Path(hdr_lora).resolve())
+        loras = (LoraPathStrengthAndSDOps(lora_path, 1.0, LTXV_LORA_COMFY_RENAMING_MAP),)
+
+        # Load pre-computed text embeddings from safetensors.
+        emb_path = Path(text_embeddings_path)
+        logger.info("Loading text embeddings from %s", emb_path)
+        with safe_open(emb_path, framework="pt", device=str(self.device)) as f:
+            self.text_embeddings: tuple[torch.Tensor, torch.Tensor] = (
+                f.get_tensor("video_context"),
+                f.get_tensor("audio_context"),
+            )
+
+        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=loras,
+            quantization=quantization,
+            registry=registry,
+            offload_mode=offload_mode,
+        )
+        self.stage_2 = DiffusionStage(
+            distilled_checkpoint_path,
+            self.dtype,
+            self.device,
+            loras=loras,
+            quantization=quantization,
+            registry=registry,
+            offload_mode=offload_mode,
+        )
+        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)
+
+        # HDR config: explicit override, or auto-detect from LoRA metadata.
+        if hdr_lora_config is not None:
+            self._hdr_config: HdrLoraConfig | None = hdr_lora_config
+        else:
+            self._hdr_config = read_hdr_lora_config(lora_path)
+
+        if self._hdr_config is not None:
+            logger.info("[HDR IC-LoRA] HDR mode enabled (%s decode)", self._hdr_config.hdr_transform)
+
+    @property
+    def hdr_transform(self) -> str:
+        """Active HDR transform name (defaults to 'logc3')."""
+        return self._hdr_config.hdr_transform if self._hdr_config is not None else "logc3"
+
+    @property
+    def reference_downscale_factor(self) -> int:
+        """Reference video downscale factor from HDR LoRA config."""
+        return self._hdr_config.reference_downscale_factor if self._hdr_config is not None else 1
+
+    def __call__(  # noqa: PLR0913
+        self,
+        seed: int,
+        height: int,
+        width: int,
+        num_frames: int,
+        frame_rate: float,
+        video_conditioning: list[tuple[str, float]],
+        tiling_config: TilingConfig | None = None,
+        high_quality_hdr: bool = False,
+        stage2_tilings: list[TileCountConfig] | None = None,
+        stage2_sigmas: list[list[float]] | None = None,
+        stage2_use_ic_lora: list[bool] | None = None,
+    ) -> torch.Tensor:
+        """Generate video with IC-LoRA conditioning and HDR output.
+        Returns a linear HDR float tensor ``[f, h, w, c]``.
+        Args:
+            seed: Random seed for reproducibility.
+            height: Desired output video height in pixels. Aligned internally
+                to the nearest multiple of 64 (rounded up). Decoded output is
+                cropped back to this size.
+            width: Desired output video width in pixels. Same alignment rules
+                as *height*.
+            num_frames: Number of frames to generate.
+            frame_rate: Output video frame rate.
+            video_conditioning: List of (path, strength) tuples for IC-LoRA video conditioning.
+            high_quality_hdr: High-quality HDR mode. Duplicates each conditioning
+                frame and generates at 2x frame count, then keeps every other
+                output frame. Reduces temporal artifacts at the cost of ~2x
+                generation time.
+        Returns:
+            Linear HDR float tensor ``[f, h, w, c]``.
+        """
+        # In high-quality HDR mode, generate 2*N - 1 frames internally
+        # (satisfies (n-1)%8==0 when N itself does), then keep every other frame.
+        if high_quality_hdr:
+            gen_num_frames = 2 * num_frames - 1
+            logger.info("[HDR IC-LoRA] High-quality HDR: %d -> %d internal frames", num_frames, gen_num_frames)
+        else:
+            gen_num_frames = num_frames
+        gen_w, gen_h, crop_w, crop_h = align_resolution(
+            width, height, ResizeMode.REFLECT_PAD, divisor=ALIGNMENT_DIVISOR
+        )
+        if gen_h < MIN_RESOLUTION or gen_w < MIN_RESOLUTION:
+            raise ValueError(
+                f"Resolution ({width}x{height}) is too small after alignment "
+                f"(got {gen_w}x{gen_h}, need at least {MIN_RESOLUTION}x{MIN_RESOLUTION})."
+            )
+        needs_crop = crop_w != gen_w or crop_h != gen_h
+        if needs_crop:
+            logger.info(
+                "[HDR IC-LoRA] Aligned %dx%d -> %dx%d, will crop to %dx%d",
+                width,
+                height,
+                gen_w,
+                gen_h,
+                crop_w,
+                crop_h,
+            )
+
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        noiser = GaussianNoiser(generator=generator)
+
+        video_context, _ = self.text_embeddings
+
+        # Stage 1: Initial low resolution video generation.
+        s1_w, s1_h = gen_w // 2, gen_h // 2
+
+        stage_1_conditionings = self.image_conditioner(
+            lambda enc: self._create_conditionings(
+                video_conditioning=video_conditioning,
+                height=s1_h,
+                width=s1_w,
+                video_encoder=enc,
+                num_frames=gen_num_frames,
+                tiling_config=tiling_config,
+                high_quality_hdr=high_quality_hdr,
+            )
+        )
+
+        stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
+
+        # HDR is video-only: skip the audio stream to avoid denoising 5B audio params.
+        video_state, _ = self.stage_1(
+            denoiser=SimpleDenoiser(video_context, None),
+            sigmas=stage_1_sigmas,
+            noiser=noiser,
+            width=s1_w,
+            height=s1_h,
+            frames=gen_num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(
+                context=video_context,
+                conditionings=stage_1_conditionings,
+            ),
+        )
+
+        if stage2_tilings is None:
+            stage2_tilings = list(STAGE2_TILINGS)
+        if stage2_sigmas is None:
+            stage2_sigmas = [list(s) for s in STAGE2_SIGMAS]
+        if stage2_use_ic_lora is None:
+            stage2_use_ic_lora = list(STAGE2_USE_IC_LORA)
+        if not (len(stage2_tilings) == len(stage2_sigmas) == len(stage2_use_ic_lora)):
+            raise ValueError("stage2_tilings, stage2_sigmas, and stage2_use_ic_lora must have equal length")
+
+        # Stage 2: Upsample and refine at full resolution.
+        upscaled_video_latent = self.upsampler(video_state.latent[:1])
+
+        stage_2_conditionings = self.image_conditioner(
+            lambda enc: self._create_conditionings(
+                video_conditioning=video_conditioning,
+                height=gen_h,
+                width=gen_w,
+                video_encoder=enc,
+                num_frames=gen_num_frames,
+                tiling_config=tiling_config,
+                high_quality_hdr=high_quality_hdr,
+            )
+        )
+        # video_tools is required by TiledDataParallelBuilder when stage_2 is
+        # wrapped for multi-GPU
+        stage2_video_tools = VideoLatentTools(
+            VideoLatentPatchifier(patch_size=1),
+            VideoLatentShape.from_pixel_shape(
+                VideoPixelShape(
+                    batch=1,
+                    frames=gen_num_frames,
+                    height=gen_h,
+                    width=gen_w,
+                    fps=frame_rate,
+                )
+            ),
+            frame_rate,
+        )
+        with self.stage_2.model_context(video_tools=stage2_video_tools) as transformer:
+            phase_latent = upscaled_video_latent
+            for phase_idx, (tiling, sigmas_list, use_ic) in enumerate(
+                zip(stage2_tilings, stage2_sigmas, stage2_use_ic_lora, strict=True)
+            ):
+                diffusion_tiling = _clamp_tile_to_latent(tiling, tuple(phase_latent.shape[2:5]))
+                conditionings = stage_2_conditionings if use_ic else []
+                sigma_t = torch.tensor(sigmas_list, dtype=torch.float32, device=self.device)
+                logger.info(
+                    "[Stage 2 / phase %d] sigmas=%s ic_lora=%s tiling_h=%s tiling_w=%s",
+                    phase_idx,
+                    sigmas_list,
+                    use_ic,
+                    diffusion_tiling.height,
+                    diffusion_tiling.width,
+                )
+                phase_latent = self._run_stage2_phase(
+                    transformer=transformer,
+                    latent=phase_latent,
+                    conditionings=conditionings,
+                    tiling=diffusion_tiling,
+                    sigmas=sigma_t,
+                    v_ctx=video_context,
+                    frame_rate=frame_rate,
+                    seed=seed,
+                )
+
+        final_video_latent = phase_latent
+
+        crop_size = (crop_w, crop_h) if needs_crop else None
+        return self._decode_video(
+            final_video_latent,
+            tiling_config,
+            generator,
+            crop_size,
+            high_quality_hdr=high_quality_hdr,
+        )
+
+    def _run_stage2_phase(
+        self,
+        transformer: object,
+        latent: torch.Tensor,
+        conditionings: list[ConditioningItem],
+        tiling: TileCountConfig,
+        sigmas: torch.Tensor,
+        v_ctx: torch.Tensor,
+        frame_rate: float,
+        seed: int,
+    ) -> torch.Tensor:
+        """Run one stage-2 denoising phase with optional IC-LoRA conditioning.
+        Each tile calls ``stage_2.run()`` with a tile-sized ``ModalitySpec`` for
+        video only (audio is omitted entirely for HDR). IC-LoRA conditionings
+        are sliced spatially to match each tile's extent.
+        """
+        batch, n_channels, n_frames, n_height, n_width = latent.shape
+        full_shape = VideoLatentShape(batch=batch, channels=n_channels, frames=n_frames, height=n_height, width=n_width)
+        full_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), full_shape, frame_rate)
+        helper = VideoModalityTilingHelper(tiling, full_tools)
+
+        ref_initial = full_tools.create_initial_state(device=self.device, dtype=self.dtype)
+        ref_modality = modality_from_latent_state(ref_initial, v_ctx, sigmas[0])
+        n_gen = full_tools.target_shape.token_count()
+        blend_output = torch.zeros(batch, n_gen, n_channels, device=self.device, dtype=self.dtype)
+        patchifier = VideoLatentPatchifier(patch_size=1)
+        df = self.reference_downscale_factor
+
+        for tile_idx, tile in enumerate(helper.tiles):
+            _, ctx = helper.tile_modality(ref_modality, tile, normalize_positions=True)
+            frame_s, height_s, width_s = tile.in_coords
+            tile_h = height_s.stop - height_s.start
+            tile_w = width_s.stop - width_s.start
+            tile_f = frame_s.stop - frame_s.start
+
+            tile_conditionings = [
+                VideoConditionByReferenceLatent(
+                    latent=cond.latent[
+                        :,
+                        :,
+                        frame_s,
+                        slice(height_s.start // df, height_s.stop // df),
+                        slice(width_s.start // df, width_s.stop // df),
+                    ].to(device=self.device, dtype=self.dtype),
+                    downscale_factor=cond.downscale_factor,
+                    strength=cond.strength,
+                )
+                for cond in conditionings
+            ]
+
+            tile_video_state, _ = self.stage_2.run(
+                transformer=transformer,
+                denoiser=SimpleDenoiser(v_ctx, None),
+                sigmas=sigmas,
+                noiser=GaussianNoiser(generator=torch.Generator(device=self.device).manual_seed(seed + tile_idx)),
+                width=tile_w * 32,
+                height=tile_h * 32,
+                frames=(tile_f - 1) * 8 + 1,
+                fps=frame_rate,
+                video=ModalitySpec(
+                    context=v_ctx,
+                    conditionings=tile_conditionings,
+                    noise_scale=sigmas[0].item(),
+                    initial_latent=latent[:, :, frame_s, height_s, width_s].to(device=self.device, dtype=self.dtype),
+                ),
+            )
+
+            tile_tokens = patchifier.patchify(tile_video_state.latent)
+            blend_output = helper.blend(tile_tokens, tile, ctx, blend_output)
+
+        return full_tools.unpatchify(replace(ref_initial, latent=blend_output)).latent
+
+    def _decode_video(
+        self,
+        latent: torch.Tensor,
+        tiling_config: TilingConfig | None,
+        generator: torch.Generator,
+        crop_size: tuple[int, int] | None = None,
+        *,
+        high_quality_hdr: bool = False,
+    ) -> torch.Tensor:
+        """Decode latent to HDR video, optionally cropping to target size.
+        Args:
+            crop_size: ``(width, height)`` to crop decoded frames to, or
+                ``None`` to skip cropping.
+            high_quality_hdr: When True, keep only every other frame (undoes the
+                2x generation applied during high-quality HDR mode).
+        Returns:
+            Linear HDR float tensor ``[f, h, w, c]``.
+        """
+        # Cast to float32 so tiled-decode accumulation buffers and blending
+        # masks run in full precision, avoiding bfloat16 seam artifacts.
+        # apply_hdr_decode_postprocess expects float32 [0, 1].
+        latent = latent.float()
+        decoded = torch.cat(
+            [chunk.float() for chunk in self.video_decoder(latent, tiling_config, generator)],
+            dim=0,
+        )
+        decoded = rearrange(decoded, "f h w c -> 1 c f h w")
+        hdr = apply_hdr_decode_postprocess(decoded, transform=self.hdr_transform)
+        del decoded
+        out = rearrange(hdr[0], "c f h w -> f h w c")
+        if crop_size is not None:
+            out = out[:, : crop_size[1], : crop_size[0], :]
+        if high_quality_hdr:
+            out = out[::2]
+        return out
+
+    def _create_conditionings(
+        self,
+        video_conditioning: list[tuple[str, float]],
+        height: int,
+        width: int,
+        num_frames: int,
+        video_encoder: VideoEncoder,
+        tiling_config: TilingConfig | None = None,
+        high_quality_hdr: bool = False,
+    ) -> list[ConditioningItem]:
+        """Create conditioning items for video generation."""
+        conditionings: list[ConditioningItem] = []
+
+        scale = self.reference_downscale_factor
+        if scale != 1 and (height % scale != 0 or width % scale != 0):
+            raise ValueError(
+                f"Output dimensions ({height}x{width}) must be divisible by reference_downscale_factor ({scale})"
+            )
+        ref_height = height // scale
+        ref_width = width // scale
+
+        # In high-quality HDR mode, load half the frames then duplicate each one.
+        load_frame_cap = (num_frames + 1) // 2 if high_quality_hdr else num_frames
+
+        for video_path, strength in video_conditioning:
+            video = torch.cat(
+                list(
+                    load_video_conditioning_hdr(
+                        video_path=video_path,
+                        height=ref_height,
+                        width=ref_width,
+                        frame_cap=load_frame_cap,
+                        dtype=self.dtype,
+                        device=self.device,
+                        hdr_transform=self.hdr_transform,
+                        resize_mode=ResizeMode.REFLECT_PAD,
+                    )
+                ),
+                dim=2,
+            )
+            if high_quality_hdr:
+                video = video.repeat_interleave(2, dim=2)[:, :, :num_frames, :, :]
+            if tiling_config is not None and ref_height * ref_width > self._tiled_vae_encode_threshold:
+                encoded_video = video_encoder.tiled_encode(video, tiling_config)
+            else:
+                encoded_video = video_encoder(video)
+
+            cond = VideoConditionByReferenceLatent(
+                latent=encoded_video,
+                downscale_factor=scale,
+                strength=strength,
+            )
+            conditionings.append(cond)
+
+        if video_conditioning:
+            logger.info("[HDR IC-LoRA] Added %d video conditioning(s)", len(video_conditioning))
+
+        return conditionings
+
+
+# ---------------------------------------------------------------------------
+# CLI helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_tiling_config(
+    spatial_tile: int = DEFAULT_SPATIAL_TILE,
+    spatial_overlap: int = DEFAULT_SPATIAL_OVERLAP,
+    temporal_tile: int = DEFAULT_TEMPORAL_TILE,
+    temporal_overlap: int = DEFAULT_TEMPORAL_OVERLAP,
+) -> TilingConfig:
+    """Build a TilingConfig from explicit sizes.
+    The defaults (1280 px spatial tile, 256 px overlap; 32 temporal frames,
+    16 overlap) are suitable for H100-80 GB.  On GPUs with less VRAM,
+    reduce the spatial tile size (e.g. ``spatial_tile=768``).
+    """
+    from ltx_core.model.video_vae.tiling import SpatialTilingConfig, TemporalTilingConfig  # noqa: PLC0415
+
+    return TilingConfig(
+        spatial_config=SpatialTilingConfig(tile_size_in_pixels=spatial_tile, tile_overlap_in_pixels=spatial_overlap),
+        temporal_config=TemporalTilingConfig(
+            tile_size_in_frames=temporal_tile,
+            tile_overlap_in_frames=temporal_overlap,
+        ),
+    )
+
+
+_VIDEO_SUFFIXES = {".mp4", ".mov"}
+
+
+def _collect_videos(input_path: Path) -> list[Path]:
+    """Return a list of .mp4/.mov files from *input_path* (file or directory)."""
+    if input_path.is_file():
+        return [input_path]
+    if input_path.is_dir():
+        return sorted(p for p in input_path.iterdir() if p.is_file() and p.suffix.lower() in _VIDEO_SUFFIXES)
+    logger.error("Input %s is not a file or directory", input_path)
+    return []
+
+
+def _process_single_video(  # noqa: PLR0913
+    pipeline: HDRICLoraPipeline,
+    video_path: Path,
+    vid_w: int,
+    vid_h: int,
+    num_frames: int,
+    frame_rate: float,
+    output_dir: Path,
+    tiling_config: TilingConfig,
+    seed: int,
+    skip_mp4: bool,
+    exr_half: bool,
+    exr_executor: "ThreadPoolExecutor",  # noqa: F821
+    exr_futures: list,
+    high_quality_hdr: bool = False,
+) -> None:
+    """Run inference on a single video: generate EXR frames + optional H.264 .mp4 preview."""
+    import gc  # noqa: PLC0415
+    import time  # noqa: PLC0415
+
+    from ltx_pipelines.utils.media_io import encode_exr_sequence_to_mp4, save_exr_tensor  # noqa: PLC0415
+
+    output_mp4 = output_dir / f"{video_path.stem}.mp4"
+    exr_dir = output_dir / f"{video_path.stem}_exr"
+
+    t0 = time.time()
+    hdr_video = pipeline(
+        seed=seed,
+        height=vid_h,
+        width=vid_w,
+        num_frames=num_frames,
+        frame_rate=frame_rate,
+        video_conditioning=[(str(video_path), 1.0)],
+        tiling_config=tiling_config,
+        high_quality_hdr=high_quality_hdr,
+    )
+
+    exr_dir.mkdir(parents=True, exist_ok=True)
+    for j in range(hdr_video.shape[0]):
+        frame_cpu = hdr_video[j].cpu().clone()
+        path = exr_dir / f"frame_{j:05d}.exr"
+        exr_futures.append(exr_executor.submit(save_exr_tensor, frame_cpu, str(path), exr_half))
+
+    del hdr_video
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    if not skip_mp4:
+        # Wait for EXR saves to finish before encoding.
+        for fut in exr_futures:
+            fut.result()
+        logger.info("Encoding H.264 sRGB preview: %s", video_path.name)
+        encode_exr_sequence_to_mp4(exr_dir, output_mp4, frame_rate)
+
+    elapsed = time.time() - t0
+    logger.info("Decode + encode: %.1fs | %s", elapsed, output_mp4)
+
+
+# ---------------------------------------------------------------------------
+# CLI entry point
+# ---------------------------------------------------------------------------
+
+
+def _build_arg_parser() -> "argparse.ArgumentParser":  # noqa: F821
+    """Build the argument parser for HDR IC-LoRA batch inference."""
+    import argparse  # noqa: PLC0415
+
+    parser = argparse.ArgumentParser(
+        description="HDR IC-LoRA inference: EXR frames + tonemapped ProRes .mov.",
+        epilog="""\
+Resolution & frame constraints
+------------------------------
+  * Width and height must each be divisible by 32.
+  * Frame count must satisfy (frames - 1) %% 8 == 0.
+    Valid counts: 1, 9, 17, 25, ..., 121, 129, 137, 145, 153, 161.
+
+Max frames by resolution (fp8_cast, bfloat16 VAE, tiled decode)
+---------------------------------------------------------------
+  Resolution       80 GB (H100)    48 GB (A6000)
+  ------------------------------------------------
+   720p 1280x720    161+ frames     161+ frames
+  1080p 1920x1080   161+ frames     161+ frames
+  2K    2048x1080   161+ frames     161+ frames
+  1440p 2560x1440   161+ frames     137 frames
+  4K    3840x2160   121 frames       49 frames
+  4K    4096x2160   105 frames       49 frames
+
+  Estimates from ltx_pipelines.utils.vram_budget. Run
+    python -c "from ltx_pipelines.utils.vram_budget import \\
+      max_frames_for_resolution as mf; print(mf(W, H, vram_gb=GB))"
+  to check your specific resolution and GPU.
+
+  * The tiled-encode threshold (%(tiled_threshold)s px) and the default
+    tiling config (%(stile)s px spatial tile) are tuned for 80 GB.
+    On lower-VRAM GPUs pass --spatial-tile 768 (or smaller).
+"""
+        % {"tiled_threshold": TILED_VAE_ENCODE_PIXEL_THRESHOLD, "stile": DEFAULT_SPATIAL_TILE},
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("--input", required=True, help="Single .mp4 or directory of .mp4 videos.")
+    parser.add_argument("--output-dir", required=True, help="Directory for .mov and EXR folders.")
+    parser.add_argument("--hdr-lora", required=True, help="HDR IC-LoRA .safetensors file.")
+    parser.add_argument("--text-embeddings", required=True, help="Pre-computed text embeddings (.safetensors file).")
+    parser.add_argument("--distilled-checkpoint-path", required=True, help="Distilled model checkpoint (.safetensors).")
+    parser.add_argument("--spatial-upsampler-path", required=True, help="Spatial upsampler (.safetensors).")
+    parser.add_argument(
+        "--num-frames",
+        type=int,
+        default=DEFAULT_NUM_FRAMES,
+        help=f"Number of output frames.  Must satisfy (n-1) %% 8 == 0 (default: {DEFAULT_NUM_FRAMES}).",
+    )
+    parser.add_argument(
+        "--spatial-tile",
+        type=int,
+        default=DEFAULT_SPATIAL_TILE,
+        help=f"Spatial tile size in pixels for tiled VAE decode (default: {DEFAULT_SPATIAL_TILE}). "
+        "Reduce on lower-VRAM GPUs (e.g. 768 for 48 GB).",
+    )
+    parser.add_argument("--skip-mp4", action="store_true", help="Skip H.264 MP4 encoding, only produce EXR.")
+    parser.add_argument("--exr-half", action="store_true", help="Save EXR as float16.")
+    parser.add_argument("--seed", type=int, default=10, help="Random seed (default: 10).")
+    parser.add_argument(
+        "--offload",
+        dest="offload_mode",
+        type=OffloadMode,
+        default=OffloadMode.NONE,
+        choices=list(OffloadMode),
+        help=(
+            "Weight offloading strategy. "
+            "'none' keeps all weights on GPU (default). "
+            "'cpu' pins weights in CPU RAM, streams to GPU per layer. "
+            "'disk' reads weights from disk on demand (lowest memory). "
+            "Example: --offload cpu"
+        ),
+    )
+    parser.add_argument(
+        "--high-quality",
+        action="store_true",
+        help="High-quality HDR mode. Generates at 2x frame count internally "
+        "and keeps every other frame for smoother output. ~2x slower.",
+    )
+    return parser
+@torch.inference_mode()
+def main() -> None:
+    """Batch HDR IC-LoRA inference: per-frame EXR + tonemapped ProRes .mov."""
+    import time  # noqa: PLC0415
+    from concurrent.futures import ThreadPoolExecutor  # noqa: PLC0415
+
+    from ltx_pipelines.utils.media_io import get_videostream_metadata  # noqa: PLC0415
+
+    logging.basicConfig(level=logging.INFO)
+
+    args = _build_arg_parser().parse_args()
+    high_quality = args.high_quality
+    num_frames = args.num_frames
+
+    tiling_config = _make_tiling_config(spatial_tile=args.spatial_tile)
+
+    input_path = Path(args.input)
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    videos = _collect_videos(input_path)
+    if not videos:
+        logger.error("No valid videos to process.")
+        return
+    logger.info("Found %d video(s), generating %d frames each", len(videos), num_frames)
+
+    logger.info("Loading pipeline...")
+    pipeline = HDRICLoraPipeline(
+        distilled_checkpoint_path=args.distilled_checkpoint_path,
+        spatial_upsampler_path=args.spatial_upsampler_path,
+        hdr_lora=args.hdr_lora,
+        text_embeddings_path=args.text_embeddings,
+        offload_mode=args.offload_mode,
+    )
+    logger.info("Pipeline loaded.")
+
+    exr_executor = ThreadPoolExecutor(max_workers=4)
+    exr_futures: list = []
+
+    total_t0 = time.time()
+    successes = 0
+
+    for i, video_path in enumerate(videos, 1):
+        meta = get_videostream_metadata(str(video_path))
+        vid_w, vid_h = meta.width, meta.height
+        logger.info("%s", "=" * 60)
+        logger.info("[%d/%d] %s  (%dx%d, %df)", i, len(videos), video_path.name, vid_w, vid_h, num_frames)
+
+        _process_single_video(
+            pipeline=pipeline,
+            video_path=video_path,
+            vid_w=vid_w,
+            vid_h=vid_h,
+            num_frames=num_frames,
+            frame_rate=meta.fps,
+            output_dir=output_dir,
+            tiling_config=tiling_config,
+            seed=args.seed,
+            skip_mp4=args.skip_mp4,
+            exr_half=args.exr_half,
+            exr_executor=exr_executor,
+            exr_futures=exr_futures,
+            high_quality_hdr=high_quality,
+        )
+        successes += 1
+
+    infer_elapsed = time.time() - total_t0
+    logger.info("%s", "=" * 60)
+    logger.info("All inference done in %.0fs  (%d/%d OK)", infer_elapsed, successes, len(videos))
+
+    if exr_futures:
+        t0 = time.time()
+        logger.info("Waiting for %d EXR saves...", len(exr_futures))
+        for fut in exr_futures:
+            fut.result()
+        exr_wait = time.time() - t0
+        if exr_wait > 0.1:
+            logger.info("EXR save wait: %.1fs", exr_wait)
+
+    logger.info("Total wall time: %.0fs", time.time() - total_t0)
+
+
+if __name__ == "__main__":
+    main()

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

@@ -2,20 +2,18 @@ import logging
 from collections.abc import Iterator
 
 import torch
-from einops import rearrange
-from safetensors import safe_open
 
 from ltx_core.components.noisers import GaussianNoiser
-from ltx_core.conditioning import (
+from ltx_core.conditioning import ConditioningItem
-    ConditioningItem,
-    ConditioningItemAttentionStrengthWrapper,
-    VideoConditionByReferenceLatent,
-)
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
 from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
-from ltx_core.types import Audio, VideoLatentShape, VideoPixelShape
+from ltx_core.types import Audio, VideoPixelShape
+from ltx_pipelines.iclora_utils import (
+    append_ic_lora_reference_video_conditionings,
+    read_lora_reference_downscale_factor,
+)
 from ltx_pipelines.utils.args import (
     ImageConditioningInput,
     VideoConditioningAction,
@@ -39,7 +37,7 @@ from ltx_pipelines.utils.constants import (
 from ltx_pipelines.utils.denoisers import SimpleDenoiser
 from ltx_pipelines.utils.helpers import assert_resolution, combined_image_conditionings, get_device
 from ltx_pipelines.utils.media_io import decode_video_by_frame, encode_video, video_preprocess
-from ltx_pipelines.utils.types import ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class ICLoraPipeline:
@@ -63,12 +61,18 @@ class ICLoraPipeline:
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
         torch_compile: bool = False,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
         self.dtype = torch.bfloat16
 
         self.prompt_encoder = PromptEncoder(
-            distilled_checkpoint_path, gemma_root, self.dtype, self.device, registry=registry
+            distilled_checkpoint_path,
+            gemma_root,
+            self.dtype,
+            self.device,
+            registry=registry,
+            offload_mode=offload_mode,
         )
         self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
         self.stage_1 = DiffusionStage(
@@ -79,6 +83,7 @@ class ICLoraPipeline:
             quantization=quantization,
             registry=registry,
             torch_compile=torch_compile,
+            offload_mode=offload_mode,
         )
         self.stage_2 = DiffusionStage(
             distilled_checkpoint_path,
@@ -88,6 +93,7 @@ class ICLoraPipeline:
             quantization=quantization,
             registry=registry,
             torch_compile=torch_compile,
+            offload_mode=offload_mode,
         )
         self.upsampler = VideoUpsampler(
             distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -100,7 +106,7 @@ class ICLoraPipeline:
         # so inference can resize reference videos to match training conditions.
         self.reference_downscale_factor = 1
         for lora in loras:
-            scale = _read_lora_reference_downscale_factor(lora.path)
+            scale = read_lora_reference_downscale_factor(lora.path)
             if scale != 1:
                 if self.reference_downscale_factor not in (1, scale):
                     raise ValueError(
@@ -125,7 +131,6 @@ class ICLoraPipeline:
         conditioning_attention_strength: float = 1.0,
         skip_stage_2: bool = False,
         conditioning_attention_mask: torch.Tensor | None = None,
-        streaming_prefetch_count: int | None = None,
         stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS,
         stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
     ) -> tuple[Iterator[torch.Tensor], Audio]:
@@ -175,7 +180,6 @@ class ICLoraPipeline:
             enhance_first_prompt=enhance_prompt,
             enhance_prompt_image=images[0][0] if len(images) > 0 else None,
             enhance_prompt_seed=seed,
-            streaming_prefetch_count=streaming_prefetch_count,
         )
         video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
 
@@ -219,7 +223,6 @@ class ICLoraPipeline:
             audio=ModalitySpec(
                 context=audio_context,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         if skip_stage_2:
@@ -264,7 +267,6 @@ class ICLoraPipeline:
                 noise_scale=stage_2_sigmas[0].item(),
                 initial_latent=audio_state.latent,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -305,104 +307,26 @@ class ICLoraPipeline:
             device=self.device,
         )
 
-        # Calculate scaled dimensions for reference video conditioning.
+        append_ic_lora_reference_video_conditionings(
-        # IC-LoRAs trained with downscaled reference videos expect the same ratio at inference.
+            conditionings,
-        scale = self.reference_downscale_factor
+            video_conditioning,
-        if scale != 1 and (height % scale != 0 or width % scale != 0):
+            height=height,
-            raise ValueError(
+            width=width,
-                f"Output dimensions ({height}x{width}) must be divisible by reference_downscale_factor ({scale})"
+            num_frames=num_frames,
-            )
+            video_encoder=video_encoder,
-        ref_height = height // scale
+            dtype=self.dtype,
-        ref_width = width // scale
+            device=self.device,
-
+            reference_downscale_factor=self.reference_downscale_factor,
-        for video_path, strength in video_conditioning:
+            conditioning_attention_strength=conditioning_attention_strength,
-            # Load video at scaled-down resolution (if scale > 1)
+            conditioning_attention_mask=conditioning_attention_mask,
-            frame_gen = decode_video_by_frame(path=video_path, frame_cap=num_frames, device=self.device)
+            tiling_config=None,
-            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)
-
-            # Build attention_mask for ConditioningItemAttentionStrengthWrapper
-            if conditioning_attention_mask is not None:
-                # Downsample pixel-space mask to latent space, then scale by strength
-                latent_mask = self._downsample_mask_to_latent(
-                    mask=conditioning_attention_mask,
-                    target_latent_shape=reference_video_shape,
-                )
-                attn_mask = latent_mask * conditioning_attention_strength
-            elif conditioning_attention_strength < 1.0:
-                # Use scalar strength only
-                attn_mask = conditioning_attention_strength
-            else:
-                attn_mask = None
-
-            cond = VideoConditionByReferenceLatent(
-                latent=encoded_video,
-                downscale_factor=scale,
-                strength=strength,
-            )
-            if attn_mask is not None:
-                cond = ConditioningItemAttentionStrengthWrapper(cond, attention_mask=attn_mask)
-            conditionings.append(cond)
 
         if video_conditioning:
-            logging.info(f"[IC-LoRA] Added {len(video_conditioning)} video conditioning(s)")
+            logging.info("[IC-LoRA] Added %d video conditioning(s)", len(video_conditioning))
 
         return conditionings
 
-    @staticmethod
-    def _downsample_mask_to_latent(
-        mask: torch.Tensor,
-        target_latent_shape: VideoLatentShape,
-    ) -> torch.Tensor:
-        """
-        Downsample a pixel-space mask to latent space using VAE scale factors.
-        Handles causal temporal downsampling: the first frame is kept separately
-        (temporal scale factor = 1 for the first frame), while the remaining
-        frames are downsampled by the VAE's temporal scale factor.
-        Args:
-            mask: Pixel-space mask of shape (B, 1, F_pixel, H_pixel, W_pixel).
-                Values in [0, 1].
-            target_latent_shape: Expected latent shape after VAE encoding.
-                Used to determine the target (F_latent, H_latent, W_latent).
-        Returns:
-            Flattened latent-space mask of shape (B, F_lat * H_lat * W_lat),
-            matching the patchifier's token ordering (f, h, w).
-        """
-        b = mask.shape[0]
-        f_lat = target_latent_shape.frames
-        h_lat = target_latent_shape.height
-        w_lat = target_latent_shape.width
-
-        # Step 1: Spatial downsampling (area interpolation per frame)
-        f_pix = mask.shape[2]
-        spatial_down = torch.nn.functional.interpolate(
-            rearrange(mask, "b 1 f h w -> (b f) 1 h w"),
-            size=(h_lat, w_lat),
-            mode="area",
-        )
-        spatial_down = rearrange(spatial_down, "(b f) 1 h w -> b 1 f h w", b=b)
-
-        # Step 2: Causal temporal downsampling
-        # First frame: kept as-is (causal VAE encodes first frame independently)
-        first_frame = spatial_down[:, :, :1, :, :]  # (B, 1, 1, H_lat, W_lat)
-
-        if f_pix > 1 and f_lat > 1:
-            # Remaining frames: downsample by temporal factor via group-mean
-            t = (f_pix - 1) // (f_lat - 1)  # temporal downscale factor
-            assert (f_pix - 1) % (f_lat - 1) == 0, (
-                f"Pixel frames ({f_pix}) not compatible with latent frames ({f_lat}): "
-                f"(f_pix - 1) must be divisible by (f_lat - 1)"
-            )
-            rest = rearrange(spatial_down[:, :, 1:, :, :], "b 1 (f t) h w -> b 1 f t h w", t=t)
-            rest = rest.mean(dim=3)  # (B, 1, F_lat-1, H_lat, W_lat)
-            latent_mask = torch.cat([first_frame, rest], dim=2)  # (B, 1, F_lat, H_lat, W_lat)
-        else:
-            latent_mask = first_frame
-
-        # Flatten to (B, F_lat * H_lat * W_lat) matching patchifier token order (f, h, w)
-        return rearrange(latent_mask, "b 1 f h w -> b (f h w)")
-
 
 @torch.inference_mode()
 def main() -> None:
@@ -463,6 +387,7 @@ def main() -> None:
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
         torch_compile=args.compile,
+        offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -479,7 +404,6 @@ def main() -> None:
         conditioning_attention_strength=conditioning_attention_strength,
         skip_stage_2=args.skip_stage_2,
         conditioning_attention_mask=conditioning_attention_mask,
-        streaming_prefetch_count=args.streaming_prefetch_count,
     )
 
     encode_video(
@@ -519,26 +443,5 @@ def _load_mask_video(
     return mask.clamp(0.0, 1.0)
 
 
-def _read_lora_reference_downscale_factor(lora_path: str) -> int:
-    """Read reference_downscale_factor from LoRA safetensors metadata.
-    Some IC-LoRA models are trained with reference videos at lower resolution than
-    the target output. This allows for more efficient training and can improve
-    generalization. The downscale factor indicates the ratio between target and
-    reference resolutions (e.g., factor=2 means reference is half the resolution).
-    Args:
-        lora_path: Path to the LoRA .safetensors file
-    Returns:
-        The reference downscale factor (1 if not specified in metadata, meaning
-        reference and target have the same resolution)
-    """
-    try:
-        with safe_open(lora_path, framework="pt") as f:
-            metadata = f.metadata() or {}
-            return int(metadata.get("reference_downscale_factor", 1))
-    except Exception as e:
-        logging.warning(f"Failed to read metadata from LoRA file '{lora_path}': {e}")
-        return 1
-
-
 if __name__ == "__main__":
     main()

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

@@ -0,0 +1,120 @@
+"""Shared IC-LoRA helpers: LoRA metadata, mask downsampling, reference-video conditioning.
+Used by ``ic_lora`` and ``lipdub`` (video reference path only). LipDub audio helpers live in ``lipdub.py``.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+from einops import rearrange
+from safetensors import safe_open
+
+from ltx_core.conditioning import (
+    ConditioningItem,
+    ConditioningItemAttentionStrengthWrapper,
+    VideoConditionByReferenceLatent,
+)
+from ltx_core.model.video_vae import TilingConfig, VideoEncoder
+from ltx_core.types import VideoLatentShape
+from ltx_pipelines.utils.media_io import decode_video_by_frame, video_preprocess
+
+
+def read_lora_reference_downscale_factor(lora_path: str) -> int:
+    """Read ``reference_downscale_factor`` from LoRA safetensors metadata (default 1)."""
+    try:
+        with safe_open(lora_path, framework="pt") as f:
+            metadata = f.metadata() or {}
+            return int(metadata.get("reference_downscale_factor", 1))
+    except Exception as e:
+        logging.warning("Failed to read metadata from LoRA file '%s': %s", lora_path, e)
+        return 1
+
+
+def downsample_mask_video_to_latent(
+    mask: torch.Tensor,
+    target_latent_shape: VideoLatentShape,
+) -> torch.Tensor:
+    """Downsample a pixel-space mask video to flattened latent token weights."""
+    b = mask.shape[0]
+    f_lat = target_latent_shape.frames
+    h_lat = target_latent_shape.height
+    w_lat = target_latent_shape.width
+
+    f_pix = mask.shape[2]
+    spatial_down = torch.nn.functional.interpolate(
+        rearrange(mask, "b 1 f h w -> (b f) 1 h w"),
+        size=(h_lat, w_lat),
+        mode="area",
+    )
+    spatial_down = rearrange(spatial_down, "(b f) 1 h w -> b 1 f h w", b=b)
+
+    first_frame = spatial_down[:, :, :1, :, :]
+
+    if f_pix > 1 and f_lat > 1:
+        t = (f_pix - 1) // (f_lat - 1)
+        assert (f_pix - 1) % (f_lat - 1) == 0, (
+            f"Pixel frames ({f_pix}) not compatible with latent frames ({f_lat}): "
+            f"(f_pix - 1) must be divisible by (f_lat - 1)"
+        )
+        rest = rearrange(spatial_down[:, :, 1:, :, :], "b 1 (f t) h w -> b 1 f t h w", t=t)
+        rest = rest.mean(dim=3)
+        latent_mask = torch.cat([first_frame, rest], dim=2)
+    else:
+        latent_mask = first_frame
+
+    return rearrange(latent_mask, "b 1 f h w -> b (f h w)")
+
+
+def append_ic_lora_reference_video_conditionings(  # noqa: PLR0913
+    conditionings: list[ConditioningItem],
+    video_conditioning: list[tuple[str, float]],
+    *,
+    height: int,
+    width: int,
+    num_frames: int,
+    video_encoder: VideoEncoder,
+    dtype: torch.dtype,
+    device: torch.device,
+    reference_downscale_factor: int,
+    conditioning_attention_strength: float,
+    conditioning_attention_mask: torch.Tensor | None,
+    tiling_config: TilingConfig | None = None,
+) -> None:
+    """Append :class:`VideoConditionByReferenceLatent` items for each reference path."""
+    scale = reference_downscale_factor
+    if scale != 1 and (height % scale != 0 or width % scale != 0):
+        raise ValueError(
+            f"Output dimensions ({height}x{width}) must be divisible by reference_downscale_factor ({scale})"
+        )
+    ref_height = height // scale
+    ref_width = width // scale
+
+    for video_path, strength in video_conditioning:
+        frame_gen = decode_video_by_frame(path=video_path, frame_cap=num_frames, device=device)
+        video = video_preprocess(frame_gen, ref_height, ref_width, dtype, device)
+        if tiling_config is not None:
+            encoded_video = video_encoder.tiled_encode(video, tiling_config)
+        else:
+            encoded_video = video_encoder(video)
+        reference_video_shape = VideoLatentShape.from_torch_shape(encoded_video.shape)
+
+        if conditioning_attention_mask is not None:
+            latent_mask = downsample_mask_video_to_latent(
+                mask=conditioning_attention_mask,
+                target_latent_shape=reference_video_shape,
+            )
+            attn_mask = latent_mask * conditioning_attention_strength
+        elif conditioning_attention_strength < 1.0:
+            attn_mask = conditioning_attention_strength
+        else:
+            attn_mask = None
+
+        cond = VideoConditionByReferenceLatent(
+            latent=encoded_video,
+            downscale_factor=scale,
+            strength=strength,
+        )
+        if attn_mask is not None:
+            cond = ConditioningItemAttentionStrengthWrapper(cond, attention_mask=attn_mask)
+        conditionings.append(cond)

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

@@ -15,7 +15,11 @@ from ltx_core.loader.registry import Registry
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
 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.args import (
+    ImageConditioningInput,
+    default_2_stage_arg_parser,
+    detect_checkpoint_path,
+)
 from ltx_pipelines.utils.blocks import (
     AudioDecoder,
     DiffusionStage,
@@ -35,7 +39,7 @@ from ltx_pipelines.utils.helpers import (
     image_conditionings_by_adding_guiding_latent,
 )
 from ltx_pipelines.utils.media_io import encode_video
-from ltx_pipelines.utils.types import ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class KeyframeInterpolationPipeline:
@@ -59,12 +63,15 @@ class KeyframeInterpolationPipeline:
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
         torch_compile: bool = False,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
         self.dtype = torch.bfloat16
         self._scheduler = LTX2Scheduler()
 
-        self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
+        self.prompt_encoder = PromptEncoder(
+            checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
+        )
         self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
         self.stage_1 = DiffusionStage(
             checkpoint_path,
@@ -74,6 +81,7 @@ class KeyframeInterpolationPipeline:
             quantization=quantization,
             registry=registry,
             torch_compile=torch_compile,
+            offload_mode=offload_mode,
         )
         stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
         self.stage_2 = DiffusionStage(
@@ -84,6 +92,7 @@ class KeyframeInterpolationPipeline:
             quantization=quantization,
             registry=registry,
             torch_compile=torch_compile,
+            offload_mode=offload_mode,
         )
         self.upsampler = VideoUpsampler(
             checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -106,7 +115,6 @@ 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,
         stage_1_sigmas: torch.Tensor | None = None,
         stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
@@ -122,7 +130,6 @@ class KeyframeInterpolationPipeline:
             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
@@ -179,7 +186,6 @@ class KeyframeInterpolationPipeline:
             audio=ModalitySpec(
                 context=a_context_p,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
             max_batch_size=max_batch_size,
         )
 
@@ -218,7 +224,6 @@ class KeyframeInterpolationPipeline:
                 noise_scale=stage_2_sigmas[0].item(),
                 initial_latent=audio_state.latent,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -241,6 +246,7 @@ def main() -> None:
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
         torch_compile=args.compile,
+        offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -271,7 +277,6 @@ def main() -> None:
         ),
         images=args.images,
         tiling_config=tiling_config,
-        streaming_prefetch_count=args.streaming_prefetch_count,
         max_batch_size=args.max_batch_size,
     )
 

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

@@ -0,0 +1,334 @@
+"""Two-stage lip-dubbing pipeline with IC-LoRA and appended audio reference conditioning."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator
+
+import torch
+
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.patchifiers import AudioPatchifier
+from ltx_core.conditioning import AudioConditionByReferenceLatent
+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.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
+from ltx_core.quantization import QuantizationPolicy
+from ltx_core.types import Audio, AudioLatentShape, SpatioTemporalScaleFactors, VideoPixelShape
+from ltx_pipelines.iclora_utils import (
+    append_ic_lora_reference_video_conditionings,
+    read_lora_reference_downscale_factor,
+)
+from ltx_pipelines.utils.args import (
+    ImageConditioningInput,
+    detect_checkpoint_path,
+    lipdub_arg_parser,
+)
+from ltx_pipelines.utils.blocks import (
+    AudioConditioner,
+    AudioDecoder,
+    DiffusionStage,
+    ImageConditioner,
+    PromptEncoder,
+    VideoDecoder,
+    VideoUpsampler,
+)
+from ltx_pipelines.utils.constants import DISTILLED_SIGMAS, STAGE_2_DISTILLED_SIGMAS, detect_params
+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_audio_from_file, encode_video, get_videostream_metadata
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
+
+
+def _snap_frames_to_8k1(frames: int) -> int:
+    """Round ``frames`` down to the nearest ``8k+1`` (the model's required frame count)."""
+    time_scale = SpatioTemporalScaleFactors.default().time
+    return ((frames - 1) // time_scale) * time_scale + 1
+
+
+class LipDubPipeline:
+    """Two-stage lip-dubbing with IC-LoRA video reference and appended audio reference tokens."""
+
+    def __init__(
+        self,
+        distilled_checkpoint_path: str,
+        spatial_upsampler_path: str,
+        gemma_root: str,
+        ic_lora: LoraPathStrengthAndSDOps,
+        device: torch.device | None = None,
+        quantization: QuantizationPolicy | None = None,
+        registry: Registry | None = None,
+        torch_compile: bool = False,
+        offload_mode: OffloadMode = OffloadMode.NONE,
+    ) -> None:
+        self.device = device or get_device()
+        self.dtype = torch.bfloat16
+        self.ic_lora = ic_lora
+        loras = (ic_lora,)
+
+        self.prompt_encoder = PromptEncoder(
+            distilled_checkpoint_path,
+            gemma_root,
+            self.dtype,
+            self.device,
+            registry=registry,
+            offload_mode=offload_mode,
+        )
+        self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
+        self.audio_conditioner = AudioConditioner(
+            distilled_checkpoint_path,
+            self.dtype,
+            self.device,
+            registry=registry,
+        )
+        self.stage = DiffusionStage(
+            distilled_checkpoint_path,
+            self.dtype,
+            self.device,
+            loras=loras,
+            quantization=quantization,
+            registry=registry,
+            torch_compile=torch_compile,
+            offload_mode=offload_mode,
+        )
+        self.upsampler = VideoUpsampler(
+            distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
+        )
+        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)
+        self.reference_downscale_factor = read_lora_reference_downscale_factor(ic_lora.path)
+
+    def _create_stage_conditionings(
+        self,
+        images: list[ImageConditioningInput],
+        reference_video_path: str,
+        reference_strength: float,
+        height: int,
+        width: int,
+        num_frames: int,
+        video_encoder: VideoEncoder,
+        encode_tiling: TilingConfig | None,
+    ) -> list:
+        conditionings = combined_image_conditionings(
+            images=images,
+            height=height,
+            width=width,
+            video_encoder=video_encoder,
+            dtype=self.dtype,
+            device=self.device,
+        )
+        append_ic_lora_reference_video_conditionings(
+            conditionings,
+            [(reference_video_path, reference_strength)],
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            video_encoder=video_encoder,
+            dtype=self.dtype,
+            device=self.device,
+            reference_downscale_factor=self.reference_downscale_factor,
+            conditioning_attention_strength=1.0,
+            conditioning_attention_mask=None,
+            tiling_config=encode_tiling,
+        )
+        return conditionings
+
+    def _encode_reference_audio_vae_latent(self, video_path: str) -> torch.Tensor:
+        audio = decode_audio_from_file(video_path, self.device)
+        if audio is None:
+            msg = f"No audio stream found in {video_path}"
+            raise ValueError(msg)
+        return self.audio_conditioner(lambda enc: vae_encode_audio(audio, enc, None))
+
+    @torch.inference_mode()
+    def __call__(  # noqa: PLR0913
+        self,
+        prompt: str,
+        seed: int,
+        height: int,
+        width: int,
+        images: list[ImageConditioningInput],
+        reference_video_path: str,
+        reference_strength: float = 1.0,
+        enhance_prompt: bool = False,
+        tiling_config: TilingConfig | None = None,
+        stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS,
+        stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
+    ) -> tuple[Iterator[torch.Tensor], Audio]:
+        assert_resolution(height=height, width=width, is_two_stage=True)
+
+        meta = get_videostream_metadata(reference_video_path)
+        num_frames = _snap_frames_to_8k1(meta.frames)
+        frame_rate = float(meta.fps)
+
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        noiser = GaussianNoiser(generator=generator)
+
+        (ctx_p,) = self.prompt_encoder(
+            [prompt],
+            enhance_first_prompt=enhance_prompt,
+            enhance_prompt_image=images[0][0] if len(images) > 0 else None,
+            enhance_prompt_seed=seed,
+        )
+        video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
+
+        stage_1_output_shape = VideoPixelShape(
+            batch=1,
+            frames=num_frames,
+            width=width // 2,
+            height=height // 2,
+            fps=frame_rate,
+        )
+        encode_tiling = TilingConfig.default()
+
+        def build_image_conditionings(output_shape: VideoPixelShape) -> list:
+            return self.image_conditioner(
+                lambda enc: self._create_stage_conditionings(
+                    images=images,
+                    reference_video_path=reference_video_path,
+                    reference_strength=reference_strength,
+                    height=output_shape.height,
+                    width=output_shape.width,
+                    num_frames=num_frames,
+                    video_encoder=enc,
+                    encode_tiling=encode_tiling,
+                )
+            )
+
+        def build_audio_ref_conditioning(audio_latent: torch.Tensor) -> AudioConditionByReferenceLatent:
+            ref_patch, ref_pos = patchify_lipdub_audio_reference_latent(
+                audio_latent,
+                negative_positions=True,
+                device=self.device,
+            )
+            return AudioConditionByReferenceLatent(ref_patch, ref_pos, strength=1.0)
+
+        stage_1_conditionings = build_image_conditionings(stage_1_output_shape)
+
+        ref_vae = self._encode_reference_audio_vae_latent(reference_video_path)
+        audio_conditionings = [build_audio_ref_conditioning(ref_vae)]
+
+        stage_1_sigmas_tensor = stage_1_sigmas.to(dtype=torch.float32, device=self.device)
+        video_state, audio_state = self.stage(
+            denoiser=SimpleDenoiser(video_context, audio_context),
+            sigmas=stage_1_sigmas_tensor,
+            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,
+                conditionings=audio_conditionings,
+            ),
+        )
+
+        s1_audio_latent = audio_state.latent.clone()
+
+        upscaled_video_latent = self.upsampler(video_state.latent[:1])
+        stage_2_sigmas_tensor = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
+        stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
+        stage_2_conditionings = build_image_conditionings(stage_2_output_shape)
+
+        stage_2_audio_conditionings = [build_audio_ref_conditioning(s1_audio_latent)]
+
+        video_state, _audio_unused = self.stage(
+            denoiser=SimpleDenoiser(video_context, audio_context),
+            sigmas=stage_2_sigmas_tensor,
+            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_tensor[0].item(),
+                initial_latent=upscaled_video_latent,
+            ),
+            audio=ModalitySpec(
+                context=audio_context,
+                conditionings=stage_2_audio_conditionings,
+                frozen=True,
+                noise_scale=0.0,
+                initial_latent=s1_audio_latent,
+            ),
+        )
+
+        decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
+        decoded_audio = self.audio_decoder(s1_audio_latent)
+        return decoded_video, decoded_audio
+
+
+def patchify_lipdub_audio_reference_latent(
+    vae_latents: torch.Tensor,
+    *,
+    negative_positions: bool,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Patchify audio VAE latents and build RoPE positions (optional negative shift for reference)."""
+    patchifier = AudioPatchifier(patch_size=1)
+    patchified = patchifier.patchify(vae_latents)
+    b, c, _t, mel_bins = vae_latents.shape
+    seq_len = patchified.shape[1]
+    latent_coords = patchifier.get_patch_grid_bounds(
+        output_shape=AudioLatentShape(batch=b, channels=c, frames=seq_len, mel_bins=mel_bins),
+        device=device,
+    )
+    positions = latent_coords.to(dtype=torch.float32)
+    if negative_positions:
+        aud_dur = positions[:, :, -1, 1].max().item()
+        positions = positions - aud_dur - 0.04
+    return patchified, positions
+
+
+@torch.inference_mode()
+def main() -> None:
+    logging.getLogger().setLevel(logging.INFO)
+    checkpoint_path = detect_checkpoint_path(distilled=True)
+    params = detect_params(checkpoint_path)
+    parser = lipdub_arg_parser(params=params)
+    args = parser.parse_args()
+
+    if not args.lora or len(args.lora) != 1:
+        raise ValueError("LipDub requires exactly one --lora (the lip-dub IC-LoRA).")
+
+    pipeline = LipDubPipeline(
+        distilled_checkpoint_path=args.distilled_checkpoint_path,
+        spatial_upsampler_path=args.spatial_upsampler_path,
+        gemma_root=args.gemma_root,
+        ic_lora=args.lora[0],
+        quantization=args.quantization,
+        torch_compile=args.compile,
+        offload_mode=args.offload_mode,
+    )
+    tiling_config = TilingConfig.default()
+    src = get_videostream_metadata(args.reference_video)
+    video_chunks_number = get_video_chunks_number(_snap_frames_to_8k1(src.frames), tiling_config)
+    video, audio = pipeline(
+        prompt=args.prompt,
+        seed=args.seed,
+        height=args.height,
+        width=args.width,
+        images=[],
+        reference_video_path=args.reference_video,
+        reference_strength=args.reference_strength,
+        tiling_config=tiling_config,
+        enhance_prompt=args.enhance_prompt,
+    )
+    encode_video(
+        video=video,
+        fps=int(src.fps),
+        audio=audio,
+        output_path=args.output_path,
+        video_chunks_number=video_chunks_number,
+    )
+
+
+if __name__ == "__main__":
+    main()

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

@@ -36,7 +36,7 @@ from ltx_pipelines.utils.media_io import (
     encode_video,
     get_videostream_metadata,
 )
-from ltx_pipelines.utils.types import ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class RetakePipeline:
@@ -74,6 +74,7 @@ class RetakePipeline:
         registry: Registry | None = None,
         distilled: bool = True,
         torch_compile: bool = False,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
         self.dtype = torch.bfloat16
@@ -86,6 +87,7 @@ class RetakePipeline:
             dtype=self.dtype,
             device=self.device,
             registry=registry,
+            offload_mode=offload_mode,
         )
         self.image_conditioner = ImageConditioner(
             checkpoint_path=checkpoint_path,
@@ -107,6 +109,7 @@ class RetakePipeline:
             quantization=quantization,
             registry=registry,
             torch_compile=torch_compile,
+            offload_mode=offload_mode,
         )
         self.video_decoder = VideoDecoder(
             checkpoint_path=checkpoint_path,
@@ -141,7 +144,6 @@ class RetakePipeline:
         regenerate_audio: bool = True,
         enhance_prompt: bool = False,
         tiling_config: TilingConfig | None = None,
-        streaming_prefetch_count: int | None = None,
         max_batch_size: int = 1,
         sigmas: torch.Tensor | None = None,
     ) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
@@ -210,7 +212,6 @@ class RetakePipeline:
             prompts_to_encode,
             enhance_first_prompt=enhance_prompt,
             enhance_prompt_seed=seed,
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         v_context_p, a_context_p = contexts[0].video_encoding, contexts[0].audio_encoding
@@ -269,7 +270,6 @@ class RetakePipeline:
             fps=output_shape.fps,
             video=video_modality_spec,
             audio=audio_modality_spec,
-            streaming_prefetch_count=streaming_prefetch_count,
             max_batch_size=max_batch_size,
         )
 
@@ -307,8 +307,9 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
-        distilled=args.distilled,
+        distilled=True,
         torch_compile=args.compile,
+        offload_mode=args.offload_mode,
     )
     params = detect_params(args.distilled_checkpoint_path)
     tiling_config = TilingConfig.default()
@@ -321,7 +322,6 @@ 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(src.frames, tiling_config)

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

@@ -20,7 +20,11 @@ from ltx_pipelines.utils import (
     combined_image_conditionings,
     get_device,
 )
-from ltx_pipelines.utils.args import ImageConditioningInput, default_1_stage_arg_parser, detect_checkpoint_path
+from ltx_pipelines.utils.args import (
+    ImageConditioningInput,
+    default_1_stage_arg_parser,
+    detect_checkpoint_path,
+)
 from ltx_pipelines.utils.blocks import (
     AudioDecoder,
     DiffusionStage,
@@ -31,7 +35,7 @@ from ltx_pipelines.utils.blocks import (
 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 ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class TI2VidOneStagePipeline:
@@ -52,6 +56,7 @@ class TI2VidOneStagePipeline:
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
         torch_compile: bool = False,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.dtype = torch.bfloat16
         self.device = device or get_device()
@@ -62,6 +67,7 @@ class TI2VidOneStagePipeline:
             dtype=self.dtype,
             device=self.device,
             registry=registry,
+            offload_mode=offload_mode,
         )
         self.image_conditioner = ImageConditioner(
             checkpoint_path=checkpoint_path,
@@ -77,6 +83,7 @@ class TI2VidOneStagePipeline:
             quantization=quantization,
             registry=registry,
             torch_compile=torch_compile,
+            offload_mode=offload_mode,
         )
         self.video_decoder = VideoDecoder(
             checkpoint_path=checkpoint_path,
@@ -105,7 +112,6 @@ 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,
         sigmas: torch.Tensor | None = None,
@@ -121,7 +127,6 @@ class TI2VidOneStagePipeline:
             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
@@ -170,7 +175,6 @@ class TI2VidOneStagePipeline:
             audio=ModalitySpec(
                 context=a_context_p,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
             max_batch_size=max_batch_size,
         )
 
@@ -192,6 +196,7 @@ def main() -> None:
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
         torch_compile=args.compile,
+        offload_mode=args.offload_mode,
     )
     video, audio = pipeline(
         prompt=args.prompt,
@@ -219,7 +224,6 @@ 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,
     )
 

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

@@ -15,7 +15,11 @@ from ltx_core.loader.registry import Registry
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
 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.args import (
+    ImageConditioningInput,
+    default_2_stage_arg_parser,
+    detect_checkpoint_path,
+)
 from ltx_pipelines.utils.blocks import (
     AudioDecoder,
     DiffusionStage,
@@ -35,7 +39,7 @@ from ltx_pipelines.utils.helpers import (
     get_device,
 )
 from ltx_pipelines.utils.media_io import encode_video
-from ltx_pipelines.utils.types import ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class TI2VidTwoStagesPipeline:
@@ -58,12 +62,15 @@ class TI2VidTwoStagesPipeline:
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
         torch_compile: bool = False,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
         self.dtype = torch.bfloat16
         self._scheduler = LTX2Scheduler()
 
-        self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
+        self.prompt_encoder = PromptEncoder(
+            checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
+        )
         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
@@ -79,6 +86,7 @@ class TI2VidTwoStagesPipeline:
             quantization=quantization,
             registry=registry,
             torch_compile=torch_compile,
+            offload_mode=offload_mode,
         )
         self.stage_2 = DiffusionStage(
             checkpoint_path,
@@ -88,6 +96,7 @@ class TI2VidTwoStagesPipeline:
             quantization=quantization,
             registry=registry,
             torch_compile=torch_compile,
+            offload_mode=offload_mode,
         )
 
     def __call__(  # noqa: PLR0913
@@ -105,7 +114,6 @@ 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,
         stage_1_sigmas: torch.Tensor | None = None,
         stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
@@ -121,7 +129,6 @@ class TI2VidTwoStagesPipeline:
             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
@@ -170,7 +177,6 @@ class TI2VidTwoStagesPipeline:
             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,
         )
 
@@ -208,7 +214,6 @@ class TI2VidTwoStagesPipeline:
                 noise_scale=stage_2_sigmas[0].item(),
                 initial_latent=audio_state.latent,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -231,6 +236,7 @@ def main() -> None:
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
         torch_compile=args.compile,
+        offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -261,7 +267,6 @@ def main() -> None:
         ),
         images=args.images,
         tiling_config=tiling_config,
-        streaming_prefetch_count=args.streaming_prefetch_count,
         max_batch_size=args.max_batch_size,
     )
 

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

@@ -33,7 +33,7 @@ from ltx_pipelines.utils.helpers import (
 )
 from ltx_pipelines.utils.media_io import encode_video
 from ltx_pipelines.utils.samplers import res2s_audio_video_denoising_loop
-from ltx_pipelines.utils.types import ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class TI2VidTwoStagesHQPipeline:
@@ -61,6 +61,7 @@ class TI2VidTwoStagesHQPipeline:
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
         torch_compile: bool = False,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
         self.dtype = torch.bfloat16
@@ -77,7 +78,9 @@ class TI2VidTwoStagesHQPipeline:
             sd_ops=distilled_lora[0].sd_ops,
         )
 
-        self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
+        self.prompt_encoder = PromptEncoder(
+            checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
+        )
         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
@@ -93,6 +96,7 @@ class TI2VidTwoStagesHQPipeline:
             quantization=quantization,
             registry=registry,
             torch_compile=torch_compile,
+            offload_mode=offload_mode,
         )
         self.stage_2 = DiffusionStage(
             checkpoint_path,
@@ -102,6 +106,7 @@ class TI2VidTwoStagesHQPipeline:
             quantization=quantization,
             registry=registry,
             torch_compile=torch_compile,
+            offload_mode=offload_mode,
         )
 
     @torch.inference_mode()
@@ -120,7 +125,6 @@ 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,
         stage_1_sigmas: torch.Tensor | None = None,
         stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
@@ -136,7 +140,6 @@ class TI2VidTwoStagesHQPipeline:
             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
@@ -190,7 +193,6 @@ class TI2VidTwoStagesHQPipeline:
             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,
         )
 
@@ -231,7 +233,6 @@ class TI2VidTwoStagesHQPipeline:
                 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)
@@ -254,6 +255,7 @@ def main() -> None:
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
         torch_compile=args.compile,
+        offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -284,7 +286,6 @@ def main() -> None:
         ),
         images=args.images,
         tiling_config=tiling_config,
-        streaming_prefetch_count=args.streaming_prefetch_count,
         max_batch_size=args.max_batch_size,
     )
 

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

@@ -16,15 +16,17 @@ from ltx_pipelines.utils.helpers import (
     image_conditionings_by_adding_guiding_latent,
 )
 from ltx_pipelines.utils.samplers import (
+    euler_cfg_pp_denoising_loop,
     euler_denoising_loop,
     gradient_estimating_euler_denoising_loop,
     res2s_audio_video_denoising_loop,
 )
-from ltx_pipelines.utils.types import Denoiser, ModalitySpec
+from ltx_pipelines.utils.types import DenoisedLatentResult, Denoiser, ModalitySpec
 
 __all__ = [
     "AudioConditioner",
     "AudioDecoder",
+    "DenoisedLatentResult",
     "Denoiser",
     "DiffusionStage",
     "FactoryGuidedDenoiser",
@@ -38,6 +40,7 @@ __all__ = [
     "assert_resolution",
     "cleanup_memory",
     "combined_image_conditionings",
+    "euler_cfg_pp_denoising_loop",
     "euler_denoising_loop",
     "get_device",
     "gradient_estimating_euler_denoising_loop",

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

@@ -1,4 +1,5 @@
 import argparse
+from collections.abc import Sequence
 from pathlib import Path
 from typing import NamedTuple
 
@@ -12,6 +13,7 @@ from ltx_pipelines.utils.constants import (
     LTX_2_3_PARAMS,
     PipelineParams,
 )
+from ltx_pipelines.utils.types import OffloadMode
 
 
 class ImageConditioningInput(NamedTuple):
@@ -114,35 +116,34 @@ def resolve_path(path: str) -> str:
 QUANTIZATION_POLICIES = ("fp8-cast", "fp8-scaled-mm")
 
 
-class QuantizationAction(argparse.Action):
+def _resolve_quantization(namespace: argparse.Namespace) -> None:
-    def __call__(
+    # Resolution is deferred until after parse_args because fp8-scaled-mm needs the
-        self,
+    # checkpoint path, which isn't on the namespace when the --quantization argument
-        parser: argparse.ArgumentParser,  # noqa: ARG002
+    # is parsed.
-        namespace: argparse.Namespace,
+    name = getattr(namespace, "quantization", None)
-        values: list[str],
+    if name is None or isinstance(name, QuantizationPolicy):
-        option_string: str | None = None,
+        return
-    ) -> None:
+    if name == "fp8-cast":
-        if len(values) > 2:
+        namespace.quantization = QuantizationPolicy.fp8_cast()
-            msg = (
+        return
-                f"{option_string} accepts at most 2 arguments (POLICY and optional AMAX_PATH), got {len(values)} values"
+    if name == "fp8-scaled-mm":
+        ckpt = getattr(namespace, "checkpoint_path", None) or getattr(namespace, "distilled_checkpoint_path", None)
+        if ckpt is None:
+            raise SystemExit(
+                "--quantization fp8-scaled-mm requires --checkpoint-path (or --distilled-checkpoint-path)."
             )
-            raise argparse.ArgumentError(self, msg)
+        namespace.quantization = QuantizationPolicy.fp8_scaled_mm(ckpt)
 
-        policy_name = values[0]
-        if policy_name not in QUANTIZATION_POLICIES:
-            msg = f"Unknown quantization policy '{policy_name}'. Choose from: {', '.join(QUANTIZATION_POLICIES)}"
-            raise argparse.ArgumentError(self, msg)
 
-        if policy_name == "fp8-cast":
+class _PipelineArgumentParser(argparse.ArgumentParser):
-            if len(values) > 1:
+    def parse_args(  # type: ignore[override]
-                msg = f"{option_string} fp8-cast does not accept additional arguments"
+        self,
-                raise argparse.ArgumentError(self, msg)
+        args: Sequence[str] | None = None,
-            policy = QuantizationPolicy.fp8_cast()
+        namespace: argparse.Namespace | None = None,
-        elif policy_name == "fp8-scaled-mm":
+    ) -> argparse.Namespace:
-            amax_path = resolve_path(values[1]) if len(values) > 1 else None
+        ns = super().parse_args(args, namespace)
-            policy = QuantizationPolicy.fp8_scaled_mm(amax_path)
+        _resolve_quantization(ns)
-
+        return ns
-        setattr(namespace, self.dest, policy)
 
 
 def detect_checkpoint_path(distilled: bool = False) -> str:
@@ -158,7 +159,7 @@ def basic_arg_parser(
     params: PipelineParams = LTX_2_3_PARAMS,
     distilled: bool = False,
 ) -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser()
+    parser = _PipelineArgumentParser()
     if distilled:
         parser.add_argument(
             "--distilled-checkpoint-path",
@@ -231,16 +232,19 @@ def basic_arg_parser(
         except ValueError as e:
             raise argparse.ArgumentTypeError(f"must be an integer, got {value}") from e
 
-    # Layer streaming
+    # Weight offloading
     parser.add_argument(
-        "--streaming-prefetch-count",
+        "--offload",
-        type=_positive_int,
+        dest="offload_mode",
-        default=None,
+        type=OffloadMode,
-        metavar="N",
+        default=OffloadMode.NONE,
+        choices=list(OffloadMode),
         help=(
-            "Enable layer streaming prefetching N layers ahead. "
+            "Weight offloading strategy. "
-            "At most 1 + N layers reside on GPU at once. "
+            "'none' keeps all weights on GPU (default). "
-            "Must be >= 1. Example: --streaming-prefetch-count 2"
+            "'cpu' pins weights in CPU RAM, streams to GPU per layer. "
+            "'disk' reads weights from disk on demand (lowest memory). "
+            "Example: --offload cpu"
         ),
     )
 
@@ -260,16 +264,14 @@ def basic_arg_parser(
 
     parser.add_argument(
         "--quantization",
-        dest="quantization",
+        choices=QUANTIZATION_POLICIES,
-        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). "
+            "fp8-scaled-mm uses FP8 scaled matrix multiplication; the layer set is auto-discovered "
-            "Example: --quantization fp8-cast or --quantization fp8-scaled-mm /path/to/amax.json"
+            "from the checkpoint's .weight_scale tensors. "
+            "Example: --quantization fp8-cast or --quantization fp8-scaled-mm"
         ),
     )
     parser.add_argument(
@@ -344,6 +346,53 @@ def video_editing_arg_parser(
     return parser
 
 
+def lipdub_arg_parser(
+    params: PipelineParams = LTX_2_3_PARAMS,
+) -> argparse.ArgumentParser:
+    """Argument parser for the lip-dub pipeline.
+    Frame count and frame rate are derived from the reference video at runtime (the frame count
+    is silently snapped down to the nearest 8k+1), so this parser intentionally omits
+    --num-frames, --frame-rate, and --image. Distilled checkpoint only.
+    """
+    parser = basic_arg_parser(params=params, distilled=True)
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=params.stage_2_height,
+        help=(
+            f"Height of the generated video in pixels, should be divisible by 64 (default: {params.stage_2_height})."
+        ),
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=params.stage_2_width,
+        help=f"Width of the generated video in pixels, should be divisible by 64 (default: {params.stage_2_width}).",
+    )
+    parser.add_argument(
+        "--spatial-upsampler-path",
+        type=resolve_path,
+        required=True,
+        help=(
+            "Path to the spatial upsampler model used to increase the resolution "
+            "of the generated video in the latent space."
+        ),
+    )
+    parser.add_argument(
+        "--reference-video",
+        type=resolve_path,
+        required=True,
+        help="Reference video file (video + audio track used for IC-LoRA and audio identity).",
+    )
+    parser.add_argument(
+        "--reference-strength",
+        type=float,
+        default=1.0,
+        help="Strength for IC-LoRA video reference conditioning (default: 1.0).",
+    )
+    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

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

@@ -15,13 +15,14 @@ from typing import Callable, TypeVar
 import torch
 
 from ltx_core.batch_split import BatchSplitAdapter
+from ltx_core.block_streaming import DISK_CPU_SLOTS, StreamingModelBuilder
 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.module_ops import ModuleOps
+from ltx_core.loader.primitives import BuilderProtocol, LoraPathStrengthAndSDOps, ModelBuilderProtocol
 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 (
@@ -37,12 +38,14 @@ from ltx_core.model.audio_vae import (
 )
 from ltx_core.model.transformer import (
     LTXV_MODEL_COMFY_RENAMING_MAP,
+    LTXModel,
     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 (
+    MEMORY_EFFICIENT_DECODE,
     VAE_DECODER_COMFY_KEYS_FILTER,
     VAE_ENCODER_COMFY_KEYS_FILTER,
     TilingConfig,
@@ -59,10 +62,11 @@ from ltx_core.text_encoders.gemma import (
     GemmaTextEncoderConfigurator,
     module_ops_from_gemma_root,
 )
-from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessorOutput
+from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessor, 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.multigpu.delegating_builder import DelegatingBuilder
 from ltx_pipelines.utils.gpu_model import gpu_model
 from ltx_pipelines.utils.helpers import (
     cleanup_memory,
@@ -70,7 +74,7 @@ from ltx_pipelines.utils.helpers import (
     generate_enhanced_prompt,
 )
 from ltx_pipelines.utils.samplers import euler_denoising_loop
-from ltx_pipelines.utils.types import Denoiser, ModalitySpec
+from ltx_pipelines.utils.types import Denoiser, ModalitySpec, OffloadMode
 
 logger = logging.getLogger(__name__)
 
@@ -83,36 +87,40 @@ _M = TypeVar("_M", bound=torch.nn.Module)
 # ---------------------------------------------------------------------------
 
 
+def _chain_quantization(
+    sd_ops: SDOps,
+    module_ops: tuple[ModuleOps, ...],
+    quantization: QuantizationPolicy,
+) -> tuple[SDOps, tuple[ModuleOps, ...]]:
+    chained_sd_ops = sd_ops
+    if quantization.sd_ops is not None:
+        chained_sd_ops = SDOps(
+            name=f"sd_ops_chain_{sd_ops.name}+{quantization.sd_ops.name}",
+            mapping=(*sd_ops.mapping, *quantization.sd_ops.mapping),
+        )
+    return chained_sd_ops, (*module_ops, *quantization.module_ops)
+
+
 @contextmanager
 def _streaming_model(
-    model: _M,
+    builder: StreamingModelBuilder,
-    layers_attr: str,
+    offload_mode: OffloadMode,
     target_device: torch.device,
-    prefetch_count: int,
+    dtype: torch.dtype,
-) -> Iterator[_M]:
+) -> Iterator:
-    """Wrap *model* with :class:`LayerStreamingWrapper`, yield it, then tear down."""
+    """Build a streaming wrapper, yield it, then tear down and free memory."""
-    wrapped = LayerStreamingWrapper(
+    cpu_slots_count = DISK_CPU_SLOTS if offload_mode == OffloadMode.DISK else None
-        model,
+    wrapped = builder.build(
-        layers_attr=layers_attr,
         target_device=target_device,
-        prefetch_count=prefetch_count,
+        dtype=dtype,
+        cpu_slots_count=cpu_slots_count,
     )
     try:
-        yield wrapped  # type: ignore[misc]
+        yield wrapped
     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(
@@ -163,18 +171,51 @@ class DiffusionStage:
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
         torch_compile: bool = False,
+        offload_mode: OffloadMode = OffloadMode.NONE,
+        transformer_builder: ModelBuilderProtocol[LTXModel] | DelegatingBuilder[LTXModel] | None = None,
     ) -> None:
         self._dtype = dtype
         self._device = device
         self._quantization = quantization
         self._torch_compile = torch_compile
-        self._transformer_builder = Builder(
+        self._offload_mode = offload_mode
-            model_path=checkpoint_path,
+        if transformer_builder is not None:
-            model_class_configurator=LTXModelConfigurator,
+            self._transformer_builder = transformer_builder
-            model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
+        else:
-            loras=tuple(loras),
+            self._transformer_builder = Builder(
-            registry=registry or DummyRegistry(),
+                model_path=checkpoint_path,
-        )
+                model_class_configurator=LTXModelConfigurator,
+                model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
+                loras=tuple(loras),
+                registry=registry or DummyRegistry(),
+            )
+
+        if offload_mode != OffloadMode.NONE:
+            if torch_compile:
+                raise ValueError("torch.compile is not supported with layer streaming")
+            streaming_sd_ops: SDOps = LTXV_MODEL_COMFY_RENAMING_MAP
+            streaming_module_ops: tuple[ModuleOps, ...] = ()
+            if quantization is not None:
+                if quantization.kind != QuantizationPolicy.Kind.FP8_CAST:
+                    raise ValueError(
+                        f"Layer streaming supports only QuantizationPolicy.fp8_cast(); "
+                        f"got kind={quantization.kind!r} which produces heterogeneous block layouts."
+                    )
+                streaming_sd_ops, streaming_module_ops = _chain_quantization(
+                    streaming_sd_ops, streaming_module_ops, quantization
+                )
+            self._streaming_builder = StreamingModelBuilder(
+                model_class_configurator=LTXModelConfigurator,
+                model_path=checkpoint_path,
+                model_sd_ops=streaming_sd_ops,
+                module_ops=streaming_module_ops,
+                loras=tuple(loras),
+                registry=registry or DummyRegistry(),
+                blocks_attr="velocity_model.transformer_blocks",
+                blocks_prefix="transformer_blocks",
+                state_dict_prefix="velocity_model.",
+                model_wrapper=lambda m: X0Model(m).eval(),
+            )
 
     def _build_transformer(self, *, device: torch.device | None = None, **kwargs: object) -> X0Model:
         target = device or self._device
@@ -189,38 +230,31 @@ class DiffusionStage:
                 LoraPathStrengthAndSDOps(
                     lora.path,
                     lora.strength,
-                    modify_sd_ops_for_compilation(
+                    modify_sd_ops_for_compilation(lora.sd_ops, number_of_layers),
-                        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, module_ops = _chain_quantization(sd_ops, module_ops, self._quantization)
-            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(
+    def _transformer_ctx(self, **kwargs: object) -> AbstractContextManager:
-        self,
+        if self._offload_mode != OffloadMode.NONE:
-        streaming_prefetch_count: int | None,
+            return _streaming_model(self._streaming_builder, self._offload_mode, self._device, self._dtype)
-        **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
+    def model_context(self, **kwargs: object) -> AbstractContextManager:
+        """Build the transformer, yield it, then free its memory on exit.
+        Keyword arguments are forwarded to the underlying builder (e.g.
+        ``video_tools`` required by ``TiledDataParallelBuilder``).
+        """
+        return self._transformer_ctx(**kwargs)
+
+    def run(  # noqa: PLR0913
         self,
+        transformer: object,
         denoiser: Denoiser,
         sigmas: torch.Tensor,
         noiser: Noiser,
@@ -232,27 +266,14 @@ class DiffusionStage:
         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.
+        """Run denoising with a pre-built transformer.
-        Args:
+        Same semantics as ``__call__`` but accepts a pre-built transformer so
-            width: Output width in pixels.
+        the model can be shared across multiple calls (e.g. tiled inference
-            height: Output height in pixels.
+        inside a single ``model_context()`` block). Audio supports
-            frames: Number of output frames.
+        ``ModalitySpec(frozen=True)`` to keep the latent unchanged throughout
-            fps: Frame rate.
+        denoising while still providing cross-modal context to the transformer.
-            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.
         """
@@ -261,7 +282,6 @@ class DiffusionStage:
 
         if loop is None:
             loop = euler_denoising_loop
-
         if stepper is None:
             stepper = EulerDiffusionStep()
 
@@ -281,28 +301,70 @@ class DiffusionStage:
             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:
+        wrapped = BatchSplitAdapter(transformer, max_batch_size=max_batch_size)  # type: ignore[arg-type]
-            transformer = BatchSplitAdapter(base_transformer, max_batch_size=max_batch_size)
+        video_state, audio_state = loop(
-            video_state, audio_state = loop(
+            sigmas=sigmas,
-                sigmas=sigmas,
+            video_state=video_state,
-                video_state=video_state,
+            audio_state=audio_state,
-                audio_state=audio_state,
+            stepper=stepper,
-                stepper=stepper,
+            transformer=wrapped,
-                transformer=transformer,
+            denoiser=denoiser,
-                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
 
+    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,
+        max_batch_size: int = 1,
+    ) -> tuple[LatentState | None, LatentState | None]:
+        """Build transformer -> run denoising loop -> free transformer.
+        Returns ``(video_state | None, audio_state | None)`` with cleared
+        conditionings and unpatchified latents for present modalities.
+        """
+        # Build video_tools up front so it can be forwarded to the transformer
+        # context (required by TiledDataParallelBuilder in multi-GPU mode).
+        # `run()` rebuilds its own tools internally; the duplication is cheap.
+        video_tools: LatentTools | None = None
+        if video is not None:
+            pixel_shape = VideoPixelShape(batch=1, frames=frames, height=height, width=width, fps=fps)
+            v_shape = VideoLatentShape.from_pixel_shape(pixel_shape)
+            video_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape, fps)
+
+        with self._transformer_ctx(video_tools=video_tools) as transformer:
+            return self.run(
+                transformer,
+                denoiser,
+                sigmas,
+                noiser,
+                width,
+                height,
+                frames,
+                fps,
+                video,
+                audio,
+                stepper,
+                loop,
+                max_batch_size,
+            )
+
 
 # ---------------------------------------------------------------------------
 # PromptEncoder
@@ -322,21 +384,41 @@ class PromptEncoder:
         dtype: torch.dtype,
         device: torch.device,
         registry: Registry | None = None,
+        offload_mode: OffloadMode = OffloadMode.NONE,
+        text_encoder_builder: BuilderProtocol | None = None,
     ) -> None:
         self._dtype = dtype
         self._device = device
+        self._offload_mode = offload_mode
 
-        module_ops = module_ops_from_gemma_root(gemma_root)
+        if text_encoder_builder is not None:
-        model_folder = find_matching_file(gemma_root, "model*.safetensors").parent
+            if offload_mode != OffloadMode.NONE:
-        weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")]
+                raise ValueError(
-
+                    "text_encoder_builder cannot be used with offload_mode != OffloadMode.NONE "
-        self._text_encoder_builder = Builder(
+                    "because no streaming text encoder builder is available."
-            model_path=tuple(weight_paths),
+                )
-            model_class_configurator=GemmaTextEncoderConfigurator,
+            self._text_encoder_builder = text_encoder_builder
-            model_sd_ops=GEMMA_LLM_KEY_OPS,
+            self._streaming_text_encoder_builder = None
-            module_ops=(GEMMA_MODEL_OPS, *module_ops),
+        else:
-            registry=registry or DummyRegistry(),
+            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._streaming_text_encoder_builder = StreamingModelBuilder(
+                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(),
+                blocks_attr="model.model.language_model.layers",
+                blocks_prefix="model.model.language_model.layers",
+            )
         self._embeddings_processor_builder = Builder(
             model_path=checkpoint_path,
             model_class_configurator=EmbeddingsProcessorConfigurator,
@@ -344,18 +426,18 @@ class PromptEncoder:
             registry=registry or DummyRegistry(),
         )
 
-    def _text_encoder_ctx(
+    def _build_text_encoder(self) -> torch.nn.Module:
-        self,
+        """Build the Gemma text encoder (non-streaming path)."""
-        streaming_prefetch_count: int | None,
+        return self._text_encoder_builder.build(device=self._device, dtype=self._dtype).eval()
-    ) -> AbstractContextManager:
+
-        if streaming_prefetch_count is not None:
+    def _build_embeddings_processor(self) -> EmbeddingsProcessor:
-            return _streaming_model(
+        """Build the embeddings processor on the target device."""
-                self._text_encoder_builder.build(device=torch.device("cpu"), dtype=self._dtype).eval(),
+        return self._embeddings_processor_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
-                layers_attr="model.model.language_model.layers",
+
-                target_device=self._device,
+    def _text_encoder_ctx(self) -> AbstractContextManager:
-                prefetch_count=streaming_prefetch_count,
+        if self._offload_mode != OffloadMode.NONE:
-            )
+            return _streaming_model(self._streaming_text_encoder_builder, self._offload_mode, self._device, self._dtype)
-        return gpu_model(self._text_encoder_builder.build(device=self._device, dtype=self._dtype).eval())
+        return gpu_model(self._build_text_encoder())
 
     def __call__(
         self,
@@ -364,10 +446,9 @@ class PromptEncoder:
         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."""
+        """Encode *prompts* through Gemma -> embeddings processor, freeing each model after use."""
-        with self._text_encoder_ctx(streaming_prefetch_count) as text_encoder:
+        with self._text_encoder_ctx() as text_encoder:
             if enhance_first_prompt:
                 prompts = list(prompts)
                 prompts[0] = generate_enhanced_prompt(
@@ -375,9 +456,7 @@ class PromptEncoder:
                 )
             raw_outputs = [text_encoder.encode(p) for p in prompts]
 
-        with gpu_model(
+        with gpu_model(self._build_embeddings_processor()) as embeddings_processor:
-            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]
 
 
@@ -475,15 +554,21 @@ class VideoDecoder:
         dtype: torch.dtype,
         device: torch.device,
         registry: Registry | None = None,
+        memory_efficient: bool = True,
+        decoder_builder: BuilderProtocol | None = None,
     ) -> None:
         self._dtype = dtype
         self._device = device
-        self._decoder_builder = Builder(
+        if decoder_builder is not None:
-            model_path=checkpoint_path,
+            self._decoder_builder = decoder_builder
-            model_class_configurator=VideoDecoderConfigurator,
+        else:
-            model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
+            self._decoder_builder = Builder(
-            registry=registry or DummyRegistry(),
+                model_path=checkpoint_path,
-        )
+                model_class_configurator=VideoDecoderConfigurator,
+                model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
+                registry=registry or DummyRegistry(),
+                module_ops=(MEMORY_EFFICIENT_DECODE,) if memory_efficient else (),
+            )
 
     def __call__(
         self,

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

@@ -0,0 +1,224 @@
+"""Color space conversion utilities for video encoding.
+Provides GPU-accelerated RGB to YUV420 conversion that runs between the
+VAE decoder (which yields float RGB chunks) and ``encode_video``, bypassing
+pyav's CPU-side libswscale conversion. The ``FrameConverter`` also carries
+the codec metadata (pixel format, colour space, colour range) that
+``encode_video`` needs to tag the output stream.
+"""
+
+from __future__ import annotations
+
+import enum
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+
+class ColorSpace(enum.Enum):
+    """YUV color space standard."""
+
+    BT_709 = "bt709"
+    BT_2020_NCL = "bt2020ncl"
+
+    @property
+    def av_colorspace(self) -> int:
+        """FFmpeg ``AVCOL_SPC_*`` constant for ``codec_context.colorspace``."""
+        return _AV_COLORSPACE[self]
+
+
+class ColorRange(enum.Enum):
+    """YUV color range."""
+
+    MPEG = "mpeg"
+    JPEG = "jpeg"
+
+    @property
+    def av_color_range(self) -> int:
+        """FFmpeg ``AVCOL_RANGE_*`` constant for ``codec_context.color_range``."""
+        return _AV_COLOR_RANGE[self]
+
+
+class PixelFormat(enum.Enum):
+    """Pixel format for video frames."""
+
+    RGB24 = "rgb24"
+    YUV420P = "yuv420p"
+
+    @property
+    def av_format(self) -> str:
+        """PyAV format string for ``VideoFrame.from_ndarray``."""
+        return self.value
+
+
+_AV_COLORSPACE = {
+    ColorSpace.BT_709: 1,  # AVCOL_SPC_BT709
+    ColorSpace.BT_2020_NCL: 9,  # AVCOL_SPC_BT2020_NCL
+}
+
+_AV_COLOR_RANGE = {
+    ColorRange.MPEG: 1,  # AVCOL_RANGE_MPEG (limited)
+    ColorRange.JPEG: 2,  # AVCOL_RANGE_JPEG (full)
+}
+
+# BT.709 RGB->YUV matrix (row-major: each row produces one of Y, U, V)
+_BT709_MATRIX = torch.tensor(
+    [
+        [0.2126, 0.7152, 0.0722],
+        [-0.1146, -0.3854, 0.5],
+        [0.5, -0.4542, -0.0458],
+    ],
+    dtype=torch.float32,
+)
+
+# BT.2020 NCL RGB->YUV matrix
+_KR_2020 = 0.2627
+_KG_2020 = 0.6780
+_KB_2020 = 0.0593
+_BT2020_MATRIX = torch.tensor(
+    [
+        [_KR_2020, _KG_2020, _KB_2020],
+        [-_KR_2020 / 1.8814, -_KG_2020 / 1.8814, 0.5],
+        [0.5, -_KG_2020 / 1.4746, -_KB_2020 / 1.4746],
+    ],
+    dtype=torch.float32,
+)
+
+_COLOR_SPACE_MATRICES = {
+    ColorSpace.BT_709: _BT709_MATRIX,
+    ColorSpace.BT_2020_NCL: _BT2020_MATRIX,
+}
+
+
+@dataclass(frozen=True)
+class FrameConverter:
+    """Converts ``[*, C, H, W]`` float ``[0, 1]`` frames to uint8.
+    Carries encoding metadata so ``encode_video`` can derive pixel format,
+    color space, and color range from the converter itself.
+    The ``fn_`` callable **may mutate its input** (PyTorch trailing-underscore
+    convention).  Callers that need to keep the original ``frames`` afterwards
+    must pass ``frames.clone()``.  Inside ``encode_video``'s per-chunk
+    generator each chunk is consumed once, so direct passthrough is safe.
+    """
+
+    pixel_format: PixelFormat
+    fn_: Callable[[torch.Tensor], torch.Tensor] = field(repr=False)
+    color_space: ColorSpace | None = None
+    color_range: ColorRange | None = None
+
+    def __call__(self, frames: torch.Tensor) -> torch.Tensor:
+        return self.fn_(frames)
+
+
+def rgb_to_yuv(image: torch.Tensor, color_space: ColorSpace) -> torch.Tensor:
+    """Convert an RGB image to YUV.
+    The image data is assumed to be in the range of ``[0, 1]``.
+    Uses a single matrix multiply for better memory locality.
+    Args:
+        image: RGB image with shape ``(*, 3, H, W)``.
+        color_space: Color space standard for the conversion matrix.
+    Returns:
+        YUV image with shape ``(*, 3, H, W)``.
+    """
+    if len(image.shape) < 3 or image.shape[-3] != 3:
+        raise ValueError(f"Input size must have a shape of (*, 3, H, W). Got {image.shape}")
+
+    mat = _COLOR_SPACE_MATRICES[color_space].to(device=image.device, dtype=image.dtype)
+    # [*, 3, H, W] -> [*, H, W, 3] @ [3, 3]^T -> [*, H, W, 3] -> [*, 3, H, W]
+    pixels = image.movedim(-3, -1)  # [*, H, W, 3]
+    yuv = pixels @ mat.T  # [*, H, W, 3]
+    return yuv.movedim(-1, -3)  # [*, 3, H, W]
+
+
+def apply_color_range_(y: torch.Tensor, uv: torch.Tensor, color_range: ColorRange) -> tuple[torch.Tensor, torch.Tensor]:
+    """Scale Y and UV planes to the specified color range, in-place.
+    Args:
+        y: Luma plane in ``[0, 1]``.
+        uv: Chroma planes centered at 0.
+        color_range: Target color range.
+    Returns:
+        Scaled ``(Y, UV)`` tensors (modified in-place).
+    """
+    if color_range == ColorRange.MPEG:
+        y.mul_(219).add_(16)
+        uv.mul_(224).add_(128)
+    elif color_range == ColorRange.JPEG:
+        y.mul_(255)
+        uv.add_(0.5).mul_(255)
+    else:
+        raise ValueError(f"Unsupported color range: {color_range}")
+    return y, uv
+
+
+def rgb_to_yuv420(
+    image: torch.Tensor, color_space: ColorSpace, color_range: ColorRange
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Convert an RGB image to YUV 4:2:0 with chroma subsampling.
+    Chroma is subsampled by averaging 2x2 pixel blocks (chroma siting
+    ``(128, 128)``).
+    Args:
+        image: RGB image with shape ``(*, 3, H, W)`` in ``[0, 1]``.
+            H and W must be divisible by 2.
+        color_space: Color space standard.
+        color_range: Color range for the output.
+    Returns:
+        ``(Y, UV)`` where Y has shape ``(*, 1, H, W)`` and UV has shape
+        ``(*, 2, H//2, W//2)``.
+    """
+    if len(image.shape) < 3 or image.shape[-3] != 3:
+        raise ValueError(f"Input size must have a shape of (*, 3, H, W). Got {image.shape}")
+    if image.shape[-2] % 2 != 0 or image.shape[-1] % 2 != 0:
+        raise ValueError(f"Input H and W must be divisible by 2. Got {image.shape}")
+
+    yuv = rgb_to_yuv(image, color_space)
+    y = yuv[..., :1, :, :]
+    # Subsample chroma: average 2x2 blocks via avg_pool2d (contiguous, fused kernel)
+    uv_full = yuv[..., 1:3, :, :].contiguous()
+    # Flatten leading dims for avg_pool2d which expects [N, C, H, W]
+    lead = uv_full.shape[:-3]
+    uv_flat = uv_full.reshape(-1, 2, uv_full.shape[-2], uv_full.shape[-1])
+    uv = torch.nn.functional.avg_pool2d(uv_flat, kernel_size=2, stride=2)
+    uv = uv.reshape(*lead, 2, uv.shape[-2], uv.shape[-1])
+
+    return apply_color_range_(y, uv, color_range)
+
+
+def pack_i420(y: torch.Tensor, uv: torch.Tensor) -> torch.Tensor:
+    """Pack Y and UV planes into I420 layout for pyav.
+    I420 packs the three planes into a single 2D array of height ``H * 3 // 2``
+    and width ``W``.  The Y plane occupies the first ``H`` rows.  The UV tensor
+    ``(*, 2, H//2, W//2)`` is reshaped to ``(*, H//2, W)`` -- U rows packed
+    two-by-two followed by V rows packed two-by-two -- and appended below.
+    Args:
+        y: Luma with shape ``(*, 1, H, W)``.
+        uv: Chroma with shape ``(*, 2, H//2, W//2)``.
+    Returns:
+        Packed tensor with shape ``(*, H*3//2, W)`` uint8.
+    """
+    y_plane = y[..., 0, :, :]  # [*, H, W]
+    uv_packed = uv.reshape(*uv.shape[:-3], uv.shape[-2], uv.shape[-1] * 2)  # [*, H//2, W]
+    packed = torch.cat([y_plane, uv_packed], dim=-2)  # [*, H*3//2, W]
+    return packed.clamp_(0, 255).to(torch.uint8)
+
+
+def _rgb_uint8_fn_(frames: torch.Tensor) -> torch.Tensor:
+    """In-place: mutates ``frames`` via ``clamp_`` + ``mul_``, returns a uint8 view."""
+    return frames.clamp_(0.0, 1.0).mul_(255.0).to(torch.uint8).movedim(-3, -1)
+
+
+rgb_uint8_converter_ = FrameConverter(pixel_format=PixelFormat.RGB24, fn_=_rgb_uint8_fn_)
+"""``(*, 3, H, W)`` float ``[0, 1]`` to ``(*, H, W, 3)`` uint8.  Mutates input."""
+
+
+def _yuv420p_bt709_fn_(frames: torch.Tensor) -> torch.Tensor:
+    y, uv = rgb_to_yuv420(frames, ColorSpace.BT_709, ColorRange.MPEG)
+    return pack_i420(y, uv)
+
+
+yuv420p_bt709_converter_ = FrameConverter(
+    pixel_format=PixelFormat.YUV420P,
+    fn_=_yuv420p_bt709_fn_,
+    color_space=ColorSpace.BT_709,
+    color_range=ColorRange.MPEG,
+)
+"""``(*, 3, H, W)`` float ``[0, 1]`` to ``(*, H*3//2, W)`` uint8 YUV420p BT.709 MPEG."""

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

@@ -20,6 +20,7 @@ from ltx_core.guidance.perturbations import (
 from ltx_core.model.transformer import X0Model
 from ltx_core.types import LatentState
 from ltx_pipelines.utils.helpers import modality_from_latent_state
+from ltx_pipelines.utils.types import DenoisedLatentResult
 
 _POSITIVE_ONLY_GUIDER = MultiModalGuider(
     params=MultiModalGuiderParams(cfg_scale=1.0, stg_scale=0.0, modality_scale=1.0),
@@ -53,7 +54,7 @@ def _repeat_state(state: LatentState, n: int) -> LatentState:
     )
 
 
-def _guided_denoise(  # noqa: PLR0913
+def _guided_denoise(  # noqa: PLR0913,PLR0915
     transformer: X0Model,
     video_state: LatentState | None,
     audio_state: LatentState | None,
@@ -66,7 +67,8 @@ def _guided_denoise(  # noqa: PLR0913
     last_denoised_video: torch.Tensor | None,
     last_denoised_audio: torch.Tensor | None,
     step_index: int,
-) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    force_uncond_pass: bool = False,
+) -> tuple[DenoisedLatentResult | None, DenoisedLatentResult | 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
@@ -80,7 +82,9 @@ def _guided_denoise(  # noqa: PLR0913
     a_skip = audio_guider.should_skip_step(step_index)
 
     if v_skip and a_skip:
-        return last_denoised_video, last_denoised_audio
+        video_result = DenoisedLatentResult.result_or_none(denoised=last_denoised_video)
+        audio_result = DenoisedLatentResult.result_or_none(denoised=last_denoised_audio)
+        return video_result, audio_result
 
     if video_state is not None and v_context is None:
         raise ValueError("v_context is required when video_state is provided")
@@ -91,10 +95,12 @@ def _guided_denoise(  # noqa: PLR0913
     _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():
+    v_needs_neg = video_guider.do_unconditional_generation() or (force_uncond_pass and video_state is not None)
-        if video_guider.do_unconditional_generation() and video_guider.negative_context is None:
+    a_needs_neg = audio_guider.do_unconditional_generation() or (force_uncond_pass and audio_state is not None)
+    if v_needs_neg or a_needs_neg:
+        if v_needs_neg 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:
+        if a_needs_neg 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
@@ -172,7 +178,14 @@ def _guided_denoise(  # noqa: PLR0913
 
     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
+    return (
+        DenoisedLatentResult.result_or_none(
+            denoised=denoised_video, uncond=uncond_v, cond=cond_v, ptb=ptb_v, mod=mod_v
+        ),
+        DenoisedLatentResult.result_or_none(
+            denoised=denoised_audio, uncond=uncond_a, cond=cond_a, ptb=ptb_a, mod=mod_a
+        ),
+    )
 
 
 class SimpleDenoiser:
@@ -195,11 +208,15 @@ class SimpleDenoiser:
         audio_state: LatentState | None,
         sigmas: torch.Tensor,
         step_index: int,
-    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    ) -> tuple[DenoisedLatentResult | None, DenoisedLatentResult | 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)
+        denoised_video, denoised_audio = transformer(video=pos_video, audio=pos_audio, perturbations=None)
+        return (
+            DenoisedLatentResult.result_or_none(denoised=denoised_video),
+            DenoisedLatentResult.result_or_none(denoised=denoised_audio),
+        )
 
 
 class GuidedDenoiser:
@@ -214,11 +231,13 @@ class GuidedDenoiser:
         a_context: torch.Tensor | None,
         video_guider: MultiModalGuider | None = None,
         audio_guider: MultiModalGuider | None = None,
+        force_uncond_pass: bool = False,
     ) -> None:
         self.v_context = v_context
         self.a_context = a_context
         self.video_guider = video_guider
         self.audio_guider = audio_guider
+        self.force_uncond_pass = force_uncond_pass
         self._last_denoised_video: torch.Tensor | None = None
         self._last_denoised_audio: torch.Tensor | None = None
 
@@ -229,8 +248,8 @@ class GuidedDenoiser:
         audio_state: LatentState | None,
         sigmas: torch.Tensor,
         step_index: int,
-    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    ) -> tuple[DenoisedLatentResult | None, DenoisedLatentResult | None]:
-        denoised_video, denoised_audio = _guided_denoise(
+        guided_denoise_result_v, guided_denoise_result_a = _guided_denoise(
             transformer=transformer,
             video_state=video_state,
             audio_state=audio_state,
@@ -242,10 +261,11 @@ class GuidedDenoiser:
             last_denoised_video=self._last_denoised_video,
             last_denoised_audio=self._last_denoised_audio,
             step_index=step_index,
+            force_uncond_pass=self.force_uncond_pass,
         )
-        self._last_denoised_video = denoised_video
+        self._last_denoised_video = guided_denoise_result_v.denoised
-        self._last_denoised_audio = denoised_audio
+        self._last_denoised_audio = guided_denoise_result_a.denoised
-        return denoised_video, denoised_audio
+        return guided_denoise_result_v, guided_denoise_result_a
 
 
 class FactoryGuidedDenoiser:
@@ -257,11 +277,13 @@ class FactoryGuidedDenoiser:
         a_context: torch.Tensor | None,
         video_guider_factory: MultiModalGuiderFactory | None = None,
         audio_guider_factory: MultiModalGuiderFactory | None = None,
+        force_uncond_pass: bool = False,
     ) -> 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.force_uncond_pass = force_uncond_pass
         self._last_denoised_video: torch.Tensor | None = None
         self._last_denoised_audio: torch.Tensor | None = None
         self._sigma_vals_cached: list[float] | None = None
@@ -273,7 +295,7 @@ class FactoryGuidedDenoiser:
         audio_state: LatentState | None,
         sigmas: torch.Tensor,
         step_index: int,
-    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    ) -> tuple[DenoisedLatentResult | None, DenoisedLatentResult | None]:
         if self._sigma_vals_cached is None:
             self._sigma_vals_cached = sigmas.detach().cpu().tolist()
         sigma_val = self._sigma_vals_cached[step_index]
@@ -287,7 +309,7 @@ class FactoryGuidedDenoiser:
             else None
         )
 
-        denoised_video, denoised_audio = _guided_denoise(
+        guided_denoise_result_v, guided_denoise_result_a = _guided_denoise(
             transformer=transformer,
             video_state=video_state,
             audio_state=audio_state,
@@ -299,7 +321,8 @@ class FactoryGuidedDenoiser:
             last_denoised_video=self._last_denoised_video,
             last_denoised_audio=self._last_denoised_audio,
             step_index=step_index,
+            force_uncond_pass=self.force_uncond_pass,
         )
-        self._last_denoised_video = denoised_video
+        self._last_denoised_video = guided_denoise_result_v.denoised
-        self._last_denoised_audio = denoised_audio
+        self._last_denoised_audio = guided_denoise_result_a.denoised
-        return denoised_video, denoised_audio
+        return guided_denoise_result_v, guided_denoise_result_a

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

@@ -37,6 +37,11 @@ def cleanup_memory() -> None:
     gc.collect()
     torch.cuda.empty_cache()
     torch.cuda.synchronize()
+    try:
+        if hasattr(torch._C, "_host_emptyCache"):
+            torch._C._host_emptyCache()
+    except Exception:
+        logging.warning("Host empty cache cleanup failed; ignoring.", exc_info=True)
 
 
 def _conform_latent_length(latent: torch.Tensor, expected_frames_count: int) -> torch.Tensor:

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

@@ -1,23 +1,37 @@
+import enum
 import logging
 import math
+import threading
 from collections.abc import Generator, Iterator
 from fractions import Fraction
 from io import BytesIO
+from pathlib import Path
+from queue import Queue
 
 import av
 import numpy as np
+import OpenImageIO
 import torch
 from einops import rearrange
 from PIL import Image
 from torch._prims_common import DeviceLikeType
 from tqdm import tqdm
 
+from ltx_core.hdr import LogC3
 from ltx_core.types import Audio, VideoPixelShape
+from ltx_pipelines.utils.color_conversion import FrameConverter, PixelFormat, yuv420p_bt709_converter_
 from ltx_pipelines.utils.constants import DEFAULT_IMAGE_CRF
 
 logger = logging.getLogger(__name__)
 
 
+class ResizeMode(enum.Enum):
+    """How to fit a conditioning video to the target resolution."""
+
+    CENTER_CROP = "center_crop"
+    REFLECT_PAD = "reflect_pad"
+
+
 def resize_aspect_ratio_preserving(image: torch.Tensor, long_side: int) -> torch.Tensor:
     """
     Resize image preserving aspect ratio (filling target long side).
@@ -75,8 +89,18 @@ def resize_and_center_crop(tensor: torch.Tensor, height: int, width: int) -> tor
     return tensor
 
 
-def normalize_latent(latent: torch.Tensor, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+def normalize_images(images: torch.Tensor, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
-    return (latent / 127.5 - 1.0).to(device=device, dtype=dtype)
+    return (images / 127.5 - 1.0).to(device=device, dtype=dtype)
+
+
+def to_vae_range(x: torch.Tensor) -> torch.Tensor:
+    """Map [0, 1] to [-1, 1] (VAE input convention)."""
+    return torch.clamp(x, 0.0, 1.0) * 2.0 - 1.0
+
+
+def from_vae_range(z: torch.Tensor) -> torch.Tensor:
+    """Map [-1, 1] (VAE output convention) to [0, 1]."""
+    return torch.clamp((z + 1.0) / 2.0, 0.0, 1.0)
 
 
 def load_image_and_preprocess(
@@ -95,7 +119,7 @@ def load_image_and_preprocess(
     image = preprocess(image=image, crf=crf)
     image = torch.tensor(image, dtype=torch.float32, device=device)
     image = resize_and_center_crop(image, height, width)
-    image = normalize_latent(image, device, dtype)
+    image = normalize_images(image, device, dtype)
     return image
 
 
@@ -116,14 +140,118 @@ def video_preprocess(
     Returns:
         Tensor of shape (1, C, F, height, width) with values in [-1, 1].
     """
-    result = None
+    result: torch.Tensor | None = None
     for f in frames:
         frame = resize_and_center_crop(f.to(torch.float32), height, width)
-        frame = normalize_latent(frame, device, dtype)
+        frame = normalize_images(frame, device, dtype)
         result = frame if result is None else torch.cat([result, frame], dim=2)
+    if result is None:
+        raise ValueError("video_preprocess received an empty frame generator; no frames were decoded from the source.")
     return result
 
 
+def align_resolution(
+    width: int,
+    height: int,
+    resize_mode: ResizeMode,
+    divisor: int = 64,
+) -> tuple[int, int, int, int]:
+    """Compute aligned generation dimensions and crop-back size.
+    Args:
+        width: Source video width (need not be aligned).
+        height: Source video height (need not be aligned).
+        resize_mode: CENTER_CROP rounds down; REFLECT_PAD rounds up.
+        divisor: Alignment divisor (default 64 for two-stage pipelines).
+    Returns:
+        ``(gen_width, gen_height, crop_width, crop_height)`` where
+        ``gen_*`` are multiples of *divisor* and ``crop_*`` are the
+        original dimensions to trim back to after decoding.  When no
+        cropping is needed ``crop_*`` equals ``gen_*``.
+    """
+    if resize_mode is ResizeMode.REFLECT_PAD:
+        gen_w = ((width + divisor - 1) // divisor) * divisor
+        gen_h = ((height + divisor - 1) // divisor) * divisor
+    else:
+        gen_w = (width // divisor) * divisor
+        gen_h = (height // divisor) * divisor
+
+    crop_w = width if gen_w != width else gen_w
+    crop_h = height if gen_h != height else gen_h
+    return gen_w, gen_h, crop_w, crop_h
+
+
+def resize_and_reflect_pad(tensor: torch.Tensor, height: int, width: int) -> torch.Tensor:
+    """Resize tensor to fit within target, then reflect-pad to exact dimensions.
+    Unlike resize_and_center_crop which stretches and crops, this preserves the
+    original aspect ratio and pads the shorter dimension with reflected pixels.
+    When the target is already >= the source in both dimensions, interpolation
+    is skipped entirely to preserve original pixels.
+    Args:
+        tensor: Input with shape (H, W, C) or (F, H, W, C)
+        height: Target height
+        width: Target width
+    Returns:
+        Tensor with shape (1, C, 1, height, width) for 3D or (1, C, F, height, width) for 4D
+    """
+    if tensor.ndim == 3:
+        tensor = rearrange(tensor, "h w c -> 1 c h w")
+    elif tensor.ndim == 4:
+        tensor = rearrange(tensor, "f h w c -> f c h w")
+    else:
+        raise ValueError(f"Expected input with 3 or 4 dimensions; got shape {tensor.shape}.")
+
+    _, _, src_h, src_w = tensor.shape
+
+    if height >= src_h and width >= src_w:
+        new_h, new_w = src_h, src_w
+    else:
+        scale = min(height / src_h, width / src_w)
+        new_h = round(src_h * scale)
+        new_w = round(src_w * scale)
+        tensor = torch.nn.functional.interpolate(tensor, size=(new_h, new_w), mode="bilinear", align_corners=False)
+
+    pad_bottom = height - new_h
+    pad_right = width - new_w
+    if pad_bottom > 0 or pad_right > 0:
+        pad_mode = "reflect" if pad_bottom < new_h and pad_right < new_w else "replicate"
+        tensor = torch.nn.functional.pad(tensor, (0, pad_right, 0, pad_bottom), mode=pad_mode)
+
+    tensor = rearrange(tensor, "f c h w -> 1 c f h w")
+    return tensor
+
+
+def load_video_conditioning_hdr(
+    video_path: str,
+    height: int,
+    width: int,
+    frame_cap: int,
+    dtype: torch.dtype,
+    device: torch.device,
+    hdr_transform: str = "logc3",
+    resize_mode: ResizeMode = ResizeMode.CENTER_CROP,
+) -> Iterator[torch.Tensor]:
+    """Load a video and yield preprocessed frames for HDR IC-LoRA conditioning.
+    Decodes through the standard path and applies the LDR compression that
+    matches training. Callers are responsible for providing Rec.709 SDR
+    input — the HDR IC-LoRA was trained on that color space.
+    Args:
+        hdr_transform: LDR-compression name (currently only ``logc3``).
+        resize_mode: How to fit the video to the target resolution.
+    Yields:
+        Per-frame tensors of shape ``(1, C, 1, height, width)``.
+    """
+    if hdr_transform != "logc3":
+        raise ValueError(f"Unsupported HDR transform: {hdr_transform}")
+
+    resize_fn = resize_and_reflect_pad if resize_mode is ResizeMode.REFLECT_PAD else resize_and_center_crop
+
+    for f in decode_video_by_frame(path=video_path, frame_cap=frame_cap, device=device):
+        frame = resize_fn(f.to(torch.float32), height, width)
+        ldr = (frame / 255.0).clamp(0.0, 1.0)
+        compressed = LogC3().compress_ldr(ldr)
+        yield to_vae_range(compressed).to(device=device, dtype=dtype)
+
+
 def decode_image(image_path: str) -> np.ndarray:
     image = Image.open(image_path)
     np_array = np.array(image)[..., :3]
@@ -202,47 +330,120 @@ def encode_video(
     audio: Audio | None,
     output_path: str,
     video_chunks_number: int,
+    frame_converter: FrameConverter = yuv420p_bt709_converter_,
+    crf: int = 19,
+    preset: str = "veryfast",
+    thread_count: int = 0,
 ) -> None:
     if isinstance(video, torch.Tensor):
         video = iter([video])
 
-    first_chunk = next(video)
+    def convert(chunk: torch.Tensor) -> torch.Tensor:
+        return frame_converter(chunk.movedim(-1, -3))
 
-    _, height, width, _ = first_chunk.shape
+    first_chunk = convert(next(video))
+
+    if frame_converter.pixel_format == PixelFormat.RGB24:
+        height, width = first_chunk.shape[-3], first_chunk.shape[-2]
+    else:
+        height = first_chunk.shape[-2] * 2 // 3
+        width = first_chunk.shape[-1]
 
     container = av.open(output_path, mode="w")
-    stream = container.add_stream("libx264", rate=int(fps))
+    success = False
-    stream.width = width
+    try:
-    stream.height = height
+        stream = container.add_stream("libx264", rate=int(fps), options={"crf": str(crf), "preset": preset})
-    stream.pix_fmt = "yuv420p"
+        stream.width = width
+        stream.height = height
+        stream.pix_fmt = "yuv420p"
+        stream.codec_context.thread_count = thread_count
+        stream.codec_context.thread_type = "FRAME"
+        if frame_converter.color_space is not None:
+            stream.codec_context.colorspace = frame_converter.color_space.av_colorspace
+        if frame_converter.color_range is not None:
+            stream.codec_context.color_range = frame_converter.color_range.av_color_range
 
-    if audio is not None:
+        if audio is not None:
-        audio_stream = _prepare_audio_stream(container, audio.sampling_rate)
+            audio_stream = _prepare_audio_stream(container, audio.sampling_rate)
 
-    def all_tiles(
+        av_format = frame_converter.pixel_format.av_format
-        first_chunk: torch.Tensor, tiles_generator: Generator[tuple[torch.Tensor, int], None, None]
-    ) -> Generator[tuple[torch.Tensor, int], None, None]:
-        yield first_chunk
-        yield from tiles_generator
 
-    for video_chunk in tqdm(all_tiles(first_chunk, video), total=video_chunks_number):
+        def cpu_chunks() -> Generator[np.ndarray, None, None]:
-        video_chunk_cpu = video_chunk.to("cpu").numpy()
+            yield first_chunk.to("cpu").numpy()
-        for frame_array in video_chunk_cpu:
+            for chunk in video:
-            frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24")
+                yield convert(chunk).to("cpu").numpy()
-            for packet in stream.encode(frame):
-                container.mux(packet)
 
-    # Flush encoder
+        _encode_chunks_threaded(
-    for packet in stream.encode():
+            container=container,
-        container.mux(packet)
+            stream=stream,
+            av_format=av_format,
+            chunks=cpu_chunks(),
+            progress_total=video_chunks_number,
+        )
 
-    if audio is not None:
+        if audio is not None:
-        _write_audio(container, audio_stream, audio)
+            _write_audio(container, audio_stream, audio)
-
+        success = True
-    container.close()
+    finally:
+        container.close()
+        if not success:
+            Path(output_path).unlink(missing_ok=True)
     logger.info(f"Video saved to {output_path}")
 
 
+def _encode_chunks_threaded(
+    container: av.container.Container,
+    stream: av.video.stream.VideoStream,
+    av_format: str,
+    chunks: Iterator[np.ndarray],
+    progress_total: int,
+) -> None:
+    """Run libx264 frame.encode + container.mux on a background thread while
+    the caller produces numpy chunks on the current thread. The 1-slot queue
+    lets the producer get one chunk ahead (so the next VAE/gather chunk
+    overlaps with libx264 encoding the previous chunk) without buffering more
+    than one chunk in CPU memory.
+    """
+    chunk_queue: Queue[np.ndarray | None] = Queue(maxsize=1)
+    encoder_error: list[BaseException] = []
+
+    def encoder_worker() -> None:
+        error: BaseException | None = None
+        while True:
+            arr = chunk_queue.get()
+            if arr is None:
+                break
+            if error is not None:
+                continue
+            try:
+                for frame_array in arr:
+                    frame = av.VideoFrame.from_ndarray(frame_array, format=av_format)
+                    for packet in stream.encode(frame):
+                        container.mux(packet)
+            except Exception as e:
+                error = e
+        if error is None:
+            try:
+                for packet in stream.encode():
+                    container.mux(packet)
+            except Exception as e:
+                error = e
+        if error is not None:
+            encoder_error.append(error)
+
+    encoder_thread = threading.Thread(target=encoder_worker, name="h264-encoder")
+    encoder_thread.start()
+    try:
+        for arr in tqdm(chunks, total=progress_total):
+            chunk_queue.put(arr)
+    finally:
+        chunk_queue.put(None)
+        encoder_thread.join()
+
+    if encoder_error:
+        raise encoder_error[0]
+
+
 _INT_FORMAT_MAX: dict[str, float] = {
     "u8": 128.0,
     "u8p": 128.0,
@@ -481,3 +682,85 @@ def preprocess(image: np.array, crf: float = DEFAULT_IMAGE_CRF) -> np.array:
     with BytesIO(video_bytes) as video_file:
         image_array = decode_single_frame(video_file)
     return image_array
+
+
+def save_exr_tensor(tensor: torch.Tensor, file_path: str | Path, half: bool = False) -> None:
+    """Save a single tensor frame as EXR with linear sRGB colorspace metadata.
+    Args:
+        tensor: ``[H, W, C]`` or ``[C, H, W]`` float tensor.
+        file_path: Output path (e.g. ``frame_0000.exr``).
+        half: Force float16 output with ZIP compression.
+    """
+    if tensor.dim() == 3 and tensor.shape[0] == 3:
+        tensor = tensor.permute(1, 2, 0)
+    use_half = half or tensor.dtype in (torch.float16, torch.half)
+    img_np = np.ascontiguousarray(tensor.cpu().numpy().astype(np.float32))
+    file_path = str(file_path)
+
+    h, w = img_np.shape[:2]
+    fmt = OpenImageIO.HALF if use_half else OpenImageIO.FLOAT
+    spec = OpenImageIO.ImageSpec(w, h, 3, fmt)
+    spec.channelnames = ("R", "G", "B")
+    spec.attribute("compression", "zip")
+    spec.attribute("chromaticities", "float[8]", (0.64, 0.33, 0.30, 0.60, 0.15, 0.06, 0.3127, 0.3290))
+    spec.attribute("colorSpace", "sRGB")
+
+    out = OpenImageIO.ImageOutput.create(file_path)
+    if out is None:
+        raise RuntimeError(
+            f"Failed to create EXR writer for '{file_path}'. Ensure OpenImageIO is built with OpenEXR support."
+        )
+    try:
+        if not out.open(file_path, spec):
+            raise RuntimeError(f"Failed to open EXR file '{file_path}': {out.geterror()}")
+        if not out.write_image(img_np):
+            raise RuntimeError(f"Failed to write EXR image '{file_path}': {out.geterror()}")
+    finally:
+        out.close()
+
+
+def _linear_to_srgb(x: np.ndarray) -> np.ndarray:
+    """Linear -> sRGB OETF per IEC 61966-2-1. Input assumed in [0, 1]."""
+    x = np.clip(x, 0.0, 1.0)
+    return np.where(x <= 0.0031308, x * 12.92, 1.055 * np.power(x, 1.0 / 2.4) - 0.055)
+
+
+def encode_exr_sequence_to_mp4(exr_dir: Path, output_mp4: Path, frame_rate: float) -> None:
+    """Convert a linear EXR frame sequence to sRGB and encode to H.264 .mp4 via PyAV.
+    Exposure is fixed at EV=0 (no gain). Each EXR frame is clamped to [0, 1],
+    passed through the sRGB OETF, quantised to 8-bit BGR, and fed to a libx264
+    stream (crf 18, yuv420p). ``frame_rate`` is the original source video's
+    frame rate so playback matches the input timing.
+    """
+    import os  # noqa: PLC0415
+
+    os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+    import cv2  # noqa: PLC0415
+
+    exr_files = sorted(exr_dir.glob("frame_*.exr"))
+    if not exr_files:
+        raise FileNotFoundError(f"No EXR frames found in {exr_dir}")
+
+    container = av.open(str(output_mp4), mode="w")
+    stream = container.add_stream("libx264", rate=Fraction(frame_rate).limit_denominator(1000))
+    stream.pix_fmt = "yuv420p"
+    stream.options = {"crf": "18", "movflags": "+faststart"}
+
+    try:
+        for i, exr_path in enumerate(exr_files):
+            hdr = cv2.imread(str(exr_path), cv2.IMREAD_UNCHANGED).astype(np.float32)
+            sdr = _linear_to_srgb(np.maximum(hdr, 0.0))
+            bgr8 = (sdr * 255.0 + 0.5).astype(np.uint8)
+
+            if i == 0:
+                stream.height = bgr8.shape[0]
+                stream.width = bgr8.shape[1]
+
+            frame = av.VideoFrame.from_ndarray(bgr8, format="bgr24")
+            for packet in stream.encode(frame):
+                container.mux(packet)
+
+        for packet in stream.encode():
+            container.mux(packet)
+    finally:
+        container.close()

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

@@ -6,7 +6,7 @@ from typing import Callable
 import torch
 from tqdm import tqdm
 
-from ltx_core.components.diffusion_steps import Res2sDiffusionStep
+from ltx_core.components.diffusion_steps import EulerCfgPpDiffusionStep, Res2sDiffusionStep
 from ltx_core.components.protocols import DiffusionStepProtocol
 from ltx_core.model.transformer import X0Model
 from ltx_core.utils import to_denoised, to_velocity
@@ -60,13 +60,15 @@ def euler_denoising_loop(
     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)``.
+        and must return a :class:`~ltx_pipelines.utils.types.DenoisedLatentResult`.
     ### Returns
     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 = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
+        video_result, audio_result = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
+        denoised_video = video_result.denoised if video_result is not None else None
+        denoised_audio = audio_result.denoised if audio_result is not None else None
 
         video_state = _step_state(video_state, denoised_video, stepper, sigmas, step_idx)
         audio_state = _step_state(audio_state, denoised_audio, stepper, sigmas, step_idx)
@@ -110,7 +112,9 @@ def gradient_estimating_euler_denoising_loop(
         return current_velocity, denoised_sample
 
     for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
-        denoised_video, denoised_audio = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
+        video_result, audio_result = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
+        denoised_video = video_result.denoised if video_result is not None else None
+        denoised_audio = audio_result.denoised if audio_result is not None else None
 
         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)
@@ -143,6 +147,11 @@ def gradient_estimating_euler_denoising_loop(
     return (video_state, audio_state)
 
 
+def _get_plain_noise(x: torch.Tensor, generator: torch.Generator) -> torch.Tensor:
+    """Draw standard Gaussian noise matching the shape, dtype, and device of ``x``."""
+    return torch.randn(x.shape, generator=generator, dtype=x.dtype, device=x.device)
+
+
 def _channelwise_normalize(x: torch.Tensor) -> torch.Tensor:
     return x.sub_(x.mean(dim=(-2, -1), keepdim=True)).div_(x.std(dim=(-2, -1), keepdim=True))
 
@@ -278,7 +287,9 @@ def res2s_audio_video_denoising_loop(  # noqa: PLR0913,PLR0915,PLR0912
         # ====================================================================
         # STAGE 1: Evaluate at current point
         # ====================================================================
-        denoised_video_1, denoised_audio_1 = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
+        video_result, audio_result = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
+        denoised_video_1 = video_result.denoised if video_result is not None else None
+        denoised_audio_1 = audio_result.denoised if audio_result is not None else None
         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:
@@ -355,13 +366,15 @@ def res2s_audio_video_denoising_loop(  # noqa: PLR0913,PLR0915,PLR0912
             else None
         )
 
-        denoised_video_2, denoised_audio_2 = denoiser(
+        video_result_2, audio_result_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 = video_result_2.denoised if video_result_2 is not None else None
+        denoised_audio_2 = audio_result_2.denoised if audio_result_2 is not None else None
         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:
@@ -410,7 +423,9 @@ def res2s_audio_video_denoising_loop(  # noqa: PLR0913,PLR0915,PLR0912
 
     # Final step if we need to fully remove the noise
     if sigmas[-1] == 0:
-        denoised_video_1, denoised_audio_1 = denoiser(transformer, video_state, audio_state, sigmas, n_full_steps)
+        video_result_final, audio_result_final = denoiser(transformer, video_state, audio_state, sigmas, n_full_steps)
+        denoised_video_1 = video_result_final.denoised if video_result_final is not None else None
+        denoised_audio_1 = audio_result_final.denoised if audio_result_final is not None else None
         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))
@@ -419,3 +434,121 @@ def res2s_audio_video_denoising_loop(  # noqa: PLR0913,PLR0915,PLR0912
             audio_state = replace(audio_state, latent=denoised_audio_1.to(model_dtype))
 
     return video_state, audio_state
+
+
+def euler_cfg_pp_denoising_loop(
+    sigmas: torch.Tensor,
+    video_state: LatentState | None,
+    audio_state: LatentState | None,
+    stepper: EulerCfgPpDiffusionStep,
+    transformer: X0Model,
+    denoiser: Denoiser,
+    noise_seed: int = -1,
+    new_noise_fn: Callable[[torch.Tensor, torch.Generator], torch.Tensor] = _get_plain_noise,
+    model_dtype: torch.dtype = torch.bfloat16,
+) -> tuple[LatentState | None, LatentState | None]:
+    """
+    Joint audio-video denoising loop using the CFG++ corrected Euler sampler.
+    Applies the CFG++ update rule at each step: the ODE derivative is computed
+    from the unconditioned denoised prediction rather than the standard velocity,
+    and an ancestral DDIM noise injection is applied in the rescaled sigma space.
+    Requires a guided denoiser whose :class:`~ltx_pipelines.utils.types.DenoisedLatentResult`
+    carries ``uncond`` tensors (i.e. CFG must be enabled).
+    Either ``video_state`` or ``audio_state`` may be ``None`` for absent modalities.
+    When both are present, noise is drawn from the same seeded generator (video
+    first, audio second) to produce a consistent random sequence.
+    ### Parameters
+    sigmas:
+        1-D tensor of noise levels defining the sampling schedule.
+    video_state:
+        Current video :class:`~ltx_core.types.LatentState`, or ``None``.
+    audio_state:
+        Current audio :class:`~ltx_core.types.LatentState`, or ``None``.
+    stepper:
+        :class:`~ltx_core.components.diffusion_steps.EulerCfgPpDiffusionStep`
+        instance carrying ``eta`` and ``s_noise`` parameters.
+    transformer:
+        The diffusion model passed to the denoiser at each step.
+    denoiser:
+        Callable implementing :class:`~ltx_pipelines.utils.types.Denoiser`.
+    noise_seed:
+        Integer seed for the noise generator. Default ``-1``.
+    new_noise_fn:
+        ``(latent, generator) -> noise`` callable. Defaults to plain
+        ``torch.randn`` (no channel-wise normalization). Pass
+        :func:`_get_new_noise` for the normalized variant used in res2s.
+    model_dtype:
+        Dtype for latent state updates. Default ``bfloat16``.
+    ### Returns
+    tuple[LatentState | None, LatentState | None]
+        Final ``(video_state, audio_state)`` after the denoising loop.
+    """
+    if not isinstance(stepper, EulerCfgPpDiffusionStep):
+        raise ValueError(f"stepper must be an instance of EulerCfgPpDiffusionStep, got {type(stepper).__name__}")
+
+    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")
+
+    generator = torch.Generator(device=present_state.latent.device).manual_seed(noise_seed)
+    draw_noise = stepper.eta > 0 and stepper.s_noise > 0
+
+    for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
+        video_result, audio_result = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
+        denoised_video = video_result.denoised if video_result is not None else None
+        denoised_audio = audio_result.denoised if audio_result is not None else None
+        uncond_video = video_result.uncond if video_result is not None else None
+        uncond_audio = audio_result.uncond if audio_result is not None else None
+
+        if video_state is not None and not isinstance(uncond_video, torch.Tensor):
+            raise ValueError(
+                "euler_cfg_pp_denoising_loop requires video DenoisedLatentResult.uncond to be a tensor. "
+                "Use GuidedDenoiser or FactoryGuidedDenoiser with cfg_scale != 1 "
+                "or force_uncond_pass=True and a negative_context."
+            )
+        if audio_state is not None and not isinstance(uncond_audio, torch.Tensor):
+            raise ValueError(
+                "euler_cfg_pp_denoising_loop requires audio DenoisedLatentResult.uncond to be a tensor. "
+                "Use GuidedDenoiser or FactoryGuidedDenoiser with cfg_scale != 1 "
+                "or force_uncond_pass=True and a negative_context."
+            )
+
+        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:
+            if video_state is not None and denoised_video is not None:
+                video_state = replace(video_state, latent=denoised_video.to(model_dtype))
+            if audio_state is not None and denoised_audio is not None:
+                audio_state = replace(audio_state, latent=denoised_audio.to(model_dtype))
+            return video_state, audio_state
+
+        # Draw noise consecutively from the same generator: video first, audio second.
+        noise_video = new_noise_fn(video_state.latent, generator) if (video_state is not None and draw_noise) else None
+        noise_audio = new_noise_fn(audio_state.latent, generator) if (audio_state is not None and draw_noise) else None
+
+        if video_state is not None and denoised_video is not None:
+            x_next = stepper.step(
+                sample=video_state.latent,
+                denoised_sample=denoised_video,
+                sigmas=sigmas,
+                step_index=step_idx,
+                uncond_denoised=uncond_video,
+                noise=noise_video,
+            )
+            video_state = replace(video_state, latent=x_next.to(model_dtype))
+
+        if audio_state is not None and denoised_audio is not None:
+            x_next = stepper.step(
+                sample=audio_state.latent,
+                denoised_sample=denoised_audio,
+                sigmas=sigmas,
+                step_index=step_idx,
+                uncond_denoised=uncond_audio,
+                noise=noise_audio,
+            )
+            audio_state = replace(audio_state, latent=x_next.to(model_dtype))
+
+    return video_state, audio_state

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

@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 from dataclasses import dataclass, field
+from enum import Enum
 from typing import Protocol
 
 import torch
@@ -37,6 +40,36 @@ class PipelineComponents:
         self.audio_patchifier = AudioPatchifier(patch_size=1)
 
 
+@dataclass(frozen=True)
+class DenoisedLatentResult:
+    """Output of one denoiser call for a single modality.
+    ``denoised`` is the final blended prediction for this modality.
+    The remaining fields carry the per-pass raw outputs from ``_guided_denoise``
+    (all ``None`` for ``SimpleDenoiser``).  Denoisers return a
+    ``(video_result, audio_result)`` tuple; either element may be ``None``
+    for absent modalities.
+    """
+
+    denoised: torch.Tensor
+    uncond: torch.Tensor | None = None
+    cond: torch.Tensor | None = None
+    ptb: torch.Tensor | None = None
+    mod: torch.Tensor | None = None
+
+    @classmethod
+    def result_or_none(
+        cls,
+        denoised: torch.Tensor | None,
+        uncond: torch.Tensor | None = None,
+        cond: torch.Tensor | None = None,
+        ptb: torch.Tensor | None = None,
+        mod: torch.Tensor | None = None,
+    ) -> DenoisedLatentResult | None:
+        if denoised is None:
+            return None
+        return cls(denoised=denoised, uncond=uncond, cond=cond, ptb=ptb, mod=mod)
+
+
 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
@@ -48,7 +81,8 @@ class Denoiser(Protocol):
         sigmas: 1-D tensor of sigma values for each diffusion step.
         step_index: Index of the current denoising step.
     Returns:
-        ``(denoised_video, denoised_audio)`` tensors (either may be ``None``).
+        A ``(video_result, audio_result)`` tuple of :class:`DenoisedLatentResult`,
+        either may be ``None`` for absent modalities.
     """
 
     def __call__(
@@ -58,7 +92,7 @@ class Denoiser(Protocol):
         audio_state: LatentState | None,
         sigmas: torch.Tensor,
         step_index: int,
-    ) -> tuple[torch.Tensor | None, torch.Tensor | None]: ...
+    ) -> tuple[DenoisedLatentResult | None, DenoisedLatentResult | None]: ...
 
 
 @dataclass(frozen=True)
@@ -74,3 +108,21 @@ class ModalitySpec:
     noise_scale: float = 1.0
     frozen: bool = False
     initial_latent: torch.Tensor | None = None
+
+
+class OffloadMode(Enum):
+    """Weight offloading strategy.
+    Controls where model weights reside during inference:
+    - ``NONE``: All weights on GPU (no streaming). Fastest inference,
+      requires enough VRAM for the full model (~28 GB for LTX-2).
+    - ``CPU``: Weights pinned in CPU RAM, streamed layer-by-layer to a
+      small GPU buffer. First pass reads from disk; subsequent passes
+      reuse the CPU cache. Requires ~36 GB RAM + ~5 GB VRAM.
+    - ``DISK``: Weights read from disk on demand through a small CPU
+      buffer, then streamed to GPU. Every pass re-reads from disk.
+      Lowest memory: ~5 GB RAM + ~5 GB VRAM.
+    """
+
+    NONE = "none"
+    CPU = "cpu"
+    DISK = "disk"

packages/ltx-trainer/configs/ltx2_av_lora.yaml

@@ -166,6 +166,11 @@ acceleration:
   # Useful when GPU memory is limited
   load_text_encoder_in_8bit: false
 
+  # Offload optimizer state to CPU during validation video sampling.
+  # Helps avoid OOM when VAE decoder + transformer + optimizer state can't coexist
+  # on the GPU (full fine-tune, high-rank LoRA). No effect for FSDP.
+  offload_optimizer_during_validation: false
+
 # -----------------------------------------------------------------------------
 # Data Configuration
 # -----------------------------------------------------------------------------
@@ -219,9 +224,6 @@ validation:
   # Set to null to disable validation during training
   interval: 100
 
-  # Number of videos to generate per prompt
-  videos_per_prompt: 1
-
   # Classifier-free guidance scale
   # Higher values = stronger adherence to prompt but may introduce artifacts
   guidance_scale: 4.0

packages/ltx-trainer/configs/ltx2_av_lora_low_vram.yaml

@@ -178,6 +178,11 @@ acceleration:
   # Useful when GPU memory is limited
   load_text_encoder_in_8bit: true
 
+  # Offload optimizer state to CPU during validation video sampling.
+  # Helps avoid OOM when VAE decoder + transformer + optimizer state can't coexist
+  # on the GPU (full fine-tune, high-rank LoRA). No effect for FSDP.
+  offload_optimizer_during_validation: true
+
 # -----------------------------------------------------------------------------
 # Data Configuration
 # -----------------------------------------------------------------------------
@@ -231,9 +236,6 @@ validation:
   # Set to null to disable validation during training
   interval: 100
 
-  # Number of videos to generate per prompt
-  videos_per_prompt: 1
-
   # Classifier-free guidance scale
   # Higher values = stronger adherence to prompt but may introduce artifacts
   guidance_scale: 4.0

packages/ltx-trainer/configs/ltx2_v2v_ic_lora.yaml

@@ -166,6 +166,11 @@ acceleration:
   # Useful when GPU memory is limited
   load_text_encoder_in_8bit: false
 
+  # Offload optimizer state to CPU during validation video sampling.
+  # Helps avoid OOM when VAE decoder + transformer + optimizer state can't coexist
+  # on the GPU (full fine-tune, high-rank LoRA). No effect for FSDP.
+  offload_optimizer_during_validation: false
+
 # -----------------------------------------------------------------------------
 # Data Configuration
 # -----------------------------------------------------------------------------
@@ -232,9 +237,6 @@ validation:
   # Set to null to disable validation during training
   interval: 100
 
-  # Number of videos to generate per prompt
-  videos_per_prompt: 1
-
   # Classifier-free guidance scale
   # Higher values = stronger adherence to prompt but may introduce artifacts
   guidance_scale: 4.0

packages/ltx-trainer/docs/configuration-reference.md

@@ -215,18 +215,20 @@ Hardware acceleration and compute optimization settings.
 
 ```yaml
 acceleration:
-  mixed_precision_mode: "bf16"      # "no", "fp16", or "bf16"
+  mixed_precision_mode: "bf16"                  # "no", "fp16", or "bf16"
-  quantization: null                # Quantization options
+  quantization: null                            # Quantization options
-  load_text_encoder_in_8bit: false  # Load text encoder in 8-bit
+  load_text_encoder_in_8bit: false              # Load text encoder in 8-bit
+  offload_optimizer_during_validation: false    # Offload optimizer state to CPU during validation
 ```
 
 **Key parameters:**
 
-| Parameter                   | Description                                                                        |
+| Parameter                             | Description                                                                                                                                                                              |
-|-----------------------------|------------------------------------------------------------------------------------|
+|---------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| `mixed_precision_mode`      | Precision mode - `"bf16"` recommended for modern GPUs                              |
+| `mixed_precision_mode`                | Precision mode - `"bf16"` recommended for modern GPUs                                                                                                                                    |
-| `quantization`              | Model quantization: `null`, `"int8-quanto"`, `"int4-quanto"`, `"fp8-quanto"`, etc. |
+| `quantization`                        | Model quantization: `null`, `"int8-quanto"`, `"int4-quanto"`, `"fp8-quanto"`, etc.                                                                                                       |
-| `load_text_encoder_in_8bit` | Load the Gemma text encoder in 8-bit to save GPU memory                            |
+| `load_text_encoder_in_8bit`           | Load the Gemma text encoder in 8-bit to save GPU memory                                                                                                                                  |
+| `offload_optimizer_during_validation` | Move optimizer state to CPU before validation video sampling and back afterwards. Useful when validation OOMs because VAE decoder + transformer + optimizer state can't coexist on the GPU (full fine-tune, high-rank LoRA). No effect for FSDP. |
 
 ### DataConfig
 
@@ -262,7 +264,6 @@ validation:
   seed: 42                            # Random seed for reproducibility
   inference_steps: 30                 # Number of inference steps
   interval: 100                       # Steps between validation runs
-  videos_per_prompt: 1                # Videos generated per prompt
   guidance_scale: 4.0                 # CFG guidance strength
   stg_scale: 1.0                      # STG guidance strength (0.0 to disable)
   stg_blocks: [ 29 ]                  # Transformer blocks to perturb for STG

packages/ltx-trainer/docs/dataset-preparation.md

@@ -50,17 +50,20 @@ This will create a `dataset.json` file containing video paths and their captions
 
 **Captioning options:**
 
-| Option | Description |
+
-|--------|-------------|
+| Option             | Description                                                |
-| `--captioner-type` | `qwen_omni` (default, local) or `gemini_flash` (API) |
+| ------------------ | ---------------------------------------------------------- |
-| `--use-8bit` | Enable 8-bit quantization for lower VRAM usage |
+| `--captioner-type` | `qwen_omni` (default, local) or `gemini_flash` (API)       |
-| `--no-audio` | Disable audio processing (video-only captions) |
+| `--use-8bit`       | Enable 8-bit quantization for lower VRAM usage             |
-| `--override` | Re-caption files that already have captions |
+| `--no-audio`       | Disable audio processing (video-only captions)             |
-| `--api-key` | API key for Gemini Flash (or set `GOOGLE_API_KEY` env var) |
+| `--override`       | Re-caption files that already have captions                |
+| `--api-key`        | API key for Gemini Flash (or set `GOOGLE_API_KEY` env var) |
+
 
 **Caption format:**
 
 The captioner produces structured captions with sections for:
+
 - **Visual content**: People, objects, actions, settings, colors, movements
 - **Speech transcription**: Word-for-word transcription of spoken content
 - **Sounds**: Music, ambient sounds, sound effects
@@ -106,15 +109,57 @@ uv run python scripts/process_dataset.py dataset.json \
     --with-audio
 ```
 
+### 🚀 Multi-GPU Preprocessing
+
+Preprocessing large datasets can take a while. To run it across multiple GPUs in parallel, wrap the command with
+`accelerate launch` (for example `--num_processes 4`). Each process handles an interleaved slice of the dataset.
+The same approach applies to `process_videos.py` and `process_captions.py` when you run them standalone.
+
+```bash
+uv run accelerate launch --num_processes 4 scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x49" \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --text-encoder-path /path/to/gemma-model
+```
+
+Outputs are written atomically (via a per-process temporary file, then renamed), so an interrupted run leaves no
+corrupt files. By default a rerun **resumes** — items whose output `.pt` already exists are skipped.
+
+> [!IMPORTANT]
+> Pass `**--overwrite`** when rerunning with changed parameters (different model checkpoint, resolution buckets,
+> text encoder, `--lora-trigger`, etc.). Without it the script keeps the stale outputs from the previous run.
+>
+> ```bash
+> uv run accelerate launch --num_processes 4 scripts/process_dataset.py dataset.json \
+>     --resolution-buckets "960x544x49" \
+>     --model-path /path/to/ltx-2.3-model.safetensors \
+>     --text-encoder-path /path/to/gemma-model \
+>     --overwrite
+> ```
+
 ### 📊 Dataset Format
 
-The trainer supports either videos or single images.
+The trainer supports videos, single images, or a mix of both in the same dataset.
-Note that your dataset must be homogeneous - either all videos or all images, mixing is not supported.
 
 > [!TIP]
 > **Image Datasets:** When using images, follow the same preprocessing steps and format requirements as with videos,
 > but use `1` for the frame count in the resolution bucket (e.g., `960x544x1`).
 
+> [!NOTE]
+> **Mixed image + video datasets:** Mixing stills and videos in a single dataset is supported, but requires some care:
+>
+> - Preprocess with **multiple resolution buckets** covering both frame counts — e.g.
+> `--resolution-buckets "960x544x1;960x544x49"`. Images are automatically assigned to the `F=1` bucket and
+> videos to an `F>1` bucket.
+> - You **must** set `optimization.batch_size: 1` in your training config (see the warning under
+> [Resolution Buckets](#-resolution-buckets)), since samples with different shapes cannot be collated into a
+> single batch. Use `gradient_accumulation_steps` if you need a larger effective batch.
+> - Per-step cost differs substantially between a single-frame sample and a many-frame sample, which can lead to
+> uneven gradient magnitudes across steps. Consider weighting the two subsets or tuning the learning rate if
+> you observe instability.
+> - If you prefer a fully officially-supported path, train two separate LoRAs (one on stills, one on video) and
+> stack them at inference.
+
 The dataset must be a CSV, JSON, or JSONL metadata file with columns for captions and video paths:
 
 **JSON format example:**
@@ -197,6 +242,7 @@ uv run python scripts/process_dataset.py dataset.json \
 > ```
 >
 > Where:
+>
 > - H = Height of video
 > - W = Width of video
 > - F = Number of frames
@@ -204,6 +250,7 @@ uv run python scripts/process_dataset.py dataset.json \
 > - 8 = VAE's temporal downsampling factor
 >
 > For example, a 768×448×89 video would have sequence length:
+>
 > ```
 > (768/32) * (448/32) * ((89-1)/8 + 1) = 24 * 14 * 12 = 4,032
 > ```
@@ -268,7 +315,6 @@ uv run python scripts/process_dataset.py dataset.json \
 
 This will create an additional `reference_latents/` directory containing the preprocessed reference video latents.
 
-
 ### Generating Reference Videos
 
 **Dataset Requirements for IC-LoRA:**
@@ -277,7 +323,7 @@ This will create an additional `reference_latents/` directory containing the pre
 - Reference and target videos must have *identical* resolution and length
 - Both reference and target videos should be preprocessed together using the same resolution buckets
 
-We provide an example script, [`scripts/compute_reference.py`](../scripts/compute_reference.py), to generate reference
+We provide an example script, `[scripts/compute_reference.py](../scripts/compute_reference.py)`, to generate reference
 videos for a given dataset. The default implementation generates Canny edge reference videos.
 
 ```bash
@@ -293,7 +339,6 @@ If you want to generate a different type of condition (depth maps, pose skeleton
 
 For reference, see our **[Canny Control Dataset](https://huggingface.co/datasets/Lightricks/Canny-Control-Dataset)** which demonstrates proper IC-LoRA dataset structure with paired videos and Canny edge maps.
 
-
 ## 🎯 LoRA Trigger Words
 
 When training a LoRA, you can specify a trigger token that will be prepended to all captions:

packages/ltx-trainer/docs/troubleshooting.md

@@ -84,6 +84,20 @@ optimization:
   optimizer_type: "adamw8bit"
 ```
 
+#### 7. Offload Optimizer State During Validation
+
+If you OOM specifically during validation video sampling — typically in
+full fine-tunes or high-rank LoRA runs where AdamW state and the VAE decoder
+can't coexist on the GPU — offload optimizer state to CPU during sampling:
+
+```yaml
+acceleration:
+  offload_optimizer_during_validation: true
+```
+
+The offload + reload happens once per validation interval, not per step.
+No effect for FSDP (sharded state).
+
 ---
 
 ## ⚠️ Common Usage Issues

packages/ltx-trainer/docs/utility-scripts.md

@@ -143,6 +143,27 @@ uv run python scripts/process_dataset.py dataset.json \
 > [!NOTE]
 > When training with multiple resolution buckets, set `optimization.batch_size: 1`.
 
+**Multi-GPU preprocessing.** Launch with `accelerate launch` to shard the dataset across processes. Reruns resume
+by default (existing `.pt` outputs are skipped); writes are atomic so interrupted runs are safe. Pass `--overwrite`
+when rerunning with changed parameters (different model, resolution buckets, text encoder, `--lora-trigger`, etc.)
+so stale outputs are replaced. Use the same `accelerate launch` pattern (and `--overwrite` when needed) with
+`process_videos.py` or `process_captions.py` when you run those scripts standalone.
+
+```bash
+# Multi-GPU preprocessing
+uv run accelerate launch --num_processes 4 scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x49" \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --text-encoder-path /path/to/gemma-model
+
+# Force re-encoding of all items (e.g. after switching model or resolution)
+uv run accelerate launch --num_processes 4 scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x49" \
+    --model-path /path/to/ltx-2.3-model.safetensors \
+    --text-encoder-path /path/to/gemma-model \
+    --overwrite
+```
+
 For detailed usage, see the [Dataset Preparation Guide](dataset-preparation.md).
 
 ### Reference Video Generation

packages/ltx-trainer/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ltx-trainer"
-version = "1.1.1"
+version = "1.1.3"
 description = "LTX-2 training, democratized."
 readme = "README.md"
 authors = [
@@ -48,7 +48,7 @@ build-backend = "hatchling.build"
 
 
 [tool.ruff]
-target-version = "1.1.1"
+target-version = "1.1.3"
 line-length = 120
 
 [tool.ruff.lint]

packages/ltx-trainer/scripts/compute_reference.py

@@ -241,8 +241,8 @@ def main(
         help="Path to input video/image file or directory containing media files",
         exists=True,
     ),
-    output: Path | None = typer.Option(  # noqa: B008
+    output: Path = typer.Option(  # noqa: B008
-        None,
+        ...,
         "--output",
         "-o",
         help="Path to json output file for reference video paths. "

packages/ltx-trainer/scripts/process_captions.py

@@ -13,12 +13,14 @@ Can be used as a standalone script:
 
 import json
 import os
+from collections.abc import Callable
 from pathlib import Path
 from typing import Any
 
 import pandas as pd
 import torch
 import typer
+from accelerate import PartialState
 from rich.console import Console
 from rich.progress import (
     BarColumn,
@@ -30,7 +32,7 @@ from rich.progress import (
     TimeElapsedColumn,
     TimeRemainingColumn,
 )
-from torch.utils.data import DataLoader, Dataset
+from torch.utils.data import DataLoader, Dataset, Subset
 from transformers.utils.logging import disable_progress_bar
 
 from ltx_trainer import logger
@@ -232,9 +234,14 @@ def compute_captions_embeddings(  # noqa: PLR0913
     batch_size: int = 8,
     device: str = "cuda",
     load_in_8bit: bool = False,
+    overwrite: bool = False,
 ) -> None:
     """
     Process captions and save text embeddings.
+    Under ``accelerate launch``, each process handles an interleaved shard of
+    the dataset (rank/world read from ``accelerate.PartialState``). Already-
+    computed ``.pt`` outputs are skipped unless ``overwrite=True``; writes are
+    atomic so an interrupted run is safe to resume.
     Args:
         dataset_file: Path to metadata file (CSV/JSON/JSONL) containing captions and media paths
         output_dir: Directory to save embeddings
@@ -247,11 +254,12 @@ def compute_captions_embeddings(  # noqa: PLR0913
         batch_size: Batch size for processing
         device: Device to use for computation
         load_in_8bit: Whether to load the Gemma text encoder in 8-bit precision
+        overwrite: Re-encode every item even if its output exists. Use when rerunning with
+            changed parameters (different text encoder, lora_trigger, etc.) so stale
+            outputs are replaced.
     """
-
     console = Console()
 
-    # Create dataset
     dataset = CaptionsDataset(
         dataset_file=dataset_file,
         caption_column=caption_column,
@@ -264,6 +272,24 @@ def compute_captions_embeddings(  # noqa: PLR0913
     output_path = Path(output_dir)
     output_path.mkdir(parents=True, exist_ok=True)
 
+    # TODO(batch-tokenization): The current Gemma tokenizer doesn't support batched tokenization.
+    if batch_size > 1:
+        logger.warning(
+            "Batch size greater than 1 is not currently supported with the Gemma tokenizer. "
+            "Overriding batch_size to 1. This will be fixed in a future update."
+        )
+        batch_size = 1
+
+    dataloader = _build_sharded_dataloader(
+        dataset,
+        batch_size=batch_size,
+        num_workers=2,
+        is_done=lambda idx: (output_path / dataset.output_paths[idx]).is_file(),
+        overwrite=overwrite,
+    )
+    if dataloader is None:
+        return
+
     # Load text encoder and embeddings processor
     with console.status("[bold]Loading Gemma text encoder...", spinner="dots"):
         text_encoder = load_text_encoder(
@@ -279,21 +305,7 @@ def compute_captions_embeddings(  # noqa: PLR0913
         )
 
     logger.info("Text encoder and embeddings processor loaded successfully")
-
+    logger.info(f"Processing captions in {len(dataloader):,} batches...")
-    # TODO(batch-tokenization): The current Gemma tokenizer doesn't support batched tokenization.
-    if batch_size > 1:
-        logger.warning(
-            "Batch size greater than 1 is not currently supported with the Gemma tokenizer. "
-            "Overriding batch_size to 1. This will be fixed in a future update."
-        )
-        batch_size = 1
-
-    # Create dataloader
-    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=2)
-
-    # Process batches
-    total_batches = len(dataloader)
-    logger.info(f"Processing captions in {total_batches:,} batches...")
 
     with Progress(
         SpinnerColumn(),
@@ -333,11 +345,44 @@ def compute_captions_embeddings(  # noqa: PLR0913
                         embedding_data["audio_prompt_embeds"] = audio_prompt_embeds[0].cpu().contiguous()
 
                     output_file = output_path / output_rel_path
-                    torch.save(embedding_data, output_file)
+                    _atomic_save(embedding_data, output_file)
 
             progress.advance(task)
 
-    logger.info(f"Processed {len(dataset):,} captions. Embeddings saved to {output_path}")
+    logger.info(f"Processed {len(dataloader.dataset):,} captions -> {output_path}")  # type: ignore[arg-type]
+
+
+def _atomic_save(data: Any, out: Path) -> None:  # noqa: ANN401
+    """Save to ``out`` atomically via per-PID temp file + replace.
+    Crash mid-write leaves an orphan ``.tmp.<pid>`` file that the skip logic
+    ignores. The per-PID suffix makes concurrent writes from multiple ranks
+    collision-free.
+    """
+    tmp = out.with_suffix(f"{out.suffix}.tmp.{os.getpid()}")
+    torch.save(data, tmp)
+    tmp.replace(out)
+
+
+def _build_sharded_dataloader(
+    dataset: Dataset,
+    *,
+    batch_size: int,
+    num_workers: int,
+    is_done: Callable[[int], bool],
+    overwrite: bool,
+) -> DataLoader | None:
+    """Return a DataLoader over this rank's interleaved shard of ``dataset``.
+    When ``overwrite`` is False, items whose outputs already exist (per
+    ``is_done``) are filtered out. Returns ``None`` if this rank has nothing
+    to do, so the caller can early-return without loading any models.
+    """
+    state = PartialState()
+    todo = [i for i in range(state.process_index, len(dataset), state.num_processes) if overwrite or not is_done(i)]
+    if not todo:
+        logger.info(f"Rank {state.process_index}/{state.num_processes}: nothing to do")
+        return None
+    logger.info(f"Rank {state.process_index}/{state.num_processes}: processing {len(todo):,} of {len(dataset):,} items")
+    return DataLoader(Subset(dataset, todo), batch_size=batch_size, shuffle=False, num_workers=num_workers)
 
 
 @app.command()
@@ -387,8 +432,15 @@ def main(  # noqa: PLR0913
         default=False,
         help="Load the Gemma text encoder in 8-bit precision to save GPU memory (requires bitsandbytes)",
     ),
+    overwrite: bool = typer.Option(
+        default=False,
+        help="Re-encode every caption even if its output exists. Use when rerunning with "
+        "changed parameters (different text encoder, lora_trigger, etc.) so stale outputs are replaced.",
+    ),
 ) -> None:
     """Process text captions and save embeddings for video generation training.
+    For multi-GPU preprocessing, invoke under ``accelerate launch`` - each process
+    will handle an interleaved shard of the dataset.
     This script processes captions from metadata files and saves text embeddings
     that can be used for training video generation models. The output embeddings
     will maintain the same folder structure and naming as the corresponding media files.
@@ -428,6 +480,7 @@ def main(  # noqa: PLR0913
         batch_size=batch_size,
         device=device,
         load_in_8bit=load_text_encoder_in_8bit,
+        overwrite=overwrite,
     )
 
 

packages/ltx-trainer/scripts/process_dataset.py

@@ -50,6 +50,7 @@ def preprocess_dataset(  # noqa: PLR0913
     reference_downscale_factor: int = 1,
     with_audio: bool = False,
     load_text_encoder_in_8bit: bool = False,
+    overwrite: bool = False,
 ) -> None:
     """Run the preprocessing pipeline with the given arguments."""
     # Validate dataset file
@@ -77,6 +78,7 @@ def preprocess_dataset(  # noqa: PLR0913
             batch_size=batch_size,
             device=device,
             load_in_8bit=load_text_encoder_in_8bit,
+            overwrite=overwrite,
         )
 
     # Process videos using the dedicated function
@@ -97,6 +99,7 @@ def preprocess_dataset(  # noqa: PLR0913
             vae_tiling=vae_tiling,
             with_audio=with_audio,
             audio_output_dir=str(audio_latents_dir) if audio_latents_dir else None,
+            overwrite=overwrite,
         )
 
         # Process reference videos if reference_column is provided
@@ -133,6 +136,7 @@ def preprocess_dataset(  # noqa: PLR0913
                 batch_size=batch_size,
                 device=device,
                 vae_tiling=vae_tiling,
+                overwrite=overwrite,
             )
 
     # Handle decoding if requested (for verification)
@@ -252,8 +256,15 @@ def main(  # noqa: PLR0913
         help="Downscale factor for reference video resolution. When > 1, reference videos are processed at "
         "1/n resolution (e.g., 2 means half resolution). Used for efficient IC-LoRA training.",
     ),
+    overwrite: bool = typer.Option(
+        default=False,
+        help="Re-compute every item even if its output exists. Use when rerunning with "
+        "changed parameters (different model, resolution, etc.) so stale outputs are replaced.",
+    ),
 ) -> None:
     """Preprocess a video dataset by computing and saving latents and text embeddings.
+    For multi-GPU preprocessing, invoke under ``accelerate launch`` - each process
+    will handle an interleaved shard of the dataset.
     The dataset must be a CSV, JSON, or JSONL file with columns for captions and video paths.
     This script is designed for LTX-2 models which use the Gemma text encoder.
     Examples:
@@ -310,6 +321,7 @@ def main(  # noqa: PLR0913
         reference_downscale_factor=reference_downscale_factor,
         with_audio=with_audio,
         load_text_encoder_in_8bit=load_text_encoder_in_8bit,
+        overwrite=overwrite,
     )
 
 

packages/ltx-trainer/scripts/process_videos.py

@@ -14,6 +14,8 @@ Can be used as a standalone script:
 
 import json
 import math
+import os
+from collections.abc import Callable
 from pathlib import Path
 from typing import Any
 
@@ -22,6 +24,7 @@ import pandas as pd
 import torch
 import torchaudio
 import typer
+from accelerate import PartialState
 from pillow_heif import register_heif_opener
 from rich.console import Console
 from rich.progress import (
@@ -34,7 +37,7 @@ from rich.progress import (
     TimeElapsedColumn,
     TimeRemainingColumn,
 )
-from torch.utils.data import DataLoader, Dataset
+from torch.utils.data import DataLoader, Dataset, Subset
 from torchvision import transforms
 from torchvision.transforms import InterpolationMode
 from torchvision.transforms.functional import crop, resize, to_tensor
@@ -444,9 +447,14 @@ def compute_latents(  # noqa: PLR0913, PLR0915
     vae_tiling: bool = False,
     with_audio: bool = False,
     audio_output_dir: str | None = None,
+    overwrite: bool = False,
 ) -> None:
     """
     Process videos and save latent representations.
+    Under ``accelerate launch``, each process handles an interleaved shard of
+    the dataset (rank/world read from ``accelerate.PartialState``). Already-
+    computed ``.pt`` outputs are skipped unless ``overwrite=True``; writes are
+    atomic so an interrupted run is safe to resume.
     Args:
         dataset_file: Path to metadata file (CSV/JSON/JSONL) containing video paths
         video_column: Column name for video paths in the metadata file
@@ -460,15 +468,15 @@ def compute_latents(  # noqa: PLR0913, PLR0915
         vae_tiling: Whether to enable VAE tiling
         with_audio: Whether to extract and encode audio from videos
         audio_output_dir: Directory to save audio latents (required if with_audio=True)
+        overwrite: Re-process every item even if its output exists. Use when rerunning with
+            changed parameters (different model, resolution, etc.) so stale outputs are replaced.
     """
-    # Validate audio parameters
     if with_audio and audio_output_dir is None:
         raise ValueError("audio_output_dir must be provided when with_audio=True")
 
     console = Console()
     torch_device = torch.device(device)
 
-    # Create dataset
     dataset = MediaDataset(
         dataset_file=dataset_file,
         main_media_column=main_media_column or video_column,
@@ -481,13 +489,34 @@ def compute_latents(  # noqa: PLR0913, PLR0915
 
     output_path = Path(output_dir)
     output_path.mkdir(parents=True, exist_ok=True)
-
+    audio_output_path: Path | None = None
-    # Set up audio output directory if needed
-    audio_output_path = None
     if with_audio:
         audio_output_path = Path(audio_output_dir)
         audio_output_path.mkdir(parents=True, exist_ok=True)
 
+    # Audio processing requires batch_size=1; must be applied before the dataloader is built.
+    if with_audio and batch_size > 1:
+        logger.warning("Audio processing requires batch_size=1. Overriding batch_size to 1.")
+        batch_size = 1
+
+    data_root = dataset.dataset_file.parent
+
+    def _is_done(idx: int) -> bool:
+        rel = dataset.main_media_paths[idx].relative_to(data_root).with_suffix(".pt")
+        if not (output_path / rel).is_file():
+            return False
+        return audio_output_path is None or (audio_output_path / rel).is_file()
+
+    dataloader = _build_sharded_dataloader(
+        dataset,
+        batch_size=batch_size,
+        num_workers=4,
+        is_done=_is_done,
+        overwrite=overwrite,
+    )
+    if dataloader is None:
+        return
+
     # Load video VAE encoder
     with console.status(f"[bold]Loading video VAE encoder from [cyan]{model_path}[/]...", spinner="dots"):
         vae = load_video_vae_encoder(model_path, device=torch_device, dtype=torch.bfloat16)
@@ -510,14 +539,6 @@ def compute_latents(  # noqa: PLR0913, PLR0915
                 n_fft=audio_vae_encoder.n_fft,
             ).to(torch_device)
 
-    # Create dataloader
-    # Note: batch_size=1 required when with_audio because audio extraction can fail for some videos,
-    # and the default collate function can't handle mixed None/dict values across a batch.
-    if with_audio and batch_size > 1:
-        logger.warning("Audio processing requires batch_size=1. Overriding batch_size to 1.")
-        batch_size = 1
-    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4)
-
     # Track audio statistics
     audio_success_count = 0
     audio_skip_count = 0
@@ -560,7 +581,7 @@ def compute_latents(  # noqa: PLR0913, PLR0915
                     "fps": batch["video_metadata"]["fps"][i].item(),
                 }
 
-                torch.save(latent_data, output_file)
+                _atomic_save(latent_data, output_file)
 
                 # Process audio if enabled (audio is already extracted by the dataset)
                 if with_audio:
@@ -588,7 +609,7 @@ def compute_latents(  # noqa: PLR0913, PLR0915
                             "duration": audio_latents["duration"],
                         }
 
-                        torch.save(audio_save_data, audio_output_file)
+                        _atomic_save(audio_save_data, audio_output_file)
                         audio_success_count += 1
                     else:
                         # Video has no audio track
@@ -596,8 +617,7 @@ def compute_latents(  # noqa: PLR0913, PLR0915
 
             progress.advance(task)
 
-    # Log summary
+    logger.info(f"Processed {len(dataloader.dataset)} videos -> {output_path}")  # type: ignore[arg-type]
-    logger.info(f"Processed {len(dataset)} videos. Latents saved to {output_path}")
     if with_audio:
         logger.info(
             f"Audio processing: {audio_success_count} videos with audio, "
@@ -935,6 +955,39 @@ def compute_scaled_resolution_buckets(
     return scaled_buckets
 
 
+def _atomic_save(data: Any, out: Path) -> None:  # noqa: ANN401
+    """Save to ``out`` atomically via per-PID temp file + replace.
+    Crash mid-write leaves an orphan ``.tmp.<pid>`` file that the skip logic
+    ignores. The per-PID suffix makes concurrent writes from multiple ranks
+    collision-free.
+    """
+    tmp = out.with_suffix(f"{out.suffix}.tmp.{os.getpid()}")
+    torch.save(data, tmp)
+    tmp.replace(out)
+
+
+def _build_sharded_dataloader(
+    dataset: Dataset,
+    *,
+    batch_size: int,
+    num_workers: int,
+    is_done: Callable[[int], bool],
+    overwrite: bool,
+) -> DataLoader | None:
+    """Return a DataLoader over this rank's interleaved shard of ``dataset``.
+    When ``overwrite`` is False, items whose outputs already exist (per
+    ``is_done``) are filtered out. Returns ``None`` if this rank has nothing
+    to do, so the caller can early-return without loading any models.
+    """
+    state = PartialState()
+    todo = [i for i in range(state.process_index, len(dataset), state.num_processes) if overwrite or not is_done(i)]
+    if not todo:
+        logger.info(f"Rank {state.process_index}/{state.num_processes}: nothing to do")
+        return None
+    logger.info(f"Rank {state.process_index}/{state.num_processes}: processing {len(todo):,} of {len(dataset):,} items")
+    return DataLoader(Subset(dataset, todo), batch_size=batch_size, shuffle=False, num_workers=num_workers)
+
+
 @app.command()
 def main(  # noqa: PLR0913
     dataset_file: str = typer.Argument(
@@ -981,8 +1034,15 @@ def main(  # noqa: PLR0913
         default=None,
         help="Output directory for audio latents (required if --with-audio is set)",
     ),
+    overwrite: bool = typer.Option(
+        default=False,
+        help="Re-encode every item even if its output exists. Use when rerunning with "
+        "changed parameters (different model, resolution, etc.) so stale outputs are replaced.",
+    ),
 ) -> None:
     """Process videos/images and save latent representations for video generation training.
+    For multi-GPU preprocessing, invoke under ``accelerate launch`` - each process
+    will handle an interleaved shard of the dataset.
     This script processes videos and images from metadata files and saves latent representations
     that can be used for training video generation models. The output latents will maintain
     the same folder structure and naming as the corresponding media files.
@@ -1032,6 +1092,7 @@ def main(  # noqa: PLR0913
         vae_tiling=vae_tiling,
         with_audio=with_audio,
         audio_output_dir=audio_output_dir,
+        overwrite=overwrite,
     )
 
 

packages/ltx-trainer/src/ltx_trainer/config.py

@@ -168,6 +168,15 @@ class AccelerationConfig(ConfigBaseModel):
         description="Whether to load the text encoder in 8-bit precision to save memory",
     )
 
+    offload_optimizer_during_validation: bool = Field(
+        default=False,
+        description="Offload optimizer state to CPU before validation video sampling and reload "
+        "it afterwards, to free VRAM for inference. Useful when optimizer state is large "
+        "(e.g. AdamW for full fine-tuning or high-rank LoRA) and validation OOMs because the "
+        "VAE decoder + transformer + optimizer state cannot coexist on the GPU. Has no effect "
+        "for FSDP (sharded state). Disabled by default.",
+    )
+
 
 class DataConfig(ConfigBaseModel):
     """Configuration for data loading and processing"""
@@ -260,12 +269,6 @@ class ValidationConfig(ConfigBaseModel):
         gt=0,
     )
 
-    videos_per_prompt: int = Field(
-        default=1,
-        description="Number of videos to generate per validation prompt",
-        gt=0,
-    )
-
     guidance_scale: float = Field(
         default=4.0,
         description="CFG guidance scale to use during validation",

packages/ltx-trainer/src/ltx_trainer/config_display.py

@@ -85,6 +85,7 @@ def print_config(config: LtxTrainerConfig) -> None:
                     ("Mixed Precision", accel.mixed_precision_mode or "[dim]—[/]"),
                     ("Quantization", str(accel.quantization) if accel.quantization else "[dim]—[/]"),
                     ("Text Encoder 8bit", fmt(accel.load_text_encoder_in_8bit)),
+                    ("Optimizer CPU Offload", fmt(accel.offload_optimizer_during_validation)),
                 ],
             ),
             (

packages/ltx-trainer/src/ltx_trainer/datasets.py

@@ -1,3 +1,4 @@
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 
 import torch
@@ -155,40 +156,77 @@ class PrecomputedDataset(Dataset):
         return source_paths
 
     def _discover_samples(self) -> dict[str, list[Path]]:
-        """Discover all valid sample files across all data sources."""
+        """Discover all valid sample files across all data sources.
-        # Use first data source as the reference to discover samples
+        Uses a fast two-pass approach: first globs all sources in parallel to build
+        full-path sets in memory, then checks expected paths via set membership.
+        This avoids O(N * num_sources) stat calls on networked filesystems while
+        correctly handling path remapping (e.g. latent_X.pt -> condition_X.pt).
+        """
+        if not self.data_sources:
+            raise ValueError("No data sources configured")
+
         data_key = "latents" if "latents" in self.data_sources else next(iter(self.data_sources.keys()))
         data_path = self.source_paths[data_key]
-        data_files = list(data_path.glob("**/*.pt"))
 
+        # Pass 1: Glob all sources in parallel, build full-path sets
+        def _glob_source(dir_name: str) -> tuple[list[Path], set[str]]:
+            source_path = self.source_paths[dir_name]
+            paths = list(source_path.glob("**/*.pt"))
+            path_set = {str(p) for p in paths}
+            return paths, path_set
+
+        with ThreadPoolExecutor(max_workers=len(self.data_sources)) as executor:
+            glob_results = dict(
+                zip(
+                    self.data_sources.keys(),
+                    executor.map(_glob_source, self.data_sources.keys()),
+                    strict=True,
+                )
+            )
+
+        # Get primary source files (cached from glob, no second scan)
+        data_files, _ = glob_results[data_key]
         if not data_files:
             raise ValueError(f"No data files found in {data_path}")
+        data_files.sort()
 
-        # Initialize sample files dict
+        # Log source sizes
-        sample_files = {output_key: [] for output_key in self.data_sources.values()}
+        for dir_name, (paths, _) in glob_results.items():
+            logger.debug(f"Source {dir_name}: {len(paths)} files")
+
+        # Build path sets for non-primary sources
+        other_path_sets = {
+            dir_name: path_set for dir_name, (_, path_set) in glob_results.items() if dir_name != data_key
+        }
+
+        # Pass 2: For each primary file, check if expected paths exist in other sources' sets
+        sample_files: dict[str, list[Path]] = {output_key: [] for output_key in self.data_sources.values()}
+        valid_count = 0
 
-        # For each data file, find corresponding files in other sources
         for data_file in data_files:
             rel_path = data_file.relative_to(data_path)
 
-            # Check if corresponding files exist in ALL sources
+            # Check all other sources via set lookup (O(1) per source, no stat calls)
-            if self._all_source_files_exist(data_file, rel_path):
+            all_exist = True
+            for dir_name, path_set in other_path_sets.items():
+                expected = self._get_expected_file_path(dir_name, data_file, rel_path)
+                if str(expected) not in path_set:
+                    logger.debug(f"Skipping {data_file.name}: no matching {dir_name} file at {expected}")
+                    all_exist = False
+                    break
+
+            if all_exist:
                 self._fill_sample_data_files(data_file, rel_path, sample_files)
+                valid_count += 1
+
+        skipped = len(data_files) - valid_count
+        if skipped > 0:
+            logger.info(f"Fast index: {valid_count} valid samples from {len(data_files)} total ({skipped} skipped)")
+        else:
+            logger.debug(f"Fast index: {valid_count} valid samples from {len(data_files)} total")
 
         return sample_files
 
-    def _all_source_files_exist(self, data_file: Path, rel_path: Path) -> bool:
-        """Check if corresponding files exist in all data sources."""
-        for dir_name in self.data_sources:
-            expected_path = self._get_expected_file_path(dir_name, data_file, rel_path)
-            if not expected_path.exists():
-                logger.warning(
-                    f"No matching {dir_name} file found for: {data_file.name} (expected in: {expected_path})"
-                )
-                return False
-
-        return True
-
     def _get_expected_file_path(self, dir_name: str, data_file: Path, rel_path: Path) -> Path:
         """Get the expected file path for a given data source."""
         source_path = self.source_paths[dir_name]
@@ -207,11 +245,14 @@ class PrecomputedDataset(Dataset):
 
     def _validate_setup(self) -> None:
         """Validate that the dataset setup is correct."""
-        if not self.sample_files:
+        sample_counts = {key: len(files) for key, files in self.sample_files.items()}
-            raise ValueError("No valid samples found - all data sources must have matching files")
+        if not sample_counts or all(count == 0 for count in sample_counts.values()):
+            raise ValueError(
+                f"No valid samples found in {self.data_root} - all configured data sources "
+                f"({list(self.data_sources)}) must have matching files (per-source counts: {sample_counts})"
+            )
 
         # Verify all output keys have the same number of samples
-        sample_counts = {key: len(files) for key, files in self.sample_files.items()}
         if len(set(sample_counts.values())) > 1:
             raise ValueError(f"Mismatched sample counts across sources: {sample_counts}")
 

packages/ltx-trainer/src/ltx_trainer/gemma_8bit.py

@@ -12,6 +12,7 @@ Example usage:
 from __future__ import annotations
 
 import logging
+import os
 from collections.abc import Generator
 from contextlib import contextmanager
 from pathlib import Path
@@ -22,7 +23,11 @@ from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder
 from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
 
 
-def load_8bit_gemma(gemma_model_path: str | Path, dtype: torch.dtype = torch.bfloat16) -> GemmaTextEncoder:
+def load_8bit_gemma(
+    gemma_model_path: str | Path,
+    dtype: torch.dtype = torch.bfloat16,
+    device: torch.device | str | int | None = None,
+) -> GemmaTextEncoder:
     """Load the Gemma text encoder in 8-bit precision using bitsandbytes.
     Only the Gemma LLM backbone is loaded here.  The embeddings processor
     (feature extractor + connectors) should be loaded separately via
@@ -30,6 +35,10 @@ def load_8bit_gemma(gemma_model_path: str | Path, dtype: torch.dtype = torch.bfl
     Args:
         gemma_model_path: Path to Gemma model directory
         dtype: Data type for non-quantized model weights
+        device: Device to place the quantized model on. When ``None`` (default),
+            the device is inferred from ``LOCAL_RANK`` if CUDA is available, so
+            multi-process launches put each rank's encoder on its own GPU
+            instead of all colliding on ``cuda:0``.
     Returns:
         GemmaTextEncoder with 8-bit quantized Gemma backbone
     Raises:
@@ -46,13 +55,23 @@ def load_8bit_gemma(gemma_model_path: str | Path, dtype: torch.dtype = torch.bfl
     gemma_path = _find_gemma_subpath(gemma_model_path, "model*.safetensors")
     tokenizer_path = _find_gemma_subpath(gemma_model_path, "tokenizer.model")
 
+    # Pin the entire model to a single device. `device_map="auto"` collides on cuda:0
+    # in multi-process launches because every rank picks the same default device.
+    device_map: str | dict[str, int | str | torch.device]
+    if device is not None:
+        device_map = {"": device}
+    elif torch.cuda.is_available():
+        device_map = {"": int(os.environ.get("LOCAL_RANK", "0"))}
+    else:
+        device_map = "auto"
+
     quantization_config = BitsAndBytesConfig(load_in_8bit=True)
     with _suppress_accelerate_memory_warnings():
         gemma_model = Gemma3ForConditionalGeneration.from_pretrained(
             gemma_path,
             quantization_config=quantization_config,
             torch_dtype=torch.bfloat16,
-            device_map="auto",
+            device_map=device_map,
             local_files_only=True,
         )
 

packages/ltx-trainer/src/ltx_trainer/hf_hub_utils.py

@@ -1,7 +1,7 @@
 import shutil
 import tempfile
 from pathlib import Path
-from typing import List, Union
+from typing import List, Optional, Union
 
 import imageio
 from huggingface_hub import HfApi, create_repo
@@ -12,7 +12,11 @@ from ltx_trainer import logger
 from ltx_trainer.config import LtxTrainerConfig
 
 
-def push_to_hub(weights_path: Path, sampled_videos_paths: List[Path], config: LtxTrainerConfig) -> None:
+def push_to_hub(
+    weights_path: Path,
+    sampled_videos_paths: Optional[List[Path]],
+    config: LtxTrainerConfig,
+) -> None:
     """Push the trained LoRA weights to HuggingFace Hub."""
     if not config.hub.hub_model_id:
         logger.warning("⚠️ HuggingFace hub_model_id not specified, skipping push to hub")
@@ -108,7 +112,7 @@ def convert_video_to_gif(video_path: Path, output_path: Path) -> None:
 
 def _create_model_card(
     output_dir: Union[str, Path],
-    videos: List[Path],
+    videos: Optional[List[Path]],
     config: LtxTrainerConfig,
 ) -> Path:
     """Generate and save a model card for the trained model."""

packages/ltx-trainer/src/ltx_trainer/model_loader.py

@@ -199,8 +199,6 @@ def load_text_encoder(
         device: Device to load model on
         dtype: Data type for model weights
         load_in_8bit: Whether to load the Gemma model in 8-bit precision using bitsandbytes.
-            When True, the model is loaded with device_map="auto" and the device argument
-            is ignored for the Gemma backbone.
     Returns:
         Loaded GemmaTextEncoder
     """
@@ -211,7 +209,7 @@ def load_text_encoder(
     if load_in_8bit:
         from ltx_trainer.gemma_8bit import load_8bit_gemma
 
-        return load_8bit_gemma(gemma_model_path, dtype)
+        return load_8bit_gemma(gemma_model_path, dtype, device=device)
 
     # Standard loading path
     from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder

packages/ltx-trainer/src/ltx_trainer/trainer.py

@@ -1,7 +1,10 @@
+import contextlib
+import math
 import os
 import re
 import time
 import warnings
+from collections.abc import Iterator
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Callable
@@ -9,8 +12,8 @@ from typing import Any, Callable
 import torch
 import wandb
 import yaml
-from accelerate import Accelerator, DistributedType
+from accelerate import Accelerator, DistributedDataParallelKwargs, DistributedType
-from accelerate.utils import set_seed
+from accelerate.utils import gather_object, set_seed
 from peft import LoraConfig, get_peft_model, get_peft_model_state_dict, set_peft_model_state_dict
 from peft.tuners.tuners_utils import BaseTunerLayer
 from peft.utils import ModulesToSaveWrapper
@@ -63,7 +66,7 @@ if not IS_MAIN_PROCESS:
 
     disable_progress_bar()
 
-StepCallback = Callable[[int, int, list[Path]], None]  # (step, total, list[sampled_video_path]) -> None
+StepCallback = Callable[[int, int, list[Path] | None], None]  # (step, total, sampled paths or None) -> None
 
 MEMORY_CHECK_INTERVAL = 200
 
@@ -186,9 +189,8 @@ class LtxvTrainer:
 
         with progress:
             if cfg.validation.interval and not cfg.validation.skip_initial_validation:
-                sampled_videos_paths = self._sample_videos(progress)
+                with self._offloaded_optimizer_state():
-                if IS_MAIN_PROCESS and sampled_videos_paths and self._config.wandb.log_validation_videos:
+                    sampled_videos_paths = self._run_distributed_validation(progress)
-                    self._log_validation_samples(sampled_videos_paths, cfg.validation.prompts)
 
             self._accelerator.wait_for_everyone()
 
@@ -228,16 +230,8 @@ class LtxvTrainer:
                         and self._global_step % cfg.validation.interval == 0
                         and is_optimization_step
                     ):
-                        if self._accelerator.distributed_type == DistributedType.FSDP:
+                        with self._offloaded_optimizer_state():
-                            # FSDP: All processes must participate in validation
+                            sampled_videos_paths = self._run_distributed_validation(progress)
-                            sampled_videos_paths = self._sample_videos(progress)
-                            if IS_MAIN_PROCESS and sampled_videos_paths and self._config.wandb.log_validation_videos:
-                                self._log_validation_samples(sampled_videos_paths, cfg.validation.prompts)
-                        # DDP: Only main process runs validation
-                        elif IS_MAIN_PROCESS:
-                            sampled_videos_paths = self._sample_videos(progress)
-                            if sampled_videos_paths and self._config.wandb.log_validation_videos:
-                                self._log_validation_samples(sampled_videos_paths, cfg.validation.prompts)
 
                     # Save checkpoint if needed
                     if (
@@ -398,11 +392,14 @@ class LtxvTrainer:
         #   3. If validation prompts are configured, computes and caches their embeddings
         #   4. Unloads the Gemma model entirely, keeps the embeddings processor for training
 
-        # Load text encoder (pure Gemma LLM) on GPU
+        # Load text encoder (pure Gemma LLM) on GPU — LOCAL_RANK before Accelerator exists
+        local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+        init_device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
+
         logger.debug("Loading text encoder...")
         text_encoder = load_text_encoder(
             gemma_model_path=self._config.model.text_encoder_path,
-            device="cuda",
+            device=init_device,
             dtype=torch.bfloat16,
             load_in_8bit=self._config.acceleration.load_text_encoder_in_8bit,
         )
@@ -411,7 +408,7 @@ class LtxvTrainer:
         logger.debug("Loading embeddings processor...")
         self._embeddings_processor = load_embeddings_processor(
             checkpoint_path=self._config.model.model_path,
-            device="cuda",
+            device=init_device,
             dtype=torch.bfloat16,
         )
 
@@ -788,6 +785,41 @@ class LtxvTrainer:
         # noinspection PyTypeChecker
         self._optimizer, self._lr_scheduler = self._accelerator.prepare(optimizer, lr_scheduler)
 
+    @contextlib.contextmanager
+    def _offloaded_optimizer_state(self) -> Iterator[None]:
+        """Context manager that offloads optimizer state to CPU during validation.
+        Opt-in via `acceleration.offload_optimizer_during_validation`. Frees VRAM for
+        validation video generation when optimizer state is large (e.g. full fine-tune
+        AdamW, high-rank LoRA). No-op for FSDP (sharded state -- manual `.cpu()` breaks
+        metadata).
+        """
+        enabled = (
+            self._config.acceleration.offload_optimizer_during_validation
+            and self._accelerator.distributed_type != DistributedType.FSDP
+        )
+
+        # Track exactly which tensors we move so we don't promote ones that were
+        # intentionally on CPU (e.g. AdamW's `step` scalar on recent PyTorch).
+        offloaded: list[tuple[dict, str]] = []
+        if enabled:
+            offloaded_bytes = 0
+            for state in self._optimizer.state.values():
+                for k, v in state.items():
+                    if isinstance(v, torch.Tensor) and v.is_cuda:
+                        offloaded.append((state, k))
+                        offloaded_bytes += v.nbytes
+            if offloaded:
+                logger.info(f"Offloading optimizer state to CPU ({offloaded_bytes / 1e9:.1f} GB)")
+                for state, k in offloaded:
+                    state[k] = state[k].cpu()
+
+        try:
+            yield
+        finally:
+            device = self._accelerator.device
+            for state, k in offloaded:
+                state[k] = state[k].to(device)
+
     def _create_scheduler(self, optimizer: torch.optim.Optimizer) -> LRScheduler | None:
         """Create learning rate scheduler based on config."""
         scheduler_type = self._config.optimization.scheduler_type
@@ -844,11 +876,18 @@ class LtxvTrainer:
     def _setup_accelerator(self) -> None:
         """Initialize the Accelerator with the appropriate settings."""
 
+        # find_unused_parameters=True keeps DDP happy when LoRA targets a branch the forward
+        # pass skips (e.g. audio LoRA with `with_audio: false`, or short module patterns like
+        # "to_k" that match the audio branch unintentionally). It's a no-op for FSDP and
+        # single-GPU runs. The probing cost is paid only on the first step.
+        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+
         # All distributed setup (DDP/FSDP, number of processes, etc.) is controlled by
         # the user's Accelerate configuration (accelerate config / accelerate launch).
         self._accelerator = Accelerator(
             mixed_precision=self._config.acceleration.mixed_precision_mode,
             gradient_accumulation_steps=self._config.optimization.gradient_accumulation_steps,
+            kwargs_handlers=[ddp_kwargs],
         )
 
         if self._accelerator.num_processes > 1:
@@ -881,11 +920,42 @@ class LtxvTrainer:
                 "Monitor training stability and consider disabling quantization if issues arise."
             )
 
+    def _run_distributed_validation(self, progress: TrainingProgress) -> list[Path]:
+        """Run validation across all ranks and log gathered results on rank 0.
+        Each rank generates only its assigned subset of prompts (see `_sample_videos`),
+        so all GPUs stay busy and no rank idles long enough to trigger NCCL timeouts.
+        Paths are gathered across ranks so rank 0 has the full list for W&B logging.
+        Note: Multi-node training requires a shared filesystem so rank 0 can read
+        videos written by other ranks.
+        """
+        sampled = self._sample_videos(progress)
+
+        if self._accelerator.num_processes > 1:
+            # gather_object returns a flat list from all ranks
+            sampled = sorted(gather_object(sampled), key=lambda x: x[0])
+
+        paths = [p for _, p in sampled]
+
+        if self._accelerator.is_main_process and paths:
+            self._log_validation_samples(paths, self._config.validation.prompts)
+
+        # Non-main ranks must not reach checkpoint collectives while main is still logging to W&B.
+        self._accelerator.wait_for_everyone()
+
+        return paths
+
     # Note: Use @torch.no_grad() instead of @torch.inference_mode() to avoid FSDP inplace update errors after validation
     @torch.no_grad()
     @free_gpu_memory_context(after=True)
-    def _sample_videos(self, progress: TrainingProgress) -> list[Path] | None:
+    def _sample_videos(self, progress: TrainingProgress) -> list[tuple[int, Path]]:
-        """Run validation by generating videos from validation prompts."""
+        """Run validation by generating videos from this rank's share of the validation prompts.
+        Prompts are split round-robin across ranks via `process_index` / `num_processes`,
+        which collapses to "all prompts" when running on a single GPU. Returns
+        (prompt_idx, path) tuples so the caller can reconstruct global order without
+        relying on filename conventions.
+        Under FSDP with multiple processes, ranks pad with extra generate passes (same prompt,
+        no disk write) so every rank runs the same number of forwards — avoids collective mismatch.
+        """
         use_images = self._config.validation.images is not None
         use_reference_videos = self._config.validation.reference_videos is not None
         generate_audio = self._config.validation.generate_audio
@@ -895,13 +965,24 @@ class LtxvTrainer:
         self._optimizer.zero_grad(set_to_none=True)
         free_gpu_memory()
 
-        # Start sampling progress tracking
+        prompts = self._config.validation.prompts
+        rank = self._accelerator.process_index
+        world_size = self._accelerator.num_processes
+        rank_indices = list(range(rank, len(prompts), world_size))
+
+        # FSDP: every rank must run the same number of forwards; pad with duplicate generates (no save).
+        work: list[tuple[int, bool]] = [(i, True) for i in rank_indices]
+        if self._accelerator.distributed_type == DistributedType.FSDP and world_size > 1:
+            max_per_rank = math.ceil(len(prompts) / world_size)
+            pad_seed = rank_indices[-1] if rank_indices else 0
+            work += [(pad_seed, False)] * (max_per_rank - len(work))
+
         sampling_ctx = progress.start_sampling(
-            num_prompts=len(self._config.validation.prompts),
+            num_prompts=len(work),
             num_steps=inference_steps,
         )
 
-        # Create validation sampler with loaded models and progress tracking
+        # Create a validation sampler with loaded models and progress tracking
         sampler = ValidationSampler(
             transformer=self._transformer,
             vae_decoder=self._vae_decoder,
@@ -915,12 +996,12 @@ class LtxvTrainer:
         output_dir = Path(self._config.output_dir) / "samples"
         output_dir.mkdir(exist_ok=True, parents=True)
 
-        video_paths = []
+        results: list[tuple[int, Path]] = []
         width, height, num_frames = self._config.validation.video_dims
 
-        for prompt_idx, prompt in enumerate(self._config.validation.prompts):
+        for local_i, (prompt_idx, save_output) in enumerate(work):
-            # Update progress to show current video
+            prompt = prompts[prompt_idx]
-            sampling_ctx.start_video(prompt_idx)
+            sampling_ctx.start_video(local_i)
 
             # Load conditioning image if provided
             condition_image = None
@@ -972,28 +1053,30 @@ class LtxvTrainer:
                 device=self._accelerator.device,
             )
 
+            if not save_output:
+                continue
+
             # Save output (image for single frame, video otherwise)
-            if IS_MAIN_PROCESS:
+            ext = "png" if num_frames == 1 else "mp4"
-                ext = "png" if num_frames == 1 else "mp4"
+            output_path = output_dir / f"step_{self._global_step:06d}_{prompt_idx + 1:02d}.{ext}"
-                output_path = output_dir / f"step_{self._global_step:06d}_{prompt_idx + 1}.{ext}"
+            if num_frames == 1:
-                if num_frames == 1:
+                save_image(video, output_path)
-                    save_image(video, output_path)
+            else:
-                else:
+                save_video(
-                    save_video(
+                    video_tensor=video,
-                        video_tensor=video,
+                    output_path=output_path,
-                        output_path=output_path,
+                    fps=self._config.validation.frame_rate,
-                        fps=self._config.validation.frame_rate,
+                    audio=audio,
-                        audio=audio,
+                    audio_sample_rate=self._vocoder.output_sampling_rate if audio is not None else None,
-                        audio_sample_rate=self._vocoder.output_sampling_rate if audio is not None else None,
+                )
-                    )
+            results.append((prompt_idx, output_path))
-                video_paths.append(output_path)
 
         # Clean up progress tasks
         sampling_ctx.cleanup()
 
         rel_outputs_path = output_dir.relative_to(self._config.output_dir)
         logger.info(f"🎥 Validation samples for step {self._global_step} saved in {rel_outputs_path}")
-        return video_paths
+        return results
 
     @staticmethod
     def _log_training_stats(stats: TrainingStats) -> None:

packages/ltx-trainer/src/ltx_trainer/training_strategies/base_strategy.py

@@ -3,7 +3,6 @@ This module defines the abstract base class that all training strategies must im
 along with the base configuration class.
 """
 
-import random
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import Any, Literal
@@ -251,13 +250,17 @@ class TrainingStrategy(ABC):
             device: Target device
             first_frame_conditioning_p: Probability of conditioning on the first frame
         Returns:
-            Boolean mask where True indicates first frame tokens (if conditioning is enabled)
+            Boolean mask where True indicates first frame tokens (if conditioning is enabled).
+            The conditioning decision is drawn independently per batch element so the training
+            signal across samples in a batch is i.i.d.
         """
         conditioning_mask = torch.zeros(batch_size, sequence_length, dtype=torch.bool, device=device)
 
-        if first_frame_conditioning_p > 0 and random.random() < first_frame_conditioning_p:
+        if first_frame_conditioning_p > 0:
             first_frame_end_idx = height * width
             if first_frame_end_idx < sequence_length:
-                conditioning_mask[:, :first_frame_end_idx] = True
+                # Per-sample Bernoulli draw so each batch element is independently conditioned.
+                per_sample_condition = torch.rand(batch_size, device=device) < first_frame_conditioning_p
+                conditioning_mask[per_sample_condition, :first_frame_end_idx] = True
 
         return conditioning_mask

packages/ltx-trainer/src/ltx_trainer/video_utils.py

@@ -5,12 +5,15 @@ with optional audio support.
 
 from fractions import Fraction
 from pathlib import Path
+from typing import Literal
 
 import av
 import numpy as np
 import torch
 from torch import Tensor
 
+VideoFormat = Literal["CFHW", "FCHW"]
+
 
 def get_video_frame_count(video_path: str | Path) -> int:
     """Get the number of frames in a video file.
@@ -68,6 +71,7 @@ def save_video(
     fps: float = 24.0,
     audio: torch.Tensor | None = None,
     audio_sample_rate: int | None = None,
+    video_format: VideoFormat | None = None,
 ) -> None:
     """Save a video tensor to a file using PyAV, optionally with audio.
     Args:
@@ -76,12 +80,16 @@ def save_video(
         fps: Frames per second for the output video
         audio: Optional audio tensor of shape [C, samples] or [samples, C] in range [-1, 1]
         audio_sample_rate: Sample rate for the audio (required if audio is provided)
+        video_format: Explicit layout of ``video_tensor``, either ``"CFHW"`` or ``"FCHW"``.
+            When ``None`` (default), the layout is auto-detected using a heuristic that only
+            works when ``shape[1] > 3`` — the ambiguous ``[C=3, F=3, H, W]`` / ``[F=3, C=3, H, W]``
+            case requires passing this argument explicitly.
     """
     output_path = Path(output_path)
     output_path.parent.mkdir(parents=True, exist_ok=True)
 
     # Normalize to [F, H, W, C] uint8 numpy array
-    video_np = _prepare_video_array(video_tensor)
+    video_np = _prepare_video_array(video_tensor, video_format=video_format)
     _, height, width, _ = video_np.shape
 
     with av.open(str(output_path), mode="w") as container:
@@ -113,11 +121,21 @@ def save_video(
             _write_audio(container, audio_stream, audio, audio_sample_rate)
 
 
-def _prepare_video_array(video_tensor: torch.Tensor) -> np.ndarray:
+def _prepare_video_array(
-    """Convert video tensor to [F, H, W, C] uint8 numpy array."""
+    video_tensor: torch.Tensor,
-    # Handle [C, F, H, W] vs [F, C, H, W] format
+    video_format: VideoFormat | None = None,
-    if video_tensor.shape[0] == 3 and video_tensor.shape[1] > 3:
+) -> np.ndarray:
+    """Convert video tensor to [F, H, W, C] uint8 numpy array.
+    If ``video_format`` is provided, it is trusted. Otherwise, the layout is auto-detected
+    using a heuristic that only fires when ``shape[0] == 3 and shape[1] > 3`` (CFHW). The
+    ambiguous ``[C=3, F=3, H, W]`` / ``[F=3, C=3, H, W]`` case cannot be disambiguated and
+    defaults to the FCHW interpretation — callers must pass ``video_format`` explicitly for
+    3-frame CFHW tensors.
+    """
+    if video_format == "CFHW":
         video_tensor = video_tensor.permute(1, 0, 2, 3)  # [C, F, H, W] -> [F, C, H, W]
+    elif video_format is None and video_tensor.shape[0] == 3 and video_tensor.shape[1] > 3:
+        video_tensor = video_tensor.permute(1, 0, 2, 3)
 
     # Normalize to [0, 255] uint8
     if video_tensor.max() <= 1.0:

uv.lock

@@ -2063,7 +2063,7 @@ wheels = [
 
 [[package]]
 name = "ltx-core"
-version = "1.1.1"
+version = "1.1.3"
 source = { editable = "packages/ltx-core" }
 dependencies = [
     { name = "accelerate" },
@@ -2121,11 +2121,13 @@ dev = [{ name = "scikit-image", specifier = ">=0.25.2" }]
 
 [[package]]
 name = "ltx-pipelines"
-version = "1.1.1"
+version = "1.1.3"
 source = { editable = "packages/ltx-pipelines" }
 dependencies = [
     { name = "av" },
     { name = "ltx-core" },
+    { name = "openimageio", version = "3.0.16.0", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-8-ltx-core-fp8-trtllm'" },
+    { name = "openimageio", version = "3.1.11.0", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-8-ltx-core-xformers' or extra != 'extra-8-ltx-core-fp8-trtllm'" },
     { name = "pillow", version = "10.3.0", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-8-ltx-core-fp8-trtllm'" },
     { name = "pillow", version = "12.1.0", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-8-ltx-core-xformers' or extra != 'extra-8-ltx-core-fp8-trtllm'" },
     { name = "tqdm" },
@@ -2135,13 +2137,14 @@ dependencies = [
 requires-dist = [
     { name = "av" },
     { name = "ltx-core", editable = "packages/ltx-core" },
+    { name = "openimageio" },
     { name = "pillow" },
     { name = "tqdm" },
 ]
 
 [[package]]
 name = "ltx-trainer"
-version = "1.1.1"
+version = "1.1.3"
 source = { editable = "packages/ltx-trainer" }
 dependencies = [
     { name = "accelerate" },
@@ -4044,6 +4047,103 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/86/8a/69176a64335aed183529207ba8bc3d329c2999d852b4f3818027203f50e6/opencv_python_headless-4.11.0.86-cp37-abi3-win_amd64.whl", hash = "sha256:6c304df9caa7a6a5710b91709dd4786bf20a74d57672b3c31f7033cc638174ca", size = 39402386, upload-time = "2025-01-16T13:52:56.418Z" },
 ]
 
+[[package]]
+name = "openimageio"
+version = "3.0.16.0"
+source = { registry = "https://pypi.org/simple" }
+resolution-markers = [
+    "python_full_version >= '3.13' and sys_platform == 'darwin'",
+    "python_full_version == '3.12.*' and sys_platform == 'darwin'",
+    "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'linux'",
+    "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'",
+    "python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux'",
+    "python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux'",
+    "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux'",
+    "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux'",
+    "python_full_version == '3.11.*' and sys_platform == 'darwin'",
+    "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'",
+    "python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux'",
+    "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux'",
+    "python_full_version < '3.11' and sys_platform == 'darwin'",
+    "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'",
+    "python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform == 'linux'",
+    "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux'",
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ab/45/f4fbc1f32b75493383ad6dc1820ecb789c552f2d231a7880c17a8ea7b276/openimageio-3.0.16.0.tar.gz", hash = "sha256:5308f04f30c98841012c61eba4972b4e0e94446be6ab4c2b7414904ed0c711c7", size = 6482737, upload-time = "2026-03-01T03:43:38.412Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/25/17/9160d6b4332b430a8b24bb5d0db39af2b2c3b483d36d11c5fa8bc9c9468b/openimageio-3.0.16.0-cp310-cp310-macosx_10_15_x86_64.whl", hash = "sha256:aa96e4100248ecc9bbf0b7fb06b1f6ce781592481de1ebfe2a665679c4219381", size = 6557400, upload-time = "2026-03-01T03:42:39.892Z" },
+    { url = "https://files.pythonhosted.org/packages/2f/26/cf15bcbce65ed45f64fa0806889592c18b598795b7f2461429250a066d00/openimageio-3.0.16.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8d4e65c80d3dca5e5daa1663efe620c6c94476be8c86181118f54aaae141cf79", size = 6194167, upload-time = "2026-03-01T03:42:41.685Z" },
+    { url = "https://files.pythonhosted.org/packages/71/55/660e086ee3acbbbab2b19f8cad93c0b0cc1044a660d0189ea9bf0ed20e93/openimageio-3.0.16.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:84e7c8619c1a583730c73c8c8a247584f5fca980e51d9ec5ad1fd3a34a58aa4e", size = 6391703, upload-time = "2026-03-01T03:42:43.063Z" },
+    { url = "https://files.pythonhosted.org/packages/0c/22/91cd80c88ec41d773e33790895a8b4732884bb5b571d760f0730ea9e184a/openimageio-3.0.16.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b68b7406314b45eb813da12426b3d60046813b0c32c2fadc0ce54264554cfa23", size = 6595708, upload-time = "2026-03-01T03:42:44.732Z" },
+    { url = "https://files.pythonhosted.org/packages/05/95/694fb96fd499c3765fb5239193bef4496e69b98a39eda7c0f1c2e3ebe056/openimageio-3.0.16.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fe0df5af0f255f8abaa9aeb728c4cbdea4bc144a19f510fb0323959b88ddade5", size = 6492976, upload-time = "2026-03-01T03:42:46.473Z" },
+    { url = "https://files.pythonhosted.org/packages/6b/b8/15a357fee630b0016b57bbcbc7e67565dca4d379b0a9916cc10ddf774ae2/openimageio-3.0.16.0-cp310-cp310-win_amd64.whl", hash = "sha256:58bd1f20d6e37759a3d6337f7d2a2929e47e1833ed1f439a1dc3ff841a335d7e", size = 7098417, upload-time = "2026-03-01T03:42:48.093Z" },
+    { url = "https://files.pythonhosted.org/packages/1a/18/a6711865161c82e056ce4001b859a49b1c28830ced8c14dec99bffedb50f/openimageio-3.0.16.0-cp311-cp311-macosx_10_15_x86_64.whl", hash = "sha256:f657f4541dd469e2c9d6cda8b266b33327742142616645ecf971433955ea5eb2", size = 6558184, upload-time = "2026-03-01T03:42:50.653Z" },
+    { url = "https://files.pythonhosted.org/packages/28/49/5ff1c6580ac55073ec7153d9f14b9b9fc5a0cffa3bc9188b21ed11ef6952/openimageio-3.0.16.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:442ac9f06c7592cdc01fc7c9ac827c3be5d2c885c2e71b55ee40e249554a878d", size = 6194989, upload-time = "2026-03-01T03:42:52.385Z" },
+    { url = "https://files.pythonhosted.org/packages/e6/92/25397936a4f881cd5c9731e260009240544d6a959bfd16e14174983a4cea/openimageio-3.0.16.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:55ce42888c85262c3be66133be0bc4236a0f445a3568783e57a467324bcbe03c", size = 6392269, upload-time = "2026-03-01T03:42:54.679Z" },
+    { url = "https://files.pythonhosted.org/packages/ca/32/590315956ec40443d6992a5e0ddb11be5615d6aaf38f6d9edf970ce28685/openimageio-3.0.16.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a7f97ff4ed54d63c6241d1634d2f55004b5f83cf28cdf609302523e2eaec1a09", size = 6596973, upload-time = "2026-03-01T03:42:56.24Z" },
+    { url = "https://files.pythonhosted.org/packages/07/b5/1be18ba1f2ae52598917da774820b4d0020a7ab50e1c54c33ab1820f06d0/openimageio-3.0.16.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fa87d72410f59d969d64112c9dd2b51999f6553f392f757844bdc85060092de5", size = 6493859, upload-time = "2026-03-01T03:42:57.886Z" },
+    { url = "https://files.pythonhosted.org/packages/0a/a7/a6aa117ed3fd28ed67e08a54f74c3ab10c45f1e2702dbd8f12c248b382ea/openimageio-3.0.16.0-cp311-cp311-win_amd64.whl", hash = "sha256:6af821d1371271e2072755702a7c5d073ec67cc27d1909f4e646c5a62a82879c", size = 7099505, upload-time = "2026-03-01T03:42:59.778Z" },
+    { url = "https://files.pythonhosted.org/packages/9b/2a/a58a9ee65f9d60fe8be646bc1c0fdbbbaafe4c77ad36587523535e9e4c63/openimageio-3.0.16.0-cp312-cp312-macosx_10_15_x86_64.whl", hash = "sha256:c2cf2919ef287135b0cc2113dbdc85b64dc0e2fa4bd5c3887fe09c0fd1926924", size = 6643409, upload-time = "2026-03-01T03:43:01.061Z" },
+    { url = "https://files.pythonhosted.org/packages/7a/c5/4aca1892f48a1bfed8b6434a02ad2762a927c0725cfd0196f3a8f46becbe/openimageio-3.0.16.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:732c7e76840b2df62975ccb91ab176e430a7dbf410acd993a30f6fd7a485e698", size = 6262459, upload-time = "2026-03-01T03:43:02.262Z" },
+    { url = "https://files.pythonhosted.org/packages/b0/68/a9e1d45f347ae2ca08d5a6166af615871c323501a4a92c8825a4e604d32d/openimageio-3.0.16.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7c825270d342cd071b369e69e36869271a7c811155744ebd4c9b8a6dbde24ac7", size = 6392045, upload-time = "2026-03-01T03:43:04.195Z" },
+    { url = "https://files.pythonhosted.org/packages/40/a0/554e03a80c8650be3985b4aa427a80491a8cc462c3cde5890357b5976597/openimageio-3.0.16.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:096e0a4940395f73903a295ec423a338594ce07d405206e9d775ba7168861fd8", size = 6597227, upload-time = "2026-03-01T03:43:05.814Z" },
+    { url = "https://files.pythonhosted.org/packages/cd/36/4792c357fa2b19d5ed4f5bc2b634dcb5f658c541ff33263e09002b680cd7/openimageio-3.0.16.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:331d103b35d1d402ff8c7a9ece77b7c2191f5d6e4b6d8299be08374c4c4b91f8", size = 6493718, upload-time = "2026-03-01T03:43:07.37Z" },
+    { url = "https://files.pythonhosted.org/packages/94/9f/e3da32d705a704ed7f1324345dc0a0ac108667bb16a8ba50da7f4c298c0c/openimageio-3.0.16.0-cp312-cp312-win_amd64.whl", hash = "sha256:dddcd78cb35946907eb2be607ba9d620c05f502652937c9ea250078725b1745e", size = 7103136, upload-time = "2026-03-01T03:43:08.983Z" },
+    { url = "https://files.pythonhosted.org/packages/67/61/0da72a9ed22127a563c12db20025255010661f7162f494741c74242d195d/openimageio-3.0.16.0-cp313-cp313-macosx_10_15_x86_64.whl", hash = "sha256:b160601825a8362eb586a760afabe1ebfd5d3b6869f012ccbe5cd637a621efd2", size = 6643458, upload-time = "2026-03-01T03:43:10.869Z" },
+    { url = "https://files.pythonhosted.org/packages/65/fe/5a743e939c3ce8d822f3a826dc93a6a5914265d613bd81cfddc3c6f91c9b/openimageio-3.0.16.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:4259743111db8ea8daf76d44724fcf20f0e2ce641906d36202983ec6092c24d2", size = 6262560, upload-time = "2026-03-01T03:43:12.499Z" },
+    { url = "https://files.pythonhosted.org/packages/de/ae/8c7069774dc198ffd7b193b49c8e3cf897a98bb8c5eb9de754be10957932/openimageio-3.0.16.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5a2a0d97d7091e33b542ae0fb66d7e250ad48a4417dbffa308c48eefd96dc3c6", size = 6392031, upload-time = "2026-03-01T03:43:14.256Z" },
+    { url = "https://files.pythonhosted.org/packages/29/9b/3c05dc022ae1ae9761e2a6d297b10cefaa9cbbcf2feb6a355710d1f4f045/openimageio-3.0.16.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2df45d95e50dcdf4fc07df09b3cfdfb0b6dbfab1dad86a4b744b698058bce62e", size = 6596464, upload-time = "2026-03-01T03:43:16.204Z" },
+    { url = "https://files.pythonhosted.org/packages/0d/82/04d1b969c56c4737783949ccd3446d63b4a0db0003684371b96877a54f53/openimageio-3.0.16.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b64ca87f11057bd60ee15f5280a6e83c719c01ed31fb6b220aa1df7e3966ad3b", size = 6493537, upload-time = "2026-03-01T03:43:17.486Z" },
+    { url = "https://files.pythonhosted.org/packages/e6/97/0bbee7fd1d2e028c4be6292e5f968556a11670bec358d57ba0ff0b0d77de/openimageio-3.0.16.0-cp313-cp313-win_amd64.whl", hash = "sha256:c439c4eb80037aadd34436947eebf3fc06ac61512b0263724bafef6d0878fd76", size = 7103123, upload-time = "2026-03-01T03:43:19.06Z" },
+]
+
+[[package]]
+name = "openimageio"
+version = "3.1.11.0"
+source = { registry = "https://pypi.org/simple" }
+resolution-markers = [
+    "python_full_version >= '3.13' and sys_platform == 'linux'",
+    "python_full_version == '3.12.*' and sys_platform == 'linux'",
+    "python_full_version >= '3.13' and sys_platform != 'linux'",
+    "python_full_version == '3.12.*' and sys_platform != 'linux'",
+    "python_full_version == '3.11.*' and sys_platform == 'linux'",
+    "python_full_version == '3.11.*' and sys_platform != 'linux'",
+    "python_full_version < '3.11' and sys_platform == 'linux'",
+    "python_full_version < '3.11' and sys_platform != 'linux'",
+]
+dependencies = [
+    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra != 'extra-8-ltx-core-fp8-trtllm') or (extra == 'extra-8-ltx-core-fp8-trtllm' and extra == 'extra-8-ltx-core-xformers')" },
+    { name = "numpy", version = "2.4.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra != 'extra-8-ltx-core-fp8-trtllm') or (extra == 'extra-8-ltx-core-fp8-trtllm' and extra == 'extra-8-ltx-core-xformers')" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/f6/bd/d49fe9f78b244251301da7b4caff7053b151df55ce0e04829004ac9f65ba/openimageio-3.1.11.0.tar.gz", hash = "sha256:7c096607569b0d01266da739bbf7496f1c0a8fb4449092f8f66fde2bf3a20494", size = 6555440, upload-time = "2026-03-01T03:09:00.936Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b9/cc/eb7725c3ca6c76fd5defe92b2d558773b1a7a89a69c483bb33d791b0a723/openimageio-3.1.11.0-cp310-cp310-macosx_10_15_x86_64.whl", hash = "sha256:ef9377cdd84d67c415e6fc67f75553cacad36976cfd3f5ae1f0b2b5a6c46a7ef", size = 6851102, upload-time = "2026-03-01T03:08:18.322Z" },
+    { url = "https://files.pythonhosted.org/packages/2a/54/ffc2d3b70dee7690dae808cfa8787cfa1f109fc7b4cd68632cade9fab5f3/openimageio-3.1.11.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:0c497762ec7993c9af51079edd723ca0a37faa085c6115b8e69f151460a60927", size = 6465995, upload-time = "2026-03-01T03:08:19.907Z" },
+    { url = "https://files.pythonhosted.org/packages/7f/7c/e2f32ee394db041af8477d984aef2c04067713e916830fb113ade2dd3426/openimageio-3.1.11.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1cfc57d9e790e76e76498b0e5c5a0bed83ae1be6104c635eca114265eefff34c", size = 6501598, upload-time = "2026-03-01T03:08:21.486Z" },
+    { url = "https://files.pythonhosted.org/packages/e1/4d/fa9d641d424d7e16a44a502c440165687adad13ca005507dbf29d8b7febf/openimageio-3.1.11.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0f01e9e115b1832bd27bb6a62e29bf435cb27f3f41b589ac7d2838bdda530f15", size = 6789743, upload-time = "2026-03-01T03:08:22.714Z" },
+    { url = "https://files.pythonhosted.org/packages/34/0a/cbde5147af9db425a5c692147da45f849d1203304f195bcae71875b43afb/openimageio-3.1.11.0-cp310-cp310-win_amd64.whl", hash = "sha256:4b57b599294afac5f21885949014ca1469866c3b648375b834d19723c67dc52c", size = 7459665, upload-time = "2026-03-01T03:08:24.267Z" },
+    { url = "https://files.pythonhosted.org/packages/0e/9d/6d67fb8358faeb7f6b195ea35483ffa58d9fd45806dbedea471c79ae8050/openimageio-3.1.11.0-cp311-cp311-macosx_10_15_x86_64.whl", hash = "sha256:d02021c7a4baf0418cde43ebddd0103886cc809c8e85c9c8cd4a1cbe69879c5a", size = 6852201, upload-time = "2026-03-01T03:08:25.864Z" },
+    { url = "https://files.pythonhosted.org/packages/ec/c2/1412eb64ea7a7156ebe335a03118e8d2dc116e63e354d44fceb475b520cc/openimageio-3.1.11.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:daa4396f8094ed94c9b963711b2a680dc8f34c5361be5bc06804b0fc94cb4f4b", size = 6466903, upload-time = "2026-03-01T03:08:27.124Z" },
+    { url = "https://files.pythonhosted.org/packages/75/20/16e4d4a310628d0960588f919a7fc85dd94be4eb37a043f5cacfb9e024fc/openimageio-3.1.11.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0f457ca10c45d4e03dd4d3af13456962bee2ee6abee91bf40a8d21ffa7da1f74", size = 6502239, upload-time = "2026-03-01T03:08:28.637Z" },
+    { url = "https://files.pythonhosted.org/packages/9c/6f/70f75c4fbc958e9f4f829aeb8d88558591cdcdb6dd36e97d8ee6f50c4011/openimageio-3.1.11.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:11655436841f11d268049f5494c6621dfd6370c104b4aa8c20f0724271850796", size = 6790238, upload-time = "2026-03-01T03:08:30.083Z" },
+    { url = "https://files.pythonhosted.org/packages/3b/1e/f85f0393237dbbcec0559c357795462e43eb2fead40d2954a99ebf24c310/openimageio-3.1.11.0-cp311-cp311-win_amd64.whl", hash = "sha256:58eacdc68992c49257be43a30ac424fd6cf617b9de1e516debc14c0bd93af075", size = 7460293, upload-time = "2026-03-01T03:08:31.646Z" },
+    { url = "https://files.pythonhosted.org/packages/cf/2c/24036c501b24f0d4136d9c4605a8861f5affc87df23fa78c279e0a9216c0/openimageio-3.1.11.0-cp312-cp312-macosx_10_15_x86_64.whl", hash = "sha256:53337dd2e11f7f0a07fbe664d21c0d578a0674695da1e5f327764a3f421aff83", size = 6942447, upload-time = "2026-03-01T03:08:33.179Z" },
+    { url = "https://files.pythonhosted.org/packages/69/ca/d84f062cbceb35a8006a9a3393da1893e23ffd8147c53d5f461fa99898cc/openimageio-3.1.11.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8d50f1d8159177561c6d5efcc0f2d7f9fd5393cc147a39b546849da8762bb4f8", size = 6540326, upload-time = "2026-03-01T03:08:34.542Z" },
+    { url = "https://files.pythonhosted.org/packages/6c/f8/ca6f4ca3d5dc2b2cc1411a1083b926d5fc9eb22e57d352333027b3eb4659/openimageio-3.1.11.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:212dd58364f8e11c0fbd90f82bbaf617afc68b4f3910538d0ac2324ef60a4c28", size = 6503147, upload-time = "2026-03-01T03:08:35.87Z" },
+    { url = "https://files.pythonhosted.org/packages/81/9c/5ea545f62790bc4f8b9c04d433963f9fe5f08548c8663e89a54f1db7002d/openimageio-3.1.11.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f845304bb17a1fd48f19c972572eb03ee3f96336f2d5599d043cb7ac1c348035", size = 6791359, upload-time = "2026-03-01T03:08:37.412Z" },
+    { url = "https://files.pythonhosted.org/packages/9b/ba/c64923d962dd6c2dd0421feaa6000b081bcf1e26e07a125da94488f9cd39/openimageio-3.1.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:30284fb3367e480eaa345b122ccf38df61bd3885c78b237942a585cefaf0538a", size = 7466181, upload-time = "2026-03-01T03:08:38.931Z" },
+    { url = "https://files.pythonhosted.org/packages/00/81/76375647018dce289e13f8a75c42cddbd9c688baa50d1dc99f73457508f0/openimageio-3.1.11.0-cp313-cp313-macosx_10_15_x86_64.whl", hash = "sha256:35ac810ccf3aa8a4cdc1fb19289dab57e20a9d915a78041a6fa2052d3b7a8c3c", size = 6942555, upload-time = "2026-03-01T03:08:40.539Z" },
+    { url = "https://files.pythonhosted.org/packages/6b/a1/b1a861bc03462bc5ac6e6a9eac3e4f9e28302f0cc14748e0cd3f1f1efadd/openimageio-3.1.11.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:5d294f1d631baea51cf9d940ee770f789946dcfae187cf78a5b440a3cc88c4c7", size = 6540381, upload-time = "2026-03-01T03:08:41.796Z" },
+    { url = "https://files.pythonhosted.org/packages/c2/5a/79ceb6c7cf57fdce947e9b341c5de1fbf6e219a0db5ce990e69431c54f90/openimageio-3.1.11.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:78cffd41f1d0e575c56aa5e601a0f3c98347cb9a20d3ca6210b4ef881cef256c", size = 6503197, upload-time = "2026-03-01T03:08:43.23Z" },
+    { url = "https://files.pythonhosted.org/packages/e1/49/6593704b762dfe1296229323219e216fda029c6153622a57f1e865f82160/openimageio-3.1.11.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:020e7b5a0d2898dca1c66c7c6f800aa2c307a509ffd3dbed2898562445b5f6bd", size = 6791511, upload-time = "2026-03-01T03:08:44.544Z" },
+    { url = "https://files.pythonhosted.org/packages/60/c2/43c1eeb5dcc7620a212c1705dc03547d20c917f3761fd6a2bde3fa3b9fb1/openimageio-3.1.11.0-cp313-cp313-win_amd64.whl", hash = "sha256:cb331042a297bd3de87b4eed5efed4c337e1446aaed3807ac4899c129e33d61a", size = 7466220, upload-time = "2026-03-01T03:08:45.78Z" },
+    { url = "https://files.pythonhosted.org/packages/2e/b9/6b09045f1a92cd38a59fd3c03dd59374f27c03270943ebd5844247633c70/openimageio-3.1.11.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:f7c9b14c3981cd992b33f8a47735ce099687cddc9d05ffdd03b87490c0bcdd47", size = 6940155, upload-time = "2026-03-01T03:08:47.09Z" },
+    { url = "https://files.pythonhosted.org/packages/13/d8/f10343f5ddf6bc9003f8e9c1eae3b2d4b5cfed01fccde6fffe5a7cafce33/openimageio-3.1.11.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:4f5e106054c01c5a768c6442b712493bfe21526dd4e202e3c18245a0224e6d53", size = 6536784, upload-time = "2026-03-01T03:08:48.364Z" },
+    { url = "https://files.pythonhosted.org/packages/4b/6f/4709da066a4dafef14b5bb0e7e221fa9d65e196806ce22721862627ab37c/openimageio-3.1.11.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4c2affc140802fdc42b2f01aeebedfdd64ff29891fdc2e2c708d4653852eff2b", size = 6503379, upload-time = "2026-03-01T03:08:49.947Z" },
+    { url = "https://files.pythonhosted.org/packages/0d/5b/de755dfb3a8837592263511fc14966901d1b8d880eeb3cde9d65b8f02e81/openimageio-3.1.11.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ee25034ceff73b8f08ae4b562a5a62cfcdaa0e873ce913c5df72d9dc2a0d56c7", size = 6791460, upload-time = "2026-03-01T03:08:51.483Z" },
+    { url = "https://files.pythonhosted.org/packages/20/fb/159531b67af2ddca76dc1bbfcc31e4b94a8d91215ae0fa168e0243327b3b/openimageio-3.1.11.0-cp314-cp314-win_amd64.whl", hash = "sha256:1945317959ac1905e97c8d9479efc716bb765dec82a647365d37e046e154a46c", size = 7767018, upload-time = "2026-03-01T03:08:52.707Z" },
+]
+
 [[package]]
 name = "openmpi"
 version = "5.0.9"
@@ -7005,7 +7105,7 @@ dependencies = [
     { name = "torch", version = "2.9.1", source = { registry = "https://pypi.org/simple" } },
 ]
 wheels = [
-    { url = "https://download.pytorch.org/whl/cu129/xformers-0.0.33%2B5d4b92a5.d20251029-cp39-abi3-linux_x86_64.whl" },
+    { url = "https://download.pytorch.org/whl/cu129/xformers-0.0.33%2B5d4b92a5.d20251029-cp39-abi3-linux_x86_64.whl", hash = "sha256:4e4f2dea153b60ca4f21cc44b82c072b97b096a16f1920c7539c91ec8ffd7ba4", upload-time = "2026-04-27T19:01:32Z" },
 ]
 
 [[package]]