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

Automated PR - 2026-05-11

Browse Source
This commit is contained in:
github-actions[bot]
2026-05-11 13:14:05 +00:00
parent 41d9243716
commit 7df34dfa83
72 changed files with 3299 additions and 911 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-core/README.md
+7 -3
View File
@@ -77,13 +77,17 @@ model = builder.build(device=torch.device("cuda"))
Use the `.lora()` method to attach one or more LoRA adapters before calling `.build()`:
```python
from ltx_core.loader import SDOps
lora_sd_ops = SDOps(name="identity").with_matching() # or a model-specific key-renaming SDOps
builder = (
SingleGPUModelBuilder(
model_class_configurator=MyModelConfigurator,
model_path="/path/to/model.safetensors",
)
.lora("/path/to/lora_a.safetensors", strength=0.8)
.lora("/path/to/lora_b.safetensors", strength=0.5)
.lora("/path/to/lora_a.safetensors", 0.8, lora_sd_ops)
.lora("/path/to/lora_b.safetensors", 0.5, lora_sd_ops)
)
model = builder.build(device=torch.device("cuda"))
```
@@ -103,7 +107,7 @@ builder = SingleGPUModelBuilder(
model_class_configurator=MyModelConfigurator,
model_path="/path/to/model.safetensors",
lora_load_device=torch.device("cuda"),
).lora("/path/to/lora.safetensors", strength=1.0)
).lora("/path/to/lora.safetensors", 1.0, lora_sd_ops)
model = builder.build(device=torch.device("cuda"))
```
packages/ltx-core/pyproject.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[project]
name = "ltx-core"
version = "1.1.2"
version = "1.1.3"
description = "Core implementation of Lightricks' LTX-2 model"
readme = "README.md"
requires-python = ">=3.10"
packages/ltx-core/src/ltx_core/block_streaming/builder.py
+130 -84
View File
@@ -7,16 +7,17 @@ from collections.abc import Callable
from dataclasses import dataclass, field, replace
from typing import Generic
import safetensors
import torch
from torch import nn
from ltx_core.block_streaming.disk import DiskBlockReader, DiskTensorReader, LoraSource
from ltx_core.block_streaming.pool import BlockLayout, WeightPool
from ltx_core.block_streaming.pool import WeightPool
from ltx_core.block_streaming.provider import WeightsProvider
from ltx_core.block_streaming.source import DiskWeightSource, PinnedWeightSource, WeightSource
from ltx_core.block_streaming.utils import build_pool_layout, resolve_attr
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 apply_loras
from ltx_core.loader.fuse_loras import aggregate_lora_products, fuse_lora_weights
from ltx_core.loader.helpers import create_meta_model, load_state_dict, read_model_config
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.primitives import (
@@ -53,8 +54,8 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
``"velocity_model.transformer_blocks"``).
blocks_prefix: State-dict key prefix for block weights
(e.g. ``"transformer_blocks"``).
state_dict_prefix: Key prefix for non-block weights
(e.g. ``"velocity_model."``).
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``).
"""
@@ -110,7 +111,6 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
if not self.blocks_prefix:
raise ValueError("blocks_prefix must be non-empty for streaming")
# 1. Create meta model (no weights allocated).
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:
@@ -118,22 +118,29 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
meta_model.eval()
blocks = resolve_attr(meta_model, self.blocks_attr)
layout = build_pool_layout(blocks[0], dtype)
# 2. Determine slot counts.
checkpoint_paths = list(self.model_path) if isinstance(self.model_path, tuple) else [self.model_path]
block_key_map, non_block_keys = _scan_checkpoint_keys(checkpoint_paths, self.model_sd_ops, self.blocks_prefix)
cpu_slots_count = cpu_slots_count if cpu_slots_count is not None else len(blocks)
gpu_slots_count = gpu_slots_count if gpu_slots_count is not None else _DEFAULT_GPU_SLOTS
# 3. Build source and load non-block weights.
if cpu_slots_count >= len(blocks):
source, lora_sources = self._build_pinned_source(meta_model, target_device, dtype, cpu_slots_count)
source, lora_sources = self._build_pinned_source(
meta_model, target_device, dtype, cpu_slots_count, block_key_map, non_block_keys
)
else:
source, lora_sources = self._build_disk_source(meta_model, layout, target_device, dtype, cpu_slots_count)
reader = DiskTensorReader(checkpoint_paths)
source, lora_sources = self._build_disk_source(
meta_model, target_device, dtype, cpu_slots_count, reader, block_key_map, non_block_keys
)
# 4. Create provider and wrapper.
copy_stream = torch.cuda.Stream(device=target_device)
gpu_pool = WeightPool(
layout, gpu_slots_count, target_device, reuse_barrier=lambda event: copy_stream.wait_event(event)
source.block_layout,
gpu_slots_count,
target_device,
reuse_barrier=lambda event: copy_stream.wait_event(event),
)
provider = WeightsProvider(gpu_pool, copy_stream, target_device, source, lora_sources, self.blocks_prefix)
return BlockStreamingWrapper(
@@ -149,89 +156,90 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
target_device: torch.device,
dtype: torch.dtype,
cpu_slots_count: int,
block_key_map: dict[int, list[tuple[str, str]]],
non_block_keys: list[tuple[str, str]],
) -> tuple[WeightSource, list[LoraSource]]:
"""Pre-load all blocks into pinned CPU buffers with LoRA fusion."""
model_sd = load_state_dict(
self.model_path, self.model_loader, self.registry, torch.device("cpu"), self.model_sd_ops
)
if self.loras:
lora_sds = [
load_state_dict([lora.path], self.model_loader, self.registry, torch.device("cpu"), lora.sd_ops)
for lora in self.loras
]
lora_sd_and_strengths = [
LoraStateDictWithStrength(sd, lora.strength) for sd, lora in zip(lora_sds, self.loras, strict=True)
]
model_sd = apply_loras(
model_sd=model_sd,
lora_sd_and_strengths=lora_sd_and_strengths,
dtype=dtype,
destination_sd=model_sd if isinstance(self.registry, DummyRegistry) else None,
lora_sd_and_strengths = [
LoraStateDictWithStrength(
load_state_dict([lora.path], self.model_loader, self.registry, torch.device("cpu"), lora.sd_ops),
lora.strength,
)
for lora in self.loras
]
# Partition: non-block weights go to GPU, block weights go directly
# to pinned buffers. This avoids holding the full state dict and
# pinned copies simultaneously.
non_block_sd: dict[str, torch.Tensor] = {}
block_tensors: dict[int, dict[str, torch.Tensor]] = {}
prefix_dot = self.blocks_prefix + "."
for block_idx in block_key_map:
if block_idx >= cpu_slots_count:
raise ValueError(
f"Pinned source requires one CPU slot per block; "
f"got block index {block_idx} with only {cpu_slots_count} slots."
)
for key, tensor in model_sd.sd.items():
if key.startswith(prefix_dot):
rest = key[len(prefix_dot) :]
idx_str, _, param_name = rest.partition(".")
try:
block_idx = int(idx_str)
except ValueError:
non_block_sd[self.state_dict_prefix + key] = tensor.to(device=target_device, dtype=dtype)
continue
block_tensors.setdefault(block_idx, {})[param_name] = tensor
blocks = resolve_attr(meta_model, self.blocks_attr)
block_tensors: dict[str, torch.Tensor] = {}
for block_idx, entries in block_key_map.items():
block_params = dict(blocks[block_idx].named_parameters())
for _sft_key, param_name in entries:
key = make_block_key(self.blocks_prefix, block_idx, param_name)
block_tensors[key] = block_params[param_name]
blocks_layout = derive_layout(block_tensors, dtype)
pinned_blocks = allocate_layout_views(blocks_layout, pin_memory=True)
should_sync = False
for key, fused in fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype=None, preserve_input_device=False):
if key in pinned_blocks:
pinned_blocks[key].copy_(fused, non_blocking=True)
model_sd.sd[key] = None
should_sync = True
else:
non_block_sd[self.state_dict_prefix + key] = tensor.to(device=target_device, dtype=dtype)
model_sd.sd[key] = fused
if should_sync:
torch.cuda.synchronize()
# Fill remaining pinned keys from the source state dict.
for key in blocks_layout:
if model_sd.sd[key] is None:
continue
pinned_blocks[key].copy_(model_sd.sd[key])
model_sd.sd[key] = None
pinned: dict[int, dict[str, torch.Tensor]] = {
block_idx: {
param_name: pinned_blocks[make_block_key(self.blocks_prefix, block_idx, param_name)]
for _sft_key, param_name in entries
}
for block_idx, entries in block_key_map.items()
}
non_block_sd: dict[str, torch.Tensor] = {
self.state_dict_prefix + model_key: model_sd.sd[model_key].to(device=target_device, dtype=dtype)
for _sft_key, model_key in non_block_keys
}
meta_model.load_state_dict(non_block_sd, strict=False, assign=True)
del model_sd, non_block_sd
# Pin block weights one block at a time, freeing the source tensors as we go.
pinned: dict[int, dict[str, torch.Tensor]] = {}
for idx in range(cpu_slots_count):
src = block_tensors.pop(idx)
pinned[idx] = {name: tensor.to(dtype=dtype).pin_memory() for name, tensor in src.items()}
return PinnedWeightSource(pinned), []
def _build_disk_source(
self,
meta_model: nn.Module,
layout: BlockLayout,
target_device: torch.device,
dtype: torch.dtype,
cpu_slots_count: int,
reader: DiskTensorReader,
block_key_map: dict[int, list[tuple[str, str]]],
non_block_keys: list[tuple[str, str]],
) -> tuple[WeightSource, list[LoraSource]]:
"""Create a DiskWeightSource backed by a DiskBlockReader for lazy loading."""
"""Create a DiskWeightSource backed by a DiskBlockReader for lazy loading.
Derives the shared pool layout from the meta model's block 0 — this
relies on module_ops (e.g. fp8_cast) leaving the meta param dtype in
sync with the post-sd_ops checkpoint dtype.
"""
lora_sources = [LoraSource(lora.path, lora.sd_ops, lora.strength) for lora in self.loras]
checkpoint_paths = list(self.model_path) if isinstance(self.model_path, tuple) else [self.model_path]
reader = DiskTensorReader(checkpoint_paths)
block_key_map: dict[int, list[tuple[str, str]]] = {}
non_block_keys: list[tuple[str, str]] = []
for sft_key in reader.keys(): # noqa: SIM118
model_key = self.model_sd_ops.apply_to_key(sft_key) if self.model_sd_ops else sft_key
if model_key is None:
continue
if model_key.startswith(self.blocks_prefix + "."):
rest = model_key[len(self.blocks_prefix) + 1 :]
idx_str, _, param_name = rest.partition(".")
try:
block_idx = int(idx_str)
except ValueError:
non_block_keys.append((sft_key, model_key))
continue
block_key_map.setdefault(block_idx, []).append((sft_key, param_name))
else:
non_block_keys.append((sft_key, model_key))
self._load_non_block_weights(
reader,
@@ -242,9 +250,11 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
sd_ops=self.model_sd_ops,
key_prefix=self.state_dict_prefix,
lora_sources=lora_sources,
matmul_device=target_device,
)
blocks = resolve_attr(meta_model, self.blocks_attr)
layout = derive_layout(dict(blocks[0].named_parameters()), dtype)
cpu_pool = WeightPool(
layout,
cpu_slots_count,
@@ -252,7 +262,12 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
reuse_barrier=lambda event: event.synchronize(),
pin_memory=True,
)
block_reader = DiskBlockReader(reader=reader, block_key_map=block_key_map, dtype=dtype)
block_reader = DiskBlockReader(
reader=reader,
block_key_map=block_key_map,
sd_ops=self.model_sd_ops,
blocks_prefix=self.blocks_prefix,
)
source = DiskWeightSource(cpu_pool, block_reader)
return source, lora_sources
@@ -265,17 +280,17 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
model_key: str,
tensor: torch.Tensor,
lora_sources: list[LoraSource],
matmul_device: torch.device | None = None,
) -> torch.Tensor:
"""Add all matching LoRA deltas to *tensor* in-place."""
"""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")]
device = tensor.device if tensor.device.type == "cuda" else matmul_device
for source in lora_sources:
delta = source.get_delta(prefix, device=device)
if delta is not None:
tensor = tensor.add_(delta.to(device=tensor.device, dtype=tensor.dtype))
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
@@ -289,17 +304,48 @@ class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType])
sd_ops: SDOps | None = None,
key_prefix: str = "",
lora_sources: list[LoraSource] | None = None,
matmul_device: torch.device | None = None,
) -> None:
"""Load non-block weights into *model* on *device*."""
state_dict: dict[str, torch.Tensor] = {}
sources = lora_sources or []
for sft_key, model_key in non_block_keys:
tensor = reader.get_tensor(sft_key).to(device=device, dtype=dtype)
tensor = StreamingModelBuilder._fuse_lora_delta(model_key, tensor, sources, matmul_device)
tensor = StreamingModelBuilder._fuse_lora_delta(model_key, tensor, sources)
if sd_ops is not None:
for kv in sd_ops.apply_to_key_value(model_key, tensor):
state_dict[key_prefix + kv.new_key] = kv.new_value
continue
state_dict[key_prefix + model_key] = tensor
model.load_state_dict(state_dict, strict=False, assign=True)
def _scan_checkpoint_keys(
checkpoint_paths: list[str],
sd_ops: SDOps | None,
blocks_prefix: str,
) -> tuple[dict[int, list[tuple[str, str]]], list[tuple[str, str]]]:
"""Partition checkpoint keys into per-block and non-block lists.
Opens the safetensors files for header-only key enumeration; no tensor data
is read.
"""
block_key_map: dict[int, list[tuple[str, str]]] = {}
non_block_keys: list[tuple[str, str]] = []
prefix_dot = blocks_prefix + "."
for path in checkpoint_paths:
with safetensors.safe_open(path, framework="pt", device="cpu") as handle:
for sft_key in handle.keys(): # noqa: SIM118
model_key = sd_ops.apply_to_key(sft_key) if sd_ops else sft_key
if model_key is None:
continue
if model_key.startswith(prefix_dot):
rest = model_key[len(prefix_dot) :]
idx_str, _, param_name = rest.partition(".")
try:
block_idx = int(idx_str)
except ValueError:
non_block_keys.append((sft_key, model_key))
continue
block_key_map.setdefault(block_idx, []).append((sft_key, param_name))
else:
non_block_keys.append((sft_key, model_key))
return block_key_map, non_block_keys
packages/ltx-core/src/ltx_core/block_streaming/disk.py
+71 -29
View File
@@ -2,11 +2,22 @@
from __future__ import annotations
from collections.abc import Iterator
import safetensors
import torch
from ltx_core.block_streaming.utils import allocate_layout_views, make_block_key
from ltx_core.loader.fuse_loras import LoraProduct
from ltx_core.loader.sd_ops import SDOps
_SAFETENSORS_DTYPE_TO_TORCH: dict[str, torch.dtype] = {
"F64": torch.float64,
"F32": torch.float32,
"F16": torch.float16,
"BF16": torch.bfloat16,
}
class DiskTensorReader:
"""Key-based tensor accessor over one or more safetensors files."""
@@ -21,9 +32,6 @@ class DiskTensorReader:
for sft_key in handle.keys(): # noqa: SIM118
self._key_to_handle_idx[sft_key] = handle_idx
def keys(self) -> list[str]:
return list(self._key_to_handle_idx.keys())
def get_tensor(self, key: str) -> torch.Tensor:
return self._handles[self._key_to_handle_idx[key]].get_tensor(key)
@@ -31,49 +39,58 @@ class DiskTensorReader:
self._handles.clear()
self._key_to_handle_idx.clear()
def __contains__(self, key: str) -> bool:
return key in self._key_to_handle_idx
def __iter__(self) -> Iterator[str]:
return iter(self._key_to_handle_idx)
class DiskBlockReader:
"""Reads one block at a time from safetensors into provided buffers.
Maps block indices to safetensors keys via a pre-computed key map.
"""
"""Reads one block at a time from safetensors into provided buffers."""
def __init__(
self,
reader: DiskTensorReader,
block_key_map: dict[int, list[tuple[str, str]]],
dtype: torch.dtype,
sd_ops: SDOps | None = None,
blocks_prefix: str = "",
) -> None:
self._reader = reader
self._block_key_map = block_key_map
self._dtype = dtype
self._sd_ops = sd_ops
self._blocks_prefix = blocks_prefix
def read_into(self, target: dict[str, torch.Tensor], block_idx: int) -> None:
block_prefix = make_block_key(self._blocks_prefix, block_idx, "")
for sft_key, param_name in self._block_key_map[block_idx]:
tensor = self._reader.get_tensor(sft_key)
if tensor.dtype != self._dtype:
tensor = tensor.to(self._dtype)
target[param_name].copy_(tensor)
if self._sd_ops is None:
target[param_name].copy_(tensor)
continue
full_key = make_block_key(self._blocks_prefix, block_idx, param_name)
for result in self._sd_ops.apply_to_key_value(full_key, tensor):
if not result.new_key.startswith(block_prefix):
raise ValueError(
f"SDOps output key '{result.new_key}' is outside block {block_idx} "
f"(expected prefix '{block_prefix}'); cannot route to a per-block buffer."
)
target[result.new_key[len(block_prefix) :]].copy_(result.new_value)
def cleanup(self) -> None:
self._reader.close()
class LoraSource:
"""Pinned-memory cache of LoRA A/B matrices for on-the-fly fusion.
At init, loads all matched A/B pairs into pinned CPU memory.
:meth:`get_delta` computes ``(B * strength) @ A`` on the given device.
"""
"""Pinned-memory cache of matched LoRA A/B factors backed by a single buffer."""
def __init__(self, path: str, sd_ops: SDOps | None, strength: float) -> None:
self.strength = strength
# param_prefix -> (pinned_a, pinned_b)
self._pinned_ab: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
a_keys: dict[str, str] = {}
b_keys: dict[str, str] = {}
with safetensors.safe_open(path, framework="pt", device="cpu") as handle:
# First pass: build key map.
for sft_key in handle.keys(): # noqa: SIM118
model_key = sd_ops.apply_to_key(sft_key) if sd_ops is not None else sft_key
if model_key is None:
@@ -83,24 +100,49 @@ class LoraSource:
elif model_key.endswith(".lora_B.weight"):
b_keys[model_key[: -len(".lora_B.weight")]] = sft_key
# Second pass: load and pin matched A+B pairs (orphans silently skipped).
for prefix in a_keys.keys() & b_keys.keys():
self._pinned_ab[prefix] = (
handle.get_tensor(a_keys[prefix]).pin_memory(),
handle.get_tensor(b_keys[prefix]).pin_memory(),
matched_prefixes = list(a_keys.keys() & b_keys.keys())
# Build the layout from safetensors header metadata only — no tensor data is read.
layout: dict[str, tuple[torch.Size, torch.dtype]] = {}
for prefix in matched_prefixes:
a_slice_view = handle.get_slice(a_keys[prefix])
b_slice_view = handle.get_slice(b_keys[prefix])
layout[f"{prefix}.A"] = (
torch.Size(a_slice_view.get_shape()),
_SAFETENSORS_DTYPE_TO_TORCH[a_slice_view.get_dtype()],
)
layout[f"{prefix}.B"] = (
torch.Size(b_slice_view.get_shape()),
_SAFETENSORS_DTYPE_TO_TORCH[b_slice_view.get_dtype()],
)
def get_delta(self, param_prefix: str, device: torch.device | None = None) -> torch.Tensor | None:
"""Return ``(B * strength) @ A`` for *param_prefix*, or ``None``."""
all_views = allocate_layout_views(layout, pin_memory=True)
for prefix in matched_prefixes:
a_view = all_views[f"{prefix}.A"]
b_view = all_views[f"{prefix}.B"]
a_view.copy_(handle.get_tensor(a_keys[prefix]))
b_view.copy_(handle.get_tensor(b_keys[prefix]))
self._pinned_ab[prefix] = (a_view, b_view)
def get_ab(
self,
param_prefix: str,
device: torch.device | None = None,
dtype: torch.dtype | None = None,
) -> LoraProduct | None:
"""Return the :class:`LoraProduct` for *param_prefix*, or ``None``."""
pair = self._pinned_ab.get(param_prefix)
if pair is None:
return None
a, b = pair
if device is not None and device.type == "cuda":
a = a.to(device=device)
b = b.to(device=device)
delta = torch.matmul(b * self.strength, a)
return delta
a = a.to(device=device, non_blocking=True)
b = b.to(device=device, non_blocking=True)
if dtype is not None:
a = a.to(dtype=dtype)
b = b.to(dtype=dtype)
return LoraProduct(a, b, self.strength)
def cleanup(self) -> None:
self._pinned_ab.clear()
packages/ltx-core/src/ltx_core/block_streaming/pool.py
+24 -13
View File
@@ -7,20 +7,16 @@ from typing import Callable
import torch
from ltx_core.block_streaming.utils import allocate_buffer
# Type alias for the buffer layout used by slot allocation.
BlockLayout = dict[str, tuple[torch.Size, torch.dtype]]
from ltx_core.block_streaming.utils import allocate_layout_views
from ltx_core.loader.primitives import TensorLayout
class WeightPool:
"""Fixed pool of pre-allocated weight buffers with event-based reuse safety.
Buffers are allocated once at construction. :meth:`acquire` pops a
free buffer (waiting any pending event first). :meth:`release`
returns it, optionally attaching an event that must complete before
the buffer can be reused.
"""Fixed pool of pre-allocated weight buffers with event-based reuse.
All slots share a single buffer (CPU or GPU); each slot is a
contiguous slice carved out of it via :func:`allocate_layout_views`.
Args:
layout: ``{name: (shape, dtype)}`` for each buffer.
buffer_layout: ``{name: (shape, dtype)}`` for each buffer.
capacity: Number of buffers to pre-allocate.
device: Device for allocation.
reuse_barrier: Called with the pending event before a buffer is reused.
@@ -29,23 +25,34 @@ class WeightPool:
def __init__(
self,
layout: BlockLayout,
buffer_layout: TensorLayout,
capacity: int,
device: torch.device,
reuse_barrier: Callable[[torch.cuda.Event], None],
pin_memory: bool = False,
) -> None:
self._buffer_layout = buffer_layout
self._capacity = capacity
self._free: deque[dict[str, torch.Tensor]] = deque()
self._events: dict[int, torch.cuda.Event] = {}
self._reuse_barrier = reuse_barrier
for _ in range(capacity):
self._free.append(allocate_buffer(layout, device, pin_memory))
memory_layout = {
_make_key(slot, name): (shape, dtype)
for slot in range(capacity)
for name, (shape, dtype) in buffer_layout.items()
}
all_views = allocate_layout_views(memory_layout, device=device, pin_memory=pin_memory)
for slot in range(capacity):
self._free.append({name: all_views[_make_key(slot, name)] for name in buffer_layout})
@property
def capacity(self) -> int:
return self._capacity
@property
def buffer_layout(self) -> TensorLayout:
return self._buffer_layout
def acquire(self) -> dict[str, torch.Tensor]:
"""Take a free buffer, waiting any pending event before returning."""
weights = self._free.popleft()
@@ -62,3 +69,7 @@ class WeightPool:
if event is not None:
self._events[id(weights)] = event
self._free.append(weights)
def _make_key(slot: int, name: str) -> str:
return f"{slot}/{name}"
packages/ltx-core/src/ltx_core/block_streaming/provider.py
+45 -8
View File
@@ -9,6 +9,29 @@ import torch
from ltx_core.block_streaming.disk import LoraSource
from ltx_core.block_streaming.pool import WeightPool
from ltx_core.block_streaming.source import WeightSource
from ltx_core.block_streaming.utils import FP8_DTYPES
from ltx_core.loader.fuse_loras import aggregate_lora_products, fuse_cast_fp8_weight
def _contiguous_byte_view(weights: dict[str, torch.Tensor]) -> torch.Tensor | None:
"""Return a ``uint8`` view spanning every tensor in *weights*, or ``None`` if
they don't share one contiguous storage region."""
tensors = list(weights.values())
if not tensors:
return None
storage = tensors[0].untyped_storage()
storage_ptr = storage.data_ptr()
start = end = tensors[0].storage_offset() * tensors[0].element_size()
for t in tensors:
if t.untyped_storage().data_ptr() != storage_ptr or not t.is_contiguous():
return None
offset = t.storage_offset() * t.element_size()
nbytes = t.numel() * t.element_size()
start = min(start, offset)
end = max(end, offset + nbytes)
view = torch.empty(0, dtype=torch.uint8, device=tensors[0].device)
view.set_(storage, start, (end - start,), (1,))
return view
class WeightsProvider:
@@ -70,8 +93,13 @@ class WeightsProvider:
instrumentation regions wrapping it -- observe the full transfer time.
"""
with torch.cuda.stream(self._copy_stream):
for name, gpu_tensor in gpu_weights.items():
gpu_tensor.copy_(cpu_weights[name], non_blocking=True)
gpu_view = _contiguous_byte_view(gpu_weights)
cpu_view = _contiguous_byte_view(cpu_weights)
if gpu_view is not None and cpu_view is not None and gpu_view.numel() == cpu_view.numel():
gpu_view.copy_(cpu_view, non_blocking=True)
else:
for name, gpu_tensor in gpu_weights.items():
gpu_tensor.copy_(cpu_weights[name], non_blocking=True)
if self._lora_sources:
self._fuse_block_loras(idx, gpu_weights)
h2d_event = torch.cuda.Event()
@@ -102,9 +130,18 @@ class WeightsProvider:
for name, tensor in weights.items():
if not name.endswith(".weight"):
continue
full_key = f"{self._blocks_prefix}.{idx}.{name}"
prefix = full_key[: -len(".weight")]
for source in self._lora_sources:
delta = source.get_delta(prefix, device=self._target_device)
if delta is not None:
tensor.add_(delta.to(dtype=tensor.dtype))
prefix = f"{self._blocks_prefix}.{idx}.{name}".removesuffix(".weight")
is_fp8 = tensor.dtype in FP8_DTYPES
agg_dtype = torch.bfloat16 if is_fp8 else tensor.dtype
products = (
ab
for ab in (s.get_ab(prefix, device=self._target_device, dtype=agg_dtype) for s in self._lora_sources)
if ab is not None
)
aggregated = aggregate_lora_products(products, agg_dtype)
if aggregated is None:
continue
if is_fp8:
tensor.copy_(fuse_cast_fp8_weight(aggregated, tensor, tensor.dtype))
else:
tensor.add_(aggregated)
packages/ltx-core/src/ltx_core/block_streaming/source.py
+20 -1
View File
@@ -9,10 +9,18 @@ import torch
from ltx_core.block_streaming.disk import DiskBlockReader
from ltx_core.block_streaming.pool import WeightPool
from ltx_core.loader.primitives import TensorLayout
class WeightSource(Protocol):
"""Provides pinned CPU weights for a given block index."""
"""Provides pinned CPU weights for a given block index.
Assumes all buffers share an identical layout across all block indices.
"""
@property
def block_layout(self) -> TensorLayout:
"""Shared per-block buffer layout (shape + dtype for each param)."""
...
def get(self, idx: int) -> dict[str, torch.Tensor]:
"""Return CPU weights for block *idx*."""
@@ -36,6 +44,10 @@ class DiskWeightSource(WeightSource):
self._events: dict[int, torch.cuda.Event] = {}
self._reader = reader
@property
def block_layout(self) -> TensorLayout:
return self._pool.buffer_layout
def get(self, idx: int) -> dict[str, torch.Tensor]:
"""Return CPU weights for block *idx*. Reads from disk on miss."""
if idx in self._cache:
@@ -68,8 +80,15 @@ class PinnedWeightSource(WeightSource):
"""Pre-loaded pinned CPU weights."""
def __init__(self, weights: dict[int, dict[str, torch.Tensor]]) -> None:
if not weights:
raise ValueError("PinnedWeightSource requires at least one block")
self._weights = weights
@property
def block_layout(self) -> TensorLayout:
first_block = self._weights[min(self._weights)]
return {name: (t.shape, t.dtype) for name, t in first_block.items()}
def get(self, idx: int) -> dict[str, torch.Tensor]:
return self._weights[idx]
packages/ltx-core/src/ltx_core/block_streaming/utils.py
+93 -19
View File
@@ -2,14 +2,24 @@
from __future__ import annotations
import itertools
from typing import TYPE_CHECKING, Any
import math
import weakref
from dataclasses import dataclass
from typing import Any
import torch
from torch import nn
if TYPE_CHECKING:
from ltx_core.block_streaming.pool import BlockLayout
from ltx_core.loader.primitives import TensorLayout
FP8_DTYPES = frozenset({torch.float8_e4m3fn, torch.float8_e5m2})
_BUFFER_ALIGN = 16
def make_block_key(blocks_prefix: str, block_idx: int, param_name: str) -> str:
"""Return the state-dict key for *param_name* under block *block_idx*."""
return f"{blocks_prefix}.{block_idx}.{param_name}"
def resolve_attr(module: nn.Module, dotted_path: str) -> nn.ModuleList:
@@ -40,21 +50,85 @@ def assign_tensor_to_module(root: nn.Module, dotted_name: str, tensor: torch.Ten
raise AttributeError(f"{leaf} is not a parameter or buffer of {type(parent).__name__}")
def build_pool_layout(block: nn.Module, dtype: torch.dtype) -> BlockLayout:
"""Derive a buffer layout from a block's parameters and buffers.
Works on meta-device blocks (shapes are valid regardless of device).
The *dtype* argument overrides each tensor's dtype so the pool matches
the target inference precision.
def derive_layout(tensors: dict[str, torch.Tensor], dtype: torch.dtype | None = None) -> TensorLayout:
"""Derive a layout from a ``{name: tensor}`` dict.
If ``dtype`` is given, non-FP8 dtypes are coerced to it (FP8 preserved). If
``None``, the source dtype is preserved as-is.
"""
layout: BlockLayout = {}
for name, tensor in itertools.chain(block.named_parameters(), block.named_buffers()):
layout[name] = (tensor.shape, dtype)
return layout
def allocate_buffer(layout: BlockLayout, device: torch.device, pin_memory: bool = False) -> dict[str, torch.Tensor]:
"""Allocate a single buffer dict matching *layout*."""
return {
name: torch.empty(shape, dtype=dtype, device=device, pin_memory=pin_memory)
for name, (shape, dtype) in layout.items()
name: (t.shape, t.dtype if dtype is None or t.dtype in FP8_DTYPES else dtype) for name, t in tensors.items()
}
def _align_up(offset: int, alignment: int) -> int:
return (offset + alignment - 1) & ~(alignment - 1)
def _alloc_pinned_exact(nbytes: int) -> torch.Tensor | None:
"""Allocate exactly ``nbytes`` of pinned host memory via ``cudaHostRegister``.
Bypasses PyTorch's ``CachingHostAllocator``, which rounds every
``pin_memory=True`` request up to ``PowerOf2Ceil(N)`` (see
``aten/src/ATen/core/CachingHostAllocator.h``). Returns ``None`` if
registration fails. The unregister hook is bound to the storage (not the
tensor) so views of the buffer keep the registration alive until the
memory is actually freed. Caller is responsible for ensuring CUDA is
available.
"""
cudart = torch.cuda.cudart()
buf = torch.empty(nbytes, dtype=torch.uint8)
ptr = buf.data_ptr()
err = int(cudart.cudaHostRegister(ptr, nbytes, 0))
if err != 0:
return None
weakref.finalize(buf.untyped_storage(), lambda p=ptr: cudart.cudaHostUnregister(p))
return buf
def _alloc_buffer(nbytes: int, device: torch.device | None, pin_memory: bool) -> torch.Tensor:
"""Allocate one ``uint8`` buffer for :func:`allocate_layout_views`.
For pinned host buffers, prefer ``cudaHostRegister`` to dodge the caching
allocator's power-of-2 rounding. Falls back to the caching allocator if
registration fails. Raises if pinning is requested without a CUDA runtime,
since pinning is fundamentally a CUDA driver operation.
"""
if pin_memory and (device is None or torch.device(device).type == "cpu"):
if not torch.cuda.is_available():
raise RuntimeError("pin_memory=True requires CUDA, which is not available")
buf = _alloc_pinned_exact(nbytes)
if buf is not None:
return buf
return torch.empty(nbytes, dtype=torch.uint8, device=device, pin_memory=pin_memory)
@dataclass(frozen=True)
class _TensorSlice:
"""Location of a single tensor view within the buffer."""
offset: int
shape: torch.Size
dtype: torch.dtype
def size(self) -> int:
return math.prod(self.shape) * self.dtype.itemsize
def allocate_layout_views(
layout: TensorLayout,
device: torch.device | None = None,
pin_memory: bool = False,
) -> dict[str, torch.Tensor]:
"""Allocate a single ``uint8`` buffer and return per-key tensor views into it.
All keys in *layout* live in one contiguous allocation; each returned
tensor is a non-overlapping slice of that buffer reinterpreted at the
requested shape and dtype. The views keep the underlying storage alive
via PyTorch refcounting — drop them all to release the memory.
"""
slices: dict[str, _TensorSlice] = {}
cursor = 0
for key, (shape, dtype) in layout.items():
cursor = _align_up(cursor, _BUFFER_ALIGN)
slices[key] = _TensorSlice(offset=cursor, shape=shape, dtype=dtype)
cursor += slices[key].size()
# Allocate at least one byte so empty layouts still produce a valid buffer.
buffer = _alloc_buffer(max(_align_up(cursor, _BUFFER_ALIGN), 1), device, pin_memory)
return {key: buffer[s.offset : s.offset + s.size()].view(s.dtype).view(s.shape) for key, s in slices.items()}
packages/ltx-core/src/ltx_core/components/diffusion_steps.py
+80
View File
@@ -4,6 +4,24 @@ from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.utils import to_velocity
def _get_ancestral_step(
sigma_from: torch.Tensor,
sigma_to: torch.Tensor,
eta: float = 1.0,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Compute ``(sigma_down, sigma_up)`` for one DDIM ancestral sampling step.
Both inputs are in the rescaled parameterization ``sigma / alpha``.
Returns ``sigma_down`` (deterministic component) and ``sigma_up``
(stochastic component) in the same rescaled space.
"""
if not eta:
return sigma_to, torch.zeros_like(sigma_to)
variance = sigma_to**2 * (sigma_from**2 - sigma_to**2).clamp(min=0) / sigma_from**2
sigma_up = (eta * variance**0.5).clamp(max=sigma_to)
sigma_down = (sigma_to**2 - sigma_up**2).clamp(min=0) ** 0.5
return sigma_down, sigma_up
class EulerDiffusionStep(DiffusionStepProtocol):
"""
First-order Euler method for diffusion sampling.
@@ -104,3 +122,65 @@ class Res2sDiffusionStep(DiffusionStepProtocol):
# Mix deterministic and stochastic components
x_noised = alpha_ratio * (denoised_next + sigma_down * eps_next) + sigma_up * noise
return x_noised.to(output_dtype)
class EulerCfgPpDiffusionStep(DiffusionStepProtocol):
"""Euler step using the CFG++ correction for the ODE derivative.
Instead of the standard velocity formula, the ODE derivative is computed
from the unconditioned prediction, keeping the conditioned prediction as
the target denoised state. Ancestral (DDIM) noise injection is applied
in the rescaled sigma parameterization (sigma / alpha).
All diffusion quantities (alpha, ODE derivative, ancestral coefficients)
are computed internally from ``sigmas`` and ``uncond_denoised``.
Reference: CFG++ (https://arxiv.org/abs/2406.08070).
"""
def __init__(self, eta: float = 1.0, s_noise: float = 1.0) -> None:
self.eta = eta
self.s_noise = s_noise
def step(
self,
sample: torch.Tensor,
denoised_sample: torch.Tensor,
sigmas: torch.Tensor,
step_index: int,
uncond_denoised: torch.Tensor,
noise: torch.Tensor | None = None,
**_kwargs,
) -> torch.Tensor:
"""Advance one CFG++ Euler step.
Args:
sample: Current noisy latent x_t.
denoised_sample: Conditioned denoised prediction x_0^cond.
sigmas: Full sigma schedule tensor.
step_index: Current step index.
uncond_denoised: Unconditioned denoised prediction x_0^uncond,
used to compute the ODE derivative direction.
noise: Noise tensor for stochastic injection; ignored when
``eta=0`` or ``s_noise=0``.
Returns:
Updated latent x_{t-1}.
"""
sigma_s = sigmas[step_index].to(torch.float32)
sigma_t = sigmas[step_index + 1].to(torch.float32)
_eps = torch.finfo(torch.float32).eps
# Clamp to avoid division by zero when sigma == 1.0 exactly.
alpha_s = (1.0 - sigma_s).clamp(min=_eps)
alpha_t = (1.0 - sigma_t).clamp(min=_eps)
x = sample.to(torch.float32)
denoised = denoised_sample.to(torch.float32)
uncond = uncond_denoised.to(torch.float32)
# ODE derivative: direction toward noise using uncond prediction (CFG++ correction)
d = (x - alpha_s * uncond) / sigma_s
# Ancestral step in rescaled sigma space (sigma / alpha)
sigma_down, sigma_up = _get_ancestral_step(sigma_s / alpha_s, sigma_t / alpha_t, eta=self.eta)
sigma_down = alpha_t * sigma_down
x_next = alpha_t * denoised + sigma_down * d
if noise is not None and self.eta > 0 and self.s_noise > 0:
x_next = x_next + alpha_t * noise.to(torch.float32) * self.s_noise * sigma_up
return x_next.to(sample.dtype)
packages/ltx-core/src/ltx_core/components/patchifiers.py
+8 -3
View File
@@ -151,17 +151,22 @@ def get_pixel_coords(
that treat frame zero differently still yield non-negative timestamps.
"""
# Broadcast the VAE scale factors so they align with the `(batch, axis, patch, bound)` layout.
# Axis 1 of `latent_coords` is ordered (frame/time, height, width) — match that explicitly by
# pulling fields from the NamedTuple rather than relying on tuple iteration order.
broadcast_shape = [1] * latent_coords.ndim
broadcast_shape[1] = -1 # axis dimension corresponds to (frame/time, height, width)
scale_tensor = torch.tensor(scale_factors, device=latent_coords.device).view(*broadcast_shape)
scale_tensor = torch.tensor(
[scale_factors.time, scale_factors.height, scale_factors.width],
device=latent_coords.device,
).view(*broadcast_shape)
# Apply per-axis scaling to convert latent bounds into pixel-space coordinates.
pixel_coords = latent_coords * scale_tensor
if causal_fix:
# VAE temporal stride for the very first frame is 1 instead of `scale_factors[0]`.
# VAE temporal stride for the very first frame is 1 instead of `scale_factors.time`.
# Shift and clamp to keep the first-frame timestamps causal and non-negative.
pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + 1 - scale_factors[0]).clamp(min=0)
pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + 1 - scale_factors.time).clamp(min=0)
return pixel_coords
packages/ltx-core/src/ltx_core/conditioning/__init__.py
+2
View File
@@ -3,6 +3,7 @@
from ltx_core.conditioning.exceptions import ConditioningError
from ltx_core.conditioning.item import ConditioningItem
from ltx_core.conditioning.types import (
AudioConditionByReferenceLatent,
ConditioningItemAttentionStrengthWrapper,
VideoConditionByKeyframeIndex,
VideoConditionByLatentIndex,
@@ -10,6 +11,7 @@ from ltx_core.conditioning.types import (
)
__all__ = [
"AudioConditionByReferenceLatent",
"ConditioningError",
"ConditioningItem",
"ConditioningItemAttentionStrengthWrapper",
packages/ltx-core/src/ltx_core/conditioning/types/__init__.py
+2
View File
@@ -3,9 +3,11 @@
from ltx_core.conditioning.types.attention_strength_wrapper import ConditioningItemAttentionStrengthWrapper
from ltx_core.conditioning.types.keyframe_cond import VideoConditionByKeyframeIndex
from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex
from ltx_core.conditioning.types.reference_audio_cond import AudioConditionByReferenceLatent
from ltx_core.conditioning.types.reference_video_cond import VideoConditionByReferenceLatent
__all__ = [
"AudioConditionByReferenceLatent",
"ConditioningItemAttentionStrengthWrapper",
"VideoConditionByKeyframeIndex",
"VideoConditionByLatentIndex",
packages/ltx-core/src/ltx_core/conditioning/types/reference_audio_cond.py
+59
View File
@@ -0,0 +1,59 @@
"""Audio reference conditioning items."""
from __future__ import annotations
import torch
from ltx_core.conditioning.mask_utils import update_attention_mask
from ltx_core.tools import LatentTools
from ltx_core.types import LatentState
class AudioConditionByReferenceLatent:
"""Append patchified reference audio tokens after the target audio sequence.
Mirrors :class:`ltx_core.conditioning.types.reference_video_cond.VideoConditionByReferenceLatent`
but for audio. The reference tokens are appended so the target audio tokens stay
in the first ``num_noisy_tokens`` positions and can be kept by
:meth:`ltx_core.tools.LatentTools.clear_conditioning`.
Args:
patchified: Patchified reference latent ``[B, T_ref, C]``.
positions: RoPE positions for reference tokens, ``[B, 1, T_ref, 2]``.
strength: 1.0 keeps reference clean; 0.0 would fully denoise it.
"""
def __init__(
self,
patchified: torch.Tensor,
positions: torch.Tensor,
strength: float = 1.0,
) -> None:
self.patchified = patchified
self.positions = positions.to(dtype=torch.float32)
self.strength = strength
def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
tokens = self.patchified
denoise_mask = torch.full(
size=(*tokens.shape[:2], 1),
fill_value=1.0 - self.strength,
device=tokens.device,
dtype=tokens.dtype,
)
new_attention_mask = update_attention_mask(
latent_state=latent_state,
attention_mask=None,
num_noisy_tokens=latent_tools.patchifier.get_token_count(latent_tools.target_shape),
num_new_tokens=tokens.shape[1],
batch_size=tokens.shape[0],
device=tokens.device,
dtype=tokens.dtype,
)
return LatentState(
latent=torch.cat([latent_state.latent, tokens], dim=1),
denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
positions=torch.cat([latent_state.positions, self.positions], dim=2),
clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
attention_mask=new_attention_mask,
)
packages/ltx-core/src/ltx_core/loader/fuse_loras.py
+79 -50
View File
@@ -1,26 +1,81 @@
from collections.abc import Iterator
from collections.abc import Iterable, Iterator
from typing import NamedTuple
import torch
from ltx_core.loader.kernels import TRITON_AVAILABLE
from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict
from ltx_core.quantization.fp8_cast import _fused_add_round_launch
from ltx_core.quantization.fp8_cast import fused_add_round_launch
from ltx_core.quantization.fp8_scaled_mm import quantize_weight_to_fp8_per_tensor
class LoraProduct(NamedTuple):
"""A LoRA's ``A``, ``B`` factors and its strength scalar."""
a: torch.Tensor
b: torch.Tensor
strength: float
def _get_device() -> torch.device:
if torch.cuda.is_available():
return torch.device("cuda", torch.cuda.current_device())
return torch.device("cpu")
def aggregate_lora_products(
products: Iterable[LoraProduct],
dtype: torch.dtype | None = None,
*,
out: torch.Tensor | None = None,
) -> torch.Tensor | None:
"""Accumulate ``sum((B * strength) @ A)`` across :class:`LoraProduct` items.
If ``out`` is provided, ``addmm_`` accumulates directly into it — caller
ensures A/B dtypes and devices match ``out``. Otherwise the first product
materializes the ``(out, in)``-shape aggregator at ``dtype``; subsequent
products use ``addmm_`` to avoid allocating the full intermediate delta.
Returns ``out`` (or the new aggregator), or ``None`` if ``products`` was empty
and ``out`` was not given.
"""
aggregated = out
for product in products:
if aggregated is None:
aggregated = torch.matmul(product.b * product.strength, product.a).to(dtype=dtype)
else:
aggregated.addmm_(product.b, product.a, alpha=product.strength)
return aggregated
def fuse_cast_fp8_weight(
delta_bf16: torch.Tensor,
weight_fp8: torch.Tensor,
target_dtype: torch.dtype,
) -> torch.Tensor:
"""Return ``(delta_bf16 + dequantize(weight_fp8)).to(target_dtype)``.
CUDA with Triton uses stochastic rounding; otherwise uses a deterministic bf16 add.
``delta_bf16`` is the bf16 accumulator and is mutated in place.
"""
if delta_bf16.dtype != torch.bfloat16:
raise ValueError(f"delta_bf16 must be bfloat16, got {delta_bf16.dtype}")
if str(weight_fp8.device).startswith("cuda") and TRITON_AVAILABLE:
fused_add_round_launch(delta_bf16, weight_fp8, seed=0)
else:
delta_bf16.add_(weight_fp8.to(dtype=torch.bfloat16))
return delta_bf16.to(dtype=target_dtype)
def fuse_lora_weights(
model_sd: StateDict,
lora_sd_and_strengths: list[LoraStateDictWithStrength],
dtype: torch.dtype | None = None,
preserve_input_device: bool = True,
) -> Iterator[tuple[str, torch.Tensor]]:
"""Yield ``(key, fused_tensor)`` for each weight modified by at least one LoRA.
For scaled-FP8 weights, this includes both the updated ``.weight`` tensor
and its corresponding ``.weight_scale`` tensor.
When ``preserve_input_device`` is False, fused tensors are yielded on the device
used for fusion; caller is responsible for moving them to their final
destination.
"""
for key, original_weight in model_sd.sd.items():
if original_weight is None or key.endswith(".weight_scale"):
@@ -30,7 +85,7 @@ def fuse_lora_weights(
target_dtype = dtype if dtype is not None else weight.dtype
deltas_dtype = target_dtype if target_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2] else torch.bfloat16
deltas = _prepare_deltas(lora_sd_and_strengths, key, deltas_dtype, weight.device)
deltas = _aggregate_deltas(lora_sd_and_strengths, key, deltas_dtype, weight.device)
if deltas is None:
continue
@@ -41,14 +96,15 @@ def fuse_lora_weights(
if is_scaled_fp8:
fused = _fuse_delta_with_scaled_fp8(deltas, weight, key, scale_key, model_sd)
else:
fused = _fuse_delta_with_cast_fp8(deltas, weight, key, target_dtype)
fused = {key: fuse_cast_fp8_weight(deltas, weight, target_dtype)}
elif weight.dtype == torch.bfloat16:
fused = _fuse_delta_with_bfloat16(deltas, weight, key, target_dtype)
deltas.add_(weight)
fused = {key: deltas.to(dtype=target_dtype)}
else:
raise ValueError(f"Unsupported dtype: {weight.dtype}")
for k, v in fused.items():
yield k, v.to(device=original_device)
yield k, v.to(device=original_device) if preserve_input_device else v
def apply_loras(
@@ -57,10 +113,12 @@ def apply_loras(
dtype: torch.dtype | None = None,
destination_sd: StateDict | None = None,
) -> StateDict:
"""Fuse LoRAs into ``model_sd`` and place the results in ``destination_sd``.
When ``destination_sd`` is provided, the fused tensors are placed directly into it.
"""
if destination_sd is not None:
sd = destination_sd.sd
for key, tensor in fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype):
sd[key] = tensor
for key, fused in fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype):
destination_sd.sd[key] = fused
return destination_sd
fused = dict(fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype))
@@ -68,26 +126,22 @@ def apply_loras(
return StateDict(sd, model_sd.device, model_sd.size, model_sd.dtype)
def _prepare_deltas(
def _aggregate_deltas(
lora_sd_and_strengths: list[LoraStateDictWithStrength], key: str, dtype: torch.dtype, device: torch.device
) -> torch.Tensor | None:
deltas = []
prefix = key[: -len(".weight")]
key_a = f"{prefix}.lora_A.weight"
key_b = f"{prefix}.lora_B.weight"
for lsd, coef in lora_sd_and_strengths:
if key_a not in lsd.sd or key_b not in lsd.sd:
continue
a = lsd.sd[key_a].to(device=device)
b = lsd.sd[key_b].to(device=device)
product = torch.matmul(b * coef, a)
del a, b
deltas.append(product.to(dtype=dtype))
if len(deltas) == 0:
return None
elif len(deltas) == 1:
return deltas[0]
return torch.sum(torch.stack(deltas, dim=0), dim=0)
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(
@@ -100,34 +154,9 @@ def _fuse_delta_with_scaled_fp8(
"""Dequantize scaled FP8 weight, add LoRA delta, and re-quantize."""
weight_scale = model_sd.sd[scale_key]
original_weight = weight.t().to(torch.float32) * weight_scale
original_weight = weight.to(torch.float32) * weight_scale
new_weight = original_weight + deltas.to(torch.float32)
new_fp8_weight, new_weight_scale = quantize_weight_to_fp8_per_tensor(new_weight)
return {key: new_fp8_weight, scale_key: new_weight_scale}
def _fuse_delta_with_cast_fp8(
deltas: torch.Tensor,
weight: torch.Tensor,
key: str,
target_dtype: torch.dtype,
) -> dict[str, torch.Tensor]:
"""Fuse LoRA delta with cast-only FP8 weight (no scale factor)."""
if str(weight.device).startswith("cuda"):
_fused_add_round_launch(deltas, weight, seed=0)
else:
deltas.add_(weight.to(dtype=deltas.dtype))
return {key: deltas.to(dtype=target_dtype)}
def _fuse_delta_with_bfloat16(
deltas: torch.Tensor,
weight: torch.Tensor,
key: str,
target_dtype: torch.dtype,
) -> dict[str, torch.Tensor]:
"""Fuse LoRA delta with bfloat16 weight."""
deltas.add_(weight)
return {key: deltas.to(dtype=target_dtype)}
packages/ltx-core/src/ltx_core/loader/kernels.py
+65 -58
View File
@@ -1,72 +1,79 @@
# ruff: noqa: ANN001, ANN201, ERA001, N803, N806
import triton
import triton.language as tl
try:
import triton
import triton.language as tl
TRITON_AVAILABLE = True
except (ImportError, OSError):
TRITON_AVAILABLE = False
@triton.jit
def fused_add_round_kernel(
x_ptr,
output_ptr, # contents will be added to the output
seed,
n_elements,
EXPONENT_BIAS,
MANTISSA_BITS,
BLOCK_SIZE: tl.constexpr,
):
"""
A kernel to upcast 8bit quantized weights to bfloat16 with stochastic rounding
and add them to bfloat16 output weights. Might be used to upcast original model weights
and to further add them to precalculated deltas coming from LoRAs.
"""
# Get program ID and compute offsets
pid = tl.program_id(axis=0)
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
if TRITON_AVAILABLE:
# Load data
x = tl.load(x_ptr + offsets, mask=mask)
rand_vals = tl.rand(seed, offsets) - 0.5
@triton.jit
def fused_add_round_kernel(
x_ptr,
output_ptr, # contents will be added to the output
seed,
n_elements,
EXPONENT_BIAS,
MANTISSA_BITS,
BLOCK_SIZE: tl.constexpr,
):
"""
A kernel to upcast 8bit quantized weights to bfloat16 with stochastic rounding
and add them to bfloat16 output weights. Might be used to upcast original model weights
and to further add them to precalculated deltas coming from LoRAs.
"""
# Get program ID and compute offsets
pid = tl.program_id(axis=0)
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
x = tl.cast(x, tl.float16)
delta = tl.load(output_ptr + offsets, mask=mask)
delta = tl.cast(delta, tl.float16)
x = x + delta
# Load data
x = tl.load(x_ptr + offsets, mask=mask)
rand_vals = tl.rand(seed, offsets) - 0.5
x_bits = tl.cast(x, tl.int16, bitcast=True)
x = tl.cast(x, tl.float16)
delta = tl.load(output_ptr + offsets, mask=mask)
delta = tl.cast(delta, tl.float16)
x = x + delta
# Calculate the exponent. Unbiased fp16 exponent is ((x_bits & 0x7C00) >> 10) - 15 for
# normal numbers and -14 for subnormals.
fp16_exponent_bits = (x_bits & 0x7C00) >> 10
fp16_normals = fp16_exponent_bits > 0
fp16_exponent = tl.where(fp16_normals, fp16_exponent_bits - 15, -14)
x_bits = tl.cast(x, tl.int16, bitcast=True)
# Add the target dtype's exponent bias and clamp to the target dtype's exponent range.
exponent = fp16_exponent + EXPONENT_BIAS
MAX_EXPONENT = 2 * EXPONENT_BIAS + 1
exponent = tl.where(exponent > MAX_EXPONENT, MAX_EXPONENT, exponent)
exponent = tl.where(exponent < 0, 0, exponent)
# Calculate the exponent. Unbiased fp16 exponent is ((x_bits & 0x7C00) >> 10) - 15 for
# normal numbers and -14 for subnormals.
fp16_exponent_bits = (x_bits & 0x7C00) >> 10
fp16_normals = fp16_exponent_bits > 0
fp16_exponent = tl.where(fp16_normals, fp16_exponent_bits - 15, -14)
# Normal ULP exponent, expressed as an fp16 exponent field:
# (exponent - EXPONENT_BIAS - MANTISSA_BITS) + 15
# Simplifies to: fp16_exponent - MANTISSA_BITS + 15
# See https://en.wikipedia.org/wiki/Unit_in_the_last_place
eps_exp = tl.maximum(0, tl.minimum(31, exponent - EXPONENT_BIAS - MANTISSA_BITS + 15))
# Add the target dtype's exponent bias and clamp to the target dtype's exponent range.
exponent = fp16_exponent + EXPONENT_BIAS
MAX_EXPONENT = 2 * EXPONENT_BIAS + 1
exponent = tl.where(exponent > MAX_EXPONENT, MAX_EXPONENT, exponent)
exponent = tl.where(exponent < 0, 0, exponent)
# Calculate epsilon in the target dtype
eps_normal = tl.cast(tl.cast(eps_exp << 10, tl.int16), tl.float16, bitcast=True)
# Normal ULP exponent, expressed as an fp16 exponent field:
# (exponent - EXPONENT_BIAS - MANTISSA_BITS) + 15
# Simplifies to: fp16_exponent - MANTISSA_BITS + 15
# See https://en.wikipedia.org/wiki/Unit_in_the_last_place
eps_exp = tl.maximum(0, tl.minimum(31, exponent - EXPONENT_BIAS - MANTISSA_BITS + 15))
# Subnormal ULP: 2^(1 - EXPONENT_BIAS - MANTISSA_BITS) ->
# fp16 exponent bits: (1 - EXPONENT_BIAS - MANTISSA_BITS) + 15 =
# 16 - EXPONENT_BIAS - MANTISSA_BITS
eps_subnormal = tl.cast(tl.cast((16 - EXPONENT_BIAS - MANTISSA_BITS) << 10, tl.int16), tl.float16, bitcast=True)
eps = tl.where(exponent > 0, eps_normal, eps_subnormal)
# Calculate epsilon in the target dtype
eps_normal = tl.cast(tl.cast(eps_exp << 10, tl.int16), tl.float16, bitcast=True)
# Apply zero mask to epsilon
eps = tl.where(x == 0, 0.0, eps)
# Subnormal ULP: 2^(1 - EXPONENT_BIAS - MANTISSA_BITS) ->
# fp16 exponent bits: (1 - EXPONENT_BIAS - MANTISSA_BITS) + 15 =
# 16 - EXPONENT_BIAS - MANTISSA_BITS
eps_subnormal = tl.cast(tl.cast((16 - EXPONENT_BIAS - MANTISSA_BITS) << 10, tl.int16), tl.float16, bitcast=True)
eps = tl.where(exponent > 0, eps_normal, eps_subnormal)
# Apply stochastic rounding
output = tl.cast(x + rand_vals * eps, tl.bfloat16)
# Apply zero mask to epsilon
eps = tl.where(x == 0, 0.0, eps)
# Store the result
tl.store(output_ptr + offsets, output, mask=mask)
# Apply stochastic rounding
output = tl.cast(x + rand_vals * eps, tl.bfloat16)
# Store the result
tl.store(output_ptr + offsets, output, mask=mask)
packages/ltx-core/src/ltx_core/loader/primitives.py
+13 -1
View File
@@ -13,6 +13,10 @@ if TYPE_CHECKING:
from ltx_core.loader.registry import Registry
# Per-key shape and dtype description for a flat collection of tensors.
TensorLayout = dict[str, tuple[torch.Size, torch.dtype]]
@dataclass(frozen=True)
class StateDict:
"""
@@ -52,7 +56,15 @@ class StateDictLoader(Protocol):
"""
class ModelBuilderProtocol(Protocol[ModelType]):
class BuilderProtocol(Protocol[ModelType]):
"""Protocol for model builders that produce a model via ``build()``."""
def build(
self, device: torch.device | None = None, dtype: torch.dtype | None = None, **kwargs: object
) -> ModelType: ...
class ModelBuilderProtocol(BuilderProtocol[ModelType], Protocol[ModelType]):
"""
Protocol for building PyTorch models from configuration dictionaries.
Implementations must provide:
packages/ltx-core/src/ltx_core/loader/single_gpu_model_builder.py
+1 -1
View File
@@ -102,7 +102,7 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
registry: Registry = field(default_factory=DummyRegistry)
lora_load_device: torch.device = field(default_factory=lambda: torch.device("cpu"))
def lora(self, lora_path: str, strength: float = 1.0, sd_ops: SDOps | None = None) -> "SingleGPUModelBuilder":
def lora(self, lora_path: str, strength: float, sd_ops: SDOps) -> "SingleGPUModelBuilder":
return replace(self, loras=(*self.loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops)))
def with_sd_ops(self, sd_ops: SDOps | None) -> "SingleGPUModelBuilder":
packages/ltx-core/src/ltx_core/model/transformer/attention.py
+1 -1
View File
@@ -144,7 +144,7 @@ class Attention(torch.nn.Module):
heads: int = 8,
dim_head: int = 64,
norm_eps: float = 1e-6,
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
rope_type: LTXRopeType = LTXRopeType.SPLIT,
attention_function: AttentionCallable | AttentionFunction = AttentionFunction.DEFAULT,
apply_gated_attention: bool = False,
) -> None:
packages/ltx-core/src/ltx_core/model/transformer/model.py
+1 -1
View File
@@ -57,7 +57,7 @@ class LTXModel(torch.nn.Module):
audio_cross_attention_dim: int = 2048,
audio_positional_embedding_max_pos: list[int] | None = None,
av_ca_timestep_scale_multiplier: int = 1,
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
rope_type: LTXRopeType = LTXRopeType.SPLIT,
double_precision_rope: bool = False,
apply_gated_attention: bool = False,
caption_projection: torch.nn.Module | None = None,
packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py
+2 -2
View File
@@ -62,7 +62,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
audio_cross_attention_dim=config.get("audio_cross_attention_dim", 2048),
audio_positional_embedding_max_pos=config.get("audio_positional_embedding_max_pos", [20]),
av_ca_timestep_scale_multiplier=config.get("av_ca_timestep_scale_multiplier", 1),
rope_type=LTXRopeType(config.get("rope_type", "interleaved")),
rope_type=LTXRopeType(config.get("rope_type", "split")),
double_precision_rope=config.get("frequencies_precision", False) == "float64",
apply_gated_attention=config.get("apply_gated_attention", False),
caption_projection=caption_projection,
@@ -114,7 +114,7 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]),
timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
use_middle_indices_grid=config.get("use_middle_indices_grid", True),
rope_type=LTXRopeType(config.get("rope_type", "interleaved")),
rope_type=LTXRopeType(config.get("rope_type", "split")),
double_precision_rope=config.get("frequencies_precision", False) == "float64",
apply_gated_attention=config.get("apply_gated_attention", False),
caption_projection=caption_projection,
packages/ltx-core/src/ltx_core/model/transformer/rope.py
+8 -2
View File
@@ -16,9 +16,10 @@ class LTXRopeType(Enum):
def apply_rotary_emb(
input_tensor: torch.Tensor,
freqs_cis: Tuple[torch.Tensor, torch.Tensor],
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
rope_type: LTXRopeType = LTXRopeType.SPLIT,
) -> torch.Tensor:
if rope_type == LTXRopeType.INTERLEAVED:
# Note: INTERLEAVED rope is a legacy mode. Prefer SPLIT instead.
return apply_interleaved_rotary_emb(input_tensor, *freqs_cis)
elif rope_type == LTXRopeType.SPLIT:
return apply_split_rotary_emb(input_tensor, *freqs_cis)
@@ -45,6 +46,11 @@ def apply_split_rotary_emb(
needs_reshape = False
if input_tensor.ndim != 4 and cos_freqs.ndim == 4:
b, h, t, _ = cos_freqs.shape
if input_tensor.shape[0] != b:
raise ValueError(
f"apply_split_rotary_emb: input_tensor batch ({input_tensor.shape[0]}) "
f"must equal cos_freqs batch ({b})."
)
input_tensor = input_tensor.reshape(b, t, h, -1).swapaxes(1, 2)
needs_reshape = True
@@ -183,7 +189,7 @@ def precompute_freqs_cis(
max_pos: list[int] | None = None,
use_middle_indices_grid: bool = False,
num_attention_heads: int = 32,
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
rope_type: LTXRopeType = LTXRopeType.SPLIT,
freq_grid_generator: Callable[[float, int, int, torch.device], torch.Tensor] = generate_freq_grid_pytorch,
) -> tuple[torch.Tensor, torch.Tensor]:
if max_pos is None:
packages/ltx-core/src/ltx_core/model/transformer/transformer.py
+1 -1
View File
@@ -27,7 +27,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
idx: int,
video: TransformerConfig | None = None,
audio: TransformerConfig | None = None,
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
rope_type: LTXRopeType = LTXRopeType.SPLIT,
norm_eps: float = 1e-6,
attention_function: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT,
):
packages/ltx-core/src/ltx_core/model/video_vae/__init__.py
+2
View File
@@ -1,5 +1,6 @@
"""Video VAE package."""
from ltx_core.model.video_vae.memory_efficient_decode import MEMORY_EFFICIENT_DECODE
from ltx_core.model.video_vae.model_configurator import (
VAE_DECODER_COMFY_KEYS_FILTER,
VAE_ENCODER_COMFY_KEYS_FILTER,
@@ -10,6 +11,7 @@ from ltx_core.model.video_vae.tiling import SpatialTilingConfig, TemporalTilingC
from ltx_core.model.video_vae.video_vae import VideoDecoder, VideoEncoder, get_video_chunks_number
__all__ = [
"MEMORY_EFFICIENT_DECODE",
"VAE_DECODER_COMFY_KEYS_FILTER",
"VAE_ENCODER_COMFY_KEYS_FILTER",
"SpatialTilingConfig",
packages/ltx-core/src/ltx_core/model/video_vae/memory_efficient_decode.py
+612
View File
@@ -0,0 +1,612 @@
"""Memory-efficient VAE decoder operations.
Reduces peak VRAM usage during video decoding through in-place operations
and workspace buffer reuse. The main optimizations are:
1. **Workspace buffers** -- Pre-allocated tensors with temporal padding replace
dynamic padding (``F.pad`` / ``concatenate``) in ``CausalConv3d``. A
workspace of shape ``[B, C, T+2, H, W]`` holds the data in positions
``[1:-1]`` with replicate padding at ``[0]`` and ``[-1]``.
2. **In-place temporal-chunked Conv3d** *(non-causal only)* -- The convolution
output is written back into the workspace buffer, avoiding a separate
output allocation. Temporal chunking with boundary save/restore ensures
correct reads despite in-place writes.
3. **In-place normalization and affine transforms** -- PixelNorm, scale/shift,
and SiLU are applied in-place on workspace views.
4. **Free-before-conv** -- For ``DepthToSpaceUpsample`` blocks the input
tensor is freed before the convolution runs so that peak VRAM never holds
input *and* output simultaneously.
Both causal and non-causal modes are supported. Non-causal mode benefits
from all four optimizations. Causal mode benefits from optimizations 1, 3,
and 4; in-place conv (2) is skipped because the asymmetric causal padding
layout prevents clean in-place overwrites.
Usage via the ``ModuleOps`` pattern (preferred)::
from ltx_core.model.video_vae import MEMORY_EFFICIENT_DECODE
builder = decoder_builder.with_module_ops(
(*decoder_builder.module_ops, MEMORY_EFFICIENT_DECODE)
)
Or applied directly to an existing decoder::
from ltx_core.model.video_vae.memory_efficient_decode import (
enable_memory_efficient_decode,
)
enable_memory_efficient_decode(decoder)
"""
from __future__ import annotations
import math
from typing import TYPE_CHECKING
import torch
from einops import rearrange
from torch import nn
from torch.nn import functional as F
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.model.common.normalization import PixelNorm
from ltx_core.model.video_vae.convolution import CausalConv3d
from ltx_core.model.video_vae.ops import unpatchify
from ltx_core.model.video_vae.resnet import ResnetBlock3D, UNetMidBlock3D
from ltx_core.model.video_vae.sampling import DepthToSpaceUpsample
if TYPE_CHECKING:
from ltx_core.model.video_vae.video_vae import VideoDecoder
# ---------------------------------------------------------------------------
# Low-level helpers
# ---------------------------------------------------------------------------
def _find_temporal_split_size(num_frames: int) -> int:
"""Find chunk size for in-place temporal convolution.
The chunk size ensures the last chunk has at least 3 frames
(the temporal kernel size), avoiding degenerate chunks.
"""
for s in range(16, 2, -1):
remainder = num_frames % s
if remainder == 0 or remainder >= 3:
return s
raise ValueError(
f"Unable to find a valid temporal split size for num_frames={num_frames}. "
"Expected a split size between 3 and 16 such that the final chunk is "
"either exact or has at least 3 frames."
)
def _pad_workspace_temporal(workspace: torch.Tensor) -> None:
"""Apply non-causal replicate padding to temporal boundaries.
Sets ``workspace[:, :, 0]`` to a copy of ``workspace[:, :, 1]`` and
``workspace[:, :, -1]`` to a copy of ``workspace[:, :, -2]``.
"""
workspace[:, :, 0, :, :].copy_(workspace[:, :, 1, :, :])
workspace[:, :, -1, :, :].copy_(workspace[:, :, -2, :, :])
# ---------------------------------------------------------------------------
# In-place Conv3d (non-causal only)
# ---------------------------------------------------------------------------
def inplace_conv3d_temporal_chunked(workspace: torch.Tensor, conv: nn.Conv3d) -> None:
"""Run a 3x3x3 Conv3d in-place on a temporally-padded workspace.
The workspace has shape ``[B, C, T+2, H, W]`` where positions ``[1:-1]``
hold the real data and positions ``[0]`` and ``[-1]`` are padding slots.
The convolution must have ``kernel_size=(3,3,3)``, ``stride=(1,1,1)``,
``padding=(0,1,1)`` -- no temporal padding, symmetric spatial padding.
The output (T frames) overwrites positions ``[1:-1]``. Temporal chunking
with boundary save/restore ensures each chunk reads unmodified input even
though earlier chunks already wrote to the same buffer.
Only valid for **non-causal** mode (symmetric replicate padding).
Args:
workspace: Tensor ``[B, max(C_in, C_out), T+2, H, W]``.
Modified in-place; after the call ``workspace[:, :C_out, 1:-1]``
holds the convolution result.
conv: ``nn.Conv3d`` with the constraints above.
"""
if conv.kernel_size != (3, 3, 3):
raise ValueError(f"Expected kernel_size=(3,3,3), got {conv.kernel_size}")
if conv.stride != (1, 1, 1):
raise ValueError(f"Expected stride=(1,1,1), got {conv.stride}")
if conv.padding != (0, 1, 1):
raise ValueError(f"Expected padding=(0,1,1), got {conv.padding}")
_pad_workspace_temporal(workspace)
total_frames = workspace.shape[2]
out_channels = conv.out_channels
in_channels = conv.in_channels
if total_frames > 16:
split_size = _find_temporal_split_size(total_frames)
num_splits = (total_frames + split_size - 1) // split_size
else:
split_size = total_frames - 1
num_splits = 1
# 1-frame buffers for saving / restoring boundary frames across chunks.
x_buf = torch.empty(
workspace.shape[0],
workspace.shape[1],
1,
workspace.shape[3],
workspace.shape[4],
device=workspace.device,
dtype=workspace.dtype,
)
o_buf = torch.empty_like(x_buf)
# Helper: extract a chunk and make it contiguous. Workspace views can
# inherit strides > 2^31 from the full buffer, which makes Conv3d's
# reflect-padding path (F.pad) crash with "input tensor must fit into
# 32-bit index math". A small .clone() per chunk avoids this.
needs_clone = workspace.untyped_storage().nbytes() > (2**31 - 1) * workspace.element_size()
def _chunk(t_start: int, t_end: int) -> torch.Tensor:
s = workspace[:, :in_channels, t_start:t_end]
return s.clone() if needs_clone else s
# --- First chunk ---
if num_splits > 1:
# Save the boundary now so the loop below can restore it. Skipped
# when there is only one chunk: the loop never runs, and the save
# would be a wasted full HW slice copy.
x_buf[:, :, 0] = workspace[:, :, split_size - 1].clone()
workspace[:, :out_channels, 1:split_size] = conv(_chunk(0, split_size + 1))
# --- Remaining chunks ---
for i in range(1, num_splits):
start = i * split_size
end = min((i + 1) * split_size, total_frames - 1)
# Save the value at start-1 (now holds previous chunk's output).
o_buf[:, :, 0] = workspace[:, :, start - 1].clone()
# Restore the original input value needed by this chunk's conv.
workspace[:, :, start - 1] = x_buf[:, :, 0]
# Save the boundary for the *next* chunk before we overwrite it.
x_buf[:, :, 0] = workspace[:, :, end - 1].clone()
workspace[:, :out_channels, start:end] = conv(_chunk(start - 1, end + 1))
# Put back the previous chunk's output at the boundary.
workspace[:, :, start - 1] = o_buf[:, :, 0]
# ---------------------------------------------------------------------------
# Causal conv helper (free-before-conv)
# ---------------------------------------------------------------------------
def _causal_pad(x: torch.Tensor, pad_size: int) -> torch.Tensor:
"""Build a causal-padded buffer of shape ``[B, C, T+pad_size, H, W]``.
Copies ``x`` into ``padded[:, :, pad_size:]`` and replicates the first
real frame into the leading ``pad_size`` slots. The caller still owns
``x`` after this returns.
"""
padded = torch.empty(
x.shape[0],
x.shape[1],
x.shape[2] + pad_size,
x.shape[3],
x.shape[4],
device=x.device,
dtype=x.dtype,
)
padded[:, :, pad_size:].copy_(x)
for i in range(pad_size):
padded[:, :, i] = padded[:, :, pad_size]
return padded
def _causal_pad_free_and_conv(x: torch.Tensor, causal_conv: CausalConv3d) -> torch.Tensor:
"""Causal-pad *x*, free it, then run the raw ``nn.Conv3d``.
This avoids the peak where both the original and padded tensors are
live simultaneously (as happens inside ``CausalConv3d.forward``).
Args:
x: Input ``[B, C_in, T, H, W]``. **Deleted** inside this function;
the caller must not use it afterwards.
Returns:
Convolution output ``[B, C_out, T, H, W]``.
"""
padded = _causal_pad(x, causal_conv.time_kernel_size - 1)
del x
result = causal_conv.conv(padded)
del padded
return result
# ---------------------------------------------------------------------------
# In-place normalization
# ---------------------------------------------------------------------------
def _pixel_norm_inplace(x: torch.Tensor, eps: float = 1e-8) -> None:
"""In-place RMS (pixel) normalization along the channel dimension."""
rms = torch.sqrt(torch.mean(x**2, dim=1, keepdim=True) + eps)
x.div_(rms)
def _norm_inplace(norm: nn.Module, x: torch.Tensor) -> None:
"""Apply *norm* in-place, using an optimised path for ``PixelNorm``."""
if isinstance(norm, PixelNorm):
_pixel_norm_inplace(x, eps=norm.eps)
else:
# GroupNorm or other -- fall back to allocating a temporary.
result = norm(x)
x.copy_(result)
del result
# ---------------------------------------------------------------------------
# Per-block efficient forwards
# ---------------------------------------------------------------------------
def _resnet_block_forward_inplace(
resnet: ResnetBlock3D,
workspace: torch.Tensor,
causal: bool,
timestep: torch.Tensor | None,
generator: torch.Generator | None,
) -> None:
"""Run a ``ResnetBlock3D`` in-place on a workspace buffer.
The workspace has shape ``[B, C, T+2, H, W]`` with real data in
``[1:-1]``. After this call ``workspace[:, :, 1:-1]`` holds the
residual-branch output ``F(x)`` (without the skip connection --
the caller adds it back to the hidden state).
Only valid when ``in_channels == out_channels`` (true for all
``ResnetBlock3D`` instances inside a ``UNetMidBlock3D``).
"""
if resnet.in_channels != resnet.out_channels:
raise ValueError(
"In-place resnet forward requires in_channels == out_channels, "
f"got {resnet.in_channels} != {resnet.out_channels}"
)
interior = workspace[:, :, 1:-1]
# --- norm1 + [ada scaling] + SiLU + conv1 ---
_norm_inplace(resnet.norm1, interior)
if resnet.timestep_conditioning and timestep is not None:
ada = resnet.scale_shift_table[None, ..., None, None, None].to(
device=interior.device, dtype=interior.dtype
) + timestep.reshape(
interior.shape[0],
4,
-1,
timestep.shape[-3],
timestep.shape[-2],
timestep.shape[-1],
)
shift1, scale1, shift2, scale2 = ada.unbind(dim=1)
interior.mul_(1 + scale1).add_(shift1)
F.silu(interior, inplace=True)
if causal:
result = resnet.conv1(interior, causal=True)
interior.copy_(result)
del result
else:
inplace_conv3d_temporal_chunked(workspace, resnet.conv1.conv)
if resnet.inject_noise:
spatial_shape = interior.shape[-2:]
scale = resnet.per_channel_scale1.to(device=interior.device, dtype=interior.dtype)
noise = torch.randn(spatial_shape, device=interior.device, dtype=interior.dtype, generator=generator)
interior.add_((noise * scale)[None, :, None, ...])
# --- norm2 + [ada scaling] + SiLU + conv2 ---
_norm_inplace(resnet.norm2, interior)
if resnet.timestep_conditioning and timestep is not None:
interior.mul_(1 + scale2).add_(shift2) # type: ignore[possibly-undefined]
F.silu(interior, inplace=True)
# dropout is always 0.0 during inference -- skip.
if causal:
result = resnet.conv2(interior, causal=True)
interior.copy_(result)
del result
else:
inplace_conv3d_temporal_chunked(workspace, resnet.conv2.conv)
if resnet.inject_noise:
spatial_shape = interior.shape[-2:]
scale = resnet.per_channel_scale2.to(device=interior.device, dtype=interior.dtype)
noise = torch.randn(spatial_shape, device=interior.device, dtype=interior.dtype, generator=generator)
interior.add_((noise * scale)[None, :, None, ...])
def _midblock_forward_efficient(
block: UNetMidBlock3D,
hidden_states: torch.Tensor,
causal: bool,
timestep: torch.Tensor | None,
generator: torch.Generator | None,
) -> torch.Tensor:
"""Memory-efficient ``UNetMidBlock3D`` forward.
Allocates a single workspace buffer that is reused across all
``ResnetBlock3D`` iterations. For each block the workspace is
populated with the current hidden state, processed in-place, and
the result is added back (residual connection).
"""
timestep_embed = None
if block.timestep_conditioning:
if timestep is None:
raise ValueError("'timestep' required when timestep_conditioning=True")
batch_size = hidden_states.shape[0]
timestep_embed = block.time_embedder(
timestep=timestep.flatten(),
hidden_dtype=hidden_states.dtype,
)
timestep_embed = timestep_embed.view(batch_size, timestep_embed.shape[-1], 1, 1, 1)
workspace = torch.empty(
hidden_states.shape[0],
hidden_states.shape[1],
hidden_states.shape[2] + 2,
hidden_states.shape[3],
hidden_states.shape[4],
device=hidden_states.device,
dtype=hidden_states.dtype,
)
for resnet in block.res_blocks:
workspace[:, :, 1:-1].copy_(hidden_states)
_resnet_block_forward_inplace(resnet, workspace, causal, timestep_embed, generator)
hidden_states.add_(workspace[:, :, 1:-1])
del workspace
return hidden_states
def _upsample_forward_efficient(
block: DepthToSpaceUpsample,
x: torch.Tensor,
causal: bool,
) -> torch.Tensor:
"""Memory-efficient ``DepthToSpaceUpsample`` forward.
For non-causal mode the input is copied into a workspace and the
convolution runs in-place. For causal mode the input is manually
padded and freed before the convolution runs. Both paths avoid
the peak where input *and* output coexist.
"""
if block.residual:
x_in = rearrange(
x,
"b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
p1=block.stride[0],
p2=block.stride[1],
p3=block.stride[2],
)
num_repeat = math.prod(block.stride) // block.out_channels_reduction_factor
x_in = x_in.repeat(1, num_repeat, 1, 1, 1)
if block.stride[0] == 2:
x_in = x_in[:, :, 1:, :, :]
conv = block.conv.conv # underlying nn.Conv3d inside CausalConv3d
in_channels = x.shape[1]
out_channels = conv.out_channels
if causal:
x = _causal_pad_free_and_conv(x, block.conv)
else:
workspace = torch.empty(
x.shape[0],
max(in_channels, out_channels),
x.shape[2] + 2,
x.shape[3],
x.shape[4],
device=x.device,
dtype=x.dtype,
)
workspace[:, :in_channels, 1:-1].copy_(x)
del x
inplace_conv3d_temporal_chunked(workspace, conv)
x = workspace[:, :out_channels, 1:-1].contiguous()
del workspace
x = rearrange(
x,
"b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
p1=block.stride[0],
p2=block.stride[1],
p3=block.stride[2],
)
if block.stride[0] == 2:
x = x[:, :, 1:, :, :]
if block.residual:
x = x + x_in
del x_in
return x
# ---------------------------------------------------------------------------
# Final norm + conv_out
# ---------------------------------------------------------------------------
def _final_norm_and_conv_out(
decoder: VideoDecoder,
sample: torch.Tensor,
causal: bool,
scaled_timestep: torch.Tensor | None,
batch_size: int,
) -> torch.Tensor:
"""Workspace-based final norm + [ada] + SiLU + conv_out + unpatchify."""
conv_out_mod: CausalConv3d = decoder.conv_out # type: ignore[assignment]
conv_out = conv_out_mod.conv
feature_channels = sample.shape[1]
workspace = torch.empty(
sample.shape[0],
max(feature_channels, conv_out.out_channels),
sample.shape[2] + 2,
sample.shape[3],
sample.shape[4],
device=sample.device,
dtype=sample.dtype,
)
workspace[:, :feature_channels, 1:-1].copy_(sample)
del sample
interior = workspace[:, :feature_channels, 1:-1]
_norm_inplace(decoder.conv_norm_out, interior)
if decoder.timestep_conditioning:
embedded_timestep = decoder.last_time_embedder(
timestep=scaled_timestep.flatten(),
hidden_dtype=interior.dtype,
)
embedded_timestep = embedded_timestep.view(batch_size, embedded_timestep.shape[-1], 1, 1, 1)
ada_values = decoder.last_scale_shift_table[None, ..., None, None, None].to(
device=interior.device, dtype=interior.dtype
) + embedded_timestep.reshape(
batch_size,
2,
-1,
embedded_timestep.shape[-3],
embedded_timestep.shape[-2],
embedded_timestep.shape[-1],
)
shift, scale = ada_values.unbind(dim=1)
interior.mul_(1 + scale).add_(shift)
F.silu(interior, inplace=True)
if causal:
# Causal: build padded tensor directly from the interior view,
# then free the workspace before running the conv.
padded = _causal_pad(interior, conv_out_mod.time_kernel_size - 1)
del workspace, interior
result = conv_out(padded)
del padded
else:
inplace_conv3d_temporal_chunked(workspace, conv_out)
result = workspace[:, : conv_out.out_channels, 1:-1].contiguous()
del workspace, interior
return unpatchify(result, patch_size_hw=decoder.patch_size, patch_size_t=1)
# ---------------------------------------------------------------------------
# Top-level efficient decoder forward
# ---------------------------------------------------------------------------
def _memory_efficient_forward(
decoder: VideoDecoder,
sample: torch.Tensor,
timestep: torch.Tensor | None = None,
generator: torch.Generator | None = None,
) -> torch.Tensor:
"""Full memory-efficient ``VideoDecoder.forward`` replacement.
Orchestrates the entire decode through workspace-based operations:
``UNetMidBlock3D`` and ``DepthToSpaceUpsample`` blocks use efficient
paths; standalone ``ResnetBlock3D`` blocks fall back to the standard
forward. The final norm + ada + SiLU + conv_out is also workspace-based.
"""
causal = decoder.causal
batch_size = sample.shape[0]
sample = sample.to(next(decoder.parameters()).dtype)
# --- Noise injection and de-normalisation (identical to standard path) ---
if decoder.timestep_conditioning:
noise = (
torch.randn(sample.size(), generator=generator, dtype=sample.dtype, device=sample.device)
* decoder.decode_noise_scale
)
sample = noise + (1.0 - decoder.decode_noise_scale) * sample
sample = decoder.per_channel_statistics.un_normalize(sample)
if timestep is None and decoder.timestep_conditioning:
timestep = torch.full((batch_size,), decoder.decode_timestep, device=sample.device, dtype=sample.dtype)
# --- conv_in (latent tensor is small -- standard path is fine) ---
sample = decoder.conv_in(sample, causal=causal)
upscale_dtype = next(iter(decoder.up_blocks.parameters())).dtype
sample = sample.to(upscale_dtype)
scaled_timestep = None
if decoder.timestep_conditioning:
if timestep is None:
raise ValueError("'timestep' required when timestep_conditioning=True")
scaled_timestep = timestep * decoder.timestep_scale_multiplier.to(sample)
# --- Up blocks (dispatch to efficient path per block type) ---
for up_block in decoder.up_blocks:
if isinstance(up_block, UNetMidBlock3D):
sample = _midblock_forward_efficient(
up_block,
sample,
causal=causal,
timestep=scaled_timestep if decoder.timestep_conditioning else None,
generator=generator,
)
elif isinstance(up_block, DepthToSpaceUpsample):
sample = _upsample_forward_efficient(up_block, sample, causal=causal)
elif isinstance(up_block, ResnetBlock3D):
sample = up_block(sample, causal=causal, generator=generator)
else:
sample = up_block(sample, causal=causal)
return _final_norm_and_conv_out(decoder, sample, causal, scaled_timestep, batch_size)
# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------
def enable_memory_efficient_decode(decoder: nn.Module) -> nn.Module:
"""Patch a ``VideoDecoder`` to use the memory-efficient forward path.
The original ``forward`` is saved as ``decoder._original_forward`` so
that it can be restored later with :func:`disable_memory_efficient_decode`.
"""
# Import here to avoid circular dependency at module level.
from ltx_core.model.video_vae.video_vae import VideoDecoder # noqa: PLC0415
if not isinstance(decoder, VideoDecoder):
raise TypeError(f"Expected VideoDecoder, got {type(decoder).__name__}")
if hasattr(decoder, "_original_forward"):
return decoder
original_forward = decoder.forward
def efficient_forward(
sample: torch.Tensor,
timestep: torch.Tensor | None = None,
generator: torch.Generator | None = None,
) -> torch.Tensor:
return _memory_efficient_forward(decoder, sample, timestep, generator)
decoder._original_forward = original_forward # type: ignore[attr-defined]
decoder.forward = efficient_forward # type: ignore[assignment]
return decoder
def disable_memory_efficient_decode(decoder: nn.Module) -> nn.Module:
"""Restore the original ``forward`` method on a patched ``VideoDecoder``."""
if hasattr(decoder, "_original_forward"):
decoder.forward = decoder._original_forward # type: ignore[attr-defined]
del decoder._original_forward # type: ignore[attr-defined]
return decoder
def _is_video_decoder(model: nn.Module) -> bool:
"""Matcher for the ``MEMORY_EFFICIENT_DECODE`` module op."""
from ltx_core.model.video_vae.video_vae import VideoDecoder # noqa: PLC0415
return isinstance(model, VideoDecoder)
MEMORY_EFFICIENT_DECODE = ModuleOps(
name="memory_efficient_vae_decode",
matcher=_is_video_decoder,
mutator=enable_memory_efficient_decode,
)
packages/ltx-core/src/ltx_core/model/video_vae/tiling.py
+2 -2
View File
@@ -64,6 +64,6 @@ class TilingConfig:
@classmethod
def default(cls) -> "TilingConfig":
return cls(
spatial_config=SpatialTilingConfig(tile_size_in_pixels=512, tile_overlap_in_pixels=64),
temporal_config=TemporalTilingConfig(tile_size_in_frames=64, tile_overlap_in_frames=24),
spatial_config=SpatialTilingConfig(tile_size_in_pixels=768, tile_overlap_in_pixels=64),
temporal_config=TemporalTilingConfig(tile_size_in_frames=80, tile_overlap_in_frames=24),
)
packages/ltx-core/src/ltx_core/model/video_vae/video_vae.py
+7 -18
View File
@@ -259,6 +259,7 @@ class VideoEncoder(nn.Module):
Args:
sample: Input video (B, C, F, H, W). F should be 1 + 8*k (e.g., 1, 9, 17, 25, 33...).
If not, the encoder crops the last frames to the nearest valid length.
Should be normalized to [-1, 1] range before encoding.
Returns:
Normalized latent means (B, 128, F', H', W') where F' = 1+(F-1)/8, H' = H/32, W' = W/32.
Example: (B, 3, 33, 512, 512) -> (B, 128, 5, 16, 16).
@@ -605,8 +606,8 @@ class VideoDecoder(nn.Module):
# many video frames and pixels correspond to a single latent cell.
self.video_downscale_factors = SpatioTemporalScaleFactors(
time=8,
width=32,
height=32,
width=32,
)
self.patch_size = patch_size
@@ -905,34 +906,22 @@ class VideoDecoder(nn.Module):
latent: torch.Tensor,
tiling_config: TilingConfig | None = None,
generator: torch.Generator | None = None,
*,
output_dtype: torch.dtype = torch.uint8,
) -> Iterator[torch.Tensor]:
"""Decode a video latent tensor, yielding chunks ``[f, h, w, c]``.
"""Decode a video latent tensor, yielding float chunks ``[f, h, w, c]`` in ``[0, 1]``.
Subclasses (e.g. ``DistributedVideoDecoder``) may override this to
control eagerness or distribution across ranks.
Args:
output_dtype: Target dtype for output tensors. ``torch.uint8``
(default) maps the decoder's ``[-1, 1]`` output to
``[0, 255]``. Any floating dtype returns ``[0, 1]`` cast
to that dtype.
"""
def _convert(frames: torch.Tensor) -> torch.Tensor:
# rearrange materializes a new contiguous tensor for this permutation,
# so in-place ops below do not mutate the caller's data.
def to_rgb(frames: torch.Tensor) -> torch.Tensor:
video = rearrange(frames[0], "c f h w -> f h w c")
video.add_(1.0).mul_(0.5).clamp_(0.0, 1.0)
if output_dtype == torch.uint8:
return video.mul_(255.0).to(torch.uint8)
return video.to(output_dtype)
return video.add_(1.0).mul_(0.5).clamp_(0.0, 1.0)
if tiling_config is not None:
for frames in self.tiled_decode(latent, tiling_config, generator=generator):
yield _convert(frames)
yield to_rgb(frames)
else:
decoded = self(latent, generator=generator)
yield _convert(decoded)
yield to_rgb(decoded)
def _group_tiles_by_temporal_slice(self, tiles: List[Tile]) -> List[List[Tile]]:
"""Group tiles by their temporal output slice."""
packages/ltx-core/src/ltx_core/quantization/__init__.py
-3
View File
@@ -3,12 +3,9 @@ from ltx_core.quantization.fp8_cast import (
UPCAST_DURING_INFERENCE,
UpcastWithStochasticRounding,
)
from ltx_core.quantization.fp8_scaled_mm import FP8_PREPARE_MODULE_OPS, FP8_TRANSPOSE_SD_OPS
from ltx_core.quantization.policy import QuantizationPolicy
__all__ = [
"FP8_PREPARE_MODULE_OPS",
"FP8_TRANSPOSE_SD_OPS",
"TRANSFORMER_LINEAR_DOWNCAST_MAP",
"UPCAST_DURING_INFERENCE",
"QuantizationPolicy",
packages/ltx-core/src/ltx_core/quantization/fp8_cast.py
+29 -8
View File
@@ -1,5 +1,6 @@
import torch
from ltx_core.loader.kernels import TRITON_AVAILABLE
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
from ltx_core.model.transformer.model import LTXModel
@@ -7,8 +8,13 @@ from ltx_core.model.transformer.model import LTXModel
BLOCK_SIZE = 1024
def _fused_add_round_launch(target_weight: torch.Tensor, original_weight: torch.Tensor, seed: int) -> torch.Tensor:
# Lazy import triton - only available on CUDA platforms
def fused_add_round_launch(target_weight: torch.Tensor, original_weight: torch.Tensor, seed: int) -> torch.Tensor:
if not TRITON_AVAILABLE:
raise RuntimeError(
"fused_add_round_launch requires Triton, which is not available on this platform. "
"Callers should gate on ltx_core.loader.kernels.TRITON_AVAILABLE and use a "
"deterministic-rounding fallback instead."
)
import triton # noqa: PLC0415
from ltx_core.loader.kernels import fused_add_round_kernel # noqa: PLC0415
@@ -53,10 +59,13 @@ def _upcast_and_round(
"""
Upcast the weight to the given dtype and optionally apply stochastic rounding.
Input weight needs to have float8_e4m3fn or float8_e5m2 dtype.
Stochastic rounding is implemented via a Triton kernel. When Triton is not
available (e.g., on Windows), this falls back to deterministic (nearest)
rounding via ``weight.to(dtype)``.
"""
if not with_stochastic_rounding:
if not with_stochastic_rounding or not TRITON_AVAILABLE or weight.device.type != "cuda":
return weight.to(dtype)
return _fused_add_round_launch(torch.zeros_like(weight, dtype=dtype), weight, seed)
return fused_add_round_launch(torch.zeros_like(weight, dtype=dtype), weight, seed)
class Fp8CastLinear(torch.nn.Linear):
@@ -82,14 +91,26 @@ class Fp8CastLinear(torch.nn.Linear):
def _replace_fwd_with_upcast(layer: torch.nn.Linear, with_stochastic_rounding: bool = False, seed: int = 0) -> None:
"""
Intended to be applied via __class__ reassignment to existing nn.Linear
instances so that their parameter and buffer tensors are preserved in-place,
avoiding re-instantiation. Forward remains defined at the class level, which
is required for torch.compile compatibility — instance-level closure
monkey-patches cause graph breaks.
instances. Forward remains defined at the class level, which is required for
torch.compile compatibility — instance-level closure monkey-patches cause
graph breaks.
Also retypes ``weight`` and ``bias`` to fp8 so the meta param dtype matches
the post-load tensor dtype (sd_ops downcasts checkpoint bf16 -> fp8 at load).
Block streaming relies on this to derive pool buffer layout from the meta
model without an eager checkpoint read.
"""
layer.__class__ = Fp8CastLinear
layer._with_stochastic_rounding = with_stochastic_rounding
layer._seed = seed
layer.weight = torch.nn.Parameter(
torch.empty(layer.weight.shape, dtype=torch.float8_e4m3fn, device=layer.weight.device),
requires_grad=layer.weight.requires_grad,
)
if layer.bias is not None:
layer.bias = torch.nn.Parameter(
torch.empty(layer.bias.shape, dtype=torch.float8_e4m3fn, device=layer.bias.device),
requires_grad=layer.bias.requires_grad,
)
def _amend_forward_with_upcast(
packages/ltx-core/src/ltx_core/quantization/fp8_scaled_mm.py
+73 -105
View File
@@ -1,11 +1,21 @@
import json
import struct
from typing import Callable
import torch
from torch import nn
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
from ltx_core.model.transformer import LTXModel
from ltx_core.quantization.trtllm_scaled_usable import trtllm_scaled_mm_usable
def _read_safetensors_dtypes(path: str) -> dict[str, str]:
"""Return ``{tensor_name: dtype_string}`` from the safetensors header."""
with open(path, "rb") as f:
header_size = struct.unpack("<Q", f.read(8))[0]
header = json.loads(f.read(header_size).decode("utf-8"))
return {k: v["dtype"] for k, v in header.items() if k != "__metadata__"}
class FP8Linear(nn.Module):
@@ -25,11 +35,8 @@ class FP8Linear(nn.Module):
self.in_features = in_features
self.out_features = out_features
fp8_shape = (in_features, out_features)
self.weight = nn.Parameter(torch.empty(fp8_shape, dtype=torch.float8_e4m3fn, device=device))
# Weight scale for FP8 dequantization (shape matches checkpoint format)
self.weight = nn.Parameter(torch.empty((out_features, in_features), dtype=torch.float8_e4m3fn, device=device))
self.weight_scale = nn.Parameter(torch.empty((), dtype=torch.float32, device=device))
# Input scale for static quantization (pre-quantized checkpoints)
self.input_scale = nn.Parameter(torch.empty((), dtype=torch.float32, device=device))
if bias:
@@ -40,31 +47,38 @@ class FP8Linear(nn.Module):
def forward(self, x: torch.Tensor) -> torch.Tensor:
origin_shape = x.shape
# Static quantization: use pre-computed scale
qinput, cur_input_scale = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x, self.input_scale)
if trtllm_scaled_mm_usable():
qinput, cur_input_scale = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x, self.input_scale)
if qinput.dim() == 3:
qinput = qinput.reshape(-1, qinput.shape[-1])
output = torch.ops.trtllm.cublas_scaled_mm(
qinput,
self.weight.t(),
scale_a=cur_input_scale,
scale_b=self.weight_scale,
bias=None,
out_dtype=x.dtype,
)
else:
# Clamp before cast: out-of-range values cast to NaN/saturated FP8, which
# produces black-screen output on some checkpoints (e.g. ltx-2-19b-dev-fp8).
fp8_min = torch.finfo(torch.float8_e4m3fn).min
fp8_max = torch.finfo(torch.float8_e4m3fn).max
qinput = torch.clamp(x * self.input_scale.reciprocal(), fp8_min, fp8_max).to(torch.float8_e4m3fn)
if qinput.dim() == 3:
qinput = qinput.reshape(-1, qinput.shape[-1])
output = torch._scaled_mm(
qinput,
self.weight.t(),
scale_a=self.input_scale,
scale_b=self.weight_scale,
out_dtype=x.dtype,
use_fast_accum=True,
)
# Flatten to 2D for matmul
if qinput.dim() == 3:
qinput = qinput.reshape(-1, qinput.shape[-1])
# FP8 scaled matmul
output = torch.ops.trtllm.cublas_scaled_mm(
qinput,
self.weight,
scale_a=cur_input_scale,
scale_b=self.weight_scale,
bias=None,
out_dtype=x.dtype,
)
# Add bias
if self.bias is not None:
bias = self.bias
if bias.dtype != output.dtype:
bias = bias.to(output.dtype)
output = output + bias
output = output + self.bias.to(output.dtype)
# Restore original shape
if output.dim() != len(origin_shape):
output_shape = list(origin_shape)
output_shape[-1] = output.shape[-1]
@@ -74,15 +88,7 @@ class FP8Linear(nn.Module):
def quantize_weight_to_fp8_per_tensor(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""
Quantize a weight tensor to FP8 (float8_e4m3fn) using per-tensor scaling.
Args:
weight: The weight tensor to quantize (any dtype, will be cast to float32)
Returns:
Tuple of (quantized_weight, weight_scale):
- quantized_weight: FP8 tensor, transposed for cublas_scaled_mm
- weight_scale: Per-tensor scale factor (reciprocal of quantization scale)
"""
"""Quantize a weight tensor to ``float8_e4m3fn`` with a per-tensor scale."""
weight_fp32 = weight.to(torch.float32)
fp8_min = torch.finfo(torch.float8_e4m3fn).min
@@ -96,7 +102,6 @@ def quantize_weight_to_fp8_per_tensor(weight: torch.Tensor) -> tuple[torch.Tenso
weight_fp32: torch.Tensor, scale: torch.Tensor, fp8_min: torch.Tensor, fp8_max: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
quantized_weight = torch.clamp(weight_fp32 * scale, min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
quantized_weight = quantized_weight.t()
weight_scale = scale.reciprocal()
return quantized_weight, weight_scale
@@ -104,36 +109,8 @@ def quantize_weight_to_fp8_per_tensor(weight: torch.Tensor) -> tuple[torch.Tenso
return quantized_weight, weight_scale
def _should_skip_layer(layer_name: str, excluded_layer_substrings: tuple[str, ...]) -> bool:
return any(substring in layer_name for substring in excluded_layer_substrings)
EXCLUDED_LAYER_SUBSTRINGS = (
"patchify_proj",
"adaln_single",
"av_ca_video_scale_shift_adaln_single",
"av_ca_a2v_gate_adaln_single",
"caption_projection",
"proj_out",
"audio_patchify_proj",
"audio_adaln_single",
"av_ca_audio_scale_shift_adaln_single",
"av_ca_v2a_gate_adaln_single",
"audio_caption_projection",
"audio_proj_out",
"transformer_blocks.0.",
*[f"transformer_blocks.{i}." for i in range(43, 48)],
)
def _linear_to_fp8linear(layer: nn.Linear) -> FP8Linear:
"""
Create an FP8Linear layer from an nn.Linear layer.
Args:
layer: The nn.Linear layer to convert (typically on meta device)
Returns:
A new FP8Linear with the same configuration
"""
"""Create an ``FP8Linear`` matching the shape/bias of *layer*."""
return FP8Linear(
in_features=layer.in_features,
out_features=layer.out_features,
@@ -142,15 +119,14 @@ def _linear_to_fp8linear(layer: nn.Linear) -> FP8Linear:
)
def _apply_fp8_prepare_to_model(model: nn.Module, excluded_layer_substrings: tuple[str, ...]) -> nn.Module:
"""Replace nn.Linear layers with FP8Linear in the module tree."""
def _swap_linears_to_fp8(model: nn.Module, should_swap: Callable[[str], bool]) -> nn.Module:
"""Replace nn.Linear layers with FP8Linear where ``should_swap(name)`` returns True."""
replacements: list[tuple[nn.Module, str, nn.Linear]] = []
for name, module in model.named_modules():
if not isinstance(module, nn.Linear) or isinstance(module, FP8Linear):
continue
if _should_skip_layer(name, excluded_layer_substrings):
if not should_swap(name):
continue
if "." in name:
@@ -168,40 +144,32 @@ def _apply_fp8_prepare_to_model(model: nn.Module, excluded_layer_substrings: tup
return model
def _create_transpose_kv_operation(
excluded_layer_substrings: tuple[str, ...],
) -> Callable[[str, torch.Tensor], list[KeyValueOperationResult]]:
def transpose_if_matches(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
# Only process .weight keys
if not key.endswith(".weight"):
return [KeyValueOperationResult(key, value)]
def get_fp8_swap_module_ops(checkpoint_path: str) -> tuple[ModuleOps, ...]:
"""Return the FP8 swap ``ModuleOps`` for layers whose ``.weight`` is ``F8_E4M3``
and which have a sibling ``.weight_scale`` tensor in the checkpoint.
Raises ``ValueError`` if no such layers are found — that combination is ambiguous
(a BF16 checkpoint with this policy would load as a no-op).
"""
dtypes = _read_safetensors_dtypes(checkpoint_path)
fp8_scale_paths = frozenset(
key.removesuffix(".weight_scale")
for key in dtypes
if key.endswith(".weight_scale") and dtypes.get(key.removesuffix(".weight_scale") + ".weight") == "F8_E4M3"
)
if not fp8_scale_paths:
raise ValueError(
f"fp8_scaled_mm requires a pre-quantized checkpoint with F8_E4M3 .weight + .weight_scale "
f"tensors, but {checkpoint_path!r} has none. Use QuantizationPolicy.fp8_cast() for BF16 checkpoints."
)
# Only transpose 2D FP8 tensors (Linear weights)
if value.dim() != 2 or value.dtype != torch.float8_e4m3fn:
return [KeyValueOperationResult(key, value)]
def _should_swap(name: str) -> bool:
suffix = "." + name
return any(p == name or p.endswith(suffix) for p in fp8_scale_paths)
# Check if the layer is excluded
layer_name = key.rsplit(".weight", 1)[0]
if _should_skip_layer(layer_name, excluded_layer_substrings):
return [KeyValueOperationResult(key, value)]
# Transpose to cuBLAS layout (in, out)
transposed_weight = value.t()
return [KeyValueOperationResult(key, transposed_weight)]
return transpose_if_matches
FP8_TRANSPOSE_SD_OPS = SDOps("fp8_transpose_weights").with_kv_operation(
_create_transpose_kv_operation(EXCLUDED_LAYER_SUBSTRINGS),
key_prefix="transformer_blocks.",
key_suffix=".weight",
)
FP8_PREPARE_MODULE_OPS = ModuleOps(
name="fp8_prepare_for_loading",
matcher=lambda model: isinstance(model, LTXModel),
mutator=lambda model: _apply_fp8_prepare_to_model(model, EXCLUDED_LAYER_SUBSTRINGS),
)
return (
ModuleOps(
name="fp8_swap_linears",
matcher=lambda model: isinstance(model, LTXModel),
mutator=lambda model: _swap_linears_to_fp8(model, _should_swap),
),
)
packages/ltx-core/src/ltx_core/quantization/policy.py
+22 -13
View File
@@ -1,39 +1,48 @@
from dataclasses import dataclass
from enum import Enum
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.sd_ops import SDOps
from ltx_core.quantization.fp8_cast import TRANSFORMER_LINEAR_DOWNCAST_MAP, UPCAST_DURING_INFERENCE
from ltx_core.quantization.fp8_scaled_mm import FP8_PREPARE_MODULE_OPS, FP8_TRANSPOSE_SD_OPS
from ltx_core.quantization.fp8_scaled_mm import get_fp8_swap_module_ops
@dataclass(frozen=True)
class QuantizationPolicy:
"""Configuration for model quantization during loading.
Attributes:
sd_ops: State dict operations for weight transformation.
module_ops: Post-load module transformations.
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.
"""
class Kind(str, Enum):
FP8_CAST = "fp8_cast"
FP8_SCALED_MM = "fp8_scaled_mm"
kind: Kind
sd_ops: SDOps | None = None
module_ops: tuple[ModuleOps, ...] = ()
@classmethod
def fp8_cast(cls) -> "QuantizationPolicy":
"""Create policy using FP8 casting with upcasting during inference."""
"""FP8 casting with upcasting during inference."""
return cls(
kind=cls.Kind.FP8_CAST,
sd_ops=TRANSFORMER_LINEAR_DOWNCAST_MAP,
module_ops=(UPCAST_DURING_INFERENCE,),
)
@classmethod
def fp8_scaled_mm(cls) -> "QuantizationPolicy":
"""Create policy using FP8 scaled matrix multiplication."""
try:
import tensorrt_llm # noqa: F401, PLC0415
except ImportError as e:
raise ImportError("tensorrt_llm is not installed, skipping FP8 scaled MM quantization") from e
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(
sd_ops=FP8_TRANSPOSE_SD_OPS,
module_ops=(FP8_PREPARE_MODULE_OPS,),
kind=cls.Kind.FP8_SCALED_MM,
sd_ops=None,
module_ops=get_fp8_swap_module_ops(checkpoint_path),
)
packages/ltx-core/src/ltx_core/quantization/trtllm_scaled_usable.py
+37
View File
@@ -0,0 +1,37 @@
"""Runtime detection of TensorRT-LLM FP8 scaled-matmul availability.
When the TRT-LLM ops are usable on the current host (Linux + Hopper-class CUDA
+ tensorrt_llm wheel installed) we use them since they outperform the PyTorch-native
``torch._scaled_mm`` path. Otherwise we fall back to the native implementation,
which is portable across platforms (Windows, macOS, AMD GPUs).
The check runs once and is cached.
"""
from __future__ import annotations
import platform
from functools import cache
import torch
@cache
def trtllm_scaled_mm_usable() -> bool:
if platform.system() != "Linux":
return False
if not torch.cuda.is_available():
return False
major, minor = torch.cuda.get_device_capability()
sm = major * 10 + minor
if sm < 90 or sm >= 120:
return False
# The import is load-bearing — registers the trtllm torch ops as a side effect.
try:
import tensorrt_llm # noqa: F401, PLC0415
except Exception:
return False
return True
packages/ltx-core/src/ltx_core/text_encoders/gemma/embeddings_connector.py
+29 -40
View File
@@ -18,7 +18,7 @@ class _BasicTransformerBlock1D(torch.nn.Module):
dim: int,
heads: int,
dim_head: int,
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
rope_type: LTXRopeType = LTXRopeType.SPLIT,
apply_gated_attention: bool = False,
):
super().__init__()
@@ -39,7 +39,7 @@ class _BasicTransformerBlock1D(torch.nn.Module):
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
additive_attention_mask: torch.Tensor | None = None,
pe: torch.Tensor | None = None,
) -> torch.Tensor:
# Notice that normalization is always applied before the real computation in the following blocks.
@@ -49,8 +49,8 @@ class _BasicTransformerBlock1D(torch.nn.Module):
norm_hidden_states = norm_hidden_states.squeeze(1)
# 2. Self-Attention
attn_output = self.attn1(norm_hidden_states, mask=attention_mask, pe=pe)
# 2. Self-Attention — `mask` is the kernel-boundary name for the additive mask.
attn_output = self.attn1(norm_hidden_states, mask=additive_attention_mask, pe=pe)
hidden_states = attn_output + hidden_states
if hidden_states.ndim == 4:
@@ -84,7 +84,7 @@ class Embeddings1DConnector(torch.nn.Module):
causal_temporal_positioning (bool): If True, uses causal attention (default=False).
num_learnable_registers (int | None): Number of learnable registers to replace padded tokens. If None, disables
register replacement. (default=128)
rope_type (LTXRopeType): The RoPE variant to use (default=DEFAULT_ROPE_TYPE).
rope_type (LTXRopeType): The RoPE variant to use.
double_precision_rope (bool): Use double precision rope calculation (default=False).
"""
@@ -99,7 +99,7 @@ class Embeddings1DConnector(torch.nn.Module):
positional_embedding_max_pos: list[int] | None = None,
causal_temporal_positioning: bool = False,
num_learnable_registers: int | None = 128,
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
rope_type: LTXRopeType = LTXRopeType.SPLIT,
double_precision_rope: bool = False,
apply_gated_attention: bool = False,
):
@@ -133,51 +133,40 @@ class Embeddings1DConnector(torch.nn.Module):
)
def _replace_padded_with_learnable_registers(
self, hidden_states: torch.Tensor, attention_mask: torch.Tensor
self, hidden_states: torch.Tensor, additive_attention_mask: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
assert hidden_states.shape[1] % self.num_learnable_registers == 0, (
f"Hidden states sequence length {hidden_states.shape[1]} must be divisible by num_learnable_registers "
f"{self.num_learnable_registers}."
)
batch_size, seq_len, _ = hidden_states.shape
num_registers_duplications = hidden_states.shape[1] // self.num_learnable_registers
learnable_registers = torch.tile(self.learnable_registers, (num_registers_duplications, 1))
attention_mask_binary = (attention_mask.squeeze(1).squeeze(1).unsqueeze(-1) >= -9000.0).int()
assert seq_len % self.num_learnable_registers == 0
non_zero_hidden_states = hidden_states[:, attention_mask_binary.squeeze().bool(), :]
non_zero_nums = non_zero_hidden_states.shape[1]
pad_length = hidden_states.shape[1] - non_zero_nums
adjusted_hidden_states = torch.nn.functional.pad(non_zero_hidden_states, pad=(0, 0, 0, pad_length), value=0)
flipped_mask = torch.flip(attention_mask_binary, dims=[1])
hidden_states = flipped_mask * adjusted_hidden_states + (1 - flipped_mask) * learnable_registers
registers = self.learnable_registers.repeat(seq_len // self.num_learnable_registers, 1).to(hidden_states.dtype)
registers = registers.unsqueeze(0).expand(batch_size, -1, -1) # (B, seq_len, hidden_dim)
binary_mask = additive_attention_mask[:, 0, 0, :].unsqueeze(-1) >= 0
binary_mask = binary_mask.to(hidden_states.dtype)
hidden_states = binary_mask * hidden_states + (1 - binary_mask) * registers
attention_mask = torch.full_like(
attention_mask,
0.0,
dtype=attention_mask.dtype,
device=attention_mask.device,
)
return hidden_states, attention_mask
return hidden_states, torch.zeros_like(additive_attention_mask)
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
additive_attention_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Forward pass of Embeddings1DConnector.
"""Forward pass of Embeddings1DConnector.
Args:
hidden_states (torch.Tensor): Input tensor of embeddings (shape [batch, seq_len, feature_dim]).
attention_mask (torch.Tensor|None): Optional mask for valid tokens (shape compatible with hidden_states).
hidden_states: (B, S, D) input embeddings.
additive_attention_mask: optional additive mask of shape (B, 1, 1, S), where
valid = 0.0 and padding = -torch.finfo(dtype).max.
Returns:
tuple[torch.Tensor, torch.Tensor]: Processed features and the corresponding (possibly modified) mask.
(hidden_states, additive_attention_mask)
"""
if self.num_learnable_registers:
hidden_states, attention_mask = self._replace_padded_with_learnable_registers(hidden_states, attention_mask)
hidden_states, additive_attention_mask = self._replace_padded_with_learnable_registers(
hidden_states, additive_attention_mask
)
indices_grid = torch.arange(hidden_states.shape[1], dtype=torch.float32, device=hidden_states.device)
indices_grid = indices_grid[None, None, :]
indices_grid = indices_grid[None, None, :].expand(hidden_states.shape[0], -1, -1)
freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch
freqs_cis = precompute_freqs_cis(
indices_grid=indices_grid,
@@ -191,11 +180,11 @@ class Embeddings1DConnector(torch.nn.Module):
)
for block in self.transformer_1d_blocks:
hidden_states = block(hidden_states, attention_mask=attention_mask, pe=freqs_cis)
hidden_states = block(hidden_states, additive_attention_mask=additive_attention_mask, pe=freqs_cis)
hidden_states = rms_norm(hidden_states)
return hidden_states, attention_mask
return hidden_states, additive_attention_mask
class Embeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConnector]):
@@ -204,7 +193,7 @@ class Embeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConnector]
@classmethod
def from_config(cls: type[Embeddings1DConnector], config: dict) -> Embeddings1DConnector:
transformer_config = config.get("transformer", {})
rope_type = LTXRopeType(transformer_config.get("rope_type", "interleaved"))
rope_type = LTXRopeType(transformer_config.get("rope_type", "split"))
double_precision_rope = transformer_config.get("frequencies_precision", False) == "float64"
pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
@@ -231,7 +220,7 @@ class AudioEmbeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConne
@classmethod
def from_config(cls: type[Embeddings1DConnector], config: dict) -> Embeddings1DConnector:
transformer_config = config.get("transformer", {})
rope_type = LTXRopeType(transformer_config.get("rope_type", "interleaved"))
rope_type = LTXRopeType(transformer_config.get("rope_type", "split"))
double_precision_rope = transformer_config.get("frequencies_precision", False) == "float64"
pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
packages/ltx-core/src/ltx_core/text_encoders/gemma/embeddings_processor.py
+36 -9
View File
@@ -19,12 +19,32 @@ def convert_to_additive_mask(attention_mask: torch.Tensor, dtype: torch.dtype) -
) * torch.finfo(dtype).max
def _to_binary_mask(encoded: torch.Tensor, encoded_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Convert connector output mask to binary mask and apply to encoded tensor."""
binary_mask = (encoded_mask < 0.000001).to(torch.int64)
binary_mask = binary_mask.reshape([encoded.shape[0], encoded.shape[1], 1])
encoded = encoded * binary_mask
return encoded, binary_mask
def _compute_right_pad_order(additive_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Compute the index permutation that places valid tokens before pads in each row.
Stable sort: valid tokens keep their relative order. Idempotent for inputs already
right-padded. The sort and reordered mask depend only on the mask, so they can be
computed once and reused across multiple feature tensors that share the mask.
Args:
additive_mask: (B, 1, 1, S) additive mask, ``0.0`` for valid, ``-finfo.max`` for pad.
Returns:
``(sort_idx, reordered_additive_mask)``: ``sort_idx`` is (B, S); the reordered mask
has the same shape as the input.
"""
binary = (additive_mask[:, 0, 0, :] >= 0).to(torch.int32) # (B, S)
sort_idx = torch.argsort(binary, dim=-1, descending=True, stable=True) # (B, S)
new_binary = torch.gather(binary, 1, sort_idx)
new_additive = (new_binary.to(additive_mask.dtype) - 1) * torch.finfo(additive_mask.dtype).max
return sort_idx, new_additive[:, None, None, :]
def _apply_right_pad_order(features: torch.Tensor, sort_idx: torch.Tensor) -> torch.Tensor:
"""Apply a precomputed right-pad permutation (from ``_compute_right_pad_order``) to features."""
return torch.gather(features, 1, sort_idx.unsqueeze(-1).expand_as(features))
def _to_binary_mask(encoded_mask: torch.Tensor, lead_shape: tuple[int, int]) -> torch.Tensor:
"""Convert connector output mask to a binary (0/1) mask shaped ``(B, S, 1)`` for broadcasting."""
return (encoded_mask < 0.000001).to(torch.int64).reshape([lead_shape[0], lead_shape[1], 1])
class EmbeddingsProcessor(nn.Module):
@@ -57,12 +77,19 @@ class EmbeddingsProcessor(nn.Module):
if self.audio_connector is None and audio_features is not None:
raise ValueError("Audio features were provided but no audio connector is configured.")
video_encoded, video_mask = self.video_connector(video_features, additive_attention_mask)
video_encoded, binary_mask = _to_binary_mask(video_encoded, video_mask)
# Connectors expect right-padded input ([valid, pad]). Normalize layout here so the
# upstream tokenizer can keep using either side without coupling to the connector.
# The sort index depends only on the mask, so compute it once and reuse for audio.
sort_idx, mask_for_connector = _compute_right_pad_order(additive_attention_mask)
video_features = _apply_right_pad_order(video_features, sort_idx)
video_encoded, video_mask = self.video_connector(video_features, mask_for_connector)
binary_mask = _to_binary_mask(video_mask, video_encoded.shape[:2])
video_encoded = video_encoded * binary_mask
audio_encoded = None
if self.audio_connector is not None:
audio_encoded, _ = self.audio_connector(audio_features, additive_attention_mask)
audio_features = _apply_right_pad_order(audio_features, sort_idx)
audio_encoded, _ = self.audio_connector(audio_features, mask_for_connector)
return video_encoded, audio_encoded, binary_mask.squeeze(-1)
packages/ltx-core/src/ltx_core/text_encoders/gemma/feature_extractor.py
+15 -27
View File
@@ -11,37 +11,25 @@ from torch import nn
def _norm_and_concat_padded_batch(
encoded_text: torch.Tensor,
sequence_lengths: torch.Tensor,
padding_side: str = "right",
attention_mask: torch.Tensor,
) -> torch.Tensor:
"""Normalize and flatten multi-layer hidden states, respecting padding.
Performs per-batch, per-layer normalization using masked mean and range,
then concatenates across the layer dimension.
then concatenates across the layer dimension. Padding-side agnostic: the
binary ``attention_mask`` already encodes which positions are valid.
Args:
encoded_text: Hidden states of shape [batch, seq_len, hidden_dim, num_layers].
sequence_lengths: Number of valid (non-padded) tokens per batch item.
padding_side: Whether padding is on "left" or "right".
attention_mask: Binary mask of shape [batch, seq_len], 1 for valid tokens, 0 for padding.
Returns:
Normalized tensor of shape [batch, seq_len, hidden_dim * num_layers],
with padded positions zeroed out.
"""
b, t, d, l = encoded_text.shape # noqa: E741
device = encoded_text.device
token_indices = torch.arange(t, device=device)[None, :]
if padding_side == "right":
mask = token_indices < sequence_lengths[:, None]
elif padding_side == "left":
start_indices = t - sequence_lengths[:, None]
mask = token_indices >= start_indices
else:
raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}")
mask = rearrange(mask, "b t -> b t 1 1")
b, _, d, l = encoded_text.shape # noqa: E741
eps = 1e-6
sequence_lengths = attention_mask.sum(dim=-1)
mask = rearrange(attention_mask.bool(), "b t -> b t 1 1")
masked = encoded_text.masked_fill(~mask, 0.0)
denom = (sequence_lengths * d).view(b, 1, 1, 1)
mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + eps)
@@ -51,12 +39,10 @@ def _norm_and_concat_padded_batch(
range_ = x_max - x_min
normed = 8 * (encoded_text - mean) / (range_ + eps)
normed = normed.reshape(b, t, -1)
normed = normed.reshape(b, -1, d * l)
mask_flattened = rearrange(mask, "b t 1 1 -> b t 1").expand(-1, -1, d * l)
normed = normed.masked_fill(~mask_flattened, 0.0)
return normed
return normed.masked_fill(~mask_flattened, 0.0)
def norm_and_concat_per_token_rms(
@@ -97,12 +83,14 @@ class FeatureExtractorV1(nn.Module):
self.is_av = is_av
def forward(
self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, padding_side: str = "left"
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor,
padding_side: str = "left", # noqa: ARG002 — kept for API stability; norm is layout-agnostic
) -> tuple[torch.Tensor, torch.Tensor | None]:
encoded = torch.stack(hidden_states, dim=-1) if isinstance(hidden_states, (list, tuple)) else hidden_states
dtype = encoded.dtype
sequence_lengths = attention_mask.sum(dim=-1)
normed = _norm_and_concat_padded_batch(encoded, sequence_lengths, padding_side)
normed = _norm_and_concat_padded_batch(encoded, attention_mask)
features = self.aggregate_embed(normed.to(dtype))
if self.is_av:
return features, features
packages/ltx-core/src/ltx_core/text_encoders/gemma/tokenizer.py
+10 -3
View File
@@ -1,6 +1,13 @@
from enum import Enum
from transformers import AutoTokenizer
class PaddingSide(str, Enum):
LEFT = "left"
RIGHT = "right"
class LTXVGemmaTokenizer:
"""
Tokenizer wrapper for Gemma models compatible with LTXV processes.
@@ -8,18 +15,18 @@ class LTXVGemmaTokenizer:
ensuring correct settings and output formatting for downstream consumption.
"""
def __init__(self, tokenizer_path: str, max_length: int = 256):
def __init__(self, tokenizer_path: str, max_length: int = 256, padding_side: PaddingSide = PaddingSide.LEFT):
"""
Initialize the tokenizer.
Args:
tokenizer_path (str): Path to the pretrained tokenizer files or model directory.
max_length (int, optional): Max sequence length for encoding. Defaults to 256.
padding_side (PaddingSide, optional): Side to pad on. Defaults to ``PaddingSide.LEFT``.
"""
self.tokenizer = AutoTokenizer.from_pretrained(
tokenizer_path, local_files_only=True, model_max_length=max_length
)
# Gemma expects left padding for chat-style prompts; for plain text it doesn't matter much.
self.tokenizer.padding_side = "left"
self.tokenizer.padding_side = padding_side.value
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
packages/ltx-core/src/ltx_core/tools.py
+1 -1
View File
@@ -138,7 +138,7 @@ class VideoLatentTools(LatentTools):
LatentState(
latent=initial_latent,
denoise_mask=denoise_mask,
positions=positions.to(dtype),
positions=positions,
clean_latent=clean_latent,
)
)
packages/ltx-core/src/ltx_core/types.py
+7 -5
View File
@@ -20,15 +20,17 @@ class SpatioTemporalScaleFactors(NamedTuple):
"""
Describes the spatiotemporal downscaling between decoded video space and
the corresponding VAE latent grid.
Field order matches the (frame/time, height, width) axis layout used by
latent tensors and meshgrid coordinates elsewhere in the codebase.
"""
time: int
width: int
height: int
width: int
@classmethod
def default(cls) -> "SpatioTemporalScaleFactors":
return cls(time=8, width=32, height=32)
return cls(time=8, height=32, width=32)
VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
@@ -74,9 +76,9 @@ class VideoLatentShape(NamedTuple):
latent_channels: int = 128,
scale_factors: SpatioTemporalScaleFactors = VIDEO_SCALE_FACTORS,
) -> "VideoLatentShape":
frames = (shape.frames - 1) // scale_factors[0] + 1
height = shape.height // scale_factors[1]
width = shape.width // scale_factors[2]
frames = (shape.frames - 1) // scale_factors.time + 1
height = shape.height // scale_factors.height
width = shape.width // scale_factors.width
return VideoLatentShape(
batch=shape.batch,