indigo/LTX-2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Automated PR - 2026-05-28

Browse Source
This commit is contained in:
github-actions[bot]
2026-05-28 14:26:15 +00:00
parent 1799988521
commit 203d4842d4
45 changed files with 1969 additions and 668 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1 -1
View File
@@ -77,7 +77,7 @@ Download the following models from the [LTX-2.3 HuggingFace repository](https://
* **Use DistilledPipeline** - Fastest inference with only 8 predefined sigmas (8 steps stage 1, 4 steps stage 2)
* **Enable FP8 quantization** - Enables lower memory footprint: `--quantization fp8-cast` (CLI) or `quantization=QuantizationPolicy.fp8_cast()` (Python). Fp8-cast should be used with bf16 checkpoints, it shall downcast them on the fly. For Hopper GPUs with TensorRT-LLM, use `--quantization fp8-scaled-mm` for FP8 scaled matrix multiplication. Fp8-scaled-mm should be used with fp8 checkpoints.
* **Install attention optimizations** - Use xFormers (`uv sync --extra xformers`) or [Flash Attention 3](https://github.com/Dao-AILab/flash-attention) for Hopper GPUs
* **Install attention optimizations** - On datacenter Blackwell GPUs (B200), install FlashAttention 4 manually: `uv pip install 'flash-attn-4==4.0.0b9'` (this specific revision is the one we have verified against torch 2.9.1+cu128; newer betas have known issues on consumer Blackwell). On other CUDA GPUs (including Hopper), use xFormers (`uv sync --extra xformers`).
* **Use gradient estimation** - Reduce inference steps from 40 to 20-30 while maintaining quality (see [pipeline documentation](packages/ltx-pipelines/README.md#denoising-loop-optimization))
* **Skip memory cleanup** - If you have sufficient VRAM, disable automatic memory cleanup between stages for faster processing
* **Choose single-stage pipeline** - Use `TI2VidOneStagePipeline` for faster generation when high resolution isn't required
packages/ltx-core/README.md
+13 -24
View File
@@ -125,39 +125,26 @@ Uses NVIDIA TensorRT-LLM's `cublas_scaled_mm` for efficient FP8 matrix multiplic
**Usage with QuantizationPolicy:**
```python
from ltx_core.quantization import QuantizationPolicy
from ltx_core.quantization.fp8_scaled_mm import build_policy as build_fp8_scaled_mm_policy
# Dynamic input quantization (no calibration needed)
policy = QuantizationPolicy.fp8_scaled_mm()
# Static input quantization with calibration file
policy = QuantizationPolicy.fp8_scaled_mm(calibration_amax_path="/path/to/amax.json")
# Discovers the layer set from the checkpoint's .weight_scale tensors
policy = build_fp8_scaled_mm_policy("/path/to/checkpoint.safetensors")
```
The policy provides `sd_ops` and `module_ops` that can be passed to the model builder:
The policy carries `sd_ops`, `module_ops`, and `fuse_rule` that are passed to the model builder:
```python
import torch
from ltx_core.loader import SingleGPUModelBuilder
builder = SingleGPUModelBuilder(
model=model,
device=device,
sd_ops=policy.sd_ops,
model_class_configurator=MyModelConfigurator,
model_path="/path/to/checkpoint.safetensors",
model_sd_ops=policy.sd_ops,
module_ops=policy.module_ops,
fuse_rule=policy.fuse_rule,
)
builder.load(checkpoint_path)
```
**Calibration File Format** (for static input quantization):
```json
{
"amax_values": {
"transformer_blocks.0.attn.to_q.input_quantizer": 12.5,
"transformer_blocks.0.attn.to_k.input_quantizer": 8.3,
...
}
}
model = builder.build(device=torch.device("cuda"))
```
#### FP8 Cast
@@ -165,7 +152,9 @@ builder.load(checkpoint_path)
A simpler approach that casts weights to FP8 for storage and upcasts during inference:
```python
policy = QuantizationPolicy.fp8_cast()
from ltx_core.quantization.fp8_cast import build_policy as build_fp8_cast_policy
policy = build_fp8_cast_policy("/path/to/checkpoint.safetensors")
```
For complete, production-ready pipeline implementations that combine these building blocks, see the [`ltx-pipelines`](../ltx-pipelines/) package.
packages/ltx-core/pyproject.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[project]
name = "ltx-core"
version = "1.1.3"
version = "1.1.4"
description = "Core implementation of Lightricks' LTX-2 model"
readme = "README.md"
requires-python = ">=3.10"
packages/ltx-core/src/ltx_core/block_streaming/builder.py
+49 -47
View File
@@ -3,7 +3,6 @@
from __future__ import annotations
import logging
from collections.abc import Callable
from dataclasses import dataclass, field, replace
from typing import Generic
@@ -17,13 +16,14 @@ from ltx_core.block_streaming.provider import WeightsProvider
from ltx_core.block_streaming.source import DiskWeightSource, PinnedWeightSource, WeightSource
from ltx_core.block_streaming.utils import allocate_layout_views, derive_layout, make_block_key, resolve_attr
from ltx_core.block_streaming.wrapper import BlockStreamingWrapper
from ltx_core.loader.fuse_loras import aggregate_lora_products, fuse_lora_weights
from ltx_core.loader.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
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.primitives import (
LoraPathStrengthAndSDOps,
LoraStateDictWithStrength,
ModelBuilderProtocol,
StateDict,
StateDictLoader,
)
from ltx_core.loader.registry import DummyRegistry, Registry
@@ -50,14 +50,13 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
loras: LoRA adapters fused into weights at load time.
model_loader: Strategy for reading checkpoint metadata.
registry: Shared cache for loaded state dicts.
fuse_rule: Per-policy LoRA merge rule. Defaults to ``bf16_fuse_rule``;
use ``fp8_cast_fuse_rule`` for fp8_cast streaming so the pinned
buffers receive correctly-quantized weights.
blocks_attr: Dotted path to the ``nn.ModuleList`` (e.g.
``"velocity_model.transformer_blocks"``).
``"transformer_blocks"``).
blocks_prefix: State-dict key prefix for block weights
(e.g. ``"transformer_blocks"``).
state_dict_prefix: Wrapper offset prepended to keys when loading into
the meta model (e.g. ``"velocity_model."`` when wrapped by ``X0Model``).
model_wrapper: Optional callable wrapping the model
(e.g. ``X0Model``).
"""
model_class_configurator: type[ModelConfigurator[ModelType]]
@@ -67,12 +66,11 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
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
# Streaming-specific
blocks_attr: str = ""
blocks_prefix: str = ""
state_dict_prefix: str = ""
model_wrapper: Callable[[ModelType], nn.Module] | None = None
def with_sd_ops(self, sd_ops: SDOps | None) -> StreamingModelBuilder:
return replace(self, model_sd_ops=sd_ops)
@@ -83,6 +81,9 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> StreamingModelBuilder:
return replace(self, loras=loras)
def with_fuse_rule(self, fuse_rule: FuseRule) -> StreamingModelBuilder:
return replace(self, fuse_rule=fuse_rule)
def model_config(self) -> dict:
"""Read model configuration from the checkpoint metadata."""
return read_model_config(self.model_path, self.model_loader)
@@ -113,8 +114,6 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
config = read_model_config(self.model_path, self.model_loader)
meta_model: nn.Module = create_meta_model(self.model_class_configurator, config, self.module_ops)
if self.model_wrapper is not None:
meta_model = self.model_wrapper(meta_model)
meta_model.eval()
blocks = resolve_attr(meta_model, self.blocks_attr)
@@ -142,7 +141,15 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
target_device,
reuse_barrier=lambda event: copy_stream.wait_event(event),
)
provider = WeightsProvider(gpu_pool, copy_stream, target_device, source, lora_sources, self.blocks_prefix)
provider = WeightsProvider(
gpu_pool,
copy_stream,
target_device,
source,
lora_sources,
self.blocks_prefix,
fuse_rule=self.fuse_rule,
)
return BlockStreamingWrapper(
model=meta_model,
blocks=blocks,
@@ -190,7 +197,9 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
pinned_blocks = allocate_layout_views(blocks_layout, pin_memory=True)
should_sync = False
for key, fused in fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype=None, preserve_input_device=False):
for key, fused in fuse_lora_weights(
model_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
@@ -216,7 +225,7 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
}
non_block_sd: dict[str, torch.Tensor] = {
self.state_dict_prefix + model_key: model_sd.sd[model_key].to(device=target_device, dtype=dtype)
model_key: model_sd.sd[model_key].to(device=target_device, dtype=dtype)
for _sft_key, model_key in non_block_keys
}
@@ -235,7 +244,7 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
non_block_keys: list[tuple[str, str]],
) -> tuple[WeightSource, list[LoraSource]]:
"""Create a DiskWeightSource backed by a DiskBlockReader for lazy loading.
Derives the shared pool layout from the meta model's block 0 — this
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.
"""
@@ -248,8 +257,8 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
target_device,
dtype,
sd_ops=self.model_sd_ops,
key_prefix=self.state_dict_prefix,
lora_sources=lora_sources,
fuse_rule=self.fuse_rule,
)
blocks = resolve_attr(meta_model, self.blocks_attr)
@@ -271,28 +280,6 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
source = DiskWeightSource(cpu_pool, block_reader)
return source, lora_sources
# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------
@staticmethod
def _fuse_lora_delta(
model_key: str,
tensor: torch.Tensor,
lora_sources: list[LoraSource],
) -> torch.Tensor:
"""Add all matching LoRA deltas to *tensor* in-place via ``addmm_``."""
if not lora_sources or not model_key.endswith(".weight"):
return tensor
prefix = model_key[: -len(".weight")]
products = (
ab
for ab in (s.get_ab(prefix, device=tensor.device, dtype=tensor.dtype) for s in lora_sources)
if ab is not None
)
aggregate_lora_products(products, out=tensor)
return tensor
@staticmethod
@torch.inference_mode()
def _load_non_block_weights(
@@ -302,21 +289,36 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
device: torch.device,
dtype: torch.dtype,
sd_ops: SDOps | None = None,
key_prefix: str = "",
lora_sources: list[LoraSource] | None = None,
fuse_rule: FuseRule = bf16_fuse_rule,
) -> None:
"""Load non-block weights into *model* on *device*."""
state_dict: dict[str, torch.Tensor] = {}
sources = lora_sources or []
"""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.
"""
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)
tensor = StreamingModelBuilder._fuse_lora_delta(model_key, tensor, sources)
if sd_ops is not None:
for kv in sd_ops.apply_to_key_value(model_key, tensor):
state_dict[key_prefix + kv.new_key] = kv.new_value
continue
state_dict[key_prefix + model_key] = tensor
model.load_state_dict(state_dict, strict=False, assign=True)
non_block_sd[kv.new_key] = kv.new_value
else:
non_block_sd[model_key] = tensor
if lora_sources:
lora_sd_and_strengths = [src.as_state_dict_with_strength() for src in lora_sources]
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,
preserve_input_device=True,
):
non_block_sd[key] = fused
model.load_state_dict(non_block_sd, strict=False, assign=True)
def _scan_checkpoint_keys(
packages/ltx-core/src/ltx_core/block_streaming/disk.py
+17
View File
@@ -9,6 +9,7 @@ import torch
from ltx_core.block_streaming.utils import allocate_layout_views, make_block_key
from ltx_core.loader.fuse_loras import LoraProduct
from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict
from ltx_core.loader.sd_ops import SDOps
_SAFETENSORS_DTYPE_TO_TORCH: dict[str, torch.dtype] = {
@@ -125,6 +126,22 @@ class LoraSource:
b_view.copy_(handle.get_tensor(b_keys[prefix]))
self._pinned_ab[prefix] = (a_view, b_view)
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``.
"""
sd: dict[str, torch.Tensor] = {}
for prefix, (a, b) in self._pinned_ab.items():
sd[f"{prefix}.lora_A.weight"] = a
sd[f"{prefix}.lora_B.weight"] = b
size = sum(t.numel() * t.element_size() for t in sd.values())
dtypes = {t.dtype for t in sd.values()}
return LoraStateDictWithStrength(
StateDict(sd=sd, device=torch.device("cpu"), size=size, dtype=dtypes),
self.strength,
)
def get_ab(
self,
param_prefix: str,
packages/ltx-core/src/ltx_core/block_streaming/provider.py
+14 -11
View File
@@ -9,8 +9,10 @@ import torch
from ltx_core.block_streaming.disk import LoraSource
from ltx_core.block_streaming.pool import WeightPool
from ltx_core.block_streaming.source import WeightSource
from ltx_core.block_streaming.utils import FP8_DTYPES
from ltx_core.loader.fuse_loras import aggregate_lora_products, fuse_cast_fp8_weight
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:
@@ -43,6 +45,8 @@ class WeightsProvider:
source: Pinned CPU weight source.
lora_sources: LoRA adapters fused on H2D copy.
blocks_prefix: State-dict prefix for LoRA key matching.
fuse_rule: Per-policy LoRA merge rule (must be streaming-compatible:
no companion-key emission). Defaults to ``bf16_fuse_rule``.
"""
def __init__(
@@ -53,6 +57,7 @@ class WeightsProvider:
source: WeightSource,
lora_sources: list[LoraSource] | None = None,
blocks_prefix: str = "",
fuse_rule: FuseRule = bf16_fuse_rule,
) -> None:
self._copy_stream = copy_stream
self._pool = pool
@@ -62,6 +67,7 @@ class WeightsProvider:
self._source = source
self._lora_sources = lora_sources or []
self._blocks_prefix = blocks_prefix
self._fuse_rule = fuse_rule
def get(self, idx: int) -> dict[str, torch.Tensor]:
"""Return GPU weights for block *idx*. Does H2D copy on miss."""
@@ -126,22 +132,19 @@ class WeightsProvider:
return len(self._cache)
def _fuse_block_loras(self, idx: int, weights: dict[str, torch.Tensor]) -> None:
"""Fuse LoRA deltas directly into GPU block weights."""
"""Fuse LoRA deltas directly into GPU block weights via ``fuse_rule``."""
agg_dtype = self._fuse_rule.aggregation_dtype
for name, tensor in weights.items():
if not name.endswith(".weight"):
continue
prefix = f"{self._blocks_prefix}.{idx}.{name}".removesuffix(".weight")
is_fp8 = tensor.dtype in FP8_DTYPES
agg_dtype = torch.bfloat16 if is_fp8 else tensor.dtype
products = (
ab
for ab in (s.get_ab(prefix, device=self._target_device, dtype=agg_dtype) for s in self._lora_sources)
if ab is not None
)
aggregated = aggregate_lora_products(products, agg_dtype)
if aggregated is None:
deltas = aggregate_lora_products(products, agg_dtype)
if deltas is None:
continue
if is_fp8:
tensor.copy_(fuse_cast_fp8_weight(aggregated, tensor, tensor.dtype))
else:
tensor.add_(aggregated)
fused = self._fuse_rule(name, tensor, deltas, _EMPTY_STATE_DICT)
tensor.copy_(fused[name])
packages/ltx-core/src/ltx_core/loader/attention_ops.py
+42
View File
@@ -0,0 +1,42 @@
"""Builder ops for swapping attention backends on a meta model before load."""
import torch
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.model.transformer.attention import (
Attention,
AttentionCallable,
AttentionFunction,
MaskedAttentionCallable,
MaskedAttentionFunction,
)
def set_attention_module_op(
attention: AttentionFunction | AttentionCallable | None = None,
masked_attention: MaskedAttentionFunction | MaskedAttentionCallable | None = None,
) -> ModuleOps:
"""Build a ``ModuleOps`` that overrides the attention callables on every
``Attention`` submodule of a model. Applied via ``create_meta_model`` so
the meta model is mutated before weight loading. Matcher returns False
for models with no ``Attention`` submodules, so the op is a no-op there.
Either or both slots may be supplied; *None* leaves that slot untouched.
"""
fn = attention.to_callable() if isinstance(attention, AttentionFunction) else attention
masked_fn = (
masked_attention.to_callable() if isinstance(masked_attention, MaskedAttentionFunction) else masked_attention
)
def matcher(model: torch.nn.Module) -> bool:
return any(isinstance(m, Attention) for m in model.modules())
def mutator(model: torch.nn.Module) -> torch.nn.Module:
for module in model.modules():
if isinstance(module, Attention):
if fn is not None:
module.attention_function = fn
if masked_fn is not None:
module.masked_attention_function = masked_fn
return model
return ModuleOps(name="set_attention_backend", matcher=matcher, mutator=mutator)
packages/ltx-core/src/ltx_core/loader/fuse_loras.py
+104 -88
View File
@@ -1,12 +1,10 @@
from collections.abc import Iterable, Iterator
from collections.abc import Callable, Iterable, Iterator
from dataclasses import dataclass
from typing import NamedTuple
import torch
from ltx_core.loader.kernels import TRITON_AVAILABLE
from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict
from ltx_core.quantization.fp8_cast import fused_add_round_launch
from ltx_core.quantization.fp8_scaled_mm import quantize_weight_to_fp8_per_tensor
class LoraProduct(NamedTuple):
@@ -17,6 +15,62 @@ class LoraProduct(NamedTuple):
strength: float
#: Signature for a fuse callable used by :class:`FuseRule`.
#:
#: Args:
#: key: The state-dict key being fused (e.g. ``"...layers.0.attn.q.weight"``).
#: weight: The current value at ``key`` from ``model_sd``, on the fusion device.
#: deltas: The pre-aggregated LoRA delta for ``key``, in ``aggregation_dtype``.
#: model_sd: The full state dict, for rules that need companion keys
#: (e.g. an existing ``.weight_scale``).
#:
#: Returns a dict of state-dict keys to overwrite -- at minimum ``{key: new_weight}``,
#: plus any companion keys (e.g. an updated ``.weight_scale``) the policy needs to
#: keep in sync.
FuseFn = Callable[[str, torch.Tensor, torch.Tensor, StateDict], dict[str, torch.Tensor]]
@dataclass(frozen=True)
class FuseRule:
"""Fuse an aggregated LoRA delta into one weight key.
Each policy supplies its own rule (see ``QuantizationPolicy.fuse_rule``);
``fuse_lora_weights`` is policy-agnostic boilerplate around it.
Attributes:
aggregation_dtype: Dtype callers must pre-aggregate LoRA deltas in
before invoking the rule.
fuse_fn: Callable that applies the pre-aggregated deltas to the weight
(and any companion keys) and returns a dict of keys to overwrite —
at minimum ``{key: new_weight}``, plus any companion keys (e.g. an
updated ``.weight_scale`` for scaled-FP8 layouts) the policy needs
to keep in sync.
"""
aggregation_dtype: torch.dtype
fuse_fn: FuseFn
def __call__(
self,
key: str,
weight: torch.Tensor,
deltas: torch.Tensor,
model_sd: StateDict,
) -> dict[str, torch.Tensor]:
return self.fuse_fn(key, weight, deltas, model_sd)
def _bf16_fuse(
key: str,
weight: torch.Tensor,
deltas: torch.Tensor,
model_sd: StateDict, # noqa: ARG001
) -> dict[str, torch.Tensor]:
deltas.add_(weight)
return {key: deltas.to(dtype=weight.dtype)}
bf16_fuse_rule = FuseRule(aggregation_dtype=torch.bfloat16, fuse_fn=_bf16_fuse)
def _get_device() -> torch.device:
if torch.cuda.is_available():
return torch.device("cuda", torch.cuda.current_device())
@@ -25,19 +79,16 @@ def _get_device() -> torch.device:
def aggregate_lora_products(
products: Iterable[LoraProduct],
dtype: torch.dtype | None = None,
*,
out: torch.Tensor | None = None,
dtype: torch.dtype,
) -> torch.Tensor | None:
"""Accumulate ``sum((B * strength) @ A)`` across :class:`LoraProduct` items.
If ``out`` is provided, ``addmm_`` accumulates directly into it — caller
ensures A/B dtypes and devices match ``out``. Otherwise the first product
materializes the ``(out, in)``-shape aggregator at ``dtype``; subsequent
products use ``addmm_`` to avoid allocating the full intermediate delta.
Returns ``out`` (or the new aggregator), or ``None`` if ``products`` was empty
and ``out`` was not given.
The first product materializes a freshly-allocated aggregator via
``torch.matmul(B * strength, A).to(dtype)`` -- preserving the
``(B * strength) @ A`` rounding pattern. Subsequent products use
``addmm_`` to avoid allocating the full intermediate delta.
Returns the aggregator, or ``None`` if ``products`` was empty.
"""
aggregated = out
aggregated: torch.Tensor | None = None
for product in products:
if aggregated is None:
aggregated = torch.matmul(product.b * product.strength, product.a).to(dtype=dtype)
@@ -46,62 +97,34 @@ def aggregate_lora_products(
return aggregated
def fuse_cast_fp8_weight(
delta_bf16: torch.Tensor,
weight_fp8: torch.Tensor,
target_dtype: torch.dtype,
) -> torch.Tensor:
"""Return ``(delta_bf16 + dequantize(weight_fp8)).to(target_dtype)``.
CUDA with Triton uses stochastic rounding; otherwise uses a deterministic bf16 add.
``delta_bf16`` is the bf16 accumulator and is mutated in place.
"""
if delta_bf16.dtype != torch.bfloat16:
raise ValueError(f"delta_bf16 must be bfloat16, got {delta_bf16.dtype}")
if str(weight_fp8.device).startswith("cuda") and TRITON_AVAILABLE:
fused_add_round_launch(delta_bf16, weight_fp8, seed=0)
else:
delta_bf16.add_(weight_fp8.to(dtype=torch.bfloat16))
return delta_bf16.to(dtype=target_dtype)
def fuse_lora_weights(
model_sd: StateDict,
lora_sd_and_strengths: list[LoraStateDictWithStrength],
dtype: torch.dtype | None = None,
fuse_rule: FuseRule = bf16_fuse_rule,
preserve_input_device: bool = True,
) -> Iterator[tuple[str, torch.Tensor]]:
"""Yield ``(key, fused_tensor)`` for each weight modified by at least one LoRA.
For scaled-FP8 weights, this includes both the updated ``.weight`` tensor
and its corresponding ``.weight_scale`` tensor.
The fusion math is delegated to ``fuse_rule``.
Output dtypes are the rule's responsibility.
When ``preserve_input_device`` is False, fused tensors are yielded on the device
used for fusion; caller is responsible for moving them to their final
destination.
"""
for key, original_weight in model_sd.sd.items():
if original_weight is None or key.endswith(".weight_scale"):
fusion_device = _get_device()
for key in _affected_weight_keys(lora_sd_and_strengths):
original_weight = model_sd.sd.get(key)
if original_weight is None:
continue
original_device = original_weight.device
weight = original_weight.to(device=_get_device())
target_dtype = dtype if dtype is not None else weight.dtype
deltas_dtype = target_dtype if target_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2] else torch.bfloat16
deltas = _aggregate_deltas(lora_sd_and_strengths, key, deltas_dtype, weight.device)
products = _products_for_sd_key(lora_sd_and_strengths, key, fuse_rule.aggregation_dtype, fusion_device)
deltas = aggregate_lora_products(products, fuse_rule.aggregation_dtype)
if deltas is None:
continue
scale_key = key.replace(".weight", ".weight_scale") if key.endswith(".weight") else None
is_scaled_fp8 = scale_key is not None and scale_key in model_sd.sd
original_device = original_weight.device
weight = original_weight.to(device=fusion_device)
if weight.dtype == torch.float8_e4m3fn:
if is_scaled_fp8:
fused = _fuse_delta_with_scaled_fp8(deltas, weight, key, scale_key, model_sd)
else:
fused = {key: fuse_cast_fp8_weight(deltas, weight, target_dtype)}
elif weight.dtype == torch.bfloat16:
deltas.add_(weight)
fused = {key: deltas.to(dtype=target_dtype)}
else:
raise ValueError(f"Unsupported dtype: {weight.dtype}")
fused = fuse_rule(key, weight, deltas, model_sd)
for k, v in fused.items():
yield k, v.to(device=original_device) if preserve_input_device else v
@@ -110,53 +133,46 @@ def fuse_lora_weights(
def apply_loras(
model_sd: StateDict,
lora_sd_and_strengths: list[LoraStateDictWithStrength],
dtype: torch.dtype | None = None,
fuse_rule: FuseRule = bf16_fuse_rule,
destination_sd: StateDict | None = None,
) -> StateDict:
"""Fuse LoRAs into ``model_sd`` and place the results in ``destination_sd``.
When ``destination_sd`` is provided, the fused tensors are placed directly into it.
"""
fused_iter = fuse_lora_weights(
model_sd,
lora_sd_and_strengths,
fuse_rule=fuse_rule,
)
if destination_sd is not None:
for key, fused in fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype):
for key, fused in fused_iter:
destination_sd.sd[key] = fused
return destination_sd
fused = dict(fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype))
fused = dict(fused_iter)
sd = {k: (fused[k] if k in fused else v.clone()) for k, v in model_sd.sd.items()}
return StateDict(sd, model_sd.device, model_sd.size, model_sd.dtype)
def _aggregate_deltas(
lora_sd_and_strengths: list[LoraStateDictWithStrength], key: str, dtype: torch.dtype, device: torch.device
) -> torch.Tensor | None:
def _affected_weight_keys(lora_sd_and_strengths: list[LoraStateDictWithStrength]) -> set[str]:
"""Return the set of ``.weight`` keys touched by at least one LoRA in the list."""
suffix = ".lora_A.weight"
return {k[: -len(suffix)] + ".weight" for lsd, _ in lora_sd_and_strengths for k in lsd.sd if k.endswith(suffix)}
def _products_for_sd_key(
lora_sd_and_strengths: list[LoraStateDictWithStrength],
key: str,
dtype: torch.dtype,
device: torch.device,
) -> Iterator[LoraProduct]:
"""Yield :class:`LoraProduct` items matching *key* across state-dict-backed LoRAs."""
prefix = key[: -len(".weight")]
key_a = f"{prefix}.lora_A.weight"
key_b = f"{prefix}.lora_B.weight"
def _ab_products() -> Iterator[LoraProduct]:
for lsd, coef in lora_sd_and_strengths:
if key_a not in lsd.sd or key_b not in lsd.sd:
continue
a = lsd.sd[key_a].to(device=device, dtype=dtype, non_blocking=True)
b = lsd.sd[key_b].to(device=device, dtype=dtype, non_blocking=True)
yield LoraProduct(a, b, coef)
return aggregate_lora_products(_ab_products(), dtype)
def _fuse_delta_with_scaled_fp8(
deltas: torch.Tensor,
weight: torch.Tensor,
key: str,
scale_key: str,
model_sd: StateDict,
) -> dict[str, torch.Tensor]:
"""Dequantize scaled FP8 weight, add LoRA delta, and re-quantize."""
weight_scale = model_sd.sd[scale_key]
original_weight = weight.to(torch.float32) * weight_scale
new_weight = original_weight + deltas.to(torch.float32)
new_fp8_weight, new_weight_scale = quantize_weight_to_fp8_per_tensor(new_weight)
return {key: new_fp8_weight, scale_key: new_weight_scale}
for lsd, coef in lora_sd_and_strengths:
if key_a not in lsd.sd or key_b not in lsd.sd:
continue
a = lsd.sd[key_a].to(device=device, dtype=dtype, non_blocking=True)
b = lsd.sd[key_b].to(device=device, dtype=dtype, non_blocking=True)
yield LoraProduct(a, b, coef)
packages/ltx-core/src/ltx_core/loader/primitives.py
+5
View File
@@ -10,6 +10,7 @@ from ltx_core.loader.sd_ops import SDOps
from ltx_core.model.model_protocol import ModelType
if TYPE_CHECKING:
from ltx_core.loader.fuse_loras import FuseRule
from ltx_core.loader.registry import Registry
@@ -110,6 +111,10 @@ class ModelBuilderProtocol(BuilderProtocol[ModelType], Protocol[ModelType]):
"""Return a copy of this builder that loads LoRA weights onto the given device."""
...
def with_fuse_rule(self, fuse_rule: "FuseRule") -> "ModelBuilderProtocol[ModelType]":
"""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:
packages/ltx-core/src/ltx_core/loader/single_gpu_model_builder.py
+14 -4
View File
@@ -5,7 +5,7 @@ from typing import Generic
import torch
from torch import nn
from ltx_core.loader.fuse_loras import apply_loras
from ltx_core.loader.fuse_loras import FuseRule, apply_loras, bf16_fuse_rule
from ltx_core.loader.helpers import create_meta_model, load_state_dict, read_model_config
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.primitives import (
@@ -46,6 +46,7 @@ def _load_model_weights(
dtype: torch.dtype | None,
model_sd_ops: SDOps | None = None,
lora_load_device: torch.device | None = None,
fuse_rule: FuseRule = bf16_fuse_rule,
) -> None:
"""Load base weights and fuse LoRAs into *meta_model* in-place."""
if lora_load_device is None:
@@ -57,7 +58,7 @@ def _load_model_weights(
if not lora_strengths or (min(lora_strengths) == 0 and max(lora_strengths) == 0):
sd = model_sd.sd
if dtype is not None:
sd = {key: value.to(dtype=dtype) for key, value in model_sd.sd.items()}
sd = {key: value.to(dtype=dtype) for key, value in sd.items()}
meta_model.load_state_dict(sd, strict=False, assign=True)
return
@@ -68,10 +69,13 @@ def _load_model_weights(
final_sd = apply_loras(
model_sd=model_sd,
lora_sd_and_strengths=lora_sd_and_strengths,
dtype=dtype,
fuse_rule=fuse_rule,
destination_sd=model_sd if isinstance(registry, DummyRegistry) else None,
)
meta_model.load_state_dict(final_sd.sd, strict=False, assign=True)
fused_sd = final_sd.sd
if dtype is not None:
fused_sd = {key: value.to(dtype=dtype) for key, value in fused_sd.items()}
meta_model.load_state_dict(fused_sd, strict=False, assign=True)
@dataclass(frozen=True)
@@ -91,6 +95,7 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
``torch.device("cpu")``, which keeps LoRA weights in CPU memory and transfers them to
the target GPU sequentially during fusion, reducing peak GPU memory usage compared to
loading all LoRA weights directly onto the GPU at once.
fuse_rule: Per-policy LoRA merge rule. Defaults to ``bf16_fuse_rule``;
"""
model_class_configurator: type[ModelConfigurator[ModelType]]
@@ -101,6 +106,7 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
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 lora(self, lora_path: str, strength: float, sd_ops: SDOps) -> "SingleGPUModelBuilder":
return replace(self, loras=(*self.loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops)))
@@ -120,6 +126,9 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
def with_lora_load_device(self, device: torch.device) -> "SingleGPUModelBuilder":
return replace(self, lora_load_device=device)
def with_fuse_rule(self, fuse_rule: FuseRule) -> "SingleGPUModelBuilder":
return replace(self, fuse_rule=fuse_rule)
def model_config(self) -> dict:
return read_model_config(self.model_path, self.model_loader)
@@ -158,5 +167,6 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
dtype=dtype,
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)
packages/ltx-core/src/ltx_core/modality_tiling.py
+62 -57
View File
@@ -21,11 +21,14 @@ from ltx_core.types import VideoLatentShape
@dataclass(frozen=True)
class TilingContext:
"""Opaque context produced by :meth:`VideoModalityTilingHelper.tile_modality`.
Carries the token-level keep mask and per-conditioning-token blend
Carries the token-level keep indices and per-conditioning-token blend
weights needed by :meth:`~VideoModalityTilingHelper.blend`.
"""
keep_mask: torch.Tensor
keep_indices: torch.Tensor
"""``(num_kept,)`` int64 — sorted indices of tokens the tile processes."""
num_total_tokens: int
"""Total number of tokens in the full (untiled) sequence."""
cond_blend_weights: torch.Tensor | None
"""``(num_kept_cond,)`` — weight for each kept conditioning token,
equal to ``1 / num_tiles_that_keep_this_token``. ``None`` when
@@ -81,14 +84,32 @@ class VideoModalityTilingHelper:
A ``(tiled_modality, context)`` tuple. Pass *context* to
:meth:`blend` together with the model output.
"""
keep_mask = self._keep_mask(modality, tile)
device = modality.positions.device
gen_indices = self._generated_token_indices(tile, device=device)
num_total = modality.latent.shape[1]
cond_blend_weights: torch.Tensor | None = None
if num_total > self._num_generated_tokens:
keep_per_tile_cond = self._all_tiles_cond_keep(modality) # (num_tiles, num_cond) bool
tile_idx = next((i for i, t in enumerate(self._tiles) if t.in_coords == tile.in_coords), None)
if tile_idx is None:
raise ValueError(
f"Tile with in_coords={tile.in_coords} is not in this helper's tile set; "
f"pass a tile obtained from `helper.tiles`."
)
my_cond_keep = keep_per_tile_cond[tile_idx]
cond_indices = self._num_generated_tokens + my_cond_keep.nonzero(as_tuple=False).squeeze(1)
keep_indices = torch.cat([gen_indices, cond_indices])
total_keepers = keep_per_tile_cond.sum(dim=0).float() # (num_cond,)
cond_blend_weights = 1.0 / total_keepers[my_cond_keep]
else:
keep_indices = gen_indices
tile_attention_mask = None
if modality.attention_mask is not None:
keep_indices = keep_mask.nonzero(as_tuple=False).squeeze(1)
tile_attention_mask = modality.attention_mask[:, keep_indices, :][:, :, keep_indices]
positions = modality.positions[:, :, keep_mask, :]
positions = modality.positions[:, :, keep_indices, :]
if normalize_positions:
num_tile_gen = self._tile_generated_token_count(tile)
gen_pos = positions[:, :, :num_tile_gen, :] # (B, 3, num_tile_gen, 2)
@@ -97,26 +118,15 @@ class VideoModalityTilingHelper:
tiled = replace(
modality,
latent=modality.latent[:, keep_mask, :],
timesteps=modality.timesteps[:, keep_mask],
latent=modality.latent[:, keep_indices, :],
timesteps=modality.timesteps[:, keep_indices],
positions=positions,
attention_mask=tile_attention_mask,
)
cond_blend_weights = None
num_total = modality.latent.shape[1]
if num_total > self._num_generated_tokens:
cond_keep = keep_mask[self._num_generated_tokens :]
# Count how many tiles keep each conditioning token.
cond_counts = torch.zeros(cond_keep.sum(), dtype=torch.float32)
for t in self._tiles:
other_mask = self._keep_mask(modality, t)
other_cond = other_mask[self._num_generated_tokens :]
# Map other tile's kept cond tokens into this tile's kept subset.
cond_counts += other_cond[cond_keep].float()
cond_blend_weights = 1.0 / cond_counts
return tiled, TilingContext(keep_mask=keep_mask, cond_blend_weights=cond_blend_weights)
return tiled, TilingContext(
keep_indices=keep_indices, num_total_tokens=num_total, cond_blend_weights=cond_blend_weights
)
# -- blend -------------------------------------------------------------
@@ -147,9 +157,9 @@ class VideoModalityTilingHelper:
"""
batch, _, dim = tile_to_blend.shape
num_tile_gen = self._tile_generated_token_count(tile)
gen_indices = self._generated_token_indices(tile)
gen_indices = self._generated_token_indices(tile, device=tile_to_blend.device)
num_total_tokens = context.keep_mask.shape[0]
num_total_tokens = context.num_total_tokens
expected_shape = (batch, num_total_tokens, dim)
if output is not None:
@@ -168,8 +178,7 @@ class VideoModalityTilingHelper:
# Scatter kept conditioning tokens, weighted by 1/N where N is
# the number of tiles that keep each token (so they sum to 1).
if num_total_tokens > self._num_generated_tokens and context.cond_blend_weights is not None:
cond_keep = context.keep_mask[self._num_generated_tokens :]
cond_indices = self._num_generated_tokens + cond_keep.nonzero(as_tuple=False).squeeze(1)
cond_indices = context.keep_indices[context.keep_indices >= self._num_generated_tokens]
weights = context.cond_blend_weights.to(device=tile_to_blend.device, dtype=tile_to_blend.dtype)
result[:, cond_indices, :] += tile_to_blend[:, num_tile_gen:, :] * weights[None, :, None]
@@ -189,46 +198,42 @@ class VideoModalityTilingHelper:
)
return self._patchifier.get_token_count(tile_shape)
def _generated_token_indices(self, tile: Tile) -> torch.Tensor:
def _generated_token_indices(self, tile: Tile, device: torch.device | None = None) -> torch.Tensor:
"""Flat token indices of *tile*'s generated tokens in the full sequence."""
frame_slice, height_slice, width_slice = tile.in_coords
f = torch.arange(frame_slice.start, frame_slice.stop)
h = torch.arange(height_slice.start, height_slice.stop)
w = torch.arange(width_slice.start, width_slice.stop)
f = torch.arange(frame_slice.start, frame_slice.stop, device=device)
h = torch.arange(height_slice.start, height_slice.stop, device=device)
w = torch.arange(width_slice.start, width_slice.stop, device=device)
return (
f[:, None, None] * self._latent_shape.height * self._latent_shape.width
+ h[None, :, None] * self._latent_shape.width
+ w[None, None, :]
).reshape(-1)
def _keep_mask(self, modality: Modality, tile: Tile) -> torch.Tensor:
"""Boolean mask ``(num_total_tokens,)`` — True for tokens the tile processes.
Generated tokens are selected by grid position. Conditioning
tokens are kept when their ``[start, end)`` intervals overlap
the tile in all three dimensions, or when they have a negative
time coordinate (reference tokens).
def _all_tiles_cond_keep(self, modality: Modality) -> torch.Tensor:
"""Vectorized (num_tiles, num_cond) bool: which tiles keep each conditioning token.
A conditioning token is kept by a tile when its ``[start, end)`` interval
overlaps the tile in all three dimensions, or when it has a negative time
coordinate (reference token).
"""
num_total = modality.latent.shape[1]
mask = torch.zeros(num_total, dtype=torch.bool)
cond_positions = modality.positions[:, :, self._num_generated_tokens :, :] # (B, 3, num_cond, 2)
device = cond_positions.device
gen_indices = self._generated_token_indices(tile)
mask[gen_indices] = True
# Per-tile (start, end) bounds along each axis; small Python loop (num_tiles <= ~16).
starts_list: list[torch.Tensor] = []
ends_list: list[torch.Tensor] = []
for t in self._tiles:
gen_idx = self._generated_token_indices(t, device=device)
gen_positions = modality.positions[:, :, gen_idx, :] # (B, 3, num_tile_gen, 2)
starts_list.append(gen_positions[..., 0].amin(dim=2)) # (B, 3)
ends_list.append(gen_positions[..., 1].amax(dim=2)) # (B, 3)
tile_starts = torch.stack(starts_list, dim=0) # (num_tiles, B, 3)
tile_ends = torch.stack(ends_list, dim=0) # (num_tiles, B, 3)
if num_total > self._num_generated_tokens:
gen_positions = modality.positions[:, :, gen_indices, :] # (B, 3, num_tile_gen, 2)
tile_start = gen_positions[..., 0].amin(dim=2) # (B, 3)
tile_end = gen_positions[..., 1].amax(dim=2) # (B, 3)
cond_positions = modality.positions[:, :, self._num_generated_tokens :, :] # (B, 3, num_cond, 2)
overlaps = (cond_positions[..., 0] < tile_end.unsqueeze(2)) & (
cond_positions[..., 1] > tile_start.unsqueeze(2)
) # (B, 3, num_cond)
overlaps_all_dims = overlaps.all(dim=1) # (B, num_cond)
has_negative_time = cond_positions[:, 0, :, 0] < 0 # (B, num_cond)
keep_cond = (overlaps_all_dims | has_negative_time).any(dim=0) # (num_cond,)
mask[self._num_generated_tokens :] = keep_cond
return mask
cond_starts = cond_positions[..., 0] # (B, 3, num_cond)
cond_ends = cond_positions[..., 1] # (B, 3, num_cond)
# Broadcast: (1, B, 3, num_cond) vs (num_tiles, B, 3, 1) -> (num_tiles, B, 3, num_cond).
overlaps = (cond_starts[None] < tile_ends[..., None]) & (cond_ends[None] > tile_starts[..., None])
overlaps_all_dims = overlaps.all(dim=2) # (num_tiles, B, num_cond)
has_negative_time = (cond_positions[:, 0, :, 0] < 0)[None] # (1, B, num_cond)
return (overlaps_all_dims | has_negative_time).any(dim=1) # (num_tiles, num_cond)
packages/ltx-core/src/ltx_core/model/transformer/attention.py
+339 -44
View File
@@ -1,12 +1,38 @@
import functools
import logging
from dataclasses import dataclass, field
from enum import Enum
from typing import Protocol
import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
from ltx_core.model.transformer.ops import (
GatedAttentionCallable,
PreAttentionCallable,
PytorchGatedAttention,
PytorchPreAttention,
)
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.
Used as the default for ``PytorchAttention`` so the wrapper-always
code path matches torch's native dispatch order without hard-coding it
(which would drift if torch updates the default).
``torch._C._get_sdp_priority_order`` is a private API; we accept that
risk because the project pins ``torch`` in the lockfile, so any
rename/removal surfaces on a controlled torch bump rather than silently.
"""
return [SDPBackend(p) for p in torch._C._get_sdp_priority_order()]
from ltx_core.model.transformer.rope import LTXRopeType, apply_rotary_emb
memory_efficient_attention = None
flash_attn_interface = None
flash_attn_4_func = None
try:
from xformers.ops import memory_efficient_attention
except ImportError:
@@ -17,15 +43,44 @@ try:
import flash_attn_interface
except ImportError:
flash_attn_interface = None
try:
from flash_attn.cute import flash_attn_func as flash_attn_4_func
except ImportError:
flash_attn_4_func = None
class AttentionCallable(Protocol):
"""Unmasked attention. Backends without a mask kernel (FA3/FA4) implement only
this protocol; backends that support masks too (Pytorch/SDPA, xFormers) are
structurally usable here and as :class:`MaskedAttentionCallable`."""
def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int) -> torch.Tensor: ...
class MaskedAttentionCallable(Protocol):
"""Masked attention. Mask is required (not optional) -- the caller has already
decided this is the masked path and chosen a backend that can serve it. Used
by :class:`Attention` when its forward receives a non-None ``mask``."""
def __call__(
self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor
) -> torch.Tensor: ...
class PytorchAttention(AttentionCallable):
def __init__(self, priority: list[SDPBackend] | None = None) -> None:
# priority=None -> snapshot torch's default SDPA priority at construction.
# Always passed through ``sdpa_kernel(..., set_priority=True)`` so the
# call site is uniform regardless of how the priority was chosen.
self._priority = priority if priority is not None else _torch_default_sdpa_priority()
@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."""
return f"SDPA[{'>'.join(b.name for b in self._priority)}]"
def __call__(
self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
) -> torch.Tensor:
@@ -41,12 +96,17 @@ class PytorchAttention(AttentionCallable):
if mask.ndim == 3:
mask = mask.unsqueeze(1)
out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
with sdpa_kernel(self._priority, set_priority=True):
out = torch.nn.functional.scaled_dot_product_attention(
q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False
)
out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
return out
class XFormersAttention(AttentionCallable):
label = "xFormers"
def __call__(
self,
q: torch.Tensor,
@@ -92,13 +152,14 @@ class XFormersAttention(AttentionCallable):
class FlashAttention3(AttentionCallable):
label = "FlashAttention3"
def __call__(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
heads: int,
mask: torch.Tensor | None = None,
) -> torch.Tensor:
if flash_attn_interface is None:
raise RuntimeError("FlashAttention3 was selected but `FlashAttention3` is not installed.")
@@ -108,32 +169,275 @@ class FlashAttention3(AttentionCallable):
q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v))
if mask is not None:
raise NotImplementedError("Mask is not supported for FlashAttention3")
out = flash_attn_interface.flash_attn_func(q.to(v.dtype), k.to(v.dtype), v)
out = out.reshape(b, -1, heads * dim_head)
return out
class FlashAttention4(AttentionCallable):
label = "FlashAttention4"
def __call__(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
heads: int,
) -> torch.Tensor:
if flash_attn_4_func is None:
raise RuntimeError("FlashAttention4 was selected but `flash-attn-4` is not installed.")
b, _, dim_head = q.shape
dim_head //= heads
q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v))
out, _ = flash_attn_4_func(q.to(v.dtype), k.to(v.dtype), v)
out = out.reshape(b, -1, heads * dim_head)
return out
# --- 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).
def _sdpa_can_use(backend: SDPBackend, *, with_mask: bool) -> bool:
"""Ask torch whether *backend* can run with the given mask shape.
``MATH`` is the universal SDPA fallback (pure PyTorch ops, no kernel
requirements) so it returns True everywhere, CPU included. The other
backends use ``torch.backends.cuda.can_use_*`` capability checks (no GPU
compute, no synchronization) and are False without CUDA. The probe shapes
are small but realistic enough to surface constraints (head dim, dtype)
that the per-backend rules care about.
"""
if backend is SDPBackend.MATH:
return True
if not torch.cuda.is_available():
return False
q = torch.empty(1, 4, 128, 64, device="cuda", dtype=torch.bfloat16)
k = torch.empty(1, 4, 128, 64, device="cuda", dtype=torch.bfloat16)
v = torch.empty(1, 4, 128, 64, device="cuda", dtype=torch.bfloat16)
mask = torch.zeros(1, 4, 128, 128, device="cuda", dtype=torch.bfloat16) if with_mask else None
params = torch.backends.cuda.SDPAParams(q, k, v, mask, 0.0, False, False)
if backend is SDPBackend.CUDNN_ATTENTION:
return torch.backends.cuda.can_use_cudnn_attention(params, debug=False)
if backend is SDPBackend.FLASH_ATTENTION:
return torch.backends.cuda.can_use_flash_attention(params, debug=False)
if backend is SDPBackend.EFFICIENT_ATTENTION:
return torch.backends.cuda.can_use_efficient_attention(params, debug=False)
return False
_SDPA_FULL_PRIORITY: tuple[SDPBackend, ...] = (
SDPBackend.CUDNN_ATTENTION,
SDPBackend.FLASH_ATTENTION,
SDPBackend.EFFICIENT_ATTENTION,
SDPBackend.MATH,
)
def _sdpa_full_priority() -> PytorchAttention:
"""Hand SDPA the full backend priority order; let torch's dispatcher pick at call time.
``sdpa_kernel(_SDPA_FULL_PRIORITY, set_priority=True)`` enables all four
backends and orders them; torch then walks the order at call time and picks
the first backend whose ``can_use_*`` check passes for the actual
shapes/dtype/mask. FLASH is rejected automatically when a mask is present;
CUDNN may be rejected under deterministic mode; MATH is the universal
fallback. Probing per-backend usability up front from generic probe shapes
cannot anticipate the variety of real call sites (e.g. broadcast key-only
masks, large head dim), so we defer the choice to the dispatcher.
"""
return PytorchAttention(priority=list(_SDPA_FULL_PRIORITY))
def _select_primary_attention() -> AttentionCallable:
"""Pick the fastest unmasked attention based on installed extras and GPU arch.
Priority by arch:
- Hopper (sm_90, H100): FA3 / xFormers (mutually exclusive at import) > FA4 > SDPA.
- Datacenter Blackwell (sm_100, B200): FA4 > SDPA. FA4 is intentionally *not*
picked on consumer Blackwell (sm_120) -- known regressions in newer
FA4 betas; users who want it on sm_120 must opt in explicitly.
- Everywhere else (Ada, Ampere, CPU): SDPA with the full backend priority
list -- torch's runtime dispatcher picks the best fit at call time.
"""
if torch.cuda.is_available():
major, _ = torch.cuda.get_device_capability(0)
if major == 9:
if flash_attn_interface is not None:
return FlashAttention3()
if memory_efficient_attention is not None:
return XFormersAttention()
if flash_attn_4_func is not None:
return FlashAttention4()
if major == 10 and flash_attn_4_func is not None:
return FlashAttention4()
return _sdpa_full_priority()
def _select_masked_attention() -> MaskedAttentionCallable:
"""Pick a mask-aware attention. Prefers xFormers when installed; else SDPA with
the full priority list (the dispatcher rejects FLASH automatically when a
mask is present and walks past it)."""
if memory_efficient_attention is not None:
return XFormersAttention()
return _sdpa_full_priority()
@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
@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
def _resolve_sdpa_variant(backend: SDPBackend, name: str, *, with_mask: bool) -> PytorchAttention:
"""Build a single-backend ``PytorchAttention`` pin, raising if the backend
can't actually serve the call on this machine. Used by both
:meth:`AttentionFunction.to_callable` and :meth:`MaskedAttentionFunction.to_callable`;
``with_mask`` differs between the two so the capability check considers
the protocol the caller intends to use. Not used for ``MATH`` -- MATH is
the universal fallback and would falsely fail the CUDA-only probe on CPU.
"""
if not _sdpa_can_use(backend, with_mask=with_mask):
raise RuntimeError(
f"{name} selected but the SDPA {backend.name} backend is not usable on this machine "
"(either no CUDA, the backend rejected the probe shapes, or "
"torch.use_deterministic_algorithms(True) excluded it)."
)
return PytorchAttention(priority=[backend])
class AttentionFunction(Enum):
PYTORCH = "pytorch"
XFORMERS = "xformers"
FLASH_ATTENTION_3 = "flash_attention_3"
DEFAULT = "default"
FLASH_ATTENTION_4 = "flash_attention_4"
SDPA_CUDNN = "sdpa_cudnn"
SDPA_FLASH = "sdpa_flash"
SDPA_EFFICIENT = "sdpa_efficient"
SDPA_MATH = "sdpa_math"
# Pick the fastest unmasked backend for the current GPU/extras combo; see
# :func:`automatic_attention`. Default for :class:`AttentionOps`.
AUTOMATIC = "automatic"
def to_callable(self) -> AttentionCallable:
def to_callable(self) -> AttentionCallable: # noqa: PLR0911
"""Resolve to a concrete callable. Use this at module init time so that
torch.compile can trace through the attention call without graph breaks."""
if self is AttentionFunction.PYTORCH:
return PytorchAttention()
elif self is AttentionFunction.XFORMERS:
return XFormersAttention()
elif self is AttentionFunction.FLASH_ATTENTION_3:
return FlashAttention3()
else:
# Default behavior: XFormers if installed else - PyTorch
return XFormersAttention() if memory_efficient_attention is not None else PytorchAttention()
torch.compile can trace through the attention call without graph breaks.
Every non-AUTOMATIC variant raises :class:`RuntimeError` when the backend
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.
"""
match self:
case AttentionFunction.AUTOMATIC:
return automatic_attention()
case AttentionFunction.PYTORCH:
return PytorchAttention()
case AttentionFunction.XFORMERS:
if memory_efficient_attention is None:
raise RuntimeError("AttentionFunction.XFORMERS selected but `xformers` is not installed.")
return XFormersAttention()
case AttentionFunction.FLASH_ATTENTION_3:
if flash_attn_interface is None:
raise RuntimeError(
"AttentionFunction.FLASH_ATTENTION_3 selected but `flash-attn-3` is not installed."
)
return FlashAttention3()
case AttentionFunction.FLASH_ATTENTION_4:
if flash_attn_4_func is None:
raise RuntimeError(
"AttentionFunction.FLASH_ATTENTION_4 selected but `flash-attn-4` is not installed."
)
return FlashAttention4()
case AttentionFunction.SDPA_MATH:
return PytorchAttention(priority=[SDPBackend.MATH])
case AttentionFunction.SDPA_CUDNN:
return _resolve_sdpa_variant(
SDPBackend.CUDNN_ATTENTION, "AttentionFunction.SDPA_CUDNN", with_mask=False
)
case AttentionFunction.SDPA_FLASH:
return _resolve_sdpa_variant(
SDPBackend.FLASH_ATTENTION, "AttentionFunction.SDPA_FLASH", with_mask=False
)
case AttentionFunction.SDPA_EFFICIENT:
return _resolve_sdpa_variant(
SDPBackend.EFFICIENT_ATTENTION, "AttentionFunction.SDPA_EFFICIENT", with_mask=False
)
class MaskedAttentionFunction(Enum):
"""Backends usable on the masked path. Mirrors :class:`AttentionFunction` minus
the variants the torch SDPA dispatcher (or the wrapped kernel) rejects with a
mask: ``SDPA_FLASH`` -- FLASH kernel cannot serve an additive ``attn_mask``;
``FLASH_ATTENTION_3``/``FLASH_ATTENTION_4`` -- neither has a mask kernel at all.
Keeping them out makes "this backend cannot mask" a type error, not a runtime one."""
PYTORCH = "pytorch"
XFORMERS = "xformers"
SDPA_CUDNN = "sdpa_cudnn"
SDPA_EFFICIENT = "sdpa_efficient"
SDPA_MATH = "sdpa_math"
# Pick the fastest mask-capable backend for the current extras combo; see
# :func:`automatic_masked_attention`. Default for the masked slot of
# :class:`AttentionOps`.
AUTOMATIC = "automatic"
def to_callable(self) -> MaskedAttentionCallable:
"""Resolve to a concrete masked callable. Same backend classes as
:meth:`AttentionFunction.to_callable`; the protocol returned just exposes
the masked call signature.
Non-AUTOMATIC variants raise :class:`RuntimeError` when the backend isn't
usable for the masked path on this machine. SDPA probes run with
``with_mask=True`` so the capability check considers the protocol the
caller will actually use."""
match self:
case MaskedAttentionFunction.AUTOMATIC:
return automatic_masked_attention()
case MaskedAttentionFunction.PYTORCH:
return PytorchAttention()
case MaskedAttentionFunction.XFORMERS:
if memory_efficient_attention is None:
raise RuntimeError("MaskedAttentionFunction.XFORMERS selected but `xformers` is not installed.")
return XFormersAttention()
case MaskedAttentionFunction.SDPA_MATH:
return PytorchAttention(priority=[SDPBackend.MATH])
case MaskedAttentionFunction.SDPA_CUDNN:
return _resolve_sdpa_variant(
SDPBackend.CUDNN_ATTENTION, "MaskedAttentionFunction.SDPA_CUDNN", with_mask=True
)
case MaskedAttentionFunction.SDPA_EFFICIENT:
return _resolve_sdpa_variant(
SDPBackend.EFFICIENT_ATTENTION, "MaskedAttentionFunction.SDPA_EFFICIENT", with_mask=True
)
@dataclass(frozen=True)
class AttentionOps:
"""Pluggable callables consumed by :class:`Attention`."""
attention_function: AttentionCallable = field(default_factory=lambda: AttentionFunction.AUTOMATIC.to_callable())
masked_attention_function: MaskedAttentionCallable = field(
default_factory=lambda: MaskedAttentionFunction.AUTOMATIC.to_callable()
)
preattention_function: PreAttentionCallable = field(default_factory=PytorchPreAttention)
gated_attention_function: GatedAttentionCallable = field(default_factory=PytorchGatedAttention)
class Attention(torch.nn.Module):
@@ -145,16 +449,17 @@ class Attention(torch.nn.Module):
dim_head: int = 64,
norm_eps: float = 1e-6,
rope_type: LTXRopeType = LTXRopeType.SPLIT,
attention_function: AttentionCallable | AttentionFunction = AttentionFunction.DEFAULT,
ops: AttentionOps | None = None,
apply_gated_attention: bool = False,
) -> None:
super().__init__()
if ops is None:
ops = AttentionOps()
self.rope_type = rope_type
self.attention_function = (
attention_function.to_callable()
if isinstance(attention_function, AttentionFunction)
else attention_function
)
self.attention_function = ops.attention_function
self.masked_attention_function = ops.masked_attention_function
self.preattention_function = ops.preattention_function
self.gated_attention_function = ops.gated_attention_function
inner_dim = dim_head * heads
context_dim = query_dim if context_dim is None else context_dim
@@ -196,7 +501,9 @@ class Attention(torch.nn.Module):
context: Key/value context tensor of shape ``(B, S, context_dim)``.
Falls back to ``x`` (self-attention) when *None*.
mask: Optional attention mask. Interpretation depends on the attention
backend (additive bias for xformers/PyTorch SDPA).
backend (additive bias for xformers/PyTorch SDPA). A non-None
``mask`` routes to ``masked_attention_function``; ``None`` keeps
the unmasked path.
pe: Rotary positional embeddings applied to both ``q`` and ``k``.
k_pe: Separate rotary positional embeddings for ``k`` only. When
*None*, ``pe`` is reused for keys.
@@ -221,29 +528,17 @@ class Attention(torch.nn.Module):
else:
q = self.to_q(x)
k = self.to_k(context)
q = self.q_norm(q)
k = self.k_norm(k)
if pe is not None:
q = apply_rotary_emb(q, pe, self.rope_type)
k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type)
out = self.attention_function(q, k, v, self.heads, mask) # (B, T, H*D)
q, k = self.preattention_function(q, k, self, mask, pe, k_pe)
if mask is None:
out = self.attention_function(q, k, v, self.heads) # (B, T, H*D)
else:
out = self.masked_attention_function(q, k, v, self.heads, mask)
if perturbation_mask is not None:
out = out * perturbation_mask + v * (1 - perturbation_mask)
# Apply per-head gating if enabled
if self.to_gate_logits is not None:
gate_logits = self.to_gate_logits(x) # (B, T, H)
b, t, _ = out.shape
# Reshape to (B, T, H, D) for per-head gating
out = out.view(b, t, self.heads, self.dim_head)
# Apply gating: 2 * sigmoid(x) so that zero-init gives identity (2 * 0.5 = 1.0)
gates = 2.0 * torch.sigmoid(gate_logits) # (B, T, H)
out = out * gates.unsqueeze(-1) # (B, T, H, D) * (B, T, H, 1)
# Reshape back to (B, T, H*D)
out = out.view(b, t, self.heads * self.dim_head)
out = self.gated_attention_function(x, out, self)
return self.to_out(out)
packages/ltx-core/src/ltx_core/model/transformer/compiling.py
+106 -11
View File
@@ -1,19 +1,112 @@
from dataclasses import dataclass, field
from typing import Any
import torch
from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.sd_ops import SDOps
from ltx_core.model.transformer.model import LTXModel
from ltx_core.model.transformer.transformer_args import BlockPerturbationsProcessor, TransformerArgs
# 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_DYNAMO_CONFIG: dict[str, Any] = {"inline_inbuilt_nn_modules": True, "cache_size_limit": 256}
def compile_transformer(model: LTXModel) -> LTXModel:
model.transformer_blocks = torch.nn.ModuleList(torch.compile(m) for m in model.transformer_blocks)
@dataclass(frozen=True)
class CompilationConfig:
"""``torch.compile`` configuration for transformer blocks. ``None`` keeps eager."""
mode: str | None = None
backend: str = "inductor"
fullgraph: bool = False
dynamic: bool | None = None
inductor_config: dict[str, Any] = field(default_factory=lambda: dict(_DEFAULT_INDUCTOR_CONFIG))
dynamo_config: dict[str, Any] = field(default_factory=lambda: dict(_DEFAULT_DYNAMO_CONFIG))
class _SeqDynamicMarkingProcessor:
"""Marks the per-block seq dim dynamic, then delegates to an inner processor.
Installed by ``compile_transformer`` so the per-block compile artifact stays
shape-polymorphic. Wraps whatever ``block_input_processor`` was already on
the model -- callers that customised the processor keep their customisation;
only the seq-dim marking is layered on top. Lives outside the compiled
region, so ``mark_dynamic`` runs in eager mode on the tensors that are
about to cross into the trace.
"""
def __init__(self, inner: BlockPerturbationsProcessor) -> None:
self.inner = inner
def __call__(
self,
args: TransformerArgs,
perturbations: BatchedPerturbationConfig,
block_idx: int,
self_attn_type: PerturbationType,
cross_attn_type: PerturbationType,
) -> TransformerArgs:
# Positional embeddings are second-from-last regardless of rope type:
# split rope is (B, H, T, D//2) -- dim -2 == 2; interleaved rope is (B, T, D)
# -- dim -2 == 1. Both work via the negative index.
torch._dynamo.mark_dynamic(args.x, 1)
cos, sin = args.positional_embeddings
torch._dynamo.mark_dynamic(cos, cos.ndim - 2)
torch._dynamo.mark_dynamic(sin, sin.ndim - 2)
if args.cross_positional_embeddings is not None:
cross_cos, cross_sin = args.cross_positional_embeddings
torch._dynamo.mark_dynamic(cross_cos, cross_cos.ndim - 2)
torch._dynamo.mark_dynamic(cross_sin, cross_sin.ndim - 2)
if args.self_attention_mask is not None:
# Dense form is (B, 1, T, T); key-padding form (from the SP wrapper)
# is (B, 1, 1, T) -- leave the size-1 query dim static so Dynamo
# keeps the broadcast.
if args.self_attention_mask.shape[2] > 1:
torch._dynamo.mark_dynamic(args.self_attention_mask, 2)
torch._dynamo.mark_dynamic(args.self_attention_mask, 3)
if args.context_mask is not None:
torch._dynamo.mark_dynamic(args.context_mask, 2)
# `timesteps` / `embedded_timestep` are per-token when conditioning sets a
# per-position denoise mask, in which case their dim 1 equals the seq length
# and must vary with it. When they're a single timestep broadcast across the
# sequence (dim 1 == 1), leaving them static lets Dynamo keep the size-1
# broadcast.
if args.timesteps.shape[1] > 1:
torch._dynamo.mark_dynamic(args.timesteps, 1)
if args.embedded_timestep.shape[1] > 1:
torch._dynamo.mark_dynamic(args.embedded_timestep, 1)
# `cross_scale_shift_timestep` is the cross-attn AdaLN scale/shift input
# derived from the own-modality per-token timesteps (denoise_mask * sigma),
# so its dim 1 equals the seq length when conditioning is per-token.
# `cross_gate_timestep` is the cross-modality sigma scalar -- dim 1 is 1
# and broadcasts, leave it static. Same guard pattern as `timesteps`.
if args.cross_scale_shift_timestep is not None and args.cross_scale_shift_timestep.shape[1] > 1:
torch._dynamo.mark_dynamic(args.cross_scale_shift_timestep, 1)
return self.inner(args, perturbations, block_idx, self_attn_type, cross_attn_type)
def compile_transformer(model: LTXModel, config: CompilationConfig) -> LTXModel:
"""Compile each transformer block via ``torch.compile`` with the given settings.
The patched forward emits ``torch.compiler.cudagraph_mark_step_begin()`` once
per step. Under CUDA-graph-enabling modes (``"reduce-overhead"`` /
``"max-autotune"``) this overrides Dynamo's per-invocation auto-mark
heuristic, which would otherwise fire once per compiled block call (48 per
forward) and treat each block call as a fresh iteration. Under other modes
the mark is a no-op (decrements an unread counter).
"""
model.transformer_blocks = torch.nn.ModuleList(
torch.compile(m, mode=config.mode, backend=config.backend, fullgraph=config.fullgraph, dynamic=config.dynamic)
for m in model.transformer_blocks
)
model.block_input_processor = _SeqDynamicMarkingProcessor(inner=model.block_input_processor)
def patched_dynamo_forward(*args, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
torch.compiler.cudagraph_mark_step_begin()
with (
torch._inductor.config.patch(unsafe_skip_cache_dynamic_shape_guards=True),
torch._dynamo.config.patch( # type: ignore[attr-defined]
inline_inbuilt_nn_modules=True, cache_size_limit=256, allow_unspec_int_on_nn_module=True
),
torch._inductor.config.patch(**config.inductor_config),
torch._dynamo.config.patch(**config.dynamo_config), # type: ignore[attr-defined]
):
return model.forward_without_compilation(*args, **kwargs)
@@ -22,11 +115,13 @@ def compile_transformer(model: LTXModel) -> LTXModel:
return model
COMPILE_TRANSFORMER = ModuleOps(
name="compile_transformer",
matcher=lambda model: isinstance(model, LTXModel),
mutator=lambda model: compile_transformer(model),
)
def build_compile_transformer_op(config: CompilationConfig) -> ModuleOps:
"""Build a ``ModuleOps`` that compiles transformer blocks with the given settings."""
return ModuleOps(
name="compile_transformer",
matcher=lambda model: isinstance(model, LTXModel),
mutator=lambda model: compile_transformer(model, config),
)
def modify_sd_ops_for_compilation(original_sd_ops: SDOps, number_of_blocks: int = 48) -> SDOps:
packages/ltx-core/src/ltx_core/model/transformer/modality.py
+17 -5
View File
@@ -17,8 +17,19 @@ class Modality:
the batch size, *T* is the total number of tokens (noisy +
conditioning), and *D* is the input dimension.
timesteps: Per-token timestep embeddings, shape ``(B, T)``.
positions: Positional coordinates, shape ``(B, 3, T)`` for video
(time, height, width) or ``(B, 1, T)`` for audio.
positions: Per-token patch coordinates used to build the RoPE
frequencies. With the default ``use_middle_indices_grid=True``,
shape is ``(B, n_pos_dims, T, 2)`` where ``n_pos_dims=3`` for
video (time, height, width) and ``n_pos_dims=1`` for audio
(time); the last dim of size 2 holds the ``[start, end)``
index bounds of each patch, and RoPE is evaluated at the
*middle* of that range -- hence the flag name. Taking the
patch midpoint produces a smoother and more accurate
positional signal than indexing by the patch's start when
patches span more than one spatial / temporal unit.
When ``use_middle_indices_grid=False``, the legacy 3-D form
``(B, n_pos_dims, T)`` of integer positional indices is
accepted instead and used as-is (no midpoint derivation).
context: Text conditioning embeddings from the prompt encoder.
enabled: Whether this modality is active in the current forward pass.
context_mask: Optional mask for the text context tokens.
@@ -34,9 +45,10 @@ class Modality:
) # Shape: (B, T, D) where B is the batch size, T is the number of tokens, and D is input dimension
sigma: torch.Tensor # Shape: (B,). Current sigma value, used for cross-attention timestep calculation.
timesteps: torch.Tensor # Shape: (B, T) where T is the number of timesteps
positions: (
torch.Tensor
) # Shape: (B, 3, T) for video, where 3 is the number of dimensions and T is the number of tokens
# Shape: (B, n_pos_dims, T, 2) by default (use_middle_indices_grid=True);
# n_pos_dims=3 for video, 1 for audio; last dim holds [start, end) patch bounds.
# Legacy form (B, n_pos_dims, T) when use_middle_indices_grid=False.
positions: torch.Tensor
context: torch.Tensor
enabled: bool = True
context_mask: torch.Tensor | None = None
packages/ltx-core/src/ltx_core/model/transformer/model.py
+46 -23
View File
@@ -2,13 +2,18 @@ from enum import Enum
import torch
from ltx_core.guidance.perturbations import BatchedPerturbationConfig
from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
from ltx_core.model.transformer.adaln import AdaLayerNormSingle, adaln_embedding_coefficient
from ltx_core.model.transformer.attention import AttentionCallable, AttentionFunction
from ltx_core.model.transformer.modality import Modality
from ltx_core.model.transformer.rope import LTXRopeType
from ltx_core.model.transformer.transformer import BasicAVTransformerBlock, TransformerConfig
from ltx_core.model.transformer.transformer import (
DEFAULT_TRANSFORMER_OPS,
BasicAVTransformerBlock,
TransformerConfig,
TransformerOpsConfig,
)
from ltx_core.model.transformer.transformer_args import (
BlockPerturbationsProcessor,
MultiModalTransformerArgsPreprocessor,
TransformerArgs,
TransformerArgsPreprocessor,
@@ -45,7 +50,7 @@ class LTXModel(torch.nn.Module):
num_layers: int = 48,
cross_attention_dim: int = 4096,
norm_eps: float = 1e-06,
attention_type: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT,
ops: TransformerOpsConfig = DEFAULT_TRANSFORMER_OPS,
positional_embedding_theta: float = 10000.0,
positional_embedding_max_pos: list[int] | None = None,
timestep_scale_multiplier: int = 1000,
@@ -115,9 +120,13 @@ class LTXModel(torch.nn.Module):
audio_attention_head_dim=audio_attention_head_dim if model_type.is_audio_enabled() else 0,
audio_cross_attention_dim=audio_cross_attention_dim,
norm_eps=norm_eps,
attention_type=attention_type,
ops=ops,
apply_gated_attention=apply_gated_attention,
)
# Hook for per-block input prep. Compile transforms in `compiling.py`
# wrap (not replace) this with a processor that also marks the seq dim
# dynamic, so any caller customisation here is preserved as the inner.
self.block_input_processor = BlockPerturbationsProcessor()
@property
def _adaln_embedding_coefficient(self) -> int:
@@ -284,7 +293,7 @@ class LTXModel(torch.nn.Module):
audio_attention_head_dim: int,
audio_cross_attention_dim: int,
norm_eps: float,
attention_type: AttentionFunction | AttentionCallable,
ops: TransformerOpsConfig,
apply_gated_attention: bool,
) -> None:
"""Initialize transformer blocks for LTX."""
@@ -315,14 +324,13 @@ class LTXModel(torch.nn.Module):
self.transformer_blocks = torch.nn.ModuleList(
[
BasicAVTransformerBlock(
idx=idx,
video=video_config,
audio=audio_config,
rope_type=self.rope_type,
norm_eps=norm_eps,
attention_function=attention_type,
ops=ops,
)
for idx in range(num_layers)
for _ in range(num_layers)
]
)
@@ -340,29 +348,44 @@ class LTXModel(torch.nn.Module):
self,
video: TransformerArgs | None,
audio: TransformerArgs | None,
perturbations: BatchedPerturbationConfig,
) -> tuple[TransformerArgs, TransformerArgs]:
"""Process transformer blocks for LTXAV."""
perturbations: BatchedPerturbationConfig | None,
) -> tuple[TransformerArgs | None, TransformerArgs | None]:
"""Process transformer blocks for LTXAV.
Per-block perturbation masks are precomputed here and attached to each
modality's ``TransformerArgs`` so the block forward has no per-block
identity to specialise on — all blocks share a single Dynamo cache slot.
"""
if perturbations is None:
batch_size = (video or audio).x.shape[0]
perturbations = BatchedPerturbationConfig.empty(batch_size)
for block_idx, block in enumerate(self.transformer_blocks):
if video is not None:
video = self.block_input_processor(
video,
perturbations,
block_idx,
self_attn_type=PerturbationType.SKIP_VIDEO_SELF_ATTN,
cross_attn_type=PerturbationType.SKIP_A2V_CROSS_ATTN,
)
if audio is not None:
audio = self.block_input_processor(
audio,
perturbations,
block_idx,
self_attn_type=PerturbationType.SKIP_AUDIO_SELF_ATTN,
cross_attn_type=PerturbationType.SKIP_V2A_CROSS_ATTN,
)
# Process transformer blocks
for block in self.transformer_blocks:
if self._enable_gradient_checkpointing and self.training:
# Use gradient checkpointing to save memory during training.
# With use_reentrant=False, we can pass dataclasses directly -
# PyTorch will track all tensor leaves in the computation graph.
video, audio = torch.utils.checkpoint.checkpoint(
block,
video,
audio,
perturbations,
use_reentrant=False,
)
else:
video, audio = block(
video=video,
audio=audio,
perturbations=perturbations,
)
video, audio = block(video=video, audio=audio)
return video, audio
packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py
+6 -5
View File
@@ -2,10 +2,10 @@ import torch
from ltx_core.loader.sd_ops import SDOps
from ltx_core.model.model_protocol import ModelConfigurator
from ltx_core.model.transformer.attention import AttentionFunction
from ltx_core.model.transformer.model import LTXModel, LTXModelType
from ltx_core.model.transformer.rope import LTXRopeType
from ltx_core.model.transformer.text_projection import create_caption_projection
from ltx_core.model.transformer.transformer import DEFAULT_TRANSFORMER_OPS, TransformerOpsConfig
from ltx_core.utils import check_config_value
@@ -16,7 +16,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
"""
@classmethod
def from_config(cls: type[LTXModel], config: dict) -> LTXModel:
def from_config(cls, config: dict, ops: TransformerOpsConfig = DEFAULT_TRANSFORMER_OPS) -> LTXModel:
# Build caption projections for 19B models (projection handled in transformer).
caption_projection, audio_caption_projection = _build_caption_projections(config, is_av=True)
@@ -40,6 +40,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
check_config_value(config, "share_ff", False)
check_config_value(config, "av_cross_ada_norm", True)
check_config_value(config, "use_middle_indices_grid", True)
check_config_value(config, "num_attention_heads", config.get("audio_num_attention_heads", float("nan")))
return LTXModel(
model_type=LTXModelType.AudioVideo,
@@ -50,7 +51,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
num_layers=config.get("num_layers", 48),
cross_attention_dim=config.get("cross_attention_dim", 4096),
norm_eps=config.get("norm_eps", 1e-06),
attention_type=AttentionFunction(config.get("attention_type", "default")),
ops=ops,
positional_embedding_theta=config.get("positional_embedding_theta", 10000.0),
positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]),
timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
@@ -78,7 +79,7 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
"""
@classmethod
def from_config(cls: type[LTXModel], config: dict) -> LTXModel:
def from_config(cls, config: dict, ops: TransformerOpsConfig = DEFAULT_TRANSFORMER_OPS) -> LTXModel:
# Build caption projection for 19B model (projection handled in transformer).
caption_projection, _ = _build_caption_projections(config, is_av=False)
@@ -109,7 +110,7 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
num_layers=config.get("num_layers", 48),
cross_attention_dim=config.get("cross_attention_dim", 4096),
norm_eps=config.get("norm_eps", 1e-06),
attention_type=AttentionFunction(config.get("attention_type", "default")),
ops=ops,
positional_embedding_theta=config.get("positional_embedding_theta", 10000.0),
positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]),
timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
packages/ltx-core/src/ltx_core/model/transformer/ops.py
+106
View File
@@ -0,0 +1,106 @@
from typing import List, Protocol
import torch
from torch import nn
from ltx_core.model.transformer.rope import apply_rotary_emb
from ltx_core.utils import rms_norm
class PreAttentionCallable(Protocol):
def __call__(
self,
q: torch.Tensor,
k: torch.Tensor,
attn_module: nn.Module,
mask: torch.Tensor | None,
pe: torch.Tensor | None,
k_pe: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor]: ...
class PytorchPreAttention(PreAttentionCallable):
def __call__(
self,
q: torch.Tensor,
k: torch.Tensor,
attn_module: nn.Module,
mask: torch.Tensor | None, # noqa: ARG002
pe: torch.Tensor | None,
k_pe: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor]:
q = attn_module.q_norm(q)
k = attn_module.k_norm(k)
if pe is not None:
q = apply_rotary_emb(q, pe, attn_module.rope_type)
k = apply_rotary_emb(k, pe if k_pe is None else k_pe, attn_module.rope_type)
return q, k
class AdaZeroCallable(Protocol):
def __call__(
self,
x: torch.Tensor,
eps: float,
scale: torch.Tensor,
shift: torch.Tensor,
) -> torch.Tensor: ...
class PytorchAdaZeroFunction(AdaZeroCallable):
def __call__(
self,
x: torch.Tensor,
eps: float,
scale: torch.Tensor,
shift: torch.Tensor,
) -> torch.Tensor:
return rms_norm(x, eps=eps) * (1 + scale) + shift
class PostSACallable(Protocol):
def __call__(
self,
x: torch.Tensor,
y: torch.Tensor,
norm_weights: torch.Tensor | None,
eps: float,
gate: torch.Tensor,
) -> List[torch.Tensor]: ...
class PytorchPostSAFunction(PostSACallable):
def __call__(
self,
x: torch.Tensor,
y: torch.Tensor,
norm_weights: torch.Tensor | None,
eps: float,
gate: torch.Tensor,
) -> List[torch.Tensor]:
x_fma = x + y * gate
return x_fma, rms_norm(x_fma, norm_weights, eps=eps)
class GatedAttentionCallable(Protocol):
def __call__(
self,
x: torch.Tensor,
attn_out: torch.Tensor,
attn_module: nn.Module,
) -> torch.Tensor: ...
class PytorchGatedAttention(GatedAttentionCallable):
def __call__(
self,
x: torch.Tensor,
attn_out: torch.Tensor,
attn_module: nn.Module,
) -> torch.Tensor:
gate_logits = attn_module.to_gate_logits(x) # (B, T, H)
b, t, _ = attn_out.shape
out = attn_out.view(b, t, attn_module.heads, attn_module.dim_head)
gates = 2.0 * torch.sigmoid(gate_logits) # (B, T, H)
out = out * gates.unsqueeze(-1) # (B, T, H, D) * (B, T, H, 1)
return out.view(b, t, attn_module.heads * attn_module.dim_head)
packages/ltx-core/src/ltx_core/model/transformer/rope.py
+21 -9
View File
@@ -43,16 +43,26 @@ def apply_interleaved_rotary_emb(
def apply_split_rotary_emb(
input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor
) -> torch.Tensor:
needs_reshape = False
if input_tensor.ndim != 4 and cos_freqs.ndim == 4:
b, h, t, _ = cos_freqs.shape
if input_tensor.shape[0] != b:
if sin_freqs.shape != cos_freqs.shape:
raise ValueError(
f"apply_split_rotary_emb: sin_freqs.shape {tuple(sin_freqs.shape)} must equal "
f"cos_freqs.shape {tuple(cos_freqs.shape)}."
)
needs_reshape = input_tensor.ndim != 4 and cos_freqs.ndim == 4
if needs_reshape:
b_freq = cos_freqs.shape[0]
h = cos_freqs.shape[1]
b_in = input_tensor.shape[0]
if b_freq not in (1, b_in):
raise ValueError(
f"apply_split_rotary_emb: input_tensor batch ({input_tensor.shape[0]}) "
f"must equal cos_freqs batch ({b})."
f"apply_split_rotary_emb: cos_freqs batch ({b_freq}) must be 1 "
f"(broadcast) or equal input_tensor batch ({b_in})."
)
input_tensor = input_tensor.reshape(b, t, h, -1).swapaxes(1, 2)
needs_reshape = True
# `unflatten` only touches the last dim, keeping the batch and seq dims as
# the input tensor's own symbolic ints under torch.compile. `reshape(b_in,
# t, h, -1)` would have forced Dynamo to specialise those dims because it
# cannot prove `b_in == cos_freqs.shape[0]` and `seq == t` across tensors.
input_tensor = input_tensor.unflatten(-1, (h, -1)).transpose(1, 2)
split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2)
first_half_input = split_input[..., :1, :]
@@ -67,7 +77,9 @@ def apply_split_rotary_emb(
output = rearrange(output, "... d r -> ... (d r)")
if needs_reshape:
output = output.swapaxes(1, 2).reshape(b, t, -1)
# `transpose(1, 2).flatten(-2)` keeps the batch and seq dims symbolic; using
# `reshape(b_in, t, -1)` would force Dynamo to specialise both axes.
output = output.transpose(1, 2).flatten(-2)
return output
packages/ltx-core/src/ltx_core/model/transformer/transformer.py
+116 -79
View File
@@ -1,11 +1,27 @@
from dataclasses import dataclass, replace
from dataclasses import dataclass, field, replace
import torch
from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
from ltx_core.model.transformer.adaln import adaln_embedding_coefficient
from ltx_core.model.transformer.attention import Attention, AttentionCallable, AttentionFunction
from ltx_core.model.transformer.attention import (
Attention,
AttentionCallable,
AttentionFunction,
AttentionOps,
MaskedAttentionCallable,
MaskedAttentionFunction,
)
from ltx_core.model.transformer.feed_forward import FeedForward
from ltx_core.model.transformer.ops import (
AdaZeroCallable,
GatedAttentionCallable,
PostSACallable,
PreAttentionCallable,
PytorchAdaZeroFunction,
PytorchGatedAttention,
PytorchPostSAFunction,
PytorchPreAttention,
)
from ltx_core.model.transformer.rope import LTXRopeType
from ltx_core.model.transformer.transformer_args import TransformerArgs
from ltx_core.utils import rms_norm
@@ -21,19 +37,68 @@ class TransformerConfig:
cross_attention_adaln: bool = False
@dataclass(frozen=True)
class TransformerOpsConfig:
"""Pluggable ops for :class:`BasicAVTransformerBlock`.
Use :meth:`from_functions` to construct from enum values or partial overrides
without spelling out a full :class:`AttentionOps`.
"""
attention_ops: AttentionOps = field(default_factory=AttentionOps)
ada_zero_function: AdaZeroCallable = field(default_factory=PytorchAdaZeroFunction)
post_sa_function: PostSACallable = field(default_factory=PytorchPostSAFunction)
@classmethod
def from_functions(
cls,
attention: AttentionFunction | AttentionCallable = AttentionFunction.AUTOMATIC,
masked_attention: MaskedAttentionFunction | MaskedAttentionCallable = MaskedAttentionFunction.AUTOMATIC,
preattention: PreAttentionCallable | None = None,
gated_attention: GatedAttentionCallable | None = None,
ada_zero: AdaZeroCallable | None = None,
post_sa: PostSACallable | None = None,
) -> "TransformerOpsConfig":
"""Build a config from individual functions or enums. Each *None* slot
falls back to the standard PyTorch implementation."""
attention_callable = attention.to_callable() if isinstance(attention, AttentionFunction) else attention
masked_callable = (
masked_attention.to_callable()
if isinstance(masked_attention, MaskedAttentionFunction)
else masked_attention
)
attention_ops = AttentionOps(
attention_function=attention_callable,
masked_attention_function=masked_callable,
preattention_function=preattention if preattention is not None else PytorchPreAttention(),
gated_attention_function=(gated_attention if gated_attention is not None else PytorchGatedAttention()),
)
return cls(
attention_ops=attention_ops,
ada_zero_function=ada_zero if ada_zero is not None else PytorchAdaZeroFunction(),
post_sa_function=post_sa if post_sa is not None else PytorchPostSAFunction(),
)
# Frozen, so safe to share as a default argument across callers that want the
# stock PyTorch ops without explicit construction.
DEFAULT_TRANSFORMER_OPS = TransformerOpsConfig()
class BasicAVTransformerBlock(torch.nn.Module):
def __init__(
self,
idx: int,
video: TransformerConfig | None = None,
audio: TransformerConfig | None = None,
rope_type: LTXRopeType = LTXRopeType.SPLIT,
norm_eps: float = 1e-6,
attention_function: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT,
ops: TransformerOpsConfig | None = None,
):
super().__init__()
self.idx = idx
if ops is None:
ops = TransformerOpsConfig()
self.ada_zero_function = ops.ada_zero_function
self.post_sa_function = ops.post_sa_function
if video is not None:
self.attn1 = Attention(
query_dim=video.dim,
@@ -42,7 +107,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
context_dim=None,
rope_type=rope_type,
norm_eps=norm_eps,
attention_function=attention_function,
ops=ops.attention_ops,
apply_gated_attention=video.apply_gated_attention,
)
self.attn2 = Attention(
@@ -52,7 +117,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
dim_head=video.d_head,
rope_type=rope_type,
norm_eps=norm_eps,
attention_function=attention_function,
ops=ops.attention_ops,
apply_gated_attention=video.apply_gated_attention,
)
self.ff = FeedForward(video.dim, dim_out=video.dim)
@@ -67,7 +132,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
context_dim=None,
rope_type=rope_type,
norm_eps=norm_eps,
attention_function=attention_function,
ops=ops.attention_ops,
apply_gated_attention=audio.apply_gated_attention,
)
self.audio_attn2 = Attention(
@@ -77,7 +142,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
dim_head=audio.d_head,
rope_type=rope_type,
norm_eps=norm_eps,
attention_function=attention_function,
ops=ops.attention_ops,
apply_gated_attention=audio.apply_gated_attention,
)
self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim)
@@ -93,7 +158,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
dim_head=audio.d_head,
rope_type=rope_type,
norm_eps=norm_eps,
attention_function=attention_function,
ops=ops.attention_ops,
apply_gated_attention=video.apply_gated_attention,
)
@@ -105,7 +170,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
dim_head=audio.d_head,
rope_type=rope_type,
norm_eps=norm_eps,
attention_function=attention_function,
ops=ops.attention_ops,
apply_gated_attention=audio.apply_gated_attention,
)
@@ -188,16 +253,10 @@ class BasicAVTransformerBlock(torch.nn.Module):
self,
video: TransformerArgs | None,
audio: TransformerArgs | None,
perturbations: BatchedPerturbationConfig | None = None,
) -> tuple[TransformerArgs | None, TransformerArgs | None]:
if video is None and audio is None:
raise ValueError("At least one of video or audio must be provided")
batch_size = (video or audio).x.shape[0]
if perturbations is None:
perturbations = BatchedPerturbationConfig.empty(batch_size)
vx = video.x if video is not None else None
ax = audio.x if audio is not None else None
@@ -211,28 +270,18 @@ class BasicAVTransformerBlock(torch.nn.Module):
vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(
self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3)
)
norm_vx = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa
norm_vx = self.ada_zero_function(vx, self.norm_eps, vscale_msa, vshift_msa)
del vshift_msa, vscale_msa
all_perturbed = perturbations.all_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx)
none_perturbed = not perturbations.any_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx)
v_mask = (
perturbations.mask_like(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx, vx)
if not all_perturbed and not none_perturbed
else None
vx_msa_out = self.attn1(
norm_vx,
pe=video.positional_embeddings,
mask=video.self_attention_mask,
perturbation_mask=video.self_attn_perturbation_mask,
all_perturbed=video.self_attn_all_perturbed,
)
vx = (
vx
+ self.attn1(
norm_vx,
pe=video.positional_embeddings,
mask=video.self_attention_mask,
perturbation_mask=v_mask,
all_perturbed=all_perturbed,
)
* vgate_msa
)
del vgate_msa, norm_vx, v_mask
vx = vx + vx_msa_out * vgate_msa
del vgate_msa, norm_vx, vx_msa_out
vx = vx + self._apply_text_cross_attention(
vx,
video.context,
@@ -250,27 +299,17 @@ class BasicAVTransformerBlock(torch.nn.Module):
self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3)
)
norm_ax = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa
norm_ax = self.ada_zero_function(ax, self.norm_eps, ascale_msa, ashift_msa)
del ashift_msa, ascale_msa
all_perturbed = perturbations.all_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx)
none_perturbed = not perturbations.any_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx)
a_mask = (
perturbations.mask_like(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx, ax)
if not all_perturbed and not none_perturbed
else None
ax_msa_out = self.audio_attn1(
norm_ax,
pe=audio.positional_embeddings,
mask=audio.self_attention_mask,
perturbation_mask=audio.self_attn_perturbation_mask,
all_perturbed=audio.self_attn_all_perturbed,
)
ax = (
ax
+ self.audio_attn1(
norm_ax,
pe=audio.positional_embeddings,
mask=audio.self_attention_mask,
perturbation_mask=a_mask,
all_perturbed=all_perturbed,
)
* agate_msa
)
del agate_msa, norm_ax, a_mask
ax = ax + ax_msa_out * agate_msa
del agate_msa, norm_ax, ax_msa_out
ax = ax + self._apply_text_cross_attention(
ax,
audio.context,
@@ -285,10 +324,11 @@ class BasicAVTransformerBlock(torch.nn.Module):
# Audio - Video cross attention.
if run_a2v or run_v2a:
vx_norm3 = rms_norm(vx, eps=self.norm_eps)
ax_norm3 = rms_norm(ax, eps=self.norm_eps)
if run_a2v and not perturbations.all_in_batch(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx):
# Snapshot vx/ax before A2V mutates vx; V2A's video keys/values must
# use the pre-A2V state so direction order doesn't bias the result.
vx_pre_av = vx
ax_pre_av = ax
if run_a2v and not video.cross_attn_skip_all:
scale_ca_video_a2v, shift_ca_video_a2v, gate_out_a2v = self.get_av_ca_ada_values(
self.scale_shift_table_a2v_ca_video,
vx.shape[0],
@@ -296,7 +336,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
video.cross_gate_timestep,
slice(0, 2),
)
vx_scaled = vx_norm3 * (1 + scale_ca_video_a2v) + shift_ca_video_a2v
a2v_vx_scaled = self.ada_zero_function(vx_pre_av, self.norm_eps, scale_ca_video_a2v, shift_ca_video_a2v)
del scale_ca_video_a2v, shift_ca_video_a2v
scale_ca_audio_a2v, shift_ca_audio_a2v, _ = self.get_av_ca_ada_values(
@@ -306,22 +346,21 @@ class BasicAVTransformerBlock(torch.nn.Module):
audio.cross_gate_timestep,
slice(0, 2),
)
ax_scaled = ax_norm3 * (1 + scale_ca_audio_a2v) + shift_ca_audio_a2v
a2v_ax_scaled = self.ada_zero_function(ax_pre_av, self.norm_eps, scale_ca_audio_a2v, shift_ca_audio_a2v)
del scale_ca_audio_a2v, shift_ca_audio_a2v
a2v_mask = perturbations.mask_like(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx, vx)
vx = vx + (
self.audio_to_video_attn(
vx_scaled,
context=ax_scaled,
a2v_vx_scaled,
context=a2v_ax_scaled,
pe=video.cross_positional_embeddings,
k_pe=audio.cross_positional_embeddings,
)
* gate_out_a2v
* a2v_mask
* video.cross_attn_perturbation_mask
)
del gate_out_a2v, a2v_mask, vx_scaled, ax_scaled
del gate_out_a2v, a2v_vx_scaled, a2v_ax_scaled
if run_v2a and not perturbations.all_in_batch(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx):
if run_v2a and not audio.cross_attn_skip_all:
scale_ca_audio_v2a, shift_ca_audio_v2a, gate_out_v2a = self.get_av_ca_ada_values(
self.scale_shift_table_a2v_ca_audio,
ax.shape[0],
@@ -329,7 +368,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
audio.cross_gate_timestep,
slice(2, 4),
)
ax_scaled = ax_norm3 * (1 + scale_ca_audio_v2a) + shift_ca_audio_v2a
v2a_ax_scaled = self.ada_zero_function(ax_pre_av, self.norm_eps, scale_ca_audio_v2a, shift_ca_audio_v2a)
del scale_ca_audio_v2a, shift_ca_audio_v2a
scale_ca_video_v2a, shift_ca_video_v2a, _ = self.get_av_ca_ada_values(
self.scale_shift_table_a2v_ca_video,
@@ -338,28 +377,26 @@ class BasicAVTransformerBlock(torch.nn.Module):
video.cross_gate_timestep,
slice(2, 4),
)
vx_scaled = vx_norm3 * (1 + scale_ca_video_v2a) + shift_ca_video_v2a
v2a_vx_scaled = self.ada_zero_function(vx_pre_av, self.norm_eps, scale_ca_video_v2a, shift_ca_video_v2a)
del scale_ca_video_v2a, shift_ca_video_v2a
v2a_mask = perturbations.mask_like(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx, ax)
ax = ax + (
self.video_to_audio_attn(
ax_scaled,
context=vx_scaled,
v2a_ax_scaled,
context=v2a_vx_scaled,
pe=audio.cross_positional_embeddings,
k_pe=video.cross_positional_embeddings,
)
* gate_out_v2a
* v2a_mask
* audio.cross_attn_perturbation_mask
)
del gate_out_v2a, v2a_mask, ax_scaled, vx_scaled
del vx_norm3, ax_norm3
del gate_out_v2a, v2a_vx_scaled, v2a_ax_scaled
del vx_pre_av, ax_pre_av
if run_vx:
vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(
self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, 6)
)
vx_scaled = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_mlp) + vshift_mlp
vx_scaled = self.ada_zero_function(vx, self.norm_eps, vscale_mlp, vshift_mlp)
vx = vx + self.ff(vx_scaled) * vgate_mlp
del vshift_mlp, vscale_mlp, vgate_mlp, vx_scaled
@@ -368,7 +405,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(
self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, 6)
)
ax_scaled = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_mlp) + ashift_mlp
ax_scaled = self.ada_zero_function(ax, self.norm_eps, ascale_mlp, ashift_mlp)
ax = ax + self.audio_ff(ax_scaled) * agate_mlp
del ashift_mlp, ascale_mlp, agate_mlp, ax_scaled
packages/ltx-core/src/ltx_core/model/transformer/transformer_args.py
+71 -17
View File
@@ -2,6 +2,7 @@ from dataclasses import dataclass, replace
import torch
from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
from ltx_core.model.transformer.adaln import AdaLayerNormSingle
from ltx_core.model.transformer.modality import Modality
from ltx_core.model.transformer.rope import (
@@ -19,8 +20,8 @@ class TransformerArgs:
context_mask: torch.Tensor
timesteps: torch.Tensor
embedded_timestep: torch.Tensor
positional_embeddings: torch.Tensor
cross_positional_embeddings: torch.Tensor | None
positional_embeddings: tuple[torch.Tensor, torch.Tensor]
cross_positional_embeddings: tuple[torch.Tensor, torch.Tensor] | None
cross_scale_shift_timestep: torch.Tensor | None
cross_gate_timestep: torch.Tensor | None
enabled: bool
@@ -28,6 +29,59 @@ class TransformerArgs:
self_attention_mask: torch.Tensor | None = (
None # Additive log-space self-attention bias (B, 1, T, T), None = full attention
)
# Per-block perturbation state, precomputed by `LTXModel._process_transformer_blocks`
# so the block forward needs no per-block identity. The bool shortcuts
# (`*_all_perturbed`, `cross_attn_skip_all`) are Python bools that Dynamo specialises
# on — fine because they're stable across denoising steps for a fixed perturbation
# config.
self_attn_perturbation_mask: torch.Tensor | None = None
self_attn_all_perturbed: bool = False
cross_attn_perturbation_mask: torch.Tensor | None = None
cross_attn_skip_all: bool = False
class BlockPerturbationsProcessor:
"""Per-block preparation of ``TransformerArgs``.
The base implementation returns a copy of ``args`` with this block's
precomputed perturbation flags and masks attached. Subclasses can layer in
operations that must run on each block's inputs but stay outside the
compile boundary -- e.g. ``torch._dynamo.mark_dynamic`` for
shape-polymorphic block compilation (see ``compiling.py``). Swapping the
processor on an ``LTXModel`` instance is how compile transforms opt in to
such behaviour without baking it into the model's forward.
``self_attn_perturbation_mask`` is None when all or none of the batch is
perturbed (the attention call can take the shortcut path). ``cross_attn_*``
is None when every sample skips the cross-attention entirely.
"""
def __call__(
self,
args: "TransformerArgs",
perturbations: BatchedPerturbationConfig,
block_idx: int,
self_attn_type: PerturbationType,
cross_attn_type: PerturbationType,
) -> "TransformerArgs":
device, dtype = args.x.device, args.x.dtype
all_self = perturbations.all_in_batch(self_attn_type, block_idx)
any_self = perturbations.any_in_batch(self_attn_type, block_idx)
self_mask: torch.Tensor | None = None
if any_self and not all_self:
self_mask = perturbations.mask(self_attn_type, block_idx, device, dtype).view(-1, 1, 1)
all_cross = perturbations.all_in_batch(cross_attn_type, block_idx)
cross_mask: torch.Tensor | None = None
if not all_cross:
cross_mask = perturbations.mask(cross_attn_type, block_idx, device, dtype).view(-1, 1, 1)
return replace(
args,
self_attn_perturbation_mask=self_mask,
self_attn_all_perturbed=all_self,
cross_attn_perturbation_mask=cross_mask,
cross_attn_skip_all=all_cross,
)
class TransformerArgsPreprocessor:
@@ -98,9 +152,12 @@ class TransformerArgsPreprocessor:
self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype
) -> torch.Tensor | None:
"""Prepare self-attention mask by converting [0,1] values to additive log-space bias.
Input shape: (B, T, T) with values in [0, 1].
Output shape: (B, 1, T, T) with 0.0 for full attention and a large negative value
for masked positions.
Input shape: 3D ``(B, T_q, T_k)`` with values in [0, 1]. The dense form
is ``(B, T, T)``; broadcastable forms like ``(1, 1, T)`` (key-only
padding) or ``(B, 1, T)`` are also valid and yield a correspondingly
broadcastable output.
Output shape: ``(B, 1, T_q, T_k)`` (heads dim inserted) with 0.0 for
full attention and a large negative value for masked positions.
Positions with attention_mask <= 0 are fully masked (mapped to the dtype's minimum
representable value). Strictly positive entries are converted via log-space for
smooth attenuation, with small values clamped for numerical stability.
@@ -120,7 +177,7 @@ class TransformerArgsPreprocessor:
if positive.any():
bias[positive] = torch.log(attention_mask[positive].clamp(min=eps)).to(x_dtype)
return bias.unsqueeze(1) # (B, 1, T, T) for head broadcast
return bias.unsqueeze(1) # (B, 1, T_q, T_k) for head broadcast
def _prepare_positional_embeddings(
self,
@@ -244,10 +301,6 @@ class MultiModalTransformerArgsPreprocessor:
if cross_modality.sigma.ndim != 1:
raise ValueError("Cross modality sigma must be a 1D tensor")
cross_timestep = cross_modality.sigma.view(
modality.timesteps.shape[0], 1, *[1] * len(modality.timesteps.shape[2:])
)
cross_pe = self.simple_preprocessor._prepare_positional_embeddings(
positions=modality.positions[:, 0:1, :],
inner_dim=self.audio_cross_attention_dim,
@@ -258,7 +311,8 @@ class MultiModalTransformerArgsPreprocessor:
)
cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep(
timestep=cross_timestep,
modality_timesteps=modality.timesteps,
cross_modality_sigma=cross_modality.sigma,
timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier,
batch_size=transformer_args.x.shape[0],
hidden_dtype=modality.latent.dtype,
@@ -273,23 +327,23 @@ class MultiModalTransformerArgsPreprocessor:
def _prepare_cross_attention_timestep(
self,
timestep: torch.Tensor | None,
modality_timesteps: torch.Tensor,
cross_modality_sigma: torch.Tensor,
timestep_scale_multiplier: int,
batch_size: int,
hidden_dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Prepare cross attention timestep embeddings."""
timestep = timestep * timestep_scale_multiplier
"""Prepare A-V cross-attention AdaLN inputs."""
av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier
scale_shift_timestep, _ = self.cross_scale_shift_adaln(
timestep.flatten(),
(modality_timesteps * timestep_scale_multiplier).flatten(),
hidden_dtype=hidden_dtype,
)
scale_shift_timestep = scale_shift_timestep.view(batch_size, -1, scale_shift_timestep.shape[-1])
gate_noise_timestep, _ = self.cross_gate_adaln(
timestep.flatten() * av_ca_factor,
(cross_modality_sigma * timestep_scale_multiplier * av_ca_factor).flatten(),
hidden_dtype=hidden_dtype,
)
gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1])
packages/ltx-core/src/ltx_core/model/video_vae/memory_efficient_decode.py
+55 -5
View File
@@ -56,6 +56,25 @@ if TYPE_CHECKING:
# ---------------------------------------------------------------------------
def _memory_format_of(t: torch.Tensor, prefer_channels_last_3d: bool = False) -> torch.memory_format:
"""Pick the memory format for a workspace allocation.
When ``prefer_channels_last_3d`` is True and ``t`` is 5D, return
``channels_last_3d`` regardless of ``t``'s current strides -- the
workspace's ``.copy_(t)`` will transcribe the data into the new layout.
This is needed because intermediate tensors inside the decoder (after
``rearrange`` + slice + residual add in ``_upsample_forward_efficient``)
are not NHWC-contiguous, so an auto-detect helper would silently fall
back to NCHW for every workspace after the first upsample.
Otherwise fall back to inspecting ``t``: ``channels_last_3d`` if ``t``
already uses it, else contiguous.
"""
if prefer_channels_last_3d and t.dim() == 5:
return torch.channels_last_3d
if t.dim() == 5 and t.is_contiguous(memory_format=torch.channels_last_3d):
return torch.channels_last_3d
return torch.contiguous_format
def _find_temporal_split_size(num_frames: int) -> int:
"""Find chunk size for in-place temporal convolution.
The chunk size ensures the last chunk has at least 3 frames
@@ -132,6 +151,7 @@ def inplace_conv3d_temporal_chunked(workspace: torch.Tensor, conv: nn.Conv3d) ->
workspace.shape[4],
device=workspace.device,
dtype=workspace.dtype,
memory_format=_memory_format_of(workspace),
)
o_buf = torch.empty_like(x_buf)
@@ -190,6 +210,7 @@ def _causal_pad(x: torch.Tensor, pad_size: int) -> torch.Tensor:
x.shape[4],
device=x.device,
dtype=x.dtype,
memory_format=_memory_format_of(x),
)
padded[:, :, pad_size:].copy_(x)
for i in range(pad_size):
@@ -325,6 +346,7 @@ def _midblock_forward_efficient(
causal: bool,
timestep: torch.Tensor | None,
generator: torch.Generator | None,
prefer_channels_last_3d: bool = False,
) -> torch.Tensor:
"""Memory-efficient ``UNetMidBlock3D`` forward.
Allocates a single workspace buffer that is reused across all
@@ -351,6 +373,7 @@ def _midblock_forward_efficient(
hidden_states.shape[4],
device=hidden_states.device,
dtype=hidden_states.dtype,
memory_format=_memory_format_of(hidden_states, prefer_channels_last_3d),
)
for resnet in block.res_blocks:
@@ -366,6 +389,7 @@ def _upsample_forward_efficient(
block: DepthToSpaceUpsample,
x: torch.Tensor,
causal: bool,
prefer_channels_last_3d: bool = False,
) -> torch.Tensor:
"""Memory-efficient ``DepthToSpaceUpsample`` forward.
For non-causal mode the input is copied into a workspace and the
@@ -393,6 +417,7 @@ def _upsample_forward_efficient(
if causal:
x = _causal_pad_free_and_conv(x, block.conv)
else:
mem_fmt = _memory_format_of(x, prefer_channels_last_3d)
workspace = torch.empty(
x.shape[0],
max(in_channels, out_channels),
@@ -401,11 +426,12 @@ def _upsample_forward_efficient(
x.shape[4],
device=x.device,
dtype=x.dtype,
memory_format=mem_fmt,
)
workspace[:, :in_channels, 1:-1].copy_(x)
del x
inplace_conv3d_temporal_chunked(workspace, conv)
x = workspace[:, :out_channels, 1:-1].contiguous()
x = workspace[:, :out_channels, 1:-1].contiguous(memory_format=mem_fmt)
del workspace
x = rearrange(
@@ -418,7 +444,7 @@ def _upsample_forward_efficient(
if block.stride[0] == 2:
x = x[:, :, 1:, :, :]
if block.residual:
x = x + x_in
x.add_(x_in)
del x_in
return x
@@ -434,12 +460,14 @@ def _final_norm_and_conv_out(
causal: bool,
scaled_timestep: torch.Tensor | None,
batch_size: int,
prefer_channels_last_3d: bool = False,
) -> torch.Tensor:
"""Workspace-based final norm + [ada] + SiLU + conv_out + unpatchify."""
conv_out_mod: CausalConv3d = decoder.conv_out # type: ignore[assignment]
conv_out = conv_out_mod.conv
feature_channels = sample.shape[1]
mem_fmt = _memory_format_of(sample, prefer_channels_last_3d)
workspace = torch.empty(
sample.shape[0],
max(feature_channels, conv_out.out_channels),
@@ -448,6 +476,7 @@ def _final_norm_and_conv_out(
sample.shape[4],
device=sample.device,
dtype=sample.dtype,
memory_format=mem_fmt,
)
workspace[:, :feature_channels, 1:-1].copy_(sample)
del sample
@@ -485,7 +514,7 @@ def _final_norm_and_conv_out(
del padded
else:
inplace_conv3d_temporal_chunked(workspace, conv_out)
result = workspace[:, : conv_out.out_channels, 1:-1].contiguous()
result = workspace[:, : conv_out.out_channels, 1:-1].contiguous(memory_format=mem_fmt)
del workspace, interior
return unpatchify(result, patch_size_hw=decoder.patch_size, patch_size_t=1)
@@ -507,6 +536,9 @@ def _memory_efficient_forward(
``UNetMidBlock3D`` and ``DepthToSpaceUpsample`` blocks use efficient
paths; standalone ``ResnetBlock3D`` blocks fall back to the standard
forward. The final norm + ada + SiLU + conv_out is also workspace-based.
All workspaces are allocated ``channels_last_3d`` so cuDNN's NHWC 3D
conv kernels run end-to-end. The caller (:func:`enable_memory_efficient_decode`)
is responsible for converting the input sample and decoder weights to NHWC.
"""
causal = decoder.causal
batch_size = sample.shape[0]
@@ -537,6 +569,11 @@ def _memory_efficient_forward(
raise ValueError("'timestep' required when timestep_conditioning=True")
scaled_timestep = timestep * decoder.timestep_scale_multiplier.to(sample)
# Workspaces are unconditionally NHWC: rearrange + slice + residual-add
# inside _upsample_forward_efficient produces NCHW-default output, so
# per-tensor inspection would silently fall back to NCHW for every
# workspace after the first upsample.
# --- Up blocks (dispatch to efficient path per block type) ---
for up_block in decoder.up_blocks:
if isinstance(up_block, UNetMidBlock3D):
@@ -546,15 +583,16 @@ def _memory_efficient_forward(
causal=causal,
timestep=scaled_timestep if decoder.timestep_conditioning else None,
generator=generator,
prefer_channels_last_3d=True,
)
elif isinstance(up_block, DepthToSpaceUpsample):
sample = _upsample_forward_efficient(up_block, sample, causal=causal)
sample = _upsample_forward_efficient(up_block, sample, causal=causal, prefer_channels_last_3d=True)
elif isinstance(up_block, ResnetBlock3D):
sample = up_block(sample, causal=causal, generator=generator)
else:
sample = up_block(sample, causal=causal)
return _final_norm_and_conv_out(decoder, sample, causal, scaled_timestep, batch_size)
return _final_norm_and_conv_out(decoder, sample, causal, scaled_timestep, batch_size, prefer_channels_last_3d=True)
# ---------------------------------------------------------------------------
@@ -564,6 +602,10 @@ def _memory_efficient_forward(
def enable_memory_efficient_decode(decoder: nn.Module) -> nn.Module:
"""Patch a ``VideoDecoder`` to use the memory-efficient forward path.
The mem-efficient path runs the decoder in ``channels_last_3d`` memory
format: weights and inputs are converted on first call so cuDNN's NHWC
3D conv kernels are used (~2x faster, avoids the large vol2col scratch
buffer of the NCHW path).
The original ``forward`` is saved as ``decoder._original_forward`` so
that it can be restored later with :func:`disable_memory_efficient_decode`.
"""
@@ -577,12 +619,20 @@ def enable_memory_efficient_decode(decoder: nn.Module) -> nn.Module:
return decoder
original_forward = decoder.forward
weights_converted = False
def efficient_forward(
sample: torch.Tensor,
timestep: torch.Tensor | None = None,
generator: torch.Generator | None = None,
) -> torch.Tensor:
nonlocal weights_converted
if sample.dim() == 5:
if not weights_converted:
# Lazy: weights are real by first-call time (meta -> loader -> here).
decoder.to(memory_format=torch.channels_last_3d)
weights_converted = True
sample = sample.to(memory_format=torch.channels_last_3d)
return _memory_efficient_forward(decoder, sample, timestep, generator)
decoder._original_forward = original_forward # type: ignore[attr-defined]
packages/ltx-core/src/ltx_core/quantization/__init__.py
+4
View File
@@ -2,7 +2,9 @@ from ltx_core.quantization.fp8_cast import (
TRANSFORMER_LINEAR_DOWNCAST_MAP,
UPCAST_DURING_INFERENCE,
UpcastWithStochasticRounding,
fp8_cast_fuse_rule,
)
from ltx_core.quantization.fp8_scaled_mm import fp8_scaled_mm_fuse_rule
from ltx_core.quantization.policy import QuantizationPolicy
__all__ = [
@@ -10,4 +12,6 @@ __all__ = [
"UPCAST_DURING_INFERENCE",
"QuantizationPolicy",
"UpcastWithStochasticRounding",
"fp8_cast_fuse_rule",
"fp8_scaled_mm_fuse_rule",
]
packages/ltx-core/src/ltx_core/quantization/fp8_cast.py
+192 -43
View File
@@ -1,9 +1,15 @@
from pathlib import Path
import safetensors
import torch
from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule
from ltx_core.loader.kernels import TRITON_AVAILABLE
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.primitives import StateDict
from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
from ltx_core.model.transformer.model import LTXModel
from ltx_core.quantization.policy import QuantizationPolicy
BLOCK_SIZE = 1024
@@ -113,58 +119,65 @@ def _replace_fwd_with_upcast(layer: torch.nn.Linear, with_stochastic_rounding: b
)
# Module-name suffixes for the Linears that participate in fp8 cast. Used by
# both the upcast matcher and the sd_ops downcast map so the two cannot drift.
# - ``.to_q`` / ``.to_k`` / ``.to_v`` / ``.to_out.0`` have a leading dot so they
# only match the attention Linears at ``...attnN.to_q`` etc.
# - ``ff.net.0.proj`` / ``ff.net.2`` are intentionally **dotless** so they match
# both video FF (``...ff.net.0.proj``) and audio FF (``...audio_ff.net.0.proj``).
_FP8_CAST_KEY_PREFIX = "transformer_blocks."
_FP8_CAST_LINEAR_SUFFIXES: tuple[str, ...] = (
".to_q",
".to_k",
".to_v",
".to_out.0",
"ff.net.0.proj",
"ff.net.2",
)
def _is_fp8_cast_linear(module_name: str) -> bool:
"""Return True if *module_name* names a Linear that should be fp8-cast."""
if _FP8_CAST_KEY_PREFIX not in module_name:
return False
return any(module_name.endswith(suffix) for suffix in _FP8_CAST_LINEAR_SUFFIXES)
def _amend_forward_with_upcast(
model: torch.nn.Module, with_stochastic_rounding: bool = False, seed: int = 0
) -> torch.nn.Module:
"""
Replace the forward method of the model's Linear layers to forward
with upcast and optional stochastic rounding.
Replace the forward method of the fp8-cast Linear layers (per
:data:`_FP8_CAST_LINEAR_SUFFIXES`) to forward with upcast and optional
stochastic rounding.
Only the Linears whose weights are downcast by :data:`TRANSFORMER_LINEAR_DOWNCAST_MAP`
are retyped. Linears outside that subset (e.g. ``to_gate_logits``) are left as
plain ``nn.Linear`` so the meta-model param dtype matches the loaded checkpoint
dtype.
"""
for m in model.modules():
if isinstance(m, (torch.nn.Linear)):
for name, m in model.named_modules():
if isinstance(m, torch.nn.Linear) and _is_fp8_cast_linear(name):
_replace_fwd_with_upcast(m, with_stochastic_rounding, seed)
return model
TRANSFORMER_LINEAR_DOWNCAST_MAP = (
SDOps("TRANSFORMER_LINEAR_DOWNCAST_MAP")
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".to_q.weight", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".to_q.bias", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".to_k.weight", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".to_k.bias", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".to_v.weight", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".to_v.bias", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".to_out.0.weight", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".to_out.0.bias", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix="ff.net.0.proj.weight", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix="ff.net.0.proj.bias", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix="ff.net.2.weight", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix="ff.net.2.bias", operation=_naive_weight_or_bias_downcast
)
)
def _build_transformer_linear_downcast_map() -> SDOps:
"""Build the sd_ops downcast map from the same suffix registry as the matcher."""
ops = SDOps("TRANSFORMER_LINEAR_DOWNCAST_MAP")
for suffix in _FP8_CAST_LINEAR_SUFFIXES:
ops = ops.with_kv_operation(
key_prefix=_FP8_CAST_KEY_PREFIX,
key_suffix=suffix + ".weight",
operation=_naive_weight_or_bias_downcast,
).with_kv_operation(
key_prefix=_FP8_CAST_KEY_PREFIX,
key_suffix=suffix + ".bias",
operation=_naive_weight_or_bias_downcast,
)
return ops
TRANSFORMER_LINEAR_DOWNCAST_MAP = _build_transformer_linear_downcast_map()
UPCAST_DURING_INFERENCE = ModuleOps(
name="upcast_fp8_during_linear_forward",
@@ -186,3 +199,139 @@ class UpcastWithStochasticRounding(ModuleOps):
matcher=lambda model: isinstance(model, LTXModel),
mutator=lambda model: _amend_forward_with_upcast(model, True, seed),
)
def fuse_cast_fp8_weight(
delta_bf16: torch.Tensor,
weight_fp8: torch.Tensor,
) -> torch.Tensor:
"""Return ``(delta_bf16 + dequantize(weight_fp8)).to(weight_fp8.dtype)``.
CUDA with Triton uses stochastic rounding via the fused kernel; otherwise
falls back to a deterministic bf16 add. ``delta_bf16`` is the bf16
accumulator and is mutated in place.
"""
if delta_bf16.dtype != torch.bfloat16:
raise ValueError(f"delta_bf16 must be bfloat16, got {delta_bf16.dtype}")
if str(weight_fp8.device).startswith("cuda") and TRITON_AVAILABLE:
fused_add_round_launch(delta_bf16, weight_fp8, seed=0)
else:
delta_bf16.add_(weight_fp8.to(dtype=torch.bfloat16))
return delta_bf16.to(dtype=weight_fp8.dtype)
def _fp8_cast_fuse(
key: str,
weight: torch.Tensor,
deltas: torch.Tensor,
model_sd: StateDict,
) -> dict[str, torch.Tensor]:
"""Cast the dequantized FP8 weight + BF16 deltas back to ``weight.dtype``
(FP8) via the fused-add-round kernel on CUDA.
Only a subset of linears are FP8-downcast (see ``TRANSFORMER_LINEAR_DOWNCAST_MAP``);
LoRAs may also target layers left in BF16 (e.g. audio ``add_q/k/v_proj``, cross-modal
projections). For those, fall back to a plain BF16 fuse.
"""
if weight.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2):
return bf16_fuse_rule(key, weight, deltas, model_sd)
return {key: fuse_cast_fp8_weight(deltas, weight)}
fp8_cast_fuse_rule = FuseRule(aggregation_dtype=torch.bfloat16, fuse_fn=_fp8_cast_fuse)
# Raw safetensors storage prefix shared by every diffusion-transformer
# parameter (and every prequant `*_scale` sibling). Verified against
# ltx-2.3-22b-{dev,distilled}-fp8.safetensors: 2924/2924 and 2992/2992 of
# the scale keys start with this exact prefix.
_RAW_DIFFUSION_MODEL_PREFIX = "model.diffusion_model."
def _read_scales(checkpoint_path: str | Path) -> dict[str, torch.Tensor]:
"""Return ``{post_rename_param_key: scale_tensor}`` for every prequant
``*_scale`` sibling in *checkpoint_path*.
Keys are returned in the post-rename form the loader will pass to the
sd-op (e.g. ``transformer_blocks.0.attn1.to_q.weight``) -- the raw
``model.diffusion_model.`` prefix and the ``_scale`` suffix are both
stripped. Catches both ``.weight_scale`` and ``.bias_scale``; the
latter is absent in the current LTX-2.3 prequant checkpoints but
accepted for forward compatibility.
"""
out: dict[str, torch.Tensor] = {}
with safetensors.safe_open(str(checkpoint_path), framework="pt", device="cpu") as h:
raw_keys = h.keys()
for k in raw_keys:
if not k.endswith("_scale"):
continue
if not k.startswith(_RAW_DIFFUSION_MODEL_PREFIX):
raise ValueError(
f"Scale key {k!r} does not start with the expected raw prefix {_RAW_DIFFUSION_MODEL_PREFIX!r}"
)
param_key = k.removeprefix(_RAW_DIFFUSION_MODEL_PREFIX).removesuffix("_scale")
out[param_key] = h.get_tensor(k)
return out
def _build_prequant_fold_sd_ops(scales: dict[str, torch.Tensor]) -> SDOps:
"""Build sd-ops that fold prequant ``*_scale`` siblings into their parent
tensor at load time.
*scales* is keyed by the **post-rename** param key (e.g.
``transformer_blocks.0.attn1.to_q.weight``); see :func:`_read_scales`.
Four ``with_kv_operation`` entries (symmetric for ``.weight`` and ``.bias``):
* ``.weight`` / ``.bias`` -> if a sibling scale exists in *scales*, fold;
then delegate to ``TRANSFORMER_LINEAR_DOWNCAST_MAP`` (downcast covered
Linears, pass everything else through). Without a scale, delegate
directly.
* ``.weight_scale`` / ``.bias_scale`` -> drop (the scale is consumed by
the fold). Raises if the scale key doesn't correspond to a known
entry in *scales* -- that means the file shipped a scale we didn't
pre-register, which would silently desync the fold.
"""
def _on_param(param_key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
scale = scales.get(param_key)
if scale is None:
return TRANSFORMER_LINEAR_DOWNCAST_MAP.apply_to_key_value(param_key, value)
scale = scale.to(device=value.device)
if scale.ndim != 0:
raise ValueError(f"Unsupported scale shape {tuple(scale.shape)} for {param_key}")
bf16 = (value.to(torch.float32) * scale).to(torch.bfloat16)
# Delegate the final fp8-vs-bf16 decision to the downcast map: Linears
# outside the fp8 subset (e.g. to_gate_logits) stay bf16 to match the
# plain nn.Linear that the upcast matcher leaves untouched.
return TRANSFORMER_LINEAR_DOWNCAST_MAP.apply_to_key_value(param_key, bf16)
def _drop_scale(scale_key: str, _value: torch.Tensor) -> list[KeyValueOperationResult]:
param_key = scale_key.removesuffix("_scale")
if param_key not in scales:
raise ValueError(
f"Scale key {scale_key!r} has no matching entry in the prequant scales dict; "
f"_read_scales and the loader's rename map have drifted"
)
return []
# Register the drop ops first so the dict-membership sanity check is the
# earliest sd-op that can fire on a scale key -- we crash on a stray scale
# before any silently mismatched fold has a chance to land in the state
# dict. Registration order is irrelevant for correctness (no overlap
# between matchers) but communicates intent.
return (
SDOps("FP8_CAST_PREQUANT_AWARE")
.with_kv_operation(key_suffix=".weight_scale", operation=_drop_scale)
.with_kv_operation(key_suffix=".bias_scale", operation=_drop_scale)
.with_kv_operation(key_suffix=".weight", operation=_on_param)
.with_kv_operation(key_suffix=".bias", operation=_on_param)
)
def build_policy(checkpoint_path: str | Path) -> QuantizationPolicy:
"""FP8 casting with upcasting during inference.
*checkpoint_path* is required (mirroring ``fp8_scaled_mm.build_policy``).
For prequantized fp8 checkpoints, sibling ``*_scale`` tensors (weight or
bias) are folded into the parent at load time.
"""
scales = _read_scales(checkpoint_path)
return QuantizationPolicy(
sd_ops=_build_prequant_fold_sd_ops(scales),
module_ops=(UPCAST_DURING_INFERENCE,),
fuse_rule=fp8_cast_fuse_rule,
)
packages/ltx-core/src/ltx_core/quantization/fp8_scaled_mm.py
+42
View File
@@ -5,8 +5,11 @@ from typing import Callable
import torch
from torch import nn
from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.primitives import StateDict
from ltx_core.model.transformer import LTXModel
from ltx_core.quantization.policy import QuantizationPolicy
from ltx_core.quantization.trtllm_scaled_usable import trtllm_scaled_mm_usable
@@ -173,3 +176,42 @@ def get_fp8_swap_module_ops(checkpoint_path: str) -> tuple[ModuleOps, ...]:
mutator=lambda model: _swap_linears_to_fp8(model, _should_swap),
),
)
def _fp8_scaled_mm_fuse(
key: str,
weight: torch.Tensor,
deltas: torch.Tensor,
model_sd: StateDict,
) -> dict[str, torch.Tensor]:
"""Dequantize via ``weight.float() * weight_scale``, add the BF16 delta,
and re-quantize to FP8 with a fresh per-tensor scale.
Layers that were not swapped to scaled FP8 (e.g. small embedder linears
excluded from the auto-discovered swap set) stay BF16 and have no
``.weight_scale`` companion -- for those, fall back to a plain bf16 fuse.
"""
scale_key = key.replace(".weight", ".weight_scale")
if scale_key not in model_sd.sd:
return bf16_fuse_rule(key, weight, deltas, model_sd)
weight_scale = model_sd.sd[scale_key]
original_weight = weight.to(torch.float32) * weight_scale
new_weight = original_weight + deltas.to(torch.float32)
new_fp8_weight, new_weight_scale = quantize_weight_to_fp8_per_tensor(new_weight)
return {key: new_fp8_weight, scale_key: new_weight_scale}
fp8_scaled_mm_fuse_rule = FuseRule(aggregation_dtype=torch.bfloat16, fuse_fn=_fp8_scaled_mm_fuse)
def build_policy(checkpoint_path: str) -> QuantizationPolicy:
"""FP8 scaled matmul for checkpoints pre-quantized with per-tensor scales.
The set of layers to swap to ``FP8Linear`` is discovered from the
checkpoint's ``.weight_scale`` tensors via suffix-matching against the
model's named modules. Requires a pre-quantized checkpoint; for BF16
checkpoints, use :func:`ltx_core.quantization.fp8_cast.build_policy`.
"""
return QuantizationPolicy(
sd_ops=None,
module_ops=get_fp8_swap_module_ops(checkpoint_path),
fuse_rule=fp8_scaled_mm_fuse_rule,
)
packages/ltx-core/src/ltx_core/quantization/policy.py
+8 -32
View File
@@ -1,48 +1,24 @@
from dataclasses import dataclass
from enum import Enum
from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.sd_ops import SDOps
from ltx_core.quantization.fp8_cast import TRANSFORMER_LINEAR_DOWNCAST_MAP, UPCAST_DURING_INFERENCE
from ltx_core.quantization.fp8_scaled_mm import get_fp8_swap_module_ops
from ltx_core.model.model_protocol import ModelConfigurator
from ltx_core.model.transformer.model import LTXModel
@dataclass(frozen=True)
class QuantizationPolicy:
"""Configuration for model quantization during loading.
Attributes:
kind: Discriminator for the policy variant.
sd_ops: State-dict operations applied to each tensor during load.
module_ops: Post-load module transformations applied to the meta model.
model_configurator: Configurator class to use when constructing the transformer.
fuse_rule: How LoRA deltas merge into this policy's weight layout.
Default ``bf16_fuse_rule`` is used when no policy is configured.
"""
class Kind(str, Enum):
FP8_CAST = "fp8_cast"
FP8_SCALED_MM = "fp8_scaled_mm"
kind: Kind
sd_ops: SDOps | None = None
module_ops: tuple[ModuleOps, ...] = ()
@classmethod
def fp8_cast(cls) -> "QuantizationPolicy":
"""FP8 casting with upcasting during inference."""
return cls(
kind=cls.Kind.FP8_CAST,
sd_ops=TRANSFORMER_LINEAR_DOWNCAST_MAP,
module_ops=(UPCAST_DURING_INFERENCE,),
)
@classmethod
def fp8_scaled_mm(cls, checkpoint_path: str) -> "QuantizationPolicy":
"""FP8 scaled matmul for checkpoints pre-quantized with per-tensor scales.
The set of layers to swap to ``FP8Linear`` is discovered from the
checkpoint's ``.weight_scale`` tensors via suffix-matching against the
model's named modules. Requires a pre-quantized checkpoint; for BF16
checkpoints, use :meth:`fp8_cast` instead.
"""
return cls(
kind=cls.Kind.FP8_SCALED_MM,
sd_ops=None,
module_ops=get_fp8_swap_module_ops(checkpoint_path),
)
model_configurator: type[ModelConfigurator[LTXModel]] | None = None
fuse_rule: FuseRule = bf16_fuse_rule
packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/base_encoder.py
+1 -1
View File
@@ -183,7 +183,7 @@ def module_ops_from_gemma_root(gemma_root: str) -> tuple[ModuleOps, ...]:
return module
def load_processor(module: GemmaTextEncoder) -> GemmaTextEncoder:
image_processor = AutoImageProcessor.from_pretrained(processor_root, local_files_only=True)
image_processor = AutoImageProcessor.from_pretrained(processor_root, local_files_only=True, use_fast=False)
if not module.tokenizer:
raise ValueError("Tokenizer model operation must be performed before processor model operation")
module.processor = Gemma3Processor(image_processor=image_processor, tokenizer=module.tokenizer.tokenizer)
packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/encoder_configurator.py
+110 -27
View File
@@ -1,4 +1,5 @@
import torch
import transformers
from transformers import Gemma3Config
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
from transformers.models.gemma3 import Gemma3ForConditionalGeneration
@@ -19,6 +20,8 @@ from ltx_core.text_encoders.gemma.feature_extractor import (
FeatureExtractorV2,
)
_TRANSFORMERS_V5: bool = int(transformers.__version__.split(".", 1)[0]) >= 5
class GemmaTextEncoderConfigurator(ModelConfigurator[GemmaTextEncoder]):
@classmethod
@@ -100,26 +103,45 @@ def _create_feature_extractor(transformer_config: dict) -> torch.nn.Module:
# --- Split SDOps: Gemma LLM keys vs Embeddings Processor keys ---
GEMMA_LLM_KEY_OPS = (
SDOps("GEMMA_LLM_KEY_OPS")
# 1. Map language model layers (note the double .model prefix)
.with_matching(prefix="language_model.model.")
.with_replacement("language_model.model.", "model.model.language_model.")
# 2. Map the Vision Tower
.with_matching(prefix="vision_tower.")
.with_replacement("vision_tower.", "model.model.vision_tower.")
# 3. Map the Multi-Modal Projector
.with_matching(prefix="multi_modal_projector.")
.with_replacement("multi_modal_projector.", "model.model.multi_modal_projector.")
# 4. Duplicate embed_tokens to lm_head (needed for prompt enhancement via generate())
.with_kv_operation(
operation=lambda key, value: [
KeyValueOperationResult(key, value),
KeyValueOperationResult("model.lm_head.weight", value),
],
key_prefix="model.model.language_model.embed_tokens.weight",
def _build_gemma_llm_key_ops(*, transformers_v5: bool) -> SDOps:
"""Build the checkpoint-key remapping for the Gemma multimodal encoder.
The vision-tower mapping differs between transformers <5 and >=5 because
upstream PR https://github.com/huggingface/transformers/pull/39847 flattened
``Gemma3ForConditionalGeneration.model.vision_tower.vision_model`` into
``model.vision_tower``. Checkpoints continue to ship the legacy
``vision_tower.vision_model.*`` prefix, so we strip the inner ``vision_model.``
when targeting v5 and pass it through unchanged for v4.
"""
base = (
SDOps("GEMMA_LLM_KEY_OPS")
# 1. Map language model layers (note the double .model prefix)
.with_matching(prefix="language_model.model.")
.with_replacement("language_model.model.", "model.model.language_model.")
# 2. Map the Vision Tower (version-dependent — see docstring)
.with_matching(prefix="vision_tower.")
)
)
if transformers_v5:
base = base.with_replacement("vision_tower.vision_model.", "model.model.vision_tower.")
else:
base = base.with_replacement("vision_tower.", "model.model.vision_tower.")
return (
base
# 3. Map the Multi-Modal Projector
.with_matching(prefix="multi_modal_projector.")
.with_replacement("multi_modal_projector.", "model.model.multi_modal_projector.")
# 4. Duplicate embed_tokens to lm_head (needed for prompt enhancement via generate())
.with_kv_operation(
operation=lambda key, value: [
KeyValueOperationResult(key, value),
KeyValueOperationResult("model.lm_head.weight", value),
],
key_prefix="model.model.language_model.embed_tokens.weight",
)
)
GEMMA_LLM_KEY_OPS = _build_gemma_llm_key_ops(transformers_v5=_TRANSFORMERS_V5)
EMBEDDINGS_PROCESSOR_KEY_OPS = (
SDOps("EMBEDDINGS_PROCESSOR_KEY_OPS")
@@ -152,24 +174,85 @@ VIDEO_ONLY_EMBEDDINGS_PROCESSOR_KEY_OPS = (
)
def _resolve_local_base_freq(config: object) -> float:
rope_parameters = getattr(config, "rope_parameters", None)
if isinstance(rope_parameters, dict) and "sliding_attention" in rope_parameters:
sliding = rope_parameters["sliding_attention"]
if isinstance(sliding, dict) and "rope_theta" in sliding:
return float(sliding["rope_theta"])
if hasattr(config, "rope_local_base_freq"):
return float(config.rope_local_base_freq)
raise AttributeError(
"Gemma text_config exposes neither rope_local_base_freq nor rope_parameters['sliding_attention']['rope_theta']"
)
def _resolve_full_rope_type(config: object) -> str:
rope_parameters = getattr(config, "rope_parameters", None)
if isinstance(rope_parameters, dict) and "full_attention" in rope_parameters:
full = rope_parameters["full_attention"]
if isinstance(full, dict) and "rope_type" in full:
return str(full["rope_type"])
rope_scaling = getattr(config, "rope_scaling", None)
if rope_scaling is not None:
if isinstance(rope_scaling, dict):
if "rope_type" in rope_scaling:
return str(rope_scaling["rope_type"])
elif hasattr(rope_scaling, "rope_type"):
return str(rope_scaling.rope_type)
raise AttributeError(
"Gemma text_config exposes neither rope_scaling.rope_type nor rope_parameters['full_attention']['rope_type']"
)
def _populate_rotary_v4(l_model: torch.nn.Module, config: object) -> None:
"""transformers <5 layout: separate ``rotary_emb_local`` + ``rotary_emb``."""
dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
base = _resolve_local_base_freq(config)
local_inv = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(dtype=torch.float) / dim))
full_inv, _ = ROPE_INIT_FUNCTIONS[_resolve_full_rope_type(config)](config)
l_model.rotary_emb_local.register_buffer("inv_freq", local_inv)
l_model.rotary_emb.register_buffer("inv_freq", full_inv)
def _populate_rotary_v5(l_model: torch.nn.Module, config: object) -> None:
"""transformers >=5 layout: single ``rotary_emb`` with per-layer-type buffers.
Mirrors ``Gemma3PreTrainedModel._init_weights`` for ``Gemma3RotaryEmbedding``
so meta-built models reach the same numerical state as a from_pretrained load.
"""
rope_emb = l_model.rotary_emb
for layer_type in dict.fromkeys(config.layer_types):
rope_params = config.rope_parameters[layer_type]
if rope_params is None:
continue
rope_type = rope_params["rope_type"]
if rope_type == "default":
inv_freq, attn_scaling = rope_emb.compute_default_rope_parameters(config, layer_type=layer_type)
else:
inv_freq, attn_scaling = ROPE_INIT_FUNCTIONS[rope_type](config, layer_type=layer_type)
rope_emb.register_buffer(f"{layer_type}_inv_freq", inv_freq, persistent=False)
rope_emb.register_buffer(f"{layer_type}_original_inv_freq", inv_freq.clone(), persistent=False)
setattr(rope_emb, f"{layer_type}_attention_scaling", attn_scaling)
def create_and_populate(module: GemmaTextEncoder) -> GemmaTextEncoder:
model = module.model
v_model = model.model.vision_tower.vision_model
v_tower = model.model.vision_tower
v_model = v_tower.vision_model if hasattr(v_tower, "vision_model") else v_tower
l_model = model.model.language_model
config = model.config.text_config
dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
base = config.rope_local_base_freq
local_rope_freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(dtype=torch.float) / dim))
inv_freqs, _ = ROPE_INIT_FUNCTIONS[config.rope_scaling["rope_type"]](config)
if hasattr(l_model, "rotary_emb_local"):
_populate_rotary_v4(l_model, config)
else:
_populate_rotary_v5(l_model, config)
positions_length = len(v_model.embeddings.position_ids[0])
position_ids = torch.arange(positions_length, dtype=torch.long, device="cpu").unsqueeze(0)
v_model.embeddings.register_buffer("position_ids", position_ids)
embed_scale = torch.tensor(model.config.text_config.hidden_size**0.5, device="cpu")
embed_scale = torch.tensor(config.hidden_size**0.5, device="cpu")
l_model.embed_tokens.register_buffer("embed_scale", embed_scale)
l_model.rotary_emb_local.register_buffer("inv_freq", local_rope_freqs)
l_model.rotary_emb.register_buffer("inv_freq", inv_freqs)
return module
packages/ltx-pipelines/README.md
+31 -13
View File
@@ -369,7 +369,9 @@ PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python -m ltx_pipelines.ti2vid_
When authoring custom scripts, pass a `QuantizationPolicy` to pipeline classes:
```python
from ltx_core.quantization import QuantizationPolicy
from ltx_core.quantization.fp8_cast import build_policy as build_fp8_cast_policy
# Alternative:
# from ltx_core.quantization.fp8_scaled_mm import build_policy as build_fp8_scaled_mm_policy
pipeline = TI2VidTwoStagesPipeline(
checkpoint_path=ltx_model_path,
@@ -377,7 +379,7 @@ pipeline = TI2VidTwoStagesPipeline(
spatial_upsampler_path=upsampler_path,
gemma_root=gemma_root_path,
loras=[],
quantization=QuantizationPolicy.fp8_cast(), # or QuantizationPolicy.fp8_scaled_mm()
quantization=build_fp8_cast_policy(ltx_model_path),
)
pipeline(...)
```
@@ -414,12 +416,13 @@ def denoising_loop(sigmas, video_state, audio_state, stepper):
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=your_denoise_function,
transformer=transformer,
denoiser=denoiser,
ge_gamma=2.0, # Gradient estimation coefficient
)
```
This allows you to use **20-30 steps instead of 40** while maintaining quality. The gradient estimation function is available in [`pipeline_utils.py`](src/ltx_pipelines/utils/helpers.py).
This allows you to use **20-30 steps instead of 40** while maintaining quality. The gradient estimation function is defined in [`samplers.py`](src/ltx_pipelines/utils/samplers.py).
---
@@ -435,15 +438,18 @@ This allows you to use **20-30 steps instead of 40** while maintaining quality.
## 📖 Example: Image-to-Video
```python
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
from ltx_pipelines.utils.args import ImageConditioningInput
from ltx_pipelines.utils.media_io import encode_video
distilled_lora = [
LoraPathStrengthAndSDOps(
"/path/to/distilled_lora.safetensors",
0.6,
LTXV_LORA_COMFY_RENAMING_MAP
LTXV_LORA_COMFY_RENAMING_MAP,
),
]
@@ -473,19 +479,31 @@ audio_guider_params = MultiModalGuiderParams(
stg_blocks=[29],
)
# Generate video from image
pipeline(
# Generate video from image. The pipeline returns (video_iterator, audio);
# the caller is responsible for encoding to file via encode_video().
num_frames = 121
frame_rate = 25.0
tiling_config = TilingConfig.default()
video, audio = pipeline(
prompt="A serene landscape with mountains in the background",
output_path="output.mp4",
negative_prompt="worst quality, low quality, blurry, distorted",
seed=42,
height=512,
width=768,
num_frames=121,
frame_rate=25.0,
num_frames=num_frames,
frame_rate=frame_rate,
num_inference_steps=40,
video_guider_params=video_guider_params,
audio_guider_params=audio_guider_params,
images=[ImageConditioningInput("input_image.jpg", 0, 1.0, 33)], # Image at frame 0, strength 1.0, CRF 33
images=[ImageConditioningInput("input_image.jpg", 0, 1.0, 33)], # path, frame_idx=0, strength=1.0, crf=33
tiling_config=tiling_config,
)
encode_video(
video=video,
fps=frame_rate,
audio=audio,
output_path="output.mp4",
video_chunks_number=get_video_chunks_number(num_frames, tiling_config),
)
```
packages/ltx-pipelines/pyproject.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[project]
name = "ltx-pipelines"
version = "1.1.3"
version = "1.1.4"
description = "Pipelines implementation for Lightricks' LTX-2 model"
readme = "README.md"
requires-python = ">=3.10"
packages/ltx-pipelines/src/ltx_pipelines/a2vid_two_stage.py
+6 -5
View File
@@ -9,6 +9,7 @@ from ltx_core.components.schedulers import LTX2Scheduler
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import Registry
from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, AudioLatentShape, VideoPixelShape
@@ -52,7 +53,7 @@ class A2VidPipelineTwoStage:
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
@@ -71,7 +72,7 @@ class A2VidPipelineTwoStage:
loras=tuple(loras),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
@@ -82,7 +83,7 @@ class A2VidPipelineTwoStage:
loras=stage_2_loras,
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.upsampler = VideoUpsampler(
@@ -238,7 +239,7 @@ class A2VidPipelineTwoStage:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
parser = default_2_stage_arg_parser()
parser.add_argument(
"--audio-path",
@@ -266,7 +267,7 @@ def main() -> None:
gemma_root=args.gemma_root,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/distilled.py
+5 -4
View File
@@ -6,6 +6,7 @@ import torch
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import Registry
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio
@@ -53,7 +54,7 @@ class DistilledPipeline:
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
@@ -75,7 +76,7 @@ class DistilledPipeline:
loras=tuple(loras),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.upsampler = VideoUpsampler(
@@ -180,7 +181,7 @@ class DistilledPipeline:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path(distilled=True)
params = detect_params(checkpoint_path)
parser = default_2_stage_distilled_arg_parser(params=params)
@@ -191,7 +192,7 @@ def main() -> None:
gemma_root=args.gemma_root,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/hdr_ic_lora.py
+5 -2
View File
@@ -59,6 +59,7 @@ from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES, STAGE_2_DISTIL
from ltx_pipelines.utils.denoisers import SimpleDenoiser
from ltx_pipelines.utils.helpers import get_device, modality_from_latent_state
from ltx_pipelines.utils.media_io import ResizeMode, align_resolution, load_video_conditioning_hdr
from ltx_pipelines.utils.quantization_factory import QuantizationKind
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
logger = logging.getLogger(__name__)
@@ -77,7 +78,7 @@ ALIGNMENT_DIVISOR = 64
# to the pipeline constructor.
TILED_VAE_ENCODE_PIXEL_THRESHOLD = 512 * 768
_DEFAULT_QUANTIZATION = QuantizationPolicy.fp8_cast()
_DEFAULT_QUANTIZATION = QuantizationKind.FP8_CAST
# Default stage-2 configuration: one refinement phase with modest 2-way tiling
# in every dimension and a short 2-step distilled sigma schedule.
@@ -205,7 +206,7 @@ class HDRICLoraPipeline:
hdr_lora: str | Path,
text_embeddings_path: str | Path,
device: torch.device | None = None,
quantization: QuantizationPolicy = _DEFAULT_QUANTIZATION,
quantization: QuantizationPolicy | QuantizationKind | None = _DEFAULT_QUANTIZATION,
registry: Registry | None = None,
hdr_lora_config: HdrLoraConfig | None = None,
tiled_vae_encode_pixel_threshold: int = TILED_VAE_ENCODE_PIXEL_THRESHOLD,
@@ -232,6 +233,8 @@ class HDRICLoraPipeline:
"""
self.device = device or get_device()
self._tiled_vae_encode_threshold = tiled_vae_encode_pixel_threshold
if isinstance(quantization, QuantizationKind):
quantization = quantization.to_policy(checkpoint_path=distilled_checkpoint_path)
if offload_mode != OffloadMode.NONE and quantization is not None:
logger.info("Offload mode enabled — disabling quantization (not supported with layer streaming).")
quantization = None
packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py
+6 -5
View File
@@ -7,6 +7,7 @@ from ltx_core.components.noisers import GaussianNoiser
from ltx_core.conditioning import ConditioningItem
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import Registry
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, VideoPixelShape
@@ -60,7 +61,7 @@ class ICLoraPipeline:
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
@@ -82,7 +83,7 @@ class ICLoraPipeline:
loras=tuple(loras),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.stage_2 = DiffusionStage(
@@ -92,7 +93,7 @@ class ICLoraPipeline:
loras=(),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.upsampler = VideoUpsampler(
@@ -330,7 +331,7 @@ class ICLoraPipeline:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path(distilled=True)
params = detect_params(checkpoint_path)
parser = default_2_stage_distilled_arg_parser(params=params)
@@ -386,7 +387,7 @@ def main() -> None:
gemma_root=args.gemma_root,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py
+6 -5
View File
@@ -12,6 +12,7 @@ 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.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, VideoPixelShape
@@ -62,7 +63,7 @@ class KeyframeInterpolationPipeline:
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
@@ -80,7 +81,7 @@ class KeyframeInterpolationPipeline:
loras=tuple(loras),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
@@ -91,7 +92,7 @@ class KeyframeInterpolationPipeline:
loras=stage_2_loras,
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.upsampler = VideoUpsampler(
@@ -233,7 +234,7 @@ class KeyframeInterpolationPipeline:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path()
params = detect_params(checkpoint_path)
parser = default_2_stage_arg_parser(params=params)
@@ -245,7 +246,7 @@ def main() -> None:
gemma_root=args.gemma_root,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/lipdub.py
+5 -4
View File
@@ -13,6 +13,7 @@ from ltx_core.conditioning import AudioConditionByReferenceLatent
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import Registry
from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, AudioLatentShape, SpatioTemporalScaleFactors, VideoPixelShape
@@ -59,7 +60,7 @@ class LipDubPipeline:
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
) -> None:
self.device = device or get_device()
@@ -89,7 +90,7 @@ class LipDubPipeline:
loras=loras,
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.upsampler = VideoUpsampler(
@@ -289,7 +290,7 @@ def patchify_lipdub_audio_reference_latent(
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path(distilled=True)
params = detect_params(checkpoint_path)
parser = lipdub_arg_parser(params=params)
@@ -304,7 +305,7 @@ def main() -> None:
gemma_root=args.gemma_root,
ic_lora=args.lora[0],
quantization=args.quantization,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/retake.py
+5 -4
View File
@@ -11,6 +11,7 @@ from ltx_core.components.schedulers import LTX2Scheduler
from ltx_core.conditioning.types.noise_mask_cond import TemporalRegionMask
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import Registry
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import (
@@ -73,7 +74,7 @@ class RetakePipeline:
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
distilled: bool = True,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
@@ -108,7 +109,7 @@ class RetakePipeline:
loras=tuple(loras),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.video_decoder = VideoDecoder(
@@ -283,7 +284,7 @@ class RetakePipeline:
@torch.inference_mode()
def main() -> None:
"""CLI entry point for retake (regenerate a time region)."""
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
parser = video_editing_arg_parser(distilled=True)
parser.description = "Retake: regenerate a time region of a video with LTX-2."
args = parser.parse_args()
@@ -308,7 +309,7 @@ def main() -> None:
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
distilled=True,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
params = detect_params(args.distilled_checkpoint_path)
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py
+5 -4
View File
@@ -12,6 +12,7 @@ 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.compiling import CompilationConfig
from ltx_core.model.video_vae.tiling import TilingConfig
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio
@@ -55,7 +56,7 @@ class TI2VidOneStagePipeline:
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.dtype = torch.bfloat16
@@ -82,7 +83,7 @@ class TI2VidOneStagePipeline:
loras=tuple(loras),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.video_decoder = VideoDecoder(
@@ -185,7 +186,7 @@ class TI2VidOneStagePipeline:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path()
params = detect_params(checkpoint_path)
parser = default_1_stage_arg_parser(params=params)
@@ -195,7 +196,7 @@ def main() -> None:
gemma_root=args.gemma_root,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
video, audio = pipeline(
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py
+6 -5
View File
@@ -12,6 +12,7 @@ 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.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, VideoPixelShape
@@ -61,7 +62,7 @@ class TI2VidTwoStagesPipeline:
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
@@ -85,7 +86,7 @@ class TI2VidTwoStagesPipeline:
loras=tuple(loras),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.stage_2 = DiffusionStage(
@@ -95,7 +96,7 @@ class TI2VidTwoStagesPipeline:
loras=(*tuple(loras), *distilled_lora),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
@@ -223,7 +224,7 @@ class TI2VidTwoStagesPipeline:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path()
params = detect_params(checkpoint_path)
parser = default_2_stage_arg_parser(params=params)
@@ -235,7 +236,7 @@ def main() -> None:
gemma_root=args.gemma_root,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_hq.py
+6 -5
View File
@@ -9,6 +9,7 @@ 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.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, VideoLatentShape, VideoPixelShape
@@ -60,7 +61,7 @@ class TI2VidTwoStagesHQPipeline:
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
@@ -95,7 +96,7 @@ class TI2VidTwoStagesHQPipeline:
loras=(*loras, distilled_lora_stage_1),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.stage_2 = DiffusionStage(
@@ -105,7 +106,7 @@ class TI2VidTwoStagesHQPipeline:
loras=(*loras, distilled_lora_stage_2),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
@@ -242,7 +243,7 @@ class TI2VidTwoStagesHQPipeline:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
parser = hq_2_stage_arg_parser(params=LTX_2_3_HQ_PARAMS)
args = parser.parse_args()
pipeline = TI2VidTwoStagesHQPipeline(
@@ -254,7 +255,7 @@ def main() -> None:
gemma_root=args.gemma_root,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/utils/args.py
+146 -26
View File
@@ -1,9 +1,11 @@
import argparse
import json
from collections.abc import Sequence
from pathlib import Path
from typing import NamedTuple
from typing import Any, NamedTuple
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.quantization import QuantizationPolicy
from ltx_pipelines.utils.constants import (
DEFAULT_IMAGE_CRF,
@@ -13,6 +15,7 @@ from ltx_pipelines.utils.constants import (
LTX_2_3_PARAMS,
PipelineParams,
)
from ltx_pipelines.utils.quantization_factory import QuantizationKind
from ltx_pipelines.utils.types import OffloadMode
@@ -32,7 +35,7 @@ class VideoConditioningAction(argparse.Action):
option_string: str | None = None, # noqa: ARG002
) -> None:
path, strength_str = values
resolved_path = resolve_path(path)
resolved_path = resolve_existing_path(path)
strength = float(strength_str)
current = getattr(namespace, self.dest) or []
current.append((resolved_path, strength))
@@ -58,7 +61,7 @@ class VideoMaskConditioningAction(argparse.Action):
msg = f"{option_string} requires exactly 2 arguments (MASK_PATH STRENGTH), got {len(values)}"
raise argparse.ArgumentError(self, msg)
mask_path = resolve_path(values[0])
mask_path = resolve_existing_path(values[0])
strength = float(values[1])
setattr(namespace, self.dest, (mask_path, strength))
@@ -76,7 +79,7 @@ class ImageAction(argparse.Action):
raise argparse.ArgumentError(self, msg)
conditioning = ImageConditioningInput(
path=resolve_path(values[0]),
path=resolve_existing_path(values[0]),
frame_idx=int(values[1]),
strength=float(values[2]),
crf=int(values[3]) if len(values) > 3 else DEFAULT_IMAGE_CRF,
@@ -101,7 +104,7 @@ class LoraAction(argparse.Action):
path = values[0]
strength_str = values[1] if len(values) > 1 else str(DEFAULT_LORA_STRENGTH)
resolved_path = resolve_path(path)
resolved_path = resolve_existing_path(path)
strength = float(strength_str)
current = getattr(namespace, self.dest) or []
@@ -109,11 +112,118 @@ class LoraAction(argparse.Action):
setattr(namespace, self.dest, current)
class CompileAction(argparse.Action):
"""Parse ``--compile [KEY=VALUE ...]`` into a :class:`CompilationConfig`.
The flag is absent -> ``args.compile`` stays at its default (``None``).
The flag is passed alone -> ``CompilationConfig()`` (vanilla torch defaults).
The flag is passed with args -> ``CompilationConfig`` with the given fields overridden.
Errors (unknown key, malformed value, duplicate key, empty value) raise
:class:`argparse.ArgumentError` so argparse formats them as friendly CLI
messages rather than uncaught tracebacks.
"""
_ALLOWED_KEYS = frozenset({"mode", "backend", "fullgraph", "dynamic", "inductor_config", "dynamo_config"})
def __call__(
self,
parser: argparse.ArgumentParser, # noqa: ARG002
namespace: argparse.Namespace,
values: list[str],
option_string: str | None = None, # noqa: ARG002
) -> None:
overrides: dict[str, object] = {}
for item in values:
if "=" not in item:
raise argparse.ArgumentError(self, f"expects KEY=VALUE pairs, got: {item!r}")
key, _, raw = item.partition("=")
key = key.strip()
if key not in self._ALLOWED_KEYS:
raise argparse.ArgumentError(
self,
f"{key!r} is not a CompilationConfig field; valid keys: {sorted(self._ALLOWED_KEYS)}",
)
if key in overrides:
raise argparse.ArgumentError(self, f"{key} given more than once")
if key == "mode":
overrides[key] = self._parse_mode(raw)
elif key == "backend":
overrides[key] = self._parse_non_empty(key, raw)
elif key == "fullgraph":
overrides[key] = self._parse_bool(key, raw)
elif key == "dynamic":
overrides[key] = self._parse_dynamic(raw)
elif key in ("inductor_config", "dynamo_config"):
overrides[key] = self._parse_json_dict(key, raw)
setattr(namespace, self.dest, CompilationConfig(**overrides))
def _parse_mode(self, raw: str) -> str | None:
stripped = raw.strip()
if not stripped:
raise argparse.ArgumentError(self, "mode=... value cannot be empty (use mode=none to clear)")
if stripped.lower() == "none":
return None
return stripped
def _parse_non_empty(self, key: str, raw: str) -> str:
stripped = raw.strip()
if not stripped:
raise argparse.ArgumentError(self, f"{key}=... value cannot be empty")
return stripped
def _parse_bool(self, key: str, raw: str) -> bool:
normalized = raw.strip().lower()
if normalized in ("true", "1"):
return True
if normalized in ("false", "0"):
return False
raise argparse.ArgumentError(self, f"{key}=... must be true or false; got {raw!r}")
def _parse_dynamic(self, raw: str) -> bool | None:
normalized = raw.strip().lower()
if normalized in ("auto", "none"):
return None
if normalized in ("true", "1"):
return True
if normalized in ("false", "0"):
return False
raise argparse.ArgumentError(self, f"dynamic=... must be auto/true/false; got {raw!r}")
def _parse_json_dict(self, key: str, raw: str) -> dict[str, Any]:
# Inline JSON object starts with '{'; otherwise treat the value as a path to a JSON file.
stripped = raw.strip()
if not stripped:
raise argparse.ArgumentError(self, f"{key}=... value cannot be empty")
if stripped.startswith("{"):
source = stripped
else:
path = Path(stripped).expanduser()
if not path.is_file():
raise argparse.ArgumentError(
self, f"{key}=... must be a JSON object or a path to a JSON file; got {raw!r}"
)
source = path.read_text()
try:
value = json.loads(source)
except json.JSONDecodeError as e:
raise argparse.ArgumentError(self, f"{key}=... must be a JSON object; got {raw!r} ({e.msg})") from None
if not isinstance(value, dict):
raise argparse.ArgumentError(self, f"{key}=... must decode to a JSON object; got {type(value).__name__}")
return value
def resolve_path(path: str) -> str:
return str(Path(path).expanduser().resolve().as_posix())
QUANTIZATION_POLICIES = ("fp8-cast", "fp8-scaled-mm")
def resolve_existing_path(path: str) -> str:
"""Resolve *path* and verify it exists."""
resolved = resolve_path(path)
if not Path(resolved).exists():
raise argparse.ArgumentError(None, f"Path not found: {resolved}")
return resolved
QUANTIZATION_POLICIES = tuple(k.value for k in QuantizationKind)
def _resolve_quantization(namespace: argparse.Namespace) -> None:
@@ -123,16 +233,14 @@ def _resolve_quantization(namespace: argparse.Namespace) -> None:
name = getattr(namespace, "quantization", None)
if name is None or isinstance(name, QuantizationPolicy):
return
if name == "fp8-cast":
namespace.quantization = QuantizationPolicy.fp8_cast()
try:
kind = QuantizationKind(name)
except ValueError:
return
if name == "fp8-scaled-mm":
ckpt = getattr(namespace, "checkpoint_path", None) or getattr(namespace, "distilled_checkpoint_path", None)
if ckpt is None:
raise SystemExit(
"--quantization fp8-scaled-mm requires --checkpoint-path (or --distilled-checkpoint-path)."
)
namespace.quantization = QuantizationPolicy.fp8_scaled_mm(ckpt)
ckpt = getattr(namespace, "checkpoint_path", None) or getattr(namespace, "distilled_checkpoint_path", None)
if ckpt is None:
raise SystemExit(f"--quantization {kind.value} requires --checkpoint-path (or --distilled-checkpoint-path).")
namespace.quantization = kind.to_policy(checkpoint_path=ckpt)
class _PipelineArgumentParser(argparse.ArgumentParser):
@@ -150,7 +258,7 @@ def detect_checkpoint_path(distilled: bool = False) -> str:
"""Pre-parse argv to extract the checkpoint path before building the full parser."""
pre = argparse.ArgumentParser(add_help=False)
flag = "--distilled-checkpoint-path" if distilled else "--checkpoint-path"
pre.add_argument(flag, type=resolve_path, required=True)
pre.add_argument(flag, type=resolve_existing_path, required=True)
known, _ = pre.parse_known_args()
return known.distilled_checkpoint_path if distilled else known.checkpoint_path
@@ -163,14 +271,14 @@ def basic_arg_parser(
if distilled:
parser.add_argument(
"--distilled-checkpoint-path",
type=resolve_path,
type=resolve_existing_path,
required=True,
help="Path to LTX-2 distilled model checkpoint (.safetensors file).",
)
else:
parser.add_argument(
"--checkpoint-path",
type=resolve_path,
type=resolve_existing_path,
required=True,
help="Path to LTX-2 model checkpoint (.safetensors file).",
)
@@ -185,7 +293,7 @@ def basic_arg_parser(
)
parser.add_argument(
"--gemma-root",
type=resolve_path,
type=resolve_existing_path,
required=True,
help="Path to the root directory containing the Gemma text encoder model files.",
)
@@ -276,8 +384,20 @@ def basic_arg_parser(
)
parser.add_argument(
"--compile",
action="store_true",
help="Enable torch.compile for transformer blocks to optimize performance.",
nargs="*",
action=CompileAction,
default=None,
metavar="KEY=VALUE",
help=(
"Enable torch.compile for transformer blocks. Pass alone for defaults, "
"or with KEY=VALUE overrides for any CompilationConfig field. "
"Keys: mode, backend, fullgraph, dynamic, inductor_config, dynamo_config. "
"inductor_config/dynamo_config take JSON objects (inline or a path to a .json file) "
"that fully replace the defaults. "
"Examples: --compile or --compile mode=reduce-overhead or "
"--compile mode=reduce-overhead fullgraph=true backend=eager or "
"--compile inductor_config='{\"max_autotune\": true}'"
),
)
return parser
@@ -340,7 +460,7 @@ def video_editing_arg_parser(
(no height/width/num-frames; resolution comes from input video). Default is distilled checkpoint only.
"""
parser = basic_arg_parser(distilled=distilled)
parser.add_argument("--video-path", type=resolve_path, required=True, help="Path to the source video.")
parser.add_argument("--video-path", type=resolve_existing_path, required=True, help="Path to the source video.")
parser.add_argument("--start-time", type=float, required=True, help="Start time of the region to regenerate (s).")
parser.add_argument("--end-time", type=float, required=True, help="End time of the region to regenerate (s).")
return parser
@@ -462,7 +582,7 @@ def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
default=video_guider.skip_step,
help=(
"Video skip step N controls periodic skipping during the video diffusion process: "
"only steps where step_index % (N + 1) == 0 are processed, all others are skipped "
"only steps where step_index %% (N + 1) == 0 are processed, all others are skipped "
f"(e.g., 0 = no skipping; 1 = skip every other step; 2 = skip 2 of every 3 steps; "
f"default: {video_guider.skip_step})."
),
@@ -522,7 +642,7 @@ def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
default=audio_guider.skip_step,
help=(
"Audio skip step N controls periodic skipping during the audio diffusion process: "
"only steps where step_index % (N + 1) == 0 are processed, all others are skipped "
"only steps where step_index %% (N + 1) == 0 are processed, all others are skipped "
f"(e.g., 0 = no skipping; 1 = skip every other step; 2 = skip 2 of every 3 steps; "
f"default: {audio_guider.skip_step})."
),
@@ -562,7 +682,7 @@ def default_2_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
)
parser.add_argument(
"--spatial-upsampler-path",
type=resolve_path,
type=resolve_existing_path,
required=True,
help=(
"Path to the spatial upsampler model used to increase the resolution "
@@ -605,7 +725,7 @@ def default_2_stage_distilled_arg_parser(params: PipelineParams = LTX_2_3_PARAMS
)
parser.add_argument(
"--spatial-upsampler-path",
type=resolve_path,
type=resolve_existing_path,
required=True,
help=(
"Path to the spatial upsampler model used to increase the resolution "
packages/ltx-pipelines/src/ltx_pipelines/utils/blocks.py
+122 -45
View File
@@ -6,6 +6,8 @@ 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
@@ -21,6 +23,8 @@ from ltx_core.components.noisers import Noiser
from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.loader import SDOps
from ltx_core.loader.attention_ops import set_attention_module_op
from ltx_core.loader.fuse_loras import bf16_fuse_rule
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.primitives import BuilderProtocol, LoraPathStrengthAndSDOps, ModelBuilderProtocol
from ltx_core.loader.registry import DummyRegistry, Registry
@@ -42,7 +46,15 @@ from ltx_core.model.transformer import (
LTXModelConfigurator,
X0Model,
)
from ltx_core.model.transformer.compiling import COMPILE_TRANSFORMER, modify_sd_ops_for_compilation
from ltx_core.model.transformer.attention import (
AttentionCallable,
AttentionFunction,
)
from ltx_core.model.transformer.compiling import (
CompilationConfig,
build_compile_transformer_op,
modify_sd_ops_for_compilation,
)
from ltx_core.model.upsampler import LatentUpsamplerConfigurator, upsample_video
from ltx_core.model.video_vae import (
MEMORY_EFFICIENT_DECODE,
@@ -53,7 +65,7 @@ from ltx_core.model.video_vae import (
VideoEncoder,
VideoEncoderConfigurator,
)
from ltx_core.quantization import QuantizationPolicy
from ltx_core.quantization import QuantizationPolicy, fp8_cast_fuse_rule
from ltx_core.text_encoders.gemma import (
EMBEDDINGS_PROCESSOR_KEY_OPS,
GEMMA_LLM_KEY_OPS,
@@ -101,6 +113,28 @@ def _chain_quantization(
return chained_sd_ops, (*module_ops, *quantization.module_ops)
def _apply_compile_ops(
sd_ops: SDOps,
module_ops: tuple[ModuleOps, ...],
loras: tuple[LoraPathStrengthAndSDOps, ...],
number_of_layers: int,
compilation_config: CompilationConfig,
) -> tuple[SDOps, tuple[ModuleOps, ...], tuple[LoraPathStrengthAndSDOps, ...]]:
"""Rewrite sd_ops/module_ops/LoRAs for compiled blocks (params land under ``_orig_mod``)."""
sd_ops = modify_sd_ops_for_compilation(sd_ops, number_of_layers)
compile_op = build_compile_transformer_op(compilation_config)
module_ops = (*module_ops, compile_op)
loras = tuple(
LoraPathStrengthAndSDOps(
lora.path,
lora.strength,
modify_sd_ops_for_compilation(lora.sd_ops, number_of_layers),
)
for lora in loras
)
return sd_ops, module_ops, loras
@contextmanager
def _streaming_model(
builder: StreamingModelBuilder,
@@ -170,79 +204,110 @@ class DiffusionStage:
loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
transformer_builder: ModelBuilderProtocol[LTXModel] | DelegatingBuilder[LTXModel] | None = None,
) -> None:
self._checkpoint_path = checkpoint_path
self._dtype = dtype
self._device = device
self._quantization = quantization
self._torch_compile = torch_compile
self._compilation_config = compilation_config
self._offload_mode = offload_mode
configurator = (
quantization.model_configurator
if quantization is not None and quantization.model_configurator is not None
else LTXModelConfigurator
)
if transformer_builder is not None:
self._transformer_builder = transformer_builder
else:
self._transformer_builder = Builder(
model_path=checkpoint_path,
model_class_configurator=LTXModelConfigurator,
model_class_configurator=configurator,
model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
loras=tuple(loras),
registry=registry or DummyRegistry(),
)
if offload_mode != OffloadMode.NONE:
if torch_compile:
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.
if quantization is not None and quantization.fuse_rule is not fp8_cast_fuse_rule:
raise ValueError(
"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_module_ops: tuple[ModuleOps, ...] = ()
if quantization is not None:
if quantization.kind != QuantizationPolicy.Kind.FP8_CAST:
raise ValueError(
f"Layer streaming supports only QuantizationPolicy.fp8_cast(); "
f"got kind={quantization.kind!r} which produces heterogeneous block layouts."
)
streaming_sd_ops, streaming_module_ops = _chain_quantization(
streaming_sd_ops, streaming_module_ops, quantization
)
self._streaming_builder = StreamingModelBuilder(
model_class_configurator=LTXModelConfigurator,
model_class_configurator=configurator,
model_path=checkpoint_path,
model_sd_ops=streaming_sd_ops,
module_ops=streaming_module_ops,
loras=tuple(loras),
registry=registry or DummyRegistry(),
blocks_attr="velocity_model.transformer_blocks",
fuse_rule=quantization.fuse_rule if quantization is not None else bf16_fuse_rule,
blocks_attr="transformer_blocks",
blocks_prefix="transformer_blocks",
state_dict_prefix="velocity_model.",
model_wrapper=lambda m: X0Model(m).eval(),
)
def with_attention(self, attention: AttentionFunction | AttentionCallable | None) -> "DiffusionStage":
"""Return a new ``DiffusionStage`` that pins the transformer build to ``attention``.
Functional: never mutates ``self``. The returned stage shares all other
configuration with the original; only the underlying builders' ``module_ops``
gain a ``set_attention_module_op(attention)`` entry so subsequent transformer
builds use that kernel. ``attention=None`` is a no-op (returns ``self``).
"""
if attention is None:
return self
op = set_attention_module_op(attention)
new = copy.copy(self)
new._transformer_builder = self._transformer_builder.with_module_ops(
(*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),
)
return new
def _build_transformer(self, *, device: torch.device | None = None, **kwargs: object) -> X0Model:
target = device or self._device
sd_ops = self._transformer_builder.model_sd_ops
module_ops = self._transformer_builder.module_ops
loras = self._transformer_builder.loras
if self._torch_compile:
module_ops = (*module_ops, COMPILE_TRANSFORMER)
if self._compilation_config is not None:
number_of_layers = self._transformer_builder.model_config()["transformer"]["num_layers"]
sd_ops = modify_sd_ops_for_compilation(sd_ops, number_of_layers)
loras = tuple(
LoraPathStrengthAndSDOps(
lora.path,
lora.strength,
modify_sd_ops_for_compilation(lora.sd_ops, number_of_layers),
)
for lora in loras
sd_ops, module_ops, loras = _apply_compile_ops(
sd_ops, module_ops, loras, number_of_layers, self._compilation_config
)
if self._quantization is not None:
sd_ops, module_ops = _chain_quantization(sd_ops, module_ops, self._quantization)
builder = self._transformer_builder.with_module_ops(module_ops).with_sd_ops(sd_ops).with_loras(loras)
if self._quantization is not None:
builder = builder.with_fuse_rule(self._quantization.fuse_rule)
return X0Model(builder.build(device=target, **kwargs)).to(target).eval()
@contextmanager
def _streaming_transformer_ctx(self) -> Iterator[X0Model]:
with _streaming_model(
self._streaming_builder, self._offload_mode, self._device, self._dtype
) as streaming_wrapper:
yield X0Model(streaming_wrapper).eval()
def _transformer_ctx(self, **kwargs: object) -> AbstractContextManager:
if self._offload_mode != OffloadMode.NONE:
return _streaming_model(self._streaming_builder, self._offload_mode, self._device, self._dtype)
return self._streaming_transformer_ctx()
return gpu_model(self._build_transformer(**kwargs))
def model_context(self, **kwargs: object) -> AbstractContextManager:
@@ -348,7 +413,17 @@ class DiffusionStage:
v_shape = VideoLatentShape.from_pixel_shape(pixel_shape)
video_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape, fps)
mode = "streaming" if self._offload_mode != OffloadMode.NONE else "standard"
logger.info("Building transformer (%s) from %s", mode, self._checkpoint_path)
with self._transformer_ctx(video_tools=video_tools) as transformer:
logger.info(
"Running denoising loop (%d steps, %dx%d %d frames @ %.1f fps)",
len(sigmas) - 1,
width,
height,
frames,
fps,
)
return self.run(
transformer,
denoiser,
@@ -387,6 +462,8 @@ class PromptEncoder:
offload_mode: OffloadMode = OffloadMode.NONE,
text_encoder_builder: BuilderProtocol | None = None,
) -> None:
self._gemma_root = gemma_root
self._checkpoint_path = checkpoint_path
self._dtype = dtype
self._device = device
self._offload_mode = offload_mode
@@ -432,7 +509,7 @@ class PromptEncoder:
def _build_embeddings_processor(self) -> EmbeddingsProcessor:
"""Build the embeddings processor on the target device."""
return self._embeddings_processor_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
return self._embeddings_processor_builder.build(device=self._device, dtype=self._dtype).eval()
def _text_encoder_ctx(self) -> AbstractContextManager:
if self._offload_mode != OffloadMode.NONE:
@@ -448,6 +525,7 @@ class PromptEncoder:
enhance_prompt_seed: int = 42,
) -> list[EmbeddingsProcessorOutput]:
"""Encode *prompts* through Gemma -> embeddings processor, freeing each model after use."""
logger.info("Building text encoder from %s", self._gemma_root)
with self._text_encoder_ctx() as text_encoder:
if enhance_first_prompt:
prompts = list(prompts)
@@ -455,9 +533,12 @@ class PromptEncoder:
text_encoder, prompts[0], enhance_prompt_image, seed=enhance_prompt_seed
)
raw_outputs = [text_encoder.encode(p) for p in prompts]
logger.info("Text encoder done, building embeddings processor from %s", self._checkpoint_path)
with gpu_model(self._build_embeddings_processor()) as embeddings_processor:
return [embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs]
result = [embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs]
logger.info("Prompt encoding complete")
return result
# ---------------------------------------------------------------------------
@@ -487,7 +568,7 @@ class ImageConditioner:
)
def _build_encoder(self) -> VideoEncoder:
return self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
return self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()
def __call__(self, fn: Callable[[VideoEncoder], T]) -> T:
"""Build video encoder → call *fn(encoder)* → free encoder."""
@@ -511,6 +592,7 @@ class VideoUpsampler:
device: torch.device,
registry: Registry | None = None,
) -> None:
self._upsampler_path = upsampler_path
self._dtype = dtype
self._device = device
self._encoder_builder = Builder(
@@ -527,13 +609,10 @@ class VideoUpsampler:
def __call__(self, latent: torch.Tensor) -> torch.Tensor:
"""Upsample *latent* using video encoder + spatial upsampler, then free both."""
logger.info("Building video encoder + spatial upsampler from %s", self._upsampler_path)
with (
gpu_model(
self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
) as encoder,
gpu_model(
self._upsampler_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
) as upsampler,
gpu_model(self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()) as encoder,
gpu_model(self._upsampler_builder.build(device=self._device, dtype=self._dtype).eval()) as upsampler,
):
return upsample_video(latent=latent, video_encoder=encoder, upsampler=upsampler)
@@ -557,6 +636,7 @@ class VideoDecoder:
memory_efficient: bool = True,
decoder_builder: BuilderProtocol | None = None,
) -> None:
self._checkpoint_path = checkpoint_path
self._dtype = dtype
self._device = device
if decoder_builder is not None:
@@ -577,7 +657,8 @@ class VideoDecoder:
generator: torch.Generator | None = None,
) -> Iterator[torch.Tensor]:
"""Decode *latent* to pixel-space video chunks. Decoder freed after exhaustion."""
decoder = self._decoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
logger.info("Building video decoder from %s", self._checkpoint_path)
decoder = self._decoder_builder.build(device=self._device, dtype=self._dtype).eval()
return _cleanup_iter(decoder.decode_video(latent, tiling_config, generator), decoder)
@@ -596,6 +677,7 @@ class AudioDecoder:
device: torch.device,
registry: Registry | None = None,
) -> None:
self._checkpoint_path = checkpoint_path
self._dtype = dtype
self._device = device
self._decoder_builder = Builder(
@@ -613,13 +695,10 @@ class AudioDecoder:
def __call__(self, latent: torch.Tensor) -> Audio:
"""Decode audio *latent* through VAE decoder + vocoder, then free both."""
logger.info("Building audio decoder + vocoder from %s", self._checkpoint_path)
with (
gpu_model(
self._decoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
) as decoder,
gpu_model(
self._vocoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
) as vocoder,
gpu_model(self._decoder_builder.build(device=self._device, dtype=self._dtype).eval()) as decoder,
gpu_model(self._vocoder_builder.build(device=self._device, dtype=self._dtype).eval()) as vocoder,
):
return vae_decode_audio(latent, decoder, vocoder)
@@ -653,7 +732,5 @@ class AudioConditioner:
def __call__(self, fn: Callable[[torch.nn.Module], T]) -> T:
"""Build audio encoder → call *fn(encoder)* → free encoder."""
with gpu_model(
self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
) as encoder:
with gpu_model(self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()) as encoder:
return fn(encoder)
packages/ltx-pipelines/src/ltx_pipelines/utils/denoisers.py
+11 -1
View File
@@ -138,6 +138,8 @@ def _guided_denoise( # noqa: PLR0913,PLR0915
ptb_configs = [ptb for _, _, _, ptb in passes]
n = len(passes)
orig_b = (video_state or audio_state).latent.shape[0]
def _batched_sigma(state: LatentState) -> torch.Tensor:
"""Expand scalar sigma to (n * B,) matching the repeated state."""
return sigma.expand(state.latent.shape[0] * n)
@@ -162,8 +164,16 @@ def _guided_denoise( # noqa: PLR0913,PLR0915
enabled=not a_skip,
)
# Replicate each pass's PerturbationConfig to all `orig_b` samples it
# carries, so `BatchedPerturbationConfig.mask_like` returns a per-sample
# mask (length n*orig_b) instead of a per-pass mask (length n). Without
# this expansion the mask is broadcast against a (n*orig_b, T, D) tensor
# and the multiplication fails with a batch-dim mismatch whenever
# `orig_b > 1` (e.g. multi-prompt benchmark panels).
batched_ptb_configs = [ptb for ptb in ptb_configs for _ in range(orig_b)]
all_v, all_a = transformer(
video=batched_video, audio=batched_audio, perturbations=BatchedPerturbationConfig(ptb_configs)
video=batched_video, audio=batched_audio, perturbations=BatchedPerturbationConfig(batched_ptb_configs)
)
# Split results back and combine via guiders.
packages/ltx-pipelines/src/ltx_pipelines/utils/quantization_factory.py
+36
View File
@@ -0,0 +1,36 @@
"""User-facing quantization-policy dispatch.
``ltx-core`` exposes one ``build_policy`` factory per backend. This module
provides the user-facing string-keyed dispatch used by CLI args and pipeline
defaults — keeping the enum out of ``ltx-core`` so adding/removing backends is
a single-file change here.
"""
from enum import Enum
from typing_extensions import assert_never
from ltx_core.quantization import QuantizationPolicy
from ltx_core.quantization.fp8_cast import build_policy as _build_fp8_cast_policy
from ltx_core.quantization.fp8_scaled_mm import build_policy as _build_fp8_scaled_mm_policy
class QuantizationKind(str, Enum):
FP8_CAST = "fp8-cast"
FP8_SCALED_MM = "fp8-scaled-mm"
def to_policy(self, checkpoint_path: str | None = None) -> QuantizationPolicy:
"""Build the :class:`QuantizationPolicy` for this kind.
``checkpoint_path`` is required for both backends: ``FP8_SCALED_MM``
uses it to discover the layer set from ``.weight_scale`` tensors,
and ``FP8_CAST`` uses it to fold any prequant scales into the fp8
weight at load time.
"""
if checkpoint_path is None:
raise ValueError(f"{self.value} quantization requires checkpoint_path.")
match self:
case QuantizationKind.FP8_CAST:
return _build_fp8_cast_policy(checkpoint_path)
case QuantizationKind.FP8_SCALED_MM:
return _build_fp8_scaled_mm_policy(checkpoint_path)
case _:
assert_never(self)
packages/ltx-trainer/pyproject.toml
+2 -2
View File
@@ -1,6 +1,6 @@
[project]
name = "ltx-trainer"
version = "1.1.3"
version = "1.1.4"
description = "LTX-2 training, democratized."
readme = "README.md"
authors = [
@@ -48,7 +48,7 @@ build-backend = "hatchling.build"
[tool.ruff]
target-version = "1.1.3"
target-version = "1.1.4"
line-length = 120
[tool.ruff.lint]
uv.lock
Generated
+3 -3
View File
@@ -2063,7 +2063,7 @@ wheels = [
[[package]]
name = "ltx-core"
version = "1.1.3"
version = "1.1.4"
source = { editable = "packages/ltx-core" }
dependencies = [
{ name = "accelerate" },
@@ -2121,7 +2121,7 @@ dev = [{ name = "scikit-image", specifier = ">=0.25.2" }]
[[package]]
name = "ltx-pipelines"
version = "1.1.3"
version = "1.1.4"
source = { editable = "packages/ltx-pipelines" }
dependencies = [
{ name = "av" },
@@ -2144,7 +2144,7 @@ requires-dist = [
[[package]]
name = "ltx-trainer"
version = "1.1.3"
version = "1.1.4"
source = { editable = "packages/ltx-trainer" }
dependencies = [
{ name = "accelerate" },