Automated PR - 2026-06-17
This commit is contained in:
@@ -8,6 +8,7 @@ The foundational library for the LTX-2 Audio-Video generation model. This packag
|
||||
- **`conditioning/`**: Tools for preparing latent states and applying conditioning (image, video, keyframes)
|
||||
- **`guidance/`**: Perturbation system for fine-grained control over attention mechanisms
|
||||
- **`loader/`**: Utilities for loading weights from `.safetensors`, fusing LoRAs, and managing memory
|
||||
- **`block_streaming/`**: Memory-efficient inference that streams transformer blocks through the GPU one at a time (from pinned CPU buffers or directly from disk)
|
||||
- **`model/`**: PyTorch implementations of the LTX-2 Transformer, Video VAE, Audio VAE, Vocoder and Upscaler
|
||||
- **`text_encoders/gemma`**: Gemma text encoder implementation with tokenizers, feature extractors, and separate encoders for audio-video and video-only generation
|
||||
- **`quantization/`**: FP8 quantization backends (FP8-TensorRT-LLM scaled MM, FP8 cast) for reduced memory footprint.
|
||||
@@ -55,6 +56,7 @@ pip install -e packages/ltx-core
|
||||
|
||||
- **Loader** ([`loader/`](src/ltx_core/loader/)): Model loading from `.safetensors`, LoRA fusion, weight remapping, and memory management
|
||||
- **Quantization** ([`quantization/`](src/ltx_core/quantization/)): FP8 quantization backends for reduced memory footprint and faster inference
|
||||
- **Block Streaming** ([`block_streaming/`](src/ltx_core/block_streaming/)): Streams transformer blocks through the GPU one block at a time, so the full model runs on machines without enough memory to hold all its weights at once
|
||||
|
||||
### Loader
|
||||
|
||||
@@ -157,6 +159,39 @@ from ltx_core.quantization.fp8_cast import build_policy as build_fp8_cast_policy
|
||||
policy = build_fp8_cast_policy("/path/to/checkpoint.safetensors")
|
||||
```
|
||||
|
||||
### Block Streaming
|
||||
|
||||
The `block_streaming/` module ([`src/ltx_core/block_streaming/`](src/ltx_core/block_streaming/)) lets the full model run on machines that lack the memory to hold all of its weights at once. It streams the transformer's blocks through a small rolling set of GPU buffers, loading each block's weights just before it runs and recycling them afterwards, so only a few blocks are resident on the GPU at any moment. Construct it with `StreamingModelBuilder`, which returns a `BlockStreamingWrapper` -- an `nn.Module` drop-in for the wrapped model.
|
||||
|
||||
#### Strategies
|
||||
|
||||
The strategy is chosen automatically from `cpu_slots_count` relative to the number of blocks:
|
||||
|
||||
- **RAM streaming** (default, `cpu_slots_count` omitted or `>= num_blocks`): all blocks are pre-loaded into pinned CPU buffers (with LoRA fusion) at build time, then copied to the GPU on demand. Fast; higher CPU memory.
|
||||
- **Disk streaming** (`cpu_slots_count < num_blocks`): blocks are read from the `.safetensors` file on demand on a background worker thread. Slower; lowest CPU memory.
|
||||
|
||||
#### Basic usage
|
||||
|
||||
```python
|
||||
import torch
|
||||
from ltx_core.block_streaming import StreamingModelBuilder
|
||||
|
||||
builder = StreamingModelBuilder(
|
||||
model_class_configurator=MyModelConfigurator,
|
||||
model_path="/path/to/model.safetensors",
|
||||
blocks_attr="transformer_blocks", # dotted path to the nn.ModuleList
|
||||
blocks_prefix="transformer_blocks", # state-dict key prefix for block weights
|
||||
)
|
||||
|
||||
# Omit cpu_slots_count for RAM streaming; pass a value < num_blocks for disk streaming.
|
||||
model = builder.build(
|
||||
device=torch.device("cuda"),
|
||||
dtype=torch.bfloat16,
|
||||
cpu_slots_count=4,
|
||||
gpu_slots_count=2,
|
||||
)
|
||||
```
|
||||
|
||||
For complete, production-ready pipeline implementations that combine these building blocks, see the [`ltx-pipelines`](../ltx-pipelines/) package.
|
||||
|
||||
---
|
||||
|
||||
Reference in New Issue
Block a user