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

Automated PR - 2026-04-23

Browse Source
This commit is contained in:
github-actions[bot]
2026-04-23 12:43:54 +00:00
parent a2c3f24078
commit b604d3fab3
49 changed files with 2664 additions and 568 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.1"
version = "1.1.2"
description = "Core implementation of Lightricks' LTX-2 model"
readme = "README.md"
requires-python = ">=3.10"
packages/ltx-core/src/ltx_core/batch_split.py
+2 -2
View File
@@ -1,5 +1,5 @@
"""Batch-splitting adapter for the transformer.
Wraps an ``X0Model`` (or ``LayerStreamingWrapper``) and splits batched inputs
Wraps an ``X0Model`` (or ``BlockStreamingWrapper``) and splits batched inputs
into smaller chunks before forwarding, then concatenates the results. This
controls peak activation memory at the cost of more forward passes.
The adapter is transparent — it has the same ``forward`` signature as
@@ -42,7 +42,7 @@ class BatchSplitAdapter(nn.Module):
Has the same ``forward`` signature as ``X0Model``:
``(video, audio, perturbations) -> (denoised_video, denoised_audio)``.
Args:
model: The model to wrap (``X0Model``, ``LayerStreamingWrapper``, etc.).
model: The model to wrap (``X0Model``, ``BlockStreamingWrapper``, etc.).
max_batch_size: Maximum batch size per forward pass. Input batches
larger than this are split into sequential chunks.
"""
packages/ltx-core/src/ltx_core/block_streaming/__init__.py
+19
View File
@@ -0,0 +1,19 @@
"""Block streaming: memory-efficient sequential-block inference.
Streams transformer blocks from safetensors to GPU one at a time.
Block weights are provided by a :class:`WeightsProvider` which handles
CPU-to-GPU copies, caching, and stream synchronization. Two weight
source strategies are available:
- **RAM streaming** (default): all blocks pre-loaded into pinned CPU
buffers with LoRA fusion at build time. Fast, higher CPU memory.
- **Disk streaming** (``cpu_slots < num_blocks``): blocks read from
disk on demand with FIFO eviction. Slower, lower CPU memory.
"""
from ltx_core.block_streaming.builder import DISK_CPU_SLOTS, StreamingModelBuilder
from ltx_core.block_streaming.wrapper import BlockStreamingWrapper
__all__ = [
"DISK_CPU_SLOTS",
"BlockStreamingWrapper",
"StreamingModelBuilder",
]
packages/ltx-core/src/ltx_core/block_streaming/builder.py
+305
View File
@@ -0,0 +1,305 @@
"""Builder that constructs a BlockStreamingWrapper from safetensors checkpoints."""
from __future__ import annotations
import logging
from collections.abc import Callable
from dataclasses import dataclass, field, replace
from typing import Generic
import 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.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.wrapper import BlockStreamingWrapper
from ltx_core.loader.fuse_loras import apply_loras
from ltx_core.loader.helpers import create_meta_model, load_state_dict, read_model_config
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.primitives import (
LoraPathStrengthAndSDOps,
LoraStateDictWithStrength,
ModelBuilderProtocol,
StateDictLoader,
)
from ltx_core.loader.registry import DummyRegistry, Registry
from ltx_core.loader.sd_ops import SDOps
from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
from ltx_core.model.model_protocol import ModelConfigurator, ModelType
logger = logging.getLogger(__name__)
DISK_CPU_SLOTS = 2
_DEFAULT_GPU_SLOTS = 2
@dataclass(frozen=True)
class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType]):
"""Immutable builder for :class:`BlockStreamingWrapper`.
Reads block weights from safetensors on demand. ``cpu_slots`` and
``gpu_slots`` control the memory/speed trade-off (see :meth:`build`).
Args:
model_class_configurator: Creates the model from a config dict.
model_path: One or more ``.safetensors`` checkpoint paths.
model_sd_ops: Key remapping applied to safetensors keys.
module_ops: Module-level mutations for the meta model.
loras: LoRA adapters fused into weights at load time.
model_loader: Strategy for reading checkpoint metadata.
registry: Shared cache for loaded state dicts.
blocks_attr: Dotted path to the ``nn.ModuleList`` (e.g.
``"velocity_model.transformer_blocks"``).
blocks_prefix: State-dict key prefix for block weights
(e.g. ``"transformer_blocks"``).
state_dict_prefix: Key prefix for non-block weights
(e.g. ``"velocity_model."``).
model_wrapper: Optional callable wrapping the model
(e.g. ``X0Model``).
"""
model_class_configurator: type[ModelConfigurator[ModelType]]
model_path: str | tuple[str, ...]
model_sd_ops: SDOps | None = None
module_ops: tuple[ModuleOps, ...] = field(default_factory=tuple)
loras: tuple[LoraPathStrengthAndSDOps, ...] = field(default_factory=tuple)
model_loader: StateDictLoader = field(default_factory=SafetensorsModelStateDictLoader)
registry: Registry = field(default_factory=DummyRegistry)
# Streaming-specific
blocks_attr: str = ""
blocks_prefix: str = ""
state_dict_prefix: str = ""
model_wrapper: Callable[[ModelType], nn.Module] | None = None
def with_sd_ops(self, sd_ops: SDOps | None) -> StreamingModelBuilder:
return replace(self, model_sd_ops=sd_ops)
def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> StreamingModelBuilder:
return replace(self, module_ops=module_ops)
def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> StreamingModelBuilder:
return replace(self, loras=loras)
def model_config(self) -> dict:
"""Read model configuration from the checkpoint metadata."""
return read_model_config(self.model_path, self.model_loader)
def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType:
"""Create a model on the meta device and apply module operations."""
return create_meta_model(self.model_class_configurator, config, module_ops)
def build(
self,
target_device: torch.device,
dtype: torch.dtype,
cpu_slots_count: int | None = None,
gpu_slots_count: int | None = None,
**_kwargs: object,
) -> BlockStreamingWrapper:
"""Build and return a ready-to-use :class:`BlockStreamingWrapper`.
Args:
target_device: GPU device for compute.
dtype: Weight dtype (e.g. ``torch.bfloat16``).
cpu_slots_count: Number of pinned CPU buffer slots.
``None`` = RAM streaming (all blocks pre-loaded with LoRA fusion).
gpu_slots_count: Number of GPU buffer slots.
``None`` = ``_DEFAULT_GPU_SLOTS`` (2).
"""
if not self.blocks_prefix:
raise ValueError("blocks_prefix must be non-empty for streaming")
# 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:
meta_model = self.model_wrapper(meta_model)
meta_model.eval()
blocks = resolve_attr(meta_model, self.blocks_attr)
layout = build_pool_layout(blocks[0], dtype)
# 2. Determine slot counts.
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)
else:
source, lora_sources = self._build_disk_source(meta_model, layout, target_device, dtype, cpu_slots_count)
# 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)
)
provider = WeightsProvider(gpu_pool, copy_stream, target_device, source, lora_sources, self.blocks_prefix)
return BlockStreamingWrapper(
model=meta_model,
blocks=blocks,
provider=provider,
target_device=target_device,
)
def _build_pinned_source(
self,
meta_model: nn.Module,
target_device: torch.device,
dtype: torch.dtype,
cpu_slots_count: int,
) -> 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,
)
# 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 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
else:
non_block_sd[self.state_dict_prefix + key] = tensor.to(device=target_device, dtype=dtype)
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,
) -> tuple[WeightSource, list[LoraSource]]:
"""Create a DiskWeightSource backed by a DiskBlockReader for lazy loading."""
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,
non_block_keys,
meta_model,
target_device,
dtype,
sd_ops=self.model_sd_ops,
key_prefix=self.state_dict_prefix,
lora_sources=lora_sources,
matmul_device=target_device,
)
cpu_pool = WeightPool(
layout,
cpu_slots_count,
torch.device("cpu"),
reuse_barrier=lambda event: event.synchronize(),
pin_memory=True,
)
block_reader = DiskBlockReader(reader=reader, block_key_map=block_key_map, dtype=dtype)
source = DiskWeightSource(cpu_pool, block_reader)
return source, lora_sources
# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------
@staticmethod
def _fuse_lora_delta(
model_key: str,
tensor: torch.Tensor,
lora_sources: list[LoraSource],
matmul_device: torch.device | None = None,
) -> torch.Tensor:
"""Add all matching LoRA deltas to *tensor* in-place."""
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))
return tensor
@staticmethod
@torch.inference_mode()
def _load_non_block_weights(
reader: DiskTensorReader,
non_block_keys: list[tuple[str, str]],
model: nn.Module,
device: torch.device,
dtype: torch.dtype,
sd_ops: SDOps | None = None,
key_prefix: str = "",
lora_sources: list[LoraSource] | None = None,
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)
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)
packages/ltx-core/src/ltx_core/block_streaming/disk.py
+106
View File
@@ -0,0 +1,106 @@
"""Safetensors I/O and LoRA fusion for block streaming."""
from __future__ import annotations
import safetensors
import torch
from ltx_core.loader.sd_ops import SDOps
class DiskTensorReader:
"""Key-based tensor accessor over one or more safetensors files."""
def __init__(self, paths: list[str]) -> None:
self._handles: list[safetensors.safe_open] = []
self._key_to_handle_idx: dict[str, int] = {}
for path in paths:
handle = safetensors.safe_open(path, framework="pt", device="cpu")
handle_idx = len(self._handles)
self._handles.append(handle)
for sft_key in handle.keys(): # noqa: SIM118
self._key_to_handle_idx[sft_key] = handle_idx
def 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)
def close(self) -> None:
self._handles.clear()
self._key_to_handle_idx.clear()
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.
"""
def __init__(
self,
reader: DiskTensorReader,
block_key_map: dict[int, list[tuple[str, str]]],
dtype: torch.dtype,
) -> None:
self._reader = reader
self._block_key_map = block_key_map
self._dtype = dtype
def read_into(self, target: dict[str, torch.Tensor], block_idx: int) -> None:
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)
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.
"""
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:
continue
if model_key.endswith(".lora_A.weight"):
a_keys[model_key[: -len(".lora_A.weight")]] = sft_key
elif model_key.endswith(".lora_B.weight"):
b_keys[model_key[: -len(".lora_B.weight")]] = sft_key
# 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(),
)
def get_delta(self, param_prefix: str, device: torch.device | None = None) -> torch.Tensor | None:
"""Return ``(B * strength) @ A`` 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
def cleanup(self) -> None:
self._pinned_ab.clear()
packages/ltx-core/src/ltx_core/block_streaming/pool.py
+64
View File
@@ -0,0 +1,64 @@
"""Weight buffer pool for block streaming."""
from __future__ import annotations
from collections import deque
from typing import Callable
import torch
from ltx_core.block_streaming.utils import allocate_buffer
# Type alias for the buffer layout used by slot allocation.
BlockLayout = dict[str, tuple[torch.Size, torch.dtype]]
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.
Args:
layout: ``{name: (shape, dtype)}`` for each buffer.
capacity: Number of buffers to pre-allocate.
device: Device for allocation.
reuse_barrier: Called with the pending event before a buffer is reused.
pin_memory: Pin buffers (for async H2D copies from CPU).
"""
def __init__(
self,
layout: BlockLayout,
capacity: int,
device: torch.device,
reuse_barrier: Callable[[torch.cuda.Event], None],
pin_memory: bool = False,
) -> None:
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))
@property
def capacity(self) -> int:
return self._capacity
def acquire(self) -> dict[str, torch.Tensor]:
"""Take a free buffer, waiting any pending event before returning."""
weights = self._free.popleft()
event = self._events.pop(id(weights), None)
if event is not None:
self._reuse_barrier(event)
return weights
def release(self, weights: dict[str, torch.Tensor], event: torch.cuda.Event | None = None) -> None:
"""Return a buffer to the free list.
If *event* is given it is waited on the next :meth:`acquire`
of this buffer, ensuring the prior operation has completed.
"""
if event is not None:
self._events[id(weights)] = event
self._free.append(weights)
packages/ltx-core/src/ltx_core/block_streaming/provider.py
+110
View File
@@ -0,0 +1,110 @@
"""GPU weights provider for block streaming."""
from __future__ import annotations
from collections import OrderedDict
import torch
from ltx_core.block_streaming.disk import LoraSource
from ltx_core.block_streaming.pool import WeightPool
from ltx_core.block_streaming.source import WeightSource
class WeightsProvider:
"""Provides GPU-ready block weights via H2D copy from a pinned CPU weight source.
Args:
pool: Pre-allocated GPU weight buffer pool.
copy_stream: Dedicated CUDA stream for async H2D copies.
target_device: GPU device for compute.
source: Pinned CPU weight source.
lora_sources: LoRA adapters fused on H2D copy.
blocks_prefix: State-dict prefix for LoRA key matching.
"""
def __init__(
self,
pool: WeightPool,
copy_stream: torch.cuda.Stream,
target_device: torch.device,
source: WeightSource,
lora_sources: list[LoraSource] | None = None,
blocks_prefix: str = "",
) -> None:
self._copy_stream = copy_stream
self._pool = pool
self._cache: OrderedDict[int, dict[str, torch.Tensor]] = OrderedDict()
self._events: dict[int, torch.cuda.Event] = {}
self._target_device = target_device
self._source = source
self._lora_sources = lora_sources or []
self._blocks_prefix = blocks_prefix
def get(self, idx: int) -> dict[str, torch.Tensor]:
"""Return GPU weights for block *idx*. Does H2D copy on miss."""
if idx in self._cache:
return self._cache[idx]
# Evict oldest GPU buffer if at capacity.
if len(self._cache) >= self._pool.capacity:
evicted_idx, evicted_weights = self._cache.popitem(last=False)
self._pool.release(evicted_weights, event=self._events.pop(evicted_idx, None))
gpu_weights = self._pool.acquire()
cpu_weights = self._source.get(idx)
h2d_event = self._copy_to_gpu(idx, gpu_weights, cpu_weights)
self._source.release(idx, event=h2d_event)
self._cache[idx] = gpu_weights
return gpu_weights
def _copy_to_gpu(
self,
idx: int,
gpu_weights: dict[str, torch.Tensor],
cpu_weights: dict[str, torch.Tensor],
) -> torch.cuda.Event:
"""Enqueue H2D copy + LoRA fusion on the copy stream and wait on compute.
The wait is intentionally inside this method so callers -- and
instrumentation regions wrapping it -- observe the full transfer time.
"""
with torch.cuda.stream(self._copy_stream):
for name, gpu_tensor in gpu_weights.items():
gpu_tensor.copy_(cpu_weights[name], non_blocking=True)
if self._lora_sources:
self._fuse_block_loras(idx, gpu_weights)
h2d_event = torch.cuda.Event()
h2d_event.record(self._copy_stream)
torch.cuda.current_stream(self._target_device).wait_event(h2d_event)
return h2d_event
def release(self, idx: int, event: torch.cuda.Event) -> None:
"""Attach a compute-done event -- waited before this buffer is recycled."""
self._events[idx] = event
def cleanup(self) -> None:
"""Synchronize streams and release all resources."""
self._copy_stream.synchronize()
torch.cuda.current_stream(self._target_device).synchronize()
self._cache.clear()
self._events.clear()
self._source.cleanup()
for lora in self._lora_sources:
lora.cleanup()
def __len__(self) -> int:
return len(self._cache)
def _fuse_block_loras(self, idx: int, weights: dict[str, torch.Tensor]) -> None:
"""Fuse LoRA deltas directly into GPU block weights."""
for name, tensor in weights.items():
if not name.endswith(".weight"):
continue
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))
packages/ltx-core/src/ltx_core/block_streaming/source.py
+83
View File
@@ -0,0 +1,83 @@
"""Weight sources for block streaming: protocol and implementations."""
from __future__ import annotations
from collections import OrderedDict
from typing import Protocol
import torch
from ltx_core.block_streaming.disk import DiskBlockReader
from ltx_core.block_streaming.pool import WeightPool
class WeightSource(Protocol):
"""Provides pinned CPU weights for a given block index."""
def get(self, idx: int) -> dict[str, torch.Tensor]:
"""Return CPU weights for block *idx*."""
...
def release(self, idx: int, event: torch.cuda.Event) -> None:
"""Signal that an async operation using these weights is guarded by *event*."""
...
def cleanup(self) -> None:
"""Release all resources (buffers, readers, events)."""
...
class DiskWeightSource(WeightSource):
"""Reads block weights from disk into pinned CPU buffers on demand."""
def __init__(self, pool: WeightPool, reader: DiskBlockReader) -> None:
self._pool = pool
self._cache: OrderedDict[int, dict[str, torch.Tensor]] = OrderedDict()
self._events: dict[int, torch.cuda.Event] = {}
self._reader = reader
def get(self, idx: int) -> dict[str, torch.Tensor]:
"""Return CPU weights for block *idx*. Reads from disk on miss."""
if idx in self._cache:
return self._cache[idx]
if len(self._cache) >= self._pool.capacity:
evicted_idx, evicted_weights = self._cache.popitem(last=False)
self._pool.release(evicted_weights, event=self._events.pop(evicted_idx, None))
weights = self._pool.acquire()
self._reader.read_into(weights, idx)
self._cache[idx] = weights
return weights
def release(self, idx: int, event: torch.cuda.Event) -> None:
"""Attach an H2D event -- waited before this buffer is recycled."""
self._events[idx] = event
def cleanup(self) -> None:
"""Clear cache and close the disk reader."""
self._cache.clear()
self._events.clear()
self._reader.cleanup()
def __len__(self) -> int:
return len(self._cache)
class PinnedWeightSource(WeightSource):
"""Pre-loaded pinned CPU weights."""
def __init__(self, weights: dict[int, dict[str, torch.Tensor]]) -> None:
self._weights = weights
def get(self, idx: int) -> dict[str, torch.Tensor]:
return self._weights[idx]
def release(self, idx: int, event: torch.cuda.Event) -> None:
pass
def cleanup(self) -> None:
self._weights.clear()
def __len__(self) -> int:
return len(self._weights)
packages/ltx-core/src/ltx_core/block_streaming/utils.py
+60
View File
@@ -0,0 +1,60 @@
"""Shared utilities for the block_streaming package."""
from __future__ import annotations
import itertools
from typing import TYPE_CHECKING, Any
import torch
from torch import nn
if TYPE_CHECKING:
from ltx_core.block_streaming.pool import BlockLayout
def resolve_attr(module: nn.Module, dotted_path: str) -> nn.ModuleList:
"""Resolve a dotted attribute path like ``'model.language_model.layers'``."""
obj: Any = module
for part in dotted_path.split("."):
obj = getattr(obj, part)
if not isinstance(obj, nn.ModuleList):
raise TypeError(f"Expected nn.ModuleList at '{dotted_path}', got {type(obj).__name__}")
return obj
def assign_tensor_to_module(root: nn.Module, dotted_name: str, tensor: torch.Tensor) -> None:
"""Assign *tensor* to the parameter/buffer at *dotted_name* inside *root*.
Unlike ``param.data = tensor``, this works even when the existing parameter
lives on the ``meta`` device (which has an incompatible storage type).
"""
parts = dotted_name.split(".")
parent = root
for part in parts[:-1]:
parent = getattr(parent, part)
leaf = parts[-1]
if leaf in parent._parameters:
parent._parameters[leaf] = nn.Parameter(tensor, requires_grad=False)
elif leaf in parent._buffers:
parent._buffers[leaf] = tensor
else:
raise AttributeError(f"{leaf} is not a parameter or buffer of {type(parent).__name__}")
def 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.
"""
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()
}
packages/ltx-core/src/ltx_core/block_streaming/wrapper.py
+96
View File
@@ -0,0 +1,96 @@
"""Block streaming wrapper: streams transformer blocks through a WeightsProvider."""
from __future__ import annotations
import itertools
from typing import Any
import torch
from torch import nn
from ltx_core.block_streaming.provider import WeightsProvider
from ltx_core.block_streaming.utils import assign_tensor_to_module
class BlockStreamingWrapper(nn.Module):
"""Streams sequential model blocks through GPU buffer caches.
The wrapper delegates all weight management to a :class:`WeightsProvider`
which handles CPU-to-GPU copies, caching, LoRA fusion, and stream
synchronization internally.
Use :class:`StreamingModelBuilder` to construct this wrapper -- it
handles checkpoint parsing, source selection, and provider creation.
Args:
model: The wrapped model (non-block params already on GPU).
blocks: Sequential blocks to stream (``nn.ModuleList``).
provider: Provides GPU-ready weights on demand.
target_device: GPU device for compute.
"""
def __init__(
self,
model: nn.Module,
blocks: nn.ModuleList,
provider: WeightsProvider,
target_device: torch.device,
) -> None:
super().__init__()
self._model = model
self._blocks = blocks
self._target_device = target_device
self._provider = provider
self._hooks: list[torch.utils.hooks.RemovableHandle] = []
self._register_hooks()
# ------------------------------------------------------------------
# Hook registration
# ------------------------------------------------------------------
def _pre_hook(self, block_idx: int) -> None:
"""Load GPU weights for a block and inject them into its parameters."""
gpu_weights = self._provider.get(block_idx)
block = self._blocks[block_idx]
for name, _param in itertools.chain(block.named_parameters(), block.named_buffers()):
assign_tensor_to_module(block, name, gpu_weights[name])
def _post_hook(self, block_idx: int) -> None:
"""Record a compute-done event and release the block weights."""
compute_done = torch.cuda.Event()
compute_done.record(torch.cuda.current_stream(self._target_device))
self._provider.release(block_idx, event=compute_done)
def _register_hooks(self) -> None:
for idx, block in enumerate(self._blocks):
pre = block.register_forward_pre_hook(
lambda _mod, _args, *, idx=idx: self._pre_hook(idx),
)
post = block.register_forward_hook(
lambda _mod, _args, _out, *, idx=idx: self._post_hook(idx),
)
self._hooks.extend([pre, post])
# ------------------------------------------------------------------
# Teardown
# ------------------------------------------------------------------
def teardown(self) -> None:
"""Remove hooks and release all resources."""
for h in self._hooks:
h.remove()
self._hooks.clear()
self._provider.cleanup()
# ------------------------------------------------------------------
# Forward and attribute delegation
# ------------------------------------------------------------------
def forward(self, *args: Any, **kwargs: Any) -> Any: # noqa: ANN401
return self._model(*args, **kwargs)
def __getattr__(self, name: str) -> Any: # noqa: ANN401
"""Proxy attribute access to the wrapped model."""
try:
return super().__getattr__(name)
except AttributeError:
return getattr(self._model, name)
packages/ltx-core/src/ltx_core/conditioning/types/keyframe_cond.py
+14 -1
View File
@@ -17,12 +17,20 @@ class VideoConditionByKeyframeIndex(ConditioningItem):
keyframes: Keyframe latents [B, C, F, H, W].
frame_idx: Frame index offset for positional encoding.
strength: Conditioning strength (1.0 = clean, 0.0 = fully denoised).
num_pixel_frames: Number of pixel frames the keyframe latent originally encodes.
"""
def __init__(self, keyframes: torch.Tensor, frame_idx: int, strength: float):
def __init__(
self,
keyframes: torch.Tensor,
frame_idx: int,
strength: float,
num_pixel_frames: int = 1,
):
self.keyframes = keyframes
self.frame_idx = frame_idx
self.strength = strength
self.num_pixel_frames = num_pixel_frames
def apply_to(
self,
@@ -41,6 +49,11 @@ class VideoConditionByKeyframeIndex(ConditioningItem):
)
positions[:, 0, ...] += self.frame_idx
# If the keyframe latent encodes a single pixel frame,
# narrow the temporal end to [start, start + 1) instead of the
# VAE-scaled range.
if self.num_pixel_frames == 1:
positions[:, 0, ..., 1:] = positions[:, 0, ..., :1] + 1
positions = positions.to(dtype=torch.float32)
positions[:, 0, ...] /= latent_tools.fps
packages/ltx-core/src/ltx_core/hdr.py
+71
View File
@@ -0,0 +1,71 @@
"""HDR utilities: LogC3 compression for HDR IC-LoRA training and inference.
Provides compress/decompress and postprocess helpers for HDR video generation.
Used by ltx-pipelines for HDR IC-LoRA and by ltx-trainer for HDR validation.
"""
from __future__ import annotations
from typing import Literal
import torch
from torch import Tensor
class LogC3:
"""ARRI LogC3 (EI 800) HDR compression.
Maps linear [0, ∞) <-> LogC3 [0, 1] via the camera log curve. The log
curve allocates more precision to shadows/midtones and compresses
highlights smoothly. Callers are responsible for mapping the [0, 1]
output to the VAE's [-1, 1] input range.
"""
name = "LogC3"
A = 5.555556
B = 0.052272
C = 0.247190
D = 0.385537
E = 5.367655
F = 0.092809
CUT = 0.010591
def compress(self, hdr: Tensor) -> Tensor:
"""Compress linear HDR [0, ∞) → LogC3 [0, 1]."""
x = torch.clamp(hdr, min=0.0)
log_part = self.C * torch.log10(self.A * x + self.B) + self.D
lin_part = self.E * x + self.F
logc = torch.where(x >= self.CUT, log_part, lin_part)
return torch.clamp(logc, 0.0, 1.0)
def compress_ldr(self, ldr: Tensor) -> Tensor:
"""Compress LDR [0, 1] → [0, 1] (no log curve, just clamp)."""
return torch.clamp(ldr, 0.0, 1.0)
def decompress(self, logc: Tensor) -> Tensor:
"""Decompress LogC3 [0, 1] → linear HDR [0, ∞)."""
logc = torch.clamp(logc, 0.0, 1.0)
cut_log = self.E * self.CUT + self.F
lin_from_log = (torch.pow(10.0, (logc - self.D) / self.C) - self.B) / self.A
lin_from_lin = (logc - self.F) / self.E
return torch.where(logc >= cut_log, lin_from_log, lin_from_lin)
def decompress_ldr(self, logc: Tensor) -> Tensor:
"""Decompress [0, 1] → LDR [0, 1] (identity clamp)."""
return torch.clamp(logc, 0.0, 1.0)
def apply_hdr_decode_postprocess(
decoded_video: Tensor,
transform: Literal["logc3"] = "logc3",
) -> Tensor:
"""Apply HDR decompress to VAE decode output for HDR recovery.
Args:
decoded_video: Tensor from VAE decode in [0, 1], shape [B, C, F, H, W].
Must be float32 for sufficient color resolution.
transform: "logc3".
Returns:
HDR video tensor float32.
"""
decoded_video = decoded_video.float()
if transform == "logc3":
return LogC3().decompress(decoded_video)
raise ValueError(f"Unsupported HDR transform: {transform}")
packages/ltx-core/src/ltx_core/layer_streaming.py
-306
View File
@@ -1,306 +0,0 @@
"""Layer streaming wrapper for memory-efficient inference.
Keeps most transformer/decoder layers on CPU pinned memory and streams them
to GPU on demand, using a secondary CUDA stream to prefetch upcoming layers
so that data transfer overlaps with compute.
General-purpose: works with any ``nn.Module`` whose forward iterates over a
``nn.ModuleList`` attribute (e.g. ``transformer_blocks``, ``layers``).
Each layer is evicted back to CPU immediately after its forward completes,
and prefetch uses modular indexing so the last layer's prefetch wraps around
to prepare early layers for the next forward pass.
Example
-------
>>> model = build_my_model(device=torch.device("cpu"))
>>> model = LayerStreamingWrapper(
... model,
... layers_attr="transformer_blocks",
... target_device=torch.device("cuda:0"),
... prefetch_count=2,
... )
>>> out = model(inputs) # hooks handle layer streaming
>>> model.teardown() # move everything back to CPU
"""
from __future__ import annotations
import functools
import itertools
import logging
from typing import Any
import torch
from torch import nn
logger = logging.getLogger(__name__)
def _resolve_attr(module: nn.Module, dotted_path: str) -> nn.ModuleList:
"""Resolve a dotted attribute path like ``'model.language_model.layers'``."""
obj: Any = module
for part in dotted_path.split("."):
obj = getattr(obj, part)
if not isinstance(obj, nn.ModuleList):
raise TypeError(f"Expected nn.ModuleList at '{dotted_path}', got {type(obj).__name__}")
return obj
class _LayerStore:
"""Manages CPU-pinned copies of layer parameters/buffers.
Tracks which layers currently reside on GPU so the prefetcher and evictor
can make correct decisions.
"""
def __init__(self, layers: nn.ModuleList, target_device: torch.device) -> None:
self.target_device = target_device
self.num_layers = len(layers)
# CPU-pinned copies keyed by (layer_idx, param_name)
self._pinned: list[dict[str, torch.Tensor]] = []
self._on_gpu: set[int] = set()
for layer in layers:
pinned: dict[str, torch.Tensor] = {}
for name, tensor in itertools.chain(layer.named_parameters(), layer.named_buffers()):
pinned_tensor = tensor.data.pin_memory()
tensor.data = pinned_tensor
pinned[name] = pinned_tensor
self._pinned.append(pinned)
def _check_idx(self, idx: int) -> None:
if idx < 0 or idx >= self.num_layers:
raise IndexError(f"Layer index {idx} out of range [0, {self.num_layers})")
def is_on_gpu(self, idx: int) -> bool:
return idx in self._on_gpu
def move_to_gpu(self, idx: int, layer: nn.Module, *, non_blocking: bool = False) -> None:
"""Move layer *idx* parameters from pinned CPU to *target_device*."""
self._check_idx(idx)
if idx in self._on_gpu:
return
pinned = self._pinned[idx]
for name, param in itertools.chain(layer.named_parameters(), layer.named_buffers()):
param.data = pinned[name].to(self.target_device, non_blocking=non_blocking)
self._on_gpu.add(idx)
def evict_to_cpu(self, idx: int, layer: nn.Module) -> None:
"""Swap layer *idx* parameters back to their pinned CPU copies."""
self._check_idx(idx)
if idx not in self._on_gpu:
return
pinned = self._pinned[idx]
for name, param in itertools.chain(layer.named_parameters(), layer.named_buffers()):
param.data = pinned[name]
self._on_gpu.discard(idx)
def cleanup(self) -> None:
"""Release all pinned memory references.
After this call, the pinned tensors can be garbage-collected once
the layer parameters (which still reference them via ``.data``) are
also released (e.g. via ``.to("meta")``).
"""
for pinned_dict in self._pinned:
pinned_dict.clear()
self._pinned.clear()
class _AsyncPrefetcher:
"""Issues H2D transfers on a dedicated CUDA stream.
Uses per-layer CUDA events so that the compute stream only waits for the
specific layer it needs, not all pending transfers.
"""
def __init__(self, store: _LayerStore, layers: nn.ModuleList) -> None:
self._store = store
self._layers = layers
self._stream = torch.cuda.Stream(device=store.target_device)
self._events: dict[int, torch.cuda.Event] = {}
def prefetch(self, idx: int) -> None:
"""Begin async transfer of layer *idx* to GPU (no-op if already there)."""
if self._store.is_on_gpu(idx) or idx in self._events:
return
with torch.cuda.stream(self._stream):
self._store.move_to_gpu(idx, self._layers[idx], non_blocking=True)
event = torch.cuda.Event()
event.record(self._stream)
self._events[idx] = event
def wait(self, idx: int) -> None:
"""Block the compute stream until layer *idx* transfer is complete."""
event = self._events.pop(idx, None)
if event is not None:
torch.cuda.current_stream(self._store.target_device).wait_event(event)
def cleanup(self) -> None:
"""Drain pending work and release CUDA stream/event resources."""
self._events.clear()
self._stream = None
self._layers = None
self._store = None
class LayerStreamingWrapper(nn.Module):
"""Wraps a model to stream its sequential layers between CPU and GPU.
Each layer is evicted immediately after its forward completes, and
prefetch wraps around using modular indexing so the end of one forward
pass prepares early layers for the next.
Parameters
----------
model:
The model to wrap, with all parameters on **CPU**.
layers_attr:
Dotted attribute path to the ``nn.ModuleList`` of sequential layers
(e.g. ``"transformer_blocks"`` or ``"model.language_model.layers"``).
target_device:
The GPU device to use for compute.
prefetch_count:
How many layers ahead to prefetch. The maximum number of layers on
GPU at once is ``1 + prefetch_count``. Must be >= 1.
"""
def __init__(
self,
model: nn.Module,
layers_attr: str,
target_device: torch.device,
prefetch_count: int = 2,
) -> None:
if prefetch_count < 1:
raise ValueError("prefetch_count must be >= 1")
super().__init__()
# Store the wrapped model as a submodule so parameters are discoverable.
self._model = model
self._layers = _resolve_attr(model, layers_attr)
self._target_device = target_device
# Clamp: no point prefetching more than num_layers - 1 (the rest are evicted).
self._prefetch_count = min(prefetch_count, len(self._layers) - 1)
self._hooks: list[torch.utils.hooks.RemovableHandle] = []
self._setup()
# ------------------------------------------------------------------
# Setup / teardown
# ------------------------------------------------------------------
def _setup(self) -> None:
# 1. Build the pinned CPU store (copies all layer tensors to pinned memory).
self._store = _LayerStore(self._layers, self._target_device)
# 2. Move all NON-layer params/buffers to GPU.
layer_tensor_ids: set[int] = set()
for layer in self._layers:
for t in itertools.chain(layer.parameters(), layer.buffers()):
layer_tensor_ids.add(id(t))
for p in self._model.parameters():
if id(p) not in layer_tensor_ids:
p.data = p.data.to(self._target_device)
for b in self._model.buffers():
if id(b) not in layer_tensor_ids:
b.data = b.data.to(self._target_device)
# 3. Pre-load the first (1 + prefetch_count) layers synchronously.
for idx in range(min(self._prefetch_count + 1, len(self._layers))):
self._store.move_to_gpu(idx, self._layers[idx])
# 4. Create the async prefetcher and register hooks.
self._prefetcher = _AsyncPrefetcher(self._store, self._layers)
self._register_hooks()
def _register_hooks(self) -> None:
idx_map: dict[int, int] = {id(layer): idx for idx, layer in enumerate(self._layers)}
num_layers = len(self._layers)
def _pre_hook(
module: nn.Module,
_args: Any, # noqa: ANN401
*,
idx: int,
) -> None:
# Wait only for THIS layer's H2D transfer (not all pending ones).
self._prefetcher.wait(idx)
if not self._store.is_on_gpu(idx):
self._store.move_to_gpu(idx, module)
# Record that the compute stream will read these weight tensors.
# They were allocated on the prefetch stream, so without this the
# caching allocator would allow the prefetch stream to reuse their
# memory immediately after eviction — even if the compute kernel
# that reads them hasn't finished yet.
compute_stream = torch.cuda.current_stream(self._target_device)
for param in itertools.chain(module.parameters(), module.buffers()):
param.data.record_stream(compute_stream)
# Kick off prefetch for upcoming layers (wraps around for next pass).
for offset in range(1, self._prefetch_count + 1):
self._prefetcher.prefetch((idx + offset) % num_layers)
def _post_hook(
module: nn.Module,
_args: Any, # noqa: ANN401
_output: Any, # noqa: ANN401
*,
idx: int,
) -> None:
# Evict this layer immediately — its computation is done.
self._store.evict_to_cpu(idx, module)
for layer in self._layers:
idx = idx_map[id(layer)]
h1 = layer.register_forward_pre_hook(functools.partial(_pre_hook, idx=idx))
h2 = layer.register_forward_hook(functools.partial(_post_hook, idx=idx))
self._hooks.extend([h1, h2])
def teardown(self) -> None:
"""Remove hooks, release pinned memory, and move parameters back to CPU.
After this call the wrapper is inert: hooks are removed, the prefetch
stream is drained and destroyed, all parameters reside on regular
(non-pinned) CPU memory, and the ``_LayerStore`` pinned-tensor cache is
cleared. Callers should still follow up with ``.to("meta")`` to release
the CPU copies if the model is no longer needed.
"""
for h in self._hooks:
h.remove()
self._hooks.clear()
# Drain all in-flight async H2D copies, then release stream resources.
# Without the synchronize, clearing the stream/events can trigger
# use-after-free at the CUDA driver level.
torch.cuda.synchronize(device=self._target_device)
if self._prefetcher is not None:
self._prefetcher.cleanup()
self._prefetcher = None
# Move everything to CPU.
for idx, layer in enumerate(self._layers):
self._store.evict_to_cpu(idx, layer)
for p in self._model.parameters():
p.data = p.data.to("cpu")
for b in self._model.buffers():
b.data = b.data.to("cpu")
# Release pinned memory. After evict_to_cpu() the layer parameters
# still reference the pinned tensors (since .to("cpu") on a pinned
# tensor is a no-op). The caller is expected to follow up with
# .to("meta") to drop the param refs; cleanup() drops the store's refs.
self._store.cleanup()
# ------------------------------------------------------------------
# Forward and attribute delegation
# ------------------------------------------------------------------
def forward(self, *args: Any, **kwargs: Any) -> Any: # noqa: ANN401
return self._model(*args, **kwargs)
def __getattr__(self, name: str) -> Any: # noqa: ANN401
"""Proxy attribute access to the wrapped model.
This allows calling methods like ``encode()`` on a wrapped
GemmaTextEncoder without the caller needing to know about the wrapper.
``nn.Module.__getattr__`` is only called when normal attribute lookup
fails, so ``_model``, ``_store``, etc. are found first via ``__dict__``.
"""
try:
return super().__getattr__(name)
except AttributeError:
return getattr(self._model, name)
packages/ltx-core/src/ltx_core/loader/__init__.py
+8
View File
@@ -1,6 +1,11 @@
"""Loader utilities for model weights, LoRAs, and safetensor operations."""
from ltx_core.loader.fuse_loras import apply_loras
from ltx_core.loader.helpers import (
create_meta_model,
load_state_dict,
read_model_config,
)
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.primitives import (
LoRAAdaptableProtocol,
@@ -45,4 +50,7 @@ __all__ = [
"StateDictLoader",
"StateDictRegistry",
"apply_loras",
"create_meta_model",
"load_state_dict",
"read_model_config",
]
packages/ltx-core/src/ltx_core/loader/helpers.py
+61
View File
@@ -0,0 +1,61 @@
"""Shared model-construction helpers used by both SingleGPUModelBuilder and StreamingModelBuilder."""
from __future__ import annotations
from typing import TypeVar
import torch
from torch import nn
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.primitives import StateDict, StateDictLoader
from ltx_core.loader.registry import Registry
from ltx_core.loader.sd_ops import SDOps
from ltx_core.model.model_protocol import ModelConfigurator
_M = TypeVar("_M", bound=nn.Module)
def load_state_dict(
paths: str | tuple[str, ...] | list[str],
loader: StateDictLoader,
registry: Registry,
device: torch.device | None,
sd_ops: SDOps | None = None,
) -> StateDict:
"""Load a state dict from disk, using registry caching."""
if isinstance(paths, str):
path_list = [paths]
elif isinstance(paths, tuple):
path_list = list(paths)
else:
path_list = paths
cached = registry.get(path_list, sd_ops)
if cached is not None:
return cached
result = loader.load(path_list, sd_ops=sd_ops, device=device)
registry.add(path_list, sd_ops=sd_ops, state_dict=result)
return result
def read_model_config(
model_path: str | tuple[str, ...],
loader: StateDictLoader,
) -> dict:
"""Read metadata from the first shard of a checkpoint."""
first = model_path[0] if isinstance(model_path, tuple) else model_path
return loader.metadata(first)
def create_meta_model(
configurator: type[ModelConfigurator[_M]],
config: dict,
module_ops: tuple[ModuleOps, ...] = (),
) -> _M:
"""Create a model on the meta device and apply module operations."""
with torch.device("meta"):
model = configurator.from_config(config)
for op in module_ops:
if op.matcher(model):
model = op.mutator(model)
return model
packages/ltx-core/src/ltx_core/loader/single_gpu_model_builder.py
+68 -42
View File
@@ -3,8 +3,10 @@ from dataclasses import dataclass, field, replace
from typing import Generic
import torch
from torch import nn
from ltx_core.loader.fuse_loras import apply_loras
from ltx_core.loader.helpers import create_meta_model, load_state_dict, read_model_config
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.primitives import (
LoRAAdaptableProtocol,
@@ -22,6 +24,56 @@ from ltx_core.model.model_protocol import ModelConfigurator, ModelType
logger: logging.Logger = logging.getLogger(__name__)
def _check_uninitialized(model: nn.Module) -> list[str]:
"""Return names of any parameters/buffers still on meta device."""
names = []
for name, param in model.named_parameters():
if str(param.device) == "meta":
names.append(name)
for name, buf in model.named_buffers():
if str(buf.device) == "meta":
names.append(name)
return names
def _load_model_weights(
meta_model: nn.Module,
model_path: str | tuple[str, ...],
loras: tuple[LoraPathStrengthAndSDOps, ...],
loader: StateDictLoader,
registry: Registry,
device: torch.device,
dtype: torch.dtype | None,
model_sd_ops: SDOps | None = None,
lora_load_device: torch.device | None = None,
) -> None:
"""Load base weights and fuse LoRAs into *meta_model* in-place."""
if lora_load_device is None:
lora_load_device = device
model_sd = load_state_dict(model_path, loader, registry, device, model_sd_ops)
lora_strengths = [lora.strength for lora in loras]
if not lora_strengths or (min(lora_strengths) == 0 and max(lora_strengths) == 0):
sd = model_sd.sd
if dtype is not None:
sd = {key: value.to(dtype=dtype) for key, value in model_sd.sd.items()}
meta_model.load_state_dict(sd, strict=False, assign=True)
return
lora_state_dicts = [load_state_dict([lora.path], loader, registry, lora_load_device, lora.sd_ops) for lora in loras]
lora_sd_and_strengths = [
LoraStateDictWithStrength(sd, strength) for sd, strength in zip(lora_state_dicts, lora_strengths, strict=True)
]
final_sd = apply_loras(
model_sd=model_sd,
lora_sd_and_strengths=lora_sd_and_strengths,
dtype=dtype,
destination_sd=model_sd if isinstance(registry, DummyRegistry) else None,
)
meta_model.load_state_dict(final_sd.sd, strict=False, assign=True)
@dataclass(frozen=True)
class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType], LoRAAdaptableProtocol):
"""
@@ -69,34 +121,22 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
return replace(self, lora_load_device=device)
def model_config(self) -> dict:
first_shard_path = self.model_path[0] if isinstance(self.model_path, tuple) else self.model_path
return self.model_loader.metadata(first_shard_path)
return read_model_config(self.model_path, self.model_loader)
def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType:
with torch.device("meta"):
model = self.model_class_configurator.from_config(config)
for module_op in module_ops:
if module_op.matcher(model):
model = module_op.mutator(model)
return model
return create_meta_model(self.model_class_configurator, config, module_ops)
def load_sd(
self, paths: list[str], registry: Registry, device: torch.device | None, sd_ops: SDOps | None = None
) -> StateDict:
state_dict = registry.get(paths, sd_ops)
if state_dict is None:
state_dict = self.model_loader.load(paths, sd_ops=sd_ops, device=device)
registry.add(paths, sd_ops=sd_ops, state_dict=state_dict)
return state_dict
return load_state_dict(paths, self.model_loader, registry, device, sd_ops)
def _return_model(self, meta_model: ModelType, device: torch.device) -> ModelType:
uninitialized_params = [name for name, param in meta_model.named_parameters() if str(param.device) == "meta"]
uninitialized_buffers = [name for name, buffer in meta_model.named_buffers() if str(buffer.device) == "meta"]
if uninitialized_params or uninitialized_buffers:
logger.warning(f"Uninitialized parameters or buffers: {uninitialized_params + uninitialized_buffers}")
uninitialized = _check_uninitialized(meta_model)
if uninitialized:
logger.warning(f"Uninitialized parameters or buffers: {uninitialized}")
return meta_model
retval = meta_model.to(device)
return retval
return meta_model.to(device)
def build(
self,
@@ -107,30 +147,16 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
device = torch.device("cuda") if device is None else device
config = self.model_config()
meta_model = self.meta_model(config, self.module_ops)
model_paths = list(self.model_path) if isinstance(self.model_path, tuple) else [self.model_path]
model_state_dict = self.load_sd(model_paths, sd_ops=self.model_sd_ops, registry=self.registry, device=device)
lora_strengths = [lora.strength for lora in self.loras]
if not lora_strengths or (min(lora_strengths) == 0 and max(lora_strengths) == 0):
sd = model_state_dict.sd
if dtype is not None:
sd = {key: value.to(dtype=dtype) for key, value in model_state_dict.sd.items()}
meta_model.load_state_dict(sd, strict=False, assign=True)
return self._return_model(meta_model, device)
lora_state_dicts = [
self.load_sd([lora.path], sd_ops=lora.sd_ops, registry=self.registry, device=self.lora_load_device)
for lora in self.loras
]
lora_sd_and_strengths = [
LoraStateDictWithStrength(sd, strength)
for sd, strength in zip(lora_state_dicts, lora_strengths, strict=True)
]
final_sd = apply_loras(
model_sd=model_state_dict,
lora_sd_and_strengths=lora_sd_and_strengths,
_load_model_weights(
meta_model=meta_model,
model_path=self.model_path,
loras=self.loras,
loader=self.model_loader,
registry=self.registry,
device=device,
dtype=dtype,
destination_sd=model_state_dict if isinstance(self.registry, DummyRegistry) else None,
model_sd_ops=self.model_sd_ops,
lora_load_device=self.lora_load_device,
)
meta_model.load_state_dict(final_sd.sd, strict=False, assign=True)
return self._return_model(meta_model, device)
packages/ltx-core/src/ltx_core/model/video_vae/video_vae.py
+22 -8
View File
@@ -698,6 +698,9 @@ class VideoDecoder(nn.Module):
When causal=False, allows future frame dependencies in convolutions but maintains same output shape.
"""
batch_size = sample.shape[0]
output_dtype = sample.dtype
weights_dtype = next(self.parameters()).dtype
sample = sample.to(weights_dtype)
# Add noise if timestep conditioning is enabled
if self.timestep_conditioning:
@@ -770,7 +773,7 @@ class VideoDecoder(nn.Module):
# Example: (B, 48, F, 128, 128) -> (B, 3, F, 512, 512) with patch_size=4
sample = unpatchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)
return sample
return sample.to(output_dtype)
def _prepare_tiles(
self,
@@ -902,23 +905,34 @@ 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 uint8 chunks ``[f, h, w, c]``.
"""Decode a video latent tensor, yielding chunks ``[f, h, w, c]``.
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_to_uint8(frames: torch.Tensor) -> torch.Tensor:
frames = (((frames + 1.0) / 2.0).clamp(0.0, 1.0) * 255.0).to(torch.uint8)
frames = rearrange(frames[0], "c f h w -> f h w c")
return frames
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.
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)
if tiling_config is not None:
for frames in self.tiled_decode(latent, tiling_config, generator=generator):
yield convert_to_uint8(frames)
yield _convert(frames)
else:
decoded = self(latent, generator=generator)
yield convert_to_uint8(decoded)
yield _convert(decoded)
def _group_tiles_by_temporal_slice(self, tiles: List[Tile]) -> List[List[Tile]]:
"""Group tiles by their temporal output slice."""