Automated PR - 2026-06-17

packages/ltx-core/README.md

@@ -8,6 +8,7 @@ The foundational library for the LTX-2 Audio-Video generation model. This packag
 - **`conditioning/`**: Tools for preparing latent states and applying conditioning (image, video, keyframes)
 - **`guidance/`**: Perturbation system for fine-grained control over attention mechanisms
 - **`loader/`**: Utilities for loading weights from `.safetensors`, fusing LoRAs, and managing memory
+- **`block_streaming/`**: Memory-efficient inference that streams transformer blocks through the GPU one at a time (from pinned CPU buffers or directly from disk)
 - **`model/`**: PyTorch implementations of the LTX-2 Transformer, Video VAE, Audio VAE, Vocoder and Upscaler
 - **`text_encoders/gemma`**: Gemma text encoder implementation with tokenizers, feature extractors, and separate encoders for audio-video and video-only generation
 - **`quantization/`**: FP8 quantization backends (FP8-TensorRT-LLM scaled MM, FP8 cast) for reduced memory footprint.
@@ -55,6 +56,7 @@ pip install -e packages/ltx-core
 
 - **Loader** ([`loader/`](src/ltx_core/loader/)): Model loading from `.safetensors`, LoRA fusion, weight remapping, and memory management
 - **Quantization** ([`quantization/`](src/ltx_core/quantization/)): FP8 quantization backends for reduced memory footprint and faster inference
+- **Block Streaming** ([`block_streaming/`](src/ltx_core/block_streaming/)): Streams transformer blocks through the GPU one block at a time, so the full model runs on machines without enough memory to hold all its weights at once
 
 ### Loader
 
@@ -157,6 +159,39 @@ from ltx_core.quantization.fp8_cast import build_policy as build_fp8_cast_policy
 policy = build_fp8_cast_policy("/path/to/checkpoint.safetensors")
 ```
 
+### Block Streaming
+
+The `block_streaming/` module ([`src/ltx_core/block_streaming/`](src/ltx_core/block_streaming/)) lets the full model run on machines that lack the memory to hold all of its weights at once. It streams the transformer's blocks through a small rolling set of GPU buffers, loading each block's weights just before it runs and recycling them afterwards, so only a few blocks are resident on the GPU at any moment. Construct it with `StreamingModelBuilder`, which returns a `BlockStreamingWrapper` -- an `nn.Module` drop-in for the wrapped model.
+
+#### Strategies
+
+The strategy is chosen automatically from `cpu_slots_count` relative to the number of blocks:
+
+- **RAM streaming** (default, `cpu_slots_count` omitted or `>= num_blocks`): all blocks are pre-loaded into pinned CPU buffers (with LoRA fusion) at build time, then copied to the GPU on demand. Fast; higher CPU memory.
+- **Disk streaming** (`cpu_slots_count < num_blocks`): blocks are read from the `.safetensors` file on demand on a background worker thread. Slower; lowest CPU memory.
+
+#### Basic usage
+
+```python
+import torch
+from ltx_core.block_streaming import StreamingModelBuilder
+
+builder = StreamingModelBuilder(
+    model_class_configurator=MyModelConfigurator,
+    model_path="/path/to/model.safetensors",
+    blocks_attr="transformer_blocks",      # dotted path to the nn.ModuleList
+    blocks_prefix="transformer_blocks",    # state-dict key prefix for block weights
+)
+
+# Omit cpu_slots_count for RAM streaming; pass a value < num_blocks for disk streaming.
+model = builder.build(
+    device=torch.device("cuda"),
+    dtype=torch.bfloat16,
+    cpu_slots_count=4,
+    gpu_slots_count=2,
+)
+```
+
 For complete, production-ready pipeline implementations that combine these building blocks, see the [`ltx-pipelines`](../ltx-pipelines/) package.
 
 ---

packages/ltx-core/pyproject.toml

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

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

@@ -5,8 +5,9 @@ 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.
+- **Disk streaming** (``cpu_slots < blocks_number``): blocks are read from
+  disk on demand by a :class:`DiskWeightSource`, on a background worker
+  thread.  Slower, lower CPU memory.
 """
 
 from ltx_core.block_streaming.builder import DISK_CPU_SLOTS, StreamingModelBuilder

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

@@ -0,0 +1,83 @@
+"""BlockFetcher: async disk reads on a worker thread."""
+
+from __future__ import annotations
+
+import logging
+import queue
+import threading
+from dataclasses import dataclass
+
+import torch
+
+from ltx_core.block_streaming.disk import DiskBlockReader
+
+logger = logging.getLogger(__name__)
+
+Buffer = dict[str, torch.Tensor]
+
+
+@dataclass(slots=True)
+class FetchHandle:
+    """Caller-facing handle for an outstanding read returned by :meth:`BlockFetcher.submit`.
+    Carries only what the caller needs: a completion event the worker sets and the
+    read's error. The worker updates these once the read finishes.
+    """
+
+    _done: threading.Event
+    _error: BaseException | None = None
+
+    def wait(self) -> BaseException | None:
+        """Block until the read finishes; return its error, or ``None`` on success."""
+        self._done.wait()
+        return self._error
+
+
+@dataclass(slots=True)
+class _ReadRequest:
+    """One outstanding read, internal to :class:`BlockFetcher`.
+    The fetcher's worker reads block ``idx`` into the caller-carved ``buffer`` and
+    updates ``handle`` (its error, then its event) once the read has finished.
+    """
+
+    idx: int
+    buffer: Buffer
+    handle: FetchHandle
+
+
+class BlockFetcher:
+    """Fills caller-supplied buffers on a worker thread."""
+
+    def __init__(self, reader: DiskBlockReader) -> None:
+        self._reader = reader
+        self._request_queue: queue.SimpleQueue[_ReadRequest | None] = queue.SimpleQueue()
+        self._worker = threading.Thread(target=self._run, name="BlockFetcher-IO", daemon=True)
+        self._worker.start()
+
+    def submit(self, idx: int, buffer: Buffer) -> FetchHandle:
+        """Enqueue a read of block *idx* into the caller-carved *buffer*, return its handle."""
+        handle = FetchHandle(_done=threading.Event())
+        request = _ReadRequest(idx=idx, buffer=buffer, handle=handle)
+        self._request_queue.put(request)
+        return handle
+
+    def cleanup(self) -> None:
+        """Drain pending reads, join the worker, close the reader."""
+        self._request_queue.put(None)
+        self._worker.join()
+        self._reader.cleanup()
+
+    def _run(self) -> None:
+        # Pinned buffers are allocated under the caller's inference_mode, so
+        # in-place copy_ from this thread requires inference_mode here too.
+        with torch.inference_mode():
+            while True:
+                request = self._request_queue.get()
+                if request is None:
+                    return
+
+                try:
+                    self._reader.read_into(request.buffer, request.idx)
+                except Exception as exc:
+                    logger.exception("BlockFetcher: fetch failed for item %d", request.idx)
+                    request.handle._error = exc
+                request.handle._done.set()

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

@@ -2,19 +2,28 @@
 
 from __future__ import annotations
 
+import copy
 import logging
-from dataclasses import dataclass, field, replace
-from typing import Generic
+from dataclasses import replace
+from typing import TYPE_CHECKING, Final, Generic
 
 import safetensors
 import torch
 from torch import nn
 
+from ltx_core.block_streaming import utils as bs_utils
+from ltx_core.block_streaming.block_fetcher import BlockFetcher
 from ltx_core.block_streaming.disk import DiskBlockReader, DiskTensorReader, LoraSource
-from ltx_core.block_streaming.pool import WeightPool
+from ltx_core.block_streaming.pool import BufferPool
 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.source import DiskWeightSource, PinnedBlock, PinnedWeightSource, WeightSource
+from ltx_core.block_streaming.utils import (
+    carve_buffer,
+    derive_layout,
+    layout_nbytes,
+    make_block_key,
+    resolve_attr,
+)
 from ltx_core.block_streaming.wrapper import BlockStreamingWrapper
 from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule, fuse_lora_weights
 from ltx_core.loader.helpers import create_meta_model, load_state_dict, read_model_config
@@ -25,23 +34,29 @@ from ltx_core.loader.primitives import (
     ModelBuilderProtocol,
     StateDict,
     StateDictLoader,
+    TensorLayout,
 )
 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
 
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
 logger = logging.getLogger(__name__)
 
 DISK_CPU_SLOTS = 2
 _DEFAULT_GPU_SLOTS = 2
+_PREFETCH_DEPTH = 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`).
+    The builder is immutable (``with_*`` return modified copies) and exposes
+    its state via read-only properties backed by private attributes.
     Args:
         model_class_configurator: Creates the model from a config dict.
         model_path: One or more ``.safetensors`` checkpoint paths.
@@ -59,30 +74,99 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
             (e.g. ``"transformer_blocks"``).
     """
 
-    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)
-    fuse_rule: FuseRule = bf16_fuse_rule
+    def __init__(
+        self,
+        model_class_configurator: type[ModelConfigurator[ModelType]],
+        model_path: str | tuple[str, ...],
+        model_sd_ops: SDOps | None = None,
+        module_ops: tuple[ModuleOps, ...] = (),
+        loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
+        model_loader: StateDictLoader | None = None,
+        registry: Registry | None = None,
+        fuse_rule: FuseRule = bf16_fuse_rule,
+        blocks_attr: str = "",
+        blocks_prefix: str = "",
+    ) -> None:
+        # Read-only: typed with the covariant ModelType, so it must not be a mutable attribute.
+        self._model_class_configurator: Final = model_class_configurator
+        self._model_path = model_path
+        self._model_sd_ops = model_sd_ops
+        self._module_ops = module_ops
+        self._loras = loras
+        self._model_loader = model_loader if model_loader is not None else SafetensorsModelStateDictLoader()
+        self._registry = registry if registry is not None else DummyRegistry()
+        self._fuse_rule = fuse_rule
+        self._blocks_attr = blocks_attr
+        self._blocks_prefix = blocks_prefix
 
-    # Streaming-specific
-    blocks_attr: str = ""
-    blocks_prefix: str = ""
+    @property
+    def model_class_configurator(self) -> type[ModelConfigurator[ModelType]]:
+        return self._model_class_configurator
 
-    def with_sd_ops(self, sd_ops: SDOps | None) -> StreamingModelBuilder:
-        return replace(self, model_sd_ops=sd_ops)
+    @property
+    def model_path(self) -> str | tuple[str, ...]:
+        return self._model_path
 
-    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> StreamingModelBuilder:
-        return replace(self, module_ops=module_ops)
+    @property
+    def model_sd_ops(self) -> SDOps | None:
+        return self._model_sd_ops
 
-    def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> StreamingModelBuilder:
-        return replace(self, loras=loras)
+    @property
+    def module_ops(self) -> tuple[ModuleOps, ...]:
+        return self._module_ops
 
-    def with_fuse_rule(self, fuse_rule: FuseRule) -> StreamingModelBuilder:
-        return replace(self, fuse_rule=fuse_rule)
+    @property
+    def loras(self) -> tuple[LoraPathStrengthAndSDOps, ...]:
+        return self._loras
+
+    @property
+    def model_loader(self) -> StateDictLoader:
+        return self._model_loader
+
+    @property
+    def registry(self) -> Registry:
+        return self._registry
+
+    @property
+    def fuse_rule(self) -> FuseRule:
+        return self._fuse_rule
+
+    @property
+    def blocks_attr(self) -> str:
+        return self._blocks_attr
+
+    @property
+    def blocks_prefix(self) -> str:
+        return self._blocks_prefix
+
+    def with_sd_ops(self, sd_ops: SDOps | None) -> Self:
+        clone = copy.copy(self)
+        clone._model_sd_ops = sd_ops
+        return clone
+
+    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> Self:
+        clone = copy.copy(self)
+        clone._module_ops = module_ops
+        return clone
+
+    def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> Self:
+        clone = copy.copy(self)
+        clone._loras = loras
+        return clone
+
+    def with_registry(self, registry: Registry) -> Self:
+        clone = copy.copy(self)
+        clone._registry = registry
+        return clone
+
+    def with_lora_load_device(self, device: torch.device) -> Self:
+        # Streaming fuses LoRAs into pinned CPU buffers; no other staging device is meaningful.
+        raise NotImplementedError("StreamingModelBuilder loads LoRA weights on CPU only.")
+
+    def with_fuse_rule(self, fuse_rule: FuseRule) -> Self:
+        clone = copy.copy(self)
+        clone._fuse_rule = fuse_rule
+        return clone
 
     def model_config(self) -> dict:
         """Read model configuration from the checkpoint metadata."""
@@ -94,16 +178,16 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
 
     def build(
         self,
-        target_device: torch.device,
-        dtype: torch.dtype,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
         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``).
+            device: GPU device for compute. ``None`` defaults to ``cuda``.
+            dtype: Weight dtype (e.g. ``torch.bfloat16``). Required.
             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.
@@ -111,6 +195,9 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
         """
         if not self.blocks_prefix:
             raise ValueError("blocks_prefix must be non-empty for streaming")
+        if dtype is None:
+            raise ValueError("StreamingModelBuilder.build requires an explicit dtype")
+        device = device if device is not None else torch.device("cuda")
 
         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)
@@ -120,31 +207,44 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
 
         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)
+        expected_indices = set(range(len(blocks)))
+        if set(block_key_map) != expected_indices:
+            missing = sorted(expected_indices - set(block_key_map))
+            extra = sorted(set(block_key_map) - expected_indices)
+            raise ValueError(
+                f"Block weights under prefix '{self.blocks_prefix}.' do not match the {len(blocks)} model blocks: "
+                f"missing indices {missing}, unexpected indices {extra}"
+            )
 
         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):
+            lora_sd_and_strengths = self._load_lora_sds()
             source, lora_sources = self._build_pinned_source(
-                meta_model, target_device, dtype, cpu_slots_count, block_key_map, non_block_keys
+                blocks, dtype, cpu_slots_count, block_key_map, lora_sd_and_strengths
             )
+            non_block_loras = lora_sd_and_strengths
         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
+                blocks, dtype, cpu_slots_count, reader, block_key_map, prefetch_depth=_PREFETCH_DEPTH
             )
+            non_block_loras = [src.as_state_dict_with_strength() for src in lora_sources]
 
-        copy_stream = torch.cuda.Stream(device=target_device)
-        gpu_pool = WeightPool(
-            source.block_layout,
+        self._load_non_block_weights(meta_model, non_block_keys, device, dtype, non_block_loras)
+
+        copy_stream = torch.cuda.Stream(device=device)
+        gpu_pool = BufferPool(
+            source.slot_nbytes,
             gpu_slots_count,
-            target_device,
+            device,
             reuse_barrier=lambda event: copy_stream.wait_event(event),
         )
         provider = WeightsProvider(
             gpu_pool,
             copy_stream,
-            target_device,
+            device,
             source,
             lora_sources,
             self.blocks_prefix,
@@ -154,24 +254,12 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
             model=meta_model,
             blocks=blocks,
             provider=provider,
-            target_device=target_device,
+            target_device=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 = [
+    def _load_lora_sds(self) -> list[LoraStateDictWithStrength]:
+        """Load each configured LoRA into a state dict for fusion (pinned path)."""
+        return [
             LoraStateDictWithStrength(
                 load_state_dict([lora.path], self.model_loader, self.registry, torch.device("cpu"), lora.sd_ops),
                 lora.strength,
@@ -179,6 +267,25 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
             for lora in self.loras
         ]
 
+    def _filtered_sd_ops(self, name_suffix: str, allowed_model_keys: frozenset[str]) -> SDOps:
+        """``model_sd_ops`` restricted to *allowed_model_keys* (post-rename keys).
+        The loader skips keys filtered to None before reading them, so a restricted
+        load never materializes the excluded partition. The distinct ``name`` avoids
+        a registry cache-id collision with the other partition.
+        """
+        base = self.model_sd_ops if self.model_sd_ops is not None else SDOps("streaming").with_matching()
+        allowed = allowed_model_keys if base.allowed_keys is None else (allowed_model_keys & base.allowed_keys)
+        return replace(base, name=f"{base.name}__{name_suffix}", allowed_keys=allowed)
+
+    def _build_pinned_source(
+        self,
+        blocks: nn.ModuleList,
+        dtype: torch.dtype,
+        cpu_slots_count: int,
+        block_key_map: dict[int, list[tuple[str, str]]],
+        lora_sd_and_strengths: list[LoraStateDictWithStrength],
+    ) -> tuple[WeightSource, list[LoraSource]]:
+        """Pre-load each block into its own contiguous pinned CPU buffer with LoRA fusion."""
         for block_idx in block_key_map:
             if block_idx >= cpu_slots_count:
                 raise ValueError(
@@ -186,87 +293,75 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
                     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] = {}
+        # One contiguous pinned buffer per block, carved into per-param views. The
+        # views (flattened by full key) are filled in place; the source then keeps
+        # only the contiguous buffer and the layout to re-carve it on read.
+        pinned_buffers: dict[int, torch.Tensor] = {}
+        block_layouts: dict[int, TensorLayout] = {}
+        fill_views: 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)
+            block_state = _block_state(blocks[block_idx])
+            layout = derive_layout({param_name: block_state[param_name] for _sft_key, param_name in entries}, dtype)
+            buffer = bs_utils.alloc_buffer(layout_nbytes(layout), torch.device("cpu"), pin_memory=True)
+            views = carve_buffer(buffer, layout)
+            pinned_buffers[block_idx] = buffer
+            block_layouts[block_idx] = layout
+            for param_name, view in views.items():
+                fill_views[make_block_key(self.blocks_prefix, block_idx, param_name)] = view
+
+        block_sd = load_state_dict(
+            self.model_path,
+            self.model_loader,
+            self.registry,
+            torch.device("cpu"),
+            self._filtered_sd_ops("blocks", frozenset(fill_views)),
+        )
 
         should_sync = False
         for key, fused in fuse_lora_weights(
-            model_sd, lora_sd_and_strengths, fuse_rule=self.fuse_rule, preserve_input_device=False
+            block_sd, lora_sd_and_strengths, fuse_rule=self.fuse_rule, 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 key not in fill_views:
+                raise ValueError(f"Block-restricted load produced {key!r}, which is not a pinned block weight")
+            fill_views[key].copy_(fused, non_blocking=True)
+            block_sd.sd[key] = None
+            should_sync = True
         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:
+        for key, view in fill_views.items():
+            if block_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] = {
-            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)
+            view.copy_(block_sd.sd[key])
+            block_sd.sd[key] = None
 
+        pinned = {idx: PinnedBlock(pinned_buffers[idx], block_layouts[idx]) for idx in pinned_buffers}
         return PinnedWeightSource(pinned), []
 
     def _build_disk_source(
         self,
-        meta_model: nn.Module,
-        target_device: torch.device,
+        blocks: nn.ModuleList,
         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]],
+        prefetch_depth: int,
     ) -> 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.
+        """Create a DiskWeightSource backed by a DiskBlockReader.
+        Pool slots are sized to the largest block and carved per block on read, so
+        heterogeneous blocks (e.g. layers with differing attention layouts) share
+        one pool. Pool capacity is ``cpu_slots_count + prefetch_depth`` so the
+        lookahead loop in ``DiskWeightSource.get`` never evicts its own target.
+        Layouts come from the meta model; assumes module_ops keep 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]
+        block_layouts = _block_layouts(blocks, block_key_map, dtype)
+        slot_nbytes = max(layout_nbytes(layout) for layout in block_layouts.values())
 
-        self._load_non_block_weights(
-            reader,
-            non_block_keys,
-            meta_model,
-            target_device,
-            dtype,
-            sd_ops=self.model_sd_ops,
-            lora_sources=lora_sources,
-            fuse_rule=self.fuse_rule,
-        )
-
-        blocks = resolve_attr(meta_model, self.blocks_attr)
-        layout = derive_layout(dict(blocks[0].named_parameters()), dtype)
-
-        cpu_pool = WeightPool(
-            layout,
-            cpu_slots_count,
+        cpu_pool = BufferPool(
+            slot_nbytes,
+            cpu_slots_count + prefetch_depth,
             torch.device("cpu"),
             reuse_barrier=lambda event: event.synchronize(),
             pin_memory=True,
@@ -277,43 +372,42 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
             sd_ops=self.model_sd_ops,
             blocks_prefix=self.blocks_prefix,
         )
-        source = DiskWeightSource(cpu_pool, block_reader)
+        fetcher = BlockFetcher(block_reader)
+        source = DiskWeightSource(
+            cpu_pool,
+            fetcher,
+            block_layouts,
+            blocks_number=len(blocks),
+            prefetch_depth=prefetch_depth,
+        )
+        lora_sources = [LoraSource(lora.path, lora.sd_ops, lora.strength) for lora in self.loras]
+
         return source, lora_sources
 
-    @staticmethod
     @torch.inference_mode()
     def _load_non_block_weights(
-        reader: DiskTensorReader,
-        non_block_keys: list[tuple[str, str]],
+        self,
         model: nn.Module,
+        non_block_keys: list[tuple[str, str]],
         device: torch.device,
         dtype: torch.dtype,
-        sd_ops: SDOps | None = None,
-        lora_sources: list[LoraSource] | None = None,
-        fuse_rule: FuseRule = bf16_fuse_rule,
+        lora_sd_and_strengths: list[LoraStateDictWithStrength],
     ) -> None:
-        """Load non-block weights into *model* on *device*, applying ``sd_ops`` and fusing LoRAs.
-        Fusion goes through :func:`fuse_lora_weights` under *fuse_rule* so the
-        bf16 rounding pattern matches the non-streaming and block-streaming
-        paths, and any quantization-specific fuse rule the builder configures
-        is honored here as well.
+        """Load the non-block weights onto *device* and fuse LoRAs -- both paths.
+        Reads through the loader with ``model_sd_ops`` restricted to the
+        non-block keys, so ``sd_ops`` (incl. kv-ops such as Gemma's ``lm_head``
+        duplication) is applied exactly once and block tensors are never read.
         """
-        non_block_sd: dict[str, torch.Tensor] = {}
-        for sft_key, model_key in non_block_keys:
-            tensor = reader.get_tensor(sft_key).to(device=device, dtype=dtype)
-            if sd_ops is not None:
-                for kv in sd_ops.apply_to_key_value(model_key, tensor):
-                    non_block_sd[kv.new_key] = kv.new_value
-            else:
-                non_block_sd[model_key] = tensor
+        non_block_sd_ops = self._filtered_sd_ops("non_block", frozenset(mk for _sft_key, mk in non_block_keys))
+        loaded = load_state_dict(self.model_path, self.model_loader, self.registry, device, non_block_sd_ops)
+        non_block_sd = {key: tensor.to(dtype=dtype) for key, tensor in loaded.sd.items()}
 
-        if lora_sources:
-            lora_sd_and_strengths = [src.as_state_dict_with_strength() for src in lora_sources]
+        if lora_sd_and_strengths:
             non_block_state = StateDict(sd=non_block_sd, device=device, size=0, dtype={dtype})
             for key, fused in fuse_lora_weights(
                 non_block_state,
                 lora_sd_and_strengths,
-                fuse_rule=fuse_rule,
+                fuse_rule=self.fuse_rule,
                 preserve_input_device=True,
             ):
                 non_block_sd[key] = fused
@@ -321,6 +415,35 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
         model.load_state_dict(non_block_sd, strict=False, assign=True)
 
 
+def _block_state(block: nn.Module) -> dict[str, torch.Tensor]:
+    """Streamed-eligible tensors of a block: parameters then buffers.
+    Block streaming swaps both params and checkpoint-backed buffers (e.g. Gemma4's
+    per-layer ``layer_scalar``), so the layout, pinned packing, and meta-ordering
+    all consult parameters and buffers together. Non-checkpoint (computed) buffers
+    are harmless here -- only keys present in ``block_key_map`` are ever streamed.
+    """
+    return {**dict(block.named_parameters()), **dict(block.named_buffers())}
+
+
+def _block_layouts(
+    blocks: nn.ModuleList,
+    block_key_map: dict[int, list[tuple[str, str]]],
+    dtype: torch.dtype,
+) -> dict[int, TensorLayout]:
+    """Per-block layout of the streamed tensors, taken from the meta model.
+    Blocks may differ in shape and even in which tensors they have (e.g. Gemma4's
+    full-attention layers drop ``v_proj``), so each block gets its own layout in
+    ``block_key_map`` order. The pinned packing, the disk reader, and the GPU carve
+    all key off this same per-block layout, so the provider's contiguous H2D copy
+    is valid for any entry order (no cross-block ordering required).
+    """
+    layouts: dict[int, TensorLayout] = {}
+    for idx, entries in block_key_map.items():
+        state = _block_state(blocks[idx])
+        layouts[idx] = derive_layout({param_name: state[param_name] for _sft_key, param_name in entries}, dtype)
+    return layouts
+
+
 def _scan_checkpoint_keys(
     checkpoint_paths: list[str],
     sd_ops: SDOps | None,

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

@@ -128,8 +128,8 @@ class LoraSource:
 
     def as_state_dict_with_strength(self) -> LoraStateDictWithStrength:
         """Return a :class:`LoraStateDictWithStrength` view of the pinned A/B factors.
-        Lets :func:`fuse_lora_weights` consume disk-streaming LoRAs without
-        re-reading the safetensors file or re-applying ``sd_ops``.
+        Lets non-block fusion consume the already-loaded disk-streaming LoRA
+        without re-reading the safetensors file or re-applying ``sd_ops``.
         """
         sd: dict[str, torch.Tensor] = {}
         for prefix, (a, b) in self._pinned_ab.items():

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

@@ -1,4 +1,4 @@
-"""Weight buffer pool for block streaming."""
+"""Raw buffer pool for block streaming."""
 
 from __future__ import annotations
 
@@ -7,69 +7,64 @@ from typing import Callable
 
 import torch
 
-from ltx_core.block_streaming.utils import allocate_layout_views
-from ltx_core.loader.primitives import TensorLayout
+from ltx_core.block_streaming import utils
 
 
-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`.
+class BufferPool:
+    """Fixed pool of pre-allocated raw buffer slots with event-based reuse.
+    Slots are carved from a single contiguous ``uint8`` buffer; each is
+    ``slot_nbytes`` long and handed out as a raw 1-D ``uint8`` tensor.
     Args:
-        buffer_layout: ``{name: (shape, dtype)}`` for each buffer.
-        capacity: Number of buffers to pre-allocate.
+        slot_nbytes: Byte size of each slot.
+        capacity: Number of slots to pre-allocate.
         device: Device for allocation.
-        reuse_barrier: Called with the pending event before a buffer is reused.
+        reuse_barrier: Called with the pending event before a slot is reused.
         pin_memory: Pin buffers (for async H2D copies from CPU).
     """
 
     def __init__(
         self,
-        buffer_layout: TensorLayout,
+        slot_nbytes: int,
         capacity: int,
         device: torch.device,
         reuse_barrier: Callable[[torch.cuda.Event], None],
         pin_memory: bool = False,
     ) -> None:
-        self._buffer_layout = buffer_layout
+        self._slot_nbytes = slot_nbytes
         self._capacity = capacity
-        self._free: deque[dict[str, torch.Tensor]] = deque()
+        self._free: deque[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)
+        buffer = utils.alloc_buffer(max(slot_nbytes * capacity, 1), device, pin_memory)
         for slot in range(capacity):
-            self._free.append({name: all_views[_make_key(slot, name)] for name in buffer_layout})
+            self._free.append(buffer[slot * slot_nbytes : (slot + 1) * slot_nbytes])
 
     @property
     def capacity(self) -> int:
         return self._capacity
 
     @property
-    def buffer_layout(self) -> TensorLayout:
-        return self._buffer_layout
+    def slot_nbytes(self) -> int:
+        return self._slot_nbytes
 
-    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)
+    def acquire(self) -> torch.Tensor:
+        """Take a free raw slot, waiting any pending event before returning.
+        Raises :class:`RuntimeError` if every slot is currently in use.
+        """
+        if not self._free:
+            raise RuntimeError(f"BufferPool exhausted: all {self._capacity} buffers are in use")
+        buffer = self._free.popleft()
+        event = self._events.pop(id(buffer), None)
         if event is not None:
             self._reuse_barrier(event)
-        return weights
+        return buffer
 
-    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.
+    def release(self, buffer: torch.Tensor, event: torch.cuda.Event | None = None) -> None:
+        """Return a raw slot to the free list.
+        The *buffer* must be the exact tensor object returned by :meth:`acquire`
+        (reuse is keyed on its identity). If *event* is given it is waited on the
+        next :meth:`acquire` of this slot, ensuring the prior operation finished.
         """
         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}"
+            self._events[id(buffer)] = event
+        self._free.append(buffer)

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

@@ -3,37 +3,28 @@
 from __future__ import annotations
 
 from collections import OrderedDict
+from typing import NamedTuple
 
 import torch
 
 from ltx_core.block_streaming.disk import LoraSource
-from ltx_core.block_streaming.pool import WeightPool
+from ltx_core.block_streaming.pool import BufferPool
 from ltx_core.block_streaming.source import WeightSource
+from ltx_core.block_streaming.utils import carve_buffer, layout_nbytes
 from ltx_core.loader.fuse_loras import FuseRule, aggregate_lora_products, bf16_fuse_rule
 from ltx_core.loader.primitives import StateDict
 
 _EMPTY_STATE_DICT = StateDict(sd={}, device=torch.device("cpu"), size=0, dtype=set())
 
 
-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 CachedBlock(NamedTuple):
+    """A cached GPU block: the raw pool slot plus the carved per-key views.
+    The raw slot is what is returned to the pool on eviction; the views are
+    what callers consume.
+    """
+
+    raw: torch.Tensor
+    views: dict[str, torch.Tensor]
 
 
 class WeightsProvider:
@@ -51,7 +42,7 @@ class WeightsProvider:
 
     def __init__(
         self,
-        pool: WeightPool,
+        pool: BufferPool,
         copy_stream: torch.cuda.Stream,
         target_device: torch.device,
         source: WeightSource,
@@ -61,7 +52,7 @@ class WeightsProvider:
     ) -> None:
         self._copy_stream = copy_stream
         self._pool = pool
-        self._cache: OrderedDict[int, dict[str, torch.Tensor]] = OrderedDict()
+        self._cache: OrderedDict[int, CachedBlock] = OrderedDict()
         self._events: dict[int, torch.cuda.Event] = {}
         self._target_device = target_device
         self._source = source
@@ -72,40 +63,45 @@ class WeightsProvider:
     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]
+            return self._cache[idx].views
 
         # 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))
+            evicted_idx, evicted = self._cache.popitem(last=False)
+            self._pool.release(evicted.raw, event=self._events.pop(evicted_idx, None))
 
-        gpu_weights = self._pool.acquire()
-        cpu_weights = self._source.get(idx)
+        layout = self._source.block_layout(idx)
+        raw = self._pool.acquire()
+        gpu_weights = carve_buffer(raw, layout)
+        cpu_buffer = self._source.get(idx)
 
-        h2d_event = self._copy_to_gpu(idx, gpu_weights, cpu_weights)
+        h2d_event = self._copy_to_gpu(idx, raw, gpu_weights, cpu_buffer, layout_nbytes(layout))
         self._source.release(idx, event=h2d_event)
 
-        self._cache[idx] = gpu_weights
+        self._cache[idx] = CachedBlock(raw, gpu_weights)
         return gpu_weights
 
     def _copy_to_gpu(
         self,
         idx: int,
+        raw: torch.Tensor,
         gpu_weights: dict[str, torch.Tensor],
-        cpu_weights: dict[str, torch.Tensor],
+        cpu_buffer: torch.Tensor,
+        nbytes: int,
     ) -> 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.
+        *cpu_buffer* is one contiguous source buffer carved by the same layout as
+        *raw*, so a single byte copy of its leading *nbytes* reproduces every view
+        in *gpu_weights*. The wait is intentionally inside this method so callers --
+        and instrumentation regions wrapping it -- observe the full transfer time.
         """
+        if not cpu_buffer.is_contiguous() or cpu_buffer.dtype != torch.uint8 or cpu_buffer.numel() < nbytes:
+            raise ValueError(
+                f"source buffer for block {idx} must be a contiguous uint8 buffer of >= {nbytes} bytes, "
+                f"got {cpu_buffer.dim()}-D {cpu_buffer.dtype} with {cpu_buffer.numel()} elements"
+            )
         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)
+            raw[:nbytes].copy_(cpu_buffer[:nbytes], non_blocking=True)
             if self._lora_sources:
                 self._fuse_block_loras(idx, gpu_weights)
             h2d_event = torch.cuda.Event()

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

@@ -2,31 +2,38 @@
 
 from __future__ import annotations
 
-from collections import OrderedDict
-from typing import Protocol
+from typing import NamedTuple, Protocol
 
 import torch
 
-from ltx_core.block_streaming.disk import DiskBlockReader
-from ltx_core.block_streaming.pool import WeightPool
+from ltx_core.block_streaming.block_fetcher import BlockFetcher, FetchHandle
+from ltx_core.block_streaming.pool import BufferPool
+from ltx_core.block_streaming.utils import carve_buffer, layout_nbytes
 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.
+    Blocks may be heterogeneous: each has its own layout, so the source exposes a
+    per-block layout and the byte size of the largest block (which sizes the
+    pool slots -- a smaller block is carved into the front of a max-sized slot).
+    The source is the single source of truth for each block's layout.
     """
 
+    def block_layout(self, idx: int) -> TensorLayout:
+        """Per-block buffer layout (shape + dtype for each param)."""
+        ...
+
     @property
-    def block_layout(self) -> TensorLayout:
-        """Shared per-block buffer layout (shape + dtype for each param)."""
+    def slot_nbytes(self) -> int:
+        """Byte size of the largest block; sizes a pool slot (16-byte aligned)."""
         ...
 
-    def get(self, idx: int) -> dict[str, torch.Tensor]:
-        """Return CPU weights for block *idx*."""
+    def get(self, idx: int) -> torch.Tensor:
+        """Return one contiguous CPU buffer for block *idx*."""
         ...
 
-    def release(self, idx: int, event: torch.cuda.Event) -> None:
+    def release(self, idx: int, event: torch.cuda.Event | None) -> None:
         """Signal that an async operation using these weights is guarded by *event*."""
         ...
 
@@ -35,68 +42,133 @@ class WeightSource(Protocol):
         ...
 
 
-class DiskWeightSource(WeightSource):
-    """Reads block weights from disk into pinned CPU buffers on demand."""
+class _Scheduled(NamedTuple):
+    """A scheduled (possibly in-flight) read: the raw pool slot and its fetch handle."""
+
+    raw: torch.Tensor
+    status: FetchHandle
+
+
+class DiskWeightSource(WeightSource):
+    """WeightSource that streams blocks from disk via a :class:`BlockFetcher`.
+    ``get(idx)`` must be paired with a ``release(idx)`` before *idx* is fetched
+    again; getting a block that is already in flight raises. Each read acquires a
+    raw pool slot, carves it to block *idx*'s layout, and hands the carved views to
+    the fetcher to fill, so one max-sized slot can serve blocks of differing shapes.
+    ``get`` returns that same contiguous slot.
+    """
+
+    def __init__(
+        self,
+        pool: BufferPool,
+        fetcher: BlockFetcher,
+        block_layouts: dict[int, TensorLayout],
+        blocks_number: int,
+        prefetch_depth: int = 0,
+    ) -> None:
+        if blocks_number <= 0:
+            raise ValueError(f"blocks_number must be > 0, got {blocks_number}")
+        if prefetch_depth < 0:
+            raise ValueError(f"prefetch_depth must be >= 0, got {prefetch_depth}")
+        max_layout_nbytes = max((layout_nbytes(layout) for layout in block_layouts.values()), default=0)
+        if pool.slot_nbytes < max_layout_nbytes:
+            raise ValueError(
+                f"pool slot is too small for the largest block: slot {pool.slot_nbytes} bytes < {max_layout_nbytes}"
+            )
 
-    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
+        self._blocks_number = blocks_number
+        self._prefetch_depth = prefetch_depth
+        self._fetcher = fetcher
+        self._block_layouts = block_layouts
+        self._scheduled: dict[int, _Scheduled] = {}
+        self._in_flight: dict[int, torch.Tensor] = {}
+
+    def block_layout(self, idx: int) -> TensorLayout:
+        return self._block_layouts[idx]
 
     @property
-    def block_layout(self) -> TensorLayout:
-        return self._pool.buffer_layout
+    def slot_nbytes(self) -> int:
+        return self._pool.slot_nbytes
 
-    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]
+    def get(self, idx: int) -> torch.Tensor:
+        if idx in self._in_flight:
+            raise RuntimeError(f"Block {idx} is already in flight; release it before getting it again")
 
-        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))
+        scheduled = self._scheduled.pop(idx, None)
+        if scheduled is None:
+            scheduled = self._schedule(idx)
+        error = scheduled.status.wait()
+        if error is not None:
+            self._pool.release(scheduled.raw)
+            raise error
 
-        weights = self._pool.acquire()
-        self._reader.read_into(weights, idx)
-        self._cache[idx] = weights
-        return weights
+        self._in_flight[idx] = scheduled.raw
+        for k in range(1, self._prefetch_depth + 1):
+            self._ensure_scheduled((idx + k) % self._blocks_number)
+        return scheduled.raw
 
-    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 release(self, idx: int, event: torch.cuda.Event | None) -> None:
+        raw_buffer = self._in_flight.pop(idx)
+        self._pool.release(raw_buffer, event=event)
 
     def cleanup(self) -> None:
-        """Clear cache and close the disk reader."""
-        self._cache.clear()
-        self._events.clear()
-        self._reader.cleanup()
+        self._fetcher.cleanup()
+        while self._in_flight:
+            _, raw_buffer = self._in_flight.popitem()
+            self._pool.release(raw_buffer)
+        while self._scheduled:
+            _, scheduled = self._scheduled.popitem()
+            scheduled.status.wait()
+            self._pool.release(scheduled.raw)
 
-    def __len__(self) -> int:
-        return len(self._cache)
+    def _ensure_scheduled(self, idx: int) -> None:
+        """Schedule a read for *idx* if one is not already pending."""
+        if idx not in self._scheduled:
+            self._scheduled[idx] = self._schedule(idx)
+
+    def _schedule(self, idx: int) -> _Scheduled:
+        """Acquire a raw slot, carve it to block *idx*, enqueue a read, return the handle.
+        The raw slot and its fetch status are returned together so the caller can
+        track both as one unit; the fetcher only receives the carved views to fill.
+        """
+        raw_buffer = self._pool.acquire()
+        carved = carve_buffer(raw_buffer, self._block_layouts[idx])
+        status = self._fetcher.submit(idx, carved)
+        return _Scheduled(raw_buffer, status)
+
+
+class PinnedBlock(NamedTuple):
+    """A pre-loaded pinned block: its single contiguous buffer and the layout to carve it with."""
+
+    buffer: torch.Tensor
+    layout: TensorLayout
 
 
 class PinnedWeightSource(WeightSource):
-    """Pre-loaded pinned CPU weights."""
+    """Pre-loaded pinned CPU weights, one contiguous (possibly heterogeneous) buffer per block."""
 
-    def __init__(self, weights: dict[int, dict[str, torch.Tensor]]) -> None:
-        if not weights:
+    def __init__(self, blocks: dict[int, PinnedBlock]) -> None:
+        if not blocks:
             raise ValueError("PinnedWeightSource requires at least one block")
-        self._weights = weights
+        self._blocks = blocks
+        self._slot_nbytes = max(layout_nbytes(block.layout) for block in blocks.values())
+
+    def block_layout(self, idx: int) -> TensorLayout:
+        return self._blocks[idx].layout
 
     @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 slot_nbytes(self) -> int:
+        return self._slot_nbytes
 
-    def get(self, idx: int) -> dict[str, torch.Tensor]:
-        return self._weights[idx]
+    def get(self, idx: int) -> torch.Tensor:
+        return self._blocks[idx].buffer
 
-    def release(self, idx: int, event: torch.cuda.Event) -> None:
+    def release(self, idx: int, event: torch.cuda.Event | None) -> None:
         pass
 
     def cleanup(self) -> None:
-        self._weights.clear()
+        self._blocks.clear()
 
     def __len__(self) -> int:
-        return len(self._weights)
+        return len(self._blocks)

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

@@ -5,7 +5,7 @@ from __future__ import annotations
 import math
 import weakref
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, NamedTuple
 
 import torch
 from torch import nn
@@ -84,7 +84,7 @@ def _alloc_pinned_exact(nbytes: int) -> torch.Tensor | None:
     return buf
 
 
-def _alloc_buffer(nbytes: int, device: torch.device | None, pin_memory: bool) -> torch.Tensor:
+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
@@ -112,6 +112,47 @@ class _TensorSlice:
         return math.prod(self.shape) * self.dtype.itemsize
 
 
+class LayoutSlices(NamedTuple):
+    """Per-key tensor slices of a layout plus the total aligned buffer size."""
+
+    slices: dict[str, _TensorSlice]
+    nbytes: int
+
+
+def _layout_slices(layout: TensorLayout) -> LayoutSlices:
+    """Compute the byte offset of each key in *layout* and the total aligned size.
+    The size is at least one byte so empty layouts still produce a valid buffer.
+    """
+    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()
+    return LayoutSlices(slices, max(_align_up(cursor, _BUFFER_ALIGN), 1))
+
+
+def layout_nbytes(layout: TensorLayout) -> int:
+    """Byte size of one contiguous, 16-byte-aligned buffer holding *layout* (>= 1)."""
+    return _layout_slices(layout).nbytes
+
+
+def carve_buffer(buffer: torch.Tensor, layout: TensorLayout) -> dict[str, torch.Tensor]:
+    """Carve per-key tensor views for *layout* into the front of *buffer*.
+    *buffer* is a 1-D ``uint8`` tensor at least :func:`layout_nbytes` long. Each
+    returned tensor is a non-overlapping slice of its leading bytes reinterpreted
+    at the requested shape and dtype; any trailing bytes are left unused. That
+    slack is what lets one max-sized pool slot hold a smaller (heterogeneous)
+    block. The views keep *buffer*'s storage alive via PyTorch refcounting.
+    """
+    if buffer.dtype != torch.uint8 or buffer.dim() != 1:
+        raise ValueError(f"carve_buffer expects a 1-D uint8 buffer, got {buffer.dim()}-D {buffer.dtype}")
+    slices, nbytes = _layout_slices(layout)
+    if buffer.numel() < nbytes:
+        raise ValueError(f"buffer too small to carve layout: need {nbytes} bytes, got {buffer.numel()}")
+    return {key: buffer[s.offset : s.offset + s.size()].view(s.dtype).view(s.shape) for key, s in slices.items()}
+
+
 def allocate_layout_views(
     layout: TensorLayout,
     device: torch.device | None = None,
@@ -123,12 +164,5 @@ def allocate_layout_views(
     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()}
+    buffer = alloc_buffer(layout_nbytes(layout), device, pin_memory)
+    return carve_buffer(buffer, layout)

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

@@ -253,6 +253,11 @@ class MultiModalGuider:
         and as scale * (cond - uncond) for stg, steering the denoising process away from the unconditioned
         prediction.
         """
+        dtype = cond.dtype
+        cond = cond.float()
+        uncond_text = uncond_text.float() if isinstance(uncond_text, torch.Tensor) else uncond_text
+        uncond_perturbed = uncond_perturbed.float() if isinstance(uncond_perturbed, torch.Tensor) else uncond_perturbed
+        uncond_modality = uncond_modality.float() if isinstance(uncond_modality, torch.Tensor) else uncond_modality
         pred = (
             cond
             + (self.params.cfg_scale - 1) * (cond - uncond_text)
@@ -265,7 +270,7 @@ class MultiModalGuider:
             factor = self.params.rescale_scale * factor + (1 - self.params.rescale_scale)
             pred = pred * factor
 
-        return pred
+        return pred.to(dtype)
 
     def do_unconditional_generation(self) -> bool:
         """Returns True if the guider is doing unconditional generation."""

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

@@ -17,18 +17,20 @@ class GaussianNoiser(Noiser):
 
     def __init__(self, generator: torch.Generator):
         super().__init__()
-
         self.generator = generator
 
-    def __call__(self, latent_state: LatentState, noise_scale: float = 1.0) -> LatentState:
-        noise = torch.randn(
+    def _sample_noise(self, latent_state: LatentState) -> torch.Tensor:
+        return torch.randn(
             *latent_state.latent.shape,
             device=latent_state.latent.device,
             dtype=latent_state.latent.dtype,
             generator=self.generator,
         )
-        scaled_mask = latent_state.denoise_mask * noise_scale
-        latent = noise * scaled_mask + latent_state.latent * (1 - scaled_mask)
+
+    def __call__(self, latent_state: LatentState, noise_scale: float = 1.0) -> LatentState:
+        noise = self._sample_noise(latent_state)
+        latent = torch.lerp(latent_state.latent.float(), noise.float(), noise_scale)
+        latent = torch.lerp(latent_state.clean_latent.float(), latent, latent_state.denoise_mask)
         return replace(
             latent_state,
             latent=latent.to(latent_state.latent.dtype),

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

@@ -7,6 +7,7 @@ from ltx_core.conditioning.types import (
     ConditioningItemAttentionStrengthWrapper,
     VideoConditionByKeyframeIndex,
     VideoConditionByLatentIndex,
+    VideoConditionByMask,
     VideoConditionByReferenceLatent,
 )
 
@@ -17,5 +18,6 @@ __all__ = [
     "ConditioningItemAttentionStrengthWrapper",
     "VideoConditionByKeyframeIndex",
     "VideoConditionByLatentIndex",
+    "VideoConditionByMask",
     "VideoConditionByReferenceLatent",
 ]

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

@@ -3,6 +3,7 @@
 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.mask_cond import VideoConditionByMask
 from ltx_core.conditioning.types.reference_audio_cond import AudioConditionByReferenceLatent
 from ltx_core.conditioning.types.reference_video_cond import VideoConditionByReferenceLatent
 
@@ -11,5 +12,6 @@ __all__ = [
     "ConditioningItemAttentionStrengthWrapper",
     "VideoConditionByKeyframeIndex",
     "VideoConditionByLatentIndex",
+    "VideoConditionByMask",
     "VideoConditionByReferenceLatent",
 ]

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

@@ -10,8 +10,9 @@ from ltx_core.types import LatentState, VideoLatentShape
 class VideoConditionByKeyframeIndex(ConditioningItem):
     """
     Conditions video generation on keyframe latents at a specific frame index.
-    Appends keyframe tokens to the latent state with positions offset by frame_idx,
-    and sets denoise strength according to the strength parameter.
+    Appends keyframe tokens to the sequence with positions offset by frame_idx: the keyframe
+    latents become clean-latent tokens (placeholder zeros in the noisy latent) and the denoise
+    mask is set from the strength parameter.
     To add attention masking, wrap with :class:`ConditioningItemAttentionStrengthWrapper`.
     Args:
         keyframes: Keyframe latents [B, C, F, H, W].
@@ -75,7 +76,7 @@ class VideoConditionByKeyframeIndex(ConditioningItem):
         )
 
         return LatentState(
-            latent=torch.cat([latent_state.latent, tokens], dim=1),
+            latent=torch.cat([latent_state.latent, torch.zeros_like(tokens)], dim=1),
             denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
             positions=torch.cat([latent_state.positions, positions], dim=2),
             clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),

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

@@ -9,8 +9,8 @@ from ltx_core.types import LatentState
 class VideoConditionByLatentIndex(ConditioningItem):
     """
     Conditions video generation by injecting latents at a specific latent frame index.
-    Replaces tokens in the latent state at positions corresponding to latent_idx,
-    and sets denoise strength according to the strength parameter.
+    Sets the clean latents at positions corresponding to latent_idx to the injected latents,
+    sets denoise strength according to the strength parameter.
     """
 
     def __init__(self, latent: torch.Tensor, strength: float, latent_idx: int):
@@ -37,7 +37,6 @@ class VideoConditionByLatentIndex(ConditioningItem):
 
         latent_state = latent_state.clone()
 
-        latent_state.latent[:, start_token:stop_token] = tokens
         latent_state.clean_latent[:, start_token:stop_token] = tokens
         latent_state.denoise_mask[:, start_token:stop_token] = 1.0 - self.strength
 

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

@@ -0,0 +1,49 @@
+"""Mask-based conditioning for inpainting and spatial conditioning."""
+
+from dataclasses import replace
+
+import torch
+
+from ltx_core.conditioning.item import ConditioningItem
+from ltx_core.tools import LatentTools
+from ltx_core.types import LatentState
+
+
+class VideoConditionByMask(ConditioningItem):
+    """Condition video generation using a binary mask over latent frames.
+    Masked positions (mask=1) receive the provided clean latent values and are
+    excluded from denoising (denoise_mask set to ``1 - strength``). Unmasked
+    positions (mask=0) are left unchanged and denoised normally.
+    The mask operates in **unpatchified latent** space — it should have shape
+    ``[B, F, H, W]`` matching the latent dimensions (after VAE encoding,
+    before patchification). This is consistent with the latent input format
+    used by all other conditioning items.
+    Args:
+        latent: Clean conditioning latents in unpatchified format [B, C, F, H, W].
+            Must match the target shape of the latent tools.
+        mask: Binary mask [B, F, H, W] in unpatchified latent space.
+            1 = conditioning position (clean, excluded from denoising),
+            0 = generated position (noised, denoised normally).
+        strength: Conditioning strength for masked positions. 1.0 = fully clean
+            (no denoising), 0.0 = no conditioning effect. Default 1.0.
+    """
+
+    def __init__(self, latent: torch.Tensor, mask: torch.Tensor, strength: float = 1.0):
+        self.latent = latent
+        self.mask = mask
+        self.strength = strength
+
+    def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
+        """Apply mask-based conditioning to the latent state."""
+        tokens = latent_tools.patchifier.patchify(self.latent)
+
+        mask = latent_tools.patchifier.patchify(self.mask.unsqueeze(1))
+
+        m = mask.to(dtype=latent_state.latent.dtype)
+        inv = 1 - m
+
+        return replace(
+            latent_state,
+            clean_latent=latent_state.clean_latent * inv + tokens * m,
+            denoise_mask=latent_state.denoise_mask * inv + (1.0 - self.strength) * m,
+        )

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

@@ -51,7 +51,7 @@ class AudioConditionByReferenceLatent:
         )
 
         return LatentState(
-            latent=torch.cat([latent_state.latent, tokens], dim=1),
+            latent=torch.cat([latent_state.latent, torch.zeros_like(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),

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

@@ -14,7 +14,8 @@ class VideoConditionByReferenceLatent(ConditioningItem):
     Conditions video generation on a reference video latent for IC-LoRA inference.
     IC-LoRAs are trained by concatenating reference (control signal) and target tokens,
     learning to attend across both. This class replicates that setup at inference by
-    appending reference tokens to the latent sequence.
+    appending the reference tokens to the sequence as clean latents (with placeholder zeros
+    in the noisy latent).
     IC-LoRAs can be trained with lower-resolution references than the target (e.g., 384px
     reference for 768px output) for efficiency and better generalization. The
     `downscale_factor` scales reference positions to match target coordinates, preserving
@@ -22,9 +23,9 @@ class VideoConditionByReferenceLatent(ConditioningItem):
     (stored in LoRA metadata).
     To add attention masking, wrap with :class:`ConditioningItemAttentionStrengthWrapper`.
     Args:
-        latent: Reference video latents [B, C, F, H, W]
-        downscale_factor: Target/reference resolution ratio (e.g., 2 = half-resolution
-            reference). Spatial positions are scaled by this factor.
+        latent: Reference video latents [B, C, F, H, W].
+        downscale_factor: Target/reference spatial ratio (e.g. 2 = half-res ref).
+        temporal_scale_factor: Target/reference temporal ratio S (e.g. 4 = ref at 1/4 fps).
         strength: Conditioning strength. 1.0 = full (reference kept clean),
             0.0 = none (reference denoised). Default 1.0.
     """
@@ -33,10 +34,12 @@ class VideoConditionByReferenceLatent(ConditioningItem):
         self,
         latent: torch.Tensor,
         downscale_factor: int = 1,
+        temporal_scale_factor: int = 1,
         strength: float = 1.0,
     ):
         self.latent = latent
         self.downscale_factor = downscale_factor
+        self.temporal_scale_factor = temporal_scale_factor
         self.strength = strength
 
     def apply_to(
@@ -44,10 +47,9 @@ class VideoConditionByReferenceLatent(ConditioningItem):
         latent_state: LatentState,
         latent_tools: VideoLatentTools,
     ) -> LatentState:
-        """Append reference video tokens with scaled positions."""
+        """Append reference video tokens with positions translated into the target frame."""
         tokens = latent_tools.patchifier.patchify(self.latent)
 
-        # Compute positions for the reference video's actual dimensions
         latent_coords = latent_tools.patchifier.get_patch_grid_bounds(
             output_shape=VideoLatentShape.from_torch_shape(self.latent.shape),
             device=self.latent.device,
@@ -58,9 +60,18 @@ class VideoConditionByReferenceLatent(ConditioningItem):
             causal_fix=latent_tools.causal_fix,
         )
         positions = positions.to(dtype=torch.float32)
-        positions[:, 0, ...] /= latent_tools.fps
 
-        # Scale spatial positions to match target coordinate space
+        # Place ref tokens on their own time spacing (= target_fps / S).
+        positions[:, 0, ...] /= latent_tools.fps / self.temporal_scale_factor
+
+        # Translate into the target's frame so ref's last patch ends with target's last
+        # patch; clamp the causal patch's negative start back to [0, 1/target_fps).
+        if self.temporal_scale_factor != 1:
+            t_target = latent_state.positions[:, 0, 0:1, 1:2].to(dtype=torch.float32)  # = 1/target_fps
+            positions[:, 0, ...] = torch.clamp(
+                positions[:, 0, ...] - (self.temporal_scale_factor - 1) * t_target,
+                min=0,
+            )
         if self.downscale_factor != 1:
             positions[:, 1, ...] *= self.downscale_factor  # height axis
             positions[:, 2, ...] *= self.downscale_factor  # width axis
@@ -83,7 +94,7 @@ class VideoConditionByReferenceLatent(ConditioningItem):
         )
 
         return LatentState(
-            latent=torch.cat([latent_state.latent, tokens], dim=1),
+            latent=torch.cat([latent_state.latent, torch.zeros_like(tokens)], dim=1),
             denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
             positions=torch.cat([latent_state.positions, positions], dim=2),
             clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),

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

@@ -1,18 +1,21 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, NamedTuple, Protocol
+from typing import TYPE_CHECKING, Any, NamedTuple, Protocol, TypeVar
 
 import torch
 
 from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.sd_ops import SDOps
-from ltx_core.model.model_protocol import ModelType
 
 if TYPE_CHECKING:
+    from typing_extensions import Self
+
     from ltx_core.loader.fuse_loras import FuseRule
     from ltx_core.loader.registry import Registry
 
+BuiltType = TypeVar("BuiltType", covariant=True)  # noqa: PLC0105
+
 
 # Per-key shape and dtype description for a flat collection of tensors.
 TensorLayout = dict[str, tuple[torch.Size, torch.dtype]]
@@ -50,74 +53,75 @@ class StateDictLoader(Protocol):
         """
         Load metadata from path
         """
+        ...
 
     def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict:
         """
         Load state dict from path or paths (for sharded model storage) and apply sd_ops
         """
+        ...
 
 
-class BuilderProtocol(Protocol[ModelType]):
+class BuilderProtocol(Protocol[BuiltType]):
     """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: ...
+        self,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+        **kwargs: Any,  # noqa: ANN401
+    ) -> BuiltType: ...
 
+    @property
+    def registry(self) -> "Registry": ...
 
-class ModelBuilderProtocol(BuilderProtocol[ModelType], Protocol[ModelType]):
-    """
-    Protocol for building PyTorch models from configuration dictionaries.
-    Implementations must provide:
-    - meta_model: Create a model from configuration dictionary and apply module operations
-    - build: Create and initialize a model from state dictionary and apply dtype transformations
-    """
-
-    model_sd_ops: SDOps | None
-    module_ops: tuple[ModuleOps, ...]
-    loras: tuple["LoraPathStrengthAndSDOps", ...]
-    registry: "Registry"
-
-    def meta_model(self, config: dict, module_ops: list[ModuleOps] | None = None) -> ModelType:
-        """
-        Create a model on the meta device from a configuration dictionary.
-        This decouples model creation from weight loading, allowing the model
-        architecture to be instantiated without allocating memory for parameters.
-        Args:
-            config: Model configuration dictionary.
-            module_ops: Optional list of module operations to apply (e.g., quantization).
-        Returns:
-            Model instance on meta device (no actual memory allocated for parameters).
-        """
-        ...
-
-    def with_sd_ops(self, sd_ops: SDOps | None) -> "ModelBuilderProtocol[ModelType]":
-        """Return a copy of this builder with the given state-dict key remapping ops."""
-        ...
-
-    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> "ModelBuilderProtocol[ModelType]":
-        """Return a copy of this builder with the given module operations (e.g. quantization)."""
-        ...
-
-    def with_loras(self, loras: tuple["LoraPathStrengthAndSDOps", ...]) -> "ModelBuilderProtocol[ModelType]":
-        """Return a copy of this builder with the given LoRAs to fuse at build time."""
-        ...
-
-    def with_registry(self, registry: "Registry") -> "ModelBuilderProtocol[ModelType]":
+    def with_registry(self, registry: "Registry") -> "Self":
         """Return a copy of this builder using the given weight registry for allocation."""
         ...
 
-    def with_lora_load_device(self, device: torch.device) -> "ModelBuilderProtocol[ModelType]":
+
+class ModelBuilderProtocol(BuilderProtocol[BuiltType], Protocol[BuiltType]):
+    """
+    Protocol for building PyTorch models from configuration dictionaries.
+    Implementations must provide:
+    - build: Create and initialize a model from state dictionary and apply dtype transformations
+    """
+
+    @property
+    def model_sd_ops(self) -> SDOps | None: ...
+
+    @property
+    def module_ops(self) -> tuple[ModuleOps, ...]: ...
+
+    @property
+    def loras(self) -> tuple["LoraPathStrengthAndSDOps", ...]: ...
+
+    def with_sd_ops(self, sd_ops: SDOps | None) -> "Self":
+        """Return a copy of this builder with the given state-dict key remapping ops."""
+        ...
+
+    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> "Self":
+        """Return a copy of this builder with the given module operations (e.g. quantization)."""
+        ...
+
+    def with_loras(self, loras: tuple["LoraPathStrengthAndSDOps", ...]) -> "Self":
+        """Return a copy of this builder with the given LoRAs to fuse at build time."""
+        ...
+
+    def with_lora_load_device(self, device: torch.device) -> "Self":
         """Return a copy of this builder that loads LoRA weights onto the given device."""
         ...
 
-    def with_fuse_rule(self, fuse_rule: "FuseRule") -> "ModelBuilderProtocol[ModelType]":
+    def with_fuse_rule(self, fuse_rule: "FuseRule") -> "Self":
         """Return a copy of this builder with the given LoRA fuse rule (e.g. from a quantization policy)."""
         ...
 
     def build(
-        self, device: torch.device | None = None, dtype: torch.dtype | None = None, **kwargs: object
-    ) -> ModelType:
+        self,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+        **kwargs: Any,  # noqa: ANN401
+    ) -> BuiltType:
         """
         Build the model
         Args:
@@ -140,8 +144,7 @@ class LoRAAdaptableProtocol(Protocol):
     - lora: Add a LoRA to the model
     """
 
-    def lora(self, lora_path: str, strength: float) -> "LoRAAdaptableProtocol":
-        pass
+    def lora(self, lora_path: str, strength: float, sd_ops: SDOps) -> "LoRAAdaptableProtocol": ...
 
 
 class LoraPathStrengthAndSDOps(NamedTuple):

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

@@ -24,6 +24,7 @@ class ContentMatching:
 
     prefix: str = ""
     suffix: str = ""
+    contains: str = ""
 
 
 class KeyValueOperationResult(NamedTuple):
@@ -72,10 +73,10 @@ class SDOps:
         new_mapping = (*self.mapping, ContentReplacement(content, replacement))
         return replace(self, mapping=new_mapping)
 
-    def with_matching(self, prefix: str = "", suffix: str = "") -> "SDOps":
-        """Create a new SDOps instance with the specified prefix and suffix matching added to the mapping."""
+    def with_matching(self, prefix: str = "", suffix: str = "", contains: str = "") -> "SDOps":
+        """Create a new SDOps instance with the specified prefix, suffix and contains matching added to the mapping."""
 
-        new_mapping = (*self.mapping, ContentMatching(prefix, suffix))
+        new_mapping = (*self.mapping, ContentMatching(prefix, suffix, contains))
         return replace(self, mapping=new_mapping)
 
     def with_additional_allowed_keys(self, keys: frozenset[str]) -> "SDOps":
@@ -100,7 +101,12 @@ class SDOps:
     def apply_to_key(self, key: str) -> str | None:
         """Apply the mapping to the given name."""
         matchers = [content for content in self.mapping if isinstance(content, ContentMatching)]
-        valid = any(key.startswith(f.prefix) and key.endswith(f.suffix) for f in matchers)
+        valid = any(
+            key.startswith(matcher.prefix)
+            and key.endswith(matcher.suffix)
+            and (not matcher.contains or matcher.contains in key)
+            for matcher in matchers
+        )
         if not valid:
             return None
 

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

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
+import copy
 import logging
-from dataclasses import dataclass, field, replace
-from typing import Generic
+from typing import TYPE_CHECKING, Final, Generic
 
 import torch
 from torch import nn
@@ -21,6 +23,9 @@ 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
 
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 
@@ -78,10 +83,12 @@ def _load_model_weights(
     meta_model.load_state_dict(fused_sd, strict=False, assign=True)
 
 
-@dataclass(frozen=True)
 class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType], LoRAAdaptableProtocol):
     """
     Builder for PyTorch models residing on a single GPU.
+    The builder is immutable: ``with_*``/``lora`` return modified copies. The
+    ``ModelBuilderProtocol`` surface is exposed via read-only properties backed
+    by private attributes.
     Attributes:
         model_class_configurator: Class responsible for constructing the model from a config dict.
         model_path: Path (or tuple of shard paths) to the model's `.safetensors` checkpoint(s).
@@ -98,54 +105,106 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
         fuse_rule: Per-policy LoRA merge rule. Defaults to ``bf16_fuse_rule``;
     """
 
-    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)
-    lora_load_device: torch.device = field(default_factory=lambda: torch.device("cpu"))
-    fuse_rule: FuseRule = bf16_fuse_rule
+    def __init__(
+        self,
+        model_class_configurator: type[ModelConfigurator[ModelType]],
+        model_path: str | tuple[str, ...],
+        model_sd_ops: SDOps | None = None,
+        module_ops: tuple[ModuleOps, ...] = (),
+        loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
+        model_loader: StateDictLoader | None = None,
+        registry: Registry | None = None,
+        lora_load_device: torch.device | None = None,
+        fuse_rule: FuseRule = bf16_fuse_rule,
+    ) -> None:
+        # Read-only: typed with the covariant ModelType, so it must not be a mutable attribute.
+        self._model_class_configurator: Final = model_class_configurator
+        self._model_path = model_path
+        self._model_sd_ops = model_sd_ops
+        self._module_ops = module_ops
+        self._loras = loras
+        self._model_loader = model_loader if model_loader is not None else SafetensorsModelStateDictLoader()
+        self._registry = registry if registry is not None else DummyRegistry()
+        self._lora_load_device = lora_load_device if lora_load_device is not None else torch.device("cpu")
+        self._fuse_rule = fuse_rule
 
-    def lora(self, lora_path: str, strength: float, sd_ops: SDOps) -> "SingleGPUModelBuilder":
-        return replace(self, loras=(*self.loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops)))
+    @property
+    def model_sd_ops(self) -> SDOps | None:
+        return self._model_sd_ops
 
-    def with_sd_ops(self, sd_ops: SDOps | None) -> "SingleGPUModelBuilder":
-        return replace(self, model_sd_ops=sd_ops)
+    @property
+    def module_ops(self) -> tuple[ModuleOps, ...]:
+        return self._module_ops
 
-    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> "SingleGPUModelBuilder":
-        return replace(self, module_ops=module_ops)
+    @property
+    def loras(self) -> tuple[LoraPathStrengthAndSDOps, ...]:
+        return self._loras
 
-    def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "SingleGPUModelBuilder":
-        return replace(self, loras=loras)
+    @property
+    def registry(self) -> Registry:
+        return self._registry
 
-    def with_registry(self, registry: Registry) -> "SingleGPUModelBuilder":
-        return replace(self, registry=registry)
+    @property
+    def model_path(self) -> str | tuple[str, ...]:
+        return self._model_path
 
-    def with_lora_load_device(self, device: torch.device) -> "SingleGPUModelBuilder":
-        return replace(self, lora_load_device=device)
+    @property
+    def model_loader(self) -> StateDictLoader:
+        return self._model_loader
 
-    def with_fuse_rule(self, fuse_rule: FuseRule) -> "SingleGPUModelBuilder":
-        return replace(self, fuse_rule=fuse_rule)
+    @property
+    def lora_load_device(self) -> torch.device:
+        return self._lora_load_device
+
+    @property
+    def fuse_rule(self) -> FuseRule:
+        return self._fuse_rule
+
+    def lora(self, lora_path: str, strength: float, sd_ops: SDOps) -> Self:
+        clone = copy.copy(self)
+        clone._loras = (*self._loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops))
+        return clone
+
+    def with_sd_ops(self, sd_ops: SDOps | None) -> Self:
+        clone = copy.copy(self)
+        clone._model_sd_ops = sd_ops
+        return clone
+
+    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> Self:
+        clone = copy.copy(self)
+        clone._module_ops = module_ops
+        return clone
+
+    def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> Self:
+        clone = copy.copy(self)
+        clone._loras = loras
+        return clone
+
+    def with_registry(self, registry: Registry) -> Self:
+        clone = copy.copy(self)
+        clone._registry = registry
+        return clone
+
+    def with_lora_load_device(self, device: torch.device) -> Self:
+        clone = copy.copy(self)
+        clone._lora_load_device = device
+        return clone
+
+    def with_fuse_rule(self, fuse_rule: FuseRule) -> Self:
+        clone = copy.copy(self)
+        clone._fuse_rule = fuse_rule
+        return clone
 
     def model_config(self) -> dict:
-        return read_model_config(self.model_path, self.model_loader)
+        return read_model_config(self._model_path, self._model_loader)
 
     def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType:
-        return create_meta_model(self.model_class_configurator, config, module_ops)
+        return create_meta_model(self._model_class_configurator, config, module_ops)
 
     def load_sd(
         self, paths: list[str], registry: Registry, device: torch.device | None, sd_ops: SDOps | None = None
     ) -> StateDict:
-        return load_state_dict(paths, self.model_loader, registry, device, sd_ops)
-
-    def _return_model(self, meta_model: ModelType, device: torch.device) -> ModelType:
-        uninitialized = _check_uninitialized(meta_model)
-        if uninitialized:
-            logger.warning(f"Uninitialized parameters or buffers: {uninitialized}")
-            return meta_model
-        return meta_model.to(device)
+        return load_state_dict(paths, self._model_loader, registry, device, sd_ops)
 
     def build(
         self,
@@ -155,18 +214,23 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
     ) -> ModelType:
         device = torch.device("cuda") if device is None else device
         config = self.model_config()
-        meta_model = self.meta_model(config, self.module_ops)
+        meta_model = self.meta_model(config, self._module_ops)
 
         _load_model_weights(
             meta_model=meta_model,
-            model_path=self.model_path,
-            loras=self.loras,
-            loader=self.model_loader,
-            registry=self.registry,
+            model_path=self._model_path,
+            loras=self._loras,
+            loader=self._model_loader,
+            registry=self._registry,
             device=device,
             dtype=dtype,
-            model_sd_ops=self.model_sd_ops,
-            lora_load_device=self.lora_load_device,
-            fuse_rule=self.fuse_rule,
+            model_sd_ops=self._model_sd_ops,
+            lora_load_device=self._lora_load_device,
+            fuse_rule=self._fuse_rule,
         )
-        return self._return_model(meta_model, device)
+
+        uninitialized = _check_uninitialized(meta_model)
+        if uninitialized:
+            logger.warning(f"Uninitialized parameters or buffers: {uninitialized}")
+            return meta_model
+        return meta_model.to(device)

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

@@ -1,6 +1,14 @@
-from typing import Protocol, TypeVar
+from __future__ import annotations
 
-ModelType = TypeVar("ModelType")
+from typing import TYPE_CHECKING, Protocol, TypeVar
+
+import torch
+
+if TYPE_CHECKING:
+    from ltx_core.guidance.perturbations import BatchedPerturbationConfig
+    from ltx_core.model.transformer.modality import Modality
+
+ModelType = TypeVar("ModelType", covariant=True, bound=torch.nn.Module)  # noqa: PLC0105
 
 
 class ModelConfigurator(Protocol[ModelType]):
@@ -8,3 +16,24 @@ class ModelConfigurator(Protocol[ModelType]):
 
     @classmethod
     def from_config(cls, config: dict) -> ModelType: ...
+
+
+class LTXModelProtocol(Protocol):
+    """Velocity-model forward interface shared by ``LTXModel`` and its multi-GPU wrappers.
+    ``forward`` pins the real signature (enforced structurally); ``__call__`` mirrors it
+    so protocol-typed values stay callable via ``model(...)``.
+    """
+
+    def forward(
+        self,
+        video: Modality | None,
+        audio: Modality | None,
+        perturbations: BatchedPerturbationConfig,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]: ...
+
+    def __call__(
+        self,
+        video: Modality | None,
+        audio: Modality | None,
+        perturbations: BatchedPerturbationConfig,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]: ...

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

@@ -3,13 +3,17 @@
 from ltx_core.model.transformer.modality import Modality
 from ltx_core.model.transformer.model import LTXModel, X0Model
 from ltx_core.model.transformer.model_configurator import (
+    LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP,
     LTXV_MODEL_COMFY_RENAMING_MAP,
+    LTXAudioOnlyModelConfigurator,
     LTXModelConfigurator,
     LTXVideoOnlyModelConfigurator,
 )
 
 __all__ = [
+    "LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP",
     "LTXV_MODEL_COMFY_RENAMING_MAP",
+    "LTXAudioOnlyModelConfigurator",
     "LTXModel",
     "LTXModelConfigurator",
     "LTXVideoOnlyModelConfigurator",

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

@@ -1,5 +1,4 @@
 import functools
-import logging
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import Protocol
@@ -15,8 +14,6 @@ from ltx_core.model.transformer.ops import (
 )
 from ltx_core.model.transformer.rope import LTXRopeType
 
-logger = logging.getLogger(__name__)
-
 
 def _torch_default_sdpa_priority() -> list[SDPBackend]:
     """Fetch torch's current default SDPA priority order at runtime.
@@ -76,9 +73,9 @@ class PytorchAttention(AttentionCallable):
 
     @property
     def label(self) -> str:
-        """Human-readable identifier (used in the AUTOMATIC selection log).
-        Encodes the SDPA priority list so a single-backend pin reads differently
-        from the full-priority dispatcher walk."""
+        """Human-readable identifier for this backend. Encodes the SDPA priority
+        list so a single-backend pin reads differently from the full-priority
+        dispatcher walk."""
         return f"SDPA[{'>'.join(b.name for b in self._priority)}]"
 
     def __call__(
@@ -199,10 +196,9 @@ class FlashAttention4(AttentionCallable):
 
 # --- Automatic selection -----------------------------------------------------
 # AUTOMATIC inspects installed extras and the GPU arch and returns the fastest
-# usable callable for each path. The selection runs once per process (cached)
-# and logs the resulting label once. The unmasked and masked picks are
-# independent: each calls its own helper and may end up on different backends
-# (e.g. FA3 unmasked + xFormers masked on H100).
+# usable callable for each path. The selection runs once per process (cached).
+# The unmasked and masked picks are independent: each calls its own helper and
+# may end up on different backends (e.g. FA3 unmasked + xFormers masked on H100).
 
 
 def _sdpa_can_use(backend: SDPBackend, *, with_mask: bool) -> bool:
@@ -289,20 +285,23 @@ def _select_masked_attention() -> MaskedAttentionCallable:
 
 @functools.cache
 def automatic_attention() -> AttentionCallable:
-    """Cached AUTOMATIC pick for the unmasked path. Logs the chosen label once
-    per process."""
-    fn = _select_primary_attention()
-    logger.info("Automatic attention selected: %s", fn.label)
-    return fn
+    """Cached AUTOMATIC pick for the unmasked path.
+    Cached so every ``AttentionOps`` in the process shares one instance."""
+    return _select_primary_attention()
 
 
 @functools.cache
 def automatic_masked_attention() -> MaskedAttentionCallable:
-    """Cached AUTOMATIC pick for the masked path. Logs the chosen label once
-    per process."""
-    fn = _select_masked_attention()
-    logger.info("Automatic masked attention selected: %s", fn.label)
-    return fn
+    """Cached AUTOMATIC pick for the masked path. See :func:`automatic_attention`."""
+    return _select_masked_attention()
+
+
+def attention_label(fn: AttentionCallable | MaskedAttentionCallable) -> str:
+    """Best-effort human-readable backend name.
+    Built-in callables expose ``.label`` (encoding the SDPA priority list for the
+    Pytorch backends); fall back to the class name for custom or wrapped callables
+    (e.g. the multi-GPU All2All wrappers) that don't define one."""
+    return getattr(fn, "label", type(fn).__name__)
 
 
 def _resolve_sdpa_variant(backend: SDPBackend, name: str, *, with_mask: bool) -> PytorchAttention:
@@ -342,8 +341,7 @@ class AttentionFunction(Enum):
         isn't usable on this machine -- missing package or SDPA backend rejected
         on this hardware (e.g. cuDNN under ``torch.use_deterministic_algorithms``).
         Opting in means "this kernel or fail loudly". ``AUTOMATIC`` returns the
-        cached :func:`automatic_attention` instance so the once-per-process log
-        fires only on the first resolution.
+        cached :func:`automatic_attention` instance so every build shares one callable.
         """
         match self:
             case AttentionFunction.AUTOMATIC:

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

@@ -11,7 +11,7 @@ from ltx_core.model.transformer.transformer_args import BlockPerturbationsProces
 
 # Defaults applied inside the patched forward. Overriding via CompilationConfig
 # replaces these wholesale; it does not merge.
-_DEFAULT_INDUCTOR_CONFIG: dict[str, Any] = {"unsafe_skip_cache_dynamic_shape_guards": True}
+_DEFAULT_INDUCTOR_CONFIG: dict[str, Any] = {}
 _DEFAULT_DYNAMO_CONFIG: dict[str, Any] = {"inline_inbuilt_nn_modules": True, "cache_size_limit": 256}
 
 

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

@@ -1,9 +1,12 @@
+import logging
 from enum import Enum
 
 import torch
 
 from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
+from ltx_core.model.model_protocol import LTXModelProtocol
 from ltx_core.model.transformer.adaln import AdaLayerNormSingle, adaln_embedding_coefficient
+from ltx_core.model.transformer.attention import attention_label
 from ltx_core.model.transformer.modality import Modality
 from ltx_core.model.transformer.rope import LTXRopeType
 from ltx_core.model.transformer.transformer import (
@@ -20,6 +23,8 @@ from ltx_core.model.transformer.transformer_args import (
 )
 from ltx_core.utils import to_denoised
 
+logger = logging.getLogger(__name__)
+
 
 class LTXModelType(Enum):
     AudioVideo = "ltx av model"
@@ -70,6 +75,15 @@ class LTXModel(torch.nn.Module):
         cross_attention_adaln: bool = False,
     ):
         super().__init__()
+        # Log the attention backends this transformer is built with. Reading the resolved
+        # ``label`` off the ops reports whatever was selected -- AUTOMATIC, an explicit pin
+        # (PYTORCH/XFORMERS/FA3/FA4/SDPA_*), or a directly supplied callable -- so this is the
+        # single source of truth for which kernel a build uses. Fires once per build.
+        logger.info(
+            "Building transformer with attention backends -- self: %s, masked: %s",
+            attention_label(ops.attention_ops.attention_function),
+            attention_label(ops.attention_ops.masked_attention_function),
+        )
         self._enable_gradient_checkpointing = False
         self.cross_attention_adaln = cross_attention_adaln
         self.use_middle_indices_grid = use_middle_indices_grid
@@ -411,7 +425,7 @@ class LTXModel(torch.nn.Module):
 
     def forward(
         self, video: Modality | None, audio: Modality | None, perturbations: BatchedPerturbationConfig
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         """
         Forward pass for LTX models.
         Returns:
@@ -459,7 +473,7 @@ class LegacyX0Model(torch.nn.Module):
     Returns fully denoised output based on the velocities produced by the base model.
     """
 
-    def __init__(self, velocity_model: LTXModel):
+    def __init__(self, velocity_model: LTXModelProtocol):
         super().__init__()
         self.velocity_model = velocity_model
 
@@ -488,7 +502,7 @@ class X0Model(torch.nn.Module):
     Applies scaled denoising to the video and audio according to the timesteps = sigma * denoising_mask.
     """
 
-    def __init__(self, velocity_model: LTXModel):
+    def __init__(self, velocity_model: LTXModelProtocol):
         super().__init__()
         self.velocity_model = velocity_model
 

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

@@ -123,6 +123,58 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
         )
 
 
+class LTXAudioOnlyModelConfigurator(ModelConfigurator[LTXModel]):
+    """
+    Configurator for LTX audio only model.
+    Builds an audio-only LTX model (``model_type=AudioOnly``) so the video
+    transformer weights are never instantiated or loaded. Useful for
+    text-to-audio inference where the video branch is unused.
+    """
+
+    @classmethod
+    def from_config(cls, config: dict, ops: TransformerOpsConfig = DEFAULT_TRANSFORMER_OPS) -> LTXModel:
+        # Build audio caption projection for 19B models (projection handled in transformer).
+        _, audio_caption_projection = _build_caption_projections(config, is_av=True)
+
+        config = config.get("transformer", {})
+
+        check_config_value(config, "dropout", 0.0)
+        check_config_value(config, "attention_bias", True)
+        check_config_value(config, "num_vector_embeds", None)
+        check_config_value(config, "activation_fn", "gelu-approximate")
+        check_config_value(config, "num_embeds_ada_norm", 1000)
+        check_config_value(config, "use_linear_projection", False)
+        check_config_value(config, "only_cross_attention", False)
+        check_config_value(config, "cross_attention_norm", True)
+        check_config_value(config, "double_self_attention", False)
+        check_config_value(config, "upcast_attention", False)
+        check_config_value(config, "standardization_norm", "rms_norm")
+        check_config_value(config, "norm_elementwise_affine", False)
+        check_config_value(config, "qk_norm", "rms_norm")
+        check_config_value(config, "positional_embedding_type", "rope")
+        check_config_value(config, "use_middle_indices_grid", True)
+
+        return LTXModel(
+            model_type=LTXModelType.AudioOnly,
+            num_layers=config.get("num_layers", 48),
+            norm_eps=config.get("norm_eps", 1e-06),
+            ops=ops,
+            timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
+            use_middle_indices_grid=config.get("use_middle_indices_grid", True),
+            audio_num_attention_heads=config.get("audio_num_attention_heads", 32),
+            audio_attention_head_dim=config.get("audio_attention_head_dim", 64),
+            audio_in_channels=config.get("audio_in_channels", 128),
+            audio_out_channels=config.get("audio_out_channels", 128),
+            audio_cross_attention_dim=config.get("audio_cross_attention_dim", 2048),
+            audio_positional_embedding_max_pos=config.get("audio_positional_embedding_max_pos", [20]),
+            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),
+            audio_caption_projection=audio_caption_projection,
+            cross_attention_adaln=config.get("cross_attention_adaln", False),
+        )
+
+
 def _build_caption_projections(
     config: dict,
     is_av: bool,
@@ -151,3 +203,16 @@ LTXV_MODEL_COMFY_RENAMING_MAP = (
     .with_matching(prefix="model.diffusion_model.")
     .with_replacement("model.diffusion_model.", "")
 )
+
+LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP = (
+    SDOps("LTXV_AUDIO_ONLY_MODEL_COMFY_MAP")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_attn1")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_attn2")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_ff")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_patchify")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_proj_out")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_adaln_single")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_prompt")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_scale_shift_table")
+    .with_replacement("model.diffusion_model.", "")
+)

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

@@ -24,7 +24,6 @@ from ltx_core.model.transformer.ops import (
 )
 from ltx_core.model.transformer.rope import LTXRopeType
 from ltx_core.model.transformer.transformer_args import TransformerArgs
-from ltx_core.utils import rms_norm
 
 
 @dataclass
@@ -222,7 +221,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
 
     def _apply_text_cross_attention(
         self,
-        x: torch.Tensor,
+        x_normed: torch.Tensor,
         context: torch.Tensor,
         attn: AttentionCallable,
         scale_shift_table: torch.Tensor,
@@ -232,11 +231,14 @@ class BasicAVTransformerBlock(torch.nn.Module):
         context_mask: torch.Tensor | None,
         cross_attention_adaln: bool = False,
     ) -> torch.Tensor:
-        """Apply text cross-attention, with optional AdaLN modulation."""
+        """Apply text cross-attention, with optional AdaLN modulation.
+        ``x_normed`` is the RMS-normalized self-attention output produced by
+        ``post_sa_function`` -- this method does not normalize again.
+        """
         if cross_attention_adaln:
-            shift_q, scale_q, gate = self.get_ada_values(scale_shift_table, x.shape[0], timestep, slice(6, 9))
+            shift_q, scale_q, gate = self.get_ada_values(scale_shift_table, x_normed.shape[0], timestep, slice(6, 9))
             return apply_cross_attention_adaln(
-                x,
+                x_normed,
                 context,
                 attn,
                 shift_q,
@@ -245,9 +247,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 prompt_scale_shift_table,
                 prompt_timestep,
                 context_mask,
-                self.norm_eps,
             )
-        return attn(rms_norm(x, eps=self.norm_eps), context=context, mask=context_mask)
+        return attn(x_normed, context=context, mask=context_mask)
 
     def forward(  # noqa: PLR0915
         self,
@@ -280,10 +281,10 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 perturbation_mask=video.self_attn_perturbation_mask,
                 all_perturbed=video.self_attn_all_perturbed,
             )
-            vx = vx + vx_msa_out * vgate_msa
+            vx, vx_normed = self.post_sa_function(vx, vx_msa_out, None, self.norm_eps, vgate_msa)
             del vgate_msa, norm_vx, vx_msa_out
             vx = vx + self._apply_text_cross_attention(
-                vx,
+                vx_normed,
                 video.context,
                 self.attn2,
                 self.scale_shift_table,
@@ -293,6 +294,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 video.context_mask,
                 cross_attention_adaln=self.cross_attention_adaln,
             )
+            del vx_normed
 
         if run_ax:
             ashift_msa, ascale_msa, agate_msa = self.get_ada_values(
@@ -308,10 +310,10 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 perturbation_mask=audio.self_attn_perturbation_mask,
                 all_perturbed=audio.self_attn_all_perturbed,
             )
-            ax = ax + ax_msa_out * agate_msa
+            ax, ax_normed = self.post_sa_function(ax, ax_msa_out, None, self.norm_eps, agate_msa)
             del agate_msa, norm_ax, ax_msa_out
             ax = ax + self._apply_text_cross_attention(
-                ax,
+                ax_normed,
                 audio.context,
                 self.audio_attn2,
                 self.audio_scale_shift_table,
@@ -321,6 +323,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 audio.context_mask,
                 cross_attention_adaln=self.cross_attention_adaln,
             )
+            del ax_normed
 
         # Audio - Video cross attention.
         if run_a2v or run_v2a:
@@ -414,7 +417,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
 
 
 def apply_cross_attention_adaln(
-    x: torch.Tensor,
+    x_normed: torch.Tensor,
     context: torch.Tensor,
     attn: AttentionCallable,
     q_shift: torch.Tensor,
@@ -423,13 +426,17 @@ def apply_cross_attention_adaln(
     prompt_scale_shift_table: torch.Tensor,
     prompt_timestep: torch.Tensor,
     context_mask: torch.Tensor | None = None,
-    norm_eps: float = 1e-6,
 ) -> torch.Tensor:
-    batch_size = x.shape[0]
+    """Apply query/key AdaLN modulation then cross-attention.
+    ``x_normed`` is already RMS-normalized by ``post_sa_function``; this only
+    applies the affine (scale/shift) modulation, so the normalization is not
+    repeated here.
+    """
+    batch_size = x_normed.shape[0]
     shift_kv, scale_kv = (
-        prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype)
+        prompt_scale_shift_table[None, None].to(device=x_normed.device, dtype=x_normed.dtype)
         + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1)
     ).unbind(dim=2)
-    attn_input = rms_norm(x, eps=norm_eps) * (1 + q_scale) + q_shift
+    attn_input = x_normed * (1 + q_scale) + q_shift
     encoder_hidden_states = context * (1 + scale_kv) + shift_kv
     return attn(attn_input, context=encoder_hidden_states, mask=context_mask) * q_gate

packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/base_encoder.py

@@ -30,21 +30,31 @@ class GemmaTextEncoder(torch.nn.Module):
 
     def encode(
         self,
-        text: str,
+        prompts: list[str],
         padding_side: str = "left",  # noqa: ARG002
-    ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
-        """Run Gemma LLM and return raw hidden states + attention mask.
-        Calls the inner model (self.model.model) to skip lm_head logits computation (~500 MiB saving).
-        Returns:
-            (hidden_states, attention_mask) where hidden_states is a tuple of per-layer tensors.
+    ) -> list[tuple[tuple[torch.Tensor, ...], torch.Tensor]]:
+        """Run a single fused Gemma forward over a batch of prompts.
+        Calls the inner model (self.model.model) to skip lm_head logits computation
+        (~500 MiB saving). The tokenizer pads every prompt to ``max_length`` (1024),
+        so the inputs stack into a single ``[N, 1024]`` batch with no further padding
+        logic; per-prompt outputs are sliced back to ``[1, 1024, D]`` / ``[1, 1024]``
+        in the original order.
         """
-        token_pairs = self.tokenizer.tokenize_with_weights(text)["gemma"]
-        input_ids = torch.tensor([[t[0] for t in token_pairs]], device=self.model.device)
-        attention_mask = torch.tensor([[w[1] for w in token_pairs]], device=self.model.device)
+        if not prompts:
+            return []
+        tokenized = [self.tokenizer.tokenize_with_weights(t)["gemma"] for t in prompts]
+        input_ids = torch.tensor(
+            [[tok for tok, _ in pairs] for pairs in tokenized],
+            device=self.model.device,
+        )
+        attention_mask = torch.tensor(
+            [[w for _, w in pairs] for pairs in tokenized],
+            device=self.model.device,
+        )
         outputs = self.model.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
         hidden_states = outputs.hidden_states
         del outputs
-        return hidden_states, attention_mask
+        return [(tuple(h[i : i + 1] for h in hidden_states), attention_mask[i : i + 1]) for i in range(len(prompts))]
 
     # --- Prompt enhancement methods ---
 

packages/ltx-pipelines/CLAUDE.md

@@ -17,6 +17,7 @@ Inference pipelines for LTX-2 audio-video generation. Depends on `ltx-core` for
 | Pipeline | File | Stages | Model | Sampler | Use case |
 |----------|------|--------|-------|---------|----------|
 | `TI2VidOneStagePipeline` | `ti2vid_one_stage.py` | 1 | Full | Euler | Simple text/image-to-video |
+| `T2AOneStagePipeline` | `t2a_one_stage.py` | 1 | Full | Euler | Text-to-audio (audio-only output, no video branch) |
 | `TI2VidTwoStagesPipeline` | `ti2vid_two_stages.py` | 2 | Full + distilled LoRA | Euler | Production quality |
 | `TI2VidTwoStagesHQPipeline` | `ti2vid_two_stages_hq.py` | 2 | Full + distilled LoRA (both stages) | Res2s | Highest quality, fewer steps |
 | `A2VidPipelineTwoStage` | `a2vid_two_stage.py` | 2 | Full + distilled LoRA | Euler | Audio-conditioned video |

packages/ltx-pipelines/README.md

@@ -58,6 +58,7 @@ Available pipeline modules:
 - `ltx_pipelines.ti2vid_two_stages` - Two-stage text/image-to-video (recommended).
 - `ltx_pipelines.ti2vid_two_stages_hq` - Two-stage text/image-to-video (different sampler, better quality).
 - `ltx_pipelines.ti2vid_one_stage` - Single-stage text/image-to-video.
+- `ltx_pipelines.t2a_one_stage` - Single-stage text-to-audio (audio-only output).
 - `ltx_pipelines.distilled` - Fast text/image-to-video pipeline using only the distilled model.
 - `ltx_pipelines.ic_lora` - Video-to-video with IC-LoRA.
 - `ltx_pipelines.keyframe_interpolation` - Keyframe interpolation.
@@ -254,6 +255,20 @@ Uses IC-LoRA on a **distilled** checkpoint with a **single** lip-dub IC-LoRA app
 
 ---
 
+### 11. T2AOneStagePipeline
+
+**Best for:** Text-to-audio — generating speech/audio only (no video) from a text prompt, e.g. driving an audio-style LoRA such as an accent LoRA.
+
+**Source**: [`src/ltx_pipelines/t2a_one_stage.py`](src/ltx_pipelines/t2a_one_stage.py)
+
+Single-stage, **audio-only** generation: the video branch is absent (`video=None`), so only the audio modality is denoised and decoded through the audio VAE + vocoder, producing a wave file. Audio duration is derived from `--num-frames` / `--frame-rate` (the same `8k+1` frame convention as video). Audio guidance (CFG/STG) is optional — the `--audio-*` flags default to the model's values; the video→audio cross-modal guidance is disabled since there is no video modality.
+
+**Extra CLI arguments (all optional, with sensible defaults):** `--num-frames`, `--frame-rate`, `--negative-prompt`, `--audio-cfg-guidance-scale`, `--audio-stg-guidance-scale`, `--audio-stg-blocks`, `--audio-rescale-scale`, `--audio-skip-step`. No `--height/--width/--image` (audio has no spatial dimensions).
+
+**Use when:** You need speech/audio from text alone, or to evaluate an audio-only LoRA (accent, voice style) without generating video.
+
+---
+
 ## 🎨 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.
@@ -400,6 +415,69 @@ By default, pipelines clean GPU memory (especially transformer weights) between
 # utils.cleanup_memory()  # Comment out if you have enough VRAM
 ```
 
+### Compilation (`torch.compile`)
+
+Compiling the transformer blocks with `torch.compile` speeds up inference. It is **opt-in and off by default**. The blocks are compiled shape-polymorphically (the sequence dimension is marked dynamic), so one compiled artifact serves any token count without recompiling.
+
+**CLI** — the `--compile` flag maps directly to `CompilationConfig`:
+
+| Form | Result |
+| ---- | ------ |
+| *(flag absent)* | eager, no compilation |
+| `--compile` | compile with defaults |
+| `--compile KEY=VALUE ...` | compile, overriding individual fields |
+
+```bash
+# Defaults
+python -m ltx_pipelines.ti2vid_two_stages --compile --checkpoint-path=...
+
+# reduce-overhead captures CUDA graphs -- the main latency lever for the denoising loop.
+# Off by default because graph capture reserves static memory pools (extra VRAM), so it
+# trades memory for speed; enable it when you have headroom.
+python -m ltx_pipelines.ti2vid_two_stages --compile mode=reduce-overhead --checkpoint-path=...
+
+# Several overrides at once
+python -m ltx_pipelines.ti2vid_two_stages \
+    --compile mode=max-autotune fullgraph=true dynamic=true --checkpoint-path=...
+```
+
+| Field | Values | Default | Notes |
+| ----- | ------ | ------- | ----- |
+| `mode` | `none`, `reduce-overhead`, `max-autotune`, … | `none` | `reduce-overhead`/`max-autotune` enable CUDA graphs |
+| `backend` | `inductor`, `eager`, … | `inductor` | |
+| `fullgraph` | `true`/`false` | `false` | |
+| `dynamic` | `auto`/`true`/`false` | `auto` | the seq dim is marked dynamic regardless |
+| `inductor_config` | JSON object or path to a `.json` | `{}` | `torch._inductor.config` overrides |
+| `dynamo_config` | JSON object or path to a `.json` | `{"inline_inbuilt_nn_modules": true, "cache_size_limit": 256}` | `torch._dynamo.config` overrides |
+
+**Controlling inductor / dynamo configs.** `inductor_config` and `dynamo_config` take either an inline JSON object or a path to a `.json` file, applied via `torch._inductor.config.patch(...)` / `torch._dynamo.config.patch(...)` around the compiled forward. They **replace the defaults wholesale — they do not merge**, so when overriding `dynamo_config` re-include any defaults you want to keep:
+
+```bash
+python -m ltx_pipelines.ti2vid_two_stages \
+    --compile 'inductor_config={"max_autotune": true}' \
+              'dynamo_config={"inline_inbuilt_nn_modules": true, "cache_size_limit": 256, "recompile_limit": 32}' \
+    --checkpoint-path=...
+```
+
+**Programmatically**, pass a `CompilationConfig` to the pipeline:
+
+```python
+from ltx_core.model.transformer.compiling import CompilationConfig
+
+pipeline = TI2VidTwoStagesPipeline(
+    ...,
+    compilation_config=CompilationConfig(mode="reduce-overhead"),
+)
+```
+
+**Faster cache loads: `unsafe_skip_cache_dynamic_shape_guards` (unsafe, opt-in).** Inductor's FX-graph cache re-checks the dynamic-shape guards stored with each entry on every lookup. Setting this flag skips that re-check (every entry is treated as a guard hit), which speeds up warm and cross-process cache loads. It is **not enabled by default** because it is a correctness hazard: a kernel first compiled at a small sequence length keeps int32 address arithmetic, and reusing it at a larger sequence length (roughly **>58k tokens/rank**) overflows int32 and reads out of bounds — surfacing as a CUDA illegal memory access or silently corrupted output. Only enable it when your token counts stay within the range the cached kernels were compiled for:
+
+```bash
+python -m ltx_pipelines.ti2vid_two_stages \
+    --compile 'inductor_config={"unsafe_skip_cache_dynamic_shape_guards": true}' \
+    --checkpoint-path=...
+```
+
 ### Denoising Loop Optimization
 
 **Gradient Estimation Denoising Loop:**

packages/ltx-pipelines/pyproject.toml

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

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

@@ -2,6 +2,7 @@
 LTX-2 Pipelines: High-level video generation pipelines and utilities.
 This package provides ready-to-use pipelines for video generation:
 - TI2VidOneStagePipeline: Text/image-to-video in a single stage
+- T2AOneStagePipeline: Text-to-audio in a single stage (audio-only output)
 - TI2VidTwoStagesPipeline: Two-stage generation with upsampling
 - DistilledPipeline: Fast distilled two-stage generation
 - ICLoraPipeline: Image/video conditioning with distilled LoRA
@@ -18,6 +19,7 @@ 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.t2a_one_stage import T2AOneStagePipeline
 from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
 from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
 
@@ -28,6 +30,7 @@ __all__ = [
     "KeyframeInterpolationPipeline",
     "LipDubPipeline",
     "RetakePipeline",
+    "T2AOneStagePipeline",
     "TI2VidOneStagePipeline",
     "TI2VidTwoStagesPipeline",
 ]

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

@@ -14,6 +14,7 @@ from ltx_core.types import Audio, VideoPixelShape
 from ltx_pipelines.iclora_utils import (
     append_ic_lora_reference_video_conditionings,
     read_lora_reference_downscale_factor,
+    read_lora_reference_temporal_scale_factor,
 )
 from ltx_pipelines.utils.args import (
     ImageConditioningInput,
@@ -102,10 +103,11 @@ class ICLoraPipeline:
         self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
         self.audio_decoder = AudioDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
 
-        # Read reference downscale factor from LoRA metadata.
-        # IC-LoRAs trained with low-resolution reference videos store this factor
-        # so inference can resize reference videos to match training conditions.
+        # Read reference scale factors from LoRA metadata.
+        # IC-LoRAs trained with scaled reference videos store these factors
+        # so inference can resize/subsample reference videos to match training conditions.
         self.reference_downscale_factor = 1
+        self.reference_temporal_scale_factor = 1
         for lora in loras:
             scale = read_lora_reference_downscale_factor(lora.path)
             if scale != 1:
@@ -116,6 +118,15 @@ class ICLoraPipeline:
                         f"specifies {scale}. Cannot combine LoRAs with different reference scales."
                     )
                 self.reference_downscale_factor = scale
+            temporal = read_lora_reference_temporal_scale_factor(lora.path)
+            if temporal != 1:
+                if self.reference_temporal_scale_factor not in (1, temporal):
+                    raise ValueError(
+                        f"Conflicting reference_temporal_scale_factor values in LoRAs: "
+                        f"already have {self.reference_temporal_scale_factor}, but {lora.path} "
+                        f"specifies {temporal}. Cannot combine LoRAs with different temporal scales."
+                    )
+                self.reference_temporal_scale_factor = temporal
 
     def __call__(  # noqa: PLR0913
         self,
@@ -318,6 +329,7 @@ class ICLoraPipeline:
             dtype=self.dtype,
             device=self.device,
             reference_downscale_factor=self.reference_downscale_factor,
+            reference_temporal_scale_factor=self.reference_temporal_scale_factor,
             conditioning_attention_strength=conditioning_attention_strength,
             conditioning_attention_mask=conditioning_attention_mask,
             tiling_config=None,

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

@@ -31,6 +31,17 @@ def read_lora_reference_downscale_factor(lora_path: str) -> int:
         return 1
 
 
+def read_lora_reference_temporal_scale_factor(lora_path: str) -> int:
+    """Read ``reference_temporal_scale_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_temporal_scale_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,
@@ -66,6 +77,12 @@ def downsample_mask_video_to_latent(
     return rearrange(latent_mask, "b 1 f h w -> b (f h w)")
 
 
+def temporal_subsample(video: torch.Tensor, temporal_scale_factor: int) -> torch.Tensor:
+    """VAE-aligned temporal subsampling: keep frame 0, then every Nth frame."""
+    indices = [0, *list(range(1, video.shape[2], temporal_scale_factor))]
+    return video[:, :, indices]
+
+
 def append_ic_lora_reference_video_conditionings(  # noqa: PLR0913
     conditionings: list[ConditioningItem],
     video_conditioning: list[tuple[str, float]],
@@ -77,6 +94,7 @@ def append_ic_lora_reference_video_conditionings(  # noqa: PLR0913
     dtype: torch.dtype,
     device: torch.device,
     reference_downscale_factor: int,
+    reference_temporal_scale_factor: int = 1,
     conditioning_attention_strength: float,
     conditioning_attention_mask: torch.Tensor | None,
     tiling_config: TilingConfig | None = None,
@@ -93,6 +111,8 @@ def append_ic_lora_reference_video_conditionings(  # noqa: PLR0913
     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 reference_temporal_scale_factor > 1:
+            video = temporal_subsample(video, reference_temporal_scale_factor)
         if tiling_config is not None:
             encoded_video = video_encoder.tiled_encode(video, tiling_config)
         else:
@@ -113,6 +133,7 @@ def append_ic_lora_reference_video_conditionings(  # noqa: PLR0913
         cond = VideoConditionByReferenceLatent(
             latent=encoded_video,
             downscale_factor=scale,
+            temporal_scale_factor=reference_temporal_scale_factor,
             strength=strength,
         )
         if attn_mask is not None:

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

@@ -0,0 +1,192 @@
+import logging
+
+import torch
+
+from ltx_core.components.guiders import (
+    MultiModalGuiderFactory,
+    MultiModalGuiderParams,
+    create_multimodal_guider_factory,
+)
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.schedulers import LTX2Scheduler
+from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer import LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP, LTXAudioOnlyModelConfigurator
+from ltx_core.model.transformer.compiling import CompilationConfig
+from ltx_core.quantization import QuantizationPolicy
+from ltx_core.types import Audio
+from ltx_pipelines.utils import get_device
+from ltx_pipelines.utils.args import (
+    default_1_stage_t2a_arg_parser,
+    detect_checkpoint_path,
+)
+from ltx_pipelines.utils.blocks import (
+    AudioDecoder,
+    DiffusionStage,
+    PromptEncoder,
+)
+from ltx_pipelines.utils.constants import detect_params
+from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser
+from ltx_pipelines.utils.media_io import encode_audio
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
+
+# Placeholder pixel dimensions used for ``VideoPixelShape`` construction.
+# Audio-only generation reads ``frames`` and ``fps`` from the pixel shape via
+# ``AudioLatentShape.from_video_pixel_shape`` (height/width are unused).
+_AUDIO_ONLY_PLACEHOLDER_RES = 512
+
+
+class T2AOneStagePipeline:
+    """
+    Single-stage text-to-audio generation pipeline.
+    Generates audio at the target duration in a single diffusion pass with
+    classifier-free guidance (CFG) on the audio modality only. The video
+    modality is fully absent — the transformer runs audio-only by passing
+    ``video=None`` to the ``DiffusionStage``.
+    Assumes full non distilled model is provided in the checkpoint_path.
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        gemma_root: str,
+        loras: list[LoraPathStrengthAndSDOps],
+        device: torch.device | None = None,
+        quantization: QuantizationPolicy | None = None,
+        registry: Registry | None = None,
+        compilation_config: CompilationConfig | None = None,
+        offload_mode: OffloadMode = OffloadMode.NONE,
+    ):
+        self.dtype = torch.bfloat16
+        self.device = device or get_device()
+        self._scheduler = LTX2Scheduler()
+        self.prompt_encoder = PromptEncoder(
+            checkpoint_path=checkpoint_path,
+            gemma_root=gemma_root,
+            dtype=self.dtype,
+            device=self.device,
+            registry=registry,
+            offload_mode=offload_mode,
+        )
+        # Audio-only: build an audio-only transformer (model_configurator) so the video
+        # weights are never instantiated, plus a use-case-specific SDOps that restricts
+        # checkpoint reads to the audio model's keys, so the video weights are never even
+        # read from disk (the loader skips any key the SDOps maps to None).
+        self.stage = DiffusionStage(
+            checkpoint_path=checkpoint_path,
+            dtype=self.dtype,
+            device=self.device,
+            loras=tuple(loras),
+            quantization=quantization,
+            registry=registry,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
+            model_configurator=LTXAudioOnlyModelConfigurator,
+            model_sd_ops=LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP,
+        )
+        self.audio_decoder = AudioDecoder(
+            checkpoint_path=checkpoint_path,
+            dtype=self.dtype,
+            device=self.device,
+            registry=registry,
+        )
+
+    def __call__(
+        self,
+        prompt: str,
+        negative_prompt: str,
+        seed: int,
+        num_frames: int,
+        frame_rate: float,
+        num_inference_steps: int,
+        audio_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
+        enhance_prompt: bool = False,
+        max_batch_size: int = 1,
+        sigmas: torch.Tensor | None = None,
+    ) -> Audio:
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        noiser = GaussianNoiser(generator=generator)
+
+        ctx_p, ctx_n = self.prompt_encoder(
+            [prompt, negative_prompt],
+            enhance_first_prompt=enhance_prompt,
+            enhance_prompt_image=None,
+            enhance_prompt_seed=seed,
+        )
+        a_context_p = ctx_p.audio_encoding
+        a_context_n = ctx_n.audio_encoding
+
+        sigmas = (sigmas if sigmas is not None else self._scheduler.execute(steps=num_inference_steps)).to(
+            dtype=torch.float32, device=self.device
+        )
+
+        # Normalize to a guider factory. Plain ``MultiModalGuiderParams`` (the default /
+        # CLI case) becomes a simple sigma-independent guider, but callers may also pass
+        # their own factory for sigma-dependent guidance; ``FactoryGuidedDenoiser`` always
+        # consumes a factory.
+        audio_guider_factory = create_multimodal_guider_factory(
+            params=audio_guider_params,
+            negative_context=a_context_n,
+        )
+
+        _, audio_state = self.stage(
+            denoiser=FactoryGuidedDenoiser(
+                v_context=None,
+                a_context=a_context_p,
+                video_guider_factory=None,
+                audio_guider_factory=audio_guider_factory,
+            ),
+            sigmas=sigmas,
+            noiser=noiser,
+            width=_AUDIO_ONLY_PLACEHOLDER_RES,
+            height=_AUDIO_ONLY_PLACEHOLDER_RES,
+            frames=num_frames,
+            fps=frame_rate,
+            video=None,
+            audio=ModalitySpec(context=a_context_p),
+            max_batch_size=max_batch_size,
+        )
+
+        return self.audio_decoder(audio_state.latent)
+
+
+@torch.inference_mode()
+def main() -> None:
+    logging.basicConfig(level=logging.INFO)
+    checkpoint_path = detect_checkpoint_path()
+    params = detect_params(checkpoint_path)
+    parser = default_1_stage_t2a_arg_parser(params=params)
+    args = parser.parse_args()
+    pipeline = T2AOneStagePipeline(
+        checkpoint_path=args.checkpoint_path,
+        gemma_root=args.gemma_root,
+        loras=tuple(args.lora) if args.lora else (),
+        quantization=args.quantization,
+        compilation_config=args.compile,
+        offload_mode=args.offload_mode,
+    )
+    audio = pipeline(
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        seed=args.seed,
+        num_frames=args.num_frames,
+        frame_rate=args.frame_rate,
+        num_inference_steps=args.num_inference_steps,
+        audio_guider_params=MultiModalGuiderParams(
+            cfg_scale=args.audio_cfg_guidance_scale,
+            stg_scale=args.audio_stg_guidance_scale,
+            rescale_scale=args.audio_rescale_scale,
+            # Audio-only generation has no video modality, so the video->audio
+            # (v2a) cross-modal guidance is meaningless here. 1.0 disables it.
+            modality_scale=1.0,
+            skip_step=args.audio_skip_step,
+            stg_blocks=args.audio_stg_blocks,
+        ),
+        max_batch_size=args.max_batch_size,
+    )
+
+    encode_audio(audio=audio, output_path=args.output_path)
+
+
+if __name__ == "__main__":
+    main()

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

@@ -650,6 +650,62 @@ def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
     return parser
 
 
+def default_1_stage_t2a_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
+    """Argument parser for single-stage text-to-audio pipelines (audio-only)."""
+    audio_guider = params.audio_guider_params
+    parser = basic_arg_parser(params=params)
+    parser.add_argument(
+        "--num-frames",
+        type=int,
+        default=params.num_frames,
+        help="Number of frames used to derive audio duration (num-frames / frame-rate).",
+    )
+    parser.add_argument(
+        "--frame-rate",
+        type=float,
+        default=params.frame_rate,
+        help="Frame rate used with --num-frames to derive the audio duration.",
+    )
+    parser.add_argument(
+        "--negative-prompt",
+        type=str,
+        default=DEFAULT_NEGATIVE_PROMPT,
+        help="Negative prompt to steer audio generation away from artifacts.",
+    )
+    parser.add_argument(
+        "--audio-cfg-guidance-scale",
+        type=float,
+        default=audio_guider.cfg_scale,
+        help=f"Audio CFG scale (default: {audio_guider.cfg_scale}).",
+    )
+    parser.add_argument(
+        "--audio-stg-guidance-scale",
+        type=float,
+        default=audio_guider.stg_scale,
+        help=f"Audio STG scale (default: {audio_guider.stg_scale}).",
+    )
+    parser.add_argument(
+        "--audio-rescale-scale",
+        type=float,
+        default=audio_guider.rescale_scale,
+        help=f"Audio rescale scale (default: {audio_guider.rescale_scale}).",
+    )
+    parser.add_argument(
+        "--audio-stg-blocks",
+        type=int,
+        nargs="*",
+        default=audio_guider.stg_blocks,
+        help=f"Blocks to perturb for Audio STG (default: {audio_guider.stg_blocks}).",
+    )
+    parser.add_argument(
+        "--audio-skip-step",
+        type=int,
+        default=audio_guider.skip_step,
+        help=f"Audio skip step (default: {audio_guider.skip_step}).",
+    )
+    return parser
+
+
 def default_2_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
     parser = default_1_stage_arg_parser(params=params)
     parser.set_defaults(height=params.stage_2_height, width=params.stage_2_width)

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

@@ -7,7 +7,6 @@ removes the need for :class:`ModelLedger`.
 from __future__ import annotations
 
 import copy
-import dataclasses
 import logging
 from collections.abc import Iterator
 from contextlib import AbstractContextManager, contextmanager
@@ -40,9 +39,9 @@ from ltx_core.model.audio_vae import (
 from ltx_core.model.audio_vae import (
     decode_audio as vae_decode_audio,
 )
+from ltx_core.model.model_protocol import LTXModelProtocol, ModelConfigurator
 from ltx_core.model.transformer import (
     LTXV_MODEL_COMFY_RENAMING_MAP,
-    LTXModel,
     LTXModelConfigurator,
     X0Model,
 )
@@ -144,7 +143,7 @@ def _streaming_model(
     """Build a streaming wrapper, yield it, then tear down and free memory."""
     cpu_slots_count = DISK_CPU_SLOTS if offload_mode == OffloadMode.DISK else None
     wrapped = builder.build(
-        target_device=target_device,
+        device=target_device,
         dtype=dtype,
         cpu_slots_count=cpu_slots_count,
     )
@@ -195,7 +194,7 @@ class DiffusionStage:
     pattern in every pipeline.
     """
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         checkpoint_path: str,
         dtype: torch.dtype,
@@ -205,7 +204,9 @@ class DiffusionStage:
         registry: Registry | None = None,
         compilation_config: CompilationConfig | None = None,
         offload_mode: OffloadMode = OffloadMode.NONE,
-        transformer_builder: ModelBuilderProtocol[LTXModel] | None = None,
+        transformer_builder: ModelBuilderProtocol[LTXModelProtocol] | None = None,
+        model_configurator: type[ModelConfigurator] = LTXModelConfigurator,
+        model_sd_ops: SDOps = LTXV_MODEL_COMFY_RENAMING_MAP,
     ) -> None:
         self._checkpoint_path = checkpoint_path
         self._dtype = dtype
@@ -213,10 +214,12 @@ class DiffusionStage:
         self._quantization = quantization
         self._compilation_config = compilation_config
         self._offload_mode = offload_mode
+        # A quantization policy may pin its own configurator; otherwise use the one
+        # provided by the caller (defaults to the audio-video LTXModelConfigurator).
         configurator = (
             quantization.model_configurator
             if quantization is not None and quantization.model_configurator is not None
-            else LTXModelConfigurator
+            else model_configurator
         )
         if transformer_builder is not None:
             self._transformer_builder = transformer_builder
@@ -224,14 +227,12 @@ class DiffusionStage:
             self._transformer_builder = Builder(
                 model_path=checkpoint_path,
                 model_class_configurator=configurator,
-                model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
+                model_sd_ops=model_sd_ops,
                 loras=tuple(loras),
                 registry=registry or DummyRegistry(),
             )
 
         if offload_mode != OffloadMode.NONE:
-            if compilation_config is not None:
-                raise ValueError("torch.compile is not supported with layer streaming")
             # WeightsProvider currently only supports plain bf16 + fp8_cast LoRA fusion
             # (no companion-key emission). Quantization policies that emit
             # companion keys (e.g. ``.weight_scale``) cannot be streamed yet.
@@ -240,8 +241,15 @@ class DiffusionStage:
                     "Block streaming is not supported with this quantization policy "
                     "(only bf16 and fp8_cast are currently supported)."
                 )
-            streaming_sd_ops: SDOps = LTXV_MODEL_COMFY_RENAMING_MAP
+            streaming_sd_ops: SDOps = model_sd_ops
             streaming_module_ops: tuple[ModuleOps, ...] = ()
+            streaming_loras = tuple(loras)
+
+            if compilation_config:
+                number_of_layers = self._transformer_builder.model_config()["transformer"]["num_layers"]
+                streaming_sd_ops, streaming_module_ops, streaming_loras = _apply_compile_ops(
+                    streaming_sd_ops, streaming_module_ops, streaming_loras, number_of_layers
+                )
             if quantization is not None:
                 streaming_sd_ops, streaming_module_ops = _chain_quantization(
                     streaming_sd_ops, streaming_module_ops, quantization
@@ -251,7 +259,7 @@ class DiffusionStage:
                 model_path=checkpoint_path,
                 model_sd_ops=streaming_sd_ops,
                 module_ops=streaming_module_ops,
-                loras=tuple(loras),
+                loras=streaming_loras,
                 registry=registry or DummyRegistry(),
                 fuse_rule=quantization.fuse_rule if quantization is not None else bf16_fuse_rule,
                 blocks_attr="transformer_blocks",
@@ -273,9 +281,8 @@ class DiffusionStage:
             (*self._transformer_builder.module_ops, op),
         )
         if self._offload_mode != OffloadMode.NONE:
-            new._streaming_builder = dataclasses.replace(
-                self._streaming_builder,
-                module_ops=(*self._streaming_builder.module_ops, op),
+            new._streaming_builder = self._streaming_builder.with_module_ops(
+                (*self._streaming_builder.module_ops, op),
             )
         return new
 
@@ -531,7 +538,7 @@ class PromptEncoder:
                 prompts[0] = generate_enhanced_prompt(
                     text_encoder, prompts[0], enhance_prompt_image, seed=enhance_prompt_seed
                 )
-            raw_outputs = [text_encoder.encode(p) for p in prompts]
+            raw_outputs = text_encoder.encode(prompts)
         logger.info("Text encoder done, building embeddings processor from %s", self._checkpoint_path)
 
         with gpu_model(self._build_embeddings_processor()) as embeddings_processor:

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

@@ -182,6 +182,8 @@ def _guided_denoise(  # noqa: PLR0913,PLR0915
     r = dict(zip(pass_names, zip(splits_v, splits_a, strict=True), strict=True))
 
     cond_v, cond_a = r["cond"]
+    cond_v = cond_v if isinstance(cond_v, torch.Tensor) else torch.tensor(cond_v)
+    cond_a = cond_a if isinstance(cond_a, torch.Tensor) else torch.tensor(cond_a)
     uncond_v, uncond_a = r.get("uncond", (0.0, 0.0))
     ptb_v, ptb_a = r.get("ptb", (0.0, 0.0))
     mod_v, mod_a = r.get("mod", (0.0, 0.0))

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

@@ -391,6 +391,24 @@ def encode_video(
     logger.info(f"Video saved to {output_path}")
 
 
+def encode_audio(audio: Audio, output_path: str) -> None:
+    """Save an audio waveform as a 16-bit PCM ``.wav`` file at the source sampling rate.
+    Reuses :func:`_write_audio` (the same muxing path used by :func:`encode_video`);
+    the only difference is a PCM (``pcm_s16le``) stream in a WAV container instead of
+    the AAC stream used for muxed video.
+    """
+    container = av.open(output_path, mode="w")
+    audio_stream = container.add_stream("pcm_s16le", rate=audio.sampling_rate)
+    audio_stream.codec_context.sample_rate = audio.sampling_rate
+    audio_stream.codec_context.layout = "stereo"
+    audio_stream.codec_context.time_base = Fraction(1, audio.sampling_rate)
+    try:
+        _write_audio(container, audio_stream, audio)
+    finally:
+        container.close()
+    logger.info(f"Audio saved to {output_path}")
+
+
 def _encode_chunks_threaded(
     container: av.container.Container,
     stream: av.video.stream.VideoStream,

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

@@ -436,7 +436,7 @@ def res2s_audio_video_denoising_loop(  # noqa: PLR0913,PLR0915,PLR0912
     return video_state, audio_state
 
 
-def euler_cfg_pp_denoising_loop(
+def euler_cfg_pp_denoising_loop(  # noqa: PLR0912
     sigmas: torch.Tensor,
     video_state: LatentState | None,
     audio_state: LatentState | None,
@@ -514,9 +514,15 @@ def euler_cfg_pp_denoising_loop(
             )
 
         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)
+            denoised_video = post_process_latent(
+                denoised_video.float(), video_state.denoise_mask, video_state.clean_latent
+            )
+            noisy_video = video_state.latent.float()
         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)
+            denoised_audio = post_process_latent(
+                denoised_audio.float(), audio_state.denoise_mask, audio_state.clean_latent
+            )
+            noisy_audio = audio_state.latent.float()
 
         if sigmas[step_idx + 1] == 0:
             if video_state is not None and denoised_video is not None:
@@ -525,30 +531,31 @@ def euler_cfg_pp_denoising_loop(
                 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:
+            video_noise = new_noise_fn(video_state.latent, generator) if draw_noise else None
             x_next = stepper.step(
-                sample=video_state.latent,
+                sample=noisy_video,
                 denoised_sample=denoised_video,
                 sigmas=sigmas,
                 step_index=step_idx,
                 uncond_denoised=uncond_video,
-                noise=noise_video,
+                noise=video_noise,
             )
+            if draw_noise:
+                x_next = post_process_latent(x_next, video_state.denoise_mask, video_state.clean_latent)
             video_state = replace(video_state, latent=x_next.to(model_dtype))
 
         if audio_state is not None and denoised_audio is not None:
+            audio_noise = new_noise_fn(audio_state.latent, generator) if draw_noise else None
             x_next = stepper.step(
-                sample=audio_state.latent,
+                sample=noisy_audio,
                 denoised_sample=denoised_audio,
                 sigmas=sigmas,
                 step_index=step_idx,
                 uncond_denoised=uncond_audio,
-                noise=noise_audio,
+                noise=audio_noise,
             )
+            if draw_noise:
+                x_next = post_process_latent(x_next, audio_state.denoise_mask, audio_state.clean_latent)
             audio_state = replace(audio_state, latent=x_next.to(model_dtype))
-
     return video_state, audio_state

packages/ltx-trainer/.gitignore

@@ -1,7 +1,3 @@
-configs/*.yaml
-!configs/ltx2_av_lora.yaml
-!configs/ltx2_av_lora_low_vram.yaml
-!configs/ltx2_v2v_ic_lora.yaml
 datasets
 outputs
 wandb

packages/ltx-trainer/AGENTS.md

@@ -6,10 +6,22 @@ This file provides guidance to AI coding assistants (Claude, Cursor, etc.) when
 
 **LTX Trainer** is a training toolkit for fine-tuning the Lightricks LTX audio-video generation models. It supports:
 
+- **Text-to-video (T2V)** - Generate video from text prompts
+- **Text-to-audio (T2A)** - Generate audio from text prompts
+- **Image-to-video (I2V)** - Generate video conditioned on a first frame
+- **Video extension** - Forward (prefix) and backward (suffix) video continuation
+- **Video inpainting** - Mask-based spatial/temporal inpainting
+- **Video outpainting** - Spatial crop-based outpainting
+- **IC-LoRA video-to-video** - In-context control adapters for style/structure transfer
+- **Audio-to-video (A2V)** and **Video-to-audio (V2A)** - Cross-modal generation with frozen conditioning
+- **Audio extension** - Forward (prefix) and backward (suffix) audio continuation
+- **Audio inpainting** - Mask-based audio inpainting
+- **IC-LoRA audio-to-audio (A2A)** - Audio reference conditioning for style transfer
+- **AV2AV IC-LoRA** - Combined video and audio reference conditioning
 - **LoRA training** - Efficient fine-tuning with adapters
 - **Full fine-tuning** - Complete model training
-- **Audio-video training** - Joint audio and video generation
-- **IC-LoRA training** - In-context control adapters for video-to-video transformations
+
+All conditioning scenarios are expressed through the unified `FlexibleStrategy` configuration.
 
 **Supported model versions:**
 
@@ -39,13 +51,14 @@ packages/ltx-trainer/
 │   ├── config_display.py         # Config pretty-printing
 │   ├── trainer.py                # Main training orchestration with Accelerate
 │   ├── model_loader.py           # Model loading using ltx-core
-│   ├── validation_sampler.py     # Inference for validation samples
+│   ├── validation_runner.py      # ValidationRunner — conditioned validation sampling
 │   ├── datasets.py               # PrecomputedDataset, DummyDataset
 │   ├── training_strategies/      # Strategy pattern for different training modes
 │   │   ├── __init__.py           # Factory function: get_training_strategy()
 │   │   ├── base_strategy.py      # TrainingStrategy ABC, ModelInputs, TrainingStrategyConfigBase
-│   │   ├── text_to_video.py      # TextToVideoStrategy, TextToVideoConfig
-│   │   └── video_to_video.py     # VideoToVideoStrategy, VideoToVideoConfig
+│   │   ├── flexible.py           # FlexibleStrategy, FlexibleStrategyConfig [RECOMMENDED]
+│   │   ├── text_to_video.py      # TextToVideoStrategy, TextToVideoConfig [DEPRECATED]
+│   │   └── video_to_video.py     # VideoToVideoStrategy, VideoToVideoConfig [DEPRECATED]
 │   ├── timestep_samplers.py      # Flow matching timestep sampling
 │   ├── gemma_8bit.py             # 8-bit Gemma text encoder loading (bitsandbytes)
 │   ├── quantization.py           # Transformer INT8/INT4/FP8 quantization
@@ -62,13 +75,25 @@ packages/ltx-trainer/
 │   ├── process_captions.py       # Text embedding computation
 │   ├── caption_videos.py         # Automatic video captioning
 │   ├── decode_latents.py         # Latent decoding for debugging
-│   ├── inference.py              # Inference with trained models
 │   ├── compute_reference.py      # Generate IC-LoRA reference videos
 │   └── split_scenes.py           # Scene detection and splitting
 ├── configs/                      # Example training configurations
-│   ├── ltx2_av_lora.yaml         # Audio-video LoRA training
-│   ├── ltx2_av_lora_low_vram.yaml
-│   ├── ltx2_v2v_ic_lora.yaml     # IC-LoRA video-to-video
+│   ├── t2v_lora.yaml             # Text-to-video LoRA
+│   ├── t2v_lora_low_vram.yaml    # Text-to-video LoRA (low VRAM)
+│   ├── i2v_lora.yaml             # Image-to-video LoRA
+│   ├── v2v_ic_lora.yaml          # IC-LoRA video-to-video
+│   ├── a2v_lora.yaml             # Audio-to-video LoRA
+│   ├── v2a_lora.yaml             # Video-to-audio LoRA
+│   ├── video_extend_lora.yaml    # Video extension (forward)
+│   ├── video_suffix_lora.yaml    # Video extension (backward)
+│   ├── video_inpainting_lora.yaml # Video inpainting
+│   ├── video_outpainting_lora.yaml # Video outpainting
+│   ├── t2a_lora.yaml             # Text-to-audio LoRA
+│   ├── audio_extend_lora.yaml    # Audio extension (forward)
+│   ├── audio_suffix_lora.yaml    # Audio extension (backward)
+│   ├── audio_inpainting_lora.yaml # Audio inpainting
+│   ├── a2a_ic_lora.yaml          # Audio-to-audio IC-LoRA
+│   ├── av2av_ic_lora.yaml        # AV2AV IC-LoRA
 │   └── accelerate/               # FSDP, DDP configs
 ├── tests/                        # Pytest tests
 └── docs/                         # Documentation
@@ -83,7 +108,8 @@ packages/ltx-trainer/
   `load_text_encoder()`, `load_embeddings_processor()`, etc.
 - Combined loader: `load_model()` returns `LtxModelComponents` dataclass
 - Uses `SingleGPUModelBuilder` from ltx-core internally
-- Text encoder and embeddings processor are loaded separately (the text encoder only needs Gemma weights; the embeddings processor only needs the LTX checkpoint)
+- Text encoder and embeddings processor are loaded separately (the text encoder only needs Gemma weights; the embeddings
+  processor only needs the LTX checkpoint)
 - 8-bit text encoder loading via `gemma_8bit.py` (bitsandbytes)
 
 **Training Flow:**
@@ -94,7 +120,7 @@ packages/ltx-trainer/
    kept)
 4. Each training step: embedding connectors applied → strategy prepares `ModelInputs` → transformer forward pass →
    strategy computes loss
-5. Training strategies (`TextToVideoStrategy`, `VideoToVideoStrategy`) handle mode-specific logic
+5. Training strategies (`FlexibleStrategy`) handle mode-specific logic including conditioning, masking, and loss computation
 6. Accelerate handles distributed training, mixed precision, and device placement
 7. Data flows as precomputed latents through `PrecomputedDataset`
 
@@ -136,7 +162,12 @@ LTX-2.3) and cross-modality (video↔audio) attention conditioning (both version
 
 - All config in `src/ltx_trainer/config.py`
 - Main class: `LtxTrainerConfig`
-- Training strategy configs: `TextToVideoConfig`, `VideoToVideoConfig`
+- `TrainingStrategyConfig` - Union of `FlexibleStrategyConfig` | `TextToVideoConfig` (deprecated) | `VideoToVideoConfig` (deprecated)
+- `FlexibleStrategyConfig` - Unified strategy config with `video`/`audio` `ModalityConfig` blocks
+- `ModalityConfig` - Per-modality config: `is_generated`, `latents_dir`, `conditions` list
+- `ConditionConfig` - Discriminated union: `FirstFrameConditionConfig`, `PrefixConditionConfig`, `SuffixConditionConfig`, `SpatialCropConditionConfig`, `MaskConditionConfig`, `ReferenceConditionConfig`
+- `ValidationSample` - Per-sample validation config with `prompt`, `conditions`, optional `video_dims`/`seed` overrides
+- `ValidationCondition` - Discriminated union for validation conditions (first_frame, prefix, suffix, spatial_crop, mask, reference, video_to_audio, audio_to_video)
 - Uses Pydantic field validators and model validators
 - Config uses `extra="forbid"` — unknown fields cause validation errors
 - Config files in `configs/` directory
@@ -206,8 +237,8 @@ These values are shared across all supported model versions:
 | Video latent channels        | 128                              | VAE encoder/decoder, patchifier, `VideoLatentShape`       |
 | Spatial compression          | 32× (H and W)                    | `SpatioTemporalScaleFactors.default()`, config validators |
 | Temporal compression         | 8×                               | `SpatioTemporalScaleFactors.default()`, config validators |
-| Frame constraint             | `frames % 8 == 1`                | Config validators, validation sampler                     |
-| Resolution constraint        | Width and height divisible by 32 | Config validators, validation sampler                     |
+| Frame constraint             | `frames % 8 == 1`                | Config validators, validation runner                      |
+| Resolution constraint        | Width and height divisible by 32 | Config validators, validation runner                      |
 | Audio latent channels        | 8                                | `AudioLatentShape`, audio patchifier                      |
 | Audio mel bins               | 16                               | `AudioLatentShape`, audio patchifier                      |
 | Patchified token dim (video) | 128 (`128 × 1 × 1 × 1`)          | Transformer `in_channels`                                 |
@@ -245,12 +276,53 @@ uv run pytest
 
 ```bash
 # Single GPU
-uv run python scripts/train.py configs/ltx2_av_lora.yaml
+uv run python scripts/train.py configs/t2v_lora.yaml
 
 # Multi-GPU with Accelerate
-uv run accelerate launch scripts/train.py configs/ltx2_av_lora.yaml
+uv run accelerate launch scripts/train.py configs/t2v_lora.yaml
 ```
 
+## Testing Standards
+
+### Structure
+
+- **Flat functions only** — use `def test_*()`, never `class Test*` with methods. Pytest collects standalone functions.
+- **Only test public interfaces** — never call private methods (`_method`) directly. Verify private behavior
+  indirectly through the public API.
+
+### What to Test
+
+- **Custom validators and business logic** — cross-field validators, domain constraints, error paths. These catch real
+  bugs.
+- **Behavioral tests** — call the public method, verify the outputs have the right shape, values, and structure. One
+  behavioral test is worth ten config-only tests.
+- **Edge cases and error paths** — boundary conditions, composed behaviors, expected exceptions.
+- **Contract tests** — required fields, rejected invalid inputs, safety mechanisms like `extra="forbid"`.
+
+### What NOT to Test
+
+- **Pydantic storing a value** — `Foo(x=1); assert foo.x == 1` tests Pydantic, not your code. If a behavioral test
+  already creates the same config and uses it, the config-only test adds nothing.
+- **Pydantic Literal defaults** — `assert config.type == "first_frame"` when `type` is `Literal["first_frame"]`.
+- **Pydantic default factories** — `assert config.conditions == []` when the field has `default_factory=list`.
+- **Tests already covered by behavioral tests** — if `test_prefix_conditioning` creates a valid `PrefixConditionConfig`
+  and exercises it end-to-end, a separate `test_prefix_valid` that just creates the same config is redundant.
+- **Trivial instantiation tests** — `strategy = Strategy(config); assert strategy.config is not None` when every other
+  test creates a strategy.
+
+### Keeping Tests DRY
+
+- **Use helper functions** for repeated setup patterns (e.g., `_make_strategy(video=_video_modality(...))` instead of
+  6-8 lines of config/strategy creation per test).
+- **Use named constants** for test dimensions (e.g., `VIDEO_SEQ_LEN`, `TOKENS_PER_FRAME`) instead of magic numbers.
+- **Merge tests that share identical setup** — when 5+ tests call `prepare_training_inputs` with the exact same
+  config and batch, each checking one assertion, merge them into one test that checks all assertions. Pytest reports
+  the exact failing line anyway.
+- **Use `@pytest.mark.parametrize`** for the same logic tested with different inputs (e.g., valid/invalid values for
+  a field).
+- **Use pytest fixtures** for shared batch data and test directories, but prefer explicit helper functions over
+  fixtures for strategy/config creation (makes the test self-documenting).
+
 ## Code Standards
 
 ### Type Hints
@@ -286,7 +358,12 @@ Key classes:
 
 - `LtxTrainerConfig` - Main configuration container
 - `ModelConfig` - Model paths, training mode (`lora` | `full`), checkpoint loading
-- `TrainingStrategyConfig` - Union of `TextToVideoConfig` | `VideoToVideoConfig` (discriminated by `name`)
+- `TrainingStrategyConfig` - Union of `FlexibleStrategyConfig` | `TextToVideoConfig` (deprecated) | `VideoToVideoConfig` (deprecated)
+- `FlexibleStrategyConfig` - Unified strategy config with `video`/`audio` `ModalityConfig` blocks
+- `ModalityConfig` - Per-modality config: `is_generated`, `latents_dir`, `conditions` list
+- `ConditionConfig` - Discriminated union: `FirstFrameConditionConfig`, `PrefixConditionConfig`, `SuffixConditionConfig`, `SpatialCropConditionConfig`, `MaskConditionConfig`, `ReferenceConditionConfig`
+- `ValidationSample` - Per-sample validation config with `prompt`, `conditions`, optional `video_dims`/`seed` overrides
+- `ValidationCondition` - Discriminated union for validation conditions (first_frame, prefix, suffix, spatial_crop, mask, reference, video_to_audio, audio_to_video)
 - `LoraConfig` - Rank, alpha, dropout, target modules
 - `OptimizationConfig` - Learning rate, batch size, gradient accumulation, scheduler, gradient checkpointing
 - `AccelerationConfig` - Mixed precision, quantization, 8-bit text encoder
@@ -310,21 +387,23 @@ Key classes:
 - Implements distributed training with Accelerate
 - Handles mixed precision, gradient accumulation, checkpointing
 - `_training_step()` applies embedding connectors then delegates to strategy
-- `_load_text_encoder_and_cache_embeddings()` loads the text encoder + embeddings processor, caches validation embeddings, then unloads the Gemma LLM (keeps only the embeddings processor connectors for training)
+- `_load_text_encoder_and_cache_embeddings()` loads the text encoder + embeddings processor, caches validation
+  embeddings, then unloads the Gemma LLM (keeps only the embeddings processor connectors for training)
 - Uses training strategies for mode-specific logic
 
 **`src/ltx_trainer/training_strategies/`** - Strategy pattern
 
 - `base_strategy.py`: `TrainingStrategy` ABC, `ModelInputs` dataclass
-- `text_to_video.py`: Standard text-to-video (with optional audio)
-- `video_to_video.py`: IC-LoRA video-to-video transformations
+- `flexible.py`: FlexibleStrategy — unified conditioning framework (recommended)
+- `text_to_video.py`: TextToVideoStrategy (deprecated — use FlexibleStrategy)
+- `video_to_video.py`: VideoToVideoStrategy (deprecated — use FlexibleStrategy)
 
 Key methods each strategy implements:
 
-- `get_data_sources()` - Required data directories
 - `prepare_training_inputs()` - Convert batch to `ModelInputs` with `Modality` objects
 - `compute_loss()` - Calculate training loss (velocity prediction, MSE with masking)
-- `requires_audio` property - Whether audio components needed
+
+The strategy's **config** declares its data directories via `get_data_sources()` (single source of truth, used for both dataset wiring and existence validation).
 
 **`src/ltx_trainer/model_loader.py`** - Model loading
 
@@ -339,14 +418,13 @@ Component loaders:
 - `load_embeddings_processor(checkpoint_path)` → `EmbeddingsProcessor` (feature extractor + connectors)
 - `load_model()` → `LtxModelComponents` (convenience wrapper)
 
-**`src/ltx_trainer/validation_sampler.py`** - Inference for validation
+**`src/ltx_trainer/validation_runner.py`** - Conditioned validation sampling
 
-Uses ltx-core components for denoising:
-
-- `LTX2Scheduler` for sigma scheduling
-- `EulerDiffusionStep` for diffusion steps
-- `CFGGuider` for classifier-free guidance
-- `STGGuider` for spatio-temporal guidance
+- Manages the full validation lifecycle: embedding caching, media encoding, denoising, decoding
+- Supports all validation condition types: first_frame, prefix, suffix, spatial_crop, mask, reference, video_to_audio, audio_to_video
+- Handles frozen modality paths (sigma=0 for conditioning modality)
+- Builds conditioning items using ltx-core's `VideoConditionByLatentIndex`, `VideoConditionByReferenceLatent`, `VideoConditionByMask`
+- Optional side-by-side reference output for IC-LoRA validation
 
 **`src/ltx_trainer/timestep_samplers.py`** - Flow matching timestep sampling
 
@@ -367,12 +445,18 @@ constructs the `GemmaTextEncoder` with quantized model, feature extractor, and e
 **`src/ltx_trainer/datasets.py`** - Dataset handling
 
 - `PrecomputedDataset` loads pre-computed VAE latents and text embeddings
-- Supports video latents, audio latents, text embeddings, reference latents (for IC-LoRA)
+- Supports video latents, audio latents, text embeddings, reference video latents, reference audio latents, video masks, and audio masks
 - Handles legacy patchified format `[seq_len, C]` → automatically unpatchifies to `[C, F, H, W]`
 - `DummyDataset` for benchmarking and minimal testing
 
 ## Common Development Tasks
 
+### Agent-Assisted Training
+
+When a user asks to train, fine-tune, create a LoRA, or produce a custom LTX-2 model, use the repository skill at
+[`.claude/skills/train-model`](../../.claude/skills/train-model/SKILL.md). The skill is the orchestrator for dataset probing, mode selection, preprocessing,
+training launch, monitoring, and post-train validation; it treats `packages/ltx-trainer/docs/` as the source of truth.
+
 ### Adding a New Configuration Parameter
 
 1. Add field to appropriate config class in `src/ltx_trainer/config.py`
@@ -382,10 +466,16 @@ constructs the `GemmaTextEncoder` with quantized model, feature extractor, and e
 
 ### Implementing a New Training Strategy
 
+The `FlexibleStrategy` now covers all use cases (T2V, T2A, I2V, V2V, A2A, AV2AV, inpainting, outpainting, extension, A2V, V2A, IC-LoRA) through
+configuration alone. A new strategy is only needed for fundamentally different training paradigms that cannot be
+expressed via `ModalityConfig` + `ConditionConfig` combinations.
+
+If you do need a new strategy:
+
 1. Create new file in `src/ltx_trainer/training_strategies/`
-2. Create config class inheriting `TrainingStrategyConfigBase`
+2. Create config class inheriting `TrainingStrategyConfigBase` and implement `get_data_sources()`
 3. Create strategy class inheriting `TrainingStrategy`
-4. Implement: `get_data_sources()`, `prepare_training_inputs()`, `compute_loss()`
+4. Implement: `prepare_training_inputs()`, `compute_loss()`
 5. Add to `__init__.py`: import, add to `TrainingStrategyConfig` union, update factory
 6. Add discriminator tag to config.py's `TrainingStrategyConfig`
 7. Create example config file in `configs/`
@@ -449,8 +539,8 @@ video_embeds, audio_embeds, binary_mask = text_encoder.embeddings_processor.crea
 
 - Validation errors: Check validators in `config.py`
 - Unknown fields: Config uses `extra="forbid"` — all fields must be defined
-- Strategy validation: IC-LoRA requires `reference_videos` in validation config
-- Video-to-video strategy requires `training_mode: "lora"`
+- FlexibleStrategy requires at least one modality with `is_generated: true`
+- Audio modality cannot use `first_frame` or `spatial_crop` conditions
 
 **Precomputed Data:**
 
@@ -480,7 +570,7 @@ Width and height must be divisible by 32.
 ### Platform Requirements
 
 - Linux required (uses `triton` which is Linux-only)
-- CUDA GPU with 24GB+ VRAM recommended (80GB+ for full fine-tuning)
+- CUDA GPU with 32GB+ VRAM recommended
 
 ## Reference: ltx-core Key Components
 

packages/ltx-trainer/README.md

@@ -1,8 +1,10 @@
 # LTX-2 Trainer
 
 This package provides tools and scripts for training and fine-tuning
-Lightricks' **LTX-2** audio-video generation model. It enables LoRA training, full
-fine-tuning, and training of video-to-video transformations (IC-LoRA) on custom datasets.
+Lightricks' **LTX-2** audio-video generation model. It supports LoRA training, full
+fine-tuning, and a flexible conditioning framework covering text-to-video, text-to-audio, image-to-video,
+video extension, audio extension, video inpainting, audio inpainting, video outpainting, IC-LoRA for video, audio, and joint
+audio-video references, audio-to-video, and video-to-audio.
 
 ---
 
@@ -17,9 +19,16 @@ All detailed guides and technical documentation are in the [docs](./docs/) direc
 - [🚀 Training Guide](docs/training-guide.md)
 - [🧪 Inference Guide](../ltx-pipelines/README.md)
 - [🔧 Utility Scripts](docs/utility-scripts.md)
+- [🧩 Custom Training Strategies](docs/custom-training-strategies.md)
 - [📚 LTX-Core Documentation](../ltx-core/README.md)
 - [🛡️ Troubleshooting Guide](docs/troubleshooting.md)
 
+### 🤖 Agent-Assisted Training
+
+Use the [`train-model`](../../.claude/skills/train-model/SKILL.md) repository skill for an end-to-end guided run:
+it probes your data and hardware, chooses the matching training mode, prepares/preprocesses the dataset, launches
+training, and monitors the job while using the docs above as the source of truth.
+
 ---
 
 ## 🔧 Requirements
@@ -28,7 +37,7 @@ All detailed guides and technical documentation are in the [docs](./docs/) direc
 - **Gemma Text Encoder** - Local Gemma model directory (required for LTX-2)
 - **Linux with CUDA** - CUDA 13+ recommended for optimal performance
 - **Nvidia GPU with 80GB+ VRAM** - Recommended for the standard config. For GPUs with 32GB VRAM (e.g., RTX 5090),
-  use the [low VRAM config](configs/ltx2_av_lora_low_vram.yaml) which enables INT8 quantization and other
+  use the [low VRAM config](configs/t2v_lora_low_vram.yaml) which enables INT8 quantization and other
   memory optimizations
 
 ---

packages/ltx-trainer/configs/a2a_ic_lora.yaml

@@ -0,0 +1,340 @@
+# =============================================================================
+# LTX-2 Audio-to-Audio IC-LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training In-Context LoRA (IC-LoRA) adapters that
+# enable audio-to-audio transformations. IC-LoRA learns to apply audio
+# transformations (e.g., style transfer, voice conversion, sound effects, etc.)
+# by conditioning on reference audio.
+#
+# Key differences from text-to-video LoRA:
+#   - Uses reference audio as conditioning input alongside text prompts
+#   - Requires preprocessed reference audio latents in addition to target latents
+#   - Audio-only training (no video modality)
+#   - Validation requires reference audio to demonstrate the transformation
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── conditions/              # Text embeddings for each sample
+#   ├── audio_latents/           # Target audio latents (what the model learns to generate)
+#   └── reference_audio_latents/ # Reference audio latents (conditioning input)
+#
+# Dataset metadata columns: audio, reference_audio, caption
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 16-32 for IC-LoRA.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES (not used for audio-only IC-LoRA):
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (not used for audio-only IC-LoRA):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # For audio-only IC-LoRA, we explicitly target audio modules.
+  # Including audio FFN layers often improves transformation quality.
+  target_modules:
+    # Audio self-attention
+    - "audio_attn1.to_k"
+    - "audio_attn1.to_q"
+    - "audio_attn1.to_v"
+    - "audio_attn1.to_out.0"
+    # Audio cross-attention to text
+    - "audio_attn2.to_k"
+    - "audio_attn2.to_q"
+    - "audio_attn2.to_v"
+    - "audio_attn2.to_out.0"
+    # Audio feed-forward (often improves transformation quality)
+    - "audio_ff.net.0.proj"
+    - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the audio-to-audio (IC-LoRA) training approach using the unified
+# flexible strategy. Reference conditioning concatenates pre-encoded reference
+# audio latents to the target sequence. Reference tokens participate in
+# bidirectional self-attention but receive no noise and are excluded from loss.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Audio modality configuration (audio-only, no video)
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing target audio latents
+    latents_dir: "audio_latents"
+
+    # Conditions applied to the audio modality during training
+    conditions:
+      # Reference conditioning (IC-LoRA): concatenates pre-encoded reference audio
+      # latents to the target sequence. The model learns to transform the reference
+      # into the target audio based on the text prompt.
+      - type: reference
+        # Directory name (within preprocessed_data_root) containing reference audio latents
+        # These are the conditioning inputs that guide the transformation
+        latents_dir: "reference_audio_latents"
+        # Probability of applying reference conditioning per training sample
+        probability: 1.0
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 2e-4
+
+  # Total number of training steps
+  steps: 3000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: audio_latents/, conditions/, and reference_audio_latents/ subdirectories
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # For IC-LoRA, each sample includes a reference condition pointing to the conditioning audio.
+  samples:
+    - prompt: >-
+        A warm and soothing piano melody with soft ambient textures, reminiscent of a quiet
+        evening by the fireplace. Gentle reverb creates a sense of intimate space.
+      conditions:
+        - type: reference
+          audio: "/path/to/reference_audio_1.wav"
+    - prompt: >-
+        Energetic electronic beats with pulsing synthesizer leads and crisp hi-hat patterns.
+        Deep bass tones provide a driving rhythm underneath bright melodic arpeggios.
+      conditions:
+        - type: reference
+          audio: "/path/to/reference_audio_2.wav"
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Generation length control [width, height, frames]
+  # With generate_video=false, width/height are unused; frames and frame_rate set audio duration.
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 512, 512, 81 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" skips both audio and video self-attention (suited to audio validation)
+
+  # Whether to generate audio in validation samples
+  # Can be enabled even when not training the audio branch
+  generate_audio: true
+
+  # Whether to generate video in validation samples
+  # Disabled for audio-only IC-LoRA since no video modality is configured
+  generate_video: false
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: 3
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-ic-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "ic-lora", "a2a" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/a2a_ic_lora"

packages/ltx-trainer/configs/a2v_lora.yaml

@@ -0,0 +1,344 @@
+# =============================================================================
+# LTX-2 Audio-to-Video LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# audio-to-video generation. The model learns to generate videos conditioned
+# on a frozen audio signal via the transformer's built-in cross-modal attention.
+#
+# In this mode, audio is provided as a frozen (clean, no noise, no loss)
+# conditioning signal. The video modality is the only generated output.
+# Audio influences video generation through the transformer's audio-to-video
+# cross-attention mechanism.
+#
+# Use this configuration when you want to:
+#   - Generate videos driven by audio content (e.g., music visualizations)
+#   - Train audio-reactive video generation
+#   - Create models that synchronize video with given audio
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/        # Video latents (VAE-encoded videos)
+#   ├── conditions/     # Text embeddings for each video
+#   └── audio_latents/  # Audio latents (frozen conditioning input)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general fine-tuning.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # Using short patterns like "to_k" matches ALL attention modules (video, audio, and cross-modal).
+  # For audio-video training, this is the recommended approach.
+  target_modules:
+    # Attention layers (matches both video and audio branches)
+    - "to_k"
+    - "to_q"
+    - "to_v"
+    - "to_out.0"
+    # Uncomment below to also train feed-forward layers (can increase the LoRA's capacity):
+    # - "ff.net.0.proj"
+    # - "ff.net.2"
+    # - "audio_ff.net.0.proj"
+    # - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the audio-to-video training approach using the unified flexible strategy.
+# Audio is frozen (no noise, sigma=0, excluded from loss) and conditions video
+# generation via the transformer's built-in cross-modal attention.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  # Video is the generated (denoised) output
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
+
+  # Audio modality configuration
+  # Audio is frozen — it acts as conditioning for video generation
+  # Frozen modalities get sigma=0, timestep=0, no noise, and no loss
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    # When false, audio is passed through the transformer clean and influences
+    # video via cross-modal attention
+    is_generated: false
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 1e-4
+
+  # Total number of training steps
+  steps: 2000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, audio_latents/, and conditions/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        A musician plays an acoustic guitar in a dimly lit recording studio, fingers moving
+        across the fretboard with practiced ease. The warm amber light from a desk lamp
+        illuminates the wooden guitar body. Sound-absorbing panels line the walls, and a
+        microphone stands nearby on a boom arm.
+      conditions:
+        - type: audio_to_video
+          audio: "/path/to/conditioning_audio_1.wav"
+    - prompt: >-
+        Rain falls steadily on a cobblestone street in a European town at dusk. Puddles form
+        between the stones, creating ripples as new drops land. Old brick buildings line both
+        sides of the narrow street, their facades glistening with moisture under warm
+        streetlights.
+      conditions:
+        - type: audio_to_video
+          audio: "/path/to/conditioning_audio_2.wav"
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 576, 576, 89 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_v"  # "stg_av" perturbs both audio and video, "stg_v" video only
+
+  # Whether to generate audio in validation samples
+  # Independent of training_strategy.audio.is_generated - you can generate audio
+  # in validation even when not training the audio branch
+  # For this audio-to-video config, false matches frozen audio conditioning (video-only synthesis).
+  generate_audio: false
+
+  # Generate video from the frozen audio condition
+  generate_video: true
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: -1
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "a2v" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/a2v_lora"

packages/ltx-trainer/configs/audio_extend_lora.yaml

@@ -0,0 +1,343 @@
+# =============================================================================
+# LTX-2 Audio Extension (Forward) LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# forward audio extension. The model learns to continue audio forward in time
+# by conditioning on a prefix of existing audio latent timesteps.
+#
+# Prefix conditioning works by providing the first N audio latent timesteps as
+# clean conditioning signals during training — they receive no noise,
+# timestep=0, and are excluded from the loss. The model learns to generate
+# audio that seamlessly continues from the given prefix.
+#
+# This is an audio-only training mode — no video modality is configured.
+#
+# Use this configuration when you want to:
+#   - Train the model to extend/continue existing audio
+#   - Fine-tune audio temporal continuation capabilities
+#   - Create seamless audio extension models
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── conditions/      # Text embeddings for each sample
+#   └── audio_latents/   # Audio latents (VAE-encoded audio)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for audio extension LoRA training.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES (not used for audio-only modes):
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (not used for audio-only modes):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # For audio-only extension, we explicitly target audio modules.
+  # Including audio FFN layers can increase the LoRA's capacity.
+  target_modules:
+    # Audio self-attention
+    - "audio_attn1.to_k"
+    - "audio_attn1.to_q"
+    - "audio_attn1.to_v"
+    - "audio_attn1.to_out.0"
+    # Audio cross-attention to text
+    - "audio_attn2.to_k"
+    - "audio_attn2.to_q"
+    - "audio_attn2.to_v"
+    - "audio_attn2.to_out.0"
+    # Audio feed-forward (often improves transformation quality)
+    - "audio_ff.net.0.proj"
+    - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the audio extension training approach using the unified flexible
+# strategy. Prefix conditioning provides the first N audio latent timesteps as
+# clean conditioning, teaching the model to generate temporal continuations.
+training_strategy:
+  name: "flexible"
+
+  # Audio modality configuration (audio-only, no video)
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+    # Conditions applied to the audio modality during training
+    conditions:
+      - type: prefix
+        # Number of audio latent timesteps to use as conditioning prefix
+        # Each audio latent timestep = 1 patchified token
+        temporal_boundary: 8
+        # Probability of applying prefix conditioning per training sample
+        # At 1.0, all training samples use audio extension mode
+        probability: 1.0
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 2e-4
+
+  # Total number of training steps
+  steps: 3000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: conditions/ and audio_latents/ subdirectories
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Each sample includes an audio prefix condition for forward audio extension.
+  samples:
+    - prompt: >-
+        A warm and soothing piano melody with soft ambient textures, reminiscent of a quiet
+        evening by the fireplace. Gentle reverb creates a sense of intimate space.
+      conditions:
+        - type: prefix
+          audio: "/path/to/prefix_audio_1.wav"
+          # Duration is in seconds; choose a value that covers a comparable audio context
+          # to the training temporal_boundary measured in audio latent timesteps.
+          duration: 1.0
+    - prompt: >-
+        Energetic electronic beats with pulsing synthesizer leads and crisp hi-hat patterns.
+        Deep bass tones provide a driving rhythm underneath bright melodic arpeggios.
+      conditions:
+        - type: prefix
+          audio: "/path/to/prefix_audio_2.wav"
+          # Duration is in seconds; choose a value that covers a comparable audio context
+          # to the training temporal_boundary measured in audio latent timesteps.
+          duration: 1.0
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Generation length control [width, height, frames]
+  # With generate_video=false, width/height are unused; frames and frame_rate set audio duration.
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 512, 512, 81 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" skips both audio and video self-attention (suited to audio validation)
+
+  # Whether to generate audio in validation samples
+  # Enabled because this audio-only config generates audio
+  generate_audio: true
+
+  # Whether to generate video in validation samples
+  # Disabled because no video modality is configured
+  generate_video: false
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: 3
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-ic-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "audio-extension", "audio-only" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/audio_extend_lora"

packages/ltx-trainer/configs/audio_inpainting_lora.yaml

@@ -0,0 +1,324 @@
+# =============================================================================
+# LTX-2 Audio Inpainting LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# audio inpainting. The model learns to fill in masked regions of audio
+# using per-sample binary masks that define which regions are conditioning
+# (provided clean) and which regions the model must generate.
+#
+# Mask conditioning works by loading per-sample binary masks from disk.
+# Masked regions receive clean latents (no noise, timestep=0) and are excluded
+# from the loss. Unmasked regions are noised and trained normally.
+#
+# This is an audio-only training mode — no video modality is configured.
+#
+# Use this configuration when you want to:
+#   - Train the model to fill in or replace regions of existing audio
+#   - Fine-tune audio inpainting capabilities on custom datasets
+#   - Create audio restoration or editing models
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── conditions/      # Text embeddings for each sample
+#   ├── audio_latents/   # Audio latents (VAE-encoded audio)
+#   └── audio_masks/     # Per-sample binary masks defining conditioning regions
+#
+# Dataset metadata columns: audio, audio_mask, caption
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for audio inpainting LoRA training.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # For audio-only inpainting, we explicitly target audio modules.
+  # Including audio FFN layers can increase the LoRA's capacity.
+  target_modules:
+    # Audio self-attention
+    - "audio_attn1.to_k"
+    - "audio_attn1.to_q"
+    - "audio_attn1.to_v"
+    - "audio_attn1.to_out.0"
+    # Audio cross-attention to text
+    - "audio_attn2.to_k"
+    - "audio_attn2.to_q"
+    - "audio_attn2.to_v"
+    - "audio_attn2.to_out.0"
+    # Audio feed-forward (often improves transformation quality)
+    - "audio_ff.net.0.proj"
+    - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the audio inpainting training approach using the unified flexible
+# strategy. Per-sample binary masks define which audio regions are provided
+# as clean conditioning and which regions the model must learn to generate.
+training_strategy:
+  name: "flexible"
+
+  # Audio modality configuration (audio-only, no video)
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+    # Conditions applied to the audio modality during training
+    conditions:
+      - type: mask
+        # Directory name (within preprocessed_data_root) containing binary masks
+        # Each mask file corresponds to a training sample and defines the
+        # conditioning region (mask=1 means conditioning, mask=0 means generate)
+        mask_dir: "audio_masks"
+        # Probability of applying mask conditioning per training sample
+        # At 1.0, all training samples use inpainting mode
+        probability: 1.0
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 2e-4
+
+  # Total number of training steps
+  steps: 3000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: conditions/, audio_latents/, and audio_masks/ subdirectories
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # For audio inpainting, each sample includes mask conditioning (audio + mask paths).
+  samples:
+    - prompt: >-
+        A warm and soothing piano melody with soft ambient textures, reminiscent of a quiet
+        evening by the fireplace. Gentle reverb creates a sense of intimate space.
+      conditions:
+        - type: mask
+          audio: "/path/to/inpainting_audio_1.wav"
+          mask: "/path/to/inpainting_mask_1.pt"
+    - prompt: >-
+        Energetic electronic beats with pulsing synthesizer leads and crisp hi-hat patterns.
+        Deep bass tones provide a driving rhythm underneath bright melodic arpeggios.
+      conditions:
+        - type: mask
+          audio: "/path/to/inpainting_audio_2.wav"
+          mask: "/path/to/inpainting_mask_2.pt"
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Generation length control [width, height, frames]
+  # With generate_video=false, width/height are unused; frames and frame_rate set audio duration.
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 512, 512, 81 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" skips both audio and video self-attention (suited to audio validation)
+
+  # Whether to generate audio in validation samples
+  # Enabled because this audio-only config generates audio
+  generate_audio: true
+
+  # Whether to generate video in validation samples
+  # Disabled because no video modality is configured
+  generate_video: false
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: 3
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-ic-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "audio-inpainting", "audio-only" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/audio_inpainting_lora"

packages/ltx-trainer/configs/audio_suffix_lora.yaml

@@ -0,0 +1,343 @@
+# =============================================================================
+# LTX-2 Audio Extension (Backward) LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# backward audio extension. The model learns to generate audio content leading into existing audio
+# by conditioning on a suffix of existing audio latent timesteps.
+#
+# Suffix conditioning works by providing the last N audio latent timesteps as
+# clean conditioning signals during training — they receive no noise,
+# timestep=0, and are excluded from the loss. The model learns to generate
+# audio that seamlessly leads into the given suffix.
+#
+# This is an audio-only training mode — no video modality is configured.
+#
+# Use this configuration when you want to:
+#   - Train the model to generate preceding content for existing audio
+#   - Fine-tune audio temporal continuation capabilities
+#   - Create seamless audio extension models
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── conditions/      # Text embeddings for each sample
+#   └── audio_latents/   # Audio latents (VAE-encoded audio)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for audio extension LoRA training.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES (not used for audio-only modes):
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (not used for audio-only modes):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # For audio-only extension, we explicitly target audio modules.
+  # Including audio FFN layers can increase the LoRA's capacity.
+  target_modules:
+    # Audio self-attention
+    - "audio_attn1.to_k"
+    - "audio_attn1.to_q"
+    - "audio_attn1.to_v"
+    - "audio_attn1.to_out.0"
+    # Audio cross-attention to text
+    - "audio_attn2.to_k"
+    - "audio_attn2.to_q"
+    - "audio_attn2.to_v"
+    - "audio_attn2.to_out.0"
+    # Audio feed-forward (often improves transformation quality)
+    - "audio_ff.net.0.proj"
+    - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the backward audio extension training approach using the unified flexible
+# strategy. Suffix conditioning provides the last N audio latent timesteps as
+# clean conditioning, teaching the model to generate content leading into existing audio.
+training_strategy:
+  name: "flexible"
+
+  # Audio modality configuration (audio-only, no video)
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+    # Conditions applied to the audio modality during training
+    conditions:
+      - type: suffix
+        # Number of audio latent timesteps to use as conditioning suffix
+        # Each audio latent timestep = 1 patchified token
+        temporal_boundary: 8
+        # Probability of applying suffix conditioning per training sample
+        # At 1.0, all training samples use audio extension mode
+        probability: 1.0
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 2e-4
+
+  # Total number of training steps
+  steps: 3000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: conditions/ and audio_latents/ subdirectories
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Each sample includes an audio suffix condition for backward audio extension.
+  samples:
+    - prompt: >-
+        A warm and soothing piano melody with soft ambient textures, reminiscent of a quiet
+        evening by the fireplace. Gentle reverb creates a sense of intimate space.
+      conditions:
+        - type: suffix
+          audio: "/path/to/suffix_audio_1.wav"
+          # Duration is in seconds; choose a value that covers a comparable audio context
+          # to the training temporal_boundary measured in audio latent timesteps.
+          duration: 1.0
+    - prompt: >-
+        Energetic electronic beats with pulsing synthesizer leads and crisp hi-hat patterns.
+        Deep bass tones provide a driving rhythm underneath bright melodic arpeggios.
+      conditions:
+        - type: suffix
+          audio: "/path/to/suffix_audio_2.wav"
+          # Duration is in seconds; choose a value that covers a comparable audio context
+          # to the training temporal_boundary measured in audio latent timesteps.
+          duration: 1.0
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Generation length control [width, height, frames]
+  # With generate_video=false, width/height are unused; frames and frame_rate set audio duration.
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 512, 512, 81 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" skips both audio and video self-attention (suited to audio validation)
+
+  # Whether to generate audio in validation samples
+  # Enabled because this audio-only config generates audio
+  generate_audio: true
+
+  # Whether to generate video in validation samples
+  # Disabled because no video modality is configured
+  generate_video: false
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: 3
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-ic-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "audio-suffix", "audio-only" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/audio_suffix_lora"

packages/ltx-trainer/configs/av2av_ic_lora.yaml

@@ -0,0 +1,342 @@
+# =============================================================================
+# LTX-2 AV2AV IC-LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training In-Context LoRA (IC-LoRA) adapters that
+# enable joint audio-video-to-audio-video transformations. IC-LoRA learns to
+# apply transformations to both the video and audio modalities simultaneously
+# by conditioning on paired reference video and audio.
+#
+# Both modalities use reference conditioning: pre-encoded reference latents
+# are concatenated to each modality's target sequence. Reference tokens
+# participate in bidirectional self-attention but receive no noise and are
+# excluded from the loss.
+#
+# Key differences from video-only IC-LoRA (v2v_ic_lora.yaml):
+#   - Both video AND audio have reference conditions
+#   - Requires preprocessed reference latents for BOTH modalities
+#   - LoRA targets all modules (video, audio, and cross-modal attention)
+#   - Validation uses both video and audio reference conditions
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/                 # Target video latents
+#   ├── audio_latents/           # Target audio latents
+#   ├── conditions/              # Text embeddings
+#   ├── reference_latents/       # Reference video latents (conditioning input)
+#   └── reference_audio_latents/ # Reference audio latents (conditioning input)
+#
+# Dataset metadata columns: video, audio, reference_video, reference_audio, caption
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # IC-LoRA reference conditioning is intended for LoRA adapter training.
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 16-32 for IC-LoRA.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # For AV2AV IC-LoRA, we target ALL modules — video, audio, and cross-modal attention.
+  # Using short patterns matches all branches simultaneously.
+  target_modules:
+    # Attention layers (matches video, audio, and cross-modal branches)
+    - "to_k"
+    - "to_q"
+    - "to_v"
+    - "to_out.0"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the AV2AV IC-LoRA training approach using the unified flexible
+# strategy. Both video and audio modalities have reference conditioning,
+# enabling joint audiovisual transformations.
+training_strategy:
+  name: "flexible"
+
+  # Video modality configuration
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing target video latents
+    latents_dir: "latents"
+
+    # Conditions applied to the video modality during training
+    conditions:
+      # Reference conditioning (IC-LoRA): concatenates pre-encoded reference video
+      # latents to the target sequence
+      - type: reference
+        latents_dir: "reference_latents"
+        probability: 1.0
+
+  # Audio modality configuration
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing target audio latents
+    latents_dir: "audio_latents"
+
+    # Conditions applied to the audio modality during training
+    conditions:
+      # Reference conditioning (IC-LoRA): concatenates pre-encoded reference audio
+      # latents to the target sequence
+      - type: reference
+        latents_dir: "reference_audio_latents"
+        probability: 1.0
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 2e-4
+
+  # Total number of training steps
+  steps: 3000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, audio_latents/, conditions/, reference_latents/, and reference_audio_latents/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # For AV2AV IC-LoRA, each sample includes reference video and reference audio conditions.
+  samples:
+    - prompt: >-
+        A man in a casual blue jacket walks along a winding path through a lush green park on a
+        bright sunny afternoon. Tall oak trees line the pathway, their leaves rustling gently in
+        the breeze. Dappled sunlight creates shifting patterns on the ground as he strolls at a
+        relaxed pace, occasionally looking up at the scenery around him. The audio captures
+        footsteps on gravel, birds singing in the trees, distant children playing, and the soft
+        whisper of wind through the foliage.
+      conditions:
+        - type: reference
+          video: "/path/to/reference_video_1.mp4"
+          downscale_factor: 1
+          temporal_scale_factor: 1
+          include_in_output: true
+        - type: reference
+          audio: "/path/to/reference_audio_1.wav"
+    - prompt: >-
+        A fluffy orange tabby cat sits perfectly still on a wooden windowsill, its green eyes
+        intently tracking small birds hopping on a branch just outside the glass. The cat's ears
+        twitch and rotate, following every movement. Warm afternoon light illuminates its fur,
+        creating a soft golden glow. Behind the cat, a cozy living room with a bookshelf and
+        houseplants is visible. The audio features gentle purring, occasional soft meows, muffled
+        bird chirps through the window, and quiet ambient room sounds.
+      conditions:
+        - type: reference
+          video: "/path/to/reference_video_2.mp4"
+          downscale_factor: 1
+          temporal_scale_factor: 1
+          include_in_output: true
+        - type: reference
+          audio: "/path/to/reference_audio_2.wav"
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 512, 512, 81 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # Both video and audio modalities are trained
+
+  # Whether to generate audio in validation samples
+  # Can be enabled even when not training the audio branch
+  generate_audio: true
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: 3
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-ic-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "ic-lora", "av2av" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/av2av_ic_lora"

packages/ltx-trainer/configs/i2v_lora.yaml

@@ -0,0 +1,348 @@
+# =============================================================================
+# LTX-2 Image-to-Video LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# image-to-video generation. The model learns to generate videos conditioned
+# on a starting image (first frame), with optional audio generation.
+#
+# First-frame conditioning works by providing the first frame as a clean
+# conditioning signal during training — it receives no noise, timestep=0,
+# and is excluded from the loss. This teaches the model to animate from
+# a given image.
+#
+# Use this configuration when you want to:
+#   - Train the model to generate videos starting from a given image
+#   - Fine-tune image-to-video capabilities on custom datasets
+#   - Create image animation models with optional audio
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/        # Video latents (VAE-encoded videos)
+#   ├── conditions/     # Text embeddings for each video
+#   └── audio_latents/  # Audio latents (VAE-encoded audio)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general fine-tuning.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # Using short patterns like "to_k" matches ALL attention modules (video, audio, and cross-modal).
+  # For audio-video training, this is the recommended approach.
+  target_modules:
+    # Attention layers (matches both video and audio branches)
+    - "to_k"
+    - "to_q"
+    - "to_v"
+    - "to_out.0"
+    # Uncomment below to also train feed-forward layers (can increase the LoRA's capacity):
+    # - "ff.net.0.proj"
+    # - "ff.net.2"
+    # - "audio_ff.net.0.proj"
+    # - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the image-to-video training approach using the unified flexible strategy.
+# First-frame conditioning provides the first frame as clean conditioning signal
+# during training, teaching the model to animate from a given image.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
+
+    # Conditions applied to the video modality during training
+    # First-frame conditioning: the first frame of each video is provided as a clean
+    # conditioning signal (no noise, timestep=0, excluded from loss)
+    conditions:
+      - type: first_frame
+        # Probability of applying first-frame conditioning per training sample
+        # At 0.5, half the training samples use I2V mode, half use pure T2V
+        # Higher values improve I2V quality but may reduce T2V diversity
+        probability: 0.5
+
+  # Audio modality configuration
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 1e-4
+
+  # Total number of training steps
+  steps: 2000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, conditions/, and audio_latents/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a
+        laptop while occasionally glancing at notes beside her. Soft natural light streams through
+        a large window, casting warm shadows across the room. She pauses to take a sip from a
+        ceramic mug, then continues working with focused concentration. The audio captures the
+        gentle clicking of keyboard keys, the soft rustle of papers, and ambient room tone with
+        occasional distant bird chirps from outside.
+      conditions:
+        - type: first_frame
+          image_or_video: "/path/to/conditioning_image_1.png"
+    - prompt: >-
+        A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet
+        dish with precise movements. Steam rises from freshly cooked vegetables as he arranges
+        them with tweezers. The stainless steel surfaces gleam under bright overhead lights, and
+        various pots simmer on the stove behind him. The audio features the sizzling of pans,
+        the clinking of utensils against plates, and the ambient hum of kitchen ventilation.
+      conditions:
+        - type: first_frame
+          image_or_video: "/path/to/conditioning_image_2.png"
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 576, 576, 89 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" perturbs both audio and video, "stg_v" video only
+
+  # Whether to generate audio in validation samples
+  # Independent of training_strategy.audio.is_generated - you can generate audio
+  # in validation even when not training the audio branch
+  generate_audio: true
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: -1
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "i2v" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/i2v_lora"

packages/ltx-trainer/configs/t2a_lora.yaml

@@ -0,0 +1,317 @@
+# =============================================================================
+# LTX-2 Text-to-Audio LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# text-to-audio generation. The model learns to generate audio from text
+# prompts without any additional conditioning.
+#
+# This is the simplest audio-only training mode — no reference audio, no
+# video modality. Only the audio branch of the transformer is trained.
+#
+# Use this configuration when you want to:
+#   - Fine-tune audio generation for specific sound styles or domains
+#   - Train custom audio generation capabilities
+#   - Create audio LoRAs that can be combined with video LoRAs
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── conditions/      # Text embeddings for each sample
+#   └── audio_latents/   # Audio latents (VAE-encoded audio)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general audio LoRA training.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES (not used for audio-only modes):
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (not used for audio-only modes):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # For audio-only training, we explicitly target audio modules.
+  # Including audio FFN layers can increase the LoRA's capacity.
+  target_modules:
+    # Audio self-attention
+    - "audio_attn1.to_k"
+    - "audio_attn1.to_q"
+    - "audio_attn1.to_v"
+    - "audio_attn1.to_out.0"
+    # Audio cross-attention to text
+    - "audio_attn2.to_k"
+    - "audio_attn2.to_q"
+    - "audio_attn2.to_v"
+    - "audio_attn2.to_out.0"
+    # Audio feed-forward (often improves transformation quality)
+    - "audio_ff.net.0.proj"
+    - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the text-to-audio training approach using the unified flexible
+# strategy. Audio-only training with no additional conditions — the model
+# learns to generate audio purely from text prompts.
+training_strategy:
+  name: "flexible"
+
+  # Audio modality configuration (audio-only, no video)
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 2e-4
+
+  # Total number of training steps
+  steps: 3000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: conditions/ and audio_latents/ subdirectories
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Text-to-audio validation samples do not need additional conditions.
+  samples:
+    - prompt: >-
+        A warm and soothing piano melody with soft ambient textures, reminiscent of a quiet
+        evening by the fireplace. Gentle reverb creates a sense of intimate space.
+    - prompt: >-
+        Energetic electronic beats with pulsing synthesizer leads and crisp hi-hat patterns.
+        Deep bass tones provide a driving rhythm underneath bright melodic arpeggios.
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Generation length control [width, height, frames]
+  # With generate_video=false, width/height are unused; frames and frame_rate set audio duration.
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 512, 512, 81 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" skips both audio and video self-attention (suited to audio validation)
+
+  # Whether to generate audio in validation samples
+  # Enabled because this audio-only config generates audio
+  generate_audio: true
+
+  # Whether to generate video in validation samples
+  # Disabled because no video modality is configured
+  generate_video: false
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: 3
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-ic-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "t2a", "audio-only" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/t2a_lora"

packages/ltx-trainer/configs/t2v_lora.yaml

@@ -1,21 +1,20 @@
 # =============================================================================
-# LTX-2 Audio-Video LoRA Training Configuration
+# LTX-2 Text-to-Video LoRA Training Configuration
 # =============================================================================
 #
 # This configuration is for training LoRA adapters on the LTX-2 model for
-# text-to-video generation. It supports both video-only and joint audio-video
-# training modes.
+# text-to-video generation with joint audio-video support.
 #
 # Use this configuration when you want to:
 #   - Fine-tune LTX-2 on your own video dataset
-#   - Train with or without audio generation
+#   - Train joint audio-video generation from text prompts
 #   - Create custom video generation styles or audiovisual concepts
 #
-# Dataset structure for text-to-video training:
+# Dataset structure:
 #   preprocessed_data_root/
 #   ├── latents/        # Video latents (VAE-encoded videos)
 #   ├── conditions/     # Text embeddings for each video
-#   └── audio_latents/  # Audio latents (only if with_audio: true)
+#   └── audio_latents/  # Audio latents (VAE-encoded audio)
 #
 # =============================================================================
 
@@ -93,24 +92,29 @@ lora:
 # -----------------------------------------------------------------------------
 # Training Strategy Configuration
 # -----------------------------------------------------------------------------
-# Defines the text-to-video training approach.
+# Defines the training approach using the unified flexible strategy.
+# This configuration trains both video and audio generation from text prompts.
 training_strategy:
-  # Strategy name: "text_to_video" for standard text-to-video training
-  name: "text_to_video"
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
 
-  # Probability of conditioning on the first frame during training
-  # Higher values train the model to perform better in image-to-video (I2V) mode,
-  # where a clean first frame is provided and the model generates the rest of the video
-  # Increase this value to train the model to perform better in image-to-video (I2V) mode
-  first_frame_conditioning_p: 0.5
+  # Video modality configuration
+  # When is_generated is true, the model learns to generate (denoise) video
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
 
-  # Enable joint audio-video training
-  # Set to true if your dataset includes audio and you want to train the audio branch
-  with_audio: true
-
-  # Directory name (within preprocessed_data_root) containing audio latents
-  # Only used when with_audio is true
-  audio_latents_dir: "audio_latents"
+  # Audio modality configuration
+  # When is_generated is true, the model learns to generate (denoise) audio
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
 
 # -----------------------------------------------------------------------------
 # Optimization Configuration
@@ -166,18 +170,19 @@ 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 state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
   offload_optimizer_during_validation: false
 
+
 # -----------------------------------------------------------------------------
 # Data Configuration
 # -----------------------------------------------------------------------------
 # Specifies the training data location and loading parameters.
 data:
   # Root directory containing preprocessed training data
-  # Should contain: latents/, conditions/, and optionally audio_latents/
+  # Should contain: latents/, conditions/, and audio_latents/
   preprocessed_data_root: "/path/to/preprocessed/data"
 
   # Number of worker processes for data loading
@@ -187,24 +192,31 @@ data:
 # -----------------------------------------------------------------------------
 # Validation Configuration
 # -----------------------------------------------------------------------------
-# Controls validation video generation during training.
+# Controls validation sampling during training.
 # NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
 # maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
 validation:
-  # Text prompts for validation video generation
-  # Provide prompts representative of your training data
-  # LTX-2 prefers longer, detailed prompts that describe both visual content and audio
-  prompts:
-    - "A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a laptop while occasionally glancing at notes beside her. Soft natural light streams through a large window, casting warm shadows across the room. She pauses to take a sip from a ceramic mug, then continues working with focused concentration. The audio captures the gentle clicking of keyboard keys, the soft rustle of papers, and ambient room tone with occasional distant bird chirps from outside."
-    - "A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet dish with precise movements. Steam rises from freshly cooked vegetables as he arranges them with tweezers. The stainless steel surfaces gleam under bright overhead lights, and various pots simmer on the stove behind him. The audio features the sizzling of pans, the clinking of utensils against plates, and the ambient hum of kitchen ventilation."
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a
+        laptop while occasionally glancing at notes beside her. Soft natural light streams through
+        a large window, casting warm shadows across the room. She pauses to take a sip from a
+        ceramic mug, then continues working with focused concentration. The audio captures the
+        gentle clicking of keyboard keys, the soft rustle of papers, and ambient room tone with
+        occasional distant bird chirps from outside.
+    - prompt: >-
+        A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet
+        dish with precise movements. Steam rises from freshly cooked vegetables as he arranges
+        them with tweezers. The stainless steel surfaces gleam under bright overhead lights, and
+        various pots simmer on the stove behind him. The audio features the sizzling of pans,
+        the clinking of utensils against plates, and the ambient hum of kitchen ventilation.
 
   # Negative prompt to avoid unwanted artifacts
   negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
 
-  # Optional: First frame images for image-to-video validation
-  # If provided, must have one image per prompt
-  images: null
-
   # Output video dimensions [width, height, frames]
   # Width and height must be divisible by 32
   # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
@@ -235,7 +247,7 @@ validation:
   stg_mode: "stg_av"  # "stg_av" perturbs both audio and video, "stg_v" video only
 
   # Whether to generate audio in validation samples
-  # Independent of training_strategy.with_audio - you can generate audio
+  # Independent of training_strategy.audio.is_generated - you can generate audio
   # in validation even when not training the audio branch
   generate_audio: true
 
@@ -298,9 +310,9 @@ wandb:
   entity: null
 
   # Tags to help organize runs
-  tags: [ "ltx2", "lora" ]
+  tags: [ "ltx2", "lora", "t2v" ]
 
-  # Log validation videos to W&B
+  # Log validation media (video/audio) to W&B
   log_validation_videos: true
 
 # -----------------------------------------------------------------------------
@@ -312,4 +324,4 @@ wandb:
 seed: 42
 
 # Directory to save outputs (checkpoints, validation videos, logs)
-output_dir: "outputs/ltx2_av_lora"
+output_dir: "outputs/t2v_lora"

packages/ltx-trainer/configs/t2v_lora_low_vram.yaml

@@ -1,8 +1,8 @@
 # =============================================================================
-# LTX-2 Audio-Video LoRA Training Configuration (Low VRAM)
+# LTX-2 Text-to-Video LoRA Training Configuration (Low VRAM)
 # =============================================================================
 #
-# This is a memory-optimized variant of the standard audio-video LoRA config.
+# This is a memory-optimized variant of the standard text-to-video LoRA config.
 # It uses 8-bit optimizer, int8 quantization, and reduced LoRA rank to minimize
 # GPU memory usage while maintaining good training quality.
 #
@@ -15,15 +15,15 @@
 # Recommended for GPUs with 32GB VRAM (e.g., RTX 5090).
 #
 # Use this configuration when you want to:
-#   - Fine-tune LTX-2 on your own video dataset
-#   - Train with or without audio generation
+#   - Fine-tune LTX-2 on your own video dataset with limited GPU memory
+#   - Train joint audio-video generation from text prompts
 #   - Create custom video generation styles or audiovisual concepts
 #
-# Dataset structure for text-to-video training:
+# Dataset structure:
 #   preprocessed_data_root/
 #   ├── latents/        # Video latents (VAE-encoded videos)
 #   ├── conditions/     # Text embeddings for each video
-#   └── audio_latents/  # Audio latents (only if with_audio: true)
+#   └── audio_latents/  # Audio latents (VAE-encoded audio)
 #
 # =============================================================================
 
@@ -103,24 +103,29 @@ lora:
 # -----------------------------------------------------------------------------
 # Training Strategy Configuration
 # -----------------------------------------------------------------------------
-# Defines the text-to-video training approach.
+# Defines the training approach using the unified flexible strategy.
+# This configuration trains both video and audio generation from text prompts.
 training_strategy:
-  # Strategy name: "text_to_video" for standard text-to-video training
-  name: "text_to_video"
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
 
-  # Probability of conditioning on the first frame during training
-  # Higher values train the model to perform better in image-to-video (I2V) mode,
-  # where a clean first frame is provided and the model generates the rest of the video
-  # Increase this value to train the model to perform better in image-to-video (I2V) mode
-  first_frame_conditioning_p: 0.5
+  # Video modality configuration
+  # When is_generated is true, the model learns to generate (denoise) video
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
 
-  # Enable joint audio-video training
-  # Set to true if your dataset includes audio and you want to train the audio branch
-  with_audio: true
-
-  # Directory name (within preprocessed_data_root) containing audio latents
-  # Only used when with_audio is true
-  audio_latents_dir: "audio_latents"
+  # Audio modality configuration
+  # When is_generated is true, the model learns to generate (denoise) audio
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
 
 # -----------------------------------------------------------------------------
 # Optimization Configuration
@@ -178,9 +183,9 @@ 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 state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
   offload_optimizer_during_validation: true
 
 # -----------------------------------------------------------------------------
@@ -189,7 +194,7 @@ acceleration:
 # Specifies the training data location and loading parameters.
 data:
   # Root directory containing preprocessed training data
-  # Should contain: latents/, conditions/, and optionally audio_latents/
+  # Should contain: latents/, conditions/, and audio_latents/
   preprocessed_data_root: "/path/to/preprocessed/data"
 
   # Number of worker processes for data loading
@@ -199,24 +204,31 @@ data:
 # -----------------------------------------------------------------------------
 # Validation Configuration
 # -----------------------------------------------------------------------------
-# Controls validation video generation during training.
+# Controls validation sampling during training.
 # NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
 # maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
 validation:
-  # Text prompts for validation video generation
-  # Provide prompts representative of your training data
-  # LTX-2 prefers longer, detailed prompts that describe both visual content and audio
-  prompts:
-    - "A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a laptop while occasionally glancing at notes beside her. Soft natural light streams through a large window, casting warm shadows across the room. She pauses to take a sip from a ceramic mug, then continues working with focused concentration. The audio captures the gentle clicking of keyboard keys, the soft rustle of papers, and ambient room tone with occasional distant bird chirps from outside."
-    - "A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet dish with precise movements. Steam rises from freshly cooked vegetables as he arranges them with tweezers. The stainless steel surfaces gleam under bright overhead lights, and various pots simmer on the stove behind him. The audio features the sizzling of pans, the clinking of utensils against plates, and the ambient hum of kitchen ventilation."
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a
+        laptop while occasionally glancing at notes beside her. Soft natural light streams through
+        a large window, casting warm shadows across the room. She pauses to take a sip from a
+        ceramic mug, then continues working with focused concentration. The audio captures the
+        gentle clicking of keyboard keys, the soft rustle of papers, and ambient room tone with
+        occasional distant bird chirps from outside.
+    - prompt: >-
+        A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet
+        dish with precise movements. Steam rises from freshly cooked vegetables as he arranges
+        them with tweezers. The stainless steel surfaces gleam under bright overhead lights, and
+        various pots simmer on the stove behind him. The audio features the sizzling of pans,
+        the clinking of utensils against plates, and the ambient hum of kitchen ventilation.
 
   # Negative prompt to avoid unwanted artifacts
   negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
 
-  # Optional: First frame images for image-to-video validation
-  # If provided, must have one image per prompt
-  images: null
-
   # Output video dimensions [width, height, frames]
   # Width and height must be divisible by 32
   # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
@@ -247,7 +259,7 @@ validation:
   stg_mode: "stg_av"  # "stg_av" perturbs both audio and video, "stg_v" video only
 
   # Whether to generate audio in validation samples
-  # Independent of training_strategy.with_audio - you can generate audio
+  # Independent of training_strategy audio modality settings — you can generate audio
   # in validation even when not training the audio branch
   generate_audio: true
 
@@ -310,9 +322,9 @@ wandb:
   entity: null
 
   # Tags to help organize runs
-  tags: [ "ltx2", "lora" ]
+  tags: [ "ltx2", "lora", "t2v", "low-vram" ]
 
-  # Log validation videos to W&B
+  # Log validation media (video/audio) to W&B
   log_validation_videos: true
 
 # -----------------------------------------------------------------------------
@@ -324,4 +336,4 @@ wandb:
 seed: 42
 
 # Directory to save outputs (checkpoints, validation videos, logs)
-output_dir: "outputs/ltx2_av_lora"
+output_dir: "outputs/t2v_lora_low_vram"

packages/ltx-trainer/configs/v2a_lora.yaml

@@ -0,0 +1,349 @@
+# =============================================================================
+# LTX-2 Video-to-Audio (Foley) LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# video-to-audio (Foley) generation. The model learns to generate audio
+# conditioned on a frozen video signal via the transformer's built-in
+# cross-modal attention.
+#
+# In this mode, video is provided as a frozen (clean, no noise, no loss)
+# conditioning signal. The audio modality is the only generated output.
+# Video influences audio generation through the transformer's video-to-audio
+# cross-attention mechanism.
+#
+# Use this configuration when you want to:
+#   - Generate sound effects (Foley) for existing videos
+#   - Train audio generation conditioned on visual content
+#   - Create models that produce audio matching video content
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/        # Video latents (frozen conditioning input)
+#   ├── conditions/     # Text embeddings for each video
+#   └── audio_latents/  # Audio latents (VAE-encoded audio, generated output)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general fine-tuning.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # For video-to-audio training, we target audio modules and the cross-modal
+  # attention that allows audio to attend to video features.
+  target_modules:
+    # Audio self-attention
+    - "audio_attn1.to_k"
+    - "audio_attn1.to_q"
+    - "audio_attn1.to_v"
+    - "audio_attn1.to_out.0"
+    # Audio cross-attention to text
+    - "audio_attn2.to_k"
+    - "audio_attn2.to_q"
+    - "audio_attn2.to_v"
+    - "audio_attn2.to_out.0"
+    # Audio feed-forward
+    - "audio_ff.net.0.proj"
+    - "audio_ff.net.2"
+    # Cross-modal attention: allows audio to attend to video features
+    - "video_to_audio_attn.to_k"
+    - "video_to_audio_attn.to_q"
+    - "video_to_audio_attn.to_v"
+    - "video_to_audio_attn.to_out.0"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the video-to-audio (Foley) training approach using the unified
+# flexible strategy. Video is frozen (no noise, sigma=0, excluded from loss)
+# and conditions audio generation via the transformer's cross-modal attention.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  # Video is frozen — it acts as conditioning for audio generation
+  # Frozen modalities get sigma=0, timestep=0, no noise, and no loss
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    # When false, video is passed through the transformer clean and influences
+    # audio via cross-modal attention
+    is_generated: false
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
+
+  # Audio modality configuration
+  # Audio is the generated (denoised) output
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 1e-4
+
+  # Total number of training steps
+  steps: 2000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, audio_latents/, and conditions/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        The sound of ocean waves crashing against rocky cliffs, with seagulls calling in the
+        distance and wind whistling through coastal grass.
+      conditions:
+        - type: video_to_audio
+          video: "/path/to/conditioning_video_1.mp4"
+    - prompt: >-
+        Footsteps echo in a marble hallway as a person walks steadily, with the distant hum of
+        air conditioning and occasional door closing sounds.
+      conditions:
+        - type: video_to_audio
+          video: "/path/to/conditioning_video_2.mp4"
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Validation dimensions [width, height, frames]
+  # Width/height resize the frozen video condition; frames and frame_rate set audio duration.
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 576, 576, 89 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" perturbs both audio and video, "stg_v" video only
+
+  # Whether to generate audio in validation samples
+  # Independent of training_strategy.audio.is_generated - you can generate audio
+  # in validation even when not training the audio branch
+  generate_audio: true
+
+  # Whether to generate video in validation samples
+  # Disabled for V2A since video is frozen conditioning, not generated
+  generate_video: false
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: -1
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "v2a", "foley" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/v2a_lora"

packages/ltx-trainer/configs/v2v_ic_lora.yaml

@@ -18,6 +18,8 @@
 #   ├── conditions/        # Text embeddings for each video
 #   └── reference_latents/ # Reference video latents (conditioning input)
 #
+# Dataset metadata columns: video, reference_video, caption
+#
 # =============================================================================
 
 # -----------------------------------------------------------------------------
@@ -34,7 +36,7 @@ model:
   text_encoder_path: "path/to/gemma-text-encoder"
 
   # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
-  # Note: video_to_video strategy requires "lora" mode
+  # IC-LoRA reference conditioning is intended for LoRA adapter training.
   training_mode: "lora"
 
   # Optional: Path to resume training from a checkpoint
@@ -97,20 +99,39 @@ lora:
 # -----------------------------------------------------------------------------
 # Training Strategy Configuration
 # -----------------------------------------------------------------------------
-# Defines the video-to-video (IC-LoRA) training approach.
+# Defines the video-to-video (IC-LoRA) training approach using the unified
+# flexible strategy. Reference conditioning concatenates pre-encoded reference
+# video latents to the target sequence. Reference tokens participate in
+# bidirectional self-attention but receive no noise and are excluded from loss.
 training_strategy:
-  # Strategy name: "video_to_video" for IC-LoRA training
-  name: "video_to_video"
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
 
-  # Probability of conditioning on the first frame during training
-  # Higher values train the model to perform better in image-to-video (I2V) mode,
-  # where a clean first frame is provided and the model generates the rest of the video
-  # Increase this value to train the model to perform better in image-to-video (I2V) mode
-  first_frame_conditioning_p: 0.2
+  # Video modality configuration
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing target video latents
+    latents_dir: "latents"
 
-  # Directory name (within preprocessed_data_root) containing reference video latents
-  # These are the conditioning inputs that guide the transformation
-  reference_latents_dir: "reference_latents"
+    # Conditions applied to the video modality during training
+    conditions:
+      # Reference conditioning (IC-LoRA): concatenates pre-encoded reference video
+      # latents to the target sequence. The model learns to transform the reference
+      # into the target video based on the text prompt.
+      - type: reference
+        # Directory name (within preprocessed_data_root) containing reference video latents
+        # These are the conditioning inputs that guide the transformation
+        latents_dir: "reference_latents"
+        # Probability of applying reference conditioning per training sample
+        probability: 1.0
+
+      # Optional first-frame conditioning to improve I2V capabilities
+      # At low probability, this teaches the model to also accept first-frame input
+      - type: first_frame
+        probability: 0.2
 
 # -----------------------------------------------------------------------------
 # Optimization Configuration
@@ -166,11 +187,12 @@ 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 state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
   offload_optimizer_during_validation: false
 
+
 # -----------------------------------------------------------------------------
 # Data Configuration
 # -----------------------------------------------------------------------------
@@ -187,37 +209,47 @@ data:
 # -----------------------------------------------------------------------------
 # Validation Configuration
 # -----------------------------------------------------------------------------
-# Controls validation video generation during training.
+# Controls validation sampling during training.
 # NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
 # maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
 validation:
-  # Text prompts for validation video generation
-  # Provide prompts representative of your training data
-  # LTX-2 prefers longer, detailed prompts that describe both visual content and audio
-  prompts:
-    - "A man in a casual blue jacket walks along a winding path through a lush green park on a bright sunny afternoon. Tall oak trees line the pathway, their leaves rustling gently in the breeze. Dappled sunlight creates shifting patterns on the ground as he strolls at a relaxed pace, occasionally looking up at the scenery around him. The audio captures footsteps on gravel, birds singing in the trees, distant children playing, and the soft whisper of wind through the foliage."
-    - "A fluffy orange tabby cat sits perfectly still on a wooden windowsill, its green eyes intently tracking small birds hopping on a branch just outside the glass. The cat's ears twitch and rotate, following every movement. Warm afternoon light illuminates its fur, creating a soft golden glow. Behind the cat, a cozy living room with a bookshelf and houseplants is visible. The audio features gentle purring, occasional soft meows, muffled bird chirps through the window, and quiet ambient room sounds."
-
-  # Reference videos for validation (REQUIRED for video_to_video strategy)
-  # Must provide one reference video per prompt
-  # These are the conditioning inputs for generating validation outputs
-  reference_videos:
-    - "/path/to/reference_video_1.mp4"
-    - "/path/to/reference_video_2.mp4"
-
-  # Downscale factor for reference videos (for efficient IC-LoRA training)
-  # When > 1, reference videos are processed at 1/n resolution
-  # Must match the --reference-downscale-factor used during dataset preprocessing
-  # Examples: 1 = same resolution, 2 = half resolution (384x384 ref for 768x768 target)
-  reference_downscale_factor: 1
+  # Validation samples — each sample describes a self-contained generation request.
+  # For IC-LoRA, each sample includes a reference condition pointing to the conditioning video.
+  samples:
+    - prompt: >-
+        A man in a casual blue jacket walks along a winding path through a lush green park on a
+        bright sunny afternoon. Tall oak trees line the pathway, their leaves rustling gently in
+        the breeze. Dappled sunlight creates shifting patterns on the ground as he strolls at a
+        relaxed pace, occasionally looking up at the scenery around him. The audio captures
+        footsteps on gravel, birds singing in the trees, distant children playing, and the soft
+        whisper of wind through the foliage.
+      conditions:
+        - type: reference
+          video: "/path/to/reference_video_1.mp4"
+          # Set these to match --reference-downscale-factor / --reference-temporal-scale-factor
+          # if reference latents were preprocessed at reduced spatial or temporal resolution.
+          downscale_factor: 1
+          temporal_scale_factor: 1
+          include_in_output: true
+    - prompt: >-
+        A fluffy orange tabby cat sits perfectly still on a wooden windowsill, its green eyes
+        intently tracking small birds hopping on a branch just outside the glass. The cat's ears
+        twitch and rotate, following every movement. Warm afternoon light illuminates its fur,
+        creating a soft golden glow. Behind the cat, a cozy living room with a bookshelf and
+        houseplants is visible. The audio features gentle purring, occasional soft meows, muffled
+        bird chirps through the window, and quiet ambient room sounds.
+      conditions:
+        - type: reference
+          video: "/path/to/reference_video_2.mp4"
+          # Set these to match --reference-downscale-factor / --reference-temporal-scale-factor
+          # if reference latents were preprocessed at reduced spatial or temporal resolution.
+          downscale_factor: 1
+          temporal_scale_factor: 1
+          include_in_output: true
 
   # Negative prompt to avoid unwanted artifacts
   negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
 
-  # Optional: First frame images for additional conditioning
-  # If provided, must have one image per prompt
-  images: null
-
   # Output video dimensions [width, height, frames]
   # Width and height must be divisible by 32
   # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
@@ -254,10 +286,6 @@ validation:
   # Skip validation at the beginning of training (step 0)
   skip_initial_validation: false
 
-  # Concatenate reference video side-by-side with generated output
-  # Useful for visually comparing the transformation quality
-  include_reference_in_output: true
-
 # -----------------------------------------------------------------------------
 # Checkpoint Configuration
 # -----------------------------------------------------------------------------
@@ -314,9 +342,9 @@ wandb:
   entity: null
 
   # Tags to help organize runs
-  tags: [ "ltx2", "ic-lora", "video-to-video" ]
+  tags: [ "ltx2", "ic-lora", "v2v" ]
 
-  # Log validation videos to W&B
+  # Log validation media (video/audio) to W&B
   log_validation_videos: true
 
 # -----------------------------------------------------------------------------
@@ -328,4 +356,4 @@ wandb:
 seed: 42
 
 # Directory to save outputs (checkpoints, validation videos, logs)
-output_dir: "outputs/ltx2_v2v_ic_lora"
+output_dir: "outputs/v2v_ic_lora"

packages/ltx-trainer/configs/video_extend_lora.yaml

@@ -0,0 +1,355 @@
+# =============================================================================
+# LTX-2 Video Extension LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# video extension (temporal continuation). The model learns to extend a video
+# forward in time by conditioning on a prefix of existing frames.
+#
+# Prefix conditioning works by providing the first N temporal units as clean
+# conditioning signals during training — they receive no noise, timestep=0,
+# and are excluded from the loss. The model learns to generate continuation
+# frames that seamlessly follow the given prefix.
+#
+# Use this configuration when you want to:
+#   - Train the model to extend/continue existing videos
+#   - Fine-tune temporal continuation capabilities on custom datasets
+#   - Create seamless video extension models with optional audio
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/        # Video latents (VAE-encoded videos)
+#   ├── conditions/     # Text embeddings for each video
+#   └── audio_latents/  # Audio latents (VAE-encoded audio)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general fine-tuning.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # Using short patterns like "to_k" matches ALL attention modules (video, audio, and cross-modal).
+  # For audio-video training, this is the recommended approach.
+  target_modules:
+    # Attention layers (matches both video and audio branches)
+    - "to_k"
+    - "to_q"
+    - "to_v"
+    - "to_out.0"
+    # Uncomment below to also train feed-forward layers (can increase the LoRA's capacity):
+    # - "ff.net.0.proj"
+    # - "ff.net.2"
+    # - "audio_ff.net.0.proj"
+    # - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the video extension training approach using the unified flexible strategy.
+# Prefix conditioning provides the first N latent frames as clean conditioning,
+# teaching the model to generate temporally coherent continuations.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
+
+    # Conditions applied to the video modality during training
+    # Prefix conditioning: the first N temporal units are provided as clean
+    # conditioning signals (no noise, timestep=0, excluded from loss)
+    conditions:
+      - type: prefix
+        # Number of latent frames to use as conditioning prefix
+        # For prefix conditioning, N latent frames correspond to (N - 1) * 8 + 1 pixel frames.
+        # temporal_boundary=8 means 57 pixel frames are used as prefix.
+        temporal_boundary: 8
+        # Probability of applying prefix conditioning per training sample
+        # At 1.0, all training samples use video extension mode
+        probability: 1.0
+
+  # Audio modality configuration
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 1e-4
+
+  # Total number of training steps
+  steps: 2000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, conditions/, and audio_latents/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a
+        laptop while occasionally glancing at notes beside her. Soft natural light streams through
+        a large window, casting warm shadows across the room. She pauses to take a sip from a
+        ceramic mug, then continues working with focused concentration. The audio captures the
+        gentle clicking of keyboard keys, the soft rustle of papers, and ambient room tone with
+        occasional distant bird chirps from outside.
+      conditions:
+        - type: prefix
+          video: "/path/to/prefix_video_1.mp4"
+          # Matches temporal_boundary=8 during training: (8 - 1) * 8 + 1 = 57 frames
+          num_frames: 57
+    - prompt: >-
+        A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet
+        dish with precise movements. Steam rises from freshly cooked vegetables as he arranges
+        them with tweezers. The stainless steel surfaces gleam under bright overhead lights, and
+        various pots simmer on the stove behind him. The audio features the sizzling of pans,
+        the clinking of utensils against plates, and the ambient hum of kitchen ventilation.
+      conditions:
+        - type: prefix
+          video: "/path/to/prefix_video_2.mp4"
+          # Matches temporal_boundary=8 during training: (8 - 1) * 8 + 1 = 57 frames
+          num_frames: 57
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 576, 576, 89 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" perturbs both audio and video, "stg_v" video only
+
+  # Whether to generate audio in validation samples
+  # Independent of training_strategy.audio.is_generated - you can generate audio
+  # in validation even when not training the audio branch
+  generate_audio: true
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: -1
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "video-extension" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/video_extend_lora"

packages/ltx-trainer/configs/video_inpainting_lora.yaml

@@ -0,0 +1,343 @@
+# =============================================================================
+# LTX-2 Video Inpainting LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# video inpainting. The model learns to fill in masked regions of a video
+# using per-sample binary masks that define which regions are conditioning
+# (provided clean) and which regions the model must generate.
+#
+# Mask conditioning works by loading per-sample binary masks from disk.
+# Masked regions receive clean latents (no noise, timestep=0) and are excluded
+# from the loss. Unmasked regions are noised and trained normally.
+#
+# Use this configuration when you want to:
+#   - Train the model to fill in or replace regions of existing videos
+#   - Fine-tune video inpainting capabilities on custom datasets
+#   - Create object removal or region editing models
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/        # Video latents (VAE-encoded videos)
+#   ├── conditions/     # Text embeddings for each video
+#   └── video_masks/    # Per-sample binary masks defining conditioning regions
+#
+# Dataset metadata columns: video, video_mask, caption
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general fine-tuning.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # This is a video-only config, so explicitly target video attention modules.
+  target_modules:
+    - "attn1.to_k"
+    - "attn1.to_q"
+    - "attn1.to_v"
+    - "attn1.to_out.0"
+    - "attn2.to_k"
+    - "attn2.to_q"
+    - "attn2.to_v"
+    - "attn2.to_out.0"
+    # Uncomment below to also train feed-forward layers (can increase the LoRA's capacity):
+    # - "ff.net.0.proj"
+    # - "ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the video inpainting training approach using the unified flexible
+# strategy. Per-sample binary masks define which regions are provided as clean
+# conditioning and which regions the model must learn to generate.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
+
+    # Conditions applied to the video modality during training
+    # Mask conditioning: per-sample binary masks loaded from disk define which
+    # tokens are conditioning (clean, timestep=0, no loss) vs generated
+    conditions:
+      - type: mask
+        # Directory name (within preprocessed_data_root) containing binary masks
+        # Each mask file corresponds to a training sample and defines the
+        # conditioning region (mask=1 means conditioning, mask=0 means generate)
+        mask_dir: "video_masks"
+        # Probability of applying mask conditioning per training sample
+        # At 1.0, all training samples use inpainting mode
+        probability: 1.0
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 1e-4
+
+  # Total number of training steps
+  steps: 2000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, conditions/, and video_masks/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a
+        laptop while occasionally glancing at notes beside her. Soft natural light streams through
+        a large window, casting warm shadows across the room. She pauses to take a sip from a
+        ceramic mug, then continues working with focused concentration.
+      conditions:
+        - type: mask
+          video: "/path/to/inpainting_video_1.mp4"
+          mask: "/path/to/inpainting_mask_1.mp4"
+    - prompt: >-
+        A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet
+        dish with precise movements. Steam rises from freshly cooked vegetables as he arranges
+        them with tweezers. The stainless steel surfaces gleam under bright overhead lights, and
+        various pots simmer on the stove behind him.
+      conditions:
+        - type: mask
+          video: "/path/to/inpainting_video_2.mp4"
+          mask: "/path/to/inpainting_mask_2.mp4"
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 576, 576, 89 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_v"  # "stg_v" for video-only validation
+
+  # Video inpainting trains video only; do not generate validation audio.
+  generate_audio: false
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: -1
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "inpainting" ]
+
+  # Log validation outputs to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/video_inpainting_lora"

packages/ltx-trainer/configs/video_outpainting_lora.yaml

@@ -0,0 +1,341 @@
+# =============================================================================
+# LTX-2 Video Outpainting LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# video outpainting (spatial extension). The model learns to generate content
+# surrounding a known rectangular region of the video.
+#
+# Spatial crop conditioning works by providing a rectangular pixel region as
+# clean conditioning during training — those tokens receive no noise,
+# timestep=0, and are excluded from the loss. The model learns to generate
+# the content outside the given region, seamlessly extending the scene.
+#
+# Use this configuration when you want to:
+#   - Extend video content beyond its original boundaries
+#   - Train spatial outpainting capabilities on custom datasets
+#   - Create video aspect ratio conversion models
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/        # Video latents (VAE-encoded videos)
+#   └── conditions/     # Text embeddings for each video
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general fine-tuning.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # This is a video-only config, so explicitly target video attention modules.
+  target_modules:
+    - "attn1.to_k"
+    - "attn1.to_q"
+    - "attn1.to_v"
+    - "attn1.to_out.0"
+    - "attn2.to_k"
+    - "attn2.to_q"
+    - "attn2.to_v"
+    - "attn2.to_out.0"
+    # Uncomment below to also train feed-forward layers (can increase the LoRA's capacity):
+    # - "ff.net.0.proj"
+    # - "ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the video outpainting training approach using the unified flexible
+# strategy. Spatial crop conditioning provides a rectangular pixel region as
+# clean conditioning, teaching the model to generate surrounding content.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
+
+    # Conditions applied to the video modality during training
+    # Spatial crop conditioning provides a rectangular region as clean context;
+    # the model learns to generate the surrounding video tokens.
+    conditions:
+      - type: spatial_crop
+        # Rectangular pixel region provided as clean conditioning (y1, x1, y2, x2)
+        # Pixels within this region are conditioning (clean, timestep=0, no loss)
+        # Pixels outside this region are generated by the model
+        # Coordinates are in pixel space — automatically converted to latent space
+        spatial_region: [0, 0, 288, 576]
+        # Probability of applying spatial crop conditioning per training sample
+        # At 1.0, all training samples use outpainting mode
+        probability: 1.0
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 1e-4
+
+  # Total number of training steps
+  steps: 2000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/ and conditions/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a
+        laptop while occasionally glancing at notes beside her. Soft natural light streams through
+        a large window, casting warm shadows across the room. She pauses to take a sip from a
+        ceramic mug, then continues working with focused concentration.
+      conditions:
+        - type: spatial_crop
+          video: "/path/to/outpainting_video_1.mp4"
+          spatial_region: [0, 0, 288, 576]
+    - prompt: >-
+        A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet
+        dish with precise movements. Steam rises from freshly cooked vegetables as he arranges
+        them with tweezers. The stainless steel surfaces gleam under bright overhead lights, and
+        various pots simmer on the stove behind him.
+      conditions:
+        - type: spatial_crop
+          video: "/path/to/outpainting_video_2.mp4"
+          spatial_region: [0, 0, 288, 576]
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 576, 576, 89 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_v"  # "stg_v" for video-only validation
+
+  # Video outpainting trains video only; do not generate validation audio.
+  generate_audio: false
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: -1
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "outpainting", "spatial-crop" ]
+
+  # Log validation outputs to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/video_outpainting_lora"

packages/ltx-trainer/configs/video_suffix_lora.yaml

@@ -0,0 +1,355 @@
+# =============================================================================
+# LTX-2 Video Backward Extension (Suffix) LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# backward video extension. The model learns to generate video content that
+# leads into a given suffix of existing frames.
+#
+# Suffix conditioning works by providing the last N temporal units as clean
+# conditioning signals during training — they receive no noise, timestep=0,
+# and are excluded from the loss. The model learns to generate preceding
+# frames that seamlessly lead into the given suffix.
+#
+# Use this configuration when you want to:
+#   - Train the model to generate "prequel" content for existing videos
+#   - Fine-tune backward temporal continuation capabilities
+#   - Create video backward extension models with optional audio
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/        # Video latents (VAE-encoded videos)
+#   ├── conditions/     # Text embeddings for each video
+#   └── audio_latents/  # Audio latents (VAE-encoded audio)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general fine-tuning.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # Using short patterns like "to_k" matches ALL attention modules (video, audio, and cross-modal).
+  # For audio-video training, this is the recommended approach.
+  target_modules:
+    # Attention layers (matches both video and audio branches)
+    - "to_k"
+    - "to_q"
+    - "to_v"
+    - "to_out.0"
+    # Uncomment below to also train feed-forward layers (can increase the LoRA's capacity):
+    # - "ff.net.0.proj"
+    # - "ff.net.2"
+    # - "audio_ff.net.0.proj"
+    # - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the backward video extension training approach using the unified
+# flexible strategy. Suffix conditioning provides the last N latent frames as
+# clean conditioning, teaching the model to generate content leading into them.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
+
+    # Conditions applied to the video modality during training
+    # Suffix conditioning: the last N temporal units are provided as clean
+    # conditioning signals (no noise, timestep=0, excluded from loss)
+    conditions:
+      - type: suffix
+        # Number of latent frames to use as conditioning suffix
+        # For suffix conditioning, N latent frames correspond to N * 8 pixel frames.
+        # temporal_boundary=8 means the last 64 pixel frames are used as suffix.
+        temporal_boundary: 8
+        # Probability of applying suffix conditioning per training sample
+        # At 1.0, all training samples use backward extension mode
+        probability: 1.0
+
+  # Audio modality configuration
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 1e-4
+
+  # Total number of training steps
+  steps: 2000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, conditions/, and audio_latents/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a
+        laptop while occasionally glancing at notes beside her. Soft natural light streams through
+        a large window, casting warm shadows across the room. She pauses to take a sip from a
+        ceramic mug, then continues working with focused concentration. The audio captures the
+        gentle clicking of keyboard keys, the soft rustle of papers, and ambient room tone with
+        occasional distant bird chirps from outside.
+      conditions:
+        - type: suffix
+          video: "/path/to/suffix_video_1.mp4"
+          # Matches temporal_boundary=8 during training: 8 * 8 = 64 frames
+          num_frames: 64
+    - prompt: >-
+        A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet
+        dish with precise movements. Steam rises from freshly cooked vegetables as he arranges
+        them with tweezers. The stainless steel surfaces gleam under bright overhead lights, and
+        various pots simmer on the stove behind him. The audio features the sizzling of pans,
+        the clinking of utensils against plates, and the ambient hum of kitchen ventilation.
+      conditions:
+        - type: suffix
+          video: "/path/to/suffix_video_2.mp4"
+          # Matches temporal_boundary=8 during training: 8 * 8 = 64 frames
+          num_frames: 64
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 576, 576, 89 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" perturbs both audio and video, "stg_v" video only
+
+  # Whether to generate audio in validation samples
+  # Independent of training_strategy.audio.is_generated - you can generate audio
+  # in validation even when not training the audio branch
+  generate_audio: true
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: -1
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "video-suffix", "backward-extension" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/video_suffix_lora"

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

@@ -10,7 +10,7 @@ sub-configurations:
 
 - **ModelConfig**: Base model and training mode settings
 - **LoraConfig**: LoRA training parameters
-- **TrainingStrategyConfig**: Training strategy settings (text-to-video or video-to-video)
+- **TrainingStrategyConfig**: Training strategy settings (flexible conditioning framework)
 - **OptimizationConfig**: Learning rate, batch sizes, and scheduler settings
 - **AccelerationConfig**: Mixed precision and quantization settings
 - **DataConfig**: Data loading parameters
@@ -24,13 +24,29 @@ sub-configurations:
 
 Check out our example configurations in the `configs` directory:
 
-- 📄 [Audio-Video LoRA Training](../configs/ltx2_av_lora.yaml) - Joint audio-video generation training
-- 📄 [Audio-Video LoRA Training (Low VRAM)](../configs/ltx2_av_lora_low_vram.yaml) - Memory-optimized config for 32GB
-  GPUs (uses 8-bit optimizer, INT8 quantization, and reduced LoRA rank)
-- 📄 [IC-LoRA Training](../configs/ltx2_v2v_ic_lora.yaml) - Video-to-video transformation training
+- 📄 [Text-to-Video LoRA](../configs/t2v_lora.yaml) - Text-to-video LoRA training
+- 📄 [Image-to-Video LoRA](../configs/i2v_lora.yaml) - Image-to-video LoRA training
+- 📄 [IC-LoRA Video-to-Video](../configs/v2v_ic_lora.yaml) - IC-LoRA video-to-video training
+- 📄 [Audio-to-Video LoRA](../configs/a2v_lora.yaml) - Audio-to-video LoRA training
+- 📄 [Video-to-Audio LoRA](../configs/v2a_lora.yaml) - Video-to-audio (Foley) LoRA training
+- 📄 [Video Extension LoRA](../configs/video_extend_lora.yaml) - Video extension (forward) LoRA training
+- 📄 [Video Suffix LoRA](../configs/video_suffix_lora.yaml) - Video extension (backward) LoRA training
+- 📄 [Video Inpainting LoRA](../configs/video_inpainting_lora.yaml) - Video inpainting LoRA training
+- 📄 [Video Outpainting LoRA](../configs/video_outpainting_lora.yaml) - Video outpainting (spatial crop) LoRA training
+- 📄 [Text-to-Audio LoRA](../configs/t2a_lora.yaml) - Text-to-audio LoRA training
+- 📄 [Audio Extension LoRA](../configs/audio_extend_lora.yaml) - Audio extension (forward) LoRA training
+- 📄 [Audio Suffix LoRA](../configs/audio_suffix_lora.yaml) - Audio extension (backward) LoRA training
+- 📄 [Audio Inpainting LoRA](../configs/audio_inpainting_lora.yaml) - Audio inpainting LoRA training
+- 📄 [Audio-to-Audio IC-LoRA](../configs/a2a_ic_lora.yaml) - Audio IC-LoRA transformation training
+- 📄 [AV2AV IC-LoRA](../configs/av2av_ic_lora.yaml) - Audio+video IC-LoRA transformation training
+- 📄 [T2V LoRA (Low VRAM)](../configs/t2v_lora_low_vram.yaml) - Memory-optimized config for 32GB GPUs
 
 ## ⚙️ Configuration Sections
 
+> [!NOTE]
+> The YAML snippets below show **recommended starting values**, not necessarily the code defaults.
+> Fields you omit from your config file will use the code defaults from [`config.py`](../src/ltx_trainer/config.py).
+
 ### ModelConfig
 
 Controls the base model and training mode settings.
@@ -149,37 +165,60 @@ target_modules:
 
 ### TrainingStrategyConfig
 
-Configures the training strategy. The trainer includes two built-in strategies described below.
-For custom use cases, see [Implementing Custom Training Strategies](custom-training-strategies.md).
+Configures the training strategy. The recommended strategy is `"flexible"`, which supports all conditioning scenarios through configuration.
 
-#### Text-to-Video Strategy
+#### Flexible Strategy
+
+The flexible strategy provides a unified conditioning framework. Each modality (video, audio) is configured
+independently with its own latents directory, generation flag, and list of conditions.
 
 ```yaml
 training_strategy:
-  name: "text_to_video"
-  first_frame_conditioning_p: 0.1     # Probability of first-frame conditioning
-  with_audio: false                   # Enable joint audio-video training
-  audio_latents_dir: "audio_latents"  # Directory for audio latents (when with_audio: true)
+  name: "flexible"
+  video:
+    is_generated: true                 # Video is denoised during training
+    latents_dir: "latents"             # Directory containing precomputed video latents
+    conditions:
+      - type: first_frame              # Use first frame as conditioning
+        probability: 0.5               # Apply this condition 50% of the time
+  audio:
+    is_generated: true                 # Audio is denoised during training
+    latents_dir: "audio_latents"       # Directory containing precomputed audio latents
+    conditions: []                     # No additional audio conditions (text-only)
 ```
 
-#### Video-to-Video Strategy (IC-LoRA)
+**ModalityConfig parameters:**
 
-```yaml
-training_strategy:
-  name: "video_to_video"
-  first_frame_conditioning_p: 0.1
-  reference_latents_dir: "reference_latents"  # Directory for reference video latents
-```
+| Parameter      | Description                                                                                                      |
+|----------------|------------------------------------------------------------------------------------------------------------------|
+| `is_generated` | `true` = modality is denoised (contributes to loss). `false` = frozen conditioning (sigma=0, no loss).           |
+| `latents_dir`  | Directory name within `preprocessed_data_root` containing precomputed latents for this modality.                 |
+| `conditions`   | List of conditioning configs applied during training (see condition types below). Text conditioning is implicit. |
 
-**Key parameters:**
+**Condition types:**
 
-| Parameter                    | Description                                                      |
-|------------------------------|------------------------------------------------------------------|
-| `name`                       | Strategy type: `"text_to_video"` or `"video_to_video"`           |
-| `first_frame_conditioning_p` | Probability of using first frame as conditioning (0.0-1.0)       |
-| `with_audio`                 | (text_to_video only) Enable joint audio-video training           |
-| `audio_latents_dir`          | (text_to_video only) Directory name for audio latents            |
-| `reference_latents_dir`      | (video_to_video only) Directory name for reference video latents |
+| Type           | Parameters                                          | Description                                                                           |
+|----------------|-----------------------------------------------------|---------------------------------------------------------------------------------------|
+| `first_frame`  | `probability`                                       | First latent frame is clean, excluded from loss. **Video only.**                      |
+| `prefix`       | `temporal_boundary`, `probability`                  | First N latent temporal units are clean. For extension forward.                       |
+| `suffix`       | `temporal_boundary`, `probability`                  | Last N latent temporal units are clean. For extension backward.                       |
+| `spatial_crop` | `spatial_region` (y1, x1, y2, x2 in px), `probability` | Rectangular region is clean, excluded from loss. For outpainting. **Video only.** |
+| `mask`         | `mask_dir`, `probability`                           | Per-sample mask directory. Masks are thresholded at `0.5`; `1` means conditioning, `0` means generate. |
+| `reference`    | `latents_dir`, `probability`                        | IC-LoRA style concatenation. Reference tokens are prepended, clean (timestep=0), no loss. |
+
+> [!NOTE]
+> The `prefix`, `suffix`, `mask`, and `reference` condition types work on both video and audio modalities —
+> place them in the `video.conditions` or `audio.conditions` list as appropriate.
+> `first_frame` and `spatial_crop` are video-only conditions.
+
+> [!NOTE]
+> Training conditions reference **directories** of precomputed data (within `preprocessed_data_root`),
+> while validation conditions reference **individual files** (images, videos, masks) that are encoded
+> on-the-fly during validation. The condition `type` names are the same, but the fields differ.
+
+> [!NOTE]
+> The legacy `text_to_video` and `video_to_video` strategies are deprecated but remain forward-compatible.
+> New configs should use `name: "flexible"`.
 
 ### OptimizationConfig
 
@@ -206,7 +245,7 @@ optimization:
 | `steps`                         | Total number of training steps                                                               |
 | `batch_size`                    | Batch size per GPU (reduce if running out of memory)                                         |
 | `gradient_accumulation_steps`   | Accumulate gradients over multiple steps                                                     |
-| `scheduler_type`                | LR scheduler: `"constant"`, `"linear"`, `"cosine"`, `"cosine_with_restarts"`, `"polynomial"` |
+| `scheduler_type`                | LR scheduler: `"constant"`, `"linear"`, `"cosine"`, `"cosine_with_restarts"`, `"polynomial"`, `"step"` |
 | `enable_gradient_checkpointing` | Trade training speed for GPU memory savings (recommended for large models)                   |
 
 ### AccelerationConfig
@@ -226,7 +265,7 @@ acceleration:
 | Parameter                             | Description                                                                                                                                                                              |
 |---------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
 | `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"`, `"int2-quanto"`, `"fp8-quanto"`, or `"fp8uz-quanto"`                                                                       |
 | `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. |
 
@@ -244,50 +283,84 @@ data:
 
 | Parameter                | Description                                                                                |
 |--------------------------|--------------------------------------------------------------------------------------------|
-| `preprocessed_data_root` | Path to your preprocessed dataset (contains `latents/`, `conditions/`, etc.)               |
+| `preprocessed_data_root` | Path to your preprocessed dataset directory produced by `process_dataset.py` (contains `latents/`, `conditions/`, etc.) |
 | `num_dataloader_workers` | Number of parallel data loading processes (0 = synchronous loading, useful when debugging) |
 
 ### ValidationConfig
 
-Validation and inference settings for monitoring training progress.
+Validation and inference settings for monitoring training progress. Validation samples use a self-describing
+format where each sample specifies its own prompt and conditions.
 
 ```yaml
 validation:
-  prompts: # Validation prompts
-    - "A cat playing with a ball"
-    - "A dog running in a field"
-  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
-  images: null                        # Optional image paths for image-to-video
-  reference_videos: null              # Reference video paths (IC-LoRA only)
-  video_dims: [ 576, 576, 89 ]        # Video dimensions [width, height, frames]
-  frame_rate: 25.0                    # Frame rate for generated videos
-  seed: 42                            # Random seed for reproducibility
-  inference_steps: 30                 # Number of inference steps
-  interval: 100                       # Steps between validation runs
-  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
-  stg_mode: "stg_av"                  # "stg_av" or "stg_v" (video only)
-  generate_audio: true                # Whether to generate audio
-  skip_initial_validation: false      # Skip validation at step 0
-  include_reference_in_output: false  # Include reference video side-by-side (IC-LoRA)
+  samples:
+    - prompt: "A cat playing with a ball"
+      conditions:
+        - type: first_frame
+          image_or_video: "/path/to/image.png"
+    - prompt: "A dog running in a field"
+  video_dims: [576, 576, 89]                # Output dimensions: [width, height, frames]
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"  # Negative prompt for all samples
+  frame_rate: 25.0                          # Output video frame rate (fps)
+  seed: 42                                  # Random seed for reproducibility
+  inference_steps: 30                       # Number of denoising steps
+  interval: 100                             # Run validation every N steps (null to disable)
+  guidance_scale: 4.0                       # CFG scale (higher = stronger prompt adherence)
+  stg_scale: 1.0                            # STG scale (0.0 to disable)
+  stg_blocks: [29]                          # Transformer blocks to apply STG perturbation
+  stg_mode: "stg_av"                        # STG mode: "stg_av" (audio+video) or "stg_v" (video only)
+  generate_audio: true                      # Whether to generate audio during validation
+  generate_video: true                      # Whether to generate video during validation
+  skip_initial_validation: false            # Skip validation at step 0
 ```
 
 **Key parameters:**
 
-| Parameter                     | Description                                                                                                              |
-|-------------------------------|--------------------------------------------------------------------------------------------------------------------------|
-| `prompts`                     | List of text prompts for validation video generation                                                                     |
-| `images`                      | List of image paths for image-to-video validation (must match number of prompts)                                         |
-| `reference_videos`            | List of reference video paths for IC-LoRA validation (must match number of prompts)                                      |
-| `video_dims`                  | Output dimensions `[width, height, frames]`. Width/height must be divisible by 32, frames must satisfy `frames % 8 == 1` |
-| `interval`                    | Steps between validation runs (set to `null` to disable)                                                                 |
-| `guidance_scale`              | CFG (Classifier-Free Guidance) scale. Recommended: 4.0                                                                   |
-| `stg_scale`                   | STG (Spatio-Temporal Guidance) scale. 0.0 disables STG. Recommended: 1.0                                                 |
-| `stg_blocks`                  | Transformer blocks to perturb for STG. Recommended: `[29]` (single block)                                                |
-| `stg_mode`                    | STG mode: `"stg_av"` perturbs both audio and video, `"stg_v"` perturbs video only                                        |
-| `generate_audio`              | Whether to generate audio in validation samples                                                                          |
-| `include_reference_in_output` | For IC-LoRA: concatenate reference video side-by-side with output                                                        |
+| Parameter                | Description                                                                                                              |
+|--------------------------|--------------------------------------------------------------------------------------------------------------------------|
+| `samples`                | List of `ValidationSample` objects (see below). Replaces the legacy `prompts`/`images`/`reference_videos` fields.        |
+| `video_dims`             | Output dimensions `[width, height, frames]`. Width/height must be divisible by 32, frames must satisfy `frames % 8 == 1` |
+| `interval`               | Steps between validation runs (set to `null` to disable)                                                                 |
+| `guidance_scale`         | CFG (Classifier-Free Guidance) scale. Recommended: 4.0                                                                   |
+| `stg_scale`              | STG (Spatio-Temporal Guidance) scale. 0.0 disables STG. Recommended: 1.0                                                 |
+| `stg_blocks`             | Transformer blocks to perturb for STG. Recommended: `[29]` (single block)                                                |
+| `stg_mode`               | STG mode: `"stg_av"` perturbs both audio and video, `"stg_v"` perturbs video only                                        |
+| `generate_audio`         | Whether to generate audio in validation samples                                                                          |
+| `generate_video`         | Whether to generate video in validation samples. Set to `false` for V2A (video-to-audio) validation. Default: `true`     |
+| `skip_initial_validation`| Skip validation video sampling at step 0 (beginning of training)                                                         |
+
+#### ValidationSample
+
+Each sample in the `samples` list has:
+
+| Field        | Description                                                                                     |
+|--------------|-------------------------------------------------------------------------------------------------|
+| `prompt`     | Text prompt for this validation sample.                                                         |
+| `conditions` | List of validation conditions (see types below). Empty list = text-only generation.             |
+| `video_dims` | Optional per-sample override for `(width, height, frames)`. Inherits from `ValidationConfig` if not set. |
+| `seed`       | Optional per-sample override for random seed. Inherits from `ValidationConfig` if not set.      |
+
+#### Validation Condition Types
+
+| Type             | Parameters                                                 | Description                                                             |
+|------------------|------------------------------------------------------------|-------------------------------------------------------------------------|
+| `first_frame`    | `image_or_video` (path)                                    | Use the first frame of the image/video as conditioning.                 |
+| `prefix`         | `video` or `audio` (path), optional `num_frames`/`duration`| Use a video/audio clip as temporal prefix (for extension forward).      |
+| `suffix`         | `video` or `audio` (path), optional `num_frames`/`duration`| Use a video/audio clip as temporal suffix (for extension backward).     |
+| `spatial_crop`   | `video` (path), `spatial_region` (y1, x1, y2, x2)         | Provide spatial context for outpainting. Video only.                    |
+| `mask`           | `video` or `audio` (path), `mask` (path)                   | Mask-based inpainting with a binary mask file.                          |
+| `reference`      | `video` or `audio` (path), optional video-reference `downscale_factor`, `temporal_scale_factor`, `include_in_output` | IC-LoRA style reference conditioning. |
+| `video_to_audio` | `video` (path)                                             | Freeze video, generate audio. For Foley/V2A tasks.                      |
+| `audio_to_video` | `audio` (path)                                             | Freeze audio, generate video. For audio-driven generation.              |
+
+For video `reference` validation conditions, `downscale_factor` is the spatial reference scale and
+`temporal_scale_factor` is the temporal reference scale. Set both to match the factors used when
+preprocessing video reference latents for training; validation media is encoded on the fly and cannot infer
+those factors from the training dataset.
+
+> [!NOTE]
+> The legacy fields `prompts`, `images`, and `reference_videos` are deprecated but auto-converted to `samples`
+> internally. New configs should use the `samples` format.
 
 ### CheckpointsConfig
 
@@ -298,6 +371,8 @@ checkpoints:
   interval: 250       # Steps between checkpoint saves (null = disabled)
   keep_last_n: 3      # Number of recent checkpoints to retain
   precision: bfloat16 # Precision for saved weights (bfloat16 or float32)
+  no_resume: false            # Ignore saved state, start from step 0
+  save_training_state: "minimal"  # "full", "minimal", or "off"
 ```
 
 **Key parameters:**
@@ -307,6 +382,8 @@ checkpoints:
 | `interval`    | Steps between intermediate checkpoint saves (set to `null` to disable)        |
 | `keep_last_n` | Number of most recent checkpoints to keep (-1 = keep all)                     |
 | `precision`   | Precision for saved checkpoint weights: `"bfloat16"` (default) or `"float32"` |
+| `no_resume` | When `true`, ignore saved training state and start from step 0. Model weights from `load_checkpoint` are still loaded. |
+| `save_training_state` | Save training state for resume: `"full"` (optimizer + scheduler + RNG), `"minimal"` (scheduler + RNG only, sufficient for LoRA), `"off"` (no resume). |
 
 ### HubConfig
 
@@ -364,6 +441,20 @@ flow_matching:
 | `timestep_sampling_mode`   | Sampling strategy: `"uniform"` or `"shifted_logit_normal"` |
 | `timestep_sampling_params` | Additional parameters for the sampling strategy            |
 
+### General Configuration
+
+Top-level settings for the training run.
+
+```yaml
+seed: 42                                  # Random seed for reproducibility
+output_dir: "outputs/my_training_run"     # Directory for outputs (checkpoints, validation videos, logs)
+```
+
+| Parameter    | Description                                              |
+|--------------|----------------------------------------------------------|
+| `seed`       | Random seed for reproducibility (default: `42`)          |
+| `output_dir` | Directory to save outputs (default: `"outputs"`)         |
+
 ## 🚀 Next Steps
 
 Once you've configured your training parameters:

packages/ltx-trainer/docs/custom-training-strategies.md

@@ -1,7 +1,7 @@
 # Implementing Custom Training Strategies
 
-This guide explains how to implement your own training strategy for specialized use cases like audio-only training,
-video inpainting, or other custom training recipes.
+This guide explains how to implement your own training strategy for specialized recipes that cannot be expressed with
+the built-in `flexible` strategy.
 
 ## 📋 Overview
 
@@ -15,12 +15,20 @@ This architecture lets you implement new training modes without modifying the co
 
 ### When You Need a Custom Strategy
 
+> [!NOTE]
+> The built-in `flexible` strategy already supports most conditioning scenarios out of the box:
+> first-frame conditioning, video extension (prefix/suffix), spatial crop (outpainting),
+> mask-based inpainting, IC-LoRA reference conditioning, and frozen modality cross-conditioning
+> (audio-to-video, video-to-audio). Only implement a custom strategy if your use case requires
+> fundamentally different training logic that cannot be expressed through the flexible strategy's
+> configuration.
+
 Consider implementing a custom strategy when you need:
 
-- **Different input modalities** (e.g., audio-only, audio-to-video conditioning)
-- **Additional conditioning signals** (e.g., masks for inpainting, depth maps)
-- **Custom loss computation** (e.g., weighted losses, auxiliary losses)
-- **Different noise application patterns** (e.g., partial masking)
+- **Custom loss computation** (e.g., weighted losses, auxiliary losses, perceptual losses)
+- **Non-standard noise application** (e.g., noise schedules different from flow matching)
+- **Novel conditioning mechanisms** not covered by the flexible strategy's condition types
+- **Additional model outputs** beyond the standard video/audio predictions
 
 ## 🏗️ Architecture Overview
 
@@ -28,7 +36,7 @@ Consider implementing a custom strategy when you need:
 
 The trainer delegates all training-mode-specific logic to the strategy:
 
-1. **Initialization** — The trainer calls `get_data_sources()` to determine which preprocessed data directories to load
+1. **Initialization** — The trainer calls `config.get_data_sources()` to determine which preprocessed data directories to load
 2. **Each training step:**
     - Calls `prepare_training_inputs()` to transform the raw batch into model-ready inputs
     - Runs the transformer forward pass
@@ -52,8 +60,8 @@ The trainer handles everything else: optimization, checkpointing, validation, an
 Before writing code, answer these questions:
 
 1. **What additional data does your strategy need?**
-    - Example: Inpainting needs mask latents alongside video latents
-    - Example: Audio-to-video needs reference audio embeddings
+    - Example: A perceptual-loss strategy may need auxiliary feature targets
+    - Example: A novel conditioning mechanism may need an additional precomputed directory
 
 2. **What does conditioning look like?**
     - Which tokens should be noised vs. kept clean?
@@ -164,6 +172,20 @@ class InpaintingConfig(TrainingStrategyConfigBase):
         ge=0.0,
         le=1.0,
     )
+
+    def get_data_sources(self) -> dict[str, str]:
+        """Define which data directories to load.
+
+        Returns a mapping of directory names (under preprocessed_data_root) to
+        batch keys. The trainer loads .pt files from each directory and exposes
+        them in the batch under the specified key. The trainer also uses this
+        mapping to validate that all required directories exist.
+        """
+        return {
+            "latents": "latents",  # -> batch["latents"]
+            "conditions": "conditions",  # -> batch["conditions"]
+            self.mask_latents_dir: "masks",  # -> batch["masks"]
+        }
 ```
 
 **Key points:**
@@ -171,6 +193,7 @@ class InpaintingConfig(TrainingStrategyConfigBase):
 - Inherit from `TrainingStrategyConfigBase`
 - Use `Literal["your_strategy_name"]` for the `name` field - this enables automatic strategy selection
 - Use Pydantic `Field` for validation and documentation
+- Implement `get_data_sources()` on the config — it's the single source of truth for data directories (used for both dataset wiring and existence validation)
 
 ### Step 4: Implement the Strategy Class
 
@@ -187,24 +210,6 @@ class InpaintingStrategy(TrainingStrategy):
     def __init__(self, config: InpaintingConfig):
         super().__init__(config)
 
-    @property
-    def requires_audio(self) -> bool:
-        """Whether this strategy requires audio components."""
-        return False  # Set to True if your strategy needs audio
-
-    def get_data_sources(self) -> dict[str, str]:
-        """Define which data directories to load.
-
-        Returns a mapping of directory names to batch keys.
-        The trainer will load .pt files from each directory and
-        make them available in the batch under the specified key.
-        """
-        return {
-            "latents": "latents",  # -> batch["latents"]
-            "conditions": "conditions",  # -> batch["conditions"]
-            self.config.mask_latents_dir: "masks",  # -> batch["masks"]
-        }
-
     def prepare_training_inputs(
             self,
             batch: dict[str, Any],
@@ -275,7 +280,6 @@ class InpaintingStrategy(TrainingStrategy):
             batch_size=batch_size,
             fps=24.0,  # Or get from latents_data
             device=device,
-            dtype=dtype,
         )
 
         # Create video Modality
@@ -328,7 +332,7 @@ You need to register your strategy in two places:
 from ltx_trainer.training_strategies.inpainting import InpaintingConfig, InpaintingStrategy
 
 # Add to the TrainingStrategyConfig type alias
-TrainingStrategyConfig = TextToVideoConfig | VideoToVideoConfig | InpaintingConfig
+TrainingStrategyConfig = TextToVideoConfig | VideoToVideoConfig | FlexibleStrategyConfig | InpaintingConfig
 
 # Add to __all__
 __all__ = [
@@ -356,7 +360,8 @@ from ltx_trainer.training_strategies.inpainting import InpaintingConfig
 TrainingStrategyConfig = Annotated[
     Annotated[TextToVideoConfig, Tag("text_to_video")]
     | Annotated[VideoToVideoConfig, Tag("video_to_video")]
-    | Annotated[InpaintingConfig, Tag("inpainting")],  # Add your config
+    | Annotated[FlexibleStrategyConfig, Tag("flexible")]
+    | Annotated[InpaintingConfig, Tag("inpainting")],
     Discriminator(_get_strategy_discriminator),
 ]
 ```
@@ -366,7 +371,7 @@ TrainingStrategyConfig = Annotated[
 Create an example config in `configs/`:
 
 ```yaml
-# configs/ltx2_inpainting_lora.yaml
+# configs/custom_inpainting_lora.yaml
 
 model:
   model_path: "/path/to/ltx2.safetensors"
@@ -408,8 +413,8 @@ The base `TrainingStrategy` class provides these helper methods:
 | `_audio_patchifier.patchify(latents)`        | Convert `[B, C, T, F]` → `[B, T, C*F]`          |
 | `_get_video_positions(...)`                  | Generate position embeddings for video          |
 | `_get_audio_positions(...)`                  | Generate position embeddings for audio          |
-| `_create_per_token_timesteps(mask, sigma)`   | Create timesteps with 0 for conditioning tokens |
-| `_create_first_frame_conditioning_mask(...)` | Create mask for first-frame conditioning        |
+| `_create_per_token_timesteps(conditioning_mask, sampled_sigma)` | Create timesteps with 0 for conditioning tokens |
+| `_create_first_frame_conditioning_mask(...)`                    | Create mask for first-frame conditioning        |
 
 ## 📊 Understanding ModelInputs
 
@@ -418,16 +423,14 @@ The `ModelInputs` dataclass contains everything needed for the forward pass and
 ```python
 @dataclass
 class ModelInputs:
-    video: Modality  # Video modality data
-    audio: Modality | None  # Audio modality (None if video-only)
+    video: Modality | None  # Video modality data
+    audio: Modality | None  # Audio modality data
 
-    video_targets: Tensor  # Target values for loss (velocity)
-    audio_targets: Tensor | None
+    video_targets: Tensor | None  # Target values for video loss (velocity)
+    audio_targets: Tensor | None  # Target values for audio loss (velocity)
 
-    video_loss_mask: Tensor  # Boolean: True = compute loss for this token
-    audio_loss_mask: Tensor | None
-
-    ref_seq_len: int | None = None  # For IC-LoRA: reference sequence length
+    video_loss_mask: Tensor | None  # Boolean loss mask for video tokens
+    audio_loss_mask: Tensor | None  # Boolean loss mask for audio tokens
 ```
 
 ## 📊 Understanding Modality
@@ -437,18 +440,20 @@ The `Modality` dataclass (from ltx-core) represents a single modality's data:
 ```python
 @dataclass(frozen=True)
 class Modality:
-    enabled: bool  # Whether this modality is active
-    latent: Tensor  # [B, seq_len, C] - the latent tokens
-    timesteps: Tensor  # [B, seq_len] - per-token timesteps (sigmas)
-    positions: Tensor  # [B, dims, seq_len, 2] - position bounds
-    context: Tensor  # [B, ctx_len, C] - text embeddings
-    context_mask: Tensor  # [B, ctx_len] - attention mask for context
+    latent: Tensor     # [B, T, D] — patchified latent tokens
+    sigma: Tensor      # [B,] — per-batch noise level (for cross-attn conditioning)
+    timesteps: Tensor  # [B, T] — per-token timestep embeddings
+    positions: Tensor  # [B, 3, T, 2] for video, [B, 1, T, 2] for audio — positional bounds
+    context: Tensor    # text conditioning embeddings
+    enabled: bool = True
+    context_mask: Tensor | None = None     # attention mask for text context
+    attention_mask: Tensor | None = None   # optional 2D self-attention mask [B, T, T]
 ```
 
 > [!NOTE]
 > **Per-token timesteps:** Each token in the sequence has its own timestep. Conditioning tokens—those that should remain
 > un-noised—must have `timestep=0`. This is how the model distinguishes clean reference tokens from tokens to denoise. Use
-`_create_per_token_timesteps(conditioning_mask, sigma)` to set this up correctly.
+> `_create_per_token_timesteps(conditioning_mask, sampled_sigma)` to set this up correctly.
 
 > [!NOTE]
 > `Modality` is immutable (frozen dataclass). Use `dataclasses.replace()` to create modified copies.
@@ -461,7 +466,7 @@ class Modality:
    from ltx_trainer.config import LtxTrainerConfig
    import yaml
 
-   with open('configs/ltx2_inpainting_lora.yaml') as f:
+   with open('configs/custom_inpainting_lora.yaml') as f:
        config = LtxTrainerConfig(**yaml.safe_load(f))
    print(f'Strategy: {config.training_strategy.name}')
    "
@@ -475,13 +480,13 @@ class Modality:
 
    config = InpaintingConfig()
    strategy = get_training_strategy(config)
-   print(f'Data sources: {strategy.get_data_sources()}')
+   print(f'Data sources: {config.get_data_sources()}')
    "
    ```
 
 3. **Run a short training test:**
    ```bash
-   uv run python scripts/train.py configs/ltx2_inpainting_lora.yaml
+   uv run python scripts/train.py configs/custom_inpainting_lora.yaml
    ```
 
 ## 💡 Tips and Best Practices
@@ -503,7 +508,8 @@ class Modality:
 
 Study these implementations for guidance:
 
-| Strategy                                                                           | Complexity | Key Features                                   |
-|------------------------------------------------------------------------------------|------------|------------------------------------------------|
-| [`TextToVideoStrategy`](../src/ltx_trainer/training_strategies/text_to_video.py)   | Simple     | First-frame conditioning, optional audio       |
-| [`VideoToVideoStrategy`](../src/ltx_trainer/training_strategies/video_to_video.py) | Medium     | Reference video concatenation, split loss mask |
+| Strategy | Complexity | Key Features |
+|----------|------------|--------------|
+| [`FlexibleStrategy`](../src/ltx_trainer/training_strategies/flexible.py) | Medium | Unified conditioning framework — supports all built-in modes |
+| [`TextToVideoStrategy`](../src/ltx_trainer/training_strategies/text_to_video.py) | Simple | First-frame conditioning, optional audio (deprecated) |
+| [`VideoToVideoStrategy`](../src/ltx_trainer/training_strategies/video_to_video.py) | Medium | Reference video concatenation, split loss mask (deprecated) |

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

@@ -33,41 +33,33 @@ uv run python scripts/split_scenes.py --help
 
 If your dataset doesn't include captions, you can automatically generate them using multimodal models that understand both video and audio.
 
+The default `qwen_omni` backend talks to a local vLLM server, which you launch once in a separate terminal:
+
 ```bash
-uv run python scripts/caption_videos.py scenes_output_dir/ \
-    --output scenes_output_dir/dataset.json
+# Terminal 1: start the captioner server (stays running)
+uv run python scripts/serve_captioner.py
 ```
 
-If you're running into VRAM issues, try enabling 8-bit quantization to reduce memory usage:
-
 ```bash
+# Terminal 2: caption your videos
 uv run python scripts/caption_videos.py scenes_output_dir/ \
-    --output scenes_output_dir/dataset.json \
-    --use-8bit
+    --output scenes_output_dir/dataset.json
 ```
 
 This will create a `dataset.json` file containing video paths and their captions.
 
 **Captioning options:**
 
-
-| Option             | Description                                                |
-| ------------------ | ---------------------------------------------------------- |
-| `--captioner-type` | `qwen_omni` (default, local) or `gemini_flash` (API)       |
-| `--use-8bit`       | Enable 8-bit quantization for lower VRAM usage             |
-| `--no-audio`       | Disable audio processing (video-only captions)             |
-| `--override`       | Re-caption files that already have captions                |
-| `--api-key`        | API key for Gemini Flash (or set `GOOGLE_API_KEY` env var) |
-
+| Option             | Description                                                      |
+| ------------------ | --------------------------------------------------------------- |
+| `--captioner-type` | `qwen_omni` (default, local vLLM server) or `gemini_flash` (API) |
+| `--vllm-url`       | Base URL of the vLLM server (default `http://127.0.0.1:8001/v1`) |
+| `--override`       | Re-caption files that already have captions                     |
+| `--api-key`        | Gemini API key (else `GEMINI_API_KEY`/`GOOGLE_API_KEY`; with no key, uses gcloud/Vertex AI auth) |
 
 **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
-- **On-screen text**: Any visible text overlays
+Each caption is a single, detailed paragraph describing both the visual content and the audio (speech, music, ambient sounds) of the clip. See the [Utility Scripts Reference](utility-scripts.md#automatic-video-captioning) for backend setup and the full list of options.
 
 > [!NOTE]
 > The automatically generated captions may contain inaccuracies or hallucinated content.
@@ -80,7 +72,7 @@ This step preprocesses your video dataset by:
 1. Resizing and cropping videos to fit specified resolution buckets
 2. Computing and caching video latent representations
 3. Computing and caching text embeddings for captions
-4. (Optional) Computing and caching audio latents
+4. Extracting and caching audio latents from videos (automatic, use `--skip-audio` to disable)
 
 > [!WARNING]
 > Very large videos (especially high spatial resolution and/or many frames) can cause GPU out-of-memory (OOM)
@@ -97,17 +89,9 @@ uv run python scripts/process_dataset.py dataset.json \
     --text-encoder-path /path/to/gemma-model
 ```
 
-### With Audio Processing
-
-For audio-video training, add the `--with-audio` flag:
-
-```bash
-uv run python scripts/process_dataset.py dataset.json \
-    --resolution-buckets "960x544x49" \
-    --model-path /path/to/ltx-2-model.safetensors \
-    --text-encoder-path /path/to/gemma-model \
-    --with-audio
-```
+Audio latents are automatically extracted from video files — no extra flag is needed. Use `--skip-audio`
+to disable this. For standalone audio files (`.wav`), use the `audio` column in your dataset instead
+(see [Convention-Based Column Detection](#convention-based-column-detection) below).
 
 ### 🚀 Multi-GPU Preprocessing
 
@@ -126,7 +110,7 @@ Outputs are written atomically (via a per-process temporary file, then renamed),
 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,
+> 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
@@ -152,7 +136,7 @@ The trainer supports videos, single images, or a mix of both in the same dataset
 > `--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
+> [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
@@ -160,7 +144,24 @@ The trainer supports videos, single images, or a mix of both in the same dataset
 > - 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:
+The dataset must be a CSV, JSON, or JSONL metadata file with columns for captions and media paths.
+
+#### Convention-Based Column Detection
+
+The preprocessing script automatically detects and processes columns based on their names. The following columns are recognized:
+
+| Column | Output Dir | Description |
+|--------|-----------|-------------|
+| `video` (or legacy `media_path`) | `latents/` | Target video to encode |
+| `audio` | `audio_latents/` | Explicit audio file (overrides auto-extraction from video) |
+| `caption` | `conditions/` | Text caption for the sample |
+| `reference_video` (or legacy `ref_media_path`) | `reference_latents/` | IC-LoRA reference video |
+| `reference_audio` | `reference_audio_latents/` | IC-LoRA reference audio |
+| `video_mask` | `video_masks/` | Binary mask for video inpainting |
+| `audio_mask` | `audio_masks/` | Binary mask for audio inpainting |
+
+> [!NOTE]
+> **Legacy column names:** `media_path` and `ref_media_path` are accepted as aliases for `video` and `reference_video` respectively. Existing datasets using these names will continue to work without modification.
 
 **JSON format example:**
 
@@ -168,11 +169,11 @@ The dataset must be a CSV, JSON, or JSONL metadata file with columns for caption
 [
   {
     "caption": "A cat playing with a ball of yarn",
-    "media_path": "videos/cat_playing.mp4"
+    "video": "videos/cat_playing.mp4"
   },
   {
     "caption": "A dog running in the park",
-    "media_path": "videos/dog_running.mp4"
+    "video": "videos/dog_running.mp4"
   }
 ]
 ```
@@ -180,18 +181,42 @@ The dataset must be a CSV, JSON, or JSONL metadata file with columns for caption
 **JSONL format example:**
 
 ```jsonl
-{"caption": "A cat playing with a ball of yarn", "media_path": "videos/cat_playing.mp4"}
-{"caption": "A dog running in the park", "media_path": "videos/dog_running.mp4"}
+{"caption": "A cat playing with a ball of yarn", "video": "videos/cat_playing.mp4"}
+{"caption": "A dog running in the park", "video": "videos/dog_running.mp4"}
 ```
 
 **CSV format example:**
 
 ```csv
-caption,media_path
+caption,video
 "A cat playing with a ball of yarn","videos/cat_playing.mp4"
 "A dog running in the park","videos/dog_running.mp4"
 ```
 
+**Additional dataset format examples:**
+
+Audio-only dataset:
+```json
+{"audio": "song.wav", "caption": "piano melody"}
+```
+
+V2V IC-LoRA with reference video:
+```json
+{"video": "clip.mp4", "reference_video": "depth.mp4", "caption": "depth to video"}
+```
+
+A2A IC-LoRA with reference audio:
+```json
+{"video": "clip.mp4", "reference_audio": "ref.wav", "caption": "match this style"}
+```
+This form auto-extracts the target audio from `clip.mp4`. For pure audio datasets, use `audio` plus
+`reference_audio` columns and preprocess with `--audio-durations`.
+
+Video inpainting with mask:
+```json
+{"video": "clip.mp4", "video_mask": "mask.mp4", "caption": "fill the sky"}
+```
+
 ### 📐 Resolution Buckets
 
 Videos are organized into "buckets" of specific dimensions (width × height × frames).
@@ -268,12 +293,31 @@ The preprocessed data is saved in a `.precomputed` directory:
 ```
 dataset/
 └── .precomputed/
-    ├── latents/            # Cached video latents
-    ├── conditions/         # Cached text embeddings
-    ├── audio_latents/      # (only if --with-audio) Cached audio latents
-    └── reference_latents/  # (only for IC-LoRA) Cached reference video latents
+    ├── latents/                  # Video latents
+    ├── conditions/               # Text embeddings
+    ├── audio_latents/            # Audio latents (auto-extracted or explicit)
+    ├── reference_latents/        # Reference video latents (IC-LoRA)
+    ├── reference_audio_latents/  # Reference audio latents (audio IC-LoRA)
+    ├── video_masks/              # Video masks (inpainting)
+    └── audio_masks/              # Audio masks (audio inpainting)
 ```
 
+Set `data.preprocessed_data_root` in your training config to this `.precomputed` directory — the parent directory that
+contains `latents/`, `conditions/`, and any mode-specific audio/reference/mask directories.
+
+## 🔊 Audio-Only Dataset Preprocessing
+
+For datasets containing only audio files (no `video` column), use `--audio-durations` to specify duration buckets:
+
+```bash
+uv run python scripts/process_dataset.py dataset.json \
+    --audio-durations "2.0;4.0;8.0" \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --text-encoder-path /path/to/gemma-model
+```
+
+The `--audio-durations` flag provides duration buckets (in seconds) for audio-only datasets. Since there is no video column to derive timing from, explicit duration buckets are required.
+
 ## 🪄 IC-LoRA Reference Video Preprocessing
 
 For IC-LoRA training, you need to preprocess datasets that include reference videos.
@@ -281,14 +325,16 @@ Reference videos provide the conditioning input while target videos represent th
 
 ### Dataset Format with Reference Videos
 
+The `reference_video` column is automatically detected by convention — no extra CLI flags are needed.
+
 **JSON format:**
 
 ```json
 [
   {
     "caption": "A cat playing with a ball of yarn",
-    "media_path": "videos/cat_playing.mp4",
-    "reference_path": "references/cat_playing_depth.mp4"
+    "video": "videos/cat_playing.mp4",
+    "reference_video": "references/cat_playing_depth.mp4"
   }
 ]
 ```
@@ -296,32 +342,39 @@ Reference videos provide the conditioning input while target videos represent th
 **JSONL format:**
 
 ```jsonl
-{"caption": "A cat playing with a ball of yarn", "media_path": "videos/cat_playing.mp4", "reference_path": "references/cat_playing_depth.mp4"}
-{"caption": "A dog running in the park", "media_path": "videos/dog_running.mp4", "reference_path": "references/dog_running_depth.mp4"}
+{"caption": "A cat playing with a ball of yarn", "video": "videos/cat_playing.mp4", "reference_video": "references/cat_playing_depth.mp4"}
+{"caption": "A dog running in the park", "video": "videos/dog_running.mp4", "reference_video": "references/dog_running_depth.mp4"}
 ```
 
 ### Preprocessing with Reference Videos
 
-To preprocess a dataset with reference videos, add the `--reference-column` argument specifying the name of the field
-in your dataset JSON/JSONL/CSV that contains the reference video paths:
+Convention-based detection means you just need the `reference_video` column in your dataset, and `process_dataset.py` will automatically detect and process it. No `--reference-column` flag is needed:
 
 ```bash
 uv run python scripts/process_dataset.py dataset.json \
     --resolution-buckets "960x544x49" \
     --model-path /path/to/ltx-2-model.safetensors \
     --text-encoder-path /path/to/gemma-model \
-    --reference-column "reference_path"
+    --reference-downscale-factor 2 \
+    --reference-temporal-scale-factor 1
 ```
 
 This will create an additional `reference_latents/` directory containing the preprocessed reference video latents.
+Use `--reference-downscale-factor` for spatial subsampling and `--reference-temporal-scale-factor` for temporal
+subsampling. Validation reference conditions should use matching `downscale_factor` and `temporal_scale_factor` values.
+
+> [!NOTE]
+> **Legacy column names:** If your dataset uses `ref_media_path`, it is accepted as an alias for `reference_video`.
 
 ### Generating Reference Videos
 
 **Dataset Requirements for IC-LoRA:**
 
 - Your dataset must contain paired videos where each target video has a corresponding reference video
-- Reference and target videos must have *identical* resolution and length
-- Both reference and target videos should be preprocessed together using the same resolution buckets
+- Reference and target videos should cover the same content. Reference videos can optionally be lower spatial
+  resolution or temporally subsampled (see Scaled Reference Conditioning in [Training Modes](training-modes.md)).
+- Both reference and target videos should be preprocessed together using the same target resolution buckets, plus any
+  reference scale factors you choose.
 
 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.
@@ -333,12 +386,63 @@ uv run python scripts/compute_reference.py scenes_output_dir/ \
 
 The script accepts a JSON file as the dataset configuration and updates it in-place by adding the filenames of the generated reference videos.
 
+> [!NOTE]
+> `compute_reference.py` writes generated references to the `reference_video` column, which `process_dataset.py`
+> detects automatically. The legacy `ref_media_path` column is also accepted.
+
 If you want to generate a different type of condition (depth maps, pose skeletons, etc.), modify or replace the `compute_reference()` function within this script.
 
 ### Example Dataset
 
 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.
 
+## 🎭 Mask Preprocessing for Inpainting
+
+For inpainting training with the `mask` condition type, provide `video_mask` or `audio_mask` columns in your dataset
+metadata. These columns point to mask media files (for example a mask image/video for video inpainting, or a waveform or
+`.pt` tensor for audio inpainting). `process_dataset.py` downsamples and thresholds them into per-sample `.pt` tensors
+under `video_masks/` or `audio_masks/`.
+
+### Processed Video Mask Format
+
+If you create masks manually instead of using `process_dataset.py`, save them as `.pt` files with the key `"mask"`
+containing a tensor of shape `[F, H, W]` where:
+
+- `F` = number of latent frames (temporal dimension)
+- `H` = latent height (pixel height / 32)
+- `W` = latent width (pixel width / 32)
+- Values are thresholded at `0.5`: values `> 0.5` are conditioning tokens (clean, excluded from loss),
+  and values `<= 0.5` are generated tokens (noised, contributes to loss).
+
+### Audio Mask Format
+
+Audio masks follow the same thresholding pattern as video masks but with shape `[T]` (temporal dimension only), where `T` is the number of audio latent frames. They are stored in `audio_masks/`.
+
+### Directory Structure
+
+Place masks in a directory within your preprocessed data root:
+
+```
+preprocessed_data_root/
+├── latents/          # Video latents
+├── conditions/       # Text embeddings
+├── video_masks/      # Video masks (one .pt per sample, matching latent filenames)
+└── audio_masks/      # Audio masks (one .pt per sample, matching latent filenames)
+```
+
+Then reference the mask directory in your training config:
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+    conditions:
+      - type: mask
+        mask_dir: "video_masks"
+```
+
 ## 🎯 LoRA Trigger Words
 
 When training a LoRA, you can specify a trigger token that will be prepended to all captions:
@@ -359,9 +463,9 @@ This acts as a trigger word that activates the LoRA during inference when you in
 
 ## 🔍 Decoding Videos for Verification
 
-If you add the `--decode` flag, the script will VAE-decode the precomputed latents and save the resulting videos
-in `.precomputed/decoded_videos`. When audio preprocessing is enabled (`--with-audio`), audio latents will also be
-decoded and saved to `.precomputed/decoded_audio`. This allows you to visually and audibly inspect the processed data.
+If you add the `--decode` flag, the script will VAE-decode the precomputed video latents and save the resulting videos
+in `.precomputed/decoded_videos`. Reference video latents are decoded to `.precomputed/decoded_reference_videos` when
+present. To inspect audio latents, run `scripts/decode_latents.py` with `--with-audio`.
 
 ```bash
 uv run python scripts/process_dataset.py dataset.json \
@@ -382,6 +486,4 @@ Once your dataset is preprocessed, you can proceed to:
 - Start training with the [Training Guide](training-guide.md)
 
 > [!TIP]
-> If your training recipe requires additional preprocessed data (e.g., masks, conditioning signals), see
-> [Implementing Custom Training Strategies](custom-training-strategies.md) for guidance on extending the
-> preprocessing pipeline.
+> The `flexible` strategy supports masks for inpainting (`mask` condition type) and spatial crop regions for outpainting (`spatial_crop` condition type) out of the box. For other custom preprocessing needs, see [Custom Training Strategies](custom-training-strategies.md).

packages/ltx-trainer/docs/quick-start.md

@@ -7,12 +7,14 @@ Get up and running with LTX-2 training in just a few steps!
 Before you begin, ensure you have:
 
 1. **LTX-2 Model Checkpoint** - A local `.safetensors` file containing the LTX-2 model weights.
-   Download `ltx-2-19b-dev.safetensors` from: [HuggingFace Hub](https://huggingface.co/Lightricks/LTX-2)
+   Download `ltx-2.3-22b-dev.safetensors` from: [HuggingFace Hub](https://huggingface.co/Lightricks/LTX-2.3)
+   The trainer supports LTX-2 and LTX-2.3 checkpoints through the same configuration API; version-specific components
+   are detected from the checkpoint.
 2. **Gemma Text Encoder** - A local directory containing the Gemma model (required for LTX-2).
    Download from: [HuggingFace Hub](https://huggingface.co/google/gemma-3-12b-it-qat-q4_0-unquantized/)
-3. **Linux with CUDA** - The trainer requires `triton` which is Linux-only
+3. **Linux with CUDA** - The trainer requires `triton` which is Linux-only; CUDA 13+ is recommended
 4. **GPU with sufficient VRAM** - 80GB recommended for the standard config. For GPUs with 32GB VRAM (e.g., RTX 5090),
-   use the [low VRAM config](../configs/ltx2_av_lora_low_vram.yaml) which enables INT8 quantization and other
+   use the [low VRAM config](../configs/t2v_lora_low_vram.yaml) which enables INT8 quantization and other
    memory optimizations
 
 ## ⚡ Installation
@@ -39,7 +41,18 @@ cd packages/ltx-trainer
 
 ## 🏋 Training Workflow
 
-### 1. Prepare Your Dataset
+If you are using an agent-enabled environment with repository skills, you can ask for the
+[`train-model`](../../../.claude/skills/train-model/SKILL.md) skill to run this workflow with you.
+It creates a run workspace, confirms the training mode, prepares data, preprocesses latents,
+launches training, and monitors the run while stopping for approval before expensive steps.
+
+### 1. Choose a Training Mode
+
+Start with [`t2v_lora.yaml`](../configs/t2v_lora.yaml) for a first run with videos and captions. For modes such as
+IC-LoRA, inpainting, or outpainting, check [Training Modes](training-modes.md) first because your metadata needs extra
+columns such as `reference_video`, `video_mask`, or `audio_mask` before preprocessing.
+
+### 2. Prepare Your Dataset
 
 Organize your videos and captions, then preprocess them:
 
@@ -57,15 +70,18 @@ uv run python scripts/process_dataset.py dataset.json \
     --text-encoder-path /path/to/gemma-model
 ```
 
+By default, preprocessing writes to `.precomputed/`. Use that directory as `data.preprocessed_data_root`
+in your training config.
+
 See [Dataset Preparation](dataset-preparation.md) for detailed instructions.
 
-### 2. Configure Training
+### 3. Configure Training
 
 Create or modify a configuration YAML file. Start with one of the example configs:
 
-- [`configs/ltx2_av_lora.yaml`](../configs/ltx2_av_lora.yaml) - Audio-video LoRA training
-- [`configs/ltx2_av_lora_low_vram.yaml`](../configs/ltx2_av_lora_low_vram.yaml) - Audio-video LoRA training (optimized for 32GB VRAM)
-- [`configs/ltx2_v2v_ic_lora.yaml`](../configs/ltx2_v2v_ic_lora.yaml) - IC-LoRA video-to-video
+- [`configs/t2v_lora.yaml`](../configs/t2v_lora.yaml) - Text-to-video LoRA
+- [`configs/t2v_lora_low_vram.yaml`](../configs/t2v_lora_low_vram.yaml) - Same as above, tuned for ~32GB VRAM (INT8 quantization and memory optimizations)
+- [`configs/v2v_ic_lora.yaml`](../configs/v2v_ic_lora.yaml) - IC-LoRA video-to-video
 
 Key settings to update:
 
@@ -82,33 +98,47 @@ output_dir: "outputs/my_training_run"
 
 See [Configuration Reference](configuration-reference.md) for all available options.
 
-### 3. Start Training
+### 4. Start Training
 
 ```bash
-uv run python scripts/train.py configs/ltx2_av_lora.yaml
+uv run python scripts/train.py configs/t2v_lora.yaml
 ```
 
 For multi-GPU training:
 
 ```bash
-uv run accelerate launch scripts/train.py configs/ltx2_av_lora.yaml
+uv run accelerate launch scripts/train.py configs/t2v_lora.yaml
 ```
 
 See [Training Guide](training-guide.md) for distributed training and advanced options.
 
 ## 🎯 Training Modes
 
+> [!TIP]
+> **First time?** Start with [`t2v_lora.yaml`](../configs/t2v_lora.yaml) — it's the simplest mode
+> and only requires videos with captions. You can explore other modes once you've confirmed your
+> setup works.
+
 The trainer supports several training modes:
 
-| Mode                 | Description                    | Config Example                             |
-|----------------------|--------------------------------|--------------------------------------------|
-| **LoRA**             | Efficient adapter training     | `training_strategy.name: "text_to_video"`  |
-| **Audio-Video LoRA** | Joint audio-video training     | `training_strategy.with_audio: true`       |
-| **IC-LoRA**          | Video-to-video transformations | `training_strategy.name: "video_to_video"` |
-| **Full Fine-tuning** | Full model training            | `model.training_mode: "full"`              |
+| Mode                  | Description                                | Example Config                                                    |
+|-----------------------|--------------------------------------------|-------------------------------------------------------------------|
+| **Text-to-Video**     | Generate video+audio from text prompts     | [`t2v_lora.yaml`](../configs/t2v_lora.yaml)                      |
+| **Image-to-Video**    | Animate from a starting image              | [`i2v_lora.yaml`](../configs/i2v_lora.yaml)                      |
+| **Video Extension**   | Extend videos temporally (forward/backward)| [`video_extend_lora.yaml`](../configs/video_extend_lora.yaml), [`video_suffix_lora.yaml`](../configs/video_suffix_lora.yaml) |
+| **IC-LoRA (V2V)**     | Video-to-video transformations             | [`v2v_ic_lora.yaml`](../configs/v2v_ic_lora.yaml)                |
+| **Audio-to-Video**    | Generate video conditioned on audio        | [`a2v_lora.yaml`](../configs/a2v_lora.yaml)                      |
+| **Video-to-Audio**    | Generate audio/foley from video            | [`v2a_lora.yaml`](../configs/v2a_lora.yaml)                      |
+| **Video Inpainting**  | Fill in masked regions of video            | [`video_inpainting_lora.yaml`](../configs/video_inpainting_lora.yaml) |
+| **Video Outpainting** | Extend video spatially                     | [`video_outpainting_lora.yaml`](../configs/video_outpainting_lora.yaml) |
+| **Text-to-Audio**     | Generate audio from text prompts           | [`t2a_lora.yaml`](../configs/t2a_lora.yaml)                      |
+| **Audio Extension**   | Extend audio temporally                    | [`audio_extend_lora.yaml`](../configs/audio_extend_lora.yaml), [`audio_suffix_lora.yaml`](../configs/audio_suffix_lora.yaml) |
+| **Audio Inpainting**  | Fill in masked regions of audio            | [`audio_inpainting_lora.yaml`](../configs/audio_inpainting_lora.yaml) |
+| **IC-LoRA (A2A)**     | Audio-to-audio transformations             | [`a2a_ic_lora.yaml`](../configs/a2a_ic_lora.yaml)                |
+| **AV2AV IC-LoRA**     | Audio+video IC-LoRA transformations        | [`av2av_ic_lora.yaml`](../configs/av2av_ic_lora.yaml)            |
+| **Full Fine-tuning**  | Full model training (any mode above)       | Set `model.training_mode: "full"`                                 |
 
-See [Training Modes](training-modes.md) for detailed explanations,
-or [Custom Training Strategies](custom-training-strategies.md) if you need to implement your own training recipe.
+See [Training Modes](training-modes.md) for detailed explanations of each mode.
 
 ## Next Steps
 
@@ -118,7 +148,7 @@ Once you've completed your first training run, you can:
   production-ready inference
   pipelines for various use cases (T2V, I2V, IC-LoRA, etc.). See the package documentation for details.
 - Learn more about [Dataset Preparation](dataset-preparation.md) for advanced preprocessing
-- Explore different [Training Modes](training-modes.md) (LoRA, Audio-Video, IC-LoRA)
+- Explore different [Training Modes](training-modes.md)
 - Dive deeper into [Training Configuration](configuration-reference.md)
 - Understand the model architecture in [LTX-Core Documentation](../../ltx-core/README.md)
 

packages/ltx-trainer/docs/training-guide.md

@@ -8,7 +8,7 @@ model uploads.
 After preprocessing your dataset and preparing a configuration file, you can start training using the trainer script:
 
 ```bash
-uv run python scripts/train.py configs/ltx2_av_lora.yaml
+uv run python scripts/train.py configs/t2v_lora.yaml
 ```
 
 The trainer will:
@@ -19,19 +19,31 @@ The trainer will:
 4. **Generate validation videos** (if configured)
 5. **Save the trained weights** in your output directory
 
+### Agent-Assisted Training
+
+If your environment supports repository skills, the
+[`train-model`](../../../.claude/skills/train-model/SKILL.md) skill provides an end-to-end
+orchestrator for this package. It asks what you want the model to learn, maps that intent to
+one of the documented [training modes](training-modes.md), probes your filesystem and GPU,
+prepares/preprocesses the dataset, writes a run-specific config, launches training, and
+monitors the job. It uses the trainer docs as its source of truth and stops for approval before
+captioning, preprocessing, or starting expensive training work.
+
 ### Output Files
 
 **For LoRA training:**
 
-- `lora_weights.safetensors` - Main LoRA weights file
+- `checkpoints/lora_weights_step_00000.safetensors` - LoRA checkpoint weights, with the current step in the filename
 - `training_config.yaml` - Copy of training configuration
-- `validation_samples/` - Generated validation videos (if enabled)
+- `samples/` - Generated validation samples (if enabled)
+- `checkpoints/training_state_step_00000.pt` - Optional resume state, depending on `checkpoints.save_training_state`
 
 **For full model fine-tuning:**
 
-- `model_weights.safetensors` - Full model weights
+- `checkpoints/model_weights_step_00000.safetensors` - Full model checkpoint weights, with the current step in the filename
 - `training_config.yaml` - Copy of training configuration
-- `validation_samples/` - Generated validation videos (if enabled)
+- `samples/` - Generated validation samples (if enabled)
+- `checkpoints/training_state_step_00000.pt` - Optional resume state, depending on `checkpoints.save_training_state`
 
 ## 🖥️ Distributed / Multi-GPU Training
 
@@ -62,22 +74,22 @@ Launch with a specific config using `--config_file`:
 # DDP (2 GPUs shown as example)
 CUDA_VISIBLE_DEVICES=0,1 \
 uv run accelerate launch --config_file configs/accelerate/ddp.yaml \
-  scripts/train.py configs/ltx2_av_lora.yaml
+  scripts/train.py configs/t2v_lora.yaml
 
 # DDP + torch.compile
 CUDA_VISIBLE_DEVICES=0,1 \
 uv run accelerate launch --config_file configs/accelerate/ddp_compile.yaml \
-  scripts/train.py configs/ltx2_av_lora.yaml
+  scripts/train.py configs/t2v_lora.yaml
 
 # FSDP (4 GPUs shown as example)
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
 uv run accelerate launch --config_file configs/accelerate/fsdp.yaml \
-  scripts/train.py configs/ltx2_av_lora.yaml
+  scripts/train.py configs/t2v_lora.yaml
 
 # FSDP + torch.compile
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
 uv run accelerate launch --config_file configs/accelerate/fsdp_compile.yaml \
-  scripts/train.py configs/ltx2_av_lora.yaml
+  scripts/train.py configs/t2v_lora.yaml
 ```
 
 **Notes:**
@@ -93,13 +105,13 @@ If you prefer to use your default Accelerate profile:
 
 ```bash
 # Use settings from your default accelerate config
-uv run accelerate launch scripts/train.py configs/ltx2_av_lora.yaml
+uv run accelerate launch scripts/train.py configs/t2v_lora.yaml
 
 # Override number of processes on the fly (e.g., 2 GPUs)
-uv run accelerate launch --num_processes 2 scripts/train.py configs/ltx2_av_lora.yaml
+uv run accelerate launch --num_processes 2 scripts/train.py configs/t2v_lora.yaml
 
 # Select specific GPUs
-CUDA_VISIBLE_DEVICES=0,1 uv run accelerate launch scripts/train.py configs/ltx2_av_lora.yaml
+CUDA_VISIBLE_DEVICES=0,1 uv run accelerate launch scripts/train.py configs/t2v_lora.yaml
 ```
 
 > [!TIP]

packages/ltx-trainer/docs/training-modes.md

@@ -1,165 +1,167 @@
 # Training Modes Guide
 
-The trainer supports several training modes, each suited for different use cases and requirements.
+The trainer uses the **flexible** training strategy (`name: "flexible"`) — a unified conditioning framework that
+supports all training modes through configuration. Every scenario is expressed by setting `is_generated` on each
+modality and adding optional conditions, rather than choosing a separate strategy class.
 
-## 🎯 Standard LoRA Training (Video-Only)
+## Key Concepts
 
-Standard LoRA (Low-Rank Adaptation) training fine-tunes the model by adding small, trainable adapter layers while
-keeping the base model frozen. This approach:
+Before diving into individual modes, here are the core ideas behind the flexible strategy:
 
-- **Requires significantly less memory and compute** than full fine-tuning
-- **Produces small, portable weight files** (typically a few hundred MB)
-- **Is ideal for learning specific styles, effects, or concepts**
-- **Can be easily combined with other LoRAs** during inference
+- **`is_generated: true`** — the modality is denoised during training and contributes to the loss. This is the
+  modality the model learns to generate.
+- **`is_generated: false`** — the modality is frozen (sigma=0, no noise, no loss). It passes through the transformer
+  clean and acts as cross-modal conditioning for the generated modality.
+- **At least one modality must have `is_generated: true`.**
+- **Conditions** are per-modality and can be composed (e.g., `reference` + `first_frame` together on the video
+  modality).
+- Audio does **not** support `first_frame` or `spatial_crop` conditions — only `prefix`, `suffix`, `mask`,
+  and `reference`.
 
-Configure standard LoRA training with:
+> [!TIP]
+> If you are using an agent-enabled environment with repository skills and are unsure which mode to choose,
+> ask for the [`train-model`](../../../.claude/skills/train-model/SKILL.md) skill. It maps your intent to one of
+> these configs and walks through dataset preparation, preprocessing, launch, and monitoring.
+
+## 📊 Quick Reference
+
+| Mode                  | Video     | Audio     | Conditions          | Config |
+|-----------------------|-----------|-----------|---------------------|--------|
+| **T2V**               | Generated | Generated | —                   | [`t2v_lora`](../configs/t2v_lora.yaml) |
+| **I2V**               | Generated | Generated | `first_frame`       | [`i2v_lora`](../configs/i2v_lora.yaml) |
+| **Video Extension**   | Generated | Generated | `prefix`/`suffix`   | [`video_extend_lora`](../configs/video_extend_lora.yaml) |
+| **V2V IC-LoRA**       | Generated | —         | `reference`         | [`v2v_ic_lora`](../configs/v2v_ic_lora.yaml) |
+| **A2V**               | Generated | Frozen    | —                   | [`a2v_lora`](../configs/a2v_lora.yaml) |
+| **V2A (Foley)**       | Frozen    | Generated | —                   | [`v2a_lora`](../configs/v2a_lora.yaml) |
+| **Video Inpainting**  | Generated | —         | `mask`              | [`video_inpainting_lora`](../configs/video_inpainting_lora.yaml) |
+| **Video Outpainting** | Generated | —         | `spatial_crop`      | [`video_outpainting_lora`](../configs/video_outpainting_lora.yaml) |
+| **T2A**               | —         | Generated | —                   | [`t2a_lora`](../configs/t2a_lora.yaml) |
+| **Audio Extension**   | —         | Generated | `prefix`/`suffix`   | [`audio_extend_lora`](../configs/audio_extend_lora.yaml) |
+| **Audio Inpainting**  | —         | Generated | `mask`              | [`audio_inpainting_lora`](../configs/audio_inpainting_lora.yaml) |
+| **A2A IC-LoRA**       | —         | Generated | `reference`         | [`a2a_ic_lora`](../configs/a2a_ic_lora.yaml) |
+| **AV2AV IC-LoRA**     | Generated | Generated | `reference` (both)  | [`av2av_ic_lora`](../configs/av2av_ic_lora.yaml) |
+
+---
+
+## 🎯 Text-to-Video (T2V)
+
+Generate video and audio from text prompts. Both modalities are denoised with no additional conditions.
 
 ```yaml
-model:
-  training_mode: "lora"
-
 training_strategy:
-  name: "text_to_video"
-  first_frame_conditioning_p: 0.1
-  with_audio: false  # Video-only training
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
 ```
 
-## 🔊 Audio-Video LoRA Training
+**Example config:** 📄 [t2v_lora.yaml](../configs/t2v_lora.yaml)
 
-LTX-2 supports joint audio-video generation. You can train LoRA adapters that affect both video and audio output:
+---
 
-- **Synchronized audio-video generation** - Audio matches the visual content
-- **Same efficient LoRA approach** - Just enable audio training
-- **Requires audio latents** - Dataset must include preprocessed audio
+## 🖼️ Image-to-Video (I2V)
 
-Configure audio-video training with:
+Generate video conditioned on a starting image. The first frame is provided as a clean conditioning signal — no noise,
+timestep=0, excluded from loss. The `probability` parameter controls how often first-frame conditioning is applied;
+remaining samples train in pure T2V mode.
 
 ```yaml
-model:
-  training_mode: "lora"
-
 training_strategy:
-  name: "text_to_video"
-  first_frame_conditioning_p: 0.1
-  with_audio: true  # Enable audio training
-  audio_latents_dir: "audio_latents"  # Directory containing audio latents
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+    conditions:
+      - type: first_frame
+        probability: 0.5
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
 ```
 
-**Example configuration file:**
+**Example config:** 📄 [i2v_lora.yaml](../configs/i2v_lora.yaml)
 
-- 📄 [Audio-Video LoRA Training](../configs/ltx2_av_lora.yaml)
+---
 
-**Dataset structure for audio-video training:**
+## ⏩ Video Extension
 
+Extend a video forward (or backward) in time. Prefix or suffix conditioning provides a span of existing latent frames
+as clean conditioning. The `temporal_boundary` sets the number of **latent frames** used as context (each latent frame
+= 8 pixel frames due to temporal compression).
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+    conditions:
+      - type: prefix            # or "suffix" for backward extension
+        temporal_boundary: 8    # 8 latent frames = 64 pixel frames
+        probability: 1.0
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
 ```
-preprocessed_data_root/
-├── latents/           # Video latents
-├── conditions/        # Text embeddings
-└── audio_latents/     # Audio latents (required when with_audio: true)
-```
-
-> [!IMPORTANT]
-> When training audio-video LoRAs, ensure your `target_modules` configuration captures video, audio, and
-> cross-modal attention branches. Use patterns like `"to_k"` instead of `"attn1.to_k"` to match:
-> - Video modules: `attn1.to_k`, `attn2.to_k`
-> - Audio modules: `audio_attn1.to_k`, `audio_attn2.to_k`
-> - Cross-modal modules: `audio_to_video_attn.to_k`, `video_to_audio_attn.to_k`
->
-> The cross-modal attention modules (`audio_to_video_attn` and `video_to_audio_attn`) enable bidirectional
-> information flow between audio and video, which is critical for synchronized audiovisual generation.
-> See [Understanding Target Modules](configuration-reference.md#understanding-target-modules) for detailed guidance.
 
 > [!NOTE]
-> You can generate audio during validation even if you're not training the audio branch.
-> Set `validation.generate_audio: true` independently of `training_strategy.with_audio`.
+> The `prefix` and `suffix` conditions also work on the audio modality for audio extension.
+> Set `temporal_boundary` on the audio modality's conditions list to condition on a prefix or suffix
+> of the audio latents.
 
-## 🔥 Full Model Fine-tuning
+**Example configs:** 📄 [video_extend_lora.yaml](../configs/video_extend_lora.yaml) (forward), 📄 [video_suffix_lora.yaml](../configs/video_suffix_lora.yaml) (backward)
 
-Full model fine-tuning updates all parameters of the base model, providing maximum flexibility but
-requiring substantial computational resources and larger training datasets:
+---
 
-- **Offers the highest potential quality and capability improvements**
-- **Requires multiple GPUs** and distributed training techniques (e.g., FSDP)
-- **Produces large checkpoint files** (several GB)
-- **Best for major model adaptations** or when LoRA limitations are reached
+## 🔄 IC-LoRA / Video-to-Video (V2V)
 
-Configure full fine-tuning with:
+In-Context LoRA learns transformations from paired videos. Pre-encoded reference latents are concatenated to the target
+sequence — reference tokens participate in bidirectional self-attention but receive no noise and are excluded from loss.
+This enables control adapters (depth, pose), style transfer, deblurring, colorization, and more.
 
 ```yaml
-model:
-  training_mode: "full"
-
 training_strategy:
-  name: "text_to_video"
-  first_frame_conditioning_p: 0.1
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+    conditions:
+      - type: reference
+        latents_dir: "reference_latents"
+        probability: 1.0
+      - type: first_frame       # optional — composable with reference
+        probability: 0.2
 ```
 
-> [!IMPORTANT]
-> Full fine-tuning of LTX-2 requires multiple high-end GPUs (e.g., 4-8× H100 80GB) and distributed
-> training with FSDP. See [Training Guide](training-guide.md) for multi-GPU setup instructions.
+> [!NOTE]
+> IC-LoRA is video-only by default (no audio modality block). Conditions can be composed — the example above also
+> applies first-frame conditioning with 20% probability alongside the reference.
+> Use [AV2AV IC-LoRA](#av2av-ic-lora) when both video and audio references should be trained jointly.
 
-## 🔄 In-Context LoRA (IC-LoRA) Training
+**Example config:** 📄 [v2v_ic_lora.yaml](../configs/v2v_ic_lora.yaml)
 
-IC-LoRA is a specialized training mode for video-to-video transformations.
-Unlike standard training modes that learn from individual videos, IC-LoRA learns transformations from pairs of videos.
-IC-LoRA enables a wide range of advanced video-to-video applications, such as:
+### Dataset Requirements
 
-- **Control adapters** (e.g., Depth, Pose): Learn to map from a control signal (like a depth map or pose skeleton) to a
-  target video
-- **Video deblurring**: Transform blurry input videos into sharp, high-quality outputs
-- **Style transfer**: Apply the style of a reference video to a target video sequence
-- **Colorization**: Convert grayscale reference videos into colorized outputs
-- **Restoration and enhancement**: Denoise, upscale, or restore old or degraded videos
+- **Paired videos** — each target video has a corresponding reference video
+- **Same frame count** between reference and target
+- Reference videos can optionally be at **lower spatial resolution** (see [Scaled Reference](#scaled-reference-conditioning) below)
+- Both must be **preprocessed** before training
 
-By providing paired reference and target videos, IC-LoRA can learn complex transformations that go beyond caption-based
-conditioning.
-
-IC-LoRA training fundamentally differs from standard LoRA and full fine-tuning:
-
-- **Reference videos** provide clean, unnoised conditioning input showing the "before" state
-- **Target videos** are noised during training and represent the desired "after" state
-- **The model learns transformations** from reference videos to target videos
-- **Loss is applied only to the target portion**, not the reference
-- **Training and inference time increase significantly** due to the doubled sequence length
-
-To enable IC-LoRA training, configure your YAML file with:
-
-```yaml
-model:
-  training_mode: "lora"  # Required: IC-LoRA uses LoRA mode
-
-training_strategy:
-  name: "video_to_video"
-  first_frame_conditioning_p: 0.1
-  reference_latents_dir: "reference_latents"  # Directory for reference video latents
-```
-
-**Example configuration file:**
-
-- 📄 [IC-LoRA Training](../configs/ltx2_v2v_ic_lora.yaml) - Video-to-video transformation training
-
-### Dataset Requirements for IC-LoRA
-
-- Your dataset must contain **paired videos** where each target video has a corresponding reference video
-- Reference and target videos must have the **same frame count** (length)
-- Reference videos can optionally be at **lower spatial resolution** than target videos (
-  see [Scaled Reference Conditioning](#scaled-reference-conditioning) below)
-- Both reference and target videos should be **preprocessed** before training
-
-**Dataset structure for IC-LoRA training:**
+**Dataset structure:**
 
 ```
 preprocessed_data_root/
-├── latents/            # Target video latents (what the model learns to generate)
-├── conditions/         # Text embeddings for each video
+├── latents/            # Target video latents
+├── conditions/         # Text embeddings
 └── reference_latents/  # Reference video latents (conditioning input)
 ```
 
 ### Generating Reference Videos
 
-We provide an example script to generate reference videos (e.g., Canny edge maps) for a given dataset.
-The script takes a JSON file as input (e.g., output of `caption_videos.py`) and updates it with the generated reference
-video paths.
+Use the `compute_reference.py` script to generate reference videos (e.g., Canny edge maps) for a dataset:
 
 ```bash
 uv run python scripts/compute_reference.py scenes_output_dir/ \
@@ -169,84 +171,392 @@ uv run python scripts/compute_reference.py scenes_output_dir/ \
 To compute a different condition (depth maps, pose skeletons, etc.), modify the `compute_reference()` function in the
 script.
 
-### Configuration Requirements for IC-LoRA
-
-- You **must** provide `reference_videos` in your validation configuration when using IC-LoRA training
-- The number of reference videos must match the number of validation prompts
-
-Example validation configuration for IC-LoRA:
-
-```yaml
-validation:
-  prompts:
-    - "First prompt describing the desired output"
-    - "Second prompt describing the desired output"
-  reference_videos:
-    - "/path/to/reference1.mp4"
-    - "/path/to/reference2.mp4"
-  reference_downscale_factor: 1  # Set to match preprocessing (e.g., 2 for half resolution)
-  include_reference_in_output: true  # Show reference side-by-side with output
-```
+> [!NOTE]
+> `compute_reference.py` writes generated references to the `reference_video` column, which
+> `process_dataset.py` detects automatically. The legacy `ref_media_path` column is also accepted.
 
 ### Scaled Reference Conditioning
 
-For more efficient training and inference, you can use **downscaled reference videos** while keeping target videos at
-full resolution. This reduces the number of conditioning tokens, leading to:
+For more efficient training and inference, use **downscaled reference videos** while keeping targets at full
+resolution. During training, the strategy infers the spatial and temporal scale factors from the preprocessed
+reference and target latents and adjusts positional encodings accordingly. This reduces conditioning tokens, leading to:
 
-- **Faster training** due to shorter sequence lengths
-- **Faster inference** with reduced memory usage
+- **Faster training** — shorter sequence lengths
+- **Faster inference** — reduced memory usage
 - **Same aspect ratio** maintained between reference and target
 
-#### How It Works
-
-When the reference video has resolution `H/n × W/n` and the target video has resolution `H × W`, the trainer
-automatically detects this scale factor `n` and adjusts the positional encodings so that the reference positions
-map to the correct locations in the target coordinate space.
-
-#### Preprocessing Datasets with Scaled References
-
-Use the `--reference-downscale-factor` option when running `process_dataset.py`:
+Preprocess with the `--reference-downscale-factor` option:
 
 ```bash
-# Process dataset with scaled reference videos (half resolution)
 uv run python scripts/process_dataset.py dataset.json \
     --resolution-buckets 768x768x25 \
     --model-path /path/to/ltx2.safetensors \
     --text-encoder-path /path/to/gemma \
-    --reference-column "reference_path" \
     --reference-downscale-factor 2
 ```
 
-This will:
+> [!NOTE]
+> The `reference_video` column is auto-detected by convention — no `--reference-column` flag needed.
 
-- Process target videos at 768×768 resolution
-- Process reference videos at 384×384 resolution (768 / 2)
-- The trainer will automatically infer the scale factor from the dimension ratio
-
-**Important**: Set `reference_downscale_factor: 2` in your validation configuration to match the preprocessing:
+Validation encodes reference media on the fly, so set `downscale_factor` and `temporal_scale_factor`
+on each `reference` validation condition to match the preprocessing factors:
 
 ```yaml
 validation:
-  reference_downscale_factor: 2  # Must match the preprocessing factor
-  reference_videos:
-    - "/path/to/reference1.mp4"
-    - "/path/to/reference2.mp4"
+  samples:
+    - prompt: "..."
+      conditions:
+        - type: reference
+          video: "/path/to/reference.mp4"
+          downscale_factor: 2
+          temporal_scale_factor: 1
+          include_in_output: true
 ```
 
 > [!NOTE]
 > The scale factor must be a positive integer, and all dimensions must be divisible by 32.
-> Common scale factors are 1 (no scaling), 2 (half resolution), or 4 (quarter resolution).
+> Common values are 1 (no scaling), 2 (half resolution), or 4 (quarter resolution).
 
-## 📊 Training Mode Comparison
+---
 
-| Aspect               | LoRA                           | Audio-Video LoRA               | Full Fine-tuning | IC-LoRA                        |
-|----------------------|--------------------------------|--------------------------------|------------------|--------------------------------|
-| **Memory Usage**     | Low                            | Low-Medium                     | High             | Medium                         |
-| **Training Speed**   | Fast                           | Fast                           | Slow             | Medium                         |
-| **Output Size**      | 100MB-few GB (depends on rank) | 100MB-few GB (depends on rank) | Tens of GB       | 100MB-few GB (depends on rank) |
-| **Flexibility**      | Medium                         | Medium                         | High             | Specialized                    |
-| **Audio Support**    | Optional                       | Yes                            | Optional         | No                             |
-| **Reference Videos** | No                             | No                             | No               | Yes (required)                 |
+## 🔊 Audio-to-Video (A2V)
+
+Generate video conditioned on frozen audio. Audio passes through the transformer clean (sigma=0) and influences video
+via the built-in cross-modal attention. Only video is denoised.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+  audio:
+    is_generated: false
+    latents_dir: "audio_latents"
+```
+
+**Example config:** 📄 [a2v_lora.yaml](../configs/a2v_lora.yaml)
+
+---
+
+## 🎵 Video-to-Audio / Foley (V2A)
+
+Generate audio (Foley) conditioned on frozen video. Video passes through the transformer clean (sigma=0) and
+conditions audio via cross-modal attention. Only audio is denoised.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: false
+    latents_dir: "latents"
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+```
+
+**Example config:** 📄 [v2a_lora.yaml](../configs/v2a_lora.yaml)
+
+---
+
+## 🎭 Video Inpainting
+
+Fill in masked regions of a video. Per-sample masks loaded from disk define which tokens are conditioning and which
+must be generated. Masks are thresholded at `0.5` to match validation/inference: tokens with `mask > 0.5` receive clean
+latents and timestep=0 and are excluded from loss; tokens with `mask <= 0.5` are denoised normally and contribute to
+loss.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+    conditions:
+      - type: mask
+        mask_dir: "video_masks"
+        probability: 1.0
+```
+
+**Dataset structure:**
+
+```
+preprocessed_data_root/
+├── latents/        # Video latents
+├── conditions/     # Text embeddings
+└── video_masks/    # Per-sample binary masks (1 → conditioning, 0 → generate)
+```
+
+In dataset metadata, provide mask media via the `video_mask` column; preprocessing converts it into `video_masks/`.
+
+**Example config:** 📄 [video_inpainting_lora.yaml](../configs/video_inpainting_lora.yaml)
+
+---
+
+## 🌅 Video Outpainting
+
+Extend a video spatially beyond its original boundaries. A rectangular pixel region is provided as clean conditioning
+(no noise, timestep=0, excluded from loss) — the model learns to generate the surrounding content. The `spatial_region`
+is specified in pixel coordinates `[y1, x1, y2, x2]` and automatically converted to latent space.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+    conditions:
+      - type: spatial_crop
+        spatial_region: [0, 0, 288, 576]  # y1, x1, y2, x2 in pixels
+        probability: 1.0
+```
+
+> [!NOTE]
+> `spatial_crop` is a video-only condition — it is not supported on the audio modality.
+
+**Example config:** 📄 [video_outpainting_lora.yaml](../configs/video_outpainting_lora.yaml)
+
+---
+
+## 🔈 Text-to-Audio (T2A)
+
+Generate audio from text prompts with no video modality. Only the audio branch of the transformer is denoised. Since
+no video modality is configured, this mode uses **audio-only LoRA targets** — explicitly targeting `audio_attn1`,
+`audio_attn2`, and `audio_ff` modules.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+```
+
+> [!NOTE]
+> With no `video` block in the strategy, the trainer only loads audio latents and text embeddings. LoRA adapters
+> should explicitly target audio modules (e.g., `audio_attn1.to_k`) rather than short patterns like `to_k` which
+> would also match video modules. See [LoRA Target Modules Guidance](#lora-target-modules-guidance) below.
+
+**Example config:** 📄 [t2a_lora.yaml](../configs/t2a_lora.yaml)
+
+---
+
+## 🔊 Audio Extension
+
+Extend audio forward (prefix) or backward (suffix) in time — the audio equivalent of Video Extension. A span of
+existing audio latent frames is provided as clean conditioning, and the model generates the continuation. The
+`temporal_boundary` sets the number of latent frames used as context. This mode uses **audio-only LoRA targets**.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+    conditions:
+      - type: prefix            # or "suffix" for backward extension
+        temporal_boundary: 8
+        probability: 1.0
+```
+
+**Example configs:** 📄 [audio_extend_lora.yaml](../configs/audio_extend_lora.yaml), 📄 [audio_suffix_lora.yaml](../configs/audio_suffix_lora.yaml)
+
+---
+
+## 🎭 Audio Inpainting
+
+Fill in masked regions of audio. Per-sample masks loaded from disk define which audio tokens are conditioning and
+which must be generated — the audio equivalent of Video Inpainting. Masks are thresholded at `0.5` with the same
+binary semantics as video inpainting. This mode uses **audio-only LoRA targets**.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+    conditions:
+      - type: mask
+        mask_dir: "audio_masks"
+        probability: 1.0
+```
+
+**Dataset structure:**
+
+```
+preprocessed_data_root/
+├── conditions/      # Text embeddings
+├── audio_latents/   # Audio latents
+└── audio_masks/     # Per-sample binary masks (1 → conditioning, 0 → generate)
+```
+
+In dataset metadata, provide mask media via the `audio_mask` column; preprocessing converts it into `audio_masks/`.
+
+**Example config:** 📄 [audio_inpainting_lora.yaml](../configs/audio_inpainting_lora.yaml)
+
+---
+
+## 🔄 IC-LoRA / Audio-to-Audio (A2A)
+
+In-Context LoRA for audio-to-audio transformations. Pre-encoded reference audio latents are concatenated to the target
+sequence — reference tokens participate in bidirectional self-attention but receive no noise and are excluded from loss.
+This enables audio style transfer, voice conversion, sound effect transformation, and more. This mode uses
+**audio-only LoRA targets**.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+    conditions:
+      - type: reference
+        latents_dir: "reference_audio_latents"
+        probability: 1.0
+```
+
+**Dataset structure:**
+
+```
+preprocessed_data_root/
+├── conditions/              # Text embeddings
+├── audio_latents/           # Target audio latents
+└── reference_audio_latents/ # Reference audio latents (conditioning input)
+```
+
+**Example config:** 📄 [a2a_ic_lora.yaml](../configs/a2a_ic_lora.yaml)
+
+---
+
+## 🔄 AV2AV IC-LoRA
+
+Joint audio-video In-Context LoRA — both modalities have reference conditioning. Pre-encoded reference latents are
+concatenated to each modality's target sequence independently. This enables joint audiovisual transformations such as
+synchronized style transfer across both video and audio.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+    conditions:
+      - type: reference
+        latents_dir: "reference_latents"
+        probability: 1.0
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+    conditions:
+      - type: reference
+        latents_dir: "reference_audio_latents"
+        probability: 1.0
+```
+
+> [!NOTE]
+> Unlike audio-only IC-LoRA (A2A), AV2AV uses short LoRA target patterns like `"to_k"` to match all branches
+> (video, audio, and cross-modal attention), since both modalities are trained.
+
+**Dataset structure:**
+
+```
+preprocessed_data_root/
+├── latents/                 # Target video latents
+├── audio_latents/           # Target audio latents
+├── conditions/              # Text embeddings
+├── reference_latents/       # Reference video latents (conditioning input)
+└── reference_audio_latents/ # Reference audio latents (conditioning input)
+```
+
+**Example config:** 📄 [av2av_ic_lora.yaml](../configs/av2av_ic_lora.yaml)
+
+---
+
+## 🔥 Full Model Fine-tuning
+
+All modes above default to `training_mode: "lora"`. For full fine-tuning, set `training_mode: "full"` — this updates
+all model parameters rather than adding LoRA adapters.
+
+```yaml
+model:
+  training_mode: "full"
+
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+```
+
+> [!IMPORTANT]
+> Full fine-tuning requires multiple high-end GPUs (e.g., 4-8× H100 80GB) and distributed training with FSDP.
+> See [Training Guide](training-guide.md) for multi-GPU setup instructions.
+
+---
+
+## 🎛️ LoRA Target Modules Guidance
+
+The `target_modules` configuration determines which transformer modules receive LoRA adapters. The right choice depends
+on whether your training involves cross-modal (audio ↔ video) interaction.
+
+**For T2V, I2V, A2V, V2A, or any mode involving both modalities** — use short patterns to match all branches
+(video, audio, and cross-modal attention):
+
+```yaml
+target_modules:
+  - "to_k"
+  - "to_q"
+  - "to_v"
+  - "to_out.0"
+```
+
+> [!IMPORTANT]
+> Short patterns like `"to_k"` match video modules (`attn1.to_k`, `attn2.to_k`), audio modules
+> (`audio_attn1.to_k`, `audio_attn2.to_k`), and cross-modal modules (`audio_to_video_attn.to_k`,
+> `video_to_audio_attn.to_k`). The cross-modal attention modules enable bidirectional information flow between
+> audio and video, which is critical for synchronized audiovisual generation.
+> See [Understanding Target Modules](configuration-reference.md#understanding-target-modules) for detailed guidance.
+
+**For video-only IC-LoRA** — explicitly target video modules (including FFN layers for better transformation quality):
+
+```yaml
+target_modules:
+  - "attn1.to_k"
+  - "attn1.to_q"
+  - "attn1.to_v"
+  - "attn1.to_out.0"
+  - "attn2.to_k"
+  - "attn2.to_q"
+  - "attn2.to_v"
+  - "attn2.to_out.0"
+  - "ff.net.0.proj"
+  - "ff.net.2"
+```
+
+**For audio-only modes (T2A, Audio Extension, Audio Inpainting, A2A IC-LoRA)** — explicitly target audio modules:
+
+```yaml
+target_modules:
+  - "audio_attn1.to_k"
+  - "audio_attn1.to_q"
+  - "audio_attn1.to_v"
+  - "audio_attn1.to_out.0"
+  - "audio_attn2.to_k"
+  - "audio_attn2.to_q"
+  - "audio_attn2.to_v"
+  - "audio_attn2.to_out.0"
+  - "audio_ff.net.0.proj"
+  - "audio_ff.net.2"
+```
+
+> [!NOTE]
+> Audio-only modes have no `video` block in the strategy, so there is no need to train video or cross-modal
+> attention modules. Targeting only `audio_*` modules keeps the LoRA small and focused.
+
+---
 
 ## 🎬 Using Trained Models for Inference
 
@@ -255,12 +565,25 @@ LoRAs:
 
 | Training Mode           | Recommended Pipeline                                  |
 |-------------------------|-------------------------------------------------------|
-| LoRA / Audio-Video LoRA | `TI2VidOneStagePipeline` or `TI2VidTwoStagesPipeline` |
-| IC-LoRA                 | `ICLoraPipeline`                                      |
+| T2V / I2V / A2V / Extension / Inpainting / Outpainting | `TI2VidOneStagePipeline` or `TI2VidTwoStagesPipeline` |
+| IC-LoRA (V2V / A2A / AV2AV) | `ICLoraPipeline`                                 |
+| V2A (Foley) / T2A / Audio Extension / Audio Inpainting | `TI2VidOneStagePipeline` or `TI2VidTwoStagesPipeline` |
 
 All pipelines support loading custom LoRAs via the `loras` parameter. See the [`ltx-pipelines`](../../ltx-pipelines/)
-package
-documentation for detailed usage instructions.
+package documentation for detailed usage instructions.
+
+> [!NOTE]
+> You can generate audio during validation even if you're not training the audio branch.
+> Set `validation.generate_audio: true` independently of whether audio has `is_generated: true`.
+
+---
+
+## 🔄 Migration from Legacy Strategies
+
+Legacy `text_to_video` and `video_to_video` strategy configs are forward-compatible and will continue to work (with a
+deprecation warning). We recommend migrating to `flexible` for access to all conditioning modes.
+
+---
 
 ## 🚀 Next Steps
 
@@ -272,6 +595,6 @@ Once you've chosen your training mode:
 
 > [!TIP]
 > Need a training mode that's not covered here?
-> See [Implementing Custom Training Strategies](custom-training-strategies.md)
-> to learn how to create your own strategy for specialized use cases like video inpainting, audio-only training, or
-> custom conditioning.
+> First check whether it can be expressed by composing existing `flexible` conditions. Use
+> [Implementing Custom Training Strategies](custom-training-strategies.md) only for custom losses,
+> noising rules, model outputs, or preprocessing that cannot be represented by configuration.

packages/ltx-trainer/docs/troubleshooting.md

@@ -8,7 +8,7 @@ Memory management is crucial for successful training with LTX-2.
 
 > [!TIP]
 > For GPUs with 32GB VRAM, use the pre-configured low VRAM config:
-> [`configs/ltx2_av_lora_low_vram.yaml`](../configs/ltx2_av_lora_low_vram.yaml)
+> [`configs/t2v_lora_low_vram.yaml`](../configs/t2v_lora_low_vram.yaml)
 > which combines 8-bit optimizer, INT8 quantization, and reduced LoRA rank.
 
 ### Memory Optimization Techniques
@@ -111,7 +111,7 @@ Ensure you've installed the dependencies and are using `uv run` to execute scrip
 # From the repository root
 uv sync
 cd packages/ltx-trainer
-uv run python scripts/train.py configs/ltx2_av_lora.yaml
+uv run python scripts/train.py configs/t2v_lora.yaml
 ```
 
 > [!TIP]
@@ -168,29 +168,29 @@ LTX-2 requires the number of frames to satisfy `frames % 8 == 1`:
 
    ```bash
    uv run accelerate launch --config_file configs/accelerate/ddp_compile.yaml \
-     scripts/train.py configs/ltx2_av_lora.yaml
+     scripts/train.py configs/t2v_lora.yaml
    ```
 
 ### Issue: Poor Quality Validation Outputs
 
 **Solutions:**
 
-1. **Use Image-to-Video Validation:**
-   For more reliable validation, use image-to-video (first-frame conditioning) rather than pure text-to-video:
+1. **Use conditioned validation:** For more reliable validation, use image-to-video (first-frame conditioning) rather than pure text-to-video:
 
    ```yaml
    validation:
-     prompts:
-       - "a professional portrait video of a person"
-     images:
-       - "/path/to/first_frame.png"  # One image per prompt
+     samples:
+       - prompt: "a professional portrait video of a person"
+         conditions:
+           - type: first_frame
+             image_or_video: "/path/to/first_frame.png"
    ```
 
 2. **Increase inference steps:**
 
    ```yaml
    validation:
-     inference_steps: 50  # Default is 30
+     inference_steps: 30
    ```
 
 3. **Adjust guidance settings:**

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

@@ -35,75 +35,53 @@ uv run python scripts/split_scenes.py video.mp4 scenes/ --max-scenes 50
 
 ### Automatic Video Captioning
 
-The `scripts/caption_videos.py` script generates captions for videos (with audio) using multimodal models.
+The `scripts/caption_videos.py` script generates a single, detailed combined audio-visual
+caption per video as a continuous paragraph of prose. Two backends are available:
+
+- **`qwen_omni` (default)** — Qwen3-Omni-30B-A3B-Thinking served via a local
+  [vLLM](https://docs.vllm.ai/) HTTP server (~1-3 s/video on H100). Highest quality, runs
+  fully offline once the model is downloaded.
+- **`gemini_flash`** — Google Gemini (cloud, `gemini-3.5-flash`). No GPU required. Auth is
+  automatic: set `GEMINI_API_KEY` (or `GOOGLE_API_KEY`) for the Developer API, or just have
+  Google Cloud credentials available (`gcloud auth` / an attached service account) and it
+  uses Vertex AI with no extra setup.
+
+**Step 1 — launch the captioner server** (`qwen_omni` only, one-time).
+
+`scripts/serve_captioner.py` runs vLLM in an isolated environment via `uvx`, so vLLM's heavy
+CUDA dependencies never touch the trainer's venv. It defaults to dynamic FP8 quantization
+(~31 GiB weights, fits on 40 GB GPUs, same speed as BF16 on H100):
 
 ```bash
-# Generate captions for all videos in a directory (uses Qwen2.5-Omni by default)
-uv run python scripts/caption_videos.py videos_dir/ --output dataset.json
+# Terminal 1 - stays running
+uv run python packages/ltx-trainer/scripts/serve_captioner.py
 
-# Use 8-bit quantization to reduce VRAM usage
-uv run python scripts/caption_videos.py videos_dir/ --output dataset.json --use-8bit
-
-# Use Gemini Flash API instead (requires API key)
-uv run python scripts/caption_videos.py videos_dir/ --output dataset.json \
-    --captioner-type gemini_flash --api-key YOUR_API_KEY
-
-# Use Gemini Flash with parallel workers for faster throughput
-uv run python scripts/caption_videos.py videos_dir/ --output dataset.json \
-    --captioner-type gemini_flash --num-workers 5
-
-# Caption without audio processing (video-only)
-uv run python scripts/caption_videos.py videos_dir/ --output dataset.json --no-audio
-
-# Force re-caption all files
-uv run python scripts/caption_videos.py videos_dir/ --output dataset.json --override
+# Useful variants:
+#   --print-cmd           show the vLLM command without running it
+#   --quantization bf16   use BF16 instead (needs ~66 GiB free VRAM)
+#   --hf-home /mnt/disk   override where the ~65 GB model is downloaded
 ```
 
-**Key features:**
-
-- **Audio-visual captioning**: Processes both video and audio content, including speech transcription
-- **Multiple backends**:
-  - `qwen_omni` (default): Local Qwen2.5-Omni model - processes video + audio locally
-  - `gemini_flash`: Google Gemini Flash API - cloud-based, requires API key
-- **Parallel captioning** (Gemini Flash only): Use `--num-workers` to run multiple API calls concurrently for faster throughput on large datasets
-- **Structured output**: Captions include visual description, speech transcription, sounds, and on-screen text
-- **Memory optimization**: 8-bit quantization option for limited VRAM
-- **Incremental processing**: Skips already-captioned files by default; progress is saved every 5 videos
-- **Multiple output formats**: JSON, JSONL, CSV, or TXT
-
-**Caption format:**
-
-The captioner produces structured captions with four sections:
-- `[VISUAL]`: Detailed description of visual content
-- `[SPEECH]`: Word-for-word transcription of spoken content
-- `[SOUNDS]`: Description of music, ambient sounds, sound effects
-- `[TEXT]`: Any on-screen text visible in the video
-
-**Parallel captioning with Gemini Flash:**
-
-When using `--captioner-type gemini_flash`, you can speed up large dataset captioning by running multiple API calls at the same time using `--num-workers` (accepts 1–10, default is 1):
+**Step 2 — caption your videos.**
 
 ```bash
-export GEMINI_API_KEY="your-key-here"
+# Terminal 2 - default backend talks to the server above
+uv run python packages/ltx-trainer/scripts/caption_videos.py videos_dir/ --output dataset.json
 
-# Caption a large dataset with 5 workers running concurrently
-uv run python scripts/caption_videos.py videos_dir/ \
-    --output dataset.json \
-    --captioner-type gemini_flash \
-    --num-workers 5
+# Remote server:           --vllm-url http://other-host:8001/v1
+# Gemini (gemini-3.5-flash): --captioner-type gemini_flash   (uses GEMINI_API_KEY, else gcloud/Vertex)
+# Gemini, parallel calls:  --captioner-type gemini_flash --num-workers 5
+# Re-caption everything:   --override
 ```
 
-> [!NOTE]
-> `--num-workers` is only supported with `gemini_flash`. Using it with `qwen_omni` or any other local model will raise an error, because local GPU models are not thread-safe.
+Captioning is incremental (already-captioned files are skipped, progress saves every 5 videos)
+and writes JSON, JSONL, CSV, or TXT based on the output extension.
 
-> [!TIP]
-> Keep `--num-workers` between 3–5 for most use cases. Very high values (8–10) may hit Gemini API rate limits depending on your quota tier.
+Qwen3-Omni-Thinking can optionally emit a `<think>...</think>` chain-of-thought before the
+caption (`--enable-thinking`). It is off by default, which is recommended for bulk captioning
+(thinking is slower as it generates the reasoning trace first).
 
-**Environment variables (for Gemini Flash):**
-
-Set one of these to use Gemini Flash without passing `--api-key`:
-- `GOOGLE_API_KEY`
-- `GEMINI_API_KEY`
+For Gemini, keep `--num-workers` at 3-5 (higher values may hit API rate limits).
 
 ### Dataset Preprocessing
 
@@ -116,13 +94,6 @@ uv run python scripts/process_dataset.py dataset.json \
     --model-path /path/to/ltx-2-model.safetensors \
     --text-encoder-path /path/to/gemma-model
 
-# With audio processing
-uv run python scripts/process_dataset.py dataset.json \
-    --resolution-buckets "960x544x49" \
-    --model-path /path/to/ltx-2-model.safetensors \
-    --text-encoder-path /path/to/gemma-model \
-    --with-audio
-
 # With video decoding for verification
 uv run python scripts/process_dataset.py dataset.json \
     --resolution-buckets "960x544x49" \
@@ -186,6 +157,10 @@ uv run python scripts/compute_reference.py videos_dir/ --output dataset.json
 > You can edit this script to generate other types of reference videos for IC-LoRA training,
 > such as depth maps, segmentation masks, or any custom video transformation.
 
+> [!NOTE]
+> `compute_reference.py` writes generated references to the `reference_video` column, which
+> `process_dataset.py` detects automatically.
+
 ## 🔍 Debugging and Verification Scripts
 
 ### Latents Decoding
@@ -224,73 +199,17 @@ uv run python scripts/decode_latents.py /path/to/latents/dir \
 - **Debug training data**: Visualize what the model actually sees during training
 - **Quality assessment**: Ensure latent encoding preserves important visual details
 
+### Inference with Trained Models
 
-### Inference Script
+For inference with trained LoRAs, use the [`ltx-pipelines`](../../ltx-pipelines/) package which provides
+production-ready pipelines:
 
-The `scripts/inference.py` script runs inference with a trained model.
+- **Text/Image-to-Video**: `TI2VidOneStagePipeline`, `TI2VidTwoStagesPipeline`
+- **Distilled (fast) inference**: `DistilledPipeline`
+- **IC-LoRA video-to-video**: `ICLoraPipeline`
+- **Keyframe interpolation**: `KeyframeInterpolationPipeline`
 
-> [!TIP]
-> For production inference, consider using the [`ltx-pipelines`](../../ltx-pipelines/) package which provides optimized,
-> feature-rich pipelines for various use cases:
-> - **Text/Image-to-Video**: `TI2VidOneStagePipeline`, `TI2VidTwoStagesPipeline`
-> - **Distilled (fast) inference**: `DistilledPipeline`
-> - **IC-LoRA video-to-video**: `ICLoraPipeline`
-> - **Keyframe interpolation**: `KeyframeInterpolationPipeline`
->
-> All pipelines support loading custom LoRAs trained with this trainer.
-
-```bash
-# Text-to-video inference (with audio by default)
-# By default, uses CFG scale 4.0 and STG scale 1.0 with block 29
-uv run python scripts/inference.py \
-    --checkpoint /path/to/model.safetensors \
-    --text-encoder-path /path/to/gemma \
-    --prompt "A cat playing with a ball" \
-    --output output.mp4
-
-# Video-only (skip audio generation)
-uv run python scripts/inference.py \
-    --checkpoint /path/to/model.safetensors \
-    --text-encoder-path /path/to/gemma \
-    --prompt "A cat playing with a ball" \
-    --skip-audio \
-    --output output.mp4
-
-# Image-to-video with conditioning image
-uv run python scripts/inference.py \
-    --checkpoint /path/to/model.safetensors \
-    --text-encoder-path /path/to/gemma \
-    --prompt "A cat walking" \
-    --condition-image first_frame.png \
-    --output output.mp4
-
-# Custom guidance settings
-uv run python scripts/inference.py \
-    --checkpoint /path/to/model.safetensors \
-    --text-encoder-path /path/to/gemma \
-    --prompt "A cat playing with a ball" \
-    --guidance-scale 4.0 \
-    --stg-scale 1.0 \
-    --stg-blocks 29 \
-    --output output.mp4
-
-# Disable STG (CFG only)
-uv run python scripts/inference.py \
-    --checkpoint /path/to/model.safetensors \
-    --text-encoder-path /path/to/gemma \
-    --prompt "A cat playing with a ball" \
-    --stg-scale 0.0 \
-    --output output.mp4
-```
-
-**Guidance parameters:**
-
-| Parameter | Default | Description |
-|-----------|---------|-------------|
-| `--guidance-scale` | 4.0 | CFG (Classifier-Free Guidance) scale |
-| `--stg-scale` | 1.0 | STG (Spatio-Temporal Guidance) scale. 0.0 disables STG |
-| `--stg-blocks` | 29 | Transformer block(s) to perturb for STG |
-| `--stg-mode` | stg_av | `stg_av` perturbs both audio and video, `stg_v` video only |
+All pipelines support loading custom LoRAs trained with this trainer.
 
 ## 🚀 Training Scripts
 
@@ -300,13 +219,13 @@ Use `scripts/train.py` for both single GPU and multi-GPU runs:
 
 ```bash
 # Single-GPU training
-uv run python scripts/train.py configs/ltx2_av_lora.yaml
+uv run python scripts/train.py configs/t2v_lora.yaml
 
 # Multi-GPU (uses your accelerate config)
-uv run accelerate launch scripts/train.py configs/ltx2_av_lora.yaml
+uv run accelerate launch scripts/train.py configs/t2v_lora.yaml
 
 # Override number of processes
-uv run accelerate launch --num_processes 4 scripts/train.py configs/ltx2_av_lora.yaml
+uv run accelerate launch --num_processes 4 scripts/train.py configs/t2v_lora.yaml
 ```
 
 For detailed usage, see the [Training Guide](training-guide.md).
@@ -316,5 +235,5 @@ For detailed usage, see the [Training Guide](training-guide.md).
 - **Start with `--help`**: Always check available options for each script
 - **Test on small datasets**: Verify workflows with a few files before processing large datasets
 - **Use decode verification**: Always decode a few samples to verify preprocessing quality
-- **Monitor VRAM usage**: Use `--use-8bit` or quantization flags when running into memory issues
+- **Monitor VRAM usage**: Reach for quantization or lower-memory settings (e.g. FP8 for the captioner server) when running into memory issues
 - **Keep backups**: Make copies of important dataset files before running conversion scripts

packages/ltx-trainer/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ltx-trainer"
-version = "1.1.5"
+version = "1.1.6"
 description = "LTX-2 training, democratized."
 readme = "README.md"
 authors = [
@@ -12,9 +12,11 @@ dependencies = [
     "accelerate>=1.2.1",
     "av>=14.2.1",
     "bitsandbytes >=0.45.2; sys_platform == 'linux'",
+    "google-genai>=2.0",
     "huggingface-hub[hf-xet]>=0.31.4",
     "imageio>=2.37.0",
     "imageio-ffmpeg>=0.6.0",
+    "openai>=2.0",
     "opencv-python>=4.11.0.86",
     "optimum-quanto>=0.2.6",
     "pandas>=2.2.3",
@@ -25,12 +27,16 @@ dependencies = [
     "safetensors>=0.5.0",
     "scenedetect>=0.6.5.2",
     "sentencepiece>=0.2.0",
+    "soundfile>=0.12.1",
     "torch>=2.6.0",
     "torchaudio>=2.7.0",
-    "torchcodec>=0.8.1",
+    # torchcodec must match the torch version (it ships a torch-ABI C++ extension and declares
+    # no torch pin of its own); the 0.9 line matches torch 2.9. torchaudio>=2.9 routes
+    # torchaudio.load() through torchcodec, so audio preprocessing needs it installed.
+    "torchcodec>=0.8.1,<0.10",
     "torchvision>=0.21.0",
     "typer>=0.15.1",
-    "wandb>=0.19.11",
+    "wandb>=0.27.0",
     "setuptools>=79.0.0",
 ]
 
@@ -48,8 +54,11 @@ build-backend = "hatchling.build"
 
 
 [tool.ruff]
-target-version = "1.1.5"
+target-version = "1.1.6"
 line-length = 120
+# Restrict isort first-party detection to src/ so stray dirs (e.g. wandb/ run output)
+# next to pyproject.toml don't get classified as first-party packages. See ruff#10519.
+src = ["src"]
 
 [tool.ruff.lint]
 select = [

packages/ltx-trainer/scripts/caption_videos.py

@@ -2,29 +2,28 @@
 
 """
 Auto-caption videos with audio using multimodal models.
-This script provides a command-line interface for generating captions for videos
-(including audio) using multimodal models. It supports:
-- Qwen2.5-Omni: Local model for audio-visual captioning (default)
-- Gemini Flash: Cloud-based API for audio-visual captioning
-The paths to videos in the generated dataset/captions file will be RELATIVE to the
-directory where the output file is stored. This makes the dataset more portable and
-easier to use in different environments.
+Backends:
+- Qwen3-Omni-30B-A3B-Thinking via a local vLLM HTTP server (default,
+  ``qwen_omni``). Launch the server once with ``scripts/serve_captioner.py``.
+- Gemini Flash 3.5 via Google's API (``gemini_flash``).
+The paths in the output file are RELATIVE to the output file's directory,
+making the dataset portable.
 Basic usage:
-    # Caption a single video (includes audio by default)
-    caption_videos.py video.mp4 --output captions.json
-    # Caption all videos in a directory
-    caption_videos.py videos_dir/ --output captions.csv
-    # Caption with custom instruction
-    caption_videos.py video.mp4 --instruction "Describe what happens in this video in detail."
+    # Launch the captioner server once (separate terminal)
+    uv run python scripts/serve_captioner.py
+    # Caption a directory
+    caption_videos.py videos_dir/ --output captions.json
+    # Caption a single video with a custom prompt
+    caption_videos.py video.mp4 --output cap.json --instruction "Describe in detail."
 Advanced usage:
-    # Use Gemini Flash API (requires GEMINI_API_KEY or GOOGLE_API_KEY env var)
+    # Use Gemini Flash 3.5 (cloud, requires GEMINI_API_KEY)
     caption_videos.py videos_dir/ --captioner-type gemini_flash
-    # Use Gemini Flash with parallel workers (2-10 workers, cloud API only)
+    # Gemini with parallel workers
     caption_videos.py videos_dir/ --captioner-type gemini_flash --num-workers 5
-    # Disable audio processing (video-only captions)
-    caption_videos.py videos_dir/ --no-audio
-    # Process videos with specific extensions and save as JSON
-    caption_videos.py videos_dir/ --extensions mp4,mov,avi --output captions.json
+    # Talk to a remote vLLM server
+    caption_videos.py videos_dir/ --vllm-url http://192.168.1.10:8001/v1
+    # Enable Qwen3 chain-of-thought (slower, more detail)
+    caption_videos.py videos_dir/ --enable-thinking
 """
 
 import csv
@@ -33,7 +32,6 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
 from enum import Enum
 from pathlib import Path
 
-import torch
 import typer
 from rich.console import Console
 from rich.progress import (
@@ -45,9 +43,14 @@ from rich.progress import (
     TimeElapsedColumn,
     TimeRemainingColumn,
 )
-from transformers.utils.logging import disable_progress_bar
 
-from ltx_trainer.captioning import CaptionerType, MediaCaptioningModel, create_captioner
+from ltx_trainer.captioning import (
+    DEFAULT_QWEN_MODEL,
+    DEFAULT_VLLM_BASE_URL,
+    CaptionerType,
+    MediaCaptioningModel,
+    create_captioner,
+)
 
 VIDEO_EXTENSIONS = ["mp4", "avi", "mov", "mkv", "webm"]
 IMAGE_EXTENSIONS = ["jpg", "jpeg", "png"]
@@ -61,8 +64,6 @@ app = typer.Typer(
     help="Auto-caption videos with audio using multimodal models.",
 )
 
-disable_progress_bar()
-
 
 class OutputFormat(str, Enum):
     """Available output formats for captions."""
@@ -73,15 +74,13 @@ class OutputFormat(str, Enum):
     JSONL = "jsonl"  # JSON Lines file with one JSON object per line
 
 
-def caption_media(  # noqa: PLR0913
+def caption_media(
     input_path: Path,
     output_path: Path,
     captioner: MediaCaptioningModel,
     extensions: list[str],
     recursive: bool,
     fps: int,
-    include_audio: bool,
-    clean_caption: bool,
     output_format: OutputFormat,
     override: bool,
     num_workers: int = 1,
@@ -94,8 +93,6 @@ def caption_media(  # noqa: PLR0913
         extensions: List of media file extensions to include
         recursive: Whether to search subdirectories recursively
         fps: Frames per second to sample from videos (ignored for images)
-        include_audio: Whether to include audio in captioning
-        clean_caption: Whether to clean up captions
         output_format: Format to save the captions in
         override: Whether to override existing captions
         num_workers: Number of parallel workers (only for cloud-based captioners like Gemini)
@@ -149,10 +146,10 @@ def caption_media(  # noqa: PLR0913
         caption = captioner.caption(
             path=media_file,
             fps=fps,
-            include_audio=include_audio,
-            clean_caption=clean_caption,
         )
-        rel_path = str(media_file.resolve().relative_to(base_dir))
+        # Don't resolve the file itself, so a symlinked clip keeps its logical path under the
+        # dataset dir instead of jumping to its (possibly external) link target.
+        rel_path = str((media_file.parent.resolve() / media_file.name).relative_to(base_dir))
         return rel_path, caption
 
     with progress:
@@ -371,16 +368,31 @@ def main(  # noqa: PLR0913
         help="Type of captioner to use. Valid values: 'qwen_omni' (local), 'gemini_flash' (API)",
         case_sensitive=False,
     ),
-    device: str | None = typer.Option(
-        None,
-        "--device",
-        "-d",
-        help="Device to use for inference (e.g., 'cuda', 'cuda:0', 'cpu'). Only for local models.",
+    vllm_url: str = typer.Option(
+        DEFAULT_VLLM_BASE_URL,
+        "--vllm-url",
+        help=(
+            "Base URL of the vLLM OpenAI-compatible server (qwen_omni only). "
+            "Launch the server with `uv run python scripts/serve_captioner.py`."
+        ),
     ),
-    use_8bit: bool = typer.Option(
+    vllm_model: str = typer.Option(
+        DEFAULT_QWEN_MODEL,
+        "--vllm-model",
+        help="Served model identifier on the vLLM server (qwen_omni only).",
+    ),
+    enable_thinking: bool = typer.Option(
         False,
-        "--use-8bit",
-        help="Whether to use 8-bit precision for the captioning model (reduces memory usage)",
+        "--enable-thinking/--no-thinking",
+        help=(
+            "Let Qwen3-Omni produce a <think>...</think> chain-of-thought before the caption. "
+            "Off by default: ~5x slower with marginal quality benefit and occasional hallucinations."
+        ),
+    ),
+    max_tokens: int = typer.Option(
+        4096,
+        "--max-tokens",
+        help="Maximum new tokens to generate per caption (qwen_omni only).",
     ),
     instruction: str | None = typer.Option(
         None,
@@ -401,20 +413,14 @@ def main(  # noqa: PLR0913
         help="Search for media files in subdirectories recursively",
     ),
     fps: int = typer.Option(
-        3,
+        2,
         "--fps",
         "-f",
-        help="Frames per second to sample from videos (ignored for images)",
-    ),
-    include_audio: bool = typer.Option(
-        True,
-        "--audio/--no-audio",
-        help="Whether to include audio in captioning (for videos with audio tracks)",
-    ),
-    clean_caption: bool = typer.Option(
-        True,
-        "--clean-caption/--raw-caption",
-        help="Whether to clean up captions by removing common VLM patterns",
+        help=(
+            "Frames per second to sample from videos. 2 is a typical default; "
+            "lower values use less compute per video. Ignored for images and for the "
+            "Gemini backend (which decides its own sampling rate)."
+        ),
     ),
     override: bool = typer.Option(
         False,
@@ -441,35 +447,36 @@ def main(  # noqa: PLR0913
     ),
 ) -> None:
     """Auto-caption videos with audio using multimodal models.
-    This script supports audio-visual captioning using:
-    - Qwen2.5-Omni: Local model (default) - processes both video and audio
-    - Gemini Flash: Cloud API - requires GOOGLE_API_KEY environment variable
+    Backends:
+    - ``qwen_omni`` (default): Qwen3-Omni-30B-A3B-Thinking via a local vLLM
+      HTTP server. Launch the server once in a separate terminal with
+      ``uv run python scripts/serve_captioner.py``. The server stays loaded
+      across script invocations.
+    - ``gemini_flash``: Google Gemini (``gemini-3.5-flash``) via the google-genai SDK.
+      Auth is automatic -- ``GEMINI_API_KEY``/``GOOGLE_API_KEY`` for the Developer API,
+      or Google Cloud credentials (gcloud / service account) for Vertex AI with no env vars.
     The paths in the output file will be relative to the output file's directory.
     Examples:
-        # Caption videos with audio using Qwen2.5-Omni (default)
+        # Caption videos using the local vLLM server (default)
         caption_videos.py videos_dir/ -o captions.json
-        # Caption using Gemini Flash API
+        # Point at a remote vLLM server
+        caption_videos.py videos_dir/ -o captions.json --vllm-url http://other-host:8001/v1
+        # Caption using Gemini Flash 3.5
         caption_videos.py videos_dir/ -o captions.json -c gemini_flash
-        # Caption without audio (video-only)
-        caption_videos.py videos_dir/ -o captions.json --no-audio
         # Caption with custom instruction
         caption_videos.py video.mp4 -o captions.json -i "Describe this video in detail"
     """
 
-    # Parallel workers are only safe for cloud-based (stateless) captioners.
-    # Local models like Qwen-Omni hold GPU state and are not thread-safe.
+    # Parallel workers are only supported for the cloud Gemini backend; qwen_omni
+    # drives a single shared vLLM server and is captioned serially from here.
     if num_workers > 1 and captioner_type != CaptionerType.GEMINI_FLASH:
         console.print(
-            "[bold red]Error:[/] --num-workers > 1 is only supported with [bold]--captioner-type gemini_flash[/].\n"
-            "Local models (e.g. qwen_omni) run on GPU and are not thread-safe — "
-            "parallel calls would cause memory corruption or incorrect results.\n"
-            "Either set [bold]--num-workers 1[/] (default) or switch to [bold]--captioner-type gemini_flash[/]."
+            "[bold red]Error:[/] --num-workers > 1 is only supported with "
+            "[bold]--captioner-type gemini_flash[/]. Use --num-workers 1 (default) "
+            "for the qwen_omni backend."
         )
         raise typer.Exit(code=1)
 
-    # Determine device for local models
-    device_str = device or ("cuda" if torch.cuda.is_available() else "cpu")
-
     # Parse extensions
     ext_list = [ext.strip() for ext in extensions.split(",")]
 
@@ -490,14 +497,15 @@ def main(  # noqa: PLR0913
     output = Path(output).resolve()
     console.print(f"Output will be saved to [bold blue]{output}[/]")
 
-    # Initialize captioning model
-    with console.status("Loading captioning model...", spinner="dots"):
+    with console.status("Initializing captioner...", spinner="dots"):
         if captioner_type == CaptionerType.QWEN_OMNI:
             captioner = create_captioner(
                 captioner_type=captioner_type,
-                device=device_str,
-                use_8bit=use_8bit,
+                base_url=vllm_url,
+                model=vllm_model,
                 instruction=instruction,
+                max_tokens=max_tokens,
+                enable_thinking=enable_thinking,
             )
         elif captioner_type == CaptionerType.GEMINI_FLASH:
             captioner = create_captioner(
@@ -508,7 +516,7 @@ def main(  # noqa: PLR0913
         else:
             raise ValueError(f"Unsupported captioner type: {captioner_type}")
 
-        console.print(f"[bold green]✓[/] {captioner_type.value} captioning model loaded successfully")
+        console.print(f"[bold green]✓[/] {captioner_type.value} captioner ready")
 
     # Caption media files
     caption_media(
@@ -518,8 +526,6 @@ def main(  # noqa: PLR0913
         extensions=ext_list,
         recursive=recursive,
         fps=fps,
-        include_audio=include_audio,
-        clean_caption=clean_caption,
         output_format=output_format,
         override=override,
         num_workers=num_workers,

packages/ltx-trainer/scripts/compute_reference.py

@@ -2,7 +2,7 @@
 Compute reference videos for IC-LoRA training.
 This script provides a command-line interface for generating reference videos to be used for IC-LoRA training.
 Note that it reads and writes to the same file (the output of caption_videos.py),
-where it adds the "reference_path" field to the JSON.
+where it adds the "reference_video" field to the JSON.
 Basic usage:
     # Compute reference videos for all videos in a directory
     compute_reference.py videos_dir/ --output videos_dir/captions.json
@@ -11,7 +11,7 @@ Basic usage:
 # Standard library imports
 import json
 from pathlib import Path
-from typing import Dict
+from typing import Any
 
 # Third-party imports
 import cv2
@@ -37,6 +37,10 @@ from ltx_trainer.video_utils import read_video, save_video
 console = Console()
 disable_progress_bar()
 
+VIDEO_COLUMNS = ("video", "media_path")
+REFERENCE_VIDEO_COLUMN = "reference_video"
+LEGACY_REFERENCE_COLUMN = "reference_path"
+
 
 def compute_reference(
     images: torch.Tensor,
@@ -79,15 +83,15 @@ def compute_reference(
 
 def _get_meta_data(
     output_path: Path,
-) -> Dict[str, str]:
+) -> list[dict[str, Any]]:
     """Get set of existing reference video paths without loading the actual files.
     Args:
         output_path: Path to the reference video paths file
     Returns:
-        Dictionary mapping media paths to reference video paths
+        Dataset rows with media paths and captions
     """
     if not output_path.exists():
-        return {}
+        return []
 
     console.print(f"[bold blue]Reading meta data from [cyan]{output_path}[/]...[/]")
 
@@ -98,11 +102,18 @@ def _get_meta_data(
 
     except Exception as e:
         console.print(f"[bold yellow]Warning: Could not check meta data: {e}[/]")
-        return {}
+        return []
+
+
+def _get_media_path(item: dict[str, Any]) -> str:
+    for column in VIDEO_COLUMNS:
+        if column in item:
+            return item[column]
+    raise KeyError(f"Dataset row must contain one of {VIDEO_COLUMNS}")
 
 
 def _save_dataset_json(
-    reference_paths: Dict[str, str],
+    reference_paths: dict[str, str],
     output_path: Path,
 ) -> None:
     """Save dataset json with reference video paths.
@@ -115,17 +126,17 @@ def _save_dataset_json(
         json_data = json.load(f)
         new_json_data = json_data.copy()
         for i, item in enumerate(json_data):
-            media_path = item["media_path"]
+            media_path = _get_media_path(item)
             reference_path = reference_paths[media_path]
-            new_json_data[i]["reference_path"] = reference_path
+            new_json_data[i].pop(LEGACY_REFERENCE_COLUMN, None)
+            new_json_data[i][REFERENCE_VIDEO_COLUMN] = reference_path
 
     with output_path.open("w", encoding="utf-8") as f:
         json.dump(new_json_data, f, indent=2, ensure_ascii=False)
 
     console.print(f"[bold green]✓[/] Reference video paths saved to [cyan]{output_path}[/]")
-    console.print("[bold yellow]Note:[/] Use these files with ImageOrVideoDataset by setting:")
-    console.print("  reference_column='[cyan]reference_path[/]'")
-    console.print("  video_column='[cyan]media_path[/]'")
+    console.print("[bold yellow]Note:[/] Reference videos were written to the '[cyan]reference_video[/]' column.")
+    console.print("  [cyan]process_dataset.py[/] detects this column automatically for IC-LoRA preprocessing.")
 
 
 def process_media(
@@ -158,7 +169,7 @@ def process_media(
     def media_path_to_reference_path(media_file: Path) -> Path:
         return media_file.parent / (media_file.stem + "_reference" + media_file.suffix)
 
-    media_files = [base_dir / Path(sample["media_path"]) for sample in meta_data]
+    media_files = [base_dir / Path(_get_media_path(sample)) for sample in meta_data]
     for media_file in media_files:
         reference_path = media_path_to_reference_path(media_file)
         media_to_process.append(media_file)
@@ -178,18 +189,23 @@ def process_media(
     )
 
     # Process media files
-    media_paths = [item["media_path"] for item in meta_data]
+    media_paths = [_get_media_path(item) for item in meta_data]
     reference_paths = {rel_path: str(media_path_to_reference_path(Path(rel_path))) for rel_path in media_paths}
 
     with progress:
         task = progress.add_task("Computing condition on videos", total=len(media_to_process))
 
-        for media_file in media_to_process:
+        for media_file, rel_path in zip(media_to_process, media_paths, strict=True):
             progress.update(task, description=f"Processing [bold blue]{media_file.name}[/]")
 
-            rel_path = str(media_file.resolve().relative_to(base_dir))
+            # Key by the original media-path string (matches the dict seeded above). Avoid
+            # resolve()/relative_to here — they crash on symlinked or absolute media paths.
             reference_path = media_path_to_reference_path(media_file)
-            reference_paths[rel_path] = str(reference_path.relative_to(base_dir))
+            try:
+                ref_stored = str(reference_path.relative_to(base_dir))
+            except ValueError:
+                ref_stored = str(reference_path)  # absolute/out-of-tree: keep it next to the source
+            reference_paths[rel_path] = ref_stored
 
             if not reference_path.resolve().exists() or override:
                 try:

packages/ltx-trainer/scripts/decode_latents.py

@@ -310,7 +310,7 @@ def main(
         help="Device to use for computation",
     ),
     vae_tiling: bool = typer.Option(
-        default=False,
+        default=True,
         help="Enable VAE tiling for larger video resolutions",
     ),
     seed: int | None = typer.Option(

packages/ltx-trainer/scripts/inference.py

@@ -1,443 +0,0 @@
-#!/usr/bin/env python3
-# ruff: noqa: T201
-"""
-CLI script for running LTX video/audio generation inference.
-Usage:
-    # Text-to-Video + Audio (default behavior)
-    python scripts/inference.py --checkpoint path/to/model.safetensors \
-        --text-encoder-path path/to/gemma \
-        --prompt "A cat playing with a ball" --output output.mp4
-    # Video only (skip audio)
-    python scripts/inference.py --checkpoint path/to/model.safetensors \
-        --text-encoder-path path/to/gemma \
-        --prompt "A cat playing with a ball" --skip-audio --output output.mp4
-    # Image-to-Video
-    python scripts/inference.py --checkpoint path/to/model.safetensors \
-        --text-encoder-path path/to/gemma \
-        --prompt "A cat walking" --condition-image first_frame.png --output output.mp4
-    # Video-to-Video (IC-LoRA style)
-    python scripts/inference.py --checkpoint path/to/model.safetensors \
-        --text-encoder-path path/to/gemma \
-        --prompt "A cat turning into a dog" --reference-video input.mp4 --output output.mp4
-    # With LoRA weights
-    python scripts/inference.py --checkpoint path/to/model.safetensors \
-        --text-encoder-path path/to/gemma \
-        --lora-path path/to/lora.safetensors \
-        --prompt "A cat in my custom style" --output output.mp4
-"""
-
-import argparse
-import re
-from pathlib import Path
-
-import torch
-import torchaudio
-from peft import LoraConfig, get_peft_model, set_peft_model_state_dict
-from safetensors.torch import load_file
-from torchvision import transforms
-
-from ltx_trainer.model_loader import load_model
-from ltx_trainer.progress import StandaloneSamplingProgress
-from ltx_trainer.utils import open_image_as_srgb
-from ltx_trainer.validation_sampler import GenerationConfig, ValidationSampler
-from ltx_trainer.video_utils import read_video, save_video
-
-
-def load_image(image_path: str) -> torch.Tensor:
-    """Load an image and convert to tensor [C, H, W] in [0, 1]."""
-    image = open_image_as_srgb(image_path)
-    transform = transforms.ToTensor()
-    return transform(image)
-
-
-def extract_lora_target_modules(state_dict: dict[str, torch.Tensor]) -> list[str]:
-    """Extract target module names from LoRA checkpoint keys.
-    LoRA keys follow the pattern (after removing "diffusion_model." prefix):
-    - transformer_blocks.0.attn1.to_k.lora_A.weight
-    - transformer_blocks.0.ff.net.0.proj.lora_B.weight
-    This extracts the full module path like "transformer_blocks.0.attn1.to_k".
-    Using full paths is more robust than partial patterns.
-    """
-    target_modules = set()
-    # Pattern to extract everything before .lora_A or .lora_B
-    pattern = re.compile(r"(.+)\.lora_[AB]\.")
-
-    for key in state_dict:
-        match = pattern.match(key)
-        if match:
-            module_path = match.group(1)
-            target_modules.add(module_path)
-
-    return sorted(target_modules)
-
-
-def load_lora_weights(transformer: torch.nn.Module, lora_path: str | Path) -> torch.nn.Module:
-    """Load LoRA weights into the transformer model.
-    The LoRA rank and target modules are automatically detected from the checkpoint.
-    Alpha is set equal to rank (standard practice for inference).
-    Args:
-        transformer: The base transformer model
-        lora_path: Path to the LoRA weights (.safetensors)
-    Returns:
-        The transformer model with LoRA weights applied
-    """
-    print(f"Loading LoRA weights from {lora_path}...")
-
-    # Load the LoRA state dict
-    state_dict = load_file(str(lora_path))
-
-    # Remove "diffusion_model." prefix (ComfyUI-compatible format)
-    state_dict = {k.replace("diffusion_model.", "", 1): v for k, v in state_dict.items()}
-
-    # Extract target modules from the checkpoint
-    target_modules = extract_lora_target_modules(state_dict)
-    if not target_modules:
-        raise ValueError(f"Could not extract target modules from LoRA checkpoint: {lora_path}")
-    print(f"  Detected {len(target_modules)} target modules")
-
-    # Auto-detect rank from the first lora_A weight shape
-    lora_rank = None
-    for key, value in state_dict.items():
-        if "lora_A" in key and value.ndim == 2:
-            lora_rank = value.shape[0]
-            break
-    if lora_rank is None:
-        raise ValueError("Could not auto-detect LoRA rank from weights")
-    print(f"  LoRA rank: {lora_rank}")
-
-    # Create LoRA config and wrap the model
-    # Alpha = rank is standard for inference (maintains the trained scale)
-    lora_config = LoraConfig(
-        r=lora_rank,
-        lora_alpha=lora_rank,
-        target_modules=target_modules,
-        lora_dropout=0.0,
-        init_lora_weights=True,
-    )
-
-    # Wrap the transformer with PEFT to add LoRA layers
-    transformer = get_peft_model(transformer, lora_config)
-
-    # Load the LoRA weights
-    base_model = transformer.get_base_model()
-    set_peft_model_state_dict(base_model, state_dict)
-
-    print("✓ LoRA weights loaded successfully")
-    return transformer
-
-
-def main() -> None:  # noqa: PLR0912, PLR0915
-    parser = argparse.ArgumentParser(
-        description="LTX Video/Audio Generation",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-    )
-
-    # Model arguments
-    parser.add_argument(
-        "--checkpoint",
-        type=str,
-        required=True,
-        help="Path to model checkpoint (.safetensors)",
-    )
-    parser.add_argument(
-        "--text-encoder-path",
-        type=str,
-        required=True,
-        help="Path to Gemma text encoder directory",
-    )
-
-    # LoRA arguments
-    parser.add_argument(
-        "--lora-path",
-        type=str,
-        default=None,
-        help="Path to LoRA weights (.safetensors)",
-    )
-
-    # Generation arguments
-    parser.add_argument(
-        "--prompt",
-        type=str,
-        required=True,
-        help="Text prompt for generation",
-    )
-    parser.add_argument(
-        "--negative-prompt",
-        type=str,
-        default="",
-        help="Negative prompt",
-    )
-    parser.add_argument(
-        "--height",
-        type=int,
-        default=544,
-        help="Video height (must be divisible by 32)",
-    )
-    parser.add_argument(
-        "--width",
-        type=int,
-        default=960,
-        help="Video width (must be divisible by 32)",
-    )
-    parser.add_argument(
-        "--num-frames",
-        type=int,
-        default=97,
-        help="Number of video frames (must be k*8 + 1)",
-    )
-    parser.add_argument(
-        "--frame-rate",
-        type=float,
-        default=25.0,
-        help="Video frame rate",
-    )
-    parser.add_argument(
-        "--num-inference-steps",
-        type=int,
-        default=30,
-        help="Number of denoising steps",
-    )
-    parser.add_argument(
-        "--guidance-scale",
-        type=float,
-        default=4.0,
-        help="Classifier-free guidance scale (CFG)",
-    )
-    parser.add_argument(
-        "--stg-scale",
-        type=float,
-        default=1.0,
-        help="STG (Spatio-Temporal Guidance) scale. 0.0 disables STG. Default: 1.0",
-    )
-    parser.add_argument(
-        "--stg-blocks",
-        type=int,
-        nargs="*",
-        default=[29],
-        help="Which transformer blocks to perturb for STG. Default: 29 (single block).",
-    )
-    parser.add_argument(
-        "--stg-mode",
-        type=str,
-        default="stg_av",
-        choices=["stg_av", "stg_v"],
-        help="STG mode: 'stg_av' perturbs both audio and video, 'stg_v' perturbs video only",
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=42,
-        help="Random seed for reproducibility",
-    )
-
-    # Conditioning arguments
-    parser.add_argument(
-        "--condition-image",
-        type=str,
-        default=None,
-        help="Path to conditioning image for image-to-video generation",
-    )
-    parser.add_argument(
-        "--reference-video",
-        type=str,
-        default=None,
-        help="Path to reference video for video-to-video generation (IC-LoRA style)",
-    )
-    parser.add_argument(
-        "--include-reference-in-output",
-        action="store_true",
-        help="Include reference video side-by-side with generated output (only for V2V)",
-    )
-
-    # Audio arguments
-    parser.add_argument(
-        "--skip-audio",
-        action="store_true",
-        help="Skip audio generation (by default, audio is generated alongside video)",
-    )
-
-    # Output arguments
-    parser.add_argument(
-        "--output",
-        type=str,
-        required=True,
-        help="Output video path (.mp4)",
-    )
-    parser.add_argument(
-        "--audio-output",
-        type=str,
-        default=None,
-        help="Output audio path (.wav, optional - if not provided, audio will be embedded in video)",
-    )
-
-    # Device arguments
-    parser.add_argument(
-        "--device",
-        type=str,
-        default="cuda",
-        help="Device to run on (cuda/cpu)",
-    )
-
-    args = parser.parse_args()
-
-    # Validate conditioning arguments
-    if args.include_reference_in_output and args.reference_video is None:
-        parser.error("--include-reference-in-output requires --reference-video")
-
-    # Validate arguments
-    generate_audio = not args.skip_audio
-
-    print("=" * 80)
-    print("LTX Video/Audio Generation")
-    print("=" * 80)
-
-    # Determine if we need VAE encoder (for image or video conditioning)
-    need_vae_encoder = args.condition_image is not None or args.reference_video is not None
-
-    components = load_model(
-        checkpoint_path=args.checkpoint,
-        device="cpu",  # Load to CPU first, sampler will move to device as needed
-        dtype=torch.bfloat16,
-        with_video_vae_encoder=need_vae_encoder,
-        with_video_vae_decoder=True,
-        with_audio_vae_decoder=generate_audio,
-        with_vocoder=generate_audio,
-        with_text_encoder=True,
-        text_encoder_path=args.text_encoder_path,
-    )
-
-    # Apply LoRA weights if provided
-    transformer = components.transformer
-    if args.lora_path is not None:
-        transformer = load_lora_weights(transformer, args.lora_path)
-
-    # Load conditioning image if provided
-    condition_image = None
-    if args.condition_image:
-        print(f"Loading conditioning image from {args.condition_image}...")
-        condition_image = load_image(args.condition_image)
-
-    # Load reference video if provided
-    reference_video = None
-    if args.reference_video:
-        print(f"Loading reference video from {args.reference_video}...")
-        reference_video, ref_fps = read_video(args.reference_video, max_frames=args.num_frames)
-        print(f"  Loaded {reference_video.shape[0]} frames @ {ref_fps:.1f} fps")
-
-    # Determine generation mode
-    if args.reference_video is not None and args.condition_image is not None:
-        mode = "Video-to-Video + Image Conditioning (V2V+I2V)"
-    elif args.reference_video is not None:
-        mode = "Video-to-Video (V2V)"
-    elif args.condition_image is not None:
-        mode = "Image-to-Video (I2V)"
-    else:
-        mode = "Text-to-Video (T2V)"
-
-    print("\n" + "=" * 80)
-    print("Generation Parameters")
-    print("=" * 80)
-    print(f"Mode: {mode}")
-    print(f"Prompt: {args.prompt}")
-    if args.negative_prompt:
-        print(f"Negative prompt: {args.negative_prompt}")
-    print(f"Resolution: {args.width}x{args.height}")
-    print(f"Frames: {args.num_frames} @ {args.frame_rate} fps")
-    print(f"Inference steps: {args.num_inference_steps}")
-    print(f"CFG scale: {args.guidance_scale}")
-    if args.stg_scale > 0:
-        blocks_str = args.stg_blocks if args.stg_blocks else "all"
-        print(f"STG scale: {args.stg_scale} (mode: {args.stg_mode}, blocks: {blocks_str})")
-    else:
-        print("STG: disabled")
-    print(f"Seed: {args.seed}")
-    if args.lora_path:
-        print(f"LoRA: {args.lora_path}")
-    if condition_image is not None:
-        print(f"Conditioning: Image ({args.condition_image})")
-    if reference_video is not None:
-        print(f"Reference: Video ({args.reference_video})")
-        if args.include_reference_in_output:
-            print("  → Will include reference side-by-side in output")
-    if generate_audio:
-        video_duration = args.num_frames / args.frame_rate
-        print(f"Audio: Enabled (duration will match video: {video_duration:.2f}s)")
-    print("=" * 80)
-
-    print(f"\nGenerating {'video + audio' if generate_audio else 'video'}...")
-
-    # Create generation config
-    gen_config = GenerationConfig(
-        prompt=args.prompt,
-        negative_prompt=args.negative_prompt,
-        height=args.height,
-        width=args.width,
-        num_frames=args.num_frames,
-        frame_rate=args.frame_rate,
-        num_inference_steps=args.num_inference_steps,
-        guidance_scale=args.guidance_scale,
-        seed=args.seed,
-        condition_image=condition_image,
-        reference_video=reference_video,
-        generate_audio=generate_audio,
-        include_reference_in_output=args.include_reference_in_output,
-        stg_scale=args.stg_scale,
-        stg_blocks=args.stg_blocks,
-        stg_mode=args.stg_mode,
-    )
-
-    # Generate with progress bar
-    with StandaloneSamplingProgress(num_steps=args.num_inference_steps) as progress:
-        # Create sampler with progress context
-        sampler = ValidationSampler(
-            transformer=transformer,
-            vae_decoder=components.video_vae_decoder,
-            vae_encoder=components.video_vae_encoder,
-            text_encoder=components.text_encoder,
-            audio_decoder=components.audio_vae_decoder if generate_audio else None,
-            vocoder=components.vocoder if generate_audio else None,
-            sampling_context=progress,
-        )
-        video, audio = sampler.generate(
-            config=gen_config,
-            device=args.device,
-        )
-
-    # Save video
-    output_path = Path(args.output)
-    output_path.parent.mkdir(parents=True, exist_ok=True)
-
-    # Get audio sample rate from vocoder if audio was generated
-    audio_sample_rate = None
-    if audio is not None and components.vocoder is not None:
-        audio_sample_rate = components.vocoder.output_sampling_rate
-
-    save_video(
-        video_tensor=video,
-        output_path=output_path,
-        fps=args.frame_rate,
-        audio=audio,
-        audio_sample_rate=audio_sample_rate,
-    )
-    print(f"✓ Video saved to {args.output}")
-
-    # Save separate audio file if requested
-    if audio is not None and args.audio_output is not None:
-        audio_output_path = Path(args.audio_output)
-        audio_output_path.parent.mkdir(parents=True, exist_ok=True)
-
-        torchaudio.save(
-            str(audio_output_path),
-            audio.cpu(),
-            sample_rate=audio_sample_rate,
-        )
-        duration = audio.shape[1] / audio_sample_rate
-        print(f"✓ Audio saved: {duration:.2f}s at {audio_sample_rate}Hz")
-
-    print("\n" + "=" * 80)
-    print("Generation complete!")
-    print("=" * 80)
-
-
-if __name__ == "__main__":
-    main()

packages/ltx-trainer/scripts/process_captions.py

@@ -154,6 +154,17 @@ class CaptionsDataset(Dataset):
         else:
             raise ValueError("Expected `dataset_file` to be a path to a CSV, JSON, or JSONL file.")
 
+    def _embedding_output_path(self, media_path: Path) -> str:
+        """Output `.pt` path relative to the dataset dir; mirrors `process_videos._output_relative`
+        so caption keys match video/audio latent keys (and absolute paths don't escape output_dir)."""
+        data_root = self.dataset_file.parent
+        resolved = data_root / media_path  # pathlib: an absolute media_path overrides data_root
+        try:
+            rel = resolved.relative_to(data_root)
+        except ValueError:
+            rel = Path(*resolved.parts[1:]) if resolved.is_absolute() else resolved
+        return str(rel.with_suffix(".pt"))
+
     def _load_caption_data_from_csv(self) -> dict[str, str]:
         """Load captions from a CSV file and compute output embedding paths."""
         df = pd.read_csv(self.dataset_file)
@@ -166,8 +177,7 @@ class CaptionsDataset(Dataset):
         caption_data = {}
         for _, row in df.iterrows():
             media_path = Path(row[self.media_column].strip())
-            # Convert media path to embedding output path (same structure, .pt extension)
-            output_path = str(media_path.with_suffix(".pt"))
+            output_path = self._embedding_output_path(media_path)
             caption_data[output_path] = row[self.caption_column]
 
         return caption_data
@@ -188,8 +198,7 @@ class CaptionsDataset(Dataset):
                 raise ValueError(f"Key '{self.media_column}' not found in JSON entry: {entry}")
 
             media_path = Path(entry[self.media_column].strip())
-            # Convert media path to embedding output path (same structure, .pt extension)
-            output_path = str(media_path.with_suffix(".pt"))
+            output_path = self._embedding_output_path(media_path)
             caption_data[output_path] = entry[self.caption_column]
 
         return caption_data
@@ -206,8 +215,7 @@ class CaptionsDataset(Dataset):
                     raise ValueError(f"Key '{self.media_column}' not found in JSONL entry: {entry}")
 
                 media_path = Path(entry[self.media_column].strip())
-                # Convert media path to embedding output path (same structure, .pt extension)
-                output_path = str(media_path.with_suffix(".pt"))
+                output_path = self._embedding_output_path(media_path)
                 caption_data[output_path] = entry[self.caption_column]
 
         return caption_data
@@ -326,7 +334,8 @@ def compute_captions_embeddings(  # noqa: PLR0913
                 # TODO(batch-tokenization): When tokenizer supports batching, encode all prompts at once.
                 # For now, process one at a time:
                 for i in range(len(batch["prompt"])):
-                    hidden_states, prompt_attention_mask = text_encoder.encode(batch["prompt"][i], padding_side="left")
+                    encoded = text_encoder.encode([batch["prompt"][i]], padding_side="left")
+                    hidden_states, prompt_attention_mask = encoded[0]
                     video_prompt_embeds, audio_prompt_embeds = embeddings_processor.feature_extractor(
                         hidden_states, prompt_attention_mask, "left"
                     )

packages/ltx-trainer/scripts/process_dataset.py

@@ -1,14 +1,21 @@
 #!/usr/bin/env python3
 
 """
-Preprocess a video dataset by computing video clips latents and text captions embeddings.
-This script provides a command-line interface for preprocessing video datasets by computing
-latent representations of video clips and text embeddings of their captions. The preprocessed
-data can be used to accelerate training of video generation models and to save GPU memory.
+Preprocess a media dataset for LTX-2 training.
+Automatically detects dataset columns and processes each according to a convention table.
+Column names determine what gets encoded and where outputs go — no per-role CLI flags needed.
+Convention table:
+    video           → Video VAE    → latents/
+    audio           → Audio VAE    → audio_latents/
+    reference_video → Video VAE    → reference_latents/
+    reference_audio → Audio VAE    → reference_audio_latents/
+    video_mask      → (downsample) → video_masks/
+    audio_mask      → (downsample) → audio_masks/
+    caption         → Text encoder → conditions/
+Legacy aliases: media_path → video, ref_media_path → reference_video
 Basic usage:
-    python scripts/process_dataset.py /path/to/dataset.json --resolution-buckets 768x768x49 \
+    python scripts/process_dataset.py /path/to/dataset.json --resolution-buckets 768x768x49 \\
         --model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma
-The dataset must be a CSV, JSON, or JSONL file with columns for captions and video paths.
 """
 
 from pathlib import Path
@@ -16,7 +23,15 @@ from pathlib import Path
 import typer
 from decode_latents import LatentsDecoder
 from process_captions import compute_captions_embeddings
-from process_videos import compute_latents, compute_scaled_resolution_buckets, parse_resolution_buckets
+from process_videos import (
+    compute_audio_latents,
+    compute_audio_masks,
+    compute_latents,
+    compute_scaled_resolution_buckets,
+    compute_video_masks,
+    detect_dataset_columns,
+    parse_resolution_buckets,
+)
 from rich.console import Console
 
 from ltx_trainer import logger
@@ -27,52 +42,82 @@ console = Console()
 app = typer.Typer(
     pretty_exceptions_enable=False,
     no_args_is_help=True,
-    help="Preprocess a video dataset by computing video clips latents and text captions embeddings. "
-    "The dataset must be a CSV, JSON, or JSONL file with columns for captions and video paths.",
+    help="Preprocess a media dataset for LTX-2 training. "
+    "Automatically detects columns (video, audio, reference_video, reference_audio, caption) "
+    "and processes each with the appropriate encoder.",
 )
 
+_KNOWN_ROLES = {"video", "audio", "reference_video", "reference_audio", "video_mask", "audio_mask", "caption"}
+_LEGACY_ALIASES = {"media_path": "video", "ref_media_path": "reference_video"}
 
-def preprocess_dataset(  # noqa: PLR0913
+
+def preprocess_dataset(  # noqa: PLR0912, PLR0913, PLR0915
     dataset_file: str,
-    caption_column: str,
-    video_column: str,
-    resolution_buckets: list[tuple[int, int, int]],
-    batch_size: int,
-    output_dir: str | None,
-    lora_trigger: str | None,
-    vae_tiling: bool,
-    decode: bool,
+    resolution_buckets: list[tuple[int, int, int]] | None,
     model_path: str,
     text_encoder_path: str,
     device: str,
+    output_dir: str | None = None,
+    video_column: str | None = None,
+    caption_column: str | None = None,
+    batch_size: int = 1,
+    lora_trigger: str | None = None,
+    vae_tiling: bool = False,
+    decode: bool = False,
     remove_llm_prefixes: bool = False,
-    reference_column: str | None = None,
     reference_downscale_factor: int = 1,
-    with_audio: bool = False,
+    reference_temporal_scale_factor: int = 1,
+    skip_audio: bool = False,
+    audio_durations: list[float] | None = None,
     load_text_encoder_in_8bit: bool = False,
     overwrite: bool = False,
 ) -> None:
-    """Run the preprocessing pipeline with the given arguments."""
-    # Validate dataset file
+    """Run the preprocessing pipeline with convention-based column detection."""
     _validate_dataset_file(dataset_file)
 
-    # Set up output directories
+    # Detect columns and resolve roles
+    dataset_columns = detect_dataset_columns(dataset_file)
+    roles = _resolve_columns(dataset_columns, video_column, caption_column)
+
+    # Log detected roles
+    for role, col in sorted(roles.items()):
+        alias_note = f" (alias for '{role}')" if col != role else ""
+        logger.info(f"Detected column '{col}'{alias_note} → {role}")
+
+    # Validate: need at least caption
+    if "caption" not in roles:
+        raise ValueError(
+            f"No caption column found. Dataset has columns: {dataset_columns}. "
+            f"Expected 'caption' or use --caption-column to specify."
+        )
+
+    # Validate: need video or audio
+    has_video = "video" in roles
+    has_audio = "audio" in roles
+    if not has_video and not has_audio:
+        raise ValueError(
+            f"No media column found. Dataset has columns: {dataset_columns}. "
+            f"Expected 'video', 'audio', or 'media_path' (legacy)."
+        )
+
+    # Validate: video modes need resolution buckets
+    if has_video and not resolution_buckets:
+        raise ValueError("--resolution-buckets is required when the dataset has a video column.")
+
     output_base = Path(output_dir) if output_dir else Path(dataset_file).parent / ".precomputed"
-    conditions_dir = output_base / "conditions"
-    latents_dir = output_base / "latents"
 
     if lora_trigger:
         logger.info(f'LoRA trigger word "{lora_trigger}" will be prepended to all captions')
 
+    # --- Phase 1: Text encoder ---
     with free_gpu_memory_context():
-        # Process captions using the dedicated function
         compute_captions_embeddings(
             dataset_file=dataset_file,
-            output_dir=str(conditions_dir),
+            output_dir=str(output_base / "conditions"),
             model_path=model_path,
             text_encoder_path=text_encoder_path,
-            caption_column=caption_column,
-            media_column=video_column,
+            caption_column=roles["caption"],
+            media_column=roles.get("video") or roles.get("audio") or roles["caption"],
             lora_trigger=lora_trigger,
             remove_llm_prefixes=remove_llm_prefixes,
             batch_size=batch_size,
@@ -81,119 +126,177 @@ def preprocess_dataset(  # noqa: PLR0913
             overwrite=overwrite,
         )
 
-    # Process videos using the dedicated function
-    audio_latents_dir = None
-    if with_audio:
-        logger.info("Audio preprocessing enabled - will extract and encode audio from videos")
-        audio_latents_dir = output_base / "audio_latents"
+    # --- Phase 2: Video VAE (video, reference_video) ---
+    if has_video and resolution_buckets:
+        # Determine if audio should be auto-extracted from video files
+        auto_audio = not skip_audio and "audio" not in roles
 
-    with free_gpu_memory_context():
-        compute_latents(
-            dataset_file=dataset_file,
-            video_column=video_column,
-            resolution_buckets=resolution_buckets,
-            output_dir=str(latents_dir),
-            model_path=model_path,
-            batch_size=batch_size,
-            device=device,
-            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
-        if reference_column:
-            # Validate: scaled references with multiple buckets can cause ambiguous bucket matching
-            if reference_downscale_factor > 1 and len(resolution_buckets) > 1:
-                raise ValueError(
-                    "When using --reference-downscale-factor > 1, only a single resolution bucket is supported. "
-                    "Using multiple buckets with scaled references can cause ambiguous bucket matching "
-                    "(e.g., a 512x256 reference could match either the scaled-down 1024x512 bucket or the 512x256 "
-                    "bucket). Please use a single resolution bucket or set --reference-downscale-factor to 1."
-                )
-
-            # Calculate and validate scaled resolution buckets for reference videos
-            reference_buckets = compute_scaled_resolution_buckets(resolution_buckets, reference_downscale_factor)
-
-            if reference_downscale_factor > 1:
-                logger.info(
-                    f"Processing reference videos for IC-LoRA training at 1/{reference_downscale_factor} resolution..."
-                )
-                logger.info(f"Reference resolution buckets: {reference_buckets}")
-            else:
-                logger.info("Processing reference videos for IC-LoRA training...")
-
-            reference_latents_dir = output_base / "reference_latents"
+        audio_latents_dir = str(output_base / "audio_latents") if auto_audio else None
+        if auto_audio:
+            logger.info("Audio will be auto-extracted from video files (use --skip-audio to disable)")
 
+        with free_gpu_memory_context():
             compute_latents(
                 dataset_file=dataset_file,
-                main_media_column=video_column,
-                video_column=reference_column,
-                resolution_buckets=reference_buckets,
-                output_dir=str(reference_latents_dir),
+                video_column=roles["video"],
+                resolution_buckets=resolution_buckets,
+                output_dir=str(output_base / "latents"),
                 model_path=model_path,
                 batch_size=batch_size,
                 device=device,
                 vae_tiling=vae_tiling,
+                with_audio=auto_audio,
+                audio_output_dir=audio_latents_dir,
                 overwrite=overwrite,
             )
 
-    # Handle decoding if requested (for verification)
+        # Process reference video if present
+        if "reference_video" in roles:
+            if reference_downscale_factor > 1 and len(resolution_buckets) > 1:
+                raise ValueError(
+                    "When using --reference-downscale-factor > 1, only a single resolution bucket is supported."
+                )
+            if reference_temporal_scale_factor > 1 and len(resolution_buckets) > 1:
+                raise ValueError(
+                    "When using --reference-temporal-scale-factor > 1, only a single resolution bucket is supported."
+                )
+
+            reference_buckets = compute_scaled_resolution_buckets(resolution_buckets, reference_downscale_factor)
+            if reference_downscale_factor > 1:
+                logger.info(f"Processing reference videos at 1/{reference_downscale_factor} resolution...")
+            if reference_temporal_scale_factor > 1:
+                logger.info(
+                    f"Temporally subsampling reference videos by {reference_temporal_scale_factor}x "
+                    f"(VAE-aligned pattern)..."
+                )
+
+            with free_gpu_memory_context():
+                compute_latents(
+                    dataset_file=dataset_file,
+                    main_media_column=roles["video"],
+                    video_column=roles["reference_video"],
+                    resolution_buckets=reference_buckets,
+                    output_dir=str(output_base / "reference_latents"),
+                    model_path=model_path,
+                    batch_size=batch_size,
+                    device=device,
+                    vae_tiling=vae_tiling,
+                    overwrite=overwrite,
+                    temporal_subsample_factor=reference_temporal_scale_factor,
+                )
+
+    # --- Phase 2b: Masks (video_mask, audio_mask) — processed after video latents for alignment ---
+    if "video_mask" in roles and has_video:
+        compute_video_masks(
+            dataset_file=dataset_file,
+            mask_column=roles["video_mask"],
+            latents_dir=str(output_base / "latents"),
+            output_dir=str(output_base / "video_masks"),
+            main_media_column=roles["video"],
+        )
+
+    # --- Phase 3: Audio VAE (audio, reference_audio) ---
+    audio_roles_to_process = [
+        ("audio", "audio_latents"),
+        ("reference_audio", "reference_audio_latents"),
+    ]
+    active_audio_roles = [(role, subdir) for role, subdir in audio_roles_to_process if role in roles]
+
+    if active_audio_roles:
+        # Determine audio duration constraint: video bucket → max_duration, or explicit buckets
+        max_audio_duration = None
+        audio_duration_buckets = None
+        if has_video and resolution_buckets:
+            max_audio_duration = max(f for f, _h, _w in resolution_buckets) / 25.0
+        elif audio_durations:
+            audio_duration_buckets = audio_durations
+
+        for role, output_subdir in active_audio_roles:
+            with free_gpu_memory_context():
+                compute_audio_latents(
+                    dataset_file=dataset_file,
+                    audio_column=roles[role],
+                    output_dir=str(output_base / output_subdir),
+                    model_path=model_path,
+                    main_media_column=roles.get("video"),
+                    max_duration=max_audio_duration,
+                    duration_buckets=audio_duration_buckets,
+                    device=device,
+                    overwrite=overwrite,
+                )
+
+    # --- Phase 4: Audio masks (after audio latents exist for temporal alignment) ---
+    if "audio_mask" in roles:
+        audio_latents_source = output_base / "audio_latents"
+        if audio_latents_source.exists():
+            compute_audio_masks(
+                dataset_file=dataset_file,
+                mask_column=roles["audio_mask"],
+                audio_latents_dir=str(audio_latents_source),
+                output_dir=str(output_base / "audio_masks"),
+                main_media_column=roles.get("video") or roles.get("audio"),
+            )
+        else:
+            logger.warning("audio_mask column found but no audio_latents/ — run with audio first")
+
+    # --- Decode for verification ---
     if decode:
         logger.info("Decoding latents for verification...")
+        decoder = LatentsDecoder(model_path=model_path, device=device, vae_tiling=vae_tiling, with_audio=has_audio)
+        if has_video:
+            decoder.decode(output_base / "latents", output_base / "decoded_videos")
+        if "reference_video" in roles and (output_base / "reference_latents").exists():
+            decoder.decode(output_base / "reference_latents", output_base / "decoded_reference_videos")
 
-        decoder = LatentsDecoder(
-            model_path=model_path,
-            device=device,
-            vae_tiling=vae_tiling,
-            with_audio=with_audio,
-        )
-        decoder.decode(latents_dir, output_base / "decoded_videos")
-
-        # Also decode reference videos if they exist
-        if reference_column:
-            reference_latents_dir = output_base / "reference_latents"
-            if reference_latents_dir.exists():
-                logger.info("Decoding reference videos...")
-                decoder.decode(reference_latents_dir, output_base / "decoded_reference_videos")
-
-        # Decode audio latents if they exist
-        if with_audio and audio_latents_dir and audio_latents_dir.exists():
-            logger.info("Decoding audio latents...")
-            decoder.decode_audio(audio_latents_dir, output_base / "decoded_audio")
-
-    # Print summary
+    # --- Summary ---
     logger.info(f"Dataset preprocessing complete! Results saved to {output_base}")
-    if reference_column:
-        logger.info("Reference videos processed and saved to reference_latents/ directory for IC-LoRA training")
-    if with_audio:
-        logger.info("Audio latents saved to audio_latents/ directory for audio-video training")
+    produced = [d.name for d in output_base.iterdir() if d.is_dir() and not d.name.startswith("decoded")]
+    logger.info(f"Output directories: {', '.join(sorted(produced))}")
 
 
 def _validate_dataset_file(dataset_path: str) -> None:
     """Validate that the dataset file exists and has the correct format."""
     dataset_file = Path(dataset_path)
-
     if not dataset_file.exists():
         raise FileNotFoundError(f"Dataset file does not exist: {dataset_file}")
-
     if not dataset_file.is_file():
         raise ValueError(f"Dataset path must be a file, not a directory: {dataset_file}")
-
     if dataset_file.suffix.lower() not in [".csv", ".json", ".jsonl"]:
         raise ValueError(f"Dataset file must be CSV, JSON, or JSONL format: {dataset_file}")
 
 
+def _resolve_columns(
+    dataset_columns: set[str],
+    video_column_override: str | None = None,
+    caption_column_override: str | None = None,
+) -> dict[str, str]:
+    """Map canonical role names to actual dataset column names.
+    Returns a dict of role → column_name for recognized roles found in the dataset.
+    """
+    roles: dict[str, str] = {}
+    for col in dataset_columns:
+        role = _LEGACY_ALIASES.get(col, col)
+        if role in _KNOWN_ROLES:
+            roles[role] = col
+
+    if video_column_override and video_column_override in dataset_columns:
+        roles["video"] = video_column_override
+    if caption_column_override and caption_column_override in dataset_columns:
+        roles["caption"] = caption_column_override
+
+    return roles
+
+
 @app.command()
 def main(  # noqa: PLR0913
     dataset_path: str = typer.Argument(
         ...,
-        help="Path to metadata file (CSV/JSON/JSONL) containing captions and video paths",
+        help="Path to metadata file (CSV/JSON/JSONL) with columns matching the convention table",
     ),
-    resolution_buckets: str = typer.Option(
-        ...,
-        help='Resolution buckets in format "WxHxF;WxHxF;..." (e.g. "768x768x25;512x512x49")',
+    resolution_buckets: str | None = typer.Option(
+        default=None,
+        help='Resolution buckets in format "WxHxF;WxHxF;..." (e.g. "768x768x25"). '
+        "Required when dataset has a video column.",
     ),
     model_path: str = typer.Option(
         ...,
@@ -203,13 +306,13 @@ def main(  # noqa: PLR0913
         ...,
         help="Path to Gemma text encoder directory",
     ),
-    caption_column: str = typer.Option(
-        default="caption",
-        help="Column name containing captions in the dataset JSON/JSONL/CSV file",
+    caption_column: str | None = typer.Option(
+        default=None,
+        help="Override: treat this column as 'caption' (default: auto-detect 'caption')",
     ),
-    video_column: str = typer.Option(
-        default="media_path",
-        help="Column name containing video paths in the dataset JSON/JSONL/CSV file",
+    video_column: str | None = typer.Option(
+        default=None,
+        help="Override: treat this column as 'video' (default: auto-detect 'video' or 'media_path')",
     ),
     batch_size: int = typer.Option(
         default=1,
@@ -229,32 +332,43 @@ def main(  # noqa: PLR0913
     ),
     lora_trigger: str | None = typer.Option(
         default=None,
-        help="Optional trigger word to prepend to each caption (activates the LoRA during inference)",
+        help="Optional trigger word to prepend to each caption",
     ),
     decode: bool = typer.Option(
         default=False,
-        help="Decode and save latents after encoding (videos and audio) for verification",
+        help="Decode and save latents after encoding for verification",
     ),
     remove_llm_prefixes: bool = typer.Option(
         default=False,
         help="Remove LLM prefixes from captions",
     ),
-    reference_column: str | None = typer.Option(
+    skip_audio: bool = typer.Option(
+        default=False,
+        help="Don't extract audio from video files (audio extraction is on by default)",
+    ),
+    audio_durations: str | None = typer.Option(
         default=None,
-        help="Column name containing reference video paths (for video-to-video training)",
+        help='Audio duration buckets in seconds for audio-only datasets (e.g. "2.0;4.0;8.0"). '
+        "When set, audio files are trimmed to the best matching duration. "
+        "Not needed when a video column is present (audio duration derived from video bucket).",
     ),
     with_audio: bool = typer.Option(
         default=False,
-        help="Extract and encode audio from video files",
+        hidden=True,
+        help="[DEPRECATED: audio is now on by default, use --skip-audio to disable]",
     ),
     load_text_encoder_in_8bit: bool = typer.Option(
         default=False,
-        help="Load the Gemma text encoder in 8-bit precision to save GPU memory (requires bitsandbytes)",
+        help="Load the Gemma text encoder in 8-bit precision to save GPU memory",
     ),
     reference_downscale_factor: int = typer.Option(
         default=1,
-        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.",
+        help="Downscale factor for reference video resolution (e.g., 2 = half resolution for IC-LoRA)",
+    ),
+    reference_temporal_scale_factor: int = typer.Option(
+        default=1,
+        help="Temporal subsampling factor for reference videos (e.g., 2 = half frame rate, "
+        "VAE-aligned: keeps frame 0, then every Nth frame from frame 1 onwards)",
     ),
     overwrite: bool = typer.Option(
         default=False,
@@ -262,64 +376,53 @@ def main(  # noqa: PLR0913
         "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
+    """Preprocess a media dataset for LTX-2 training.
+    See module docstring for the convention table. Audio is auto-extracted from
+    video files by default — use --skip-audio to disable.
+    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:
-        # Process a dataset with LTX-2 model
-        python scripts/process_dataset.py dataset.json --resolution-buckets 768x768x25 \\
-            --model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma
-        # Process dataset with custom column names
-        python scripts/process_dataset.py dataset.json --resolution-buckets 768x768x25 \\
-            --model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\
-            --caption-column "text" --video-column "video_path"
-        # Process dataset with reference videos for IC-LoRA training
-        python scripts/process_dataset.py dataset.json --resolution-buckets 768x768x25 \\
-            --model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\
-            --reference-column "reference_path"
-        # Process dataset with scaled reference videos (half resolution) for efficient IC-LoRA
-        python scripts/process_dataset.py dataset.json --resolution-buckets 768x768x25 \\
-            --model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\
-            --reference-column "reference_path" --reference-downscale-factor 2
-        # Process dataset with audio for audio-video training
-        python scripts/process_dataset.py dataset.json --resolution-buckets 768x512x97 \\
-            --model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\
-            --with-audio
     """
-    parsed_resolution_buckets = parse_resolution_buckets(resolution_buckets)
-
-    if len(parsed_resolution_buckets) > 1:
+    # Handle deprecated --with-audio flag
+    if with_audio:
         logger.warning(
-            "Using multiple resolution buckets. "
-            "When training with multiple resolution buckets, you must use a batch size of 1."
+            "--with-audio is deprecated. Audio extraction is now on by default. Use --skip-audio to disable."
         )
 
-    # Validate reference_downscale_factor
+    parsed_buckets = parse_resolution_buckets(resolution_buckets) if resolution_buckets else None
+
+    if parsed_buckets and len(parsed_buckets) > 1:
+        logger.warning("Using multiple resolution buckets. Training batch size must be 1.")
+
     if reference_downscale_factor < 1:
         raise typer.BadParameter("--reference-downscale-factor must be >= 1")
 
-    if reference_downscale_factor > 1 and not reference_column:
-        logger.warning("--reference-downscale-factor specified but no --reference-column provided. Ignoring.")
+    if reference_temporal_scale_factor < 1:
+        raise typer.BadParameter("--reference-temporal-scale-factor must be >= 1")
+
+    parsed_audio_durations = None
+    if audio_durations:
+        parsed_audio_durations = [float(d) for d in audio_durations.split(";")]
+        if any(d <= 0 for d in parsed_audio_durations):
+            raise typer.BadParameter("All audio durations must be positive")
 
     preprocess_dataset(
         dataset_file=dataset_path,
-        caption_column=caption_column,
-        video_column=video_column,
-        resolution_buckets=parsed_resolution_buckets,
-        batch_size=batch_size,
-        output_dir=output_dir,
-        lora_trigger=lora_trigger,
-        vae_tiling=vae_tiling,
-        decode=decode,
+        resolution_buckets=parsed_buckets,
         model_path=model_path,
         text_encoder_path=text_encoder_path,
         device=device,
+        output_dir=output_dir,
+        video_column=video_column,
+        caption_column=caption_column,
+        batch_size=batch_size,
+        lora_trigger=lora_trigger,
+        vae_tiling=vae_tiling,
+        decode=decode,
         remove_llm_prefixes=remove_llm_prefixes,
-        reference_column=reference_column,
         reference_downscale_factor=reference_downscale_factor,
-        with_audio=with_audio,
+        reference_temporal_scale_factor=reference_temporal_scale_factor,
+        skip_audio=skip_audio,
+        audio_durations=parsed_audio_durations,
         load_text_encoder_in_8bit=load_text_encoder_in_8bit,
         overwrite=overwrite,
     )

packages/ltx-trainer/scripts/process_videos.py

@@ -41,6 +41,7 @@ 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
+from torchvision.transforms.functional import resize as tv_resize
 from transformers.utils.logging import disable_progress_bar
 
 from ltx_core.model.audio_vae import AudioProcessor
@@ -73,6 +74,10 @@ app = typer.Typer(
 )
 
 
+def _clamp_01(x: torch.Tensor) -> torch.Tensor:
+    return x.clamp_(0, 1)
+
+
 class MediaDataset(Dataset):
     """
     Dataset for processing video and image files.
@@ -92,6 +97,7 @@ class MediaDataset(Dataset):
         resolution_buckets: list[tuple[int, int, int]],
         reshape_mode: str = "center",
         with_audio: bool = False,
+        temporal_subsample_factor: int = 1,
     ) -> None:
         """
         Initialize the media dataset.
@@ -101,6 +107,8 @@ class MediaDataset(Dataset):
             resolution_buckets: List of (frames, height, width) tuples
             reshape_mode: How to crop videos ("center", "random")
             with_audio: Whether to extract audio from video files
+            temporal_subsample_factor: Factor for VAE-aligned temporal subsampling.
+                When > 1, keeps frame 0 then takes every Nth frame from frame 1 onwards.
         """
         super().__init__()
 
@@ -109,6 +117,7 @@ class MediaDataset(Dataset):
         self.resolution_buckets = resolution_buckets
         self.reshape_mode = reshape_mode
         self.with_audio = with_audio
+        self.temporal_subsample_factor = temporal_subsample_factor
 
         # First load main media paths
         self.main_media_paths = self._load_video_paths(main_media_column)
@@ -124,7 +133,7 @@ class MediaDataset(Dataset):
         # Set up video transforms
         self.transforms = transforms.Compose(
             [
-                transforms.Lambda(lambda x: x.clamp_(0, 1)),
+                transforms.Lambda(_clamp_01),
                 transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
             ]
         )
@@ -142,8 +151,8 @@ class MediaDataset(Dataset):
 
         # Compute relative path of the video
         data_root = self.dataset_file.parent
-        relative_path = str(video_path.relative_to(data_root))
-        media_relative_path = str(self.main_media_paths[index].relative_to(data_root))
+        relative_path = str(_output_relative(video_path, data_root))
+        media_relative_path = str(_output_relative(self.main_media_paths[index], data_root))
 
         if video_path.suffix.lower() in [".png", ".jpg", ".jpeg"]:
             media_tensor = self._preprocess_image(video_path)
@@ -185,97 +194,29 @@ class MediaDataset(Dataset):
 
     @staticmethod
     def _extract_audio(video_path: Path, target_duration: float) -> dict[str, torch.Tensor | int] | None:
-        """Extract audio track from a video file, trimmed to match video duration."""
-        try:
-            # torchaudio can extract audio from video files directly
-            # waveform shape: [channels, samples]
-            waveform, sample_rate = torchaudio.load(str(video_path))
-
-            # Trim or pad to target duration
-            target_samples = int(target_duration * sample_rate)
-            current_samples = waveform.shape[-1]
-
-            if current_samples > target_samples:
-                # Trim to target duration
-                waveform = waveform[..., :target_samples]
-            elif current_samples < target_samples:
-                # Pad with zeros to target duration
-                padding = target_samples - current_samples
-                waveform = torch.nn.functional.pad(waveform, (0, padding))
-                logger.warning(f"Padded audio to {target_duration:.2f} seconds for {video_path}")
-
-            return {"waveform": waveform, "sample_rate": sample_rate}
-
-        except Exception as e:
-            logger.debug(f"Could not extract audio from {video_path}: {e}")
+        """Extract audio track from a video file, trimmed/padded to match video duration."""
+        audio = _load_audio_from_file(video_path, max_duration=target_duration)
+        if audio is None:
             return None
 
-    def _load_video_paths(self, column: str) -> list[Path]:
-        """Load video paths from the specified data source."""
-        if self.dataset_file.suffix == ".csv":
-            return self._load_video_paths_from_csv(column)
-        elif self.dataset_file.suffix == ".json":
-            return self._load_video_paths_from_json(column)
-        elif self.dataset_file.suffix == ".jsonl":
-            return self._load_video_paths_from_jsonl(column)
+        # Pad if shorter than target (_load_audio_from_file only trims, doesn't pad)
+        target_samples = int(target_duration * audio.sampling_rate)
+        if audio.waveform.shape[-1] < target_samples:
+            padding = target_samples - audio.waveform.shape[-1]
+            waveform = torch.nn.functional.pad(audio.waveform, (0, padding))
+            logger.warning(f"Padded audio to {target_duration:.2f} seconds for {video_path}")
         else:
-            raise ValueError("Expected `dataset_file` to be a path to a CSV, JSON, or JSONL file.")
+            waveform = audio.waveform
 
-    def _load_video_paths_from_csv(self, column: str) -> list[Path]:
-        """Load video paths from a CSV file."""
-        df = pd.read_csv(self.dataset_file)
-        if column not in df.columns:
-            raise ValueError(f"Column '{column}' not found in CSV file")
+        return {"waveform": waveform, "sample_rate": audio.sampling_rate}
 
-        data_root = self.dataset_file.parent
-        video_paths = [data_root / Path(line.strip()) for line in df[column].tolist()]
-
-        # Validate that all paths exist
-        invalid_paths = [path for path in video_paths if not path.is_file()]
-        if invalid_paths:
-            raise ValueError(f"Found {len(invalid_paths)} invalid video paths. First few: {invalid_paths[:5]}")
-
-        return video_paths
-
-    def _load_video_paths_from_json(self, column: str) -> list[Path]:
-        """Load video paths from a JSON file."""
-        with open(self.dataset_file, "r", encoding="utf-8") as file:
-            data = json.load(file)
-
-        if not isinstance(data, list):
-            raise ValueError("JSON file must contain a list of objects")
-
-        data_root = self.dataset_file.parent
-        video_paths = []
-        for entry in data:
-            if column not in entry:
-                raise ValueError(f"Key '{column}' not found in JSON entry")
-            video_paths.append(data_root / Path(entry[column].strip()))
-
-        # Validate that all paths exist
-        invalid_paths = [path for path in video_paths if not path.is_file()]
-        if invalid_paths:
-            raise ValueError(f"Found {len(invalid_paths)} invalid video paths. First few: {invalid_paths[:5]}")
-
-        return video_paths
-
-    def _load_video_paths_from_jsonl(self, column: str) -> list[Path]:
-        """Load video paths from a JSONL file."""
-        data_root = self.dataset_file.parent
-        video_paths = []
-        with open(self.dataset_file, "r", encoding="utf-8") as file:
-            for line in file:
-                entry = json.loads(line)
-                if column not in entry:
-                    raise ValueError(f"Key '{column}' not found in JSONL entry")
-                video_paths.append(data_root / Path(entry[column].strip()))
-
-        # Validate that all paths exist
-        invalid_paths = [path for path in video_paths if not path.is_file()]
-        if invalid_paths:
-            raise ValueError(f"Found {len(invalid_paths)} invalid video paths. First few: {invalid_paths[:5]}")
-
-        return video_paths
+    def _load_video_paths(self, column: str) -> list[Path]:
+        """Load video paths from the specified data source, validating existence."""
+        paths = _load_paths_from_dataset(self.dataset_file, column)
+        invalid = [p for p in paths if not p.is_file()]
+        if invalid:
+            raise ValueError(f"Found {len(invalid)} invalid paths in '{column}'. First few: {invalid[:5]}")
+        return paths
 
     def _filter_valid_videos(self) -> None:
         """Filter out videos with insufficient frames."""
@@ -348,6 +289,11 @@ class MediaDataset(Dataset):
         # Trim video to target number of frames
         frames_resized = frames_resized[:target_num_frames]
 
+        # VAE-aligned temporal subsampling: keep frame 0, then every Nth frame
+        if self.temporal_subsample_factor > 1:
+            indices = _compute_temporal_subsample_indices(target_num_frames, self.temporal_subsample_factor)
+            frames_resized = frames_resized[indices]
+
         # Apply transforms to each frame and stack
         video = torch.stack([self.transforms(frame) for frame in frames_resized], dim=0)
 
@@ -434,7 +380,18 @@ class MediaDataset(Dataset):
         return media_tensor
 
 
-def compute_latents(  # noqa: PLR0913, PLR0915
+def _compute_temporal_subsample_indices(num_frames: int, factor: int) -> list[int]:
+    """Compute VAE-aligned temporal subsample indices.
+    Keeps frame 0 (the VAE's standalone first-frame latent), then takes every
+    ``factor``-th frame from frame 1 onwards.  This ensures each resulting
+    8-frame VAE group spans ``factor`` groups of the original video.
+    """
+    if factor == 1:
+        return list(range(num_frames))
+    return [0, *list(range(1, num_frames, factor))]
+
+
+def compute_latents(  # noqa: PLR0912, PLR0913, PLR0915
     dataset_file: str | Path,
     video_column: str,
     resolution_buckets: list[tuple[int, int, int]],
@@ -447,7 +404,9 @@ def compute_latents(  # noqa: PLR0913, PLR0915
     vae_tiling: bool = False,
     with_audio: bool = False,
     audio_output_dir: str | None = None,
+    num_dataloader_workers: int = 4,
     overwrite: bool = False,
+    temporal_subsample_factor: int = 1,
 ) -> None:
     """
     Process videos and save latent representations.
@@ -468,9 +427,28 @@ 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)
+        num_dataloader_workers: Number of DataLoader worker processes (0 for in-process loading)
         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.
+        temporal_subsample_factor: Factor for VAE-aligned temporal subsampling of reference videos
     """
+    # Validate temporal subsampling compatibility with resolution buckets
+    if temporal_subsample_factor > 1:
+        for frames, _h, _w in resolution_buckets:
+            pixel_frames_minus_one = frames - 1
+            if pixel_frames_minus_one % temporal_subsample_factor != 0:
+                raise ValueError(
+                    f"Frame count {frames} is not compatible with "
+                    f"temporal_subsample_factor={temporal_subsample_factor}. "
+                    f"(frames - 1) must be divisible by the factor."
+                )
+            subsampled = 1 + pixel_frames_minus_one // temporal_subsample_factor
+            if (subsampled - 1) % VAE_TEMPORAL_FACTOR != 0:
+                raise ValueError(
+                    f"After temporal subsampling {frames} → {subsampled} frames, "
+                    f"result does not satisfy (frames - 1) % {VAE_TEMPORAL_FACTOR} == 0."
+                )
+
     if with_audio and audio_output_dir is None:
         raise ValueError("audio_output_dir must be provided when with_audio=True")
 
@@ -484,11 +462,13 @@ def compute_latents(  # noqa: PLR0913, PLR0915
         resolution_buckets=resolution_buckets,
         reshape_mode=reshape_mode,
         with_audio=with_audio,
+        temporal_subsample_factor=temporal_subsample_factor,
     )
     logger.info(f"Loaded {len(dataset)} valid media files")
 
     output_path = Path(output_dir)
     output_path.mkdir(parents=True, exist_ok=True)
+
     audio_output_path: Path | None = None
     if with_audio:
         audio_output_path = Path(audio_output_dir)
@@ -499,10 +479,10 @@ def compute_latents(  # noqa: PLR0913, PLR0915
         logger.warning("Audio processing requires batch_size=1. Overriding batch_size to 1.")
         batch_size = 1
 
-    data_root = dataset.dataset_file.parent
+    data_root = Path(dataset_file).parent
 
     def _is_done(idx: int) -> bool:
-        rel = dataset.main_media_paths[idx].relative_to(data_root).with_suffix(".pt")
+        rel = _output_relative(dataset.main_media_paths[idx], 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()
@@ -510,18 +490,16 @@ def compute_latents(  # noqa: PLR0913, PLR0915
     dataloader = _build_sharded_dataloader(
         dataset,
         batch_size=batch_size,
-        num_workers=4,
+        num_workers=num_dataloader_workers,
         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)
 
-    # Load audio VAE encoder and audio processor if needed
     audio_vae_encoder = None
     audio_processor = None
     if with_audio:
@@ -531,7 +509,6 @@ def compute_latents(  # noqa: PLR0913, PLR0915
                 device=torch_device,
                 dtype=torch.float32,  # Audio VAE needs float32 for quality. TODO: re-test with bfloat16.
             )
-            # Create audio processor for waveform-to-spectrogram conversion
             audio_processor = AudioProcessor(
                 target_sample_rate=audio_vae_encoder.sample_rate,
                 mel_bins=audio_vae_encoder.mel_bins,
@@ -562,7 +539,7 @@ def compute_latents(  # noqa: PLR0913, PLR0915
 
             # Encode video
             with torch.inference_mode():
-                video_latent_data = encode_video(vae=vae, video=video, use_tiling=vae_tiling)
+                video_latent_data = _encode_video(vae=vae, video=video, use_tiling=vae_tiling)
 
             # Save latents for each item in batch
             for i in range(len(batch["relative_path"])):
@@ -572,13 +549,15 @@ def compute_latents(  # noqa: PLR0913, PLR0915
                 # Create output directory maintaining structure
                 output_file.parent.mkdir(parents=True, exist_ok=True)
 
-                # Index into batch to get this item's latents
+                # Store the latent's effective fps (= source_fps / subsample factor).
+                # Downstream position math expects the rate the saved latents actually have.
+                effective_fps = batch["video_metadata"]["fps"][i].item() / temporal_subsample_factor
                 latent_data = {
                     "latents": video_latent_data["latents"][i].cpu().contiguous(),  # [C, F', H', W']
                     "num_frames": video_latent_data["num_frames"],
                     "height": video_latent_data["height"],
                     "width": video_latent_data["width"],
-                    "fps": batch["video_metadata"]["fps"][i].item(),
+                    "fps": effective_fps,
                 }
 
                 _atomic_save(latent_data, output_file)
@@ -596,7 +575,7 @@ def compute_latents(  # noqa: PLR0913, PLR0915
 
                         # Encode audio
                         with torch.inference_mode():
-                            audio_latents = encode_audio(audio_vae_encoder, audio_processor, audio_data)
+                            audio_latents = _encode_audio(audio_vae_encoder, audio_processor, audio_data)
 
                         # Save audio latents
                         audio_output_file = audio_output_path / output_rel_path
@@ -625,7 +604,7 @@ def compute_latents(  # noqa: PLR0913, PLR0915
         )
 
 
-def encode_video(
+def _encode_video(
     vae: torch.nn.Module,
     video: torch.Tensor,
     dtype: torch.dtype | None = None,
@@ -662,7 +641,7 @@ def encode_video(
 
     # Choose encoding method based on tiling flag
     if use_tiling:
-        latents = tiled_encode_video(
+        latents = _tiled_encode_video(
             vae=vae,
             video=video,
             tile_size=tile_size,
@@ -685,7 +664,7 @@ def encode_video(
     }
 
 
-def tiled_encode_video(  # noqa: PLR0912, PLR0915
+def _tiled_encode_video(  # noqa: PLR0912, PLR0915
     vae: torch.nn.Module,
     video: torch.Tensor,
     tile_size: int = DEFAULT_TILE_SIZE,
@@ -840,7 +819,7 @@ def tiled_encode_video(  # noqa: PLR0912, PLR0915
     return output
 
 
-def encode_audio(
+def _encode_audio(
     audio_vae_encoder: torch.nn.Module,
     audio_processor: torch.nn.Module,
     audio: Audio,
@@ -868,6 +847,35 @@ def encode_audio(
     if waveform.dim() == 2:
         waveform = waveform.unsqueeze(0)
 
+    # Convert to stereo if needed (audio VAE expects 2 channels)
+    # Channel order for surround: 5.1=[L,R,C,LFE,Ls,Rs], 7.1=[L,R,C,LFE,Ls,Rs,Lb,Rb]
+    num_channels = waveform.shape[1]
+    if num_channels == 1:
+        # Mono to stereo: duplicate the channel
+        waveform = waveform.repeat(1, 2, 1)
+    elif num_channels == 6:
+        # 5.1 downmix with normalized weights (sum to 1.0)
+        # Original: L = L + 0.707*C + 0.707*Ls, weights sum = 2.414
+        w_main = 1.0 / 2.414  # ~0.414
+        w_other = 0.707 / 2.414  # ~0.293
+        left = w_main * waveform[:, 0, :] + w_other * waveform[:, 2, :] + w_other * waveform[:, 4, :]
+        right = w_main * waveform[:, 1, :] + w_other * waveform[:, 2, :] + w_other * waveform[:, 5, :]
+        waveform = torch.stack([left, right], dim=1)
+    elif num_channels == 8:
+        # 7.1 downmix with normalized weights (sum to 1.0)
+        # Original: L = L + 0.707*C + 0.707*Ls + 0.707*Lb, weights sum = 3.121
+        w_main = 1.0 / 3.121  # ~0.320
+        w_other = 0.707 / 3.121  # ~0.227
+        center = waveform[:, 2, :]
+        left = w_main * waveform[:, 0, :] + w_other * (center + waveform[:, 4, :] + waveform[:, 6, :])
+        right = w_main * waveform[:, 1, :] + w_other * (center + waveform[:, 5, :] + waveform[:, 7, :])
+        waveform = torch.stack([left, right], dim=1)
+    elif num_channels > 2:
+        # Unknown layout: average all channels to mono, then duplicate to stereo
+        logger.warning(f"Unknown audio channel layout ({num_channels} channels), using mean downmix")
+        mono = waveform.mean(dim=1, keepdim=True)
+        waveform = mono.repeat(1, 2, 1)
+
     # Calculate duration
     duration = waveform.shape[-1] / audio.sampling_rate
 
@@ -889,6 +897,372 @@ def encode_audio(
     }
 
 
+AUDIO_FILE_EXTENSIONS = {".wav", ".mp3", ".flac", ".ogg", ".aac", ".m4a"}
+VIDEO_FILE_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".webm"}
+IMAGE_FILE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".heic", ".heif", ".bmp", ".tiff", ".webp"}
+
+
+def compute_video_masks(
+    dataset_file: str | Path,
+    mask_column: str,
+    latents_dir: str,
+    output_dir: str,
+    main_media_column: str | None = None,
+    overwrite: bool = False,
+) -> None:
+    """Preprocess video mask files to latent-space binary masks.
+    For each sample, loads the mask video/image, applies the same spatial
+    resize/crop as the target video (read from saved latent metadata), downsamples
+    to latent dimensions, binarizes, and saves as a .pt tensor.
+    Args:
+        dataset_file: Path to metadata file (CSV/JSON/JSONL).
+        mask_column: Column name containing mask video/image paths.
+        latents_dir: Directory containing the target video latents (for reading
+            spatial/temporal metadata to ensure mask alignment).
+        output_dir: Directory to save mask .pt files.
+        main_media_column: Column for output file naming (defaults to mask_column).
+    """
+    dataset_path = Path(dataset_file)
+    data_root = dataset_path.parent
+    latents_path = Path(latents_dir)
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    naming_column = main_media_column or mask_column
+    mask_paths = _load_paths_from_dataset(dataset_path, mask_column)
+    naming_paths = _load_paths_from_dataset(dataset_path, naming_column) if naming_column != mask_column else mask_paths
+
+    success = 0
+    for mask_file, naming_file in zip(mask_paths, naming_paths, strict=True):
+        rel_path = _output_relative(naming_file, data_root)
+        latent_file = latents_path / rel_path.with_suffix(".pt")
+        out_file = output_path / rel_path.with_suffix(".pt")
+
+        if not latent_file.exists():
+            logger.warning(f"No target latent found at {latent_file}, skipping mask {mask_file}")
+            continue
+
+        if not overwrite and out_file.is_file():
+            continue
+
+        target_meta = torch.load(latent_file, map_location="cpu", weights_only=True)
+        latent_f = target_meta["num_frames"]
+        latent_h = target_meta["height"]
+        latent_w = target_meta["width"]
+        pixel_h = latent_h * VAE_SPATIAL_FACTOR
+        pixel_w = latent_w * VAE_SPATIAL_FACTOR
+        pixel_f = (latent_f - 1) * VAE_TEMPORAL_FACTOR + 1
+
+        # Load mask as video or image
+        if mask_file.suffix.lower() in IMAGE_FILE_EXTENSIONS:
+            img = to_tensor(open_image_as_srgb(mask_file)).mean(dim=0, keepdim=True)  # [1, H, W]
+            img = tv_resize(img.unsqueeze(0), [pixel_h, pixel_w]).squeeze(0)  # [1, H, W]
+            mask_pixels = img.expand(pixel_f, -1, -1)  # tile across frames → [F, H, W]
+        else:
+            frames, _ = read_video(str(mask_file), max_frames=pixel_f)  # [F, C, H, W]
+            frames = frames[:pixel_f].mean(dim=1)  # grayscale → [F, H, W]
+            frames = torch.nn.functional.interpolate(
+                frames.unsqueeze(1), size=(pixel_h, pixel_w), mode="nearest"
+            ).squeeze(1)  # [F, H, W]
+            mask_pixels = frames
+
+        # Downsample to latent dims: [F, H, W] → [F', H', W']
+        mask_latent = torch.nn.functional.avg_pool2d(mask_pixels.unsqueeze(1), kernel_size=VAE_SPATIAL_FACTOR).squeeze(
+            1
+        )  # [F, H', W'] → spatial done
+        # Temporal: max-pool over groups of VAE_TEMPORAL_FACTOR frames (any masked frame masks the group)
+        f_spatial = mask_latent.shape[0]
+        pad_f = (VAE_TEMPORAL_FACTOR - f_spatial % VAE_TEMPORAL_FACTOR) % VAE_TEMPORAL_FACTOR
+        if pad_f > 0:
+            mask_latent = torch.nn.functional.pad(mask_latent, (0, 0, 0, 0, 0, pad_f))
+        h_prime, w_prime = mask_latent.shape[1], mask_latent.shape[2]
+        mask_latent = mask_latent.reshape(-1, VAE_TEMPORAL_FACTOR, h_prime, w_prime).amax(dim=1)[:latent_f]
+
+        # Binarize
+        mask_latent = (mask_latent > 0.5).float()
+
+        out_file.parent.mkdir(parents=True, exist_ok=True)
+        _atomic_save({"mask": mask_latent}, out_file)
+        success += 1
+
+    logger.info(f"Mask preprocessing complete: {success} masks saved to {output_path}")
+
+
+def compute_audio_masks(
+    dataset_file: str | Path,
+    mask_column: str,
+    audio_latents_dir: str,
+    output_dir: str,
+    main_media_column: str | None = None,
+    overwrite: bool = False,
+) -> None:
+    """Preprocess audio mask files to latent-space binary masks.
+    For each sample, loads the mask (a 1D waveform-like signal or a simple tensor),
+    resamples it to match the target audio latent temporal length, binarizes, and saves.
+    Args:
+        dataset_file: Path to metadata file (CSV/JSON/JSONL).
+        mask_column: Column name containing mask file paths (.wav or .pt).
+        audio_latents_dir: Directory containing the target audio latents (for reading
+            temporal metadata to ensure mask alignment).
+        output_dir: Directory to save mask .pt files.
+        main_media_column: Column for output file naming (defaults to mask_column).
+    """
+    dataset_path = Path(dataset_file)
+    data_root = dataset_path.parent
+    audio_latents_path = Path(audio_latents_dir)
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    naming_column = main_media_column or mask_column
+    mask_paths = _load_paths_from_dataset(dataset_path, mask_column)
+    naming_paths = _load_paths_from_dataset(dataset_path, naming_column) if naming_column != mask_column else mask_paths
+
+    success = 0
+    for mask_file, naming_file in zip(mask_paths, naming_paths, strict=True):
+        rel_path = _output_relative(naming_file, data_root)
+        latent_file = audio_latents_path / rel_path.with_suffix(".pt")
+        out_file = output_path / rel_path.with_suffix(".pt")
+
+        if not latent_file.exists():
+            logger.warning(f"No target audio latent found at {latent_file}, skipping mask {mask_file}")
+            continue
+
+        if not overwrite and out_file.is_file():
+            continue
+
+        target_meta = torch.load(latent_file, map_location="cpu", weights_only=True)
+        latent_t = target_meta["num_time_steps"]
+
+        # Load mask: .pt file (raw tensor) or .wav (use amplitude envelope)
+        if mask_file.suffix == ".pt":
+            raw_mask = torch.load(mask_file, map_location="cpu", weights_only=True)
+            if isinstance(raw_mask, dict):
+                raw_mask = raw_mask.get("mask", next(iter(raw_mask.values())))
+            raw_mask = raw_mask.float().flatten()
+        else:
+            audio = _load_audio_from_file(mask_file)
+            if audio is None:
+                logger.warning(f"Could not load audio mask from {mask_file}")
+                continue
+            raw_mask = audio.waveform.abs().mean(dim=0)  # mono amplitude envelope
+
+        # Resample to target audio latent length
+        mask_resampled = torch.nn.functional.interpolate(
+            raw_mask.unsqueeze(0).unsqueeze(0), size=latent_t, mode="nearest"
+        ).squeeze()  # [latent_t]
+
+        mask_binary = (mask_resampled > 0.5).float()
+
+        out_file.parent.mkdir(parents=True, exist_ok=True)
+        _atomic_save({"mask": mask_binary}, out_file)
+        success += 1
+
+    logger.info(f"Audio mask preprocessing complete: {success} masks saved to {output_path}")
+
+
+def compute_audio_latents(  # noqa: PLR0915
+    dataset_file: str | Path,
+    audio_column: str,
+    output_dir: str,
+    model_path: str,
+    main_media_column: str | None = None,
+    max_duration: float | None = None,
+    duration_buckets: list[float] | None = None,
+    device: str = "cuda",
+    overwrite: bool = False,
+) -> None:
+    """Encode audio files into latent representations.
+    Supports standalone audio files (.wav, .mp3, etc.) and audio tracks
+    extracted from video files (.mp4, etc.).
+    Args:
+        dataset_file: Path to metadata file (CSV/JSON/JSONL).
+        audio_column: Column name containing audio file paths.
+        output_dir: Directory to save audio latents.
+        model_path: Path to LTX-2 checkpoint (.safetensors).
+        main_media_column: Column for output file naming (defaults to audio_column).
+            Ensures alignment with other latent directories.
+        max_duration: Maximum audio duration in seconds. Audio is trimmed if longer.
+            Mutually exclusive with duration_buckets.
+        duration_buckets: List of allowed durations in seconds (e.g. [2.0, 4.0, 8.0]).
+            Each audio file is matched to the largest bucket that fits its duration,
+            then trimmed to exactly that length. Files shorter than the smallest
+            bucket are skipped. Ensures uniform lengths for batched training.
+        device: Device to use for computation.
+    """
+    console = Console()
+    torch_device = torch.device(device)
+
+    dataset_path = Path(dataset_file)
+    data_root = dataset_path.parent
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    naming_column = main_media_column or audio_column
+    audio_paths = _load_paths_from_dataset(dataset_path, audio_column)
+    naming_paths = (
+        _load_paths_from_dataset(dataset_path, naming_column) if naming_column != audio_column else audio_paths
+    )
+
+    with console.status(f"[bold]Loading audio VAE encoder from [cyan]{model_path}[/]...", spinner="dots"):
+        audio_vae_encoder = load_audio_vae_encoder(
+            checkpoint_path=model_path,
+            device=torch_device,
+            dtype=torch.float32,
+        )
+        audio_processor = AudioProcessor(
+            target_sample_rate=audio_vae_encoder.sample_rate,
+            mel_bins=audio_vae_encoder.mel_bins,
+            mel_hop_length=audio_vae_encoder.mel_hop_length,
+            n_fft=audio_vae_encoder.n_fft,
+        ).to(torch_device)
+
+    sorted_buckets = sorted(duration_buckets, reverse=True) if duration_buckets else None
+    success_count = 0
+    skip_count = 0
+
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        TaskProgressColumn(),
+        MofNCompleteColumn(),
+        TimeElapsedColumn(),
+        TimeRemainingColumn(),
+        console=console,
+    ) as progress:
+        task = progress.add_task("Encoding audio", total=len(audio_paths))
+
+        for audio_path, naming_path in zip(audio_paths, naming_paths, strict=True):
+            rel_path = _output_relative(naming_path, data_root)
+            output_file = output_path / rel_path.with_suffix(".pt")
+            output_file.parent.mkdir(parents=True, exist_ok=True)
+
+            if not overwrite and output_file.is_file():
+                success_count += 1
+                progress.advance(task)
+                continue
+
+            # Load audio (no trimming yet — need full duration for bucket matching)
+            audio = _load_audio_from_file(audio_path)
+            if audio is None:
+                skip_count += 1
+                progress.advance(task)
+                continue
+
+            file_duration = audio.waveform.shape[-1] / audio.sampling_rate
+
+            # Determine target duration: bucket matching, max_duration cap, or full file
+            target_duration = file_duration
+            if sorted_buckets:
+                bucket = next((b for b in sorted_buckets if b <= file_duration), None)
+                if bucket is None:
+                    logger.warning(
+                        f"Skipping {audio_path.name} ({file_duration:.1f}s) — shorter than "
+                        f"smallest bucket ({sorted_buckets[-1]:.1f}s)"
+                    )
+                    skip_count += 1
+                    progress.advance(task)
+                    continue
+                target_duration = bucket
+            elif max_duration is not None:
+                target_duration = min(file_duration, max_duration)
+
+            # Trim to target duration
+            target_samples = int(target_duration * audio.sampling_rate)
+            trimmed_waveform = audio.waveform[:, :target_samples]
+            audio = Audio(waveform=trimmed_waveform, sampling_rate=audio.sampling_rate)
+
+            with torch.inference_mode():
+                audio_latents = _encode_audio(audio_vae_encoder, audio_processor, audio)
+
+            _atomic_save(
+                {
+                    "latents": audio_latents["latents"].cpu().contiguous(),
+                    "num_time_steps": audio_latents["num_time_steps"],
+                    "frequency_bins": audio_latents["frequency_bins"],
+                    "duration": audio_latents["duration"],
+                },
+                output_file,
+            )
+            success_count += 1
+            progress.advance(task)
+
+    logger.info(f"Audio encoding complete: {success_count} encoded, {skip_count} skipped. Saved to {output_path}")
+
+
+def _output_relative(path: Path, data_root: Path) -> Path:
+    """Relative path used to name a sample's cached output, mirroring the input layout.
+    Normally media lives under the dataset directory and this is just the path relative to it.
+    If a media path is absolute or otherwise outside the dataset directory (e.g. a one-off
+    metadata file that references media elsewhere), mirror its absolute structure under the
+    output directory instead of raising, so out-of-tree media stays collision-free.
+    """
+    try:
+        return path.relative_to(data_root)
+    except ValueError:
+        return Path(*path.parts[1:]) if path.is_absolute() else path
+
+
+def _load_paths_from_dataset(dataset_file: Path, column: str) -> list[Path]:
+    """Load file paths from a dataset column, resolving relative to the dataset file's directory."""
+    data_root = dataset_file.parent
+
+    if dataset_file.suffix == ".csv":
+        df = pd.read_csv(dataset_file)
+        if column not in df.columns:
+            raise ValueError(f"Column '{column}' not found in CSV file")
+        return [data_root / Path(str(v).strip()) for v in df[column].tolist()]
+
+    if dataset_file.suffix == ".json":
+        with open(dataset_file, encoding="utf-8") as f:
+            data = json.load(f)
+        if not isinstance(data, list):
+            raise ValueError("JSON file must contain a list of objects")
+        return [data_root / Path(entry[column].strip()) for entry in data]
+
+    if dataset_file.suffix == ".jsonl":
+        paths = []
+        with open(dataset_file, encoding="utf-8") as f:
+            for line in f:
+                entry = json.loads(line)
+                paths.append(data_root / Path(entry[column].strip()))
+        return paths
+
+    raise ValueError(f"Unsupported dataset format: {dataset_file.suffix}")
+
+
+def _load_audio_from_file(audio_path: Path, max_duration: float | None = None) -> Audio | None:
+    """Load audio from an audio or video file, optionally trimming to max_duration."""
+    try:
+        waveform, sample_rate = torchaudio.load(str(audio_path))
+    except Exception:
+        logger.debug(f"Could not load audio from {audio_path}")
+        return None
+
+    if max_duration is not None:
+        max_samples = int(max_duration * sample_rate)
+        if waveform.shape[-1] > max_samples:
+            waveform = waveform[:, :max_samples]
+
+    return Audio(waveform=waveform, sampling_rate=sample_rate)
+
+
+def detect_dataset_columns(dataset_file: str | Path) -> set[str]:
+    """Read column names from a dataset file without loading all data."""
+    path = Path(dataset_file)
+    if path.suffix == ".csv":
+        df = pd.read_csv(path, nrows=0)
+        return set(df.columns)
+    if path.suffix == ".json":
+        with open(path, encoding="utf-8") as f:
+            data = json.load(f)
+        return set(data[0].keys()) if isinstance(data, list) and data else set()
+    if path.suffix == ".jsonl":
+        with open(path, encoding="utf-8") as f:
+            return set(json.loads(f.readline()).keys())
+    return set()
+
+
 def parse_resolution_buckets(resolution_buckets_str: str) -> list[tuple[int, int, int]]:
     """Parse resolution buckets from string format to list of tuples (frames, height, width)"""
     resolution_buckets = []
@@ -1041,11 +1415,11 @@ def main(  # noqa: PLR0913
     ),
 ) -> 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.
+    For multi-GPU preprocessing, invoke under ``accelerate launch`` -- each process
+    will handle an interleaved shard of the dataset.
     Examples:
         # Process videos from a CSV file
         python scripts/process_videos.py dataset.csv --resolution-buckets 768x768x25 \\

packages/ltx-trainer/scripts/serve_captioner.py

@@ -0,0 +1,203 @@
+#!/usr/bin/env python3
+
+"""Launch a vLLM server for Qwen3-Omni captioning.
+Runs the actual server via ``uvx`` so that vLLM and its CUDA-tied
+dependencies live in their own isolated environment (no impact on this
+package's dependency tree).
+The captioning script (``caption_videos.py``) talks to the server over its
+OpenAI-compatible HTTP API. Once the server is up it stays loaded across
+captioning runs; no per-script model warmup cost.
+Typical usage::
+    # Default: dynamic FP8 quantization, listen on 127.0.0.1:8001
+    uv run python scripts/serve_captioner.py
+    # Just print the chosen `uvx vllm serve ...` command without running it
+    uv run python scripts/serve_captioner.py --print-cmd
+    # Use full bf16 on a GPU with >= 66 GiB free VRAM (slightly more reliable
+    # numerics but 2x the weight memory)
+    uv run python scripts/serve_captioner.py --quantization bf16
+    # Use a different port or expose on all interfaces
+    uv run python scripts/serve_captioner.py --port 9000 --host 0.0.0.0
+"""
+
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+import typer
+from rich.console import Console
+
+console = Console()
+
+# Model identifier we serve. The captioner client must use the same string.
+DEFAULT_MODEL = "Qwen/Qwen3-Omni-30B-A3B-Thinking"
+
+# Pinned vLLM version known to support Qwen3-Omni on CUDA 12.x.
+# vLLM 0.20+ requires CUDA 13. Update both as the environment evolves.
+# The ``[audio]`` extra is required for Qwen3-Omni to decode audio at all.
+DEFAULT_VLLM_SPEC = "vllm[audio]==0.11.2"
+
+# Approximate disk needed for the model download (HF cache structure).
+MODEL_DISK_GIB = 65.0
+
+
+app = typer.Typer(
+    pretty_exceptions_enable=False,
+    no_args_is_help=False,
+    help="Launch a local vLLM server for Qwen3-Omni captioning.",
+)
+
+
+def _query_disk_free_gib(path: Path) -> float:
+    return shutil.disk_usage(str(path)).free / 1024**3
+
+
+def _build_vllm_args(
+    *,
+    model: str,
+    host: str,
+    port: int,
+    quantization: str,
+    max_model_len: int,
+    gpu_memory_utilization: float,
+    extra_args: list[str],
+) -> list[str]:
+    """Construct the `vllm serve ...` argv."""
+    args = [
+        "vllm",
+        "serve",
+        model,
+        "--host",
+        host,
+        "--port",
+        str(port),
+        "--dtype",
+        "bfloat16",
+        "--max-model-len",
+        str(max_model_len),
+        "--gpu-memory-utilization",
+        str(gpu_memory_utilization),
+        # Let the server accept ``file://`` URLs pointing at local videos.
+        "--allowed-local-media-path",
+        "/",
+        # The model is a multimodal MoE; cap each input to one of each
+        # modality to match what our captioner sends.
+        "--limit-mm-per-prompt",
+        '{"image": 1, "video": 1, "audio": 1}',
+        # Small concurrent-sequence cap so KV cache headroom isn't fragmented.
+        "--max-num-seqs",
+        "4",
+    ]
+    if quantization == "fp8":
+        args += ["--quantization", "fp8"]
+    args += extra_args
+    return args
+
+
+@app.command()
+def main(
+    model: str = typer.Option(DEFAULT_MODEL, "--model", help="Model identifier to serve."),
+    host: str = typer.Option("127.0.0.1", "--host", help="Listen address. Use 0.0.0.0 for remote access."),
+    port: int = typer.Option(8001, "--port", help="HTTP port."),
+    quantization: str = typer.Option(
+        "fp8",
+        "--quantization",
+        "-q",
+        help=(
+            "Weight precision. 'fp8' (default, dynamic FP8 -- ~31 GiB weights) is "
+            "the recommended choice; it fits on 40 GiB GPUs and runs at the same "
+            "speed as bf16 on H100. 'bf16' uses ~60 GiB of weights -- pick it if "
+            "you have abundant VRAM and want minimal numerical drift."
+        ),
+    ),
+    max_model_len: int = typer.Option(
+        32768,
+        "--max-model-len",
+        help="Maximum context length the server accepts (must fit input video tokens + max_tokens).",
+    ),
+    gpu_memory_utilization: float = typer.Option(
+        0.9,
+        "--gpu-memory-utilization",
+        help="Fraction of GPU memory vLLM may reserve (model + KV cache).",
+    ),
+    hf_home: Path | None = typer.Option(  # noqa: B008
+        None,
+        "--hf-home",
+        help=(
+            "Override HF_HOME (where the model is downloaded). The model is ~65 GB; "
+            "by default this follows your environment's HF_HOME or HuggingFace's default."
+        ),
+    ),
+    vllm_spec: str = typer.Option(
+        DEFAULT_VLLM_SPEC,
+        "--vllm-spec",
+        help="pip-style spec passed to `uvx --from`. Pin a version that matches your CUDA.",
+    ),
+    print_cmd: bool = typer.Option(
+        False,
+        "--print-cmd",
+        help="Print the chosen command without running it.",
+    ),
+    extra_args: list[str] | None = typer.Argument(  # noqa: B008
+        None,
+        help="Additional args passed through to `vllm serve` after `--`.",
+    ),
+) -> None:
+    """Launch the vLLM server for Qwen3-Omni."""
+    extra = extra_args or []
+
+    if quantization not in ("bf16", "fp8"):
+        console.print(f"[red]--quantization must be 'bf16' or 'fp8'; got {quantization!r}.[/]")
+        raise typer.Exit(code=1)
+
+    # Disk check (only meaningful before first download).
+    cache_root = hf_home or Path(os.environ.get("HF_HOME", str(Path.home() / ".cache" / "huggingface")))
+    cache_root.mkdir(parents=True, exist_ok=True)
+    free_disk = _query_disk_free_gib(cache_root)
+    if free_disk < MODEL_DISK_GIB:
+        console.print(
+            f"[yellow]\u26a0 Only {free_disk:.1f} GiB free on disk under {cache_root} but the "
+            f"model needs ~{MODEL_DISK_GIB:.0f} GiB. Either free up space, set --hf-home "
+            f"to a larger volume, or expect the download to fail mid-way.[/]"
+        )
+
+    vllm_args = _build_vllm_args(
+        model=model,
+        host=host,
+        port=port,
+        quantization=quantization,
+        max_model_len=max_model_len,
+        gpu_memory_utilization=gpu_memory_utilization,
+        extra_args=extra,
+    )
+
+    # Use ``uvx --from vllm==...`` so vLLM lives in its own throwaway venv
+    # (or a cached tool venv). The `--` separates uvx args from the command's.
+    uvx_cmd = ["uvx", "--from", vllm_spec, *vllm_args]
+
+    env = os.environ.copy()
+    # vLLM 0.11.x requires the V0 engine for Qwen3-Omni's multimodal pipeline.
+    env.setdefault("VLLM_USE_V1", "0")
+    if hf_home is not None:
+        env["HF_HOME"] = str(hf_home)
+
+    console.print("\n[bold]Command:[/]")
+    console.print("  " + " ".join(uvx_cmd))
+    if hf_home is not None:
+        console.print(f"  [dim](with HF_HOME={hf_home})[/]")
+
+    if print_cmd:
+        return
+
+    console.print("\n[dim]Launching... (first run downloads the model -- ~5 min on a fast link)[/]\n")
+    try:
+        completed = subprocess.run(uvx_cmd, env=env, check=False)
+    except KeyboardInterrupt:
+        console.print("\n[yellow]Interrupted.[/]")
+        return
+    sys.exit(completed.returncode)
+
+
+if __name__ == "__main__":
+    app()

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

@@ -1,59 +1,117 @@
 """
 Audio-visual media captioning using multimodal models.
 This module provides captioning capabilities for videos with audio using:
-- Qwen2.5-Omni: Local model supporting text, audio, image, and video inputs (default)
-- Gemini Flash: Cloud-based API for audio-visual captioning
-Requirements:
-- Qwen2.5-Omni: transformers>=4.50, torch
-- Gemini Flash: google-generativeai (uv pip install google-generativeai)
-  Set GEMINI_API_KEY or GOOGLE_API_KEY environment variable
+- Qwen3-Omni via a local vLLM server (default)
+- Gemini Flash 3.5 (cloud API)
+Both produce a single combined English caption per video as a single
+continuous paragraph of prose.
+The Qwen3-Omni backend runs in a separately-launched vLLM server rather than
+in-process, so vLLM's heavy CUDA dependencies stay out of this package. The
+captioner talks to it over the OpenAI-compatible HTTP API.
+Launch the server once (in an isolated environment) with:
+.. code-block:: bash
+    uv run python scripts/serve_captioner.py
+That helper picks BF16 vs FP8 dynamic quantization based on the GPU's free
+memory and forwards everything else to ``vllm serve``. To check the recommended
+command without running it, pass ``--print-cmd``.
+To use Gemini instead, install ``google-genai`` and either set ``GEMINI_API_KEY``
+(Gemini Developer API) or have Google Cloud credentials available (gcloud / an
+attached service account), in which case it uses Vertex AI automatically.
 """
 
-import itertools
+import json
+import os
 import re
+import subprocess
+import tempfile
 from abc import ABC, abstractmethod
 from enum import Enum
 from pathlib import Path
+from typing import ClassVar
 
-import torch
+DEFAULT_VIDEO_CAPTION_INSTRUCTION = """\
+Analyze this video and produce a single detailed caption covering both its visual content and its audio. Be \
+detailed enough that someone reading the caption could form an accurate mental picture of what happens on screen \
+and what can be heard. Be exhaustive: include every meaningful detail you can see and hear, including small \
+objects, textures, secondary movements, and minor background sounds.
 
-# Instruction for audio-visual captioning (default) - includes speech transcription and sounds
-DEFAULT_CAPTION_INSTRUCTION = """\
-Analyze this media and provide a detailed caption in the following EXACT format. Fill in ALL sections:
+Begin the caption directly with the action or visual detail; do not preface it with phrases like \
+"The video opens with...", "The scene shows...", "We see...", or "There is...".
 
-[VISUAL]: <Detailed description of people, objects, actions, settings, colors, and movements>
-[SPEECH]: <Word-for-word transcription of everything spoken.
-           Listen carefully and transcribe the exact words. If no speech, write "None">
-[SOUNDS]: <Description of music, ambient sounds, sound effects. If none, write "None">
-[TEXT]: <Any on-screen text visible. If none, write "None">
+For every shot, include:
+- The shot type and framing (extreme wide / wide / medium / medium close-up / close-up / extreme close-up) and any \
+camera motion.
+- Characters' clothing, appearance, posture, and movement (direction, speed, quality).
+- The environment's materials, textures, lighting, and colors.
+- All audio: spoken dialogue (quoted exactly in the original language), tone of voice, music (style, mood, \
+volume changes), and environmental sounds. If a category is absent -- for example no music is playing, or no one is \
+speaking -- state that explicitly. Do not invent specific instruments, music genres, moods, or ambient sounds \
+that are not actually present.
+- Any on-screen text (signs, titles, labels).
 
-You MUST fill in all four sections. For [SPEECH], transcribe the actual words spoken, not a summary."""
+Describe only what is visible or audible. Do not infer emotions, intentions, or anything outside the segment. \
+Refer to people descriptively (e.g., "the man in the blue jacket"). Narrate strictly in chronological order; if \
+the video contains multiple shots, describe each one in turn.
 
-# Instruction for video-only captioning (no audio processing)
-VIDEO_ONLY_CAPTION_INSTRUCTION = """\
-Analyze this media and provide a detailed caption in the following EXACT format. Fill in ALL sections:
+Write everything as a single continuous paragraph of prose. Do not use section headers, bullet points, or labels \
+like "Audio:" / "Visual:" / "Shot:". Integrate visual and audio details naturally within the same sentences.
 
-[VISUAL]: <Detailed description of people, objects, actions, settings, colors, and movements>
-[TEXT]: <Any on-screen text visible. If none, write "None">
+Return a JSON object with exactly one key:
 
-You MUST fill in both sections."""
+{"combined_caption_english": "<your caption here>"}"""
+
+
+DEFAULT_IMAGE_CAPTION_INSTRUCTION = """\
+Analyze this image and produce a single detailed caption of its visual content. Be detailed enough that \
+someone reading the caption could form an accurate mental picture of the image. Be thorough: include every meaningful \
+detail that is actually present, including small objects, textures, and background elements.
+
+Begin the caption directly with the main subject or a visual detail; do not preface it with phrases like \
+"The image shows...", "This is a photo of...", "We see...", or "There is...".
+
+Include:
+- The framing and composition (close-up / medium / wide / overhead, etc.) and the vantage point.
+- The medium or style if distinctive (photograph, illustration, 3D render, painting).
+- People's clothing, appearance, and posture, and what they are doing.
+- The setting's materials, textures, lighting, and colors.
+- Transcribe any visible text verbatim (signs, labels, titles, captions).
+
+Describe only what is visible. Do not infer emotions or intentions, and do not describe sounds, motion, or \
+events before or after the moment shown -- this is a single still image. When something is ambiguous, describe \
+the visible cue (e.g., "warm low-angle light") rather than guessing the underlying fact (e.g., "sunrise"). \
+Refer to people descriptively (e.g., "the man in the blue jacket").
+
+Only describe what is present. Never state that something is absent or missing -- do not write phrases like \
+"there is no text", "no people are present", or "no other objects". If a category such as people or text does \
+not appear, simply leave it out.
+
+Write everything as a single continuous paragraph of prose. Do not use section headers, bullet points, or \
+labels.
+
+Return a JSON object with exactly one key:
+
+{"combined_caption_english": "<your caption here>"}"""
+
+
+# Default model served by ``scripts/serve_captioner.py``. The captioner does not
+# download or load this model itself -- it just sends requests to the vLLM
+# server, which already has the model loaded.
+DEFAULT_QWEN_MODEL = "Qwen/Qwen3-Omni-30B-A3B-Thinking"
+DEFAULT_VLLM_BASE_URL = "http://127.0.0.1:8001/v1"
+
+# Key the combined-caption prompt asks the model to return its caption under.
+_CAPTION_JSON_KEY = "combined_caption_english"
 
 
 class CaptionerType(str, Enum):
     """Enum for different types of media captioners."""
 
-    QWEN_OMNI = "qwen_omni"  # Local Qwen2.5-Omni model (audio + video)
-    GEMINI_FLASH = "gemini_flash"  # Gemini Flash API (audio + video)
+    QWEN_OMNI = "qwen_omni"  # Qwen3-Omni via local vLLM HTTP server
+    GEMINI_FLASH = "gemini_flash"  # Gemini Flash 3.5 cloud API
 
 
 def create_captioner(captioner_type: CaptionerType, **kwargs) -> "MediaCaptioningModel":
-    """Factory function to create a media captioner.
-    Args:
-        captioner_type: The type of captioner to create
-        **kwargs: Additional arguments to pass to the captioner constructor
-    Returns:
-        An instance of a MediaCaptioningModel
-    """
+    """Factory function to create a media captioner."""
     match captioner_type:
         case CaptionerType.QWEN_OMNI:
             return QwenOmniCaptioner(**kwargs)
@@ -66,336 +124,332 @@ def create_captioner(captioner_type: CaptionerType, **kwargs) -> "MediaCaptionin
 class MediaCaptioningModel(ABC):
     """Abstract base class for audio-visual media captioning models."""
 
+    instruction: str | None = None
+
     @abstractmethod
     def caption(self, path: str | Path, **kwargs) -> str:
-        """Generate a caption for the given video or image.
-        Args:
-            path: Path to the video/image file to caption
-        Returns:
-            A string containing the generated caption
-        """
+        """Generate a caption for the given video or image."""
 
-    @property
-    @abstractmethod
-    def supports_audio(self) -> bool:
-        """Whether this captioner supports audio input."""
+    def _resolve_instruction(self, path: str | Path) -> str:
+        """Return the custom instruction, or the image/video default for this input."""
+        if self.instruction is not None:
+            return self.instruction
+        return DEFAULT_IMAGE_CAPTION_INSTRUCTION if self._is_image_file(path) else DEFAULT_VIDEO_CAPTION_INSTRUCTION
 
     @staticmethod
     def _is_image_file(path: str | Path) -> bool:
-        """Check if the file is an image based on extension."""
         return str(path).lower().endswith((".png", ".jpg", ".jpeg", ".heic", ".heif", ".webp"))
 
     @staticmethod
     def _is_video_file(path: str | Path) -> bool:
-        """Check if the file is a video based on extension."""
         return str(path).lower().endswith((".mp4", ".avi", ".mov", ".mkv", ".webm"))
 
-    @staticmethod
-    def _clean_raw_caption(caption: str) -> str:
-        """Clean up the raw caption by removing common VLM patterns."""
-        start = ["The", "This"]
-        kind = ["video", "image", "scene", "animated sequence", "clip", "footage"]
-        act = ["displays", "shows", "features", "depicts", "presents", "showcases", "captures", "contains"]
-
-        for x, y, z in itertools.product(start, kind, act):
-            caption = caption.replace(f"{x} {y} {z} ", "", 1)
-
-        return caption
-
 
 class QwenOmniCaptioner(MediaCaptioningModel):
-    """Audio-visual captioning using Alibaba's Qwen2.5-Omni model.
-    Qwen2.5-Omni is an end-to-end multimodal model that can perceive text, images, audio, and video.
-    It uses a Thinker-Talker architecture where the Thinker generates text and the Talker can
-    generate speech. For captioning, we use only the Thinker component for text generation.
-    Key features:
-    - Block-wise processing for streaming multimodal inputs
-    - TMRoPE (Time-aligned Multimodal RoPE) for synchronizing video and audio timestamps
-    - Can extract and process audio directly from video files
-    See: https://huggingface.co/docs/transformers/en/model_doc/qwen2_5_omni
-    Model: Qwen/Qwen2.5-Omni-7B (7B parameters)
+    """Audio-visual captioning via a local vLLM server running Qwen3-Omni.
+    The vLLM server must already be running. See ``scripts/serve_captioner.py``
+    for a helper that launches one in an isolated environment (no impact on
+    this package's dependency tree).
+    The captioner uses the OpenAI-compatible chat completions API. It sends
+    a ``file://`` URL pointing at the local video, the default combined-caption
+    prompt, and parses the JSON-wrapped response.
     """
 
-    MODEL_ID = "Qwen/Qwen2.5-Omni-7B"
-
-    # Default system prompt required by Qwen2.5-Omni for proper audio processing
-    DEFAULT_SYSTEM_PROMPT = (
-        "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, "
-        "capable of perceiving auditory and visual inputs, as well as generating text and speech."
-    )
-
     def __init__(
         self,
-        device: str | torch.device | None = None,
-        use_8bit: bool = False,
+        base_url: str = DEFAULT_VLLM_BASE_URL,
+        model: str = DEFAULT_QWEN_MODEL,
+        api_key: str = "EMPTY",
         instruction: str | None = None,
+        max_tokens: int = 4096,
+        enable_thinking: bool = False,
+        timeout_s: float = 600.0,
     ):
-        """
-        Initialize the Qwen2.5-Omni captioner.
+        """Initialize the Qwen3-Omni captioner.
         Args:
-            device: Device to use for inference (e.g., 'cuda', 'cuda:0', 'cpu')
-            use_8bit: Whether to use 8-bit quantization for reduced memory usage
-            instruction: Custom instruction prompt. If None, uses the default instruction
+            base_url: Base URL of the vLLM OpenAI-compatible server (default
+                ``http://127.0.0.1:8001/v1``).
+            model: Model identifier the server is serving. Must match the
+                server's ``--served-model-name`` (defaults to the HuggingFace
+                model ID).
+            api_key: Token sent in the ``Authorization`` header. vLLM accepts
+                any value by default.
+            instruction: Custom instruction prompt. If ``None``, uses the
+                default combined-caption prompt.
+            max_tokens: Maximum new tokens to generate per caption. 4096 leaves
+                comfortable headroom for both ``enable_thinking`` modes.
+            enable_thinking: Whether to let the Thinking model produce a
+                ``<think>...</think>`` chain-of-thought before the caption.
+                Off by default: it makes captioning ~5x slower with little
+                quality benefit and occasionally introduces hallucinations
+                (e.g., inventing dialogue or background music).
+            timeout_s: Per-request HTTP timeout.
         """
-        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
-        self.instruction = instruction
-        self._load_model(use_8bit=use_8bit)
+        from openai import OpenAI  # noqa: PLC0415
 
-    @property
-    def supports_audio(self) -> bool:
-        return True
+        self.model = model
+        self.instruction = instruction
+        self.max_tokens = max_tokens
+        self.enable_thinking = enable_thinking
+        self._client = OpenAI(base_url=base_url, api_key=api_key, timeout=timeout_s)
 
     def caption(
         self,
         path: str | Path,
-        fps: int = 1,
-        include_audio: bool = True,
-        clean_caption: bool = True,
+        fps: int = 2,
     ) -> str:
         """Generate a caption for the given video or image.
         Args:
-            path: Path to the video/image file to caption
-            fps: Frames per second to sample from videos
-            include_audio: Whether to include audio in the captioning (for videos)
-            clean_caption: Whether to clean up the raw caption by removing common VLM patterns
+            path: Path to the video/image file to caption.
+            fps: Frames per second to sample from the video. Passed through to
+                vLLM's multimodal processor (``mm_processor_kwargs.fps``).
+                Default 2 is a typical choice for video MLLMs at this resolution.
+                Ignored for image inputs.
         Returns:
-            A string containing the generated caption
+            The extracted caption string.
         """
         path = Path(path)
         is_image = self._is_image_file(path)
         is_video = self._is_video_file(path)
+        if not (is_image or is_video):
+            raise ValueError(f"Unsupported media file: {path}")
 
-        # Determine if we should process audio
-        use_audio = include_audio and is_video
-
-        # Use custom instruction if provided, otherwise pick appropriate default
-        if self.instruction is not None:
-            instruction = self.instruction
-        else:
-            instruction = DEFAULT_CAPTION_INSTRUCTION if use_audio else VIDEO_ONLY_CAPTION_INSTRUCTION
-
-        # Build the user content based on media type
-        # Based on HuggingFace docs: https://huggingface.co/docs/transformers/en/model_doc/qwen2_5_omni
-        user_content = []
+        instruction = self._resolve_instruction(path)
 
         if is_image:
-            user_content.append({"type": "image", "image": str(path)})
-        elif is_video:
-            user_content.append({"type": "video", "video": str(path)})
+            content = [
+                {"type": "image_url", "image_url": {"url": f"file://{path.resolve()}"}},
+                {"type": "text", "text": instruction},
+            ]
+            return _parse_caption_response(self._chat(content)).strip()
 
-        # Add the instruction text
-        user_content.append({"type": "text", "text": instruction})
+        return self._caption_video(path, instruction, fps)
 
-        # Build conversation - use the default system prompt required by Qwen2.5-Omni
-        # Using a custom system prompt causes warnings and may affect audio processing
-        messages = [
-            {
-                "role": "system",
-                "content": [{"type": "text", "text": self.DEFAULT_SYSTEM_PROMPT}],
-            },
-            {"role": "user", "content": user_content},
-        ]
-
-        # Process inputs using the processor's apply_chat_template
-        # For videos with audio, use load_audio_from_video=True and use_audio_in_video=True
-        inputs = self.processor.apply_chat_template(
-            messages,
-            load_audio_from_video=use_audio,
-            add_generation_prompt=True,
-            tokenize=True,
-            return_dict=True,
-            return_tensors="pt",
-            fps=fps,
-            padding=True,
-            use_audio_in_video=use_audio,
-        ).to(self.model.device)
-
-        # Generate caption (text only, using Thinker-only model)
-        # Note: For Qwen2_5OmniThinkerForConditionalGeneration, use standard generate params
-        # (not thinker_ prefixed ones, those are for the full Qwen2_5OmniForConditionalGeneration)
-        input_len = inputs["input_ids"].shape[1]
-
-        output_tokens = self.model.generate(
-            **inputs,
-            use_audio_in_video=use_audio,
-            do_sample=False,
-            max_new_tokens=1024,
+    def _chat(self, content: list[dict], mm_kwargs: dict | None = None) -> str:
+        """Send one chat-completions request and return the raw response text."""
+        extra_body: dict = {
+            "repetition_penalty": 1.05,
+            "chat_template_kwargs": {"enable_thinking": self.enable_thinking},
+        }
+        if mm_kwargs:
+            extra_body["mm_processor_kwargs"] = mm_kwargs
+        response = self._client.chat.completions.create(
+            model=self.model,
+            messages=[{"role": "user", "content": content}],
+            max_tokens=self.max_tokens,
+            temperature=0.0,
+            extra_body=extra_body,
         )
+        return response.choices[0].message.content or ""
 
-        # Extract only the generated tokens (exclude the input/prompt tokens)
-        generated_tokens = output_tokens[:, input_len:]
-
-        # Decode only the generated response
-        caption_raw = self.processor.batch_decode(
-            generated_tokens,
-            skip_special_tokens=True,
-            clean_up_tokenization_spaces=False,
-        )[0]
-
-        # Remove hallucinated conversation turns (e.g., "Human\nHuman\n..." or "Human: ...")
-        # This is a known issue with chat models continuing to generate fake turns
-        # We look for patterns that are clearly hallucinated chat turns, not legitimate uses of "human"
-
-        # Match "\nHuman" followed by ":", "\n", or end of string (chat turn patterns)
-        # This won't match "A human walks..." or "...the human body..."
-        caption_raw = re.split(r"\nHuman(?::|(?:\s*\n)|$)", caption_raw, maxsplit=1)[0]
-        caption_raw = caption_raw.strip()
-
-        # Clean up caption if requested
-        return self._clean_raw_caption(caption_raw) if clean_caption else caption_raw
-
-    def _load_model(self, use_8bit: bool) -> None:
-        """Load the Qwen2.5-Omni model and processor.
-        Uses the Thinker-only model (Qwen2_5OmniThinkerForConditionalGeneration) for text generation
-        to save compute by not loading the audio generation components.
+    def _caption_video(self, path: Path, instruction: str, fps: int) -> str:
+        """Caption a video, sending its audio track as a separate modality.
+        vLLM does not extract a video's audio on its own (and its
+        ``use_audio_in_video`` path is broken server-side), so we pull the audio
+        into a 16 kHz mono WAV and send it alongside the video -- otherwise the
+        model only sees frames and fabricates any spoken content.
         """
-        from transformers import (  # noqa: PLC0415
-            BitsAndBytesConfig,
-            Qwen2_5OmniProcessor,
-            Qwen2_5OmniThinkerForConditionalGeneration,
-        )
+        with tempfile.TemporaryDirectory(prefix="qwencap_") as tmp:
+            work = Path(tmp)
 
-        quantization_config = BitsAndBytesConfig(load_in_8bit=True) if use_8bit else None
+            # Best-effort: ffmpeg fails (and we send video only) if there's no audio.
+            audio_url: str | None = None
+            try:
+                wav = work / "audio.wav"
+                _extract_audio_wav(path, wav)
+                audio_url = f"file://{wav.resolve()}"
+            except subprocess.CalledProcessError:
+                pass
 
-        # Use Thinker-only model for text generation (saves memory by not loading Talker)
-        self.model = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(
-            self.MODEL_ID,
-            dtype=torch.bfloat16,
-            low_cpu_mem_usage=True,
-            quantization_config=quantization_config,
-            device_map="auto",
-        )
+            def content(video: Path) -> list[dict]:
+                parts: list[dict] = [{"type": "video_url", "video_url": {"url": f"file://{video.resolve()}"}}]
+                if audio_url:
+                    parts.append({"type": "audio_url", "audio_url": {"url": audio_url}})
+                parts.append({"type": "text", "text": instruction})
+                return parts
 
-        self.processor = Qwen2_5OmniProcessor.from_pretrained(self.MODEL_ID)
+            mm_kwargs = {"fps": fps}
+            try:
+                raw = self._chat(content(path), mm_kwargs)
+            except Exception as e:
+                # Raw / variable-frame-rate videos over-report their frame count, which
+                # breaks the server's frame sampler ("... frames from video"). Re-encode
+                # to a constant frame rate and retry once.
+                if "frames from video" not in str(e):
+                    raise
+                cfr = work / "video_cfr.mp4"
+                _transcode_cfr(path, cfr)
+                raw = self._chat(content(cfr), mm_kwargs)
+
+            return _parse_caption_response(raw).strip()
 
 
 class GeminiFlashCaptioner(MediaCaptioningModel):
-    """Audio-visual captioning using Google's Gemini Flash API.
-    Gemini Flash is a cloud-based multimodal model that natively supports
-    audio and video understanding. Requires a Google API key.
-    Note: This captioner requires the `google-generativeai` package and a valid API key.
-    Set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable, or pass the key directly.
+    """Audio-visual captioning using Google's Gemini via the Google Gen AI SDK.
+    Uses the ``google-genai`` package (the current SDK; ``google-generativeai``
+    is deprecated). Auth is resolved automatically:
+    1. If an API key is given (``api_key`` argument, or ``GEMINI_API_KEY`` /
+       ``GOOGLE_API_KEY`` in the environment) -> the Gemini Developer API (AI Studio).
+    2. Otherwise, if Google Cloud Application Default Credentials are available
+       (an attached service account or ``gcloud auth application-default login``)
+       -> Vertex AI. The project comes from ADC (or ``GOOGLE_CLOUD_PROJECT``) and
+       the location defaults to ``global`` (override with ``GOOGLE_CLOUD_LOCATION``).
+       This means it "just works" on a gcloud-authed GCP VM with no env vars.
+    If neither is available, a clear error explains how to authenticate.
+    Media is sent inline (``Part.from_bytes``), which works on both backends.
     """
 
-    MODEL_ID = "gemini-flash-lite-latest"
+    MODEL_ID = "gemini-3.5-flash"
+
+    _MIME_TYPES: ClassVar[dict[str, str]] = {
+        ".jpg": "image/jpeg",
+        ".jpeg": "image/jpeg",
+        ".png": "image/png",
+        ".webp": "image/webp",
+        ".heic": "image/heic",
+        ".heif": "image/heif",
+        ".mp4": "video/mp4",
+        ".mov": "video/quicktime",
+        ".avi": "video/x-msvideo",
+        ".mkv": "video/x-matroska",
+        ".webm": "video/webm",
+    }
 
     def __init__(
         self,
         api_key: str | None = None,
         instruction: str | None = None,
+        model: str | None = None,
     ):
-        """Initialize the Gemini Flash captioner.
+        """Initialize the Gemini captioner.
         Args:
-            api_key: Google API key. If not provided, will look for
-                     GEMINI_API_KEY or GOOGLE_API_KEY environment variable.
-            instruction: Custom instruction prompt. If None, uses the default instruction
+            api_key: Gemini Developer API key. If ``None``, falls back to
+                ``GEMINI_API_KEY`` / ``GOOGLE_API_KEY``; if no key is set at all,
+                uses Vertex AI via Application Default Credentials.
+            instruction: Custom instruction prompt. If ``None``, uses the default
+                image or video prompt depending on the input.
+            model: Override the served model id (defaults to ``MODEL_ID``).
         """
         self.instruction = instruction
-        self._init_client(api_key)
-
-    @property
-    def supports_audio(self) -> bool:
-        return True
+        self.model = model or self.MODEL_ID
+        self._client = self._make_client(api_key)
 
     def caption(
         self,
         path: str | Path,
-        fps: int = 3,  # noqa: ARG002 - kept for API compatibility
-        include_audio: bool = True,
-        clean_caption: bool = True,
+        fps: int = 2,  # noqa: ARG002 - kept for API compatibility
     ) -> str:
-        """Generate a caption for the given video or image.
-        Args:
-            path: Path to the video/image file to caption
-            fps: Frames per second (not used for Gemini, kept for API compatibility)
-            include_audio: Whether to include audio content in the caption
-            clean_caption: Whether to clean up the raw caption
-        Returns:
-            A string containing the generated caption
-        """
-        import time  # noqa: PLC0415
+        from google.genai import types  # noqa: PLC0415
 
         path = Path(path)
-        is_video = self._is_video_file(path)
-        use_audio = include_audio and is_video
+        instruction = self._resolve_instruction(path)
+        media = types.Part.from_bytes(data=path.read_bytes(), mime_type=self._mime_type(path))
+        response = self._client.models.generate_content(
+            model=self.model,
+            contents=[media, instruction],
+            config=types.GenerateContentConfig(temperature=0.0),
+        )
 
-        # Use custom instruction if provided, otherwise pick appropriate default
-        if self.instruction is not None:
-            instruction = self.instruction
-        else:
-            instruction = DEFAULT_CAPTION_INSTRUCTION if use_audio else VIDEO_ONLY_CAPTION_INSTRUCTION
+        # Gemini may also return JSON if it followed our prompt format.
+        return _parse_caption_response(response.text or "").strip()
 
-        # Upload the file to Gemini
-        uploaded_file = self._genai.upload_file(path)
+    @classmethod
+    def _mime_type(cls, path: Path) -> str:
+        try:
+            return cls._MIME_TYPES[path.suffix.lower()]
+        except KeyError:
+            raise ValueError(f"Unsupported media type for Gemini: {path.suffix}") from None
 
-        # Wait for processing to complete (videos need time to process)
-        while uploaded_file.state.name == "PROCESSING":
-            time.sleep(1)
-            uploaded_file = self._genai.get_file(uploaded_file.name)
+    def _make_client(self, api_key: str | None):  # noqa: ANN202 - genai.Client type is lazy-imported
+        from google import genai  # noqa: PLC0415
 
-        if uploaded_file.state.name == "FAILED":
-            raise RuntimeError(f"File processing failed: {uploaded_file.state.name}")
+        # 1. API key (explicit arg or env) -> Gemini Developer API.
+        key = api_key or os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
+        if key:
+            return genai.Client(api_key=key)
 
-        # Generate caption
-        response = self._model.generate_content([uploaded_file, instruction])
-
-        caption_raw = response.text
-
-        # Clean up the uploaded file
-        self._genai.delete_file(uploaded_file.name)
-
-        # Clean up caption if requested
-        return self._clean_raw_caption(caption_raw) if clean_caption else caption_raw
-
-    def _init_client(self, api_key: str | None) -> None:
-        """Initialize the Gemini API client."""
-        import os  # noqa: PLC0415
+        # 2. No key -> Vertex AI via Application Default Credentials (gcloud / service account).
+        import google.auth  # noqa: PLC0415
 
         try:
-            import google.generativeai as genai  # noqa: PLC0415
-        except ImportError as e:
-            raise ImportError(
-                "The `google-generativeai` package is required for Gemini Flash captioning. "
-                "Install it with: `uv pip install google-generativeai`"
+            _, adc_project = google.auth.default()
+        except Exception as e:
+            raise ValueError(
+                "No Gemini credentials found. Provide an API key (--api-key, or "
+                "GEMINI_API_KEY / GOOGLE_API_KEY), or set up Google Cloud credentials "
+                "for Vertex AI (e.g. `gcloud auth application-default login` or an "
+                "attached service account)."
             ) from e
 
-        # Get API key from argument or environment
-        # GEMINI_API_KEY is the recommended variable, GOOGLE_API_KEY also works
-        resolved_api_key = api_key or os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
-
-        if not resolved_api_key:
-            raise ValueError(
-                "Gemini API key is required. Provide it via the `api_key` argument "
-                "or set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable."
-            )
-
-        # Configure the genai library with the API key
-        genai.configure(api_key=resolved_api_key)
-
-        # Store reference to genai module for file operations
-        self._genai = genai
-
-        # Initialize the model
-        self._model = genai.GenerativeModel(self.MODEL_ID)
+        project = os.environ.get("GOOGLE_CLOUD_PROJECT") or adc_project
+        location = os.environ.get("GOOGLE_CLOUD_LOCATION", "global")
+        return genai.Client(vertexai=True, project=project, location=location)
 
 
-def example() -> None:
-    """Example usage of the captioning module."""
-    import sys  # noqa: PLC0415
+def _parse_caption_response(raw: str) -> str:
+    """Extract the caption text from a model response.
+    Backend-agnostic: works for any model that follows the combined-caption
+    prompt. Handles the formats a model may produce:
+    - Plain caption text
+    - JSON ``{"combined_caption_english": "..."}``
+    - ``<think>...</think>`` chain-of-thought followed by either of the above
+    - Truncated JSON (when generation hits a token limit mid-string)
+    """
+    text = re.sub(r"<think>[\s\S]*?</think>", "", raw).strip()
 
-    if len(sys.argv) < 2:
-        print(f"Usage: python {sys.argv[0]} <video_path> [captioner_type]")  # noqa: T201
-        print("  captioner_type: qwen_omni (default) or gemini_flash")  # noqa: T201
-        sys.exit(1)
+    # Thinking models (e.g. Qwen3-Omni-*-Thinking) emit the reasoning trace
+    # without an opening ``<think>`` tag, because the chat template injects it
+    # for them -- so the response starts mid-thought and is terminated by a lone
+    # ``</think>`` before the real answer. Drop everything up to that closer.
+    if "</think>" in text:
+        text = text.rsplit("</think>", 1)[1].strip()
 
-    video_path = sys.argv[1]
-    captioner_type = CaptionerType(sys.argv[2]) if len(sys.argv) > 2 else CaptionerType.QWEN_OMNI
+    if not text:
+        return raw.strip()
 
-    print(f"Using {captioner_type.value} captioner:")  # noqa: T201
-    captioner = create_captioner(captioner_type)
-    caption = captioner.caption(video_path)
-    print(f"CAPTION: {caption}")  # noqa: T201
+    try:
+        parsed = json.loads(text)
+        if isinstance(parsed, dict) and _CAPTION_JSON_KEY in parsed:
+            return parsed[_CAPTION_JSON_KEY]
+    except (json.JSONDecodeError, ValueError):
+        pass
+
+    match = re.search(rf"\{{[^{{}}]*\"{_CAPTION_JSON_KEY}\"[^{{}}]*\}}", text)
+    if match:
+        try:
+            parsed = json.loads(match.group())
+            if isinstance(parsed, dict) and _CAPTION_JSON_KEY in parsed:
+                return parsed[_CAPTION_JSON_KEY]
+        except (json.JSONDecodeError, ValueError):
+            pass
+
+    # Truncated JSON: extract the string value even if the closing quote/brace is missing.
+    match = re.search(rf'"{_CAPTION_JSON_KEY}"\s*:\s*"((?:[^"\\]|\\.)*)', text)
+    if match:
+        try:
+            return json.loads('"' + match.group(1) + '"')
+        except (json.JSONDecodeError, ValueError):
+            return match.group(1)
+
+    return text
 
 
-if __name__ == "__main__":
-    example()
+def _run_ffmpeg(args: list[str]) -> None:
+    """Run the ffmpeg binary bundled with ``imageio-ffmpeg`` (a dependency)."""
+    import imageio_ffmpeg  # noqa: PLC0415
+
+    cmd = [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error", *args]
+    subprocess.run(cmd, check=True, capture_output=True)
+
+
+def _extract_audio_wav(src: Path, dest: Path) -> None:
+    """Extract the audio track to a 16 kHz mono PCM WAV (matches pretraining).
+    Raises ``CalledProcessError`` when the video has no audio stream.
+    """
+    _run_ffmpeg(["-i", str(src), "-vn", "-ac", "1", "-ar", "16000", "-c:a", "pcm_s16le", str(dest)])
+
+
+def _transcode_cfr(src: Path, dest: Path) -> None:
+    """Re-encode the video to a constant frame rate so the server's frame sampler can
+    read every requested index (raw / variable-frame-rate videos over-report frames)."""
+    _run_ffmpeg(["-i", str(src), "-fps_mode", "cfr", "-c:v", "libx264", "-pix_fmt", "yuv420p", "-an", str(dest)])

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

@@ -1,10 +1,11 @@
 from pathlib import Path
-from typing import Annotated, Literal
+from typing import Annotated, Literal, Union
 
 from pydantic import BaseModel, ConfigDict, Discriminator, Field, Tag, ValidationInfo, field_validator, model_validator
 
 from ltx_trainer.quantization import QuantizationOptions
 from ltx_trainer.training_strategies.base_strategy import TrainingStrategyConfigBase
+from ltx_trainer.training_strategies.flexible import FlexibleStrategyConfig
 from ltx_trainer.training_strategies.text_to_video import TextToVideoConfig
 from ltx_trainer.training_strategies.video_to_video import VideoToVideoConfig
 
@@ -13,6 +14,226 @@ class ConfigBaseModel(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
 
+# =============================================================================
+# Validation Condition Types
+# =============================================================================
+
+
+class FirstFrameConditionConfig(ConfigBaseModel):
+    """First-frame conditioning (intrinsic, latent_idx=0). Always targets video.
+    If image_or_video points to a video file, the first frame is automatically extracted.
+    """
+
+    type: Literal["first_frame"] = "first_frame"
+    image_or_video: str | Path
+
+
+class PrefixConditionConfig(ConfigBaseModel):
+    """Prefix conditioning for temporal extension (intrinsic). Exactly one of video/audio must be set."""
+
+    type: Literal["prefix"] = "prefix"
+    video: str | None = None
+    audio: str | None = None
+    num_frames: int | None = Field(
+        default=None,
+        ge=1,
+        description="Number of pixel frames for video prefix. Must satisfy num_frames %% 8 == 1.",
+    )
+    duration: float | None = Field(default=None, gt=0, description="Duration in seconds for audio prefix")
+
+    @model_validator(mode="after")
+    def validate_exactly_one_modality(self) -> "PrefixConditionConfig":
+        if (self.video is None) == (self.audio is None):
+            raise ValueError("Exactly one of 'video' or 'audio' must be set for prefix condition")
+        return self
+
+    @model_validator(mode="after")
+    def validate_num_frames_constraint(self) -> "PrefixConditionConfig":
+        if self.video is not None and self.num_frames is not None and self.num_frames % 8 != 1:
+            raise ValueError(
+                f"num_frames ({self.num_frames}) must satisfy num_frames % 8 == 1 "
+                f"for video prefix (e.g., 1, 9, 17, 25, ...)"
+            )
+        return self
+
+
+class SuffixConditionConfig(ConfigBaseModel):
+    """Suffix conditioning for temporal extension (intrinsic). Exactly one of video/audio must be set."""
+
+    type: Literal["suffix"] = "suffix"
+    video: str | None = None
+    audio: str | None = None
+    num_frames: int | None = Field(
+        default=None,
+        ge=1,
+        description="Number of pixel frames for video suffix. Must satisfy num_frames %% 8 == 0.",
+    )
+    duration: float | None = Field(default=None, gt=0, description="Duration in seconds for audio suffix")
+
+    @model_validator(mode="after")
+    def validate_exactly_one_modality(self) -> "SuffixConditionConfig":
+        if (self.video is None) == (self.audio is None):
+            raise ValueError("Exactly one of 'video' or 'audio' must be set for suffix condition")
+        return self
+
+    @model_validator(mode="after")
+    def validate_num_frames_constraint(self) -> "SuffixConditionConfig":
+        if self.video is not None and self.num_frames is not None and self.num_frames % 8 != 0:
+            raise ValueError(
+                f"num_frames ({self.num_frames}) must satisfy num_frames % 8 == 0 "
+                f"for video suffix (e.g., 8, 16, 24, 32, ...)"
+            )
+        return self
+
+
+class SpatialCropConditionConfig(ConfigBaseModel):
+    """Spatial crop conditioning for outpainting (intrinsic, video only)."""
+
+    type: Literal["spatial_crop"] = "spatial_crop"
+    video: str
+    spatial_region: tuple[int, int, int, int] = Field(
+        ..., description="Spatial crop region as (y1, x1, y2, x2) in pixel coordinates"
+    )
+
+
+class MaskConditionConfig(ConfigBaseModel):
+    """Mask-based conditioning for inpainting (intrinsic). Exactly one of video/audio must be set."""
+
+    type: Literal["mask"] = "mask"
+    video: str | None = None
+    audio: str | None = None
+    mask: str
+
+    @model_validator(mode="after")
+    def validate_exactly_one_modality(self) -> "MaskConditionConfig":
+        if (self.video is None) == (self.audio is None):
+            raise ValueError("Exactly one of 'video' or 'audio' must be set for mask condition")
+        return self
+
+
+class ReferenceConditionConfig(ConfigBaseModel):
+    """Reference conditioning (IC-LoRA style concatenation). Exactly one of video/audio must be set."""
+
+    type: Literal["reference"] = "reference"
+    video: str | None = None
+    audio: str | None = None
+    downscale_factor: int = Field(default=1, ge=1)
+    temporal_scale_factor: int = Field(default=1, ge=1)
+    include_in_output: bool = False
+
+    @model_validator(mode="after")
+    def validate_exactly_one_modality(self) -> "ReferenceConditionConfig":
+        if (self.video is None) == (self.audio is None):
+            raise ValueError("Exactly one of 'video' or 'audio' must be set for reference condition")
+        return self
+
+
+class VideoToAudioConditionConfig(ConfigBaseModel):
+    """Video-to-audio — video is provided as frozen cross-modal conditioning.
+    The video is kept clean (sigma=0) and influences audio generation via cross-modal attention.
+    """
+
+    type: Literal["video_to_audio"] = "video_to_audio"
+    video: str
+
+
+class AudioToVideoConditionConfig(ConfigBaseModel):
+    """Audio-to-video — audio is provided as frozen cross-modal conditioning.
+    The audio is kept clean (sigma=0) and influences video generation via cross-modal attention.
+    """
+
+    type: Literal["audio_to_video"] = "audio_to_video"
+    audio: str
+
+
+ValidationCondition = Annotated[
+    Union[
+        FirstFrameConditionConfig,
+        PrefixConditionConfig,
+        SuffixConditionConfig,
+        SpatialCropConditionConfig,
+        MaskConditionConfig,
+        ReferenceConditionConfig,
+        VideoToAudioConditionConfig,
+        AudioToVideoConditionConfig,
+    ],
+    Field(discriminator="type"),
+]
+
+
+def _condition_targets_video(cond: ValidationCondition) -> bool:
+    """Check if a validation condition targets the video modality."""
+    if cond.type in ("first_frame", "spatial_crop", "video_to_audio"):
+        return True
+    if cond.type in ("prefix", "suffix", "mask", "reference"):
+        return getattr(cond, "video", None) is not None
+    return False
+
+
+def _condition_targets_audio(cond: ValidationCondition) -> bool:
+    """Check if a validation condition targets the audio modality."""
+    if cond.type == "audio_to_video":
+        return True
+    if cond.type in ("prefix", "suffix", "mask", "reference"):
+        return getattr(cond, "audio", None) is not None
+    return False
+
+
+class ValidationSample(ConfigBaseModel):
+    """Configuration for a single validation sample — fully self-describing."""
+
+    prompt: str
+    conditions: list[ValidationCondition] = Field(default_factory=list)
+
+    video_dims: tuple[int, int, int] | None = Field(
+        default=None,
+        description="Per-sample override for (width, height, frames). None = inherit from ValidationConfig.",
+    )
+    seed: int | None = Field(
+        default=None,
+        description="Per-sample override for random seed. None = inherit from ValidationConfig.",
+    )
+
+    @field_validator("video_dims")
+    @classmethod
+    def validate_video_dims(cls, v: tuple[int, int, int] | None) -> tuple[int, int, int] | None:
+        if v is None:
+            return v
+        width, height, frames = v
+        if width % 32 != 0:
+            raise ValueError(f"Width ({width}) must be divisible by 32")
+        if height % 32 != 0:
+            raise ValueError(f"Height ({height}) must be divisible by 32")
+        if frames % 8 != 1:
+            raise ValueError(f"Frames ({frames}) must satisfy frames % 8 == 1 for LTX-2 (e.g., 1, 9, 17, 25, ...)")
+        return v
+
+    @model_validator(mode="after")
+    def validate_frozen_modality_conflicts(self) -> "ValidationSample":
+        frozen_types = {c.type for c in self.conditions if c.type in ("video_to_audio", "audio_to_video")}
+
+        if "video_to_audio" in frozen_types and "audio_to_video" in frozen_types:
+            raise ValueError(
+                "Cannot have both video_to_audio and audio_to_video conditions — nothing would be generated"
+            )
+
+        if "video_to_audio" in frozen_types:
+            for c in self.conditions:
+                if c.type != "video_to_audio" and _condition_targets_video(c):
+                    raise ValueError(
+                        f"Cannot use video-targeting '{c.type}' condition when video is frozen (video_to_audio)"
+                    )
+
+        if "audio_to_video" in frozen_types:
+            for c in self.conditions:
+                if c.type != "audio_to_video" and _condition_targets_audio(c):
+                    raise ValueError(
+                        f"Cannot use audio-targeting '{c.type}' condition when audio is frozen (audio_to_video)"
+                    )
+
+        return self
+
+
 class ModelConfig(ConfigBaseModel):
     """Configuration for the base model and training mode"""
 
@@ -89,7 +310,9 @@ def _get_strategy_discriminator(v: dict | TrainingStrategyConfigBase) -> str:
 
 # Union type for all strategy configs with discriminator
 TrainingStrategyConfig = Annotated[
-    Annotated[TextToVideoConfig, Tag("text_to_video")] | Annotated[VideoToVideoConfig, Tag("video_to_video")],
+    Annotated[TextToVideoConfig, Tag("text_to_video")]
+    | Annotated[VideoToVideoConfig, Tag("video_to_video")]
+    | Annotated[FlexibleStrategyConfig, Tag("flexible")],
     Discriminator(_get_strategy_discriminator),
 ]
 
@@ -191,13 +414,32 @@ class DataConfig(ConfigBaseModel):
         ge=0,
     )
 
+    @field_validator("preprocessed_data_root")
+    @classmethod
+    def validate_preprocessed_data_root(cls, v: str) -> str:
+        """Validate that preprocessed_data_root exists."""
+        path = Path(v).expanduser().resolve()
+        if not path.exists():
+            raise ValueError(f"Dataset path does not exist: {v}")
+        if not path.is_dir():
+            raise ValueError(f"Dataset path is not a directory: {v}")
+        return str(path)
+
 
 class ValidationConfig(ConfigBaseModel):
     """Configuration for validation during training"""
 
+    # Per-sample configuration (new format — preferred)
+    samples: list[ValidationSample] = Field(
+        default_factory=list,
+        description="List of validation samples. Each sample is fully self-describing with its own "
+        "prompt, conditions, and optional overrides. Replaces prompts/images/reference_videos.",
+    )
+
+    # Legacy fields (deprecated — converted to samples internally via convert_legacy_format)
     prompts: list[str] = Field(
         default_factory=list,
-        description="List of prompts to use for validation",
+        description="[DEPRECATED: use 'samples' instead] List of prompts to use for validation",
     )
 
     negative_prompt: str = Field(
@@ -207,19 +449,22 @@ class ValidationConfig(ConfigBaseModel):
 
     images: list[str] | None = Field(
         default=None,
-        description="List of image paths to use for validation. "
+        description="[DEPRECATED: use 'samples' with first_frame conditions] "
+        "List of image paths to use for validation. "
         "One image path must be provided for each validation prompt",
     )
 
     reference_videos: list[str] | None = Field(
         default=None,
-        description="List of reference video paths to use for validation. "
+        description="[DEPRECATED: use 'samples' with reference conditions] "
+        "List of reference video paths to use for validation. "
         "One video path must be provided for each validation prompt",
     )
 
     reference_downscale_factor: int = Field(
         default=1,
-        description="Downscale factor for reference videos in IC-LoRA validation. "
+        description="[DEPRECATED: use downscale_factor on ReferenceCondition] "
+        "Downscale factor for reference videos in IC-LoRA validation. "
         "When > 1, reference videos are processed at 1/n resolution (e.g., 2 means half resolution). "
         "Must match the factor used during dataset preprocessing.",
         ge=1,
@@ -301,6 +546,13 @@ class ValidationConfig(ConfigBaseModel):
         "in validation even when not training the audio branch.",
     )
 
+    generate_video: bool = Field(
+        default=True,
+        description="Whether to generate video in validation samples. "
+        "Set to False for audio-only or v2a validation to save VRAM by skipping video VAE decoder loading. "
+        "When False, validation will only generate audio (requires generate_audio=True).",
+    )
+
     skip_initial_validation: bool = Field(
         default=False,
         description="Skip validation video sampling at step 0 (beginning of training)",
@@ -308,7 +560,8 @@ class ValidationConfig(ConfigBaseModel):
 
     include_reference_in_output: bool = Field(
         default=False,
-        description="For video-to-video training: concatenate the original reference video side-by-side "
+        description="[DEPRECATED: use include_in_output on ReferenceCondition] "
+        "For video-to-video training: concatenate the original reference video side-by-side "
         "with the generated output. The reference comes from the input video, not from the model's output.",
     )
 
@@ -346,13 +599,33 @@ class ValidationConfig(ConfigBaseModel):
 
         return v
 
+    @model_validator(mode="after")
+    def convert_legacy_format(self) -> "ValidationConfig":
+        """Convert deprecated prompts/images/reference_videos to the new samples format."""
+        if self.prompts and not self.samples:
+            samples = []
+            for i, prompt in enumerate(self.prompts):
+                conditions: list[ValidationCondition] = []
+                if self.images and i < len(self.images):
+                    conditions.append(FirstFrameConditionConfig(image_or_video=self.images[i]))
+                if self.reference_videos and i < len(self.reference_videos):
+                    conditions.append(
+                        ReferenceConditionConfig(
+                            video=self.reference_videos[i],
+                            downscale_factor=self.reference_downscale_factor,
+                            include_in_output=self.include_reference_in_output,
+                        )
+                    )
+                samples.append(ValidationSample(prompt=prompt, conditions=conditions))
+            self.samples = samples
+        return self
+
     @model_validator(mode="after")
     def validate_scaled_reference_dimensions(self) -> "ValidationConfig":
         """Validate that scaled reference dimensions are valid when reference_downscale_factor > 1."""
         if self.reference_downscale_factor > 1:
             width, height, _frames = self.video_dims
 
-            # Validate that downscale factor evenly divides the target dimensions
             if width % self.reference_downscale_factor != 0:
                 raise ValueError(
                     f"Width {width} is not evenly divisible by reference_downscale_factor "
@@ -367,7 +640,6 @@ class ValidationConfig(ConfigBaseModel):
             scaled_width = width // self.reference_downscale_factor
             scaled_height = height // self.reference_downscale_factor
 
-            # Validate scaled dimensions are divisible by 32
             if scaled_width % 32 != 0:
                 raise ValueError(
                     f"Scaled reference width {scaled_width} (from {width} / {self.reference_downscale_factor}) "
@@ -381,6 +653,16 @@ class ValidationConfig(ConfigBaseModel):
 
         return self
 
+    @model_validator(mode="after")
+    def validate_output_modality_requirements(self) -> "ValidationConfig":
+        """Validate output modality settings when validation is configured."""
+        has_validation = bool(self.prompts) or bool(self.samples)
+        if has_validation and not self.generate_video and not self.generate_audio:
+            raise ValueError(
+                "At least one of generate_video or generate_audio must be True when validation is configured."
+            )
+        return self
+
 
 class CheckpointsConfig(ConfigBaseModel):
     """Configuration for model checkpointing during training"""
@@ -514,19 +796,31 @@ class LtxTrainerConfig(ConfigBaseModel):
         """Expand user home directory in output path."""
         return str(Path(v).expanduser().resolve())
 
+    def _validate_data_dirs_exist(self) -> None:
+        """Verify that every directory declared by the training strategy exists under the data root."""
+        data_root = Path(self.data.preprocessed_data_root)
+        for dir_name in self.training_strategy.get_data_sources():
+            dir_path = data_root / dir_name
+            if not dir_path.is_dir():
+                raise ValueError(
+                    f"Required data directory '{dir_name}' does not exist under preprocessed_data_root: {dir_path}"
+                )
+
     @model_validator(mode="after")
     def validate_strategy_compatibility(self) -> "LtxTrainerConfig":
         """Validate that training strategy and other configurations are compatible."""
+        self._validate_data_dirs_exist()
 
         # Check that reference videos are provided when using video_to_video strategy
-        if (
-            self.training_strategy.name == "video_to_video"
-            and self.validation.interval
-            and not self.validation.reference_videos
-        ):
-            raise ValueError(
-                "reference_videos must be provided in validation config when using video_to_video strategy"
+        if self.training_strategy.name == "video_to_video" and self.validation.interval:
+            has_reference = bool(self.validation.reference_videos) or any(
+                cond.type == "reference" for sample in self.validation.samples for cond in sample.conditions
             )
+            if not has_reference:
+                raise ValueError(
+                    "reference_videos or samples with reference conditions must be provided "
+                    "in validation config when using video_to_video strategy"
+                )
 
         # Check that LoRA config is provided when training mode is lora
         if self.model.training_mode == "lora" and self.lora is None:

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

@@ -54,60 +54,6 @@ class SamplingContext:
         self._progress.update(self._task, visible=False)
 
 
-class StandaloneSamplingProgress:
-    """Standalone progress display for inference scripts.
-    Unlike SamplingContext (which integrates with TrainingProgress), this class
-    manages its own Rich Progress instance for use in standalone inference scripts.
-    Usage:
-        with StandaloneSamplingProgress(num_steps=30) as ctx:
-            for step in range(30):
-                # ... denoising step ...
-                ctx.advance_step()
-    """
-
-    def __init__(self, num_steps: int, description: str = "Generating"):
-        """Initialize standalone sampling progress.
-        Args:
-            num_steps: Total number of denoising steps
-            description: Description to show in progress bar
-        """
-        self._num_steps = num_steps
-        self._description = description
-        self._progress: Progress | None = None
-        self._task: TaskID | None = None
-
-    def __enter__(self) -> "StandaloneSamplingProgress":
-        """Start the progress display."""
-        self._progress = Progress(
-            TextColumn("[progress.description]{task.description}"),
-            BarColumn(bar_width=40, style="blue"),
-            TextColumn("{task.fields[info]}", style="cyan"),
-            TimeElapsedColumn(),
-            TextColumn("ETA:"),
-            TimeRemainingColumn(compact=True),
-        )
-        self._progress.__enter__()
-        self._task = self._progress.add_task(
-            self._description,
-            total=self._num_steps,
-            info=f"step 0/{self._num_steps}",
-        )
-        return self
-
-    def __exit__(self, *args) -> None:
-        """Stop the progress display."""
-        if self._progress is not None:
-            self._progress.__exit__(*args)
-
-    def advance_step(self) -> None:
-        """Advance the denoising step by one."""
-        if self._progress is None or self._task is None:
-            return
-        self._progress.advance(self._task)
-        completed = int(self._progress.tasks[self._task].completed)
-        self._progress.update(self._task, info=f"step {completed}/{self._num_steps}")
-
-
 class TrainingProgress:
     """Manages Rich progress display for training and validation.
     This class encapsulates all progress bar logic, providing a clean interface

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

@@ -12,8 +12,8 @@ from typing import Any, Callable
 import torch
 import wandb
 import yaml
-from accelerate import Accelerator, DistributedDataParallelKwargs, DistributedType
-from accelerate.utils import gather_object, set_seed
+from accelerate import Accelerator, DistributedType
+from accelerate.utils import DistributedDataParallelKwargs, 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
@@ -30,26 +30,22 @@ from torch.optim.lr_scheduler import (
     StepLR,
 )
 from torch.utils.data import DataLoader
-from torchvision.transforms import functional as F  # noqa: N812
 
 from ltx_core.text_encoders.gemma import convert_to_additive_mask
 from ltx_trainer import logger
 from ltx_trainer.config import LtxTrainerConfig
 from ltx_trainer.config_display import print_config
 from ltx_trainer.datasets import PrecomputedDataset
-from ltx_trainer.gpu_utils import free_gpu_memory, free_gpu_memory_context, get_gpu_memory_gb
+from ltx_trainer.gpu_utils import free_gpu_memory, get_gpu_memory_gb
 from ltx_trainer.hf_hub_utils import push_to_hub
-from ltx_trainer.model_loader import load_embeddings_processor, load_text_encoder
-from ltx_trainer.model_loader import load_model as load_ltx_model
+from ltx_trainer.model_loader import load_embeddings_processor, load_transformer
 from ltx_trainer.progress import TrainingProgress
 from ltx_trainer.quantization import quantize_model
 from ltx_trainer.sigma_tracker import SigmaBucketTracker
 from ltx_trainer.timestep_samplers import SAMPLERS
 from ltx_trainer.training_state import ConfigFingerprint, RngStates, TrainingState
 from ltx_trainer.training_strategies import get_training_strategy
-from ltx_trainer.utils import open_image_as_srgb, save_image
-from ltx_trainer.validation_sampler import CachedPromptEmbeddings, GenerationConfig, ValidationSampler
-from ltx_trainer.video_utils import read_video, save_video
+from ltx_trainer.validation_runner import ValidationRunner
 
 # Disable irrelevant warnings from transformers
 os.environ["TOKENIZERS_PARALLELISM"] = "true"
@@ -66,7 +62,7 @@ if not IS_MAIN_PROCESS:
 
     disable_progress_bar()
 
-StepCallback = Callable[[int, int, list[Path] | None], None]  # (step, total, sampled paths or None) -> None
+StepCallback = Callable[[int, int, list[Path]], None]  # (step, total, list[sampled_video_path]) -> None
 
 MEMORY_CHECK_INTERVAL = 200
 
@@ -96,7 +92,16 @@ class LtxvTrainer:
         if IS_MAIN_PROCESS:
             print_config(trainer_config)
         self._training_strategy = get_training_strategy(self._config.training_strategy)
-        self._cached_validation_embeddings = self._load_text_encoder_and_cache_embeddings()
+
+        # ValidationRunner loads its own models (text encoder, VAE encoder/decoder, etc.),
+        # caches prompt embeddings and conditioning media, then unloads encoders.
+        self._validation_runner = ValidationRunner(
+            config=self._config.validation,
+            model_path=self._config.model.model_path,
+            text_encoder_path=self._config.model.text_encoder_path,
+            load_text_encoder_in_8bit=self._config.acceleration.load_text_encoder_in_8bit,
+        )
+
         self._load_models()
         self._setup_accelerator()
         self._collect_trainable_params()
@@ -108,8 +113,8 @@ class LtxvTrainer:
         self._checkpoint_paths: list[Path] = []
         self._training_state_paths: list[Path] = []
         self._training_state_size_warned = False
-        self._wandb_run = None
         self._sigma_tracker = SigmaBucketTracker()
+        self._wandb_run = None
 
     def train(  # noqa: PLR0912, PLR0915
         self,
@@ -190,7 +195,7 @@ class LtxvTrainer:
         with progress:
             if cfg.validation.interval and not cfg.validation.skip_initial_validation:
                 with self._offloaded_optimizer_state():
-                    sampled_videos_paths = self._run_distributed_validation(progress)
+                    sampled_videos_paths = self._run_validation(progress)
 
             self._accelerator.wait_for_everyone()
 
@@ -223,7 +228,7 @@ class LtxvTrainer:
                     if self._lr_scheduler is not None:
                         self._lr_scheduler.step()
 
-                    # Run validation if needed
+                    # Run validation if needed (handles DDP/FSDP work distribution internally)
                     if (
                         cfg.validation.interval
                         and self._global_step > 0
@@ -231,7 +236,7 @@ class LtxvTrainer:
                         and is_optimization_step
                     ):
                         with self._offloaded_optimizer_state():
-                            sampled_videos_paths = self._run_distributed_validation(progress)
+                            sampled_videos_paths = self._run_validation(progress)
 
                     # Save checkpoint if needed
                     if (
@@ -376,111 +381,39 @@ class LtxvTrainer:
             perturbations=None,
         )
 
-        # Use strategy to compute loss
+        # Use strategy to compute loss (returns per-element [B,] for sigma-bucket tracking)
         loss = self._training_strategy.compute_loss(video_pred, audio_pred, model_inputs)
-        sigma = model_inputs.video.sigma.detach() if model_inputs.video.enabled else model_inputs.audio.sigma.detach()
+
+        # Sigma comes from whichever modality is generated (video preferred, else audio).
+        if model_inputs.video is not None and model_inputs.video.enabled:
+            sigma = model_inputs.video.sigma.detach()
+        else:
+            sigma = model_inputs.audio.sigma.detach()
 
         return TrainingStepOutput(loss=loss, sigma=sigma)
 
-    @free_gpu_memory_context(after=True)
-    def _load_text_encoder_and_cache_embeddings(self) -> list[CachedPromptEmbeddings] | None:
-        """Load text encoder + embeddings processor, compute and cache validation embeddings."""
+    def _load_models(self) -> None:
+        """Load the transformer and embeddings processor for training."""
+        logger.debug("Loading transformer...")
+        self._transformer = load_transformer(
+            checkpoint_path=self._config.model.model_path,
+            device="cpu",
+            dtype=torch.bfloat16,
+        )
 
-        # This method:
-        #   1. Loads the pure Gemma text encoder on GPU
-        #   2. Loads the embeddings processor (feature extractor + connectors)
-        #   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 — LOCAL_RANK before Accelerator exists
+        # DDP-safe: LOCAL_RANK is set by accelerate before trainer init. Loading on bare
+        # "cuda" would resolve to cuda:0 on every rank and crash with a device mismatch.
         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=init_device,
-            dtype=torch.bfloat16,
-            load_in_8bit=self._config.acceleration.load_text_encoder_in_8bit,
-        )
-
-        # Load embeddings processor (feature extractor + connectors)
         logger.debug("Loading embeddings processor...")
         self._embeddings_processor = load_embeddings_processor(
             checkpoint_path=self._config.model.model_path,
             device=init_device,
             dtype=torch.bfloat16,
         )
-
-        # Cache validation embeddings if prompts are configured
-        cached_embeddings = None
-        if self._config.validation.prompts:
-            logger.info(f"Pre-computing embeddings for {len(self._config.validation.prompts)} validation prompts...")
-            cached_embeddings = []
-            with torch.inference_mode():
-                for prompt in self._config.validation.prompts:
-                    pos_hs, pos_mask = text_encoder.encode(prompt)
-                    pos_out = self._embeddings_processor.process_hidden_states(pos_hs, pos_mask)
-
-                    neg_hs, neg_mask = text_encoder.encode(self._config.validation.negative_prompt)
-                    neg_out = self._embeddings_processor.process_hidden_states(neg_hs, neg_mask)
-
-                    cached_embeddings.append(
-                        CachedPromptEmbeddings(
-                            video_context_positive=pos_out.video_encoding.cpu(),
-                            audio_context_positive=pos_out.audio_encoding.cpu(),
-                            video_context_negative=neg_out.video_encoding.cpu(),
-                            audio_context_negative=(
-                                neg_out.audio_encoding.cpu() if neg_out.audio_encoding is not None else None
-                            ),
-                        )
-                    )
-
-        # Unload Gemma model and feature extractor, keep only connectors for training
-        del text_encoder
         self._embeddings_processor.feature_extractor = None
 
-        logger.debug("Validation prompt embeddings cached. Gemma model unloaded")
-        return cached_embeddings
-
-    def _load_models(self) -> None:
-        """Load the LTX-2 model components."""
-        # Load audio components if:
-        # 1. Training strategy requires audio (training the audio branch), OR
-        # 2. Validation is configured to generate audio (even if not training audio)
-        load_audio = self._training_strategy.requires_audio or self._config.validation.generate_audio
-
-        # Check if we need VAE encoder (for image or reference video conditioning)
-        need_vae_encoder = (
-            self._config.validation.images is not None or self._config.validation.reference_videos is not None
-        )
-
-        # Load all model components (except text encoder - already handled)
-        components = load_ltx_model(
-            checkpoint_path=self._config.model.model_path,
-            device="cpu",
-            dtype=torch.bfloat16,
-            with_video_vae_encoder=need_vae_encoder,  # Needed for image conditioning
-            with_video_vae_decoder=True,  # Needed for validation sampling
-            with_audio_vae_decoder=load_audio,
-            with_vocoder=load_audio,
-            with_text_encoder=False,  # Text encoder handled separately
-        )
-
-        # Extract components
-        self._transformer = components.transformer
-        self._vae_decoder = components.video_vae_decoder.to(dtype=torch.bfloat16)
-        self._vae_encoder = components.video_vae_encoder
-        if self._vae_encoder is not None:
-            self._vae_encoder = self._vae_encoder.to(dtype=torch.bfloat16)
-        self._scheduler = components.scheduler
-        self._audio_vae = components.audio_vae_decoder
-        self._vocoder = components.vocoder
-        # Note: self._embeddings_processor was set in _load_text_encoder_and_cache_embeddings
-
-        # Determine initial dtype based on training mode.
-        # Note: For FSDP + LoRA, we'll cast to FP32 later in _prepare_models_for_training()
-        # after the accelerator is set up, and we can detect FSDP.
         transformer_dtype = torch.bfloat16 if self._config.model.training_mode == "lora" else torch.float32
         self._transformer = self._transformer.to(dtype=transformer_dtype)
 
@@ -494,16 +427,7 @@ class LtxvTrainer:
                 precision=self._config.acceleration.quantization,
             )
 
-        # Freeze all models. We later unfreeze the transformer based on training mode.
-        # Note: embedding_connectors are already frozen (they come from the frozen text encoder)
-        self._vae_decoder.requires_grad_(False)
-        if self._vae_encoder is not None:
-            self._vae_encoder.requires_grad_(False)
         self._transformer.requires_grad_(False)
-        if self._audio_vae is not None:
-            self._audio_vae.requires_grad_(False)
-        if self._vocoder is not None:
-            self._vocoder.requires_grad_(False)
 
     def _collect_trainable_params(self) -> None:
         """Collect trainable parameters based on training mode."""
@@ -692,13 +616,6 @@ class LtxvTrainer:
 
         transformer.set_gradient_checkpointing(self._config.optimization.enable_gradient_checkpointing)
 
-        # Keep frozen models on CPU for memory efficiency
-        self._vae_decoder = self._vae_decoder.to("cpu")
-        if self._vae_encoder is not None:
-            self._vae_encoder = self._vae_encoder.to("cpu")
-
-        # Embedding connectors are already on GPU from _load_text_encoder_and_cache_embeddings
-
         # noinspection PyTypeChecker
         self._transformer = self._accelerator.prepare(self._transformer)
 
@@ -740,7 +657,7 @@ class LtxvTrainer:
         """Initialize the training data loader using the strategy's data sources."""
         if self._dataset is None:
             # Get data sources from the training strategy
-            data_sources = self._training_strategy.get_data_sources()
+            data_sources = self._config.training_strategy.get_data_sources()
 
             self._dataset = PrecomputedDataset(self._config.data.preprocessed_data_root, data_sources=data_sources)
             logger.debug(f"Loaded dataset with {len(self._dataset):,} samples from sources: {list(data_sources)}")
@@ -785,41 +702,6 @@ 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
@@ -920,164 +802,104 @@ 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.
+    @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 _run_validation(self, progress: TrainingProgress) -> list[Path]:
+        """Run distributed validation by delegating to the ValidationRunner.
+        Each rank generates its assigned subset of validation samples (round-robin by
+        `process_index`/`num_processes`), 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.
+        Under FSDP with multiple processes, ranks pad with extra generate passes
+        (same sample, no disk write) so every rank runs the same number of forwards --
+        avoids collective mismatch.
         Note: Multi-node training requires a shared filesystem so rank 0 can read
         videos written by other ranks.
         """
-        sampled = self._sample_videos(progress)
+        self._optimizer.zero_grad(set_to_none=True)
+        free_gpu_memory()
 
-        if self._accelerator.num_processes > 1:
-            # gather_object returns a flat list from all ranks
+        num_samples = len(self._config.validation.samples)
+        if num_samples == 0:
+            return []
+
+        rank = self._accelerator.process_index
+        world_size = self._accelerator.num_processes
+
+        rank_indices = list(range(rank, num_samples, world_size))
+        work_items: list[tuple[int, bool]] = [(i, True) for i in rank_indices]
+        if self._accelerator.distributed_type == DistributedType.FSDP and world_size > 1:
+            # FSDP forwards run collective ops; pad short ranks with no-save duplicates so
+            # every rank executes the same number of forwards. A rank with empty
+            # rank_indices (world_size > num_samples) still pads with sample 0 to stay in
+            # sync with the others.
+            max_per_rank = math.ceil(num_samples / world_size)
+            pad_seed = rank_indices[-1] if rank_indices else 0
+            work_items += [(pad_seed, False)] * (max_per_rank - len(work_items))
+
+        # W&B logging is handled by the trainer (after gathering across ranks),
+        # so we always pass wandb_run=None to the runner.
+        sampled = self._validation_runner.run(
+            transformer=self._transformer,
+            step=self._global_step,
+            output_dir=Path(self._config.output_dir),
+            device=self._accelerator.device,
+            progress=progress,
+            wandb_run=None,
+            work_items=work_items,
+        )
+
+        if world_size > 1:
             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)
+        if (
+            self._accelerator.is_main_process
+            and paths
+            and self._config.wandb.log_validation_videos
+            and self._wandb_run is not None
+        ):
+            self._validation_runner.log_to_wandb(self._wandb_run, paths, self._global_step)
 
         # 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[tuple[int, Path]]:
-        """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
-        inference_steps = self._config.validation.inference_steps
-
-        # Zero gradients and free GPU memory to reclaim memory before validation sampling
-        self._optimizer.zero_grad(set_to_none=True)
-        free_gpu_memory()
-
-        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(work),
-            num_steps=inference_steps,
-        )
-
-        # Create a validation sampler with loaded models and progress tracking
-        sampler = ValidationSampler(
-            transformer=self._transformer,
-            vae_decoder=self._vae_decoder,
-            vae_encoder=self._vae_encoder,
-            text_encoder=None,
-            audio_decoder=self._audio_vae if generate_audio else None,
-            vocoder=self._vocoder if generate_audio else None,
-            sampling_context=sampling_ctx,
-        )
-
-        output_dir = Path(self._config.output_dir) / "samples"
-        output_dir.mkdir(exist_ok=True, parents=True)
-
-        results: list[tuple[int, Path]] = []
-        width, height, num_frames = self._config.validation.video_dims
-
-        for local_i, (prompt_idx, save_output) in enumerate(work):
-            prompt = prompts[prompt_idx]
-            sampling_ctx.start_video(local_i)
-
-            # Load conditioning image if provided
-            condition_image = None
-            if use_images:
-                image_path = self._config.validation.images[prompt_idx]
-                image = open_image_as_srgb(image_path)
-                # Convert PIL image to tensor [C, H, W] in [0, 1]
-                condition_image = F.to_tensor(image)
-
-            # Load reference video if provided (for IC-LoRA)
-            reference_video = None
-            if use_reference_videos:
-                ref_video_path = self._config.validation.reference_videos[prompt_idx]
-                # read_video returns [F, C, H, W] in [0, 1]
-                reference_video, _ = read_video(ref_video_path, max_frames=num_frames)
-
-            # Get cached embeddings for this prompt if available
-            cached_embeddings = (
-                self._cached_validation_embeddings[prompt_idx]
-                if self._cached_validation_embeddings is not None
-                else None
-            )
-
-            # Create generation config
-            gen_config = GenerationConfig(
-                prompt=prompt,
-                negative_prompt=self._config.validation.negative_prompt,
-                height=height,
-                width=width,
-                num_frames=num_frames,
-                frame_rate=self._config.validation.frame_rate,
-                num_inference_steps=inference_steps,
-                guidance_scale=self._config.validation.guidance_scale,
-                seed=self._config.validation.seed,
-                condition_image=condition_image,
-                reference_video=reference_video,
-                reference_downscale_factor=self._config.validation.reference_downscale_factor,
-                generate_audio=generate_audio,
-                include_reference_in_output=self._config.validation.include_reference_in_output,
-                cached_embeddings=cached_embeddings,
-                stg_scale=self._config.validation.stg_scale,
-                stg_blocks=self._config.validation.stg_blocks,
-                stg_mode=self._config.validation.stg_mode,
-            )
-
-            # Generate sample
-            video, audio = sampler.generate(
-                config=gen_config,
-                device=self._accelerator.device,
-            )
-
-            if not save_output:
-                continue
-
-            # Save output (image for single frame, video otherwise)
-            ext = "png" if num_frames == 1 else "mp4"
-            output_path = output_dir / f"step_{self._global_step:06d}_{prompt_idx + 1:02d}.{ext}"
-            if num_frames == 1:
-                save_image(video, output_path)
-            else:
-                save_video(
-                    video_tensor=video,
-                    output_path=output_path,
-                    fps=self._config.validation.frame_rate,
-                    audio=audio,
-                    audio_sample_rate=self._vocoder.output_sampling_rate if audio is not None else None,
-                )
-            results.append((prompt_idx, 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 results
-
     @staticmethod
     def _log_training_stats(stats: TrainingStats) -> None:
         """Log training statistics."""
@@ -1166,8 +988,8 @@ class LtxvTrainer:
     def _save_training_state(self, save_dir: Path) -> None:
         """Save training state alongside checkpoint for resume.
         Respects checkpoints.save_training_state config:
-        - "full": optimizer + scheduler + RNG + step + wandb_run_id
-        - "minimal": scheduler + RNG + step + wandb_run_id
+        - "full": optimizer + scheduler + RNG + step
+        - "minimal": scheduler + RNG + step only
         - "off": skip entirely
         """
         if not IS_MAIN_PROCESS:
@@ -1270,7 +1092,7 @@ class LtxvTrainer:
         logger.info(f"💾 Training configuration saved to: {config_path.relative_to(self._config.output_dir)}")
 
     def _init_wandb(self, resume_run_id: str | None = None) -> None:
-        """Initialize Weights & Biases run."""
+        """Initialize Weights & Biases run, resuming an existing run if its id is provided."""
         if not self._config.wandb.enabled or not IS_MAIN_PROCESS:
             self._wandb_run = None
             return
@@ -1285,7 +1107,7 @@ class LtxvTrainer:
         }
         if resume_run_id is not None:
             init_kwargs["id"] = resume_run_id
-            init_kwargs["resume"] = "allow"
+            init_kwargs["resume"] = "must"
         run = wandb.init(**init_kwargs)
         self._wandb_run = run
 
@@ -1293,22 +1115,3 @@ class LtxvTrainer:
         """Log metrics to Weights & Biases."""
         if self._wandb_run is not None:
             self._wandb_run.log(metrics)
-
-    def _log_validation_samples(self, sample_paths: list[Path], prompts: list[str]) -> None:
-        """Log validation samples (videos or images) to Weights & Biases."""
-        if not self._config.wandb.log_validation_videos or self._wandb_run is None:
-            return
-
-        # Determine if outputs are images or videos based on file extension
-        is_image = sample_paths and sample_paths[0].suffix.lower() in (".png", ".jpg", ".jpeg", ".heic", ".webp")
-
-        if is_image:
-            samples = [
-                wandb.Image(str(path), caption=prompt) for path, prompt in zip(sample_paths, prompts, strict=True)
-            ]
-        else:
-            samples = [
-                wandb.Video(str(path), caption=prompt, format=path.suffix.lower().lstrip("."))
-                for path, prompt in zip(sample_paths, prompts, strict=True)
-            ]
-        self._wandb_run.log({"validation_samples": samples}, step=self._global_step)

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

@@ -1,10 +1,13 @@
 """Training strategies for different conditioning modes.
 This package implements the Strategy Pattern to handle different training modes:
-- Text-to-video training (standard generation, optionally with audio)
-- Video-to-video training (IC-LoRA mode with reference videos)
+- Text-to-video training (standard generation, optionally with audio) [DEPRECATED]
+- Video-to-video training (IC-LoRA mode with reference videos) [DEPRECATED]
+- Flexible training (unified conditioning framework supporting all scenarios) [RECOMMENDED]
 Each strategy encapsulates the specific logic for preparing model inputs and computing loss.
 """
 
+import warnings
+
 from ltx_trainer import logger
 from ltx_trainer.training_strategies.base_strategy import (
     DEFAULT_FPS,
@@ -13,15 +16,18 @@ from ltx_trainer.training_strategies.base_strategy import (
     TrainingStrategy,
     TrainingStrategyConfigBase,
 )
+from ltx_trainer.training_strategies.flexible import FlexibleStrategy, FlexibleStrategyConfig
 from ltx_trainer.training_strategies.text_to_video import TextToVideoConfig, TextToVideoStrategy
 from ltx_trainer.training_strategies.video_to_video import VideoToVideoConfig, VideoToVideoStrategy
 
 # Type alias for all strategy config types
-TrainingStrategyConfig = TextToVideoConfig | VideoToVideoConfig
+TrainingStrategyConfig = TextToVideoConfig | VideoToVideoConfig | FlexibleStrategyConfig
 
 __all__ = [
     "DEFAULT_FPS",
     "VIDEO_SCALE_FACTORS",
+    "FlexibleStrategy",
+    "FlexibleStrategyConfig",
     "ModelInputs",
     "TextToVideoConfig",
     "TextToVideoStrategy",
@@ -43,16 +49,42 @@ def get_training_strategy(config: TrainingStrategyConfig) -> TrainingStrategy:
         The appropriate training strategy instance
     Raises:
         ValueError: If strategy name is not supported
+    Note:
+        The `text_to_video` and `video_to_video` strategies are deprecated.
+        Please use the `flexible` strategy instead.
     """
 
     match config:
         case TextToVideoConfig():
+            warnings.warn(
+                "The 'text_to_video' training strategy is deprecated and will be removed "
+                "in a future version. Please migrate to the 'flexible' strategy. "
+                "See the migration guide in the documentation.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
             strategy = TextToVideoStrategy(config)
         case VideoToVideoConfig():
+            warnings.warn(
+                "The 'video_to_video' training strategy is deprecated and will be removed "
+                "in a future version. Please migrate to the 'flexible' strategy. "
+                "See the migration guide in the documentation.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
             strategy = VideoToVideoStrategy(config)
+        case FlexibleStrategyConfig():
+            strategy = FlexibleStrategy(config)
         case _:
             raise ValueError(f"Unknown training strategy config type: {type(config).__name__}")
 
-    audio_mode = "(audio enabled 🔈)" if getattr(config, "with_audio", False) else "(audio disabled 🔇)"
+    # Determine audio mode for logging
+    if hasattr(config, "with_audio"):
+        audio_mode = "(audio enabled 🔈)" if config.with_audio else "(audio disabled 🔇)"
+    elif hasattr(config, "audio") and config.audio is not None:
+        audio_mode = "(audio enabled 🔈)"
+    else:
+        audio_mode = "(audio disabled 🔇)"
+
     logger.debug(f"🎯 Using {strategy.__class__.__name__} training strategy {audio_mode}")
     return strategy

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

@@ -34,29 +34,36 @@ class TrainingStrategyConfigBase(BaseModel):
 
     model_config = ConfigDict(extra="forbid")
 
-    name: Literal["text_to_video", "video_to_video"] = Field(
+    name: Literal["text_to_video", "video_to_video", "flexible"] = Field(
         description="Unique name identifying the training strategy type"
     )
 
+    @abstractmethod
+    def get_data_sources(self) -> dict[str, str]:
+        """Get the required data sources for this strategy.
+        Returns a mapping of directory name (relative to ``preprocessed_data_root``)
+        to the dataset output key under which that directory's contents are exposed.
+        This is the single source of truth for which directories the strategy needs:
+        it drives both dataset wiring (in the trainer) and existence validation
+        (in ``LtxTrainerConfig``).
+        """
+
 
 @dataclass
 class ModelInputs:
     """Container for model inputs using the Modality-based interface."""
 
-    video: Modality
+    video: Modality | None
     audio: Modality | None
 
     # Training targets (for loss computation)
-    video_targets: Tensor
+    video_targets: Tensor | None
     audio_targets: Tensor | None
 
-    # Masks for loss computation
-    video_loss_mask: Tensor  # Boolean mask: True = compute loss for this token
+    # Masks for loss computation (True = compute loss for this token)
+    video_loss_mask: Tensor | None
     audio_loss_mask: Tensor | None
 
-    # Metadata needed for loss computation in some strategies
-    ref_seq_len: int | None = None  # For IC-LoRA: length of reference sequence
-
 
 class TrainingStrategy(ABC):
     """Abstract base class for training strategies.
@@ -73,24 +80,6 @@ class TrainingStrategy(ABC):
         self._video_patchifier = VideoLatentPatchifier(patch_size=1)
         self._audio_patchifier = AudioPatchifier(patch_size=1)
 
-    @property
-    def requires_audio(self) -> bool:
-        """Whether this training strategy requires audio components.
-        Override this property in subclasses that support audio training.
-        The trainer uses this to determine whether to load audio VAE and vocoder.
-        Returns:
-            True if audio components should be loaded, False otherwise.
-        """
-        return False
-
-    @abstractmethod
-    def get_data_sources(self) -> list[str] | dict[str, str]:
-        """Get the required data sources for this training strategy.
-        Returns:
-            Either a list of data directory names (where output keys match directory names)
-            or a dictionary mapping data directory names to custom output keys for the dataset
-        """
-
     @abstractmethod
     def prepare_training_inputs(
         self,
@@ -145,7 +134,6 @@ class TrainingStrategy(ABC):
         batch_size: int,
         fps: float,
         device: torch.device,
-        dtype: torch.dtype,
     ) -> Tensor:
         """Generate video position embeddings using ltx_core's native implementation.
         Args:
@@ -155,9 +143,8 @@ class TrainingStrategy(ABC):
             batch_size: Batch size
             fps: Frames per second
             device: Target device
-            dtype: Target dtype
         Returns:
-            Position tensor of shape [B, 3, seq_len, 2]
+            Position tensor of shape [B, 3, seq_len, 2] (float32)
         """
         latent_coords = self._video_patchifier.get_patch_grid_bounds(
             output_shape=VideoLatentShape(
@@ -175,7 +162,7 @@ class TrainingStrategy(ABC):
             latent_coords=latent_coords,
             scale_factors=VIDEO_SCALE_FACTORS,
             causal_fix=True,
-        ).to(dtype)
+        ).float()
 
         # Scale temporal dimension by 1/fps to get time in seconds
         pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / fps
@@ -187,14 +174,12 @@ class TrainingStrategy(ABC):
         num_time_steps: int,
         batch_size: int,
         device: torch.device,
-        dtype: torch.dtype,
     ) -> Tensor:
         """Generate audio position embeddings using ltx_core's native implementation.
         Args:
             num_time_steps: Number of audio time steps (T, not T*mel_bins)
             batch_size: Batch size
             device: Target device
-            dtype: Target dtype
         Returns:
             Position tensor of shape [B, 1, num_time_steps, 2]
         Note:
@@ -204,7 +189,7 @@ class TrainingStrategy(ABC):
         """
         mel_bins = 16
 
-        latent_coords = self._audio_patchifier.get_patch_grid_bounds(
+        return self._audio_patchifier.get_patch_grid_bounds(
             output_shape=AudioLatentShape(
                 frames=num_time_steps,
                 mel_bins=mel_bins,
@@ -214,8 +199,6 @@ class TrainingStrategy(ABC):
             device=device,
         )
 
-        return latent_coords.to(dtype)
-
     @staticmethod
     def _create_per_token_timesteps(conditioning_mask: Tensor, sampled_sigma: Tensor) -> Tensor:
         """Create per-token timesteps based on conditioning mask.

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

@@ -0,0 +1,718 @@
+"""Flexible training strategy for a unified conditioning framework.
+This strategy implements the Unified Conditioning Framework that supports:
+- Simple fine-tuning with text conditioning (text-to-video/audio)
+- Intrinsic conditioning (first_frame, prefix, suffix, spatial_crop, mask)
+- Extrinsic conditioning (concatenation-based, IC-LoRA style)
+The flexible strategy replaces TextToVideoStrategy and VideoToVideoStrategy by expressing
+all conditioning scenarios through configuration rather than code.
+"""
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Annotated, Any, Literal, Union
+
+import torch
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+from torch import Tensor
+
+from ltx_core.model.transformer.modality import Modality
+from ltx_trainer.timestep_samplers import TimestepSampler
+from ltx_trainer.training_strategies.base_strategy import (
+    DEFAULT_FPS,
+    VIDEO_SCALE_FACTORS,
+    ModelInputs,
+    TrainingStrategy,
+    TrainingStrategyConfigBase,
+)
+
+# =============================================================================
+# Configuration Classes
+# =============================================================================
+
+
+class IntrinsicConditionBase(BaseModel):
+    """Base for intrinsic conditioning — tokens get clean latents, timestep=0, no loss."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    probability: float = Field(
+        default=1.0,
+        ge=0.0,
+        le=1.0,
+        description="Probability of applying this condition",
+    )
+
+
+class FirstFrameConditionConfig(IntrinsicConditionBase):
+    """First frame conditioning — frame 0 is clean, excluded from loss."""
+
+    type: Literal["first_frame"] = "first_frame"
+
+
+class PrefixConditionConfig(IntrinsicConditionBase):
+    """Prefix conditioning — first N temporal units are clean, excluded from loss."""
+
+    type: Literal["prefix"] = "prefix"
+    temporal_boundary: int = Field(
+        ...,
+        ge=1,
+        description="Number of temporal units for prefix region. "
+        "For video: number of latent frames. For audio: number of audio latent timesteps.",
+    )
+
+
+class SuffixConditionConfig(IntrinsicConditionBase):
+    """Suffix conditioning — last N temporal units are clean, excluded from loss."""
+
+    type: Literal["suffix"] = "suffix"
+    temporal_boundary: int = Field(
+        ...,
+        ge=1,
+        description="Number of temporal units for suffix region. "
+        "For video: number of latent frames. For audio: number of audio latent timesteps.",
+    )
+
+
+class SpatialCropConditionConfig(IntrinsicConditionBase):
+    """Spatial crop conditioning — rectangular pixel region is clean, excluded from loss."""
+
+    type: Literal["spatial_crop"] = "spatial_crop"
+    spatial_region: tuple[int, int, int, int] = Field(
+        ...,
+        description="Spatial crop region as (y1, x1, y2, x2) in pixel coordinates",
+    )
+
+
+class MaskConditionConfig(IntrinsicConditionBase):
+    """Mask conditioning — per-sample binary mask determines conditioning tokens."""
+
+    type: Literal["mask"] = "mask"
+    mask_dir: str = Field(
+        ...,
+        description="Directory containing per-sample masks",
+    )
+
+
+class ReferenceConditionConfig(BaseModel):
+    """Reference conditioning (IC-LoRA style concatenation).
+    External reference latents are concatenated to the target sequence.
+    Reference tokens are clean (timestep=0), excluded from loss, and
+    participate in bidirectional self-attention.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    type: Literal["reference"] = "reference"
+    latents_dir: str = Field(..., description="Directory for reference latents")
+    probability: float = Field(default=1.0, ge=0.0, le=1.0, description="Probability of applying this condition")
+
+
+# Discriminated union for condition configs
+ConditionConfig = Annotated[
+    Union[
+        FirstFrameConditionConfig,
+        PrefixConditionConfig,
+        SuffixConditionConfig,
+        SpatialCropConditionConfig,
+        MaskConditionConfig,
+        ReferenceConditionConfig,
+    ],
+    Field(discriminator="type"),
+]
+
+
+class ModalityConfig(BaseModel):
+    """Configuration for a single modality (video or audio)."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    is_generated: bool = Field(
+        ...,
+        description="True = generated modality (denoised, contributes to loss), False = conditioning-only modality",
+    )
+
+    latents_dir: str = Field(
+        ...,
+        description="Directory for latents",
+    )
+
+    conditions: list[ConditionConfig] = Field(
+        default_factory=list,
+        description="List of conditions (e.g. first_frame, prefix, reference). Text conditioning is always applied.",
+    )
+
+
+class FlexibleStrategyConfig(TrainingStrategyConfigBase):
+    """Configuration for the flexible training strategy.
+    This strategy supports all conditioning scenarios through configuration:
+    - Text-to-video/audio with simple fine-tuning
+    - Intrinsic conditioning like first-frame, extension, outpainting
+    - Reference conditioning like IC-LoRA (concatenation-based reference)
+    """
+
+    name: Literal["flexible"] = "flexible"
+
+    video: ModalityConfig | None = Field(
+        default=None,
+        description="Video modality configuration",
+    )
+
+    audio: ModalityConfig | None = Field(
+        default=None,
+        description="Audio modality configuration",
+    )
+
+    @model_validator(mode="after")
+    def validate_at_least_one_generated(self) -> "FlexibleStrategyConfig":
+        """Ensure at least one modality has is_generated=true."""
+        has_video_target = self.video is not None and self.video.is_generated
+        has_audio_target = self.audio is not None and self.audio.is_generated
+        if not has_video_target and not has_audio_target:
+            raise ValueError("At least one modality must have is_generated=true")
+        return self
+
+    @model_validator(mode="after")
+    def validate_audio_intrinsic_regions(self) -> "FlexibleStrategyConfig":
+        """Reject video-only intrinsic regions on the audio modality."""
+        if self.audio is None:
+            return self
+        for cond in self.audio.conditions:
+            if isinstance(cond, (FirstFrameConditionConfig, SpatialCropConditionConfig)):
+                raise ValueError(
+                    f"Intrinsic condition '{cond.type}' is not supported for audio. "
+                    f"Audio supports: prefix, suffix, mask."
+                )
+        return self
+
+    def get_data_sources(self) -> dict[str, str]:
+        """Dynamically determine required data sources from config.
+        Returns a mapping of directory name (under ``preprocessed_data_root``) to
+        the dataset output key.
+        """
+        sources: dict[str, str] = {"conditions": "conditions"}
+
+        if self.video is not None:
+            sources[self.video.latents_dir] = "video_latents"
+        if self.audio is not None:
+            sources[self.audio.latents_dir] = "audio_latents"
+
+        for modality_config in (self.video, self.audio):
+            if modality_config is None:
+                continue
+            for cond in modality_config.conditions:
+                if isinstance(cond, ReferenceConditionConfig):
+                    sources[cond.latents_dir] = cond.latents_dir
+                elif isinstance(cond, MaskConditionConfig):
+                    sources[cond.mask_dir] = cond.mask_dir
+
+        return sources
+
+
+# =============================================================================
+# Helper Data Structures
+# =============================================================================
+
+
+@dataclass
+class ModalityProcessingResult:
+    """Result of processing a single modality."""
+
+    modality: Modality
+    targets: Tensor | None
+    loss_mask: Tensor | None
+
+
+@dataclass
+class LatentData:
+    """Loaded and patchified latents with metadata."""
+
+    latents: Tensor  # [B, seq_len, C]
+    num_frames: int
+    height: int
+    width: int
+    fps: float
+
+
+# =============================================================================
+# FlexibleStrategy Implementation
+# =============================================================================
+
+
+class FlexibleStrategy(TrainingStrategy):
+    """Unified training strategy supporting all conditioning scenarios.
+    This strategy implements the Unified Conditioning Framework, allowing
+    any training scenario to be expressed through configuration.
+    """
+
+    config: FlexibleStrategyConfig
+
+    def __init__(self, config: FlexibleStrategyConfig):
+        """Initialize strategy with configuration.
+        Args:
+            config: Flexible strategy configuration
+        """
+        super().__init__(config)
+        self.config = config
+        self.reference_spatial_scale_factor, self.reference_temporal_scale_factor = (
+            self._infer_reference_scale_factors_from_config()
+        )
+
+    def prepare_training_inputs(
+        self,
+        batch: dict[str, Any],
+        timestep_sampler: TimestepSampler,
+    ) -> ModelInputs:
+        """Prepare training inputs by processing video and audio modalities."""
+        video_result = self._process_modality(self.config.video, batch, "video", timestep_sampler)
+        audio_result = self._process_modality(self.config.audio, batch, "audio", timestep_sampler)
+
+        return ModelInputs(
+            video=video_result.modality if video_result else None,
+            audio=audio_result.modality if audio_result else None,
+            video_targets=video_result.targets if video_result else None,
+            audio_targets=audio_result.targets if audio_result else None,
+            video_loss_mask=video_result.loss_mask if video_result else None,
+            audio_loss_mask=audio_result.loss_mask if audio_result else None,
+        )
+
+    def compute_loss(
+        self,
+        video_pred: Tensor | None,
+        audio_pred: Tensor | None,
+        inputs: ModelInputs,
+    ) -> Tensor:
+        """Compute masked MSE loss for video and audio predictions. Returns [B,]."""
+        total_loss = None
+
+        if video_pred is not None and inputs.video_targets is not None:
+            video_loss = self._compute_modality_loss(
+                pred=video_pred,
+                targets=inputs.video_targets,
+                loss_mask=inputs.video_loss_mask,
+            )
+            total_loss = video_loss
+
+        if audio_pred is not None and inputs.audio_targets is not None:
+            audio_loss = self._compute_modality_loss(
+                pred=audio_pred,
+                targets=inputs.audio_targets,
+                loss_mask=inputs.audio_loss_mask,
+            )
+            total_loss = audio_loss if total_loss is None else total_loss + audio_loss
+
+        if total_loss is None:
+            raise ValueError("No valid predictions and targets provided for loss computation")
+
+        return total_loss
+
+    def get_checkpoint_metadata(self) -> dict[str, Any]:
+        """Include reference scale factors in checkpoint metadata for inference pipelines."""
+        metadata: dict[str, Any] = {}
+        spatial = self.reference_spatial_scale_factor
+        temporal = self.reference_temporal_scale_factor
+        if spatial is not None and spatial != 1:
+            metadata["reference_spatial_scale_factor"] = spatial
+            metadata["reference_downscale_factor"] = spatial  # backward compat
+        if temporal is not None and temporal != 1:
+            metadata["reference_temporal_scale_factor"] = temporal
+        return metadata
+
+    def _infer_reference_scale_factors_from_config(self) -> tuple[int | None, int | None]:
+        """Infer spatial and temporal scale factors by peeking at one sample pair."""
+        if self.config.video is None:
+            return None, None
+        for cond in self.config.video.conditions:
+            if not isinstance(cond, ReferenceConditionConfig):
+                continue
+            target_dir = Path(self.config.video.latents_dir)
+            ref_dir = Path(cond.latents_dir)
+            for sample_file in target_dir.rglob("*.pt"):
+                ref_file = ref_dir / sample_file.relative_to(target_dir)
+                if not ref_file.exists():
+                    continue
+                target_data = torch.load(sample_file, map_location="cpu", weights_only=True)
+                ref_data = torch.load(ref_file, map_location="cpu", weights_only=True)
+                if "height" not in ref_data or "height" not in target_data:
+                    continue
+                spatial = self._infer_scale_factor(
+                    ref_data["height"],
+                    ref_data["width"],
+                    target_data["height"],
+                    target_data["width"],
+                )
+                temporal = self._infer_temporal_scale_factor(
+                    ref_data["num_frames"],
+                    target_data["num_frames"],
+                )
+                return spatial, temporal
+        return None, None
+
+    def _process_modality(
+        self,
+        modality_config: ModalityConfig | None,
+        batch: dict[str, Any],
+        modality_key: str,
+        timestep_sampler: TimestepSampler,
+    ) -> ModalityProcessingResult | None:
+        """Process a single modality: load latents, add noise, apply conditions, build Modality."""
+        if modality_config is None:
+            return None
+
+        # Step 1: Load and patchify latents
+        data = self._patchify_latent_data(batch[f"{modality_key}_latents"], modality_key)
+        latents = data.latents
+
+        batch_size, seq_len, _ = latents.shape
+        device = latents.device
+        dtype = latents.dtype
+
+        # Step 2: Get text embeddings
+        conditions = batch["conditions"]
+        prompt_embeds = conditions[f"{modality_key}_prompt_embeds"]
+        prompt_attention_mask = conditions["prompt_attention_mask"]
+
+        # Step 3: Initialize noise, timesteps, and loss mask based on is_generated flag
+        if modality_config.is_generated:
+            noisy_latents, targets, timesteps, loss_mask, sigmas = self._initialize_noisy_target(
+                latents, timestep_sampler
+            )
+        else:
+            # Conditioning modality: keep clean (sigma=0), no loss
+            noisy_latents = latents
+            targets = None
+            timesteps = torch.zeros(batch_size, seq_len, device=device, dtype=dtype)
+            loss_mask = None
+            sigmas = torch.zeros(batch_size, device=device, dtype=dtype)
+
+        # Step 4: Generate positions
+        if modality_key == "video":
+            positions = self._get_video_positions(
+                num_frames=data.num_frames,
+                height=data.height,
+                width=data.width,
+                batch_size=batch_size,
+                fps=data.fps,
+                device=device,
+            )
+        else:
+            positions = self._get_audio_positions(
+                num_time_steps=seq_len,
+                batch_size=batch_size,
+                device=device,
+            )
+
+        # Step 5: Apply conditions (intrinsic first, then extrinsic)
+        for cond in modality_config.conditions:
+            if isinstance(cond, IntrinsicConditionBase) and modality_config.is_generated:
+                noisy_latents, timesteps, loss_mask = self._apply_intrinsic_condition(
+                    noisy_latents=noisy_latents,
+                    clean_latents=latents,
+                    timesteps=timesteps,
+                    loss_mask=loss_mask,
+                    config=cond,
+                    height=data.height,
+                    width=data.width,
+                    batch=batch,
+                )
+
+        for cond in modality_config.conditions:
+            if isinstance(cond, ReferenceConditionConfig):
+                noisy_latents, positions, timesteps, loss_mask, targets = self._apply_reference_condition(
+                    noisy_latents=noisy_latents,
+                    positions=positions,
+                    timesteps=timesteps,
+                    loss_mask=loss_mask,
+                    targets=targets,
+                    batch=batch,
+                    config=cond,
+                    modality_key=modality_key,
+                )
+
+        # Step 6: Build Modality
+        modality = Modality(
+            enabled=True,
+            latent=noisy_latents,
+            sigma=sigmas,
+            timesteps=timesteps,
+            positions=positions,
+            context=prompt_embeds,
+            context_mask=prompt_attention_mask,
+        )
+
+        return ModalityProcessingResult(
+            modality=modality,
+            targets=targets,
+            loss_mask=loss_mask,
+        )
+
+    @staticmethod
+    def _initialize_noisy_target(
+        latents: Tensor,
+        timestep_sampler: TimestepSampler,
+    ) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+        """Add noise to latents and create training targets. Returns (noisy, targets, timesteps, mask, sigmas)."""
+        batch_size, seq_len, _ = latents.shape
+        sigmas = timestep_sampler.sample_for(latents)
+        noise = torch.randn_like(latents)
+        sigmas_expanded = sigmas.view(-1, 1, 1)
+        noisy_latents = (1 - sigmas_expanded) * latents + sigmas_expanded * noise
+        targets = noise - latents  # velocity prediction
+        timesteps = sigmas.view(-1, 1).expand(batch_size, seq_len).clone()
+        loss_mask = torch.ones(batch_size, seq_len, dtype=torch.bool, device=latents.device)
+        return noisy_latents, targets, timesteps, loss_mask, sigmas
+
+    def _apply_intrinsic_condition(
+        self,
+        noisy_latents: Tensor,
+        clean_latents: Tensor,
+        timesteps: Tensor,
+        loss_mask: Tensor,
+        config: IntrinsicConditionBase,
+        height: int,
+        width: int,
+        batch: dict[str, Any],
+    ) -> tuple[Tensor, Tensor, Tensor]:
+        """Apply intrinsic conditioning using a binary mask.
+        For each token, the mask value determines conditioning strength:
+        - mask=1: conditioned (clean latent, timestep=0, excluded from loss)
+        - mask=0: generated (noisy latent, original timestep, contributes to loss)
+        The conditioning decision is drawn independently per batch element so the training
+        signal across samples in a batch is i.i.d. -- a single batch-wide draw would
+        correlate gradient updates across the batch.
+        """
+        batch_size, seq_len, _ = noisy_latents.shape
+        device = noisy_latents.device
+
+        # Per-sample Bernoulli draw -- each element is independently conditioned.
+        apply_per_sample = torch.rand(batch_size, device=device) < config.probability
+        if not apply_per_sample.any():
+            return noisy_latents, timesteps, loss_mask
+
+        if isinstance(config, FirstFrameConditionConfig):
+            mask = self._compute_temporal_mask(batch_size, seq_len, height, width, 1, False, device)
+        elif isinstance(config, PrefixConditionConfig):
+            mask = self._compute_temporal_mask(
+                batch_size, seq_len, height, width, config.temporal_boundary, False, device
+            )
+        elif isinstance(config, SuffixConditionConfig):
+            mask = self._compute_temporal_mask(
+                batch_size, seq_len, height, width, config.temporal_boundary, True, device
+            )
+        elif isinstance(config, SpatialCropConditionConfig):
+            mask = self._compute_spatial_crop_mask(batch_size, seq_len, height, width, config.spatial_region, device)
+        elif isinstance(config, MaskConditionConfig):
+            # Binarize to match inference, which thresholds masks at load time
+            # (validation_runner._load_and_downsample_mask / _load_audio_mask).
+            mask = (batch[config.mask_dir]["mask"].reshape(batch_size, seq_len) > 0.5).float()
+        else:
+            raise ValueError(f"Unknown intrinsic condition type: {type(config).__name__}")
+
+        # Zero the mask for samples the per-sample draw did not select.
+        mask = mask * apply_per_sample.view(-1, 1).to(mask.dtype)
+
+        # Apply binary mask: clean conditioned tokens, noisy generated tokens.
+        m = mask.unsqueeze(-1)
+        noisy_latents = m * clean_latents + (1 - m) * noisy_latents
+        timesteps = (1 - mask) * timesteps
+        loss_mask = loss_mask & (mask == 0)
+
+        return noisy_latents, timesteps, loss_mask
+
+    @staticmethod
+    def _compute_temporal_mask(
+        batch_size: int,
+        seq_len: int,
+        height: int,
+        width: int,
+        num_frames: int,
+        from_end: bool,
+        device: torch.device,
+    ) -> Tensor:
+        """Compute float mask for temporal region (prefix or suffix). Returns [B, seq_len] in {0, 1}."""
+        tokens_per_frame = height * width
+        num_tokens = num_frames * tokens_per_frame
+        mask = torch.zeros(batch_size, seq_len, device=device)
+        if from_end:
+            mask[:, -num_tokens:] = 1.0
+        else:
+            mask[:, :num_tokens] = 1.0
+        return mask
+
+    @staticmethod
+    def _compute_spatial_crop_mask(
+        batch_size: int,
+        seq_len: int,
+        height: int,
+        width: int,
+        region: tuple[int, int, int, int],
+        device: torch.device,
+    ) -> Tensor:
+        """Compute float mask for spatial crop region (y1, x1, y2, x2) in pixel coords.
+        Returns [B, seq_len] in {0, 1}.
+        """
+        y1, x1, y2, x2 = region
+        num_frames = seq_len // (height * width)
+
+        # Convert pixel to latent coordinates and clamp (per-axis VAE scale factor).
+        def to_latent(v: int, scale: int, max_v: int) -> int:
+            return max(0, min(v // scale, max_v))
+
+        ly1 = to_latent(y1, VIDEO_SCALE_FACTORS.height, height)
+        ly2 = to_latent(y2, VIDEO_SCALE_FACTORS.height, height)
+        lx1 = to_latent(x1, VIDEO_SCALE_FACTORS.width, width)
+        lx2 = to_latent(x2, VIDEO_SCALE_FACTORS.width, width)
+
+        # Create spatial mask and tile across frames
+        spatial_mask = torch.zeros(height, width, device=device)
+        spatial_mask[ly1:ly2, lx1:lx2] = 1.0
+        full_mask = spatial_mask.flatten().repeat(num_frames)
+
+        return full_mask.unsqueeze(0).expand(batch_size, -1)
+
+    def _patchify_latent_data(self, latent_data: dict[str, Any], modality_key: str) -> LatentData:
+        """Patchify latent data and extract metadata."""
+        latents = latent_data["latents"]
+
+        if modality_key == "video":
+            num_frames = latent_data["num_frames"][0].item()
+            height = latent_data["height"][0].item()
+            width = latent_data["width"][0].item()
+            fps = latent_data.get("fps")
+            fps = fps[0].item() if fps is not None else DEFAULT_FPS
+            latents = self._video_patchifier.patchify(latents)
+        else:
+            num_frames = latent_data.get("num_frames", [latents.shape[2]])[0]
+            if isinstance(num_frames, Tensor):
+                num_frames = num_frames.item()
+            height = 1
+            width = 1
+            fps = 1.0
+            latents = self._audio_patchifier.patchify(latents)
+
+        return LatentData(latents=latents, num_frames=num_frames, height=height, width=width, fps=fps)
+
+    def _apply_reference_condition(
+        self,
+        noisy_latents: Tensor,
+        positions: Tensor,
+        timesteps: Tensor,
+        loss_mask: Tensor | None,
+        targets: Tensor | None,
+        batch: dict[str, Any],
+        config: ReferenceConditionConfig,
+        modality_key: str,
+    ) -> tuple[Tensor, Tensor, Tensor, Tensor | None, Tensor | None]:
+        """Concatenate reference latents to target sequence for reference conditioning (IC-LoRA style).
+        The apply/skip decision is batch-wide (reference conditioning changes the sequence
+        length, so it cannot be applied to only part of a batch) but is drawn from the torch
+        RNG so runs are reproducible under ``torch.manual_seed`` — mirroring the intrinsic
+        per-sample draw rather than Python's unseeded ``random``.
+        """
+        if torch.rand((), device=noisy_latents.device).item() >= config.probability:
+            return noisy_latents, positions, timesteps, loss_mask, targets
+
+        # Load and patchify condition latents
+        cond = self._patchify_latent_data(batch[config.latents_dir], modality_key)
+        cond_latents = cond.latents
+
+        batch_size, cond_seq_len, _ = cond_latents.shape
+        device = cond_latents.device
+        dtype = cond_latents.dtype
+
+        # Generate condition positions
+        if modality_key == "video":
+            cond_positions = self._get_video_positions(
+                num_frames=cond.num_frames,
+                height=cond.height,
+                width=cond.width,
+                batch_size=batch_size,
+                fps=cond.fps,
+                device=device,
+            )
+        else:
+            cond_positions = self._get_audio_positions(
+                num_time_steps=cond_seq_len,
+                batch_size=batch_size,
+                device=device,
+            )
+
+        # Translate / rescale ref positions into the target's frame (video only).
+        if modality_key == "video":
+            spatial_sf = self.reference_spatial_scale_factor or 1
+            temporal_sf = self.reference_temporal_scale_factor or 1
+            if spatial_sf != 1 or temporal_sf != 1:
+                cond_positions = cond_positions.clone()
+                if temporal_sf != 1:
+                    # Ref positions are already at the ref's effective fps (source_fps / S,
+                    # stored by process_videos.py). Shift by (S - 1) / target_fps so ref's
+                    # last patch aligns with target's last; clamp the causal patch at 0.
+                    t_target = positions[:, 0, 0:1, 1:2]  # = 1 / target_fps
+                    cond_positions[:, 0, ...] = torch.clamp(
+                        cond_positions[:, 0, ...] - (temporal_sf - 1) * t_target, min=0
+                    )
+                if spatial_sf != 1:
+                    cond_positions[:, 1, ...] *= spatial_sf
+                    cond_positions[:, 2, ...] *= spatial_sf
+
+        # Condition tokens: clean, timestep=0, no loss
+        cond_timesteps = torch.zeros(batch_size, cond_seq_len, device=device, dtype=dtype)
+        cond_loss_mask = torch.zeros(batch_size, cond_seq_len, dtype=torch.bool, device=device)
+
+        # Concatenate condition and target sequences (condition first, then target)
+        combined_latents = torch.cat([cond_latents, noisy_latents], dim=1)
+        combined_positions = torch.cat([cond_positions, positions], dim=2)
+        combined_timesteps = torch.cat([cond_timesteps, timesteps], dim=1)
+
+        combined_loss_mask = torch.cat([cond_loss_mask, loss_mask], dim=1) if loss_mask is not None else None
+
+        # Targets remain unchanged (only for target portion, not condition portion)
+
+        return combined_latents, combined_positions, combined_timesteps, combined_loss_mask, targets
+
+    @staticmethod
+    def _compute_modality_loss(pred: Tensor, targets: Tensor, loss_mask: Tensor) -> Tensor:
+        """Compute per-element MSE loss for a single modality. Returns [B,]."""
+        # Slice prediction to match targets length (removes any prepended condition tokens)
+        target_len = targets.shape[1]
+        pred = pred[:, -target_len:, :]
+        mask = loss_mask[:, -target_len:]
+
+        # Compute masked MSE loss, reduce per-element [B,] over (seq, channels)
+        mask_expanded = mask.unsqueeze(-1).float()
+        squared_error = (pred - targets).pow(2)
+        masked_loss = squared_error * mask_expanded
+        return masked_loss.mean(dim=[-2, -1]) / mask_expanded.mean(dim=[-2, -1]).clamp(min=1e-8)
+
+    @staticmethod
+    def _infer_scale_factor(cond_height: int, cond_width: int, target_height: int, target_width: int) -> int:
+        """Infer spatial scale factor between condition and target resolutions."""
+        if target_height == cond_height and target_width == cond_width:
+            return 1
+        scale_h = target_height // cond_height if cond_height > 0 else 1
+        scale_w = target_width // cond_width if cond_width > 0 else 1
+        if scale_h != scale_w:
+            raise ValueError(
+                f"Non-uniform scale factors between condition and target: height={scale_h}, width={scale_w}. "
+                "Condition and target resolutions must scale uniformly."
+            )
+        return scale_h
+
+    @staticmethod
+    def _infer_temporal_scale_factor(cond_num_frames: int, target_num_frames: int) -> int:
+        """Infer temporal scale factor between condition and target latent frame counts.
+        The first latent frame encodes a single pixel frame (the VAE's causal structure),
+        so the temporal groups count is (num_frames - 1).  The scale factor is the ratio
+        of target groups to condition groups.
+        """
+        if target_num_frames == cond_num_frames:
+            return 1
+        target_groups = target_num_frames - 1
+        cond_groups = cond_num_frames - 1
+        if cond_groups <= 0 or target_groups <= 0:
+            return 1
+        if target_groups % cond_groups != 0:
+            raise ValueError(
+                f"Target temporal groups ({target_groups}) is not evenly divisible by "
+                f"condition temporal groups ({cond_groups})."
+            )
+        return target_groups // cond_groups

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

@@ -45,6 +45,18 @@ class TextToVideoConfig(TrainingStrategyConfigBase):
         description="Directory name for audio latents when with_audio is True",
     )
 
+    def get_data_sources(self) -> dict[str, str]:
+        """Text-to-video training requires latents and text conditions.
+        When ``with_audio`` is True, also requires audio latents.
+        """
+        sources = {
+            "latents": "latents",
+            "conditions": "conditions",
+        }
+        if self.with_audio:
+            sources[self.audio_latents_dir] = "audio_latents"
+        return sources
+
 
 class TextToVideoStrategy(TrainingStrategy):
     """Text-to-video training strategy.
@@ -64,26 +76,6 @@ class TextToVideoStrategy(TrainingStrategy):
         """
         super().__init__(config)
 
-    @property
-    def requires_audio(self) -> bool:
-        """Whether this training strategy requires audio components."""
-        return self.config.with_audio
-
-    def get_data_sources(self) -> list[str] | dict[str, str]:
-        """
-        Text-to-video training requires latents and text conditions.
-        When with_audio is True, also requires audio latents.
-        """
-        sources = {
-            "latents": "latents",
-            "conditions": "conditions",
-        }
-
-        if self.config.with_audio:
-            sources[self.config.audio_latents_dir] = "audio_latents"
-
-        return sources
-
     def prepare_training_inputs(
         self,
         batch: dict[str, Any],
@@ -119,7 +111,6 @@ class TextToVideoStrategy(TrainingStrategy):
         batch_size = video_latents.shape[0]
         video_seq_len = video_latents.shape[1]
         device = video_latents.device
-        dtype = video_latents.dtype
 
         # Create conditioning mask (first frame conditioning)
         video_conditioning_mask = self._create_first_frame_conditioning_mask(
@@ -157,7 +148,6 @@ class TextToVideoStrategy(TrainingStrategy):
             batch_size=batch_size,
             fps=fps,
             device=device,
-            dtype=dtype,
         )
 
         # Create video Modality
@@ -187,7 +177,6 @@ class TextToVideoStrategy(TrainingStrategy):
                 prompt_attention_mask=prompt_attention_mask,
                 batch_size=batch_size,
                 device=device,
-                dtype=dtype,
             )
 
         return ModelInputs(
@@ -207,7 +196,6 @@ class TextToVideoStrategy(TrainingStrategy):
         prompt_attention_mask: Tensor,
         batch_size: int,
         device: torch.device,
-        dtype: torch.dtype,
     ) -> tuple[Modality, Tensor, Tensor]:
         """Prepare audio inputs for joint audio-video training.
         Args:
@@ -217,7 +205,6 @@ class TextToVideoStrategy(TrainingStrategy):
             prompt_attention_mask: Attention mask for context
             batch_size: Batch size
             device: Target device
-            dtype: Target dtype
         Returns:
             Tuple of (audio_modality, audio_targets, audio_loss_mask)
         """
@@ -248,7 +235,6 @@ class TextToVideoStrategy(TrainingStrategy):
             num_time_steps=audio_seq_len,
             batch_size=batch_size,
             device=device,
-            dtype=dtype,
         )
 
         # Create audio Modality

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

@@ -39,6 +39,14 @@ class VideoToVideoConfig(TrainingStrategyConfigBase):
         description="Directory name for latents of reference videos",
     )
 
+    def get_data_sources(self) -> dict[str, str]:
+        """IC-LoRA training requires latents, conditions, and reference latents."""
+        return {
+            "latents": "latents",
+            "conditions": "conditions",
+            self.reference_latents_dir: "ref_latents",
+        }
+
 
 class VideoToVideoStrategy(TrainingStrategy):
     """Video-to-video training strategy for IC-LoRA.
@@ -62,14 +70,6 @@ class VideoToVideoStrategy(TrainingStrategy):
         super().__init__(config)
         self.reference_downscale_factor = None  # Will be inferred from first batch
 
-    def get_data_sources(self) -> dict[str, str]:
-        """IC-LoRA training requires latents, conditions, and reference latents."""
-        return {
-            "latents": "latents",
-            "conditions": "conditions",
-            self.config.reference_latents_dir: "ref_latents",
-        }
-
     def prepare_training_inputs(  # noqa: PLR0915
         self,
         batch: dict[str, Any],
@@ -133,7 +133,6 @@ class VideoToVideoStrategy(TrainingStrategy):
         ref_seq_len = ref_latents.shape[1]
         target_seq_len = target_latents.shape[1]
         device = target_latents.device
-        dtype = target_latents.dtype
 
         # Create conditioning mask
         # Reference tokens are always conditioning (timestep=0)
@@ -164,7 +163,7 @@ class VideoToVideoStrategy(TrainingStrategy):
         target_conditioning_mask_expanded = target_conditioning_mask.unsqueeze(-1)
         noisy_target = torch.where(target_conditioning_mask_expanded, target_latents, noisy_target)
 
-        # Targets for loss computation
+        # Targets for loss computation (velocity prediction) - only for target portion
         targets = noise - target_latents
 
         # Concatenate reference (clean) and target (noisy)
@@ -181,7 +180,6 @@ class VideoToVideoStrategy(TrainingStrategy):
             batch_size=batch_size,
             fps=fps,
             device=device,
-            dtype=dtype,
         )
 
         # Scale reference positions to match target coordinate space
@@ -200,7 +198,6 @@ class VideoToVideoStrategy(TrainingStrategy):
             batch_size=batch_size,
             fps=fps,
             device=device,
-            dtype=dtype,
         )
 
         # Concatenate positions along sequence dimension
@@ -231,7 +228,6 @@ class VideoToVideoStrategy(TrainingStrategy):
             audio_targets=None,
             video_loss_mask=video_loss_mask,
             audio_loss_mask=None,
-            ref_seq_len=ref_seq_len,
         )
 
     def compute_loss(
@@ -240,13 +236,11 @@ class VideoToVideoStrategy(TrainingStrategy):
         _audio_pred: Tensor | None,
         inputs: ModelInputs,
     ) -> Tensor:
-        """Compute masked loss only on target portion. Returns [B,]."""
-        # Extract target portion of prediction
-        ref_seq_len = inputs.ref_seq_len
-        target_pred = video_pred[:, ref_seq_len:, :]
-
-        # Get target portion of loss mask
-        target_loss_mask = inputs.video_loss_mask[:, ref_seq_len:]
+        """Compute masked loss on target portion only. Returns [B,]."""
+        # Slice prediction to match targets length (removes prepended reference tokens)
+        target_len = inputs.video_targets.shape[1]
+        target_pred = video_pred[:, -target_len:, :]
+        target_loss_mask = inputs.video_loss_mask[:, -target_len:]
 
         # Compute per-element loss [B,]
         loss = (target_pred - inputs.video_targets).pow(2)

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

@@ -1,874 +0,0 @@
-"""Validation sampling for LTX-2 training using ltx-core components.
-This module provides a simplified validation pipeline for generating samples during training,
-using the new ltx-core components (VideoLatentTools, AudioLatentTools, LatentState, etc.).
-"""
-
-from dataclasses import dataclass, replace
-from typing import TYPE_CHECKING, Literal
-
-import torch
-from einops import rearrange
-from torch import Tensor
-
-from ltx_core.components.diffusion_steps import EulerDiffusionStep
-from ltx_core.components.guiders import CFGGuider, STGGuider
-from ltx_core.components.noisers import GaussianNoiser
-from ltx_core.components.patchifiers import (
-    AudioPatchifier,
-    VideoLatentPatchifier,
-    get_pixel_coords,
-)
-from ltx_core.components.schedulers import LTX2Scheduler
-from ltx_core.guidance.perturbations import (
-    BatchedPerturbationConfig,
-    Perturbation,
-    PerturbationConfig,
-    PerturbationType,
-)
-from ltx_core.model.transformer.modality import Modality
-from ltx_core.model.transformer.model import X0Model
-from ltx_core.model.video_vae import SpatialTilingConfig, TemporalTilingConfig, TilingConfig
-from ltx_core.tools import AudioLatentTools, VideoLatentTools
-from ltx_core.types import AudioLatentShape, LatentState, SpatioTemporalScaleFactors, VideoLatentShape, VideoPixelShape
-from ltx_trainer.progress import SamplingContext
-
-if TYPE_CHECKING:
-    from ltx_core.model.audio_vae import AudioDecoder, Vocoder
-    from ltx_core.model.transformer import LTXModel
-    from ltx_core.model.video_vae import VideoDecoder, VideoEncoder
-    from ltx_core.text_encoders.gemma import GemmaTextEncoder
-    from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessor
-
-VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
-
-
-@dataclass
-class CachedPromptEmbeddings:
-    """Pre-computed text embeddings for a validation prompt.
-    These embeddings are computed once at training start and reused for all validation runs,
-    avoiding the need to load the full Gemma text encoder during validation.
-    """
-
-    video_context_positive: Tensor  # [1, seq_len, hidden_dim]
-    audio_context_positive: Tensor  # [1, seq_len, hidden_dim]
-    video_context_negative: Tensor | None = None
-    audio_context_negative: Tensor | None = None
-
-
-@dataclass
-class TiledDecodingConfig:
-    """Configuration for tiled video decoding to reduce VRAM usage.
-    Tiled decoding splits the latent tensor into overlapping tiles, decodes each
-    tile individually, and blends them together. This significantly reduces peak
-    VRAM usage at the cost of slightly slower decoding.
-    Defaults match the recommended values from ltx-core tests.
-    """
-
-    enabled: bool = True  # Whether to use tiled decoding (enabled by default)
-    tile_size_pixels: int = 192  # Spatial tile size in pixels (must be ≥64 and divisible by 32)
-    tile_overlap_pixels: int = 64  # Spatial tile overlap in pixels (must be divisible by 32)
-    tile_size_frames: int = 48  # Temporal tile size in frames (must be ≥16 and divisible by 8)
-    tile_overlap_frames: int = 24  # Temporal tile overlap in frames (must be divisible by 8)
-
-
-@dataclass
-class GenerationConfig:
-    """Configuration for video/audio generation."""
-
-    prompt: str  # Text prompt for generation
-    negative_prompt: str = ""  # Negative prompt to avoid unwanted artifacts
-    height: int = 544  # Output video height in pixels
-    width: int = 960  # Output video width in pixels
-    num_frames: int = 97  # Number of frames to generate
-    frame_rate: float = 25.0  # Frame rate for temporal position scaling
-    num_inference_steps: int = 30  # Number of denoising steps
-    guidance_scale: float = 4.0  # CFG guidance scale
-    seed: int = 42  # Random seed for reproducibility
-    condition_image: Tensor | None = None  # Optional first frame image for image-to-video
-    reference_video: Tensor | None = None  # For IC-LoRA: [F, C, H, W] in [0, 1]
-    reference_downscale_factor: int = 1  # For IC-LoRA: downscale factor (1 = same resolution, 2 = half resolution)
-    generate_audio: bool = True  # Whether to generate audio alongside video
-    include_reference_in_output: bool = False  # For IC-LoRA: concatenate original reference with generated output
-    cached_embeddings: CachedPromptEmbeddings | None = None  # Pre-computed text embeddings (avoids loading Gemma)
-    stg_scale: float = 0.0  # STG strength (0.0 = disabled, recommended: 1.0)
-    stg_blocks: list[int] | None = None  # Transformer blocks to perturb (None = all, recommended: [29])
-    stg_mode: Literal["stg_av", "stg_v"] = "stg_av"  # STG mode: "stg_av" (audio+video) or "stg_v" (video only)
-    # Tiled decoding config: None = use defaults (enabled), False = disable, or TiledDecodingConfig for custom settings
-    tiled_decoding: TiledDecodingConfig | Literal[False] | None = None
-
-    def __post_init__(self) -> None:
-        """Apply default tiled decoding config if not provided."""
-        if self.tiled_decoding is None:
-            # Use default config with tiling enabled
-            object.__setattr__(self, "tiled_decoding", TiledDecodingConfig())
-        elif self.tiled_decoding is False:
-            # Explicitly disabled - use config with enabled=False
-            object.__setattr__(self, "tiled_decoding", TiledDecodingConfig(enabled=False))
-
-
-class ValidationSampler:
-    """Generates validation samples during training using ltx-core components.
-    This class provides a simplified interface for generating video (and optionally audio)
-    samples during training validation. It supports:
-    - Text-to-video generation
-    - Image-to-video generation (first frame conditioning)
-    - Video-to-video generation (IC-LoRA reference video conditioning)
-    - Optional audio generation
-    The implementation follows the patterns from ltx_pipelines.single_stage.
-    Text embeddings can be provided either via:
-    - A full text_encoder (encodes prompts on-the-fly)
-    - Pre-computed cached_embeddings (avoids loading Gemma during validation)
-    """
-
-    def __init__(
-        self,
-        transformer: "LTXModel",
-        vae_decoder: "VideoDecoder",
-        vae_encoder: "VideoEncoder | None",
-        text_encoder: "GemmaTextEncoder | None" = None,
-        audio_decoder: "AudioDecoder | None" = None,
-        vocoder: "Vocoder | None" = None,
-        sampling_context: SamplingContext | None = None,
-        embeddings_processor: "EmbeddingsProcessor | None" = None,
-    ):
-        """Initialize the validation sampler.
-        Args:
-            transformer: LTX-2 transformer model
-            vae_decoder: Video VAE decoder
-            vae_encoder: Video VAE encoder (for image/video conditioning), can be None if not needed
-            text_encoder: Gemma text encoder (optional if cached_embeddings in config)
-            audio_decoder: Optional audio VAE decoder (for audio generation)
-            vocoder: Optional vocoder (for audio generation)
-            sampling_context: Optional SamplingContext for progress display during denoising
-            embeddings_processor: Optional embeddings processor (required if text_encoder provided)
-        """
-        self._transformer = transformer
-        self._vae_decoder = vae_decoder
-        self._vae_encoder = vae_encoder
-        self._text_encoder = text_encoder
-        self._embeddings_processor = embeddings_processor
-        self._audio_decoder = audio_decoder
-        self._vocoder = vocoder
-        self._sampling_context = sampling_context
-
-        # Patchifiers
-        self._video_patchifier = VideoLatentPatchifier(patch_size=1)
-        self._audio_patchifier = AudioPatchifier(patch_size=1)
-
-    # Note: Use @torch.no_grad() instead of @torch.inference_mode() to avoid FSDP inplace update errors after validation
-    @torch.no_grad()
-    def generate(
-        self,
-        config: GenerationConfig,
-        device: torch.device | str = "cuda",
-    ) -> tuple[Tensor, Tensor | None]:
-        """Generate a video (and optionally audio) sample.
-        Args:
-            config: Generation configuration
-            device: Device to run generation on
-        Returns:
-            Tuple of:
-                - video: Video tensor [C, F, H, W] in [0, 1] (float32)
-                - audio: Audio waveform tensor [C, samples] or None
-        """
-        device = torch.device(device) if isinstance(device, str) else device
-        self._validate_config(config)
-
-        # Route to appropriate generation method
-        if config.reference_video is not None:
-            return self._generate_with_reference(config, device)
-        return self._generate_standard(config, device)
-
-    def _generate_standard(self, config: GenerationConfig, device: torch.device) -> tuple[Tensor, Tensor | None]:
-        """Standard generation (text-to-video or image-to-video)."""
-        # Get prompt embeddings (from cache or encode on-the-fly)
-        v_ctx_pos, a_ctx_pos, v_ctx_neg, a_ctx_neg = self._get_prompt_embeddings(config, device)
-
-        # Setup generator
-        generator = torch.Generator(device=device).manual_seed(config.seed)
-
-        # Create latent tools
-        video_tools = self._create_video_latent_tools(config)
-        audio_tools = self._create_audio_latent_tools(config) if config.generate_audio else None
-
-        # Create initial states
-        video_clean_state = video_tools.create_initial_state(device=device, dtype=torch.bfloat16)
-        audio_clean_state = (
-            audio_tools.create_initial_state(device=device, dtype=torch.bfloat16) if audio_tools else None
-        )
-
-        # Apply image conditioning if provided
-        if config.condition_image is not None:
-            video_clean_state = self._apply_image_conditioning(
-                video_clean_state, config.condition_image, config, device
-            )
-
-        # Add noise
-        noiser = GaussianNoiser(generator=generator)
-        video_state = noiser(latent_state=video_clean_state, noise_scale=1.0)
-        audio_state = noiser(latent_state=audio_clean_state, noise_scale=1.0) if audio_clean_state else None
-
-        # Run denoising loop
-        video_state, audio_state = self._run_denoising(
-            config=config,
-            video_state=video_state,
-            audio_state=audio_state,
-            video_clean_state=video_clean_state,
-            audio_clean_state=audio_clean_state,
-            v_ctx_pos=v_ctx_pos,
-            a_ctx_pos=a_ctx_pos,
-            v_ctx_neg=v_ctx_neg,
-            a_ctx_neg=a_ctx_neg,
-            device=device,
-        )
-
-        # Decode outputs
-        video_state = video_tools.clear_conditioning(video_state)
-        video_state = video_tools.unpatchify(video_state)
-        video_output = self._decode_video(video_state, device, config.tiled_decoding)
-
-        audio_output = None
-        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)
-            audio_output = self._decode_audio(audio_state, device)
-
-        return video_output, audio_output
-
-    def _generate_with_reference(self, config: GenerationConfig, device: torch.device) -> tuple[Tensor, Tensor | None]:
-        """Generate with reference video conditioning (IC-LoRA style).
-        For IC-LoRA:
-        - Reference video latents are concatenated with target latents
-        - Reference latents have timestep=0 (clean, not denoised)
-        - Target latents are denoised normally
-        - If condition_image is also provided, the first frame of the target is conditioned
-        - If include_reference_in_output is True, the preprocessed reference video
-          is concatenated side-by-side with the generated video
-        """
-        # Get prompt embeddings (from cache or encode on-the-fly)
-        v_ctx_pos, a_ctx_pos, v_ctx_neg, a_ctx_neg = self._get_prompt_embeddings(config, device)
-
-        # Setup generator
-        generator = torch.Generator(device=device).manual_seed(config.seed)
-
-        # Preprocess and encode reference video
-        ref_video_preprocessed = self._preprocess_reference_video(config)
-        ref_latent, ref_positions = self._encode_video(ref_video_preprocessed, config.frame_rate, device)
-        ref_seq_len = ref_latent.shape[1]
-
-        # Scale reference positions to match target coordinate space
-        # Position tensor shape: [B, 3, seq_len, 2] where dim 1 is (time, height, width)
-        if config.reference_downscale_factor != 1:
-            ref_positions = ref_positions.clone()
-            ref_positions[:, 1, ...] *= config.reference_downscale_factor  # height axis
-            ref_positions[:, 2, ...] *= config.reference_downscale_factor  # width axis
-            # Time axis (index 0) remains unchanged
-
-        # Create target video state
-        video_tools = self._create_video_latent_tools(config)
-        target_clean_state = video_tools.create_initial_state(device=device, dtype=torch.bfloat16)
-
-        # Apply first-frame image conditioning to target if provided
-        if config.condition_image is not None:
-            target_clean_state = self._apply_image_conditioning(
-                target_clean_state, config.condition_image, config, device
-            )
-
-        # Create combined state (reference + target)
-        # denoise_mask shape is [B, seq_len, 1] after patchification
-        ref_denoise_mask = torch.zeros(1, ref_seq_len, 1, device=device, dtype=torch.float32)
-        combined_clean_state = LatentState(
-            latent=torch.cat([ref_latent, target_clean_state.latent], dim=1),
-            denoise_mask=torch.cat([ref_denoise_mask, target_clean_state.denoise_mask], dim=1),
-            positions=torch.cat([ref_positions, target_clean_state.positions], dim=2),
-            clean_latent=torch.cat([ref_latent, target_clean_state.clean_latent], dim=1),
-        )
-
-        # Add noise (only to the target portion via denoise_mask)
-        noiser = GaussianNoiser(generator=generator)
-        combined_state = noiser(latent_state=combined_clean_state, noise_scale=1.0)
-
-        # Create audio state if needed
-        audio_tools = self._create_audio_latent_tools(config) if config.generate_audio else None
-        audio_clean_state = (
-            audio_tools.create_initial_state(device=device, dtype=torch.bfloat16) if audio_tools else None
-        )
-        audio_state = noiser(latent_state=audio_clean_state, noise_scale=1.0) if audio_clean_state else None
-
-        # Run denoising loop
-        combined_state, audio_state = self._run_denoising(
-            config=config,
-            video_state=combined_state,
-            audio_state=audio_state,
-            video_clean_state=combined_clean_state,
-            audio_clean_state=audio_clean_state,
-            v_ctx_pos=v_ctx_pos,
-            a_ctx_pos=a_ctx_pos,
-            v_ctx_neg=v_ctx_neg,
-            a_ctx_neg=a_ctx_neg,
-            device=device,
-        )
-
-        # Extract target portion and decode
-        target_latent = combined_state.latent[:, ref_seq_len:]
-        video_output = self._decode_video_latent(target_latent, config, device)
-
-        # Optionally concatenate original reference video side-by-side
-        if config.include_reference_in_output:
-            # Use preprocessed reference (already resized/cropped, in pixel space)
-            # Convert from [B, C, F, H, W] to [C, F, H, W]
-            ref_video_pixels = ref_video_preprocessed[0].cpu()
-            # Normalize from [-1, 1] to [0, 1]
-            ref_video_pixels = ((ref_video_pixels + 1.0) / 2.0).clamp(0.0, 1.0)
-            video_output = self._concatenate_videos_side_by_side(ref_video_pixels, video_output)
-
-        # Decode audio
-        audio_output = None
-        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)
-            audio_output = self._decode_audio(audio_state, device)
-
-        return video_output, audio_output
-
-    def _create_video_latent_tools(self, config: GenerationConfig) -> VideoLatentTools:
-        """Create video latent tools for the given configuration."""
-        pixel_shape = VideoPixelShape(
-            batch=1,
-            frames=config.num_frames,
-            height=config.height,
-            width=config.width,
-            fps=config.frame_rate,
-        )
-        return VideoLatentTools(
-            patchifier=self._video_patchifier,
-            target_shape=VideoLatentShape.from_pixel_shape(shape=pixel_shape),
-            fps=config.frame_rate,
-            scale_factors=VIDEO_SCALE_FACTORS,
-            causal_fix=True,
-        )
-
-    def _create_audio_latent_tools(self, config: GenerationConfig) -> AudioLatentTools:
-        """Create audio latent tools for the given configuration."""
-        return AudioLatentTools(
-            patchifier=self._audio_patchifier,
-            target_shape=AudioLatentShape.from_duration(batch=1, duration=config.num_frames / config.frame_rate),
-        )
-
-    def _apply_image_conditioning(
-        self, video_state: LatentState, image: Tensor, config: GenerationConfig, device: torch.device
-    ) -> LatentState:
-        """Apply first-frame image conditioning to the video state."""
-        # Encode the image
-        encoded_image = self._encode_conditioning_image(image, config.height, config.width, device)
-
-        # Patchify the encoded image (single frame)
-        patchified_image = self._video_patchifier.patchify(encoded_image)  # [1, 1, C] -> [1, num_patches, C]
-        num_image_tokens = patchified_image.shape[1]
-
-        # Update the first frame tokens in the latent
-        new_latent = video_state.latent.clone()
-        new_latent[:, :num_image_tokens] = patchified_image.to(new_latent.dtype)
-
-        # Update clean_latent as well (conditioning image is clean)
-        new_clean_latent = video_state.clean_latent.clone()
-        new_clean_latent[:, :num_image_tokens] = patchified_image.to(new_clean_latent.dtype)
-
-        # Set denoise_mask to 0 for conditioned tokens (don't denoise them)
-        new_denoise_mask = video_state.denoise_mask.clone()
-        new_denoise_mask[:, :num_image_tokens] = 0.0
-
-        return LatentState(
-            latent=new_latent,
-            denoise_mask=new_denoise_mask,
-            positions=video_state.positions,
-            clean_latent=new_clean_latent,
-        )
-
-    @staticmethod
-    def _preprocess_reference_video(config: GenerationConfig) -> Tensor:
-        """Preprocess reference video: resize, crop, and convert to model input format.
-        When reference_downscale_factor > 1, the reference video is downscaled to a smaller
-        resolution for more efficient inference. The positions will be scaled up later
-        to match the target coordinate space.
-        Args:
-            config: Generation configuration
-        Returns:
-            Preprocessed video tensor [B, C, F, H, W] in [-1, 1] range
-        """
-        ref_video = config.reference_video  # [F, C, H, W] in [0, 1]
-        scale_factor = config.reference_downscale_factor
-
-        # Target dimensions for reference (scaled down if scale_factor > 1)
-        target_height = config.height // scale_factor
-        target_width = config.width // scale_factor
-
-        # Validate scaled dimensions
-        if target_height % 32 != 0 or target_width % 32 != 0:
-            raise ValueError(
-                f"Scaled reference dimensions ({target_height}x{target_width}) must be divisible by 32. "
-                f"Original: {config.height}x{config.width}, scale_factor: {scale_factor}"
-            )
-
-        current_height, current_width = ref_video.shape[2:]
-
-        # Resize maintaining aspect ratio and center crop if needed
-        if current_height != target_height or current_width != target_width:
-            aspect_ratio = current_width / current_height
-            target_aspect_ratio = target_width / target_height
-
-            if aspect_ratio > target_aspect_ratio:
-                resize_height, resize_width = target_height, int(target_height * aspect_ratio)
-            else:
-                resize_height, resize_width = int(target_width / aspect_ratio), target_width
-
-            ref_video = torch.nn.functional.interpolate(
-                ref_video, size=(resize_height, resize_width), mode="bilinear", align_corners=False
-            )
-
-            # Center crop
-            h_start = (resize_height - target_height) // 2
-            w_start = (resize_width - target_width) // 2
-            ref_video = ref_video[:, :, h_start : h_start + target_height, w_start : w_start + target_width]
-
-        # Convert to [B, C, F, H, W] and trim to valid frame count (k*8 + 1)
-        ref_video = rearrange(ref_video, "f c h w -> 1 c f h w")
-        valid_frames = (ref_video.shape[2] - 1) // 8 * 8 + 1
-        ref_video = ref_video[:, :, :valid_frames]
-
-        # Convert to [-1, 1] range
-        return ref_video * 2.0 - 1.0
-
-    def _encode_video(self, video: Tensor, fps: float, device: torch.device) -> tuple[Tensor, Tensor]:
-        """Encode video to patchified latents and compute positions.
-        Args:
-            video: Video tensor [B, C, F, H, W] in [-1, 1] range
-            fps: Frame rate for temporal position scaling
-            device: Device to run encoding on
-        Returns:
-            Tuple of (patchified_latents, positions)
-        """
-        video = video.to(device=device, dtype=torch.float32)
-
-        # Encode with VAE
-        self._vae_encoder.to(device)
-        with torch.autocast(device_type=str(device).split(":")[0], dtype=torch.bfloat16):
-            latents = self._vae_encoder(video)
-        self._vae_encoder.to("cpu")
-
-        latents = latents.to(torch.bfloat16)
-        patchified = self._video_patchifier.patchify(latents)
-
-        # Compute positions
-        latent_shape = VideoLatentShape(
-            batch=1,
-            channels=latents.shape[1],
-            frames=latents.shape[2],
-            height=latents.shape[3],
-            width=latents.shape[4],
-        )
-        latent_coords = self._video_patchifier.get_patch_grid_bounds(output_shape=latent_shape, device=device)
-        positions = get_pixel_coords(latent_coords, scale_factors=VIDEO_SCALE_FACTORS, causal_fix=True)
-        positions = positions.to(torch.bfloat16)
-        positions[:, 0, ...] = positions[:, 0, ...] / fps
-
-        return patchified, positions
-
-    def _run_denoising(
-        self,
-        config: GenerationConfig,
-        video_state: LatentState,
-        audio_state: LatentState | None,
-        video_clean_state: LatentState,
-        audio_clean_state: LatentState | None,
-        v_ctx_pos: Tensor,
-        a_ctx_pos: Tensor,
-        v_ctx_neg: Tensor | None,
-        a_ctx_neg: Tensor | None,
-        device: torch.device,
-    ) -> tuple[LatentState, LatentState | None]:
-        """Run the denoising loop using X0 prediction with CFG and optional STG."""
-        scheduler = LTX2Scheduler()
-        sigmas = scheduler.execute(steps=config.num_inference_steps).to(device).float()
-        stepper = EulerDiffusionStep()
-        cfg_guider = CFGGuider(config.guidance_scale)
-        stg_guider = STGGuider(config.stg_scale)
-
-        # Build STG perturbation config if STG is enabled
-        stg_perturbation_config = self._build_stg_perturbation_config(config) if stg_guider.enabled() else None
-
-        # Create initial modalities (will be updated each step via replace())
-        video = Modality(
-            enabled=True,
-            latent=video_state.latent,
-            sigma=sigmas[0].repeat(video_state.latent.shape[0]),
-            timesteps=video_state.denoise_mask,
-            positions=video_state.positions,
-            context=v_ctx_pos,
-            context_mask=None,
-        )
-
-        # Audio modality is None when not generating audio
-        audio: Modality | None = None
-        if audio_state is not None:
-            audio = Modality(
-                enabled=True,
-                latent=audio_state.latent,
-                sigma=sigmas[0].repeat(audio_state.latent.shape[0]),
-                timesteps=audio_state.denoise_mask,
-                positions=audio_state.positions,
-                context=a_ctx_pos,
-                context_mask=None,
-            )
-
-        # Wrap transformer with X0Model to convert velocity predictions to denoised outputs
-        self._transformer.to(device)
-        x0_model = X0Model(self._transformer)
-
-        with torch.autocast(device_type=str(device).split(":")[0], dtype=torch.bfloat16):
-            for step_idx, sigma in enumerate(sigmas[:-1]):
-                # Update modalities with current state and timesteps
-                video = replace(
-                    video,
-                    latent=video_state.latent,
-                    sigma=sigma.repeat(video_state.latent.shape[0]),
-                    timesteps=sigma * video_state.denoise_mask,
-                    positions=video_state.positions,
-                )
-
-                if audio is not None and audio_state is not None:
-                    audio = replace(
-                        audio,
-                        latent=audio_state.latent,
-                        sigma=sigma.repeat(audio_state.latent.shape[0]),
-                        timesteps=sigma * audio_state.denoise_mask,
-                        positions=audio_state.positions,
-                    )
-
-                # Run model (positive pass) - X0Model returns denoised outputs
-                pos_video, pos_audio = x0_model(video=video, audio=audio, perturbations=None)
-                denoised_video, denoised_audio = pos_video, pos_audio
-
-                # Apply CFG if guidance_scale != 1.0
-                if cfg_guider.enabled() and v_ctx_neg is not None:
-                    video_neg = replace(video, context=v_ctx_neg)
-                    audio_neg = replace(audio, context=a_ctx_neg) if audio is not None else None
-                    neg_video, neg_audio = x0_model(video=video_neg, audio=audio_neg, perturbations=None)
-
-                    denoised_video = denoised_video + cfg_guider.delta(pos_video, neg_video)
-                    if audio is not None and denoised_audio is not None:
-                        denoised_audio = denoised_audio + cfg_guider.delta(pos_audio, neg_audio)
-
-                # Apply STG if stg_scale != 0.0
-                if stg_guider.enabled() and stg_perturbation_config is not None:
-                    perturbed_video, perturbed_audio = x0_model(
-                        video=video, audio=audio, perturbations=stg_perturbation_config
-                    )
-                    denoised_video = denoised_video + stg_guider.delta(pos_video, perturbed_video)
-                    if audio is not None and denoised_audio is not None and perturbed_audio is not None:
-                        denoised_audio = denoised_audio + stg_guider.delta(pos_audio, perturbed_audio)
-
-                # Apply conditioning mask (keep conditioned tokens clean)
-                denoised_video = denoised_video * video_state.denoise_mask + video_clean_state.latent.float() * (
-                    1 - video_state.denoise_mask
-                )
-                if audio is not None and audio_state is not None and audio_clean_state is not None:
-                    denoised_audio = denoised_audio * audio_state.denoise_mask + audio_clean_state.latent.float() * (
-                        1 - audio_state.denoise_mask
-                    )
-
-                # Euler step
-                video_state = replace(
-                    video_state,
-                    latent=stepper.step(
-                        sample=video.latent, denoised_sample=denoised_video, sigmas=sigmas, step_index=step_idx
-                    ),
-                )
-                if audio is not None and audio_state is not None:
-                    audio_state = replace(
-                        audio_state,
-                        latent=stepper.step(
-                            sample=audio.latent, denoised_sample=denoised_audio, sigmas=sigmas, step_index=step_idx
-                        ),
-                    )
-
-                # Update progress
-                if self._sampling_context is not None:
-                    self._sampling_context.advance_step()
-
-        return video_state, audio_state
-
-    @staticmethod
-    def _build_stg_perturbation_config(config: GenerationConfig) -> BatchedPerturbationConfig:
-        """Build the perturbation config for STG based on the stg_mode."""
-        # Always skip video self-attention for STG
-        perturbations: list[Perturbation] = [
-            Perturbation(type=PerturbationType.SKIP_VIDEO_SELF_ATTN, blocks=config.stg_blocks)
-        ]
-
-        # Optionally also skip audio self-attention (stg_av mode)
-        if config.stg_mode == "stg_av":
-            perturbations.append(Perturbation(type=PerturbationType.SKIP_AUDIO_SELF_ATTN, blocks=config.stg_blocks))
-
-        perturbation_config = PerturbationConfig(perturbations=perturbations)
-        # Batch size is 1 for validation
-        return BatchedPerturbationConfig(perturbations=[perturbation_config])
-
-    def _decode_video_latent(self, latent: Tensor, config: GenerationConfig, device: torch.device) -> Tensor:
-        """Decode patchified video latent to pixel space."""
-        # Unpatchify
-        latent_frames = config.num_frames // VIDEO_SCALE_FACTORS.time + 1
-        latent_height = config.height // VIDEO_SCALE_FACTORS.height
-        latent_width = config.width // VIDEO_SCALE_FACTORS.width
-
-        unpatchified = self._video_patchifier.unpatchify(
-            latent,
-            output_shape=VideoLatentShape(
-                height=latent_height,
-                width=latent_width,
-                frames=latent_frames,
-                batch=1,
-                channels=128,
-            ),
-        )
-
-        # Decode - ensure bfloat16 to match decoder weights
-        self._vae_decoder.to(device)
-        unpatchified = unpatchified.to(dtype=torch.bfloat16)
-        tiled_config = config.tiled_decoding
-
-        if tiled_config is not None and tiled_config.enabled:
-            # Use tiled decoding for reduced VRAM
-            tiling_config = TilingConfig(
-                spatial_config=SpatialTilingConfig(
-                    tile_size_in_pixels=tiled_config.tile_size_pixels,
-                    tile_overlap_in_pixels=tiled_config.tile_overlap_pixels,
-                ),
-                temporal_config=TemporalTilingConfig(
-                    tile_size_in_frames=tiled_config.tile_size_frames,
-                    tile_overlap_in_frames=tiled_config.tile_overlap_frames,
-                ),
-            )
-            chunks = []
-            for video_chunk in self._vae_decoder.tiled_decode(
-                unpatchified,
-                tiling_config=tiling_config,
-            ):
-                chunks.append(video_chunk)
-            decoded_video = torch.cat(chunks, dim=2)
-        else:
-            # Standard full decoding
-            decoded_video = self._vae_decoder(unpatchified)
-
-        decoded_video = ((decoded_video + 1.0) / 2.0).clamp(0.0, 1.0)
-        self._vae_decoder.to("cpu")
-
-        return decoded_video[0].float().cpu()
-
-    def _validate_config(self, config: GenerationConfig) -> None:
-        """Validate generation configuration."""
-        if config.height % 32 != 0 or config.width % 32 != 0:
-            raise ValueError(f"height and width must be divisible by 32, got {config.height}x{config.width}")
-        if config.num_frames % 8 != 1:
-            raise ValueError(f"num_frames must satisfy num_frames % 8 == 1, got {config.num_frames}")
-        if config.generate_audio and (self._audio_decoder is None or self._vocoder is None):
-            raise ValueError("Audio generation requires audio_decoder and vocoder")
-        if config.condition_image is not None and self._vae_encoder is None:
-            raise ValueError("Image conditioning requires vae_encoder")
-        if config.reference_video is not None and self._vae_encoder is None:
-            raise ValueError("Reference video conditioning requires vae_encoder")
-
-        # Validate prompt embedding source
-        if config.cached_embeddings is None and self._text_encoder is None:
-            raise ValueError("Either text_encoder or config.cached_embeddings must be provided")
-        if config.cached_embeddings is None and self._embeddings_processor is None:
-            raise ValueError("embeddings_processor is required when encoding prompts on-the-fly")
-
-    def _get_prompt_embeddings(
-        self, config: GenerationConfig, device: torch.device
-    ) -> tuple[Tensor, Tensor, Tensor | None, Tensor | None]:
-        """Get prompt embeddings from config cache or encode on-the-fly."""
-        if config.cached_embeddings is not None:
-            # Use pre-computed embeddings from config
-            cached = config.cached_embeddings
-            v_ctx_pos = cached.video_context_positive.to(device)
-            a_ctx_pos = cached.audio_context_positive.to(device)
-            v_ctx_neg = cached.video_context_negative.to(device) if cached.video_context_negative is not None else None
-            a_ctx_neg = cached.audio_context_negative.to(device) if cached.audio_context_negative is not None else None
-            return v_ctx_pos, a_ctx_pos, v_ctx_neg, a_ctx_neg
-
-        # Fall back to encoding on-the-fly
-        return self._encode_prompts(config, device)
-
-    def _encode_prompts(
-        self, config: GenerationConfig, device: torch.device
-    ) -> tuple[Tensor, Tensor, Tensor | None, Tensor | None]:
-        """Encode positive and negative prompts using the text encoder + embeddings processor."""
-        self._text_encoder.to(device)
-        self._embeddings_processor.to(device)
-
-        pos_hs, pos_mask = self._text_encoder.encode(config.prompt)
-        pos_out = self._embeddings_processor.process_hidden_states(pos_hs, pos_mask)
-        v_ctx_pos, a_ctx_pos = pos_out.video_encoding, pos_out.audio_encoding
-
-        v_ctx_neg, a_ctx_neg = None, None
-        if config.guidance_scale != 1.0:
-            neg_hs, neg_mask = self._text_encoder.encode(config.negative_prompt)
-            neg_out = self._embeddings_processor.process_hidden_states(neg_hs, neg_mask)
-            v_ctx_neg, a_ctx_neg = neg_out.video_encoding, neg_out.audio_encoding
-
-        # Move the base Gemma model to CPU
-        self._text_encoder.model.to("cpu")
-
-        return v_ctx_pos, a_ctx_pos, v_ctx_neg, a_ctx_neg
-
-    def _decode_video(
-        self, video_state: LatentState, device: torch.device, tiled_config: TiledDecodingConfig | None = None
-    ) -> Tensor:
-        """Decode video latents to pixel space.
-        Args:
-            video_state: Video latent state to decode
-            device: Device to run decoding on
-            tiled_config: Optional tiled decoding configuration for reduced VRAM usage
-        Returns:
-            Decoded video tensor [C, F, H, W] in [0, 1] range
-        """
-        self._vae_decoder.to(device)
-        # Ensure latent is bfloat16 to match decoder weights
-        latent = video_state.latent.to(dtype=torch.bfloat16)
-
-        if tiled_config is not None and tiled_config.enabled:
-            # Use tiled decoding for reduced VRAM
-            tiling_config = TilingConfig(
-                spatial_config=SpatialTilingConfig(
-                    tile_size_in_pixels=tiled_config.tile_size_pixels,
-                    tile_overlap_in_pixels=tiled_config.tile_overlap_pixels,
-                ),
-                temporal_config=TemporalTilingConfig(
-                    tile_size_in_frames=tiled_config.tile_size_frames,
-                    tile_overlap_in_frames=tiled_config.tile_overlap_frames,
-                ),
-            )
-            chunks = []
-            for video_chunk in self._vae_decoder.tiled_decode(
-                latent,
-                tiling_config=tiling_config,
-            ):
-                chunks.append(video_chunk)
-            decoded_video = torch.cat(chunks, dim=2)
-        else:
-            # Standard full decoding
-            decoded_video = self._vae_decoder(latent)
-
-        decoded_video = ((decoded_video + 1.0) / 2.0).clamp(0.0, 1.0)
-        self._vae_decoder.to("cpu")
-        return decoded_video[0].float().cpu()
-
-    def _decode_audio(self, audio_state: LatentState, device: torch.device) -> Tensor:
-        """Decode audio latents to waveform."""
-        self._audio_decoder.to(device)
-        first_param = next(self._audio_decoder.parameters(), None)
-        decoder_dtype = first_param.dtype if first_param is not None else audio_state.latent.dtype
-        latent = audio_state.latent.to(dtype=decoder_dtype, device=device)
-        decoded_audio = self._audio_decoder(latent)
-        self._audio_decoder.to("cpu")
-
-        self._vocoder.to(device)
-        audio_waveform = self._vocoder(decoded_audio)
-        self._vocoder.to("cpu")
-
-        return audio_waveform.squeeze(0).float().cpu()
-
-    @staticmethod
-    def _concatenate_videos_side_by_side(left_video: Tensor, right_video: Tensor) -> Tensor:
-        """Concatenate two videos side-by-side (horizontally).
-        If the videos have different frame counts, the shorter one is padded with
-        its last frame repeated.
-        Args:
-            left_video: Left video tensor [C, F1, H1, W1] in [0, 1]
-            right_video: Right video tensor [C, F2, H2, W2] in [0, 1]
-        Returns:
-            Concatenated video tensor [C, max(F1,F2), H2, W1_scaled+W2] in [0, 1]
-        """
-        left_height, left_width = left_video.shape[2], left_video.shape[3]
-        right_height = right_video.shape[2]
-
-        # Resize left video to match right video's height if needed
-        if left_height != right_height:
-            # Scale width proportionally to maintain aspect ratio
-            scale = right_height / left_height
-            new_width = int(left_width * scale)
-            # Interpolate expects [N, C, H, W], we have [C, F, H, W]
-            # Reshape to [C*F, 1, H, W] -> interpolate -> reshape back
-            c, f, h, w = left_video.shape
-            left_video = left_video.reshape(c * f, 1, h, w)
-            left_video = torch.nn.functional.interpolate(
-                left_video, size=(right_height, new_width), mode="bilinear", align_corners=False
-            )
-            left_video = left_video.reshape(c, f, right_height, new_width)
-
-        left_frames = left_video.shape[1]
-        right_frames = right_video.shape[1]
-
-        # Pad shorter video by repeating last frame
-        if left_frames < right_frames:
-            padding = left_video[:, -1:, :, :].expand(-1, right_frames - left_frames, -1, -1)
-            left_video = torch.cat([left_video, padding], dim=1)
-        elif right_frames < left_frames:
-            padding = right_video[:, -1:, :, :].expand(-1, left_frames - right_frames, -1, -1)
-            right_video = torch.cat([right_video, padding], dim=1)
-
-        # Concatenate along width dimension
-        return torch.cat([left_video, right_video], dim=3)
-
-    def _encode_conditioning_image(
-        self,
-        image: Tensor,
-        target_height: int,
-        target_width: int,
-        device: torch.device,
-    ) -> Tensor:
-        """Encode a conditioning image to latent space.
-        The image is resized to cover the target dimensions while preserving aspect ratio,
-        then center-cropped to exactly match the target size.
-        """
-        # image is [C, H, W] in [0, 1]  # noqa: ERA001
-        current_height, current_width = image.shape[1:]
-
-        # Resize maintaining aspect ratio (cover target, then center crop)
-        if current_height != target_height or current_width != target_width:
-            aspect_ratio = current_width / current_height
-            target_aspect_ratio = target_width / target_height
-
-            if aspect_ratio > target_aspect_ratio:
-                # Image is wider than target - resize to match height, crop width
-                resize_height = target_height
-                resize_width = int(target_height * aspect_ratio)
-            else:
-                # Image is taller than target - resize to match width, crop height
-                resize_height = int(target_width / aspect_ratio)
-                resize_width = target_width
-
-            image = rearrange(image, "c h w -> 1 c h w")
-            image = torch.nn.functional.interpolate(
-                image, size=(resize_height, resize_width), mode="bilinear", align_corners=False
-            )
-
-            # Center crop to target dimensions
-            h_start = (resize_height - target_height) // 2
-            w_start = (resize_width - target_width) // 2
-            image = image[:, :, h_start : h_start + target_height, w_start : w_start + target_width]
-        else:
-            image = rearrange(image, "c h w -> 1 c h w")
-
-        # Add frame dimension and convert to [-1, 1]
-        image = rearrange(image, "b c h w -> b c 1 h w")
-        image = (image * 2.0 - 1.0).to(device=device, dtype=torch.float32)
-
-        # Encode
-        self._vae_encoder.to(device)
-        with torch.autocast(device_type=str(device).split(":")[0], dtype=torch.bfloat16):
-            encoded = self._vae_encoder(image)
-        self._vae_encoder.to("cpu")
-
-        return encoded