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

Automated PR - 2026-04-13

Browse Source
This commit is contained in:
github-actions[bot]
2026-04-13 14:29:35 +00:00
parent 59ca828d5a
commit d887bbd1e0
29 changed files with 463 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-core/pyproject.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[project]
name = "ltx-core"
version = "1.1.0"
version = "1.1.1"
description = "Core implementation of Lightricks' LTX-2 model"
readme = "README.md"
requires-python = ">=3.10"
packages/ltx-core/src/ltx_core/layer_streaming.py
+32 -50
View File
@@ -44,32 +44,26 @@ def _resolve_attr(module: nn.Module, dotted_path: str) -> nn.ModuleList:
class _LayerStore:
"""Manages on-demand pinning of layer parameters for GPU streaming.
Stores references to each layer's source data (which may be file-backed
mmap views or in-memory tensors). When a layer needs to be transferred
to GPU, its source data is pinned on demand and copied; on eviction the
pinned copy is freed and the source data is restored.
"""Manages CPU-pinned copies of layer parameters/buffers.
Tracks which layers currently reside on GPU so the prefetcher and evictor
can make correct decisions.
"""
def __init__(self, layers: nn.ModuleList, target_device: torch.device) -> None:
self.target_device = target_device
self.num_layers = len(layers)
# CPU-pinned copies keyed by (layer_idx, param_name)
self._pinned: list[dict[str, torch.Tensor]] = []
self._on_gpu: set[int] = set()
# Keep a reference to the source data for each layer so we can pin it
# on demand and restore it after eviction.
self._source_data: list[dict[str, torch.Tensor]] = []
for layer in layers:
source: dict[str, torch.Tensor] = {}
pinned: dict[str, torch.Tensor] = {}
for name, tensor in itertools.chain(layer.named_parameters(), layer.named_buffers()):
source[name] = tensor.data
self._source_data.append(source)
# Hold pinned tensors alive until the H2D transfer completes.
# Without this, the CachingHostAllocator can reclaim a pinned tensor
# as soon as its Python reference is dropped, even if an async H2D
# transfer is still reading from it.
self._pinned_in_flight: dict[int, list[torch.Tensor]] = {}
pinned_tensor = tensor.data.pin_memory()
tensor.data = pinned_tensor
pinned[name] = pinned_tensor
self._pinned.append(pinned)
def _check_idx(self, idx: int) -> None:
if idx < 0 or idx >= self.num_layers:
@@ -79,45 +73,34 @@ class _LayerStore:
return idx in self._on_gpu
def move_to_gpu(self, idx: int, layer: nn.Module, *, non_blocking: bool = False) -> None:
"""Pin layer *idx* on demand, then transfer to GPU."""
"""Move layer *idx* parameters from pinned CPU to *target_device*."""
self._check_idx(idx)
if idx in self._on_gpu:
return
source = self._source_data[idx]
pinned_refs: list[torch.Tensor] = []
pinned = self._pinned[idx]
for name, param in itertools.chain(layer.named_parameters(), layer.named_buffers()):
pinned = source[name].pin_memory()
param.data = pinned.to(self.target_device, non_blocking=non_blocking)
pinned_refs.append(pinned)
# Keep pinned tensors alive until eviction — the async H2D transfer
# may still be reading from them.
self._pinned_in_flight[idx] = pinned_refs
param.data = pinned[name].to(self.target_device, non_blocking=non_blocking)
self._on_gpu.add(idx)
def evict_to_cpu(self, idx: int, layer: nn.Module) -> None:
"""Restore source data, freeing the GPU and pinned copies."""
"""Swap layer *idx* parameters back to their pinned CPU copies."""
self._check_idx(idx)
if idx not in self._on_gpu:
return
source = self._source_data[idx]
pinned = self._pinned[idx]
for name, param in itertools.chain(layer.named_parameters(), layer.named_buffers()):
param.data = source[name]
# Release pinned tensors — the H2D transfer is complete by now
# (the compute stream waited on the prefetch event before using
# the layer, and we only evict after compute finishes).
self._pinned_in_flight.pop(idx, None)
param.data = pinned[name]
self._on_gpu.discard(idx)
def cleanup(self) -> None:
"""Release all source data and in-flight pinned references.
After this call, the source tensors can be garbage-collected once
"""Release all pinned memory references.
After this call, the pinned tensors can be garbage-collected once
the layer parameters (which still reference them via ``.data``) are
also released (e.g. via ``.to("meta")``).
"""
for source_dict in self._source_data:
source_dict.clear()
self._source_data.clear()
self._pinned_in_flight.clear()
for pinned_dict in self._pinned:
pinned_dict.clear()
self._pinned.clear()
class _AsyncPrefetcher:
@@ -228,8 +211,6 @@ class LayerStreamingWrapper(nn.Module):
idx_map: dict[int, int] = {id(layer): idx for idx, layer in enumerate(self._layers)}
num_layers = len(self._layers)
compute_stream = torch.cuda.current_stream(self._target_device)
def _pre_hook(
module: nn.Module,
_args: Any, # noqa: ANN401
@@ -246,6 +227,7 @@ class LayerStreamingWrapper(nn.Module):
# caching allocator would allow the prefetch stream to reuse their
# memory immediately after eviction — even if the compute kernel
# that reads them hasn't finished yet.
compute_stream = torch.cuda.current_stream(self._target_device)
for param in itertools.chain(module.parameters(), module.buffers()):
param.data.record_stream(compute_stream)
@@ -270,12 +252,12 @@ class LayerStreamingWrapper(nn.Module):
self._hooks.extend([h1, h2])
def teardown(self) -> None:
"""Remove hooks, release resources, and move parameters back to CPU.
"""Remove hooks, release pinned memory, and move parameters back to CPU.
After this call the wrapper is inert: hooks are removed, the prefetch
stream is drained and destroyed, all parameters reside on CPU, and the
``_LayerStore`` source data references are cleared. Callers should
still follow up with ``.to("meta")`` to release the CPU copies if the
model is no longer needed.
stream is drained and destroyed, all parameters reside on regular
(non-pinned) CPU memory, and the ``_LayerStore`` pinned-tensor cache is
cleared. Callers should still follow up with ``.to("meta")`` to release
the CPU copies if the model is no longer needed.
"""
for h in self._hooks:
h.remove()
@@ -298,10 +280,10 @@ class LayerStreamingWrapper(nn.Module):
for b in self._model.buffers():
b.data = b.data.to("cpu")
# Release source data references. After evict_to_cpu() the layer
# params point to the source data. The caller is expected to follow
# up with .to("meta") to drop the param refs; cleanup() drops the
# store's refs.
# Release pinned memory. After evict_to_cpu() the layer parameters
# still reference the pinned tensors (since .to("cpu") on a pinned
# tensor is a no-op). The caller is expected to follow up with
# .to("meta") to drop the param refs; cleanup() drops the store's refs.
self._store.cleanup()
# ------------------------------------------------------------------
packages/ltx-core/src/ltx_core/modality_tiling.py
+14 -2
View File
@@ -67,11 +67,16 @@ class VideoModalityTilingHelper:
# -- tile modality -----------------------------------------------------
def tile_modality(self, modality: Modality, tile: Tile) -> tuple[Modality, TilingContext]:
def tile_modality(
self, modality: Modality, tile: Tile, *, normalize_positions: bool = True
) -> tuple[Modality, TilingContext]:
"""Slice *modality* to the tokens covered by *tile*.
Selects generated tokens belonging to the tile's spatial region
and conditioning tokens that overlap with the tile (or have
negative time coordinates).
Args:
normalize_positions: When True, shift all positions so the
tile's generated tokens start at zero in every dimension.
Returns:
A ``(tiled_modality, context)`` tuple. Pass *context* to
:meth:`blend` together with the model output.
@@ -83,11 +88,18 @@ class VideoModalityTilingHelper:
keep_indices = keep_mask.nonzero(as_tuple=False).squeeze(1)
tile_attention_mask = modality.attention_mask[:, keep_indices, :][:, :, keep_indices]
positions = modality.positions[:, :, keep_mask, :]
if normalize_positions:
num_tile_gen = self._tile_generated_token_count(tile)
gen_pos = positions[:, :, :num_tile_gen, :] # (B, 3, num_tile_gen, 2)
offset = gen_pos[..., 0].amin(dim=2, keepdim=True).unsqueeze(-1) # (B, 3, 1, 1)
positions = positions - offset
tiled = replace(
modality,
latent=modality.latent[:, keep_mask, :],
timesteps=modality.timesteps[:, keep_mask],
positions=modality.positions[:, :, keep_mask, :],
positions=positions,
attention_mask=tile_attention_mask,
)