Automated PR - 2026-03-30
This commit is contained in:
github-actions[bot]
@@ -0,0 +1,95 @@
|
||||
"""Batch-splitting adapter for the transformer.
|
||||
Wraps an ``X0Model`` (or ``LayerStreamingWrapper``) 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
|
||||
``X0Model`` and proxies attribute access to the wrapped model.
|
||||
Example
|
||||
-------
|
||||
>>> from ltx_core.batch_split import BatchSplitAdapter
|
||||
>>> adapter = BatchSplitAdapter(model, max_batch_size=1)
|
||||
>>> # Receives B=4, runs 4xB=1 internally, returns B=4
|
||||
>>> denoised_video, denoised_audio = adapter(video=v_b4, audio=a_b4, perturbations=ptb)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from ltx_core.guidance.perturbations import BatchedPerturbationConfig
|
||||
from ltx_core.model.transformer.modality import Modality
|
||||
|
||||
|
||||
def _split_perturbations(config: BatchedPerturbationConfig, sizes: list[int]) -> list[BatchedPerturbationConfig]:
|
||||
"""Split a ``BatchedPerturbationConfig`` along the batch dimension."""
|
||||
it = iter(config.perturbations)
|
||||
return [BatchedPerturbationConfig([next(it) for _ in range(s)]) for s in sizes]
|
||||
|
||||
|
||||
def _merge_tensors(tensors: list[torch.Tensor | None]) -> torch.Tensor | None:
|
||||
"""Concatenate tensors along batch dim, or return None if all are None."""
|
||||
non_none = [t for t in tensors if t is not None]
|
||||
if not non_none:
|
||||
return None
|
||||
return torch.cat(non_none, dim=0)
|
||||
|
||||
|
||||
class BatchSplitAdapter(nn.Module):
|
||||
"""Wraps a model and splits batched forward calls into smaller chunks.
|
||||
Has the same ``forward`` signature as ``X0Model``:
|
||||
``(video, audio, perturbations) -> (denoised_video, denoised_audio)``.
|
||||
Args:
|
||||
model: The model to wrap (``X0Model``, ``LayerStreamingWrapper``, etc.).
|
||||
max_batch_size: Maximum batch size per forward pass. Input batches
|
||||
larger than this are split into sequential chunks.
|
||||
"""
|
||||
|
||||
def __init__(self, model: nn.Module, max_batch_size: int) -> None:
|
||||
if max_batch_size < 1:
|
||||
raise ValueError(f"max_batch_size must be >= 1, got {max_batch_size}")
|
||||
super().__init__()
|
||||
self._model = model
|
||||
self._max_batch_size = max_batch_size
|
||||
|
||||
def _get_chunk_sizes(self, batch_size: int) -> list[int]:
|
||||
full, remainder = divmod(batch_size, self._max_batch_size)
|
||||
sizes = [self._max_batch_size] * full
|
||||
if remainder:
|
||||
sizes.append(remainder)
|
||||
return sizes
|
||||
|
||||
def forward(
|
||||
self,
|
||||
video: Modality | None,
|
||||
audio: Modality | None,
|
||||
perturbations: BatchedPerturbationConfig,
|
||||
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
|
||||
batch_size = (video or audio).latent.shape[0]
|
||||
|
||||
if batch_size <= self._max_batch_size:
|
||||
return self._model(video=video, audio=audio, perturbations=perturbations)
|
||||
|
||||
sizes = self._get_chunk_sizes(batch_size)
|
||||
n = len(sizes)
|
||||
|
||||
v_chunks = video.split(sizes) if video is not None else [None] * n
|
||||
a_chunks = audio.split(sizes) if audio is not None else [None] * n
|
||||
p_chunks = _split_perturbations(perturbations, sizes)
|
||||
|
||||
chunk_results = [
|
||||
self._model(video=vc, audio=ac, perturbations=pc)
|
||||
for vc, ac, pc in zip(v_chunks, a_chunks, p_chunks, strict=True)
|
||||
]
|
||||
|
||||
results_v, results_a = zip(*chunk_results, strict=True)
|
||||
return _merge_tensors(list(results_v)), _merge_tensors(list(results_a))
|
||||
|
||||
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)
|
||||
@@ -77,11 +77,22 @@ class Res2sDiffusionStep(DiffusionStepProtocol):
|
||||
sigmas: torch.Tensor,
|
||||
step_index: int,
|
||||
noise: torch.Tensor,
|
||||
eta: float = 0.5,
|
||||
) -> torch.Tensor:
|
||||
"""Advance one step with SDE noise injection via get_sde_coeff."""
|
||||
"""Advance one step with SDE noise injection via get_sde_coeff.
|
||||
Args:
|
||||
sample: Current noisy sample.
|
||||
denoised_sample: Denoised prediction from the model.
|
||||
sigmas: Noise schedule tensor.
|
||||
step_index: Current step index in the schedule.
|
||||
noise: Random noise tensor for stochastic injection.
|
||||
eta: Controls stochastic noise injection strength (0=deterministic, 1=maximum). Default 0.5.
|
||||
Returns:
|
||||
Next sample with SDE noise injection applied.
|
||||
"""
|
||||
sigma = sigmas[step_index]
|
||||
sigma_next = sigmas[step_index + 1]
|
||||
alpha_ratio, sigma_down, sigma_up = self.get_sde_coeff(sigma_next, sigma_up=sigma_next * 0.5)
|
||||
alpha_ratio, sigma_down, sigma_up = self.get_sde_coeff(sigma_next, sigma_up=sigma_next * eta)
|
||||
output_dtype = denoised_sample.dtype
|
||||
if torch.any(sigma_up == 0) or torch.any(sigma_next == 0):
|
||||
return denoised_sample
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
from ltx_core.components.patchifiers import get_pixel_coords
|
||||
from ltx_core.conditioning.item import ConditioningItem
|
||||
from ltx_core.tools import LatentTools, SpatioTemporalScaleFactors
|
||||
from ltx_core.types import AudioLatentShape, LatentState, VideoLatentShape
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TemporalRegionMask(ConditioningItem):
|
||||
"""Conditioning item that sets ``denoise_mask = 0`` outside a time range
|
||||
and ``1`` inside, so only the specified temporal region is regenerated.
|
||||
Uses ``start_time`` and ``end_time`` in seconds. Works in *patchified*
|
||||
(token) space using the patchifier's ``get_patch_grid_bounds``: for video
|
||||
coords are latent frame indices (converted from seconds via ``fps``), for
|
||||
audio coords are already in seconds.
|
||||
"""
|
||||
|
||||
start_time: float # seconds, inclusive
|
||||
end_time: float # seconds, exclusive
|
||||
fps: float
|
||||
|
||||
def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
|
||||
coords = latent_tools.patchifier.get_patch_grid_bounds(
|
||||
latent_tools.target_shape, device=latent_state.denoise_mask.device
|
||||
)
|
||||
if isinstance(latent_tools.target_shape, AudioLatentShape):
|
||||
# Audio: patchifier get_patch_grid_bounds returns seconds
|
||||
t_boundaries = coords[:, 0]
|
||||
elif isinstance(latent_tools.target_shape, VideoLatentShape):
|
||||
# Video: patchifier get_patch_grid_bounds returns latent bounds, converting to frame numbers & pixel bounds
|
||||
scale_factors = getattr(latent_tools, "scale_factors", SpatioTemporalScaleFactors.default())
|
||||
pixel_bounds = get_pixel_coords(coords, scale_factors, causal_fix=getattr(latent_tools, "causal_fix", True))
|
||||
# converting frame numbers to seconds
|
||||
t_boundaries = pixel_bounds[:, 0] / self.fps
|
||||
else:
|
||||
raise ValueError("Unsupported LatentShape type, expected AudioLatentShape or VideoLatentShape")
|
||||
t_start, t_end = t_boundaries.unbind(dim=-1) # [B, N]
|
||||
in_region = (t_end > self.start_time) & (t_start < self.end_time)
|
||||
state = latent_state.clone()
|
||||
mask_val = in_region.to(state.denoise_mask.dtype)
|
||||
if state.denoise_mask.dim() == 3:
|
||||
mask_val = mask_val.unsqueeze(-1)
|
||||
state.denoise_mask.copy_(mask_val)
|
||||
return state
|
||||
@@ -0,0 +1,324 @@
|
||||
"""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 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.
|
||||
"""
|
||||
|
||||
def __init__(self, layers: nn.ModuleList, target_device: torch.device) -> None:
|
||||
self.target_device = target_device
|
||||
self.num_layers = len(layers)
|
||||
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] = {}
|
||||
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]] = {}
|
||||
|
||||
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:
|
||||
"""Pin layer *idx* on demand, then transfer to GPU."""
|
||||
self._check_idx(idx)
|
||||
if idx in self._on_gpu:
|
||||
return
|
||||
source = self._source_data[idx]
|
||||
pinned_refs: list[torch.Tensor] = []
|
||||
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
|
||||
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."""
|
||||
self._check_idx(idx)
|
||||
if idx not in self._on_gpu:
|
||||
return
|
||||
source = self._source_data[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)
|
||||
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
|
||||
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()
|
||||
|
||||
|
||||
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)
|
||||
|
||||
compute_stream = torch.cuda.current_stream(self._target_device)
|
||||
|
||||
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.
|
||||
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 resources, 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.
|
||||
"""
|
||||
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 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.
|
||||
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)
|
||||
@@ -1,46 +1,71 @@
|
||||
from collections.abc import Iterator
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict
|
||||
from ltx_core.quantization.fp8_cast import calculate_weight_float8
|
||||
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
|
||||
|
||||
|
||||
def _get_device() -> torch.device:
|
||||
if torch.cuda.is_available():
|
||||
return torch.device("cuda", torch.cuda.current_device())
|
||||
return torch.device("cpu")
|
||||
|
||||
|
||||
def fuse_lora_weights(
|
||||
model_sd: StateDict,
|
||||
lora_sd_and_strengths: list[LoraStateDictWithStrength],
|
||||
dtype: torch.dtype | None = None,
|
||||
) -> 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.
|
||||
"""
|
||||
for key, original_weight in model_sd.sd.items():
|
||||
if original_weight is None or key.endswith(".weight_scale"):
|
||||
continue
|
||||
original_device = original_weight.device
|
||||
weight = original_weight.to(device=_get_device())
|
||||
target_dtype = dtype if dtype is not None else weight.dtype
|
||||
deltas_dtype = target_dtype if target_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2] else torch.bfloat16
|
||||
|
||||
deltas = _prepare_deltas(lora_sd_and_strengths, key, deltas_dtype, weight.device)
|
||||
if deltas is None:
|
||||
continue
|
||||
|
||||
scale_key = key.replace(".weight", ".weight_scale") if key.endswith(".weight") else None
|
||||
is_scaled_fp8 = scale_key is not None and scale_key in model_sd.sd
|
||||
|
||||
if weight.dtype == torch.float8_e4m3fn:
|
||||
if is_scaled_fp8:
|
||||
fused = _fuse_delta_with_scaled_fp8(deltas, weight, key, scale_key, model_sd)
|
||||
else:
|
||||
fused = _fuse_delta_with_cast_fp8(deltas, weight, key, target_dtype)
|
||||
elif weight.dtype == torch.bfloat16:
|
||||
fused = _fuse_delta_with_bfloat16(deltas, weight, key, target_dtype)
|
||||
else:
|
||||
raise ValueError(f"Unsupported dtype: {weight.dtype}")
|
||||
|
||||
for k, v in fused.items():
|
||||
yield k, v.to(device=original_device)
|
||||
|
||||
|
||||
def apply_loras(
|
||||
model_sd: StateDict,
|
||||
lora_sd_and_strengths: list[LoraStateDictWithStrength],
|
||||
dtype: torch.dtype | None = None,
|
||||
destination_sd: StateDict | None = None,
|
||||
) -> StateDict:
|
||||
sd = {}
|
||||
if destination_sd is not None:
|
||||
sd = destination_sd.sd
|
||||
size = 0
|
||||
device = torch.device("meta")
|
||||
inner_dtypes = set()
|
||||
for key, weight in model_sd.sd.items():
|
||||
if weight is None:
|
||||
continue
|
||||
# Skip scale keys - they are handled together with their weight keys
|
||||
if key.endswith(".weight_scale"):
|
||||
continue
|
||||
device = weight.device
|
||||
target_dtype = dtype if dtype is not None else weight.dtype
|
||||
deltas_dtype = target_dtype if target_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2] else torch.bfloat16
|
||||
|
||||
scale_key = key.replace(".weight", ".weight_scale") if key.endswith(".weight") else None
|
||||
is_scaled_fp8 = scale_key is not None and scale_key in model_sd.sd
|
||||
|
||||
deltas = _prepare_deltas(lora_sd_and_strengths, key, deltas_dtype, device)
|
||||
fused = _fuse_deltas(deltas, weight, key, sd, target_dtype, device, is_scaled_fp8, scale_key, model_sd)
|
||||
|
||||
sd.update(fused)
|
||||
for tensor in fused.values():
|
||||
inner_dtypes.add(tensor.dtype)
|
||||
size += tensor.nbytes
|
||||
|
||||
if destination_sd is not None:
|
||||
for key, tensor in fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype):
|
||||
sd[key] = tensor
|
||||
return destination_sd
|
||||
return StateDict(sd, device, size, inner_dtypes)
|
||||
|
||||
fused = dict(fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype))
|
||||
sd = {k: (fused[k] if k in fused else v.clone()) for k, v in model_sd.sd.items()}
|
||||
return StateDict(sd, model_sd.device, model_sd.size, model_sd.dtype)
|
||||
|
||||
|
||||
def _prepare_deltas(
|
||||
@@ -65,50 +90,6 @@ def _prepare_deltas(
|
||||
return torch.sum(torch.stack(deltas, dim=0), dim=0)
|
||||
|
||||
|
||||
def _fuse_deltas(
|
||||
deltas: torch.Tensor | None,
|
||||
weight: torch.Tensor,
|
||||
key: str,
|
||||
sd: dict[str, torch.Tensor],
|
||||
target_dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
is_scaled_fp8: bool,
|
||||
scale_key: str | None,
|
||||
model_sd: StateDict,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
if deltas is None:
|
||||
if key in sd:
|
||||
return {}
|
||||
fused = _copy_weight_without_lora(weight, key, target_dtype, device, is_scaled_fp8, scale_key, model_sd)
|
||||
elif weight.dtype == torch.float8_e4m3fn:
|
||||
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, device)
|
||||
elif weight.dtype == torch.bfloat16:
|
||||
fused = _fuse_delta_with_bfloat16(deltas, weight, key, target_dtype)
|
||||
else:
|
||||
raise ValueError(f"Unsupported dtype: {weight.dtype}")
|
||||
|
||||
return fused
|
||||
|
||||
|
||||
def _copy_weight_without_lora(
|
||||
weight: torch.Tensor,
|
||||
key: str,
|
||||
target_dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
is_scaled_fp8: bool,
|
||||
scale_key: str | None,
|
||||
model_sd: StateDict,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Copy original weight (and scale if applicable) when no LoRA affects this key."""
|
||||
result = {key: weight.clone().to(dtype=target_dtype, device=device)}
|
||||
if is_scaled_fp8:
|
||||
result[scale_key] = model_sd.sd[scale_key].clone()
|
||||
return result
|
||||
|
||||
|
||||
def _fuse_delta_with_scaled_fp8(
|
||||
deltas: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
@@ -132,13 +113,12 @@ def _fuse_delta_with_cast_fp8(
|
||||
weight: torch.Tensor,
|
||||
key: str,
|
||||
target_dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Fuse LoRA delta with cast-only FP8 weight (no scale factor)."""
|
||||
if str(device).startswith("cuda"):
|
||||
deltas = calculate_weight_float8(deltas, weight)
|
||||
if str(weight.device).startswith("cuda"):
|
||||
_fused_add_round_launch(deltas, weight, seed=0)
|
||||
else:
|
||||
deltas.add_(weight.to(dtype=deltas.dtype, device=device))
|
||||
deltas.add_(weight.to(dtype=deltas.dtype))
|
||||
return {key: deltas.to(dtype=target_dtype)}
|
||||
|
||||
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import NamedTuple, Protocol
|
||||
from typing import TYPE_CHECKING, NamedTuple, Protocol
|
||||
|
||||
import torch
|
||||
|
||||
@@ -7,6 +9,9 @@ from ltx_core.loader.module_ops import ModuleOps
|
||||
from ltx_core.loader.sd_ops import SDOps
|
||||
from ltx_core.model.model_protocol import ModelType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ltx_core.loader.registry import Registry
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StateDict:
|
||||
@@ -55,6 +60,11 @@ class ModelBuilderProtocol(Protocol[ModelType]):
|
||||
- build: Create and initialize a model from state dictionary and apply dtype transformations
|
||||
"""
|
||||
|
||||
model_sd_ops: SDOps | None
|
||||
module_ops: tuple[ModuleOps, ...]
|
||||
loras: tuple["LoraPathStrengthAndSDOps", ...]
|
||||
registry: "Registry"
|
||||
|
||||
def meta_model(self, config: dict, module_ops: list[ModuleOps] | None = None) -> ModelType:
|
||||
"""
|
||||
Create a model on the meta device from a configuration dictionary.
|
||||
@@ -68,16 +78,43 @@ class ModelBuilderProtocol(Protocol[ModelType]):
|
||||
"""
|
||||
...
|
||||
|
||||
def build(self, dtype: torch.dtype | None = None) -> ModelType:
|
||||
def with_sd_ops(self, sd_ops: SDOps | None) -> "ModelBuilderProtocol[ModelType]":
|
||||
"""Return a copy of this builder with the given state-dict key remapping ops."""
|
||||
...
|
||||
|
||||
def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> "ModelBuilderProtocol[ModelType]":
|
||||
"""Return a copy of this builder with the given module operations (e.g. quantization)."""
|
||||
...
|
||||
|
||||
def with_loras(self, loras: tuple["LoraPathStrengthAndSDOps", ...]) -> "ModelBuilderProtocol[ModelType]":
|
||||
"""Return a copy of this builder with the given LoRAs to fuse at build time."""
|
||||
...
|
||||
|
||||
def with_registry(self, registry: "Registry") -> "ModelBuilderProtocol[ModelType]":
|
||||
"""Return a copy of this builder using the given weight registry for allocation."""
|
||||
...
|
||||
|
||||
def with_lora_load_device(self, device: torch.device) -> "ModelBuilderProtocol[ModelType]":
|
||||
"""Return a copy of this builder that loads LoRA weights onto the given device."""
|
||||
...
|
||||
|
||||
def build(
|
||||
self, device: torch.device | None = None, dtype: torch.dtype | None = None, **kwargs: object
|
||||
) -> ModelType:
|
||||
"""
|
||||
Build the model
|
||||
Args:
|
||||
device: Target device for the model
|
||||
dtype: Target dtype for the model, if None, uses the dtype of the model_path model
|
||||
Returns:
|
||||
Model instance
|
||||
"""
|
||||
...
|
||||
|
||||
def model_config(self) -> dict:
|
||||
"""Return the model configuration dictionary extracted from the checkpoint metadata."""
|
||||
...
|
||||
|
||||
|
||||
class LoRAAdaptableProtocol(Protocol):
|
||||
"""
|
||||
|
||||
@@ -64,6 +64,7 @@ class SDOps:
|
||||
mapping: tuple[
|
||||
ContentReplacement | ContentMatching | SDKeyValueOperation, ...
|
||||
] = () # Immutable tuple of (key, value) pairs
|
||||
allowed_keys: frozenset[str] | None = None
|
||||
|
||||
def with_replacement(self, content: str, replacement: str) -> "SDOps":
|
||||
"""Create a new SDOps instance with the specified replacement added to the mapping."""
|
||||
@@ -77,6 +78,13 @@ class SDOps:
|
||||
new_mapping = (*self.mapping, ContentMatching(prefix, suffix))
|
||||
return replace(self, mapping=new_mapping)
|
||||
|
||||
def with_additional_allowed_keys(self, keys: frozenset[str]) -> "SDOps":
|
||||
"""Create a new SDOps instance that only passes keys present in *keys* (post-replacement).
|
||||
If allowed_keys already exists, the sets are merged via union.
|
||||
"""
|
||||
merged = frozenset(keys) | self.allowed_keys if self.allowed_keys is not None else frozenset(keys)
|
||||
return replace(self, allowed_keys=merged)
|
||||
|
||||
def with_kv_operation(
|
||||
self,
|
||||
operation: KeyValueOperation,
|
||||
@@ -101,6 +109,10 @@ class SDOps:
|
||||
continue
|
||||
if replacement.content in key:
|
||||
key = key.replace(replacement.content, replacement.replacement)
|
||||
|
||||
if self.allowed_keys is not None and key not in self.allowed_keys:
|
||||
return None
|
||||
|
||||
return key
|
||||
|
||||
def apply_to_key_value(self, key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
|
||||
|
||||
@@ -53,6 +53,21 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
|
||||
def lora(self, lora_path: str, strength: float = 1.0, sd_ops: SDOps | None = None) -> "SingleGPUModelBuilder":
|
||||
return replace(self, loras=(*self.loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops)))
|
||||
|
||||
def with_sd_ops(self, sd_ops: SDOps | None) -> "SingleGPUModelBuilder":
|
||||
return replace(self, model_sd_ops=sd_ops)
|
||||
|
||||
def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> "SingleGPUModelBuilder":
|
||||
return replace(self, module_ops=module_ops)
|
||||
|
||||
def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "SingleGPUModelBuilder":
|
||||
return replace(self, loras=loras)
|
||||
|
||||
def with_registry(self, registry: Registry) -> "SingleGPUModelBuilder":
|
||||
return replace(self, registry=registry)
|
||||
|
||||
def with_lora_load_device(self, device: torch.device) -> "SingleGPUModelBuilder":
|
||||
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)
|
||||
@@ -83,7 +98,12 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
|
||||
retval = meta_model.to(device)
|
||||
return retval
|
||||
|
||||
def build(self, device: torch.device | None = None, dtype: torch.dtype | None = None) -> ModelType:
|
||||
def build(
|
||||
self,
|
||||
device: torch.device | None = None,
|
||||
dtype: torch.dtype | None = None,
|
||||
**kwargs: object, # noqa: ARG002
|
||||
) -> ModelType:
|
||||
device = torch.device("cuda") if device is None else device
|
||||
config = self.model_config()
|
||||
meta_model = self.meta_model(config, self.module_ops)
|
||||
|
||||
@@ -0,0 +1,222 @@
|
||||
"""Video modality tiling helpers.
|
||||
Provides :class:`VideoModalityTilingHelper` — a stateless helper that
|
||||
tiles and blends video :class:`Modality` token sequences by
|
||||
spatial/temporal region. Tile geometry is represented by the existing
|
||||
:class:`Tile` NamedTuple from :mod:`ltx_core.tiling`; no distributed
|
||||
primitives are required.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, replace
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.model.transformer.modality import Modality
|
||||
from ltx_core.tiling import Tile, TileCountConfig, create_tiles, identity_mapping_operation, split_by_count
|
||||
from ltx_core.tools import VideoLatentTools
|
||||
from ltx_core.types import VideoLatentShape
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TilingContext:
|
||||
"""Opaque context produced by :meth:`VideoModalityTilingHelper.tile_modality`.
|
||||
Carries the token-level keep mask and per-conditioning-token blend
|
||||
weights needed by :meth:`~VideoModalityTilingHelper.blend`.
|
||||
"""
|
||||
|
||||
keep_mask: torch.Tensor
|
||||
cond_blend_weights: torch.Tensor | None
|
||||
"""``(num_kept_cond,)`` — weight for each kept conditioning token,
|
||||
equal to ``1 / num_tiles_that_keep_this_token``. ``None`` when
|
||||
there are no conditioning tokens."""
|
||||
|
||||
|
||||
class VideoModalityTilingHelper:
|
||||
"""Stateless helper that tiles and blends video :class:`Modality` sequences.
|
||||
Constructed once with a :class:`TileCountConfig` and
|
||||
:class:`VideoLatentTools`. Tiles are computed at construction and
|
||||
available via the :attr:`tiles` property. Use :meth:`tile_modality`
|
||||
and :meth:`blend` with any tile from that list.
|
||||
Usage::
|
||||
helper = VideoModalityTilingHelper(tiling, video_tools)
|
||||
for tile in helper.tiles:
|
||||
tiled_mod, ctx = helper.tile_modality(modality, tile)
|
||||
result = run_model(tiled_mod)
|
||||
helper.blend(result, tile, ctx, output=output)
|
||||
"""
|
||||
|
||||
def __init__(self, tiling: TileCountConfig, video_tools: VideoLatentTools) -> None:
|
||||
self._patchifier = video_tools.patchifier
|
||||
self._latent_shape = video_tools.target_shape
|
||||
self._num_generated_tokens = self._patchifier.get_token_count(self._latent_shape)
|
||||
self._tiles = create_tiles(
|
||||
torch.Size([self._latent_shape.frames, self._latent_shape.height, self._latent_shape.width]),
|
||||
splitters=[
|
||||
split_by_count(tiling.frames.num_tiles, tiling.frames.overlap),
|
||||
split_by_count(tiling.height.num_tiles, tiling.height.overlap),
|
||||
split_by_count(tiling.width.num_tiles, tiling.width.overlap),
|
||||
],
|
||||
mappers=[identity_mapping_operation] * 3,
|
||||
)
|
||||
|
||||
@property
|
||||
def tiles(self) -> list[Tile]:
|
||||
"""All tiles for the configured tiling layout."""
|
||||
return self._tiles
|
||||
|
||||
# -- tile modality -----------------------------------------------------
|
||||
|
||||
def tile_modality(self, modality: Modality, tile: Tile) -> 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).
|
||||
Returns:
|
||||
A ``(tiled_modality, context)`` tuple. Pass *context* to
|
||||
:meth:`blend` together with the model output.
|
||||
"""
|
||||
keep_mask = self._keep_mask(modality, tile)
|
||||
|
||||
tile_attention_mask = None
|
||||
if modality.attention_mask is not None:
|
||||
keep_indices = keep_mask.nonzero(as_tuple=False).squeeze(1)
|
||||
tile_attention_mask = modality.attention_mask[:, keep_indices, :][:, :, keep_indices]
|
||||
|
||||
tiled = replace(
|
||||
modality,
|
||||
latent=modality.latent[:, keep_mask, :],
|
||||
timesteps=modality.timesteps[:, keep_mask],
|
||||
positions=modality.positions[:, :, keep_mask, :],
|
||||
attention_mask=tile_attention_mask,
|
||||
)
|
||||
|
||||
cond_blend_weights = None
|
||||
num_total = modality.latent.shape[1]
|
||||
if num_total > self._num_generated_tokens:
|
||||
cond_keep = keep_mask[self._num_generated_tokens :]
|
||||
# Count how many tiles keep each conditioning token.
|
||||
cond_counts = torch.zeros(cond_keep.sum(), dtype=torch.float32)
|
||||
for t in self._tiles:
|
||||
other_mask = self._keep_mask(modality, t)
|
||||
other_cond = other_mask[self._num_generated_tokens :]
|
||||
# Map other tile's kept cond tokens into this tile's kept subset.
|
||||
cond_counts += other_cond[cond_keep].float()
|
||||
cond_blend_weights = 1.0 / cond_counts
|
||||
|
||||
return tiled, TilingContext(keep_mask=keep_mask, cond_blend_weights=cond_blend_weights)
|
||||
|
||||
# -- blend -------------------------------------------------------------
|
||||
|
||||
def blend(
|
||||
self,
|
||||
tile_to_blend: torch.Tensor,
|
||||
tile: Tile,
|
||||
context: TilingContext,
|
||||
output: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Blend-weight tile results and accumulate into the full token space.
|
||||
Premultiplied (blend-weighted) data is **added** to *output*,
|
||||
allowing multiple tiles to be accumulated into the same buffer.
|
||||
Args:
|
||||
tile_to_blend: Denoised tile tensor ``(B, num_tile_tokens, D)``,
|
||||
where the first ``_tile_generated_token_count(tile)``
|
||||
entries are generated tokens and the remainder are
|
||||
conditioning tokens.
|
||||
tile: The :class:`Tile` that was used in :meth:`tile_modality`.
|
||||
context: The :class:`TilingContext` returned by :meth:`tile_modality`.
|
||||
output: Optional pre-allocated output tensor. When provided
|
||||
its shape must be ``(B, num_total_tokens, D)`` and the
|
||||
blended tile is **added** into it. When ``None`` a new
|
||||
zero-filled tensor is created.
|
||||
Returns:
|
||||
The output tensor with the blended tile added at the correct
|
||||
positions.
|
||||
"""
|
||||
batch, _, dim = tile_to_blend.shape
|
||||
num_tile_gen = self._tile_generated_token_count(tile)
|
||||
gen_indices = self._generated_token_indices(tile)
|
||||
|
||||
num_total_tokens = context.keep_mask.shape[0]
|
||||
expected_shape = (batch, num_total_tokens, dim)
|
||||
|
||||
if output is not None:
|
||||
if output.shape != expected_shape:
|
||||
raise ValueError(f"Expected output shape {expected_shape}, got {output.shape}")
|
||||
result = output
|
||||
else:
|
||||
result = torch.zeros(*expected_shape, device=tile_to_blend.device, dtype=tile_to_blend.dtype)
|
||||
|
||||
# Blend mask is (tile_F, tile_H, tile_W) — one weight per token in row-major order.
|
||||
blend_weights = tile.blend_mask.reshape(-1).to(device=tile_to_blend.device, dtype=tile_to_blend.dtype)
|
||||
tile_gen = tile_to_blend[:, :num_tile_gen, :] * blend_weights[None, :, None]
|
||||
|
||||
result[:, gen_indices, :] += tile_gen
|
||||
|
||||
# Scatter kept conditioning tokens, weighted by 1/N where N is
|
||||
# the number of tiles that keep each token (so they sum to 1).
|
||||
if num_total_tokens > self._num_generated_tokens and context.cond_blend_weights is not None:
|
||||
cond_keep = context.keep_mask[self._num_generated_tokens :]
|
||||
cond_indices = self._num_generated_tokens + cond_keep.nonzero(as_tuple=False).squeeze(1)
|
||||
weights = context.cond_blend_weights.to(device=tile_to_blend.device, dtype=tile_to_blend.dtype)
|
||||
result[:, cond_indices, :] += tile_to_blend[:, num_tile_gen:, :] * weights[None, :, None]
|
||||
|
||||
return result
|
||||
|
||||
# -- private -----------------------------------------------------------
|
||||
|
||||
def _tile_generated_token_count(self, tile: Tile) -> int:
|
||||
"""Number of generated tokens in *tile*."""
|
||||
frame_slice, height_slice, width_slice = tile.in_coords
|
||||
tile_shape = VideoLatentShape(
|
||||
batch=self._latent_shape.batch,
|
||||
channels=self._latent_shape.channels,
|
||||
frames=frame_slice.stop - frame_slice.start,
|
||||
height=height_slice.stop - height_slice.start,
|
||||
width=width_slice.stop - width_slice.start,
|
||||
)
|
||||
return self._patchifier.get_token_count(tile_shape)
|
||||
|
||||
def _generated_token_indices(self, tile: Tile) -> torch.Tensor:
|
||||
"""Flat token indices of *tile*'s generated tokens in the full sequence."""
|
||||
frame_slice, height_slice, width_slice = tile.in_coords
|
||||
f = torch.arange(frame_slice.start, frame_slice.stop)
|
||||
h = torch.arange(height_slice.start, height_slice.stop)
|
||||
w = torch.arange(width_slice.start, width_slice.stop)
|
||||
return (
|
||||
f[:, None, None] * self._latent_shape.height * self._latent_shape.width
|
||||
+ h[None, :, None] * self._latent_shape.width
|
||||
+ w[None, None, :]
|
||||
).reshape(-1)
|
||||
|
||||
def _keep_mask(self, modality: Modality, tile: Tile) -> torch.Tensor:
|
||||
"""Boolean mask ``(num_total_tokens,)`` — True for tokens the tile processes.
|
||||
Generated tokens are selected by grid position. Conditioning
|
||||
tokens are kept when their ``[start, end)`` intervals overlap
|
||||
the tile in all three dimensions, or when they have a negative
|
||||
time coordinate (reference tokens).
|
||||
"""
|
||||
num_total = modality.latent.shape[1]
|
||||
mask = torch.zeros(num_total, dtype=torch.bool)
|
||||
|
||||
gen_indices = self._generated_token_indices(tile)
|
||||
mask[gen_indices] = True
|
||||
|
||||
if num_total > self._num_generated_tokens:
|
||||
gen_positions = modality.positions[:, :, gen_indices, :] # (B, 3, num_tile_gen, 2)
|
||||
tile_start = gen_positions[..., 0].amin(dim=2) # (B, 3)
|
||||
tile_end = gen_positions[..., 1].amax(dim=2) # (B, 3)
|
||||
|
||||
cond_positions = modality.positions[:, :, self._num_generated_tokens :, :] # (B, 3, num_cond, 2)
|
||||
|
||||
overlaps = (cond_positions[..., 0] < tile_end.unsqueeze(2)) & (
|
||||
cond_positions[..., 1] > tile_start.unsqueeze(2)
|
||||
) # (B, 3, num_cond)
|
||||
overlaps_all_dims = overlaps.all(dim=1) # (B, num_cond)
|
||||
|
||||
has_negative_time = cond_positions[:, 0, :, 0] < 0 # (B, num_cond)
|
||||
|
||||
keep_cond = (overlaps_all_dims | has_negative_time).any(dim=0) # (num_cond,)
|
||||
mask[self._num_generated_tokens :] = keep_cond
|
||||
|
||||
return mask
|
||||
@@ -500,6 +500,8 @@ class VocoderWithBWE(nn.Module):
|
||||
to a higher sample rate. The BWE computes a mel spectrogram from the
|
||||
vocoder output, runs it through a second generator to predict a residual,
|
||||
and adds it to a sinc-resampled skip connection.
|
||||
The forward pass runs in fp32 via autocast to avoid bfloat16 accumulation
|
||||
errors that degrade spectral metrics by 40-90%.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -548,28 +550,45 @@ class VocoderWithBWE(nn.Module):
|
||||
|
||||
def forward(self, mel_spec: torch.Tensor) -> torch.Tensor:
|
||||
"""Run the full vocoder + BWE forward pass.
|
||||
Runs in float32 regardless of weight or input dtype. bfloat16 arithmetic
|
||||
causes 40-90% spectral metric degradation due to accumulation errors
|
||||
compounding through 108 sequential convolutions in the BigVGAN v2 architecture.
|
||||
Args:
|
||||
mel_spec: Mel spectrogram of shape (B, 2, T, mel_bins) for stereo
|
||||
or (B, T, mel_bins) for mono. Same format as Vocoder.forward.
|
||||
Returns:
|
||||
Waveform tensor of shape (B, out_channels, T_out) clipped to [-1, 1].
|
||||
"""
|
||||
x = self.vocoder(mel_spec)
|
||||
_, _, length_low_rate = x.shape
|
||||
output_length = length_low_rate * self.output_sampling_rate // self.input_sampling_rate
|
||||
input_dtype = mel_spec.dtype
|
||||
# Run the entire forward pass in fp32. bfloat16 accumulation errors
|
||||
# compound through 108 sequential convolutions and degrade spectral
|
||||
# metrics (mel_l1, MRSTFT) by 40-90% while perceptual quality (CDPAM)
|
||||
# is unaffected. fp32 eliminates this degradation.
|
||||
# We use autocast(dtype=float32) rather than self.float() because it
|
||||
# upcasts bf16 weights per-op at kernel level, avoiding the temporary
|
||||
# memory spike of self.float() / self.to(original_dtype).
|
||||
# Benchmarked on H100 (128.5M-param model):
|
||||
# autocast fp32: +70 MB peak VRAM, 123 ms (vs 482 MB / 95 ms for bf16)
|
||||
# model.float(): +324 MB peak VRAM, 149 ms
|
||||
# Tested: both approaches produce bit-identical output.
|
||||
|
||||
# Pad to multiple of hop_length for exact mel frame count
|
||||
remainder = length_low_rate % self.hop_length
|
||||
if remainder != 0:
|
||||
x = F.pad(x, (0, self.hop_length - remainder))
|
||||
with torch.autocast(device_type=mel_spec.device.type, dtype=torch.float32):
|
||||
x = self.vocoder(mel_spec.float())
|
||||
_, _, length_low_rate = x.shape
|
||||
output_length = length_low_rate * self.output_sampling_rate // self.input_sampling_rate
|
||||
|
||||
# Compute mel spectrogram from vocoder output: (B, C, n_mels, T_frames)
|
||||
mel = self._compute_mel(x)
|
||||
# Pad to multiple of hop_length for exact mel frame count
|
||||
remainder = length_low_rate % self.hop_length
|
||||
if remainder != 0:
|
||||
x = F.pad(x, (0, self.hop_length - remainder))
|
||||
|
||||
# Vocoder.forward expects (B, C, T, mel_bins) — transpose before calling bwe_generator
|
||||
mel_for_bwe = mel.transpose(2, 3) # (B, C, T_frames, mel_bins)
|
||||
residual = self.bwe_generator(mel_for_bwe)
|
||||
skip = self.resampler(x)
|
||||
assert residual.shape == skip.shape, f"residual {residual.shape} != skip {skip.shape}"
|
||||
# Compute mel spectrogram from vocoder output: (B, C, n_mels, T_frames)
|
||||
mel = self._compute_mel(x)
|
||||
|
||||
return torch.clamp(residual + skip, -1, 1)[..., :output_length]
|
||||
# Vocoder.forward expects (B, C, T, mel_bins) — transpose before calling bwe_generator
|
||||
mel_for_bwe = mel.transpose(2, 3) # (B, C, T_frames, mel_bins)
|
||||
residual = self.bwe_generator(mel_for_bwe)
|
||||
skip = self.resampler(x)
|
||||
assert residual.shape == skip.shape, f"residual {residual.shape} != skip {skip.shape}"
|
||||
|
||||
return torch.clamp(residual + skip, -1, 1)[..., :output_length].to(input_dtype)
|
||||
|
||||
@@ -122,22 +122,18 @@ class AttentionFunction(Enum):
|
||||
FLASH_ATTENTION_3 = "flash_attention_3"
|
||||
DEFAULT = "default"
|
||||
|
||||
def __call__(
|
||||
self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
|
||||
) -> torch.Tensor:
|
||||
def to_callable(self) -> AttentionCallable:
|
||||
"""Resolve to a concrete callable. Use this at module init time so that
|
||||
torch.compile can trace through the attention call without graph breaks."""
|
||||
if self is AttentionFunction.PYTORCH:
|
||||
return PytorchAttention()(q, k, v, heads, mask)
|
||||
return PytorchAttention()
|
||||
elif self is AttentionFunction.XFORMERS:
|
||||
return XFormersAttention()(q, k, v, heads, mask)
|
||||
return XFormersAttention()
|
||||
elif self is AttentionFunction.FLASH_ATTENTION_3:
|
||||
return FlashAttention3()(q, k, v, heads, mask)
|
||||
return FlashAttention3()
|
||||
else:
|
||||
# Default behavior: XFormers if installed else - PyTorch
|
||||
return (
|
||||
XFormersAttention()(q, k, v, heads, mask)
|
||||
if memory_efficient_attention is not None
|
||||
else PytorchAttention()(q, k, v, heads, mask)
|
||||
)
|
||||
return XFormersAttention() if memory_efficient_attention is not None else PytorchAttention()
|
||||
|
||||
|
||||
class Attention(torch.nn.Module):
|
||||
@@ -154,7 +150,11 @@ class Attention(torch.nn.Module):
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.rope_type = rope_type
|
||||
self.attention_function = attention_function
|
||||
self.attention_function = (
|
||||
attention_function.to_callable()
|
||||
if isinstance(attention_function, AttentionFunction)
|
||||
else attention_function
|
||||
)
|
||||
|
||||
inner_dim = dim_head * heads
|
||||
context_dim = query_dim if context_dim is None else context_dim
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
import torch
|
||||
|
||||
from ltx_core.loader.module_ops import ModuleOps
|
||||
from ltx_core.loader.sd_ops import SDOps
|
||||
from ltx_core.model.transformer.model import LTXModel
|
||||
|
||||
|
||||
def compile_transformer(model: LTXModel) -> LTXModel:
|
||||
model.transformer_blocks = torch.nn.ModuleList(torch.compile(m) for m in model.transformer_blocks)
|
||||
|
||||
def patched_dynamo_forward(*args, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
with (
|
||||
torch._inductor.config.patch(unsafe_skip_cache_dynamic_shape_guards=True),
|
||||
torch._dynamo.config.patch( # type: ignore[attr-defined]
|
||||
inline_inbuilt_nn_modules=True, cache_size_limit=256, allow_unspec_int_on_nn_module=True
|
||||
),
|
||||
):
|
||||
return model.forward_without_compilation(*args, **kwargs)
|
||||
|
||||
model.forward_without_compilation = model.forward
|
||||
model.forward = patched_dynamo_forward
|
||||
return model
|
||||
|
||||
|
||||
COMPILE_TRANSFORMER = ModuleOps(
|
||||
name="compile_transformer",
|
||||
matcher=lambda model: isinstance(model, LTXModel),
|
||||
mutator=lambda model: compile_transformer(model),
|
||||
)
|
||||
|
||||
|
||||
def modify_sd_ops_for_compilation(original_sd_ops: SDOps, number_of_blocks: int = 48) -> SDOps:
|
||||
for i in range(number_of_blocks):
|
||||
original_sd_ops = original_sd_ops.with_replacement(
|
||||
f"transformer_blocks.{i}.", f"transformer_blocks.{i}._orig_mod."
|
||||
)
|
||||
return original_sd_ops
|
||||
@@ -1,3 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
@@ -38,3 +41,17 @@ class Modality:
|
||||
enabled: bool = True
|
||||
context_mask: torch.Tensor | None = None
|
||||
attention_mask: torch.Tensor | None = None
|
||||
|
||||
def split(self, sizes: list[int]) -> list[Modality]:
|
||||
"""Split along the batch dimension into chunks of the given sizes."""
|
||||
n = len(sizes)
|
||||
split_fields: dict[str, list[torch.Tensor | None] | list[bool]] = {}
|
||||
for f in dataclasses.fields(self):
|
||||
value = getattr(self, f.name)
|
||||
if isinstance(value, torch.Tensor):
|
||||
split_fields[f.name] = list(value.split(sizes, dim=0))
|
||||
elif value is None or isinstance(value, bool):
|
||||
split_fields[f.name] = [value] * n
|
||||
else:
|
||||
raise TypeError(f"Cannot split field {f.name!r}: unsupported type {type(value)}")
|
||||
return [Modality(**{name: parts[i] for name, parts in split_fields.items()}) for i in range(n)]
|
||||
|
||||
@@ -7,7 +7,7 @@ from ltx_core.model.video_vae.model_configurator import (
|
||||
VideoEncoderConfigurator,
|
||||
)
|
||||
from ltx_core.model.video_vae.tiling import SpatialTilingConfig, TemporalTilingConfig, TilingConfig
|
||||
from ltx_core.model.video_vae.video_vae import VideoDecoder, VideoEncoder, decode_video, get_video_chunks_number
|
||||
from ltx_core.model.video_vae.video_vae import VideoDecoder, VideoEncoder, get_video_chunks_number
|
||||
|
||||
__all__ = [
|
||||
"VAE_DECODER_COMFY_KEYS_FILTER",
|
||||
@@ -19,6 +19,5 @@ __all__ = [
|
||||
"VideoDecoderConfigurator",
|
||||
"VideoEncoder",
|
||||
"VideoEncoderConfigurator",
|
||||
"decode_video",
|
||||
"get_video_chunks_number",
|
||||
]
|
||||
|
||||
@@ -1,72 +1,4 @@
|
||||
import itertools
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, List, NamedTuple, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def compute_trapezoidal_mask_1d(
|
||||
length: int,
|
||||
ramp_left: int,
|
||||
ramp_right: int,
|
||||
left_starts_from_0: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Generate a 1D trapezoidal blending mask with linear ramps.
|
||||
Args:
|
||||
length: Output length of the mask.
|
||||
ramp_left: Fade-in length on the left.
|
||||
ramp_right: Fade-out length on the right.
|
||||
left_starts_from_0: Whether the ramp starts from 0 or first non-zero value.
|
||||
Useful for temporal tiles where the first tile is causal.
|
||||
Returns:
|
||||
A 1D tensor of shape `(length,)` with values in [0, 1].
|
||||
"""
|
||||
if length <= 0:
|
||||
raise ValueError("Mask length must be positive.")
|
||||
|
||||
ramp_left = max(0, min(ramp_left, length))
|
||||
ramp_right = max(0, min(ramp_right, length))
|
||||
|
||||
mask = torch.ones(length)
|
||||
|
||||
if ramp_left > 0:
|
||||
interval_length = ramp_left + 1 if left_starts_from_0 else ramp_left + 2
|
||||
fade_in = torch.linspace(0.0, 1.0, interval_length)[:-1]
|
||||
if not left_starts_from_0:
|
||||
fade_in = fade_in[1:]
|
||||
mask[:ramp_left] *= fade_in
|
||||
|
||||
if ramp_right > 0:
|
||||
fade_out = torch.linspace(1.0, 0.0, steps=ramp_right + 2)[1:-1]
|
||||
mask[-ramp_right:] *= fade_out
|
||||
|
||||
return mask.clamp_(0, 1)
|
||||
|
||||
|
||||
def compute_rectangular_mask_1d(
|
||||
length: int,
|
||||
left_ramp: int,
|
||||
right_ramp: int,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Generate a 1D rectangular (pulse) mask.
|
||||
Args:
|
||||
length: Output length of the mask.
|
||||
left_ramp: Number of elements at the start of the mask to set to 0.
|
||||
right_ramp: Number of elements at the end of the mask to set to 0.
|
||||
Returns:
|
||||
A 1D tensor of shape `(length,)` with values 0 or 1.
|
||||
"""
|
||||
if length <= 0:
|
||||
raise ValueError("Mask length must be positive.")
|
||||
|
||||
mask = torch.ones(length)
|
||||
if left_ramp > 0:
|
||||
mask[:left_ramp] = 0
|
||||
if right_ramp > 0:
|
||||
mask[-right_ramp:] = 0
|
||||
return mask
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -135,157 +67,3 @@ class TilingConfig:
|
||||
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),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DimensionIntervals:
|
||||
"""Defines how a single dimension is split into overlapping intervals (tiles).
|
||||
Each list has length N where N is the number of intervals. The i-th element
|
||||
of each list describes the i-th interval.
|
||||
Attributes:
|
||||
starts: Start index of each interval (inclusive).
|
||||
ends: End index of each interval (exclusive).
|
||||
left_ramps: Length of the left blend ramp for each interval.
|
||||
Used to create masks that fade in from 0 to 1.
|
||||
right_ramps: Length of the right blend ramp for each interval.
|
||||
Used to create masks that fade out from 1 to 0.
|
||||
"""
|
||||
|
||||
starts: List[int]
|
||||
ends: List[int]
|
||||
left_ramps: List[int]
|
||||
right_ramps: List[int]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TensorTilingSpec:
|
||||
"""Specifies how a tensor of a given shape is split into intervals (tiles) along each dimension.
|
||||
Attributes:
|
||||
original_shape: Shape of the tensor being tiled.
|
||||
dimension_intervals: Per-dimension intervals (starts, ends, ramps) for each axis.
|
||||
"""
|
||||
|
||||
original_shape: torch.Size
|
||||
dimension_intervals: Tuple[DimensionIntervals, ...]
|
||||
|
||||
|
||||
# Operation to split a single dimension of the tensor into intervals based on the length along the dimension.
|
||||
SplitOperation = Callable[[int], DimensionIntervals]
|
||||
# Operation to map the intervals in input dimension to slices and masks along a corresponding output dimension.
|
||||
MappingOperation = Callable[[DimensionIntervals], tuple[list[slice], list[torch.Tensor | None]]]
|
||||
|
||||
|
||||
def default_split_operation(length: int) -> DimensionIntervals:
|
||||
return DimensionIntervals(starts=[0], ends=[length], left_ramps=[0], right_ramps=[0])
|
||||
|
||||
|
||||
DEFAULT_SPLIT_OPERATION: SplitOperation = default_split_operation
|
||||
|
||||
|
||||
def default_mapping_operation(
|
||||
_intervals: DimensionIntervals,
|
||||
) -> tuple[list[slice], list[torch.Tensor | None]]:
|
||||
return [slice(0, None)], [None]
|
||||
|
||||
|
||||
DEFAULT_MAPPING_OPERATION: MappingOperation = default_mapping_operation
|
||||
|
||||
|
||||
class Tile(NamedTuple):
|
||||
"""
|
||||
Represents a single tile.
|
||||
Attributes:
|
||||
in_coords:
|
||||
Tuple of slices specifying where to cut the tile from the INPUT tensor.
|
||||
out_coords:
|
||||
Tuple of slices specifying where this tile's OUTPUT should be placed in the reconstructed OUTPUT tensor.
|
||||
masks_1d:
|
||||
Per-dimension masks in OUTPUT units.
|
||||
These are used to create all-dimensional blending mask.
|
||||
Methods:
|
||||
blend_mask:
|
||||
Create a single N-D mask from the per-dimension masks.
|
||||
"""
|
||||
|
||||
in_coords: Tuple[slice, ...]
|
||||
out_coords: Tuple[slice, ...]
|
||||
masks_1d: Tuple[Tuple[torch.Tensor, ...]]
|
||||
|
||||
@property
|
||||
def blend_mask(self) -> torch.Tensor:
|
||||
num_dims = len(self.out_coords)
|
||||
per_dimension_masks: List[torch.Tensor] = []
|
||||
|
||||
for dim_idx in range(num_dims):
|
||||
mask_1d = self.masks_1d[dim_idx]
|
||||
view_shape = [1] * num_dims
|
||||
if mask_1d is None:
|
||||
# Broadcast mask along this dimension (length 1).
|
||||
one = torch.ones(1)
|
||||
|
||||
view_shape[dim_idx] = 1
|
||||
per_dimension_masks.append(one.view(*view_shape))
|
||||
continue
|
||||
|
||||
# Reshape (L,) -> (1, ..., L, ..., 1) so masks across dimensions broadcast-multiply.
|
||||
view_shape[dim_idx] = mask_1d.shape[0]
|
||||
per_dimension_masks.append(mask_1d.view(*view_shape))
|
||||
|
||||
# Multiply per-dimension masks to form the full N-D mask (separable blending window).
|
||||
combined_mask = per_dimension_masks[0]
|
||||
for mask in per_dimension_masks[1:]:
|
||||
combined_mask = combined_mask * mask
|
||||
|
||||
return combined_mask
|
||||
|
||||
|
||||
def create_tiles_from_intervals_and_mappers(
|
||||
intervals: TensorTilingSpec,
|
||||
mappers: List[MappingOperation],
|
||||
) -> List[Tile]:
|
||||
full_dim_input_slices = []
|
||||
full_dim_output_slices = []
|
||||
full_dim_masks_1d = []
|
||||
for axis_index in range(len(intervals.original_shape)):
|
||||
dimension_intervals = intervals.dimension_intervals[axis_index]
|
||||
starts = dimension_intervals.starts
|
||||
ends = dimension_intervals.ends
|
||||
input_slices = [slice(s, e) for s, e in zip(starts, ends, strict=True)]
|
||||
output_slices, masks_1d = mappers[axis_index](dimension_intervals)
|
||||
full_dim_input_slices.append(input_slices)
|
||||
full_dim_output_slices.append(output_slices)
|
||||
full_dim_masks_1d.append(masks_1d)
|
||||
|
||||
tiles = []
|
||||
tile_in_coords = list(itertools.product(*full_dim_input_slices))
|
||||
tile_out_coords = list(itertools.product(*full_dim_output_slices))
|
||||
tile_mask_1ds = list(itertools.product(*full_dim_masks_1d))
|
||||
for in_coord, out_coord, mask_1d in zip(tile_in_coords, tile_out_coords, tile_mask_1ds, strict=True):
|
||||
tiles.append(
|
||||
Tile(
|
||||
in_coords=in_coord,
|
||||
out_coords=out_coord,
|
||||
masks_1d=mask_1d,
|
||||
)
|
||||
)
|
||||
return tiles
|
||||
|
||||
|
||||
def create_tiles(
|
||||
tensor_shape: torch.Size,
|
||||
splitters: List[SplitOperation],
|
||||
mappers: List[MappingOperation],
|
||||
) -> List[Tile]:
|
||||
if len(splitters) != len(tensor_shape):
|
||||
raise ValueError(
|
||||
f"Number of splitters must be equal to number of dimensions in tensor shape, "
|
||||
f"got {len(splitters)} and {len(tensor_shape)}"
|
||||
)
|
||||
if len(mappers) != len(tensor_shape):
|
||||
raise ValueError(
|
||||
f"Number of mappers must be equal to number of dimensions in tensor shape, "
|
||||
f"got {len(mappers)} and {len(tensor_shape)}"
|
||||
)
|
||||
intervals = [splitter(length) for splitter, length in zip(splitters, tensor_shape, strict=True)]
|
||||
tiling_spec = TensorTilingSpec(original_shape=tensor_shape, dimension_intervals=tuple(intervals))
|
||||
return create_tiles_from_intervals_and_mappers(tiling_spec, mappers)
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
import logging
|
||||
from dataclasses import replace
|
||||
from typing import Any, Callable, Iterator, List, Tuple
|
||||
|
||||
import torch
|
||||
@@ -13,17 +12,23 @@ from ltx_core.model.video_vae.enums import LogVarianceType, NormLayerType, Paddi
|
||||
from ltx_core.model.video_vae.ops import PerChannelStatistics, patchify, unpatchify
|
||||
from ltx_core.model.video_vae.resnet import ResnetBlock3D, UNetMidBlock3D
|
||||
from ltx_core.model.video_vae.sampling import DepthToSpaceUpsample, SpaceToDepthDownsample
|
||||
from ltx_core.model.video_vae.tiling import (
|
||||
from ltx_core.model.video_vae.tiling import TilingConfig
|
||||
from ltx_core.tiling import (
|
||||
DEFAULT_MAPPING_OPERATION,
|
||||
DEFAULT_SPLIT_OPERATION,
|
||||
DimensionIntervals,
|
||||
MappingOperation,
|
||||
SplitOperation,
|
||||
Tile,
|
||||
TilingConfig,
|
||||
compute_rectangular_mask_1d,
|
||||
compute_trapezoidal_mask_1d,
|
||||
create_tiles,
|
||||
split_temporal,
|
||||
)
|
||||
from ltx_core.tiling import (
|
||||
split_by_size as split_in_spatial,
|
||||
)
|
||||
from ltx_core.tiling import (
|
||||
split_temporal_causal as split_in_temporal,
|
||||
)
|
||||
from ltx_core.types import VIDEO_SCALE_FACTORS, SpatioTemporalScaleFactors, VideoLatentShape
|
||||
|
||||
@@ -444,12 +449,12 @@ def prepare_tiles_for_encoding(
|
||||
# Define split and map operations for the spatial dimensions
|
||||
|
||||
# Height axis (H)
|
||||
splitters[3] = split_with_symmetric_overlaps(tile_size_px, overlap_px)
|
||||
mappers[3] = make_mapping_operation(map_spatial_interval_to_latent, scale=VIDEO_SCALE_FACTORS.height)
|
||||
splitters[3] = split_in_spatial(tile_size_px, overlap_px)
|
||||
mappers[3] = to_mapping_operation(map_spatial_interval_to_latent, scale=VIDEO_SCALE_FACTORS.height)
|
||||
|
||||
# Width axis (W)
|
||||
splitters[4] = split_with_symmetric_overlaps(tile_size_px, overlap_px)
|
||||
mappers[4] = make_mapping_operation(map_spatial_interval_to_latent, scale=VIDEO_SCALE_FACTORS.width)
|
||||
splitters[4] = split_in_spatial(tile_size_px, overlap_px)
|
||||
mappers[4] = to_mapping_operation(map_spatial_interval_to_latent, scale=VIDEO_SCALE_FACTORS.width)
|
||||
|
||||
if tiling_config is not None and tiling_config.temporal_config is not None:
|
||||
cfg = tiling_config.temporal_config
|
||||
@@ -460,8 +465,8 @@ def prepare_tiles_for_encoding(
|
||||
logger.warning(f"Overlap frames {overlap_frames} is less than 16, setting to minimum required 16")
|
||||
overlap_frames = minimum_temporal_overlap_frames
|
||||
|
||||
splitters[2] = split_temporal_frames(tile_size_frames, overlap_frames)
|
||||
mappers[2] = make_mapping_operation(map_temporal_interval_to_latent, scale=VIDEO_SCALE_FACTORS.time)
|
||||
splitters[2] = split_temporal(tile_size_frames, overlap_frames)
|
||||
mappers[2] = to_mapping_operation(map_temporal_interval_to_latent, scale=VIDEO_SCALE_FACTORS.time)
|
||||
|
||||
return create_tiles(video.shape, splitters, mappers)
|
||||
|
||||
@@ -784,8 +789,8 @@ class VideoDecoder(nn.Module):
|
||||
axis_length = latent.shape[axis_idx]
|
||||
lower_threshold = max(2, overlap + 1)
|
||||
tile_size = max(lower_threshold, round(size * axis_length / long_side))
|
||||
splitters[axis_idx] = split_with_symmetric_overlaps(tile_size, overlap)
|
||||
mappers[axis_idx] = make_mapping_operation(map_spatial_interval_to_pixel, scale=factor)
|
||||
splitters[axis_idx] = split_in_spatial(tile_size, overlap)
|
||||
mappers[axis_idx] = to_mapping_operation(map_spatial_slice, scale=factor)
|
||||
|
||||
enable_on_axis(3, self.video_downscale_factors.height)
|
||||
enable_on_axis(4, self.video_downscale_factors.width)
|
||||
@@ -794,8 +799,8 @@ class VideoDecoder(nn.Module):
|
||||
cfg = tiling_config.temporal_config
|
||||
tile_size = cfg.tile_size_in_frames // self.video_downscale_factors.time
|
||||
overlap = cfg.tile_overlap_in_frames // self.video_downscale_factors.time
|
||||
splitters[2] = split_temporal_latents(tile_size, overlap)
|
||||
mappers[2] = make_mapping_operation(map_temporal_interval_to_frame, scale=self.video_downscale_factors.time)
|
||||
splitters[2] = split_in_temporal(tile_size, overlap)
|
||||
mappers[2] = to_mapping_operation(map_temporal_slice, scale=self.video_downscale_factors.time)
|
||||
|
||||
return create_tiles(latent.shape, splitters, mappers)
|
||||
|
||||
@@ -892,6 +897,29 @@ class VideoDecoder(nn.Module):
|
||||
previous_weights = previous_weights.clamp(min=1e-8)
|
||||
yield previous_chunk / previous_weights
|
||||
|
||||
def decode_video(
|
||||
self,
|
||||
latent: torch.Tensor,
|
||||
tiling_config: TilingConfig | None = None,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> Iterator[torch.Tensor]:
|
||||
"""Decode a video latent tensor, yielding uint8 chunks ``[f, h, w, c]``.
|
||||
Subclasses (e.g. ``DistributedVideoDecoder``) may override this to
|
||||
control eagerness or distribution across ranks.
|
||||
"""
|
||||
|
||||
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
|
||||
|
||||
if tiling_config is not None:
|
||||
for frames in self.tiled_decode(latent, tiling_config, generator=generator):
|
||||
yield convert_to_uint8(frames)
|
||||
else:
|
||||
decoded = self(latent, generator=generator)
|
||||
yield convert_to_uint8(decoded)
|
||||
|
||||
def _group_tiles_by_temporal_slice(self, tiles: List[Tile]) -> List[List[Tile]]:
|
||||
"""Group tiles by their temporal output slice."""
|
||||
if not tiles:
|
||||
@@ -963,36 +991,6 @@ class VideoDecoder(nn.Module):
|
||||
return weights
|
||||
|
||||
|
||||
def decode_video(
|
||||
latent: torch.Tensor,
|
||||
video_decoder: VideoDecoder,
|
||||
tiling_config: TilingConfig | None = None,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> Iterator[torch.Tensor]:
|
||||
"""
|
||||
Decode a video latent tensor with the given decoder.
|
||||
Args:
|
||||
latent: Tensor [c, f, h, w]
|
||||
video_decoder: Decoder module.
|
||||
tiling_config: Optional tiling settings.
|
||||
generator: Optional random generator for deterministic decoding.
|
||||
Yields:
|
||||
Decoded chunk [f, h, w, c], uint8 in [0, 255].
|
||||
"""
|
||||
|
||||
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
|
||||
|
||||
if tiling_config is not None:
|
||||
for frames in video_decoder.tiled_decode(latent, tiling_config, generator=generator):
|
||||
yield convert_to_uint8(frames)
|
||||
else:
|
||||
decoded_video = video_decoder(latent, generator=generator)
|
||||
yield convert_to_uint8(decoded_video)
|
||||
|
||||
|
||||
def get_video_chunks_number(num_frames: int, tiling_config: TilingConfig | None = None) -> int:
|
||||
"""
|
||||
Get the number of video chunks for a given number of frames and tiling configuration.
|
||||
@@ -1009,82 +1007,7 @@ def get_video_chunks_number(num_frames: int, tiling_config: TilingConfig | None
|
||||
return (num_frames - 1 + frame_stride - 1) // frame_stride
|
||||
|
||||
|
||||
def split_with_symmetric_overlaps(size: int, overlap: int) -> SplitOperation:
|
||||
def split(dimension_size: int) -> DimensionIntervals:
|
||||
if dimension_size <= size:
|
||||
return DEFAULT_SPLIT_OPERATION(dimension_size)
|
||||
amount = (dimension_size + size - 2 * overlap - 1) // (size - overlap)
|
||||
starts = [i * (size - overlap) for i in range(amount)]
|
||||
ends = [start + size for start in starts]
|
||||
ends[-1] = dimension_size
|
||||
left_ramps = [0] + [overlap] * (amount - 1)
|
||||
right_ramps = [overlap] * (amount - 1) + [0]
|
||||
return DimensionIntervals(starts=starts, ends=ends, left_ramps=left_ramps, right_ramps=right_ramps)
|
||||
|
||||
return split
|
||||
|
||||
|
||||
def split_temporal_latents(size: int, overlap: int) -> SplitOperation:
|
||||
"""Split a temporal axis into overlapping tiles with causal handling.
|
||||
Example with size=24, overlap=8 (units are whatever axis you split):
|
||||
Non-causal split would produce:
|
||||
Tile 0: [0, 24), left_ramp=0, right_ramp=8
|
||||
Tile 1: [16, 40), left_ramp=8, right_ramp=8
|
||||
Tile 2: [32, 56), left_ramp=8, right_ramp=0
|
||||
Causal split produces:
|
||||
Tile 0: [0, 24), left_ramp=0, right_ramp=8 (unchanged - starts at anchor)
|
||||
Tile 1: [15, 40), left_ramp=9, right_ramp=8 (shifted back 1, ramp +1)
|
||||
Tile 2: [31, 56), left_ramp=9, right_ramp=0 (shifted back 1, ramp +1)
|
||||
This ensures each tile can causally depend on frames from previous tiles while maintaining
|
||||
proper temporal continuity through the blend ramps.
|
||||
Args:
|
||||
size: Tile size in *axis units* (latent steps for LTX time tiling)
|
||||
overlap: Overlap between tiles in the same units
|
||||
Returns:
|
||||
Split operation that divides temporal dimension with causal handling
|
||||
"""
|
||||
non_causal_split = split_with_symmetric_overlaps(size, overlap)
|
||||
|
||||
def split(dimension_size: int) -> DimensionIntervals:
|
||||
if dimension_size <= size:
|
||||
return DEFAULT_SPLIT_OPERATION(dimension_size)
|
||||
intervals = non_causal_split(dimension_size)
|
||||
|
||||
starts = intervals.starts
|
||||
starts[1:] = [s - 1 for s in starts[1:]]
|
||||
|
||||
# Extend blend ramps by 1 for non-first tiles to blend over the extra frame
|
||||
left_ramps = intervals.left_ramps
|
||||
left_ramps[1:] = [r + 1 for r in left_ramps[1:]]
|
||||
|
||||
return replace(intervals, starts=starts, left_ramps=left_ramps)
|
||||
|
||||
return split
|
||||
|
||||
|
||||
def split_temporal_frames(tile_size_frames: int, overlap_frames: int) -> SplitOperation:
|
||||
"""Split a temporal axis in video frame space into overlapping tiles.
|
||||
Args:
|
||||
tile_size_frames: Tile length in frames.
|
||||
overlap_frames: Overlap between consecutive tiles in frames.
|
||||
Returns:
|
||||
Split operation that takes frame count and returns DimensionIntervals in frame indices.
|
||||
"""
|
||||
non_causal_split = split_with_symmetric_overlaps(tile_size_frames, overlap_frames)
|
||||
|
||||
def split(dimension_size: int) -> DimensionIntervals:
|
||||
if dimension_size <= tile_size_frames:
|
||||
return DEFAULT_SPLIT_OPERATION(dimension_size)
|
||||
intervals = non_causal_split(dimension_size)
|
||||
ends = intervals.ends
|
||||
ends[:-1] = [e + 1 for e in ends[:-1]]
|
||||
right_ramps = [0] * len(intervals.right_ramps)
|
||||
return replace(intervals, ends=ends, right_ramps=right_ramps)
|
||||
|
||||
return split
|
||||
|
||||
|
||||
def make_mapping_operation(
|
||||
def to_mapping_operation(
|
||||
map_func: Callable[[int, int, int, int, int], Tuple[slice, torch.Tensor | None]],
|
||||
scale: int,
|
||||
) -> MappingOperation:
|
||||
@@ -1102,13 +1025,10 @@ def make_mapping_operation(
|
||||
def map_op(intervals: DimensionIntervals) -> tuple[list[slice], list[torch.Tensor | None]]:
|
||||
output_slices: list[slice] = []
|
||||
masks_1d: list[torch.Tensor | None] = []
|
||||
number_of_slices = len(intervals.starts)
|
||||
for i in range(number_of_slices):
|
||||
start = intervals.starts[i]
|
||||
end = intervals.ends[i]
|
||||
left_ramp = intervals.left_ramps[i]
|
||||
right_ramp = intervals.right_ramps[i]
|
||||
output_slice, mask_1d = map_func(start, end, left_ramp, right_ramp, scale)
|
||||
for interval in intervals.intervals:
|
||||
output_slice, mask_1d = map_func(
|
||||
interval.start, interval.end, interval.left_ramp, interval.right_ramp, scale
|
||||
)
|
||||
output_slices.append(output_slice)
|
||||
masks_1d.append(mask_1d)
|
||||
return output_slices, masks_1d
|
||||
@@ -1116,31 +1036,13 @@ def make_mapping_operation(
|
||||
return map_op
|
||||
|
||||
|
||||
def map_temporal_interval_to_frame(
|
||||
begin: int,
|
||||
end: int,
|
||||
left_ramp: int,
|
||||
right_ramp: int,
|
||||
scale: int,
|
||||
) -> Tuple[slice, torch.Tensor]:
|
||||
"""Map temporal interval in latent space to video frame space.
|
||||
Args:
|
||||
begin: Start position in latent space
|
||||
end: End position in latent space
|
||||
left_ramp: Left ramp size in latent space
|
||||
right_ramp: Right ramp size in latent space
|
||||
scale: Scale factor for transformation
|
||||
Returns:
|
||||
Tuple of (output_slice, blend_mask)
|
||||
"""
|
||||
def map_temporal_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]:
|
||||
start = begin * scale
|
||||
stop = 1 + (end - 1) * scale
|
||||
left_ramp = 0 if left_ramp == 0 else 1 + (left_ramp - 1) * scale
|
||||
right_ramp = right_ramp * scale
|
||||
|
||||
left_ramp_frames = 0 if left_ramp == 0 else 1 + (left_ramp - 1) * scale
|
||||
right_ramp_frames = right_ramp * scale
|
||||
|
||||
mask_1d = compute_trapezoidal_mask_1d(stop - start, left_ramp_frames, right_ramp_frames, True)
|
||||
return slice(start, stop), mask_1d
|
||||
return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp, right_ramp, True)
|
||||
|
||||
|
||||
def map_temporal_interval_to_latent(
|
||||
@@ -1171,25 +1073,13 @@ def map_temporal_interval_to_latent(
|
||||
return slice(start, stop), mask_1d
|
||||
|
||||
|
||||
def map_spatial_interval_to_pixel(
|
||||
begin: int,
|
||||
end: int,
|
||||
left_ramp: int,
|
||||
right_ramp: int,
|
||||
scale: int,
|
||||
) -> Tuple[slice, torch.Tensor]:
|
||||
"""Map spatial interval in latent space to pixel space.
|
||||
Args:
|
||||
begin: Start position in latent space
|
||||
end: End position in latent space
|
||||
left_ramp: Left ramp size in latent space
|
||||
right_ramp: Right ramp size in latent space
|
||||
scale: Scale factor for transformation
|
||||
"""
|
||||
def map_spatial_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]:
|
||||
start = begin * scale
|
||||
stop = end * scale
|
||||
mask_1d = compute_trapezoidal_mask_1d(stop - start, left_ramp * scale, right_ramp * scale, False)
|
||||
return slice(start, stop), mask_1d
|
||||
left_ramp = left_ramp * scale
|
||||
right_ramp = right_ramp * scale
|
||||
|
||||
return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp, right_ramp, False)
|
||||
|
||||
|
||||
def map_spatial_interval_to_latent(
|
||||
|
||||
@@ -7,12 +7,6 @@ from ltx_core.model.transformer.model import LTXModel
|
||||
BLOCK_SIZE = 1024
|
||||
|
||||
|
||||
def calculate_weight_float8(target_weights: torch.Tensor, original_weights: torch.Tensor) -> torch.Tensor:
|
||||
result = _fused_add_round_launch(target_weights, original_weights, seed=0).to(target_weights.dtype)
|
||||
target_weights.copy_(result, non_blocking=True)
|
||||
return target_weights
|
||||
|
||||
|
||||
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
|
||||
import triton # noqa: PLC0415
|
||||
@@ -65,34 +59,44 @@ def _upcast_and_round(
|
||||
return _fused_add_round_launch(torch.zeros_like(weight, dtype=dtype), weight, seed)
|
||||
|
||||
|
||||
class Fp8CastLinear(torch.nn.Linear):
|
||||
"""nn.Linear storing weights in fp8, upcasting to input dtype during forward.
|
||||
Used via __class__ reassignment (not subclassing) so existing weight tensors
|
||||
are preserved in-place. Class-level forward is required for torch.compile
|
||||
compatibility — instance-level closure monkey-patches cause graph breaks.
|
||||
"""
|
||||
|
||||
_with_stochastic_rounding: bool
|
||||
_seed: int
|
||||
|
||||
def forward(self, input: torch.Tensor) -> torch.Tensor: # noqa: A002, type: ignore[override]
|
||||
w_up = _upcast_and_round(self.weight, input.dtype, self._with_stochastic_rounding, self._seed)
|
||||
b_up = (
|
||||
_upcast_and_round(self.bias, input.dtype, self._with_stochastic_rounding, self._seed)
|
||||
if self.bias is not None
|
||||
else None
|
||||
)
|
||||
return torch.nn.functional.linear(input, w_up, b_up)
|
||||
|
||||
|
||||
def _replace_fwd_with_upcast(layer: torch.nn.Linear, with_stochastic_rounding: bool = False, seed: int = 0) -> None:
|
||||
"""
|
||||
Replace linear.forward and rms_norm.forward with a version that:
|
||||
- upcasts weight and bias to input's dtype
|
||||
- returns F.linear or F.rms_norm calculated in that dtype
|
||||
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.
|
||||
"""
|
||||
|
||||
layer.original_forward = layer.forward
|
||||
|
||||
def new_linear_forward(*args, **_kwargs) -> torch.Tensor:
|
||||
# assume first arg is the input tensor
|
||||
x = args[0]
|
||||
w_up = _upcast_and_round(layer.weight, x.dtype, with_stochastic_rounding, seed)
|
||||
b_up = None
|
||||
|
||||
if layer.bias is not None:
|
||||
b_up = _upcast_and_round(layer.bias, x.dtype, with_stochastic_rounding, seed)
|
||||
|
||||
return torch.nn.functional.linear(x, w_up, b_up)
|
||||
|
||||
layer.forward = new_linear_forward
|
||||
layer.__class__ = Fp8CastLinear
|
||||
layer._with_stochastic_rounding = with_stochastic_rounding
|
||||
layer._seed = seed
|
||||
|
||||
|
||||
def _amend_forward_with_upcast(
|
||||
model: torch.nn.Module, with_stochastic_rounding: bool = False, seed: int = 0
|
||||
) -> torch.nn.Module:
|
||||
"""
|
||||
Replace the forward method of the model's Linear and RMSNorm layers to forward
|
||||
Replace the forward method of the model's Linear layers to forward
|
||||
with upcast and optional stochastic rounding.
|
||||
"""
|
||||
for m in model.modules():
|
||||
|
||||
@@ -0,0 +1,464 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
from dataclasses import dataclass, replace
|
||||
from typing import Callable, NamedTuple
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def compute_trapezoidal_mask_1d(
|
||||
length: int,
|
||||
ramp_left: int,
|
||||
ramp_right: int,
|
||||
left_starts_from_0: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Generate a 1D trapezoidal blending mask with linear ramps.
|
||||
Args:
|
||||
length: Output length of the mask.
|
||||
ramp_left: Fade-in length on the left.
|
||||
ramp_right: Fade-out length on the right.
|
||||
left_starts_from_0: Whether the ramp starts from 0 or first non-zero value.
|
||||
Useful for temporal tiles where the first tile is causal.
|
||||
Returns:
|
||||
A 1D tensor of shape `(length,)` with values in [0, 1].
|
||||
"""
|
||||
if length <= 0:
|
||||
raise ValueError("Mask length must be positive.")
|
||||
|
||||
ramp_left = max(0, min(ramp_left, length))
|
||||
ramp_right = max(0, min(ramp_right, length))
|
||||
|
||||
mask = torch.ones(length)
|
||||
|
||||
if ramp_left > 0:
|
||||
interval_length = ramp_left + 1 if left_starts_from_0 else ramp_left + 2
|
||||
fade_in = torch.linspace(0.0, 1.0, interval_length)[:-1]
|
||||
if not left_starts_from_0:
|
||||
fade_in = fade_in[1:]
|
||||
mask[:ramp_left] *= fade_in
|
||||
|
||||
if ramp_right > 0:
|
||||
fade_out = torch.linspace(1.0, 0.0, steps=ramp_right + 2)[1:-1]
|
||||
mask[-ramp_right:] *= fade_out
|
||||
|
||||
return mask.clamp_(0, 1)
|
||||
|
||||
|
||||
def compute_rectangular_mask_1d(
|
||||
length: int,
|
||||
left_ramp: int,
|
||||
right_ramp: int,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Generate a 1D rectangular (pulse) mask.
|
||||
Args:
|
||||
length: Output length of the mask.
|
||||
left_ramp: Number of elements at the start of the mask to set to 0.
|
||||
right_ramp: Number of elements at the end of the mask to set to 0.
|
||||
Returns:
|
||||
A 1D tensor of shape `(length,)` with values 0 or 1.
|
||||
"""
|
||||
if length <= 0:
|
||||
raise ValueError("Mask length must be positive.")
|
||||
|
||||
mask = torch.ones(length)
|
||||
if left_ramp > 0:
|
||||
mask[:left_ramp] = 0
|
||||
if right_ramp > 0:
|
||||
mask[-right_ramp:] = 0
|
||||
return mask
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DimensionInterval:
|
||||
start: int
|
||||
end: int
|
||||
left_ramp: int
|
||||
right_ramp: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DimensionIntervals:
|
||||
"""Intervals which a single dimension of the latent space is split into.
|
||||
Each interval is defined by its start, end, left ramp, and right ramp.
|
||||
The start and end are the indices of the first and last element (exclusive) in the interval.
|
||||
Ramps are regions of the interval where the value of the mask tensor is
|
||||
interpolated between 0 and 1 for blending with neighboring intervals.
|
||||
The left ramp and right ramp values are the lengths of the left and right ramps.
|
||||
"""
|
||||
|
||||
intervals: list[DimensionInterval]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LatentIntervals:
|
||||
"""Intervals which the latent tensor of given shape is split into.
|
||||
Each dimension of the latent space is split into intervals based on the length along said dimension.
|
||||
"""
|
||||
|
||||
original_shape: torch.Size
|
||||
dimension_intervals: tuple[DimensionIntervals, ...]
|
||||
|
||||
|
||||
# Operation to split a single dimension of the tensor into intervals based on the length along the dimension.
|
||||
SplitOperation = Callable[[int], DimensionIntervals]
|
||||
# Operation to map the intervals in input dimension to slices and masks along a corresponding output dimension.
|
||||
MappingOperation = Callable[[DimensionIntervals], tuple[list[slice], list[torch.Tensor | None]]]
|
||||
|
||||
|
||||
def default_split_operation(length: int) -> DimensionIntervals:
|
||||
return DimensionIntervals(intervals=[DimensionInterval(start=0, end=length, left_ramp=0, right_ramp=0)])
|
||||
|
||||
|
||||
DEFAULT_SPLIT_OPERATION: SplitOperation = default_split_operation
|
||||
|
||||
|
||||
def default_mapping_operation(
|
||||
_intervals: DimensionIntervals,
|
||||
) -> tuple[list[slice], list[torch.Tensor | None]]:
|
||||
return [slice(0, None)], [None]
|
||||
|
||||
|
||||
DEFAULT_MAPPING_OPERATION: MappingOperation = default_mapping_operation
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Split functions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def split_by_size(size: int, overlap: int) -> SplitOperation:
|
||||
"""Split a dimension into overlapping tiles of a given size.
|
||||
Tiles are sized ``size`` with ``overlap`` shared elements between
|
||||
consecutive tiles. The last tile may be shorter if the dimension
|
||||
doesn't divide evenly.
|
||||
Args:
|
||||
size: Target tile size (in axis units).
|
||||
overlap: Overlap between consecutive tiles.
|
||||
Returns:
|
||||
A split operation that divides a dimension into tiles.
|
||||
"""
|
||||
if size <= 0:
|
||||
raise ValueError(f"size must be > 0, got {size}")
|
||||
if overlap < 0 or overlap >= size:
|
||||
raise ValueError(f"overlap must satisfy 0 <= overlap < size, got overlap={overlap}, size={size}")
|
||||
|
||||
def split(dimension_size: int) -> DimensionIntervals:
|
||||
if dimension_size <= size:
|
||||
return DEFAULT_SPLIT_OPERATION(dimension_size)
|
||||
amount = (dimension_size + size - 2 * overlap - 1) // (size - overlap)
|
||||
intervals = [
|
||||
DimensionInterval(start=0, end=size, left_ramp=0, right_ramp=overlap),
|
||||
*(
|
||||
DimensionInterval(
|
||||
start=i * (size - overlap),
|
||||
end=i * (size - overlap) + size,
|
||||
left_ramp=overlap,
|
||||
right_ramp=overlap,
|
||||
)
|
||||
for i in range(1, amount - 1)
|
||||
),
|
||||
DimensionInterval(
|
||||
start=(amount - 1) * (size - overlap), end=dimension_size, left_ramp=overlap, right_ramp=0
|
||||
),
|
||||
]
|
||||
return DimensionIntervals(intervals=intervals)
|
||||
|
||||
return split
|
||||
|
||||
|
||||
def split_temporal_causal(size: int, overlap: int) -> SplitOperation:
|
||||
"""Split a temporal axis into overlapping tiles with causal handling.
|
||||
Each tile after the first is shifted back by 1 and its left ramp is
|
||||
increased by 1, ensuring causal continuity through the blend ramps.
|
||||
Args:
|
||||
size: Tile size in axis units.
|
||||
overlap: Overlap between tiles in the same units.
|
||||
Returns:
|
||||
Split operation that divides temporal dimension with causal handling.
|
||||
"""
|
||||
non_causal_split = split_by_size(size, overlap)
|
||||
|
||||
def split(dimension_size: int) -> DimensionIntervals:
|
||||
if dimension_size <= size:
|
||||
return DEFAULT_SPLIT_OPERATION(dimension_size)
|
||||
dim_intervals = non_causal_split(dimension_size)
|
||||
modified_intervals = [dim_intervals.intervals[0]] + [
|
||||
replace(interval, start=interval.start - 1, left_ramp=interval.left_ramp + 1)
|
||||
for interval in dim_intervals.intervals[1:]
|
||||
]
|
||||
return DimensionIntervals(intervals=modified_intervals)
|
||||
|
||||
return split
|
||||
|
||||
|
||||
def split_temporal(tile_size_frames: int, overlap_frames: int) -> SplitOperation:
|
||||
"""Split a temporal axis in video frame space into overlapping tiles.
|
||||
Args:
|
||||
tile_size_frames: Tile length in frames.
|
||||
overlap_frames: Overlap between consecutive tiles in frames.
|
||||
Returns:
|
||||
Split operation that takes frame count and returns DimensionIntervals in frame indices.
|
||||
"""
|
||||
non_causal_split = split_by_size(tile_size_frames, overlap_frames)
|
||||
|
||||
def split(dimension_size: int) -> DimensionIntervals:
|
||||
if dimension_size <= tile_size_frames:
|
||||
return DEFAULT_SPLIT_OPERATION(dimension_size)
|
||||
dim_intervals = non_causal_split(dimension_size)
|
||||
modified_intervals = [
|
||||
replace(interval, end=interval.end + 1, right_ramp=0) for interval in dim_intervals.intervals[:-1]
|
||||
] + [replace(dim_intervals.intervals[-1], right_ramp=0)]
|
||||
return DimensionIntervals(intervals=modified_intervals)
|
||||
|
||||
return split
|
||||
|
||||
|
||||
def split_by_count_temporal_causal(num_tiles: int, overlap: int = 0) -> SplitOperation:
|
||||
"""Split a temporal dimension by count with causal handling.
|
||||
Wraps :func:`split_by_count` with the same causal adjustment as
|
||||
:func:`split_temporal_causal`: each tile after the first is shifted
|
||||
back by 1 and its left ramp is increased by 1.
|
||||
Args:
|
||||
num_tiles: Number of tiles. Must be >= 1.
|
||||
overlap: Overlap between adjacent tiles (default 0).
|
||||
Returns:
|
||||
A split operation that divides a temporal dimension into tiles.
|
||||
"""
|
||||
non_causal_split = split_by_count(num_tiles, overlap)
|
||||
|
||||
def split(dimension_size: int) -> DimensionIntervals:
|
||||
dim_intervals = non_causal_split(dimension_size)
|
||||
if len(dim_intervals.intervals) <= 1:
|
||||
return dim_intervals
|
||||
modified_intervals = [dim_intervals.intervals[0]] + [
|
||||
replace(interval, start=interval.start - 1, left_ramp=interval.left_ramp + 1)
|
||||
for interval in dim_intervals.intervals[1:]
|
||||
]
|
||||
return DimensionIntervals(intervals=modified_intervals)
|
||||
|
||||
return split
|
||||
|
||||
|
||||
def split_by_count(num_tiles: int, overlap: int = 0) -> SplitOperation:
|
||||
"""Split a dimension into a given number of tiles with overlap.
|
||||
Computes the tile size as
|
||||
``(dim_size + overlap * (num_tiles - 1)) // num_tiles`` so that
|
||||
``num_tiles`` tiles of that size with ``overlap`` shared elements
|
||||
cover the dimension evenly. Delegates to :func:`split_by_size` for
|
||||
the actual interval construction.
|
||||
When the total ``dim_size + overlap * (num_tiles - 1)`` is not evenly
|
||||
divisible by ``num_tiles``, the first ``remainder`` tiles each absorb
|
||||
one extra unit.
|
||||
Args:
|
||||
num_tiles: Number of tiles. Must be >= 1.
|
||||
overlap: Overlap between adjacent tiles (default 0). Must be >= 0
|
||||
and less than the computed tile size.
|
||||
Returns:
|
||||
A split operation that divides a dimension into tiles.
|
||||
"""
|
||||
if num_tiles < 1:
|
||||
raise ValueError(f"num_tiles must be >= 1, got {num_tiles}")
|
||||
if overlap < 0:
|
||||
raise ValueError(f"overlap must be >= 0, got {overlap}")
|
||||
|
||||
def split(dim_size: int) -> DimensionIntervals:
|
||||
if num_tiles > dim_size:
|
||||
raise ValueError(
|
||||
f"num_tiles ({num_tiles}) exceeds dim_size ({dim_size}). Cannot assign at least 1 unit per tile."
|
||||
)
|
||||
if num_tiles == 1:
|
||||
return DEFAULT_SPLIT_OPERATION(dim_size)
|
||||
|
||||
total = dim_size + overlap * (num_tiles - 1)
|
||||
tile_size = total // num_tiles
|
||||
remainder = total % num_tiles
|
||||
|
||||
base_intervals = split_by_size(tile_size, overlap)(dim_size - remainder).intervals
|
||||
|
||||
# First `remainder` tiles each absorb 1 extra unit; shift subsequent boundaries.
|
||||
intervals: list[DimensionInterval] = []
|
||||
for i, iv in enumerate(base_intervals):
|
||||
shift = min(i, remainder)
|
||||
grow = 1 if i < remainder else 0
|
||||
intervals.append(replace(iv, start=iv.start + shift, end=iv.end + shift + grow))
|
||||
|
||||
return DimensionIntervals(intervals=intervals)
|
||||
|
||||
return split
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Mapping operations
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def identity_mapping_operation(intervals: DimensionIntervals) -> tuple[list[slice], list[torch.Tensor | None]]:
|
||||
"""Map each DimensionInterval to an output region at the same position, with trapezoidal blend masks.
|
||||
For every interval the output start/end matches the input start/end and a
|
||||
1-D blending mask is built from the interval's left_ramp and right_ramp.
|
||||
"""
|
||||
out_slices: list[slice] = []
|
||||
masks: list[torch.Tensor | None] = []
|
||||
for iv in intervals.intervals:
|
||||
out_slices.append(slice(iv.start, iv.end))
|
||||
masks.append(compute_trapezoidal_mask_1d(iv.end - iv.start, iv.left_ramp, iv.right_ramp))
|
||||
return out_slices, masks
|
||||
|
||||
|
||||
class Tile(NamedTuple):
|
||||
"""
|
||||
Represents a single tile.
|
||||
Attributes:
|
||||
in_coords:
|
||||
Tuple of slices specifying where to cut the tile from the INPUT tensor.
|
||||
out_coords:
|
||||
Tuple of slices specifying where this tile's OUTPUT should be placed in the reconstructed OUTPUT tensor.
|
||||
masks_1d:
|
||||
Per-dimension masks in OUTPUT units.
|
||||
These are used to create all-dimensional blending mask.
|
||||
Methods:
|
||||
blend_mask:
|
||||
Create a single N-D mask from the per-dimension masks.
|
||||
"""
|
||||
|
||||
in_coords: tuple[slice, ...]
|
||||
out_coords: tuple[slice, ...]
|
||||
masks_1d: tuple[torch.Tensor | None, ...]
|
||||
|
||||
@property
|
||||
def blend_mask(self) -> torch.Tensor:
|
||||
num_dims = len(self.out_coords)
|
||||
per_dimension_masks: list[torch.Tensor] = []
|
||||
|
||||
for dim_idx in range(num_dims):
|
||||
mask_1d = self.masks_1d[dim_idx]
|
||||
view_shape = [1] * num_dims
|
||||
if mask_1d is None:
|
||||
# Broadcast mask along this dimension (length 1).
|
||||
one = torch.ones(1)
|
||||
|
||||
view_shape[dim_idx] = 1
|
||||
per_dimension_masks.append(one.view(*view_shape))
|
||||
continue
|
||||
|
||||
# Reshape (L,) -> (1, ..., L, ..., 1) so masks across dimensions broadcast-multiply.
|
||||
view_shape[dim_idx] = mask_1d.shape[0]
|
||||
per_dimension_masks.append(mask_1d.view(*view_shape))
|
||||
|
||||
# Multiply per-dimension masks to form the full N-D mask (separable blending window).
|
||||
combined_mask = per_dimension_masks[0]
|
||||
for mask in per_dimension_masks[1:]:
|
||||
combined_mask = combined_mask * mask
|
||||
|
||||
return combined_mask
|
||||
|
||||
|
||||
def create_tiles_from_intervals_and_mappers(
|
||||
intervals: LatentIntervals,
|
||||
mappers: list[MappingOperation],
|
||||
) -> list[Tile]:
|
||||
full_dim_input_slices: list[list[slice]] = []
|
||||
full_dim_output_slices: list[list[slice]] = []
|
||||
full_dim_masks_1d: list[list[torch.Tensor | None]] = []
|
||||
for axis_index in range(len(intervals.original_shape)):
|
||||
dimension_intervals = intervals.dimension_intervals[axis_index]
|
||||
input_slices = [slice(interval.start, interval.end) for interval in dimension_intervals.intervals]
|
||||
output_slices, masks_1d = mappers[axis_index](dimension_intervals)
|
||||
n_intervals = len(input_slices)
|
||||
if len(output_slices) != n_intervals or len(masks_1d) != n_intervals:
|
||||
raise ValueError(
|
||||
f"Axis {axis_index}: mapper produced {len(output_slices)} output slices and "
|
||||
f"{len(masks_1d)} masks for {n_intervals} input intervals"
|
||||
)
|
||||
full_dim_input_slices.append(input_slices)
|
||||
full_dim_output_slices.append(output_slices)
|
||||
full_dim_masks_1d.append(masks_1d)
|
||||
|
||||
return [
|
||||
Tile(in_coords=in_coord, out_coords=out_coord, masks_1d=mask_1d)
|
||||
for in_coord, out_coord, mask_1d in zip(
|
||||
itertools.product(*full_dim_input_slices),
|
||||
itertools.product(*full_dim_output_slices),
|
||||
itertools.product(*full_dim_masks_1d),
|
||||
strict=True,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def create_tiles(
|
||||
latent_shape: torch.Size,
|
||||
splitters: list[SplitOperation],
|
||||
mappers: list[MappingOperation],
|
||||
) -> list[Tile]:
|
||||
if len(splitters) != len(latent_shape):
|
||||
raise ValueError(
|
||||
f"Number of splitters must be equal to number of dimensions in latent shape, "
|
||||
f"got {len(splitters)} and {len(latent_shape)}"
|
||||
)
|
||||
if len(mappers) != len(latent_shape):
|
||||
raise ValueError(
|
||||
f"Number of mappers must be equal to number of dimensions in latent shape, "
|
||||
f"got {len(mappers)} and {len(latent_shape)}"
|
||||
)
|
||||
intervals = [splitter(length) for splitter, length in zip(splitters, latent_shape, strict=True)]
|
||||
latent_intervals = LatentIntervals(original_shape=latent_shape, dimension_intervals=tuple(intervals))
|
||||
return create_tiles_from_intervals_and_mappers(latent_intervals, mappers)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Video-grid tiling configs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DimensionTilingConfig:
|
||||
"""Tiling parameters for a single dimension of the patchified grid.
|
||||
Attributes:
|
||||
num_tiles: Number of tiles along this dimension.
|
||||
overlap: Overlap between adjacent tiles, in latent grid units.
|
||||
Adjacent tiles share ``overlap`` grid cells at their
|
||||
boundary, producing an overlap zone blended with
|
||||
trapezoidal masks.
|
||||
"""
|
||||
|
||||
num_tiles: int
|
||||
overlap: int = 0
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.num_tiles < 1:
|
||||
raise ValueError(f"num_tiles must be >= 1, got {self.num_tiles}")
|
||||
if self.overlap < 0:
|
||||
raise ValueError(f"overlap must be >= 0, got {self.overlap}")
|
||||
|
||||
@classmethod
|
||||
def from_tile_size(cls, dim_size: int, tile_size: int, overlap: int = 0) -> DimensionTilingConfig:
|
||||
"""Create config by computing ``num_tiles`` from dimension size and tile size.
|
||||
Args:
|
||||
dim_size: Total length of the dimension.
|
||||
tile_size: Desired tile size.
|
||||
overlap: Overlap between consecutive tiles.
|
||||
Returns:
|
||||
A ``DimensionTilingConfig`` with the computed ``num_tiles``.
|
||||
"""
|
||||
split_op = split_by_size(tile_size, overlap)
|
||||
intervals = split_op(dim_size)
|
||||
return cls(num_tiles=len(intervals.intervals), overlap=overlap)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TileCountConfig:
|
||||
"""Tiling layout for a ``(F, H, W)`` grid.
|
||||
Specifies tile *counts* per dimension (as opposed to tile *sizes*
|
||||
which are used by the single-GPU VAE ``TilingConfig``).
|
||||
Attributes:
|
||||
frames: Tiling along the temporal (frames) dimension.
|
||||
height: Tiling along the latent height dimension.
|
||||
width: Tiling along the latent width dimension.
|
||||
"""
|
||||
|
||||
frames: DimensionTilingConfig = DimensionTilingConfig(num_tiles=1, overlap=0)
|
||||
height: DimensionTilingConfig = DimensionTilingConfig(num_tiles=1, overlap=0)
|
||||
width: DimensionTilingConfig = DimensionTilingConfig(num_tiles=1, overlap=0)
|
||||
Reference in New Issue
Block a user