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

Automated PR - 2026-07-07

Browse Source
This commit is contained in:
github-actions[bot]
2026-07-07 16:57:50 +00:00
parent 780984275f
commit 63fd9a4f86
157 changed files with 15976 additions and 5043 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-pipelines/CLAUDE.md
+1 -1
View File
@@ -49,7 +49,7 @@ Inference pipelines for LTX-2 audio-video generation. Depends on `ltx-core` for
## Shared building blocks (`utils/blocks.py`)
- `DiffusionStage` -- owns transformer lifecycle; builds model on call, frees on exit via `gpu_model()` context manager (moves params to meta device to release GPU/CPU memory). Accepts optional `stepper` and `loop` overrides.
- `DiffusionStage` -- owns transformer lifecycle; builds model on call, frees on exit via `gpu_model()` context manager (moves params to meta device to release GPU/CPU memory). Accepts optional `stepper` and `loop` overrides. `__init__` takes a pre-built transformer builder; pipelines construct it via the `DiffusionStage.from_checkpoint(checkpoint_path, ..., loras=...)` classmethod, which builds the standard (and, when offloading, streaming) builders. `with_builder` / `with_loras` return a new stage with a swapped builder / LoRA set without re-specifying config.
- `PromptEncoder` -- Gemma text encoder + embeddings processor (video 4096-dim, audio 2048-dim).
- `ImageConditioner` / `AudioConditioner` -- temporary encoder scope; builds encoder, passes to callable, frees.
- `VideoUpsampler` -- 2x spatial upsampling via encoder + upsampler.
packages/ltx-pipelines/README.md
+15 -559
View File
@@ -1,15 +1,9 @@
# LTX-2 Pipelines
High-level pipeline implementations for generating audio-video content with Lightricks' **LTX-2** model. This package provides ready-to-use pipelines for text-to-video, image-to-video, video-to-video, and keyframe interpolation tasks.
High-level pipeline implementations for generating audio-video content with Lightricks' **LTX-2** model. This package provides ready-to-use pipelines for text-to-video, image-to-video, video-to-video, audio-to-video, keyframe interpolation, and retake tasks.
Pipelines are built using building blocks from [`ltx-core`](../ltx-core/) (schedulers, guiders, noisers, patchifiers) and handle the complete inference flow including model loading, encoding, decoding, and file I/O.
---
## 📋 Overview
LTX-2 Pipelines provides production-ready implementations that abstract away the complexity of the diffusion process, model loading, and memory management. Each pipeline is optimized for specific use cases and offers different trade-offs between speed, quality, and memory usage.
**Key Features:**
- 🎬 **Multiple Pipeline Types**: Text-to-video, image-to-video, video-to-video, audio-to-video, keyframe interpolation, and retake
@@ -19,27 +13,12 @@ LTX-2 Pipelines provides production-ready implementations that abstract away the
- 📦 **Self-Contained**: Handles model loading, encoding, decoding, and file I/O
- 🚀 **CLI Support**: All pipelines can be run as command-line scripts
---
## 🚀 Quick Start
`ltx-pipelines` provides ready-made inference pipelines for text-to-video, image-to-video, video-to-video, audio-to-video, keyframe interpolation, and retake. Built using building blocks from [`ltx-core`](../ltx-core/), these pipelines handle the complete inference flow including model loading, encoding, decoding, and file I/O.
## 🔧 Installation
## Quick Start
```bash
# From the repository root
uv sync --frozen
# Or install as a package
pip install -e packages/ltx-pipelines
```
### Running Pipelines
All pipelines can be run directly from the command line. Each pipeline module is executable:
```bash
# Run a pipeline (example: two-stage text-to-video)
python -m ltx_pipelines.ti2vid_two_stages \
--checkpoint-path path/to/checkpoint.safetensors \
@@ -48,544 +27,21 @@ python -m ltx_pipelines.ti2vid_two_stages \
--gemma-root path/to/gemma \
--prompt "A beautiful sunset over the ocean" \
--output-path output.mp4
# View all available options for any pipeline
python -m ltx_pipelines.ti2vid_two_stages --help
```
Available pipeline modules:
- `ltx_pipelines.ti2vid_two_stages` - Two-stage text/image-to-video (recommended).
- `ltx_pipelines.ti2vid_two_stages_hq` - Two-stage text/image-to-video (different sampler, better quality).
- `ltx_pipelines.ti2vid_one_stage` - Single-stage text/image-to-video.
- `ltx_pipelines.t2a_one_stage` - Single-stage text-to-audio (audio-only output).
- `ltx_pipelines.distilled` - Fast text/image-to-video pipeline using only the distilled model.
- `ltx_pipelines.ic_lora` - Video-to-video with IC-LoRA.
- `ltx_pipelines.keyframe_interpolation` - Keyframe interpolation.
- `ltx_pipelines.a2vid_two_stage` - Audio-to-video generation conditioned on an input audio.
- `ltx_pipelines.retake` - Regenerate a time region of an existing video.
- `ltx_pipelines.hdr_ic_lora` - Video-to-video with HDR output (linear float via LogC3 inverse decode).
- `ltx_pipelines.lipdub` - Lip dubbing / re-voicing with IC-LoRA and audio reference conditioning.
Use `--help` with any pipeline module to see all available options and parameters.
---
## 🎯 Pipeline Selection Guide
### Quick Decision Tree
```text
Do you have an existing video to modify?
├─ YES → Use RetakePipeline (regenerate a specific time region)
│
Do you have an audio file to drive generation?
├─ YES → Use A2VidPipelineTwoStage (audio-to-video)
│
Do you need HDR output (linear float frames for EXR / tonemapping)?
├─ YES → Use HDRICLoraPipeline (video-to-video with LogC3 inverse decode)
│
Do you need to condition on existing images/videos?
├─ YES → Do you have reference videos for video-to-video?
│ ├─ YES → Use ICLoraPipeline
│ └─ NO → Do you have multiple keyframe images to interpolate?
│ ├─ YES → Use KeyframeInterpolationPipeline
│ └─ NO → Use TI2VidTwoStagesPipeline (image conditioning only)
│
└─ NO → Text-to-video only
├─ Do you need best quality?
│ └─ YES → Use TI2VidTwoStagesPipeline (recommended for production)
│
└─ Do you need fastest inference?
└─ YES → Use DistilledPipeline (with 8 predefined sigmas)
```
> **Note:** [`TI2VidOneStagePipeline`](src/ltx_pipelines/ti2vid_one_stage.py) is primarily for educational purposes. For best quality, use two-stage pipelines ([`TI2VidTwoStagesPipeline`](src/ltx_pipelines/ti2vid_two_stages.py), [`TI2VidTwoStagesHQPipeline`](src/ltx_pipelines/ti2vid_two_stages_hq.py), [`ICLoraPipeline`](src/ltx_pipelines/ic_lora.py), [`KeyframeInterpolationPipeline`](src/ltx_pipelines/keyframe_interpolation.py), [`A2VidPipelineTwoStage`](src/ltx_pipelines/a2vid_two_stage.py), or [`DistilledPipeline`](src/ltx_pipelines/distilled.py)). For editing existing videos, use [`RetakePipeline`](src/ltx_pipelines/retake.py).
### Features Comparison
| Pipeline | Stages | [Multimodal Guidance](#%EF%B8%8F-multimodal-guidance) | Upsampling | Conditioning | Best For |
| -------- | ------ | --- | ---------- | ------------- | -------- |
| **TI2VidTwoStagesPipeline** | 2 | ✅ | ✅ | Image | **Production quality** (recommended) |
| **TI2VidTwoStagesHQPipeline** | 2 | ✅ | ✅ | Image | Same as above, res_2s sampler (higher quality) |
| **TI2VidOneStagePipeline** | 1 | ✅ | ❌ | Image | Educational, prototyping |
| **DistilledPipeline** | 2 | ❌ | ✅ | Image | Fastest inference (8 sigmas) |
| **ICLoraPipeline** | 2 | ✅ | ✅ | Image + Video | Video-to-video transformations |
| **KeyframeInterpolationPipeline** | 2 | ✅ | ✅ | Keyframes | Animation, interpolation |
| **A2VidPipelineTwoStage** | 2 | ✅ | ✅ | Audio + Image | Audio-driven video generation |
| **RetakePipeline** | 1 | ✅ | ❌ | Source Video | Regenerating a time region of a video |
| **HDRICLoraPipeline** | 2 | ❌ | ✅ | Video | HDR video-to-video (linear float output for EXR) |
| **LipDubPipeline** | 2 | ✅ | ✅ | Video + Audio | Lip dubbing with audio ref conditioning |
---
## 📦 Available Pipelines
### 1. TI2VidTwoStagesPipeline
**Best for:** High-quality text/image-to-video generation with upsampling. **Recommended for production use.**
**Source**: [`src/ltx_pipelines/ti2vid_two_stages.py`](src/ltx_pipelines/ti2vid_two_stages.py)
Two-stage generation: Stage 1 generates low-resolution video with [multimodal guidance](#%EF%B8%8F-multimodal-guidance), Stage 2 upsamples to 2x resolution with distilled LoRA refinement. Supports image conditioning. Highest quality output, slower than one-stage but significantly better quality.
**Use when:** Production-quality video generation, higher resolution needed, quality over speed, text-to-video with image conditioning.
---
### 2. TI2VidTwoStagesHQPipeline
**Best for:** Same two-stage text/image-to-video as TI2VidTwoStagesPipeline but with a different sampler and step count.
**Source**: [`src/ltx_pipelines/ti2vid_two_stages_hq.py`](src/ltx_pipelines/ti2vid_two_stages_hq.py)
Uses the **res_2s** second-order sampler instead of Euler. Same stage structure (stage 1 at target resolution with CFG, stage 2 upsampling with distilled LoRA) and image conditioning support. Typically allows fewer steps for comparable quality; trade-offs differ from the default Euler-based pipeline.
**Use when:** You want the same two-stage workflow with fewer steps or prefer the res_2s sampling behavior.
---
### 3. TI2VidOneStagePipeline
**Best for:** Educational purposes and quick prototyping.
**Source**: [`src/ltx_pipelines/ti2vid_one_stage.py`](src/ltx_pipelines/ti2vid_one_stage.py)
> **⚠️ Important:** This pipeline is primarily for educational purposes. For production-quality results, use `TI2VidTwoStagesPipeline` or other two-stage pipelines.
Single-stage generation (no upsampling) with [multimodal guidance](#%EF%B8%8F-multimodal-guidance) and image conditioning support. Faster inference but lower resolution output (typically 512x768).
**Use when:** Learning how the pipeline works, quick prototyping, testing, or when high resolution is not needed.
---
### 4. DistilledPipeline
**Best for:** Fastest inference with good quality using a distilled model with predefined sigma schedule.
**Source**: [`src/ltx_pipelines/distilled.py`](src/ltx_pipelines/distilled.py)
Two-stage generation with 8 predefined sigmas (8 steps in stage 1, 4 steps in stage 2). No guidance required. Fastest inference among all pipelines. Supports image conditioning. Requires spatial upsampler.
**Use when:** Fastest inference is critical, batch processing many videos, or when you have a distilled model checkpoint.
---
### 5. ICLoraPipeline
**Best for:** Video-to-video and image-to-video transformations using IC-LoRA.
**Source**: [`src/ltx_pipelines/ic_lora.py`](src/ltx_pipelines/ic_lora.py)
Two-stage generation with IC-LoRA support. Can condition on reference videos (video-to-video) or images at specific frames. CFG guidance in stage 1, upsampling in stage 2. Requires IC-LoRA trained model.
**Note:** ICLoraPipeline can only be used with a distilled model.
**Use when:** Video-to-video transformations, image-to-video with strong control, or when you have reference videos to guide generation.
---
### 6. KeyframeInterpolationPipeline
**Best for:** Generating videos by interpolating between keyframe images.
**Source**: [`src/ltx_pipelines/keyframe_interpolation.py`](src/ltx_pipelines/keyframe_interpolation.py)
Two-stage generation with keyframe interpolation. Uses guiding latents (additive conditioning) instead of replacing latents for smoother transitions. [Multimodal guidance](#%EF%B8%8F-multimodal-guidance) in stage 1, upsampling in stage 2.
**Use when:** You have keyframe images and want to interpolate between them, creating smooth transitions, or animation/motion interpolation tasks.
---
### 7. A2VidPipelineTwoStage
**Best for:** Generating video driven by an input audio.
**Source**: [`src/ltx_pipelines/a2vid_two_stage.py`](src/ltx_pipelines/a2vid_two_stage.py)
Two-stage audio-to-video generation. Stage 1 generates video at half resolution with audio conditioning (video-only denoising with the audio frozen), then Stage 2 upsamples by 2x and refines the video while keeping the audio fixed, using a distilled LoRA. The input audio is encoded via the audio VAE and used as the initial audio latent, but the original audio waveform is passed through and returned in the output to preserve fidelity. Supports image conditioning and prompt enhancement.
**Extra CLI arguments:** `--audio-path` (required), `--audio-start-time`, `--audio-max-duration`.
**Use when:** You have an audio clip and want to generate a matching video, audio-reactive video generation, or music visualization.
---
### 8. RetakePipeline
**Best for:** Regenerating a specific time region of an existing video while keeping the rest unchanged.
**Source**: [`src/ltx_pipelines/retake.py`](src/ltx_pipelines/retake.py)
Single-stage generation that encodes the source video and audio into latents, applies a temporal region mask to mark `[start_time, end_time]` for regeneration, and denoises only the masked region from a text prompt. Content outside the time window is preserved. Supports independent control over video and audio regeneration (`regenerate_video`, `regenerate_audio` flags), and can use either the full model with CFG guidance or the distilled model with a fixed sigma schedule.
**Extra CLI arguments:** `--video-path` (required), `--start-time` (required), `--end-time` (required).
**Constraints:** Source video frame count must satisfy the 8k+1 format (e.g. 97, 193) and resolution must be multiples of 32.
**Use when:** You want to re-do a specific section of a generated video (e.g. fix a bad segment), selectively regenerate audio or video in a time window, or iterate on part of a result without re-generating the entire clip.
---
### 9. HDRICLoraPipeline
**Best for:** Video-to-video generation with HDR output for EXR export and offline tonemapping.
**Source**: [`src/ltx_pipelines/hdr_ic_lora.py`](src/ltx_pipelines/hdr_ic_lora.py)
Two-stage video-to-video on the distilled model with an HDR IC-LoRA. Decoded latents pass through an HDR inverse transform (ARRI LogC3, auto-detected from LoRA metadata) to produce a **linear HDR float** tensor `[f, h, w, c]`. Video-only (audio skipped). Text embeddings are pre-computed externally and loaded from a `.safetensors` file. Tonemapping and EXR saving are the caller's responsibility. LoRA and embeddings: [`Lightricks/LTX-2.3-22b-IC-LoRA-HDR`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-HDR).
**Extra CLI arguments:** `--input` (mp4 or directory, required), `--output-dir` (required), `--hdr-lora` (required), `--text-embeddings` (pre-computed `.safetensors`, required), `--num-frames`, `--spatial-tile` (tiled VAE decode tile size; reduce on lower-VRAM GPUs), `--skip-mp4` (EXR only, no H.264 preview), `--exr-half` (float16 EXR), `--high-quality` (generates 2x frames internally for smoother output, ~2x slower), `--offload {none,cpu,disk}` (weight offloading; disables FP8 quantization when not `none`).
**Use when:** You need linear HDR float output for EXR export, color grading, or custom tonemapping workflows.
---
### 10. LipDubPipeline
**Best for:** Lip dubbing, rephrasing while keeping the same speaker identity and matching lip movements to new audio.
**Source**: [`src/ltx_pipelines/lipdub.py`](src/ltx_pipelines/lipdub.py)
Uses IC-LoRA on a **distilled** checkpoint with a **single** lip-dub IC-LoRA applied in **both** stages. The reference clip provides video and audio reference tokens whose VAE latents are appended to the target audio sequence as frozen reference tokens. The frame count and frame rate are derived from the reference video (frame count is silently snapped to the nearest `8k+1`), so the CLI does not accept `--num-frames` or `--frame-rate`. Required: `--reference-video`. Optional: `--reference-strength`. LoRA: [`Lightricks/LTX-2.3-22b-IC-LoRA-LipDub`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-LipDub).
**Note:** Requires a distilled model checkpoint and one lip-dub IC-LoRA (`--lora` exactly once).
**Use when:** Dubbing, rephrasing with matched lips and speaker identity.
---
### 11. T2AOneStagePipeline
**Best for:** Text-to-audio — generating speech/audio only (no video) from a text prompt, e.g. driving an audio-style LoRA such as an accent LoRA.
**Source**: [`src/ltx_pipelines/t2a_one_stage.py`](src/ltx_pipelines/t2a_one_stage.py)
Single-stage, **audio-only** generation: the video branch is absent (`video=None`), so only the audio modality is denoised and decoded through the audio VAE + vocoder, producing a wave file. Audio duration is derived from `--num-frames` / `--frame-rate` (the same `8k+1` frame convention as video). Audio guidance (CFG/STG) is optional — the `--audio-*` flags default to the model's values; the video→audio cross-modal guidance is disabled since there is no video modality.
**Extra CLI arguments (all optional, with sensible defaults):** `--num-frames`, `--frame-rate`, `--negative-prompt`, `--audio-cfg-guidance-scale`, `--audio-stg-guidance-scale`, `--audio-stg-blocks`, `--audio-rescale-scale`, `--audio-skip-step`. No `--height/--width/--image` (audio has no spatial dimensions).
**Use when:** You need speech/audio from text alone, or to evaluate an audio-only LoRA (accent, voice style) without generating video.
---
## 🎨 Conditioning Types
Pipelines use different conditioning methods from [`ltx-core`](../ltx-core/) for controlling generation. See the [ltx-core conditioning documentation](../ltx-core/README.md#conditioning--control) for details.
### Image Conditioning
All pipelines support image conditioning, but with different methods:
- **Replacing Latents** ([`image_conditionings_by_replacing_latent`](src/ltx_pipelines/utils/helpers.py)):
- Used by: `TI2VidOneStagePipeline`, `TI2VidTwoStagesPipeline`, `DistilledPipeline`, `ICLoraPipeline`
- Replaces the latent at a specific frame with the encoded image
- Strong control over specific frames
- **Guiding Latents** ([`image_conditionings_by_adding_guiding_latent`](src/ltx_pipelines/utils/helpers.py)):
- Used by: `KeyframeInterpolationPipeline`
- Adds the image as a guiding signal rather than replacing
- Better for smooth interpolation between keyframes
### Video Conditioning
- **Video Conditioning** (ICLoraPipeline only):
- Conditions on entire reference videos
- Useful for video-to-video transformations
- Uses `VideoConditionByKeyframeIndex` from [`ltx-core`](../ltx-core/)
---
## 🎛️ Multimodal Guidance
LTX-2 pipelines use **multimodal guidance** to steer the diffusion process for both video and audio modalities. Each modality (video, audio) has its own guider with independent parameters, allowing fine-grained control over generation quality and adherence to prompts.
### Guidance Parameters
The `MultiModalGuiderParams` dataclass controls guidance behavior:
| Parameter | Description |
| --------- | ----------- |
| `cfg_scale` | **Classifier-Free Guidance** scale. Higher values make the output adhere more strongly to the text prompt. Typical values: 2.0–5.0. Set to **1.0** to disable. |
| `stg_scale` | **Spatio-Temporal Guidance** scale. Controls perturbation-based guidance for improved temporal coherence. Typical values: 0.5–1.5. Set to **0.0** to disable. |
| `stg_blocks` | Which transformer blocks to perturb for STG (e.g., `[29]` for the last block). Set to **`[]`** to disable STG. |
| `rescale_scale` | Rescales the guided prediction to match the variance of the conditional prediction. Helps prevent over-saturation. Typical values: 0.5–0.7. Set to **0.0** to disable. |
| `modality_scale` | **Modality CFG** scale. Steers the model away from unsynced video and audio results, improving audio-visual coherence. Set to **1.0** to disable. |
| `skip_step` | Skip guidance every N steps. Can speed up inference with minimal quality loss. Set to **0** to disable (never skip). |
### How It Works
The multimodal guider combines three guidance signals during each denoising step:
1. **CFG (Text Guidance)**: Steers generation toward the text prompt by computing `(cond - uncond_text)`.
2. **STG (Perturbation Guidance)**: Improves structural coherence by perturbing specific transformer blocks and steering away from the perturbed prediction.
3. **Modality CFG**: For joint audio-video generation, steers the model away from unsynced video and audio results.
### Example Configuration
```python
from ltx_core.components.guiders import MultiModalGuiderParams
# Video guider: moderate CFG, STG enabled, modality isolation
video_guider_params = MultiModalGuiderParams(
cfg_scale=3.0,
stg_scale=1.0,
rescale_scale=0.7,
modality_scale=3.0,
stg_blocks=[29],
)
# Audio guider: higher CFG for stronger prompt adherence
audio_guider_params = MultiModalGuiderParams(
cfg_scale=7.0,
stg_scale=1.0,
rescale_scale=0.7,
modality_scale=3.0,
stg_blocks=[29],
)
```
> **Tip:** Start with the default values from [`constants.py`](src/ltx_pipelines/utils/constants.py) and adjust based on your use case. Higher `cfg_scale` = stronger prompt adherence but potentially less natural motion; higher `stg_scale` = better temporal coherence but slower inference (requires extra forward passes).
>
> **Tip:** When generating video with audio, set `modality_scale` > 1.0 (e.g., 3.0) to improve audio-visual sync. If generating video-only, set it to 1.0 to disable.
---
## ⚡ Optimization Tips
### Memory Optimization
**FP8 Quantization (Lower Memory Footprint):**
For smaller GPU memory footprint, use the `--quantization` flag and set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`.
Two quantization policies are available:
| Policy | CLI Flag | Description |
| ------ | -------- | ----------- |
| **FP8 Cast** | `--quantization fp8-cast` | Downcasts transformer linear weights to FP8 during loading; upcasts on the fly during inference. No extra dependencies. |
| **FP8 Scaled MM** | `--quantization fp8-scaled-mm` | Uses FP8 scaled matrix multiplication via TensorRT-LLM (`tensorrt_llm` must be installed). Best performance on Hopper GPUs. |
**CLI:**
```bash
# FP8 Cast (works on any GPU with FP8 support)
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python -m ltx_pipelines.ti2vid_two_stages \
--quantization fp8-cast --checkpoint-path=...
# FP8 Scaled MM (requires tensorrt_llm, best on Hopper GPUs)
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python -m ltx_pipelines.ti2vid_two_stages \
--quantization fp8-scaled-mm --checkpoint-path=...
```
**Programmatically:**
When authoring custom scripts, pass a `QuantizationPolicy` to pipeline classes:
```python
from ltx_core.quantization.fp8_cast import build_policy as build_fp8_cast_policy
# Alternative:
# from ltx_core.quantization.fp8_scaled_mm import build_policy as build_fp8_scaled_mm_policy
pipeline = TI2VidTwoStagesPipeline(
checkpoint_path=ltx_model_path,
distilled_lora=distilled_lora,
spatial_upsampler_path=upsampler_path,
gemma_root=gemma_root_path,
loras=[],
quantization=build_fp8_cast_policy(ltx_model_path),
)
pipeline(...)
```
You still need to use `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` when launching:
```bash
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python my_denoising_pipeline.py
```
**Memory Cleanup Between Stages:**
By default, pipelines clean GPU memory (especially transformer weights) between stages. If you have enough memory, you can skip this cleanup to reduce running time:
```python
# In pipeline implementations, memory cleanup happens automatically
# between stages. For custom pipelines, you can skip:
# utils.cleanup_memory() # Comment out if you have enough VRAM
```
### Compilation (`torch.compile`)
Compiling the transformer blocks with `torch.compile` speeds up inference. It is **opt-in and off by default**. The blocks are compiled shape-polymorphically (the sequence dimension is marked dynamic), so one compiled artifact serves any token count without recompiling.
**CLI** — the `--compile` flag maps directly to `CompilationConfig`:
| Form | Result |
| ---- | ------ |
| *(flag absent)* | eager, no compilation |
| `--compile` | compile with defaults |
| `--compile KEY=VALUE ...` | compile, overriding individual fields |
```bash
# Defaults
python -m ltx_pipelines.ti2vid_two_stages --compile --checkpoint-path=...
# reduce-overhead captures CUDA graphs -- the main latency lever for the denoising loop.
# Off by default because graph capture reserves static memory pools (extra VRAM), so it
# trades memory for speed; enable it when you have headroom.
python -m ltx_pipelines.ti2vid_two_stages --compile mode=reduce-overhead --checkpoint-path=...
# Several overrides at once
python -m ltx_pipelines.ti2vid_two_stages \
--compile mode=max-autotune fullgraph=true dynamic=true --checkpoint-path=...
```
| Field | Values | Default | Notes |
| ----- | ------ | ------- | ----- |
| `mode` | `none`, `reduce-overhead`, `max-autotune`, … | `none` | `reduce-overhead`/`max-autotune` enable CUDA graphs |
| `backend` | `inductor`, `eager`, … | `inductor` | |
| `fullgraph` | `true`/`false` | `false` | |
| `dynamic` | `auto`/`true`/`false` | `auto` | the seq dim is marked dynamic regardless |
| `inductor_config` | JSON object or path to a `.json` | `{}` | `torch._inductor.config` overrides |
| `dynamo_config` | JSON object or path to a `.json` | `{"inline_inbuilt_nn_modules": true, "cache_size_limit": 256}` | `torch._dynamo.config` overrides |
**Controlling inductor / dynamo configs.** `inductor_config` and `dynamo_config` take either an inline JSON object or a path to a `.json` file, applied via `torch._inductor.config.patch(...)` / `torch._dynamo.config.patch(...)` around the compiled forward. They **replace the defaults wholesale — they do not merge**, so when overriding `dynamo_config` re-include any defaults you want to keep:
```bash
python -m ltx_pipelines.ti2vid_two_stages \
--compile 'inductor_config={"max_autotune": true}' \
'dynamo_config={"inline_inbuilt_nn_modules": true, "cache_size_limit": 256, "recompile_limit": 32}' \
--checkpoint-path=...
```
**Programmatically**, pass a `CompilationConfig` to the pipeline:
```python
from ltx_core.model.transformer.compiling import CompilationConfig
pipeline = TI2VidTwoStagesPipeline(
...,
compilation_config=CompilationConfig(mode="reduce-overhead"),
)
```
**Faster cache loads: `unsafe_skip_cache_dynamic_shape_guards` (unsafe, opt-in).** Inductor's FX-graph cache re-checks the dynamic-shape guards stored with each entry on every lookup. Setting this flag skips that re-check (every entry is treated as a guard hit), which speeds up warm and cross-process cache loads. It is **not enabled by default** because it is a correctness hazard: a kernel first compiled at a small sequence length keeps int32 address arithmetic, and reusing it at a larger sequence length (roughly **>58k tokens/rank**) overflows int32 and reads out of bounds — surfacing as a CUDA illegal memory access or silently corrupted output. Only enable it when your token counts stay within the range the cached kernels were compiled for:
```bash
python -m ltx_pipelines.ti2vid_two_stages \
--compile 'inductor_config={"unsafe_skip_cache_dynamic_shape_guards": true}' \
--checkpoint-path=...
```
### Denoising Loop Optimization
**Gradient Estimation Denoising Loop:**
Instead of the standard Euler denoising loop, you can use gradient estimation for fewer steps (~20-30 instead of 40):
```python
from ltx_pipelines.utils import gradient_estimating_euler_denoising_loop
# Use gradient estimation denoising loop
def denoising_loop(sigmas, video_state, audio_state, stepper):
return gradient_estimating_euler_denoising_loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
transformer=transformer,
denoiser=denoiser,
ge_gamma=2.0, # Gradient estimation coefficient
)
```
This allows you to use **20-30 steps instead of 40** while maintaining quality. The gradient estimation function is defined in [`samplers.py`](src/ltx_pipelines/utils/samplers.py).
---
## 🔧 Requirements
- **LTX-2 Model Checkpoint** - Local `.safetensors` file
- **Gemma Text Encoder** - Local Gemma model directory
- **Spatial Upscaler** - Required for two-stage pipelines (except one-stage)
- **Distilled LoRA** - Required for two-stage pipelines (except one-stage and distilled)
---
## 📖 Example: Image-to-Video
```python
from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
from ltx_pipelines.utils.args import ImageConditioningInput
from ltx_pipelines.utils.media_io import encode_video
distilled_lora = [
LoraPathStrengthAndSDOps(
"/path/to/distilled_lora.safetensors",
0.6,
LTXV_LORA_COMFY_RENAMING_MAP,
),
]
pipeline = TI2VidTwoStagesPipeline(
checkpoint_path="/path/to/checkpoint.safetensors",
distilled_lora=distilled_lora,
spatial_upsampler_path="/path/to/upsampler.safetensors",
gemma_root="/path/to/gemma",
loras=[],
)
video_guider_params = MultiModalGuiderParams(
cfg_scale=3.0,
stg_scale=1.0,
rescale_scale=0.7,
modality_scale=3.0,
skip_step=0,
stg_blocks=[29],
)
audio_guider_params = MultiModalGuiderParams(
cfg_scale=7.0,
stg_scale=1.0,
rescale_scale=0.7,
modality_scale=3.0,
skip_step=0,
stg_blocks=[29],
)
# Generate video from image. The pipeline returns (video_iterator, audio);
# the caller is responsible for encoding to file via encode_video().
num_frames = 121
frame_rate = 25.0
tiling_config = TilingConfig.default()
video, audio = pipeline(
prompt="A serene landscape with mountains in the background",
negative_prompt="worst quality, low quality, blurry, distorted",
seed=42,
height=512,
width=768,
num_frames=num_frames,
frame_rate=frame_rate,
num_inference_steps=40,
video_guider_params=video_guider_params,
audio_guider_params=audio_guider_params,
images=[ImageConditioningInput("input_image.jpg", 0, 1.0, 33)], # path, frame_idx=0, strength=1.0, crf=33
tiling_config=tiling_config,
)
encode_video(
video=video,
fps=frame_rate,
audio=audio,
output_path="output.mp4",
video_chunks_number=get_video_chunks_number(num_frames, tiling_config),
)
```
---
See [Installation & Usage](docs/installation.md) for full setup, CLI modules, and shared flags.
## 📚 Documentation
| Topic | Description |
| ----- | ----------- |
| [Installation & Usage](docs/installation.md) | Install, requirements, running pipelines from the CLI, common flags |
| [Pipeline Selection Guide](docs/pipeline-selection.md) | Decision tree + feature comparison to pick the right pipeline |
| [Available Pipelines](docs/pipelines.md) | Full reference for all 11 pipelines |
| [Conditioning Types](docs/conditioning.md) | Image and video conditioning methods |
| [Multimodal Guidance](docs/multimodal-guidance.md) | CFG / STG / modality guidance parameters and tuning |
| [Optimization Tips](docs/optimization.md) | FP8 quantization, `torch.compile`, gradient estimation |
| [Multi-GPU Inference](docs/multigpu/README.md) | Run a single generation across GPUs for latency (SP, TDP, distributed VAE, distributed Gemma) |
## 🔗 Related Projects
packages/ltx-pipelines/docs/conditioning.md
+22
View File
@@ -0,0 +1,22 @@
# Conditioning Types
Pipelines use different conditioning methods from [`ltx-core`](../../ltx-core/) for controlling generation. See the [ltx-core conditioning documentation](../../ltx-core/README.md#conditioning--control) for details.
## Image Conditioning
All pipelines support image conditioning, but with different methods:
- **Replacing Latents** ([`image_conditionings_by_replacing_latent`](../src/ltx_pipelines/utils/helpers.py)):
- Replaces the latent at a specific frame with the encoded image
- Strong control over specific frames
- **Guiding Latents** ([`image_conditionings_by_adding_guiding_latent`](../src/ltx_pipelines/utils/helpers.py)):
- Adds the image as a guiding signal rather than replacing
- Better for smooth interpolation between keyframes
## Video Conditioning
- **Video Conditioning** (ICLoraPipeline only):
- Conditions on entire reference videos
- Useful for video-to-video transformations
- Uses `VideoConditionByKeyframeIndex` from [`ltx-core`](../../ltx-core/)
packages/ltx-pipelines/docs/installation.md
+64
View File
@@ -0,0 +1,64 @@
# Installation & Usage
## Installation
```bash
# From the repository root
uv sync --frozen
# Or install as a package
pip install -e packages/ltx-pipelines
```
## Requirements
- **LTX-2 Model Checkpoint** - Local `.safetensors` file
- **Gemma Text Encoder** - Local Gemma model directory
- **Spatial Upscaler** - Required by two-stage pipelines, for the upsampling stage
- **Distilled LoRA** - Required by two-stage non-distilled pipelines, used for the stage-2 refinement
## Running Pipelines
All pipelines can be run directly from the command line. Each pipeline module is executable:
```bash
# Run a pipeline (example: two-stage text-to-video)
python -m ltx_pipelines.ti2vid_two_stages \
--checkpoint-path path/to/checkpoint.safetensors \
--distilled-lora path/to/distilled_lora.safetensors 0.8 \
--spatial-upsampler-path path/to/upsampler.safetensors \
--gemma-root path/to/gemma \
--prompt "A beautiful sunset over the ocean" \
--output-path output.mp4
# View all available options for any pipeline
python -m ltx_pipelines.ti2vid_two_stages --help
```
### Available pipeline modules
- `ltx_pipelines.ti2vid_two_stages` - Two-stage text/image-to-video (recommended). ([docs](pipelines.md#1-ti2vidtwostagespipeline), [source](../src/ltx_pipelines/ti2vid_two_stages.py))
- `ltx_pipelines.ti2vid_two_stages_hq` - Two-stage text/image-to-video (different sampler, better quality). ([docs](pipelines.md#2-ti2vidtwostageshqpipeline), [source](../src/ltx_pipelines/ti2vid_two_stages_hq.py))
- `ltx_pipelines.ti2vid_one_stage` - Single-stage text/image-to-video. ([docs](pipelines.md#3-ti2vidonestagepipeline), [source](../src/ltx_pipelines/ti2vid_one_stage.py))
- `ltx_pipelines.t2a_one_stage` - Single-stage text-to-audio (audio-only output). ([docs](pipelines.md#11-t2aonestagepipeline), [source](../src/ltx_pipelines/t2a_one_stage.py))
- `ltx_pipelines.distilled` - Fast text/image-to-video pipeline using only the distilled model. ([docs](pipelines.md#4-distilledpipeline), [source](../src/ltx_pipelines/distilled.py))
- `ltx_pipelines.ic_lora` - Video-to-video with IC-LoRA. ([docs](pipelines.md#5-iclorapipeline), [source](../src/ltx_pipelines/ic_lora.py))
- `ltx_pipelines.keyframe_interpolation` - Keyframe interpolation. ([docs](pipelines.md#6-keyframeinterpolationpipeline), [source](../src/ltx_pipelines/keyframe_interpolation.py))
- `ltx_pipelines.a2vid_two_stage` - Audio-to-video generation conditioned on an input audio. ([docs](pipelines.md#7-a2vidpipelinetwostage), [source](../src/ltx_pipelines/a2vid_two_stage.py))
- `ltx_pipelines.retake` - Regenerate a time region of an existing video. ([docs](pipelines.md#8-retakepipeline), [source](../src/ltx_pipelines/retake.py))
- `ltx_pipelines.hdr_ic_lora` - Video-to-video with HDR output (linear float via LogC3 inverse decode). ([docs](pipelines.md#9-hdriclorapipeline), [source](../src/ltx_pipelines/hdr_ic_lora.py))
- `ltx_pipelines.lipdub` - Lip dubbing / re-voicing with IC-LoRA and audio reference conditioning. ([docs](pipelines.md#10-lipdubpipeline), [source](../src/ltx_pipelines/lipdub.py))
Use `--help` with any pipeline module to see all available options and parameters.
## Common CLI flags
These flags are shared across the pipeline CLIs (they come from a common base parser); run a module with `--help` for its full set.
- `--seed <int>` - random seed for reproducible generation (default 10).
- `--offload {none,cpu,disk}` - offload transformer weights to reduce peak GPU memory. `cpu` holds them in system RAM; `disk` streams them from disk when RAM is also limited (slower). Default `none`.
- `--quantization {fp8-cast,fp8-scaled-mm}` - run the transformer in FP8 to cut memory. `fp8-cast` downcasts a bf16 checkpoint on the fly (any FP8-capable GPU); `fp8-scaled-mm` expects an fp8 checkpoint and native FP8 support (best on Hopper+).
- `--max-batch-size <int>` - max batch per transformer forward pass (default 1). Higher values reduce layer-streaming transfers at the cost of peak memory.
- `--compile [key=value ...]` - enable `torch.compile`, optionally overriding the compilation config.
- `--lora <path> [strength]` - apply a LoRA (repeatable; default strength 1.0).
- `--enhance-prompt` - rewrite the prompt with the built-in enhancer before generation.
packages/ltx-pipelines/docs/multigpu/README.md
+78
View File
@@ -0,0 +1,78 @@
# Multi-GPU Inference
Run LTX-2 pipelines across several GPUs on a single machine.
> ## ⚠️ Important
>
> **Multi-GPU (MGPU) is a latency tool, not a memory tool.** It is designed to reduce
> the latency of a single generation on multi-GPU servers (H100, B200) by splitting
> each denoising step and the VAE decode across GPUs.
>
> **MGPU is not a way to fit a bigger model.** The mutable **working copy** of the
> transformer is a **full replica on every GPU** (each rank builds the whole model;
> LoRAs are fused into it in place). MGPU therefore cannot make a checkpoint that
> doesn't fit on one GPU suddenly fit — for that use FP8 quantization and weight
> offloading (see [Optimization Tips](../optimization.md)).
>
> Each rank *also* holds a second, immutable copy of the **clean (pre-LoRA)
> weights** — kept for LoRA hot-swap (reset + broadcast) — but that copy is
> **sharded** across GPUs (`ShardedSD`, ~1/world_size per rank), not replicated.
> Sequence parallelism additionally splits **activation** memory across ranks. See
> the [weight tracker](pipeline-setup.md#transformerweighttracker--working-copy--sharded-clean-weights)
> for the exact layout. The headline purpose is **latency**, not memory.
>
> **Single machine only.** One process per GPU, `MASTER_ADDR=localhost`, one rank
> per GPU. No multi-node.
## Requirements
- **Linux** -- NCCL and CUDA-IPC peer buffers are Linux-only (no macOS/Windows).
- **>=2 CUDA GPUs on a single node** with P2P access (NVLink/PCIe). No multi-node.
- **PyTorch with CUDA.**
- **`ltx-kernels` built** -- the SP all2all kernel is mandatory. Build with
`uv sync --group kernels` (needs a CUDA toolkit / nvcc and a C++ compiler, gcc or
clang). See the root README.
## Capabilities
| Technique | Purpose |
| --------- | ------- |
| [Sequence parallelism (SP)](sequence-parallel.md) | Split the token sequence across GPUs; faithful — numerically equivalent to single-GPU |
| [Tiled data parallelism (TDP)](tiled-data-parallel.md) | One spatial (height x width) tile per GPU; for resolutions outside the training distribution. **Upscale only** |
| [Distributed decoder](distributed-decoder.md) | Decode latent tiles in parallel, assemble on the driver |
| [Distributed Gemma](gemma.md) | Shard Gemma across GPUs via Accelerate `device_map`, or replicate + split prompts |
| [MGPU controller](controller.md) | Persistent worker fleet; dispatch a job, stream results |
| [Pipeline setup](pipeline-setup.md) | Swap single-GPU builders for MGPU builders; share one weights registry |
## Architecture overview
The [`MGPUController`](controller.md) spawns one worker process per GPU and runs a
user-defined **runner** (a subclass of `MGPURunner`) in SPMD lockstep. A runner's
`setup()` builds a standard pipeline, then **swaps** each block's builder for an MGPU
builder (SP / TDP / distributed decoder / distributed Gemma). All builders share one
`StateDictRegistry` so the checkpoint loads from disk once per process.
Two runners are provided, each with a CLI:
- [`ltx_pipelines.ti2vid_two_stages_mgpu`](../../src/ltx_pipelines/ti2vid_two_stages_mgpu.py) — SP stage 1 + TDP stage 2 + Accelerate Gemma + distributed VAE.
- [`ltx_pipelines.distilled_mgpu`](../../src/ltx_pipelines/distilled_mgpu.py) — SP (shared stage) + Accelerate Gemma + distributed VAE.
```bash
# Two-stage on all visible GPUs
python -m ltx_pipelines.ti2vid_two_stages_mgpu \
--checkpoint-path path/to/checkpoint.safetensors \
--distilled-lora path/to/distilled_lora.safetensors 1.0 \
--spatial-upsampler-path path/to/upsampler.safetensors \
--gemma-root path/to/gemma \
--prompt "A beautiful sunset over the ocean" \
--output-path output.mp4
```
## Pages
- **[Controller](controller.md)** — `MGPUController` / `MGPURunner` / `Stream`, lifecycle, one-job-at-a-time contract, threading, error handling.
- **[Pipeline setup](pipeline-setup.md)** — swapping builders, the shared weights registry, the LoRA-hot-swap weight tracker.
- **[Sequence parallelism](sequence-parallel.md)** — faithful token-dim split, the all2all kernels, `AttentionManager`, `SequenceParallelBuilder`.
- **[Tiled data parallelism](tiled-data-parallel.md)** — out-of-distribution resolutions, position normalization, shared negative (reference) positions, `TiledDataParallelBuilder`.
- **[Distributed decoder](distributed-decoder.md)** — inter-GPU vs intra-GPU tiling, `DistributedDecoderBuilder`.
- **[Gemma](gemma.md)** — `AccelerateGemmaBuilder` (Accelerate `device_map` sharding) and `BatchParallelGemmaBuilder` (replicated; not for the distilled pipeline).
packages/ltx-pipelines/docs/multigpu/controller.md
+115
View File
@@ -0,0 +1,115 @@
# MGPU Controller
**Source**: [`multigpu/controller.py`](../../src/ltx_pipelines/multigpu/controller.py), [`multigpu/runner.py`](../../src/ltx_pipelines/multigpu/runner.py)
The controller is a persistent, one-job-at-a-time GPU fleet. It spawns one worker
process per GPU, runs a user-defined **runner** in SPMD lockstep, and streams
results back.
## Public classes
### `MGPUController`
```python
MGPUController(
runner_cls: type[MGPURunner],
*,
num_gpus: int | None = None, # GPUs 0..num_gpus-1 (default: all visible)
devices: Sequence[int] | None = None, # place on specific physical GPUs, e.g. [2, 3]
logs_specs: LogsSpecs | None = None,
)
```
`num_gpus` and `devices` are mutually exclusive. `devices=[2, 3]` puts rank r on
`cuda:devices[r]`, so two controllers can share one machine on disjoint GPU sets.
Lifecycle:
| Method | What it does |
| ------ | ------------ |
| `start(*, timeout=30min, **setup_kwargs)` | Spawn the fleet, run `setup(**setup_kwargs)` on every rank, block until all report ready. `timeout` bounds NCCL init + CUDA init + `setup()`; it must exceed the slowest model load. |
| `stream(*, timeout=None, **kwargs) -> Stream` | Dispatch one job and return **immediately**. Iterate the returned `Stream` to collect. |
| `shutdown(*, graceful_timeout=60.0)` | Tear down the fleet; also force-terminates it — safe to call from another thread to recover a job that cannot be drained. |
| `is_alive` (property) | True while the fleet is up and unpoisoned. |
### `MGPURunner`
`MGPURunner` is the abstract base class implemented per pipeline. The controller
ships the subclass to every worker by value (a runner defined in `__main__` or a test
module is supported), builds one instance per worker, injects the NCCL groups, calls
`setup()` once, then invokes the instance per job.
```python
class MyRunner(MGPURunner):
@torch.inference_mode()
def setup(self, *, checkpoint_path: str, ...) -> None:
# build the pipeline + swap in MGPU builders (see pipeline-setup.md)
...
@torch.inference_mode()
def __call__(self, *, prompt: str, ...) -> Iterator[...]:
video, audio = self._pipeline(...)
yield output_path # __call__ MUST be a generator (use `yield`, even once)
```
- `setup()` and `__call__()` run on **every** rank. `self.groups` gives the
per-component `NCCLGroups` (`gemma_group`, `transformer_group`, `vae_group`).
- The framework does **not** apply inference mode — decorate `setup`/`__call__` explicitly.
## Usage
```python
from ltx_pipelines.multigpu import MGPUController
controller = MGPUController(MyRunner, num_gpus=8)
controller.start(checkpoint_path="...", gemma_root="...") # setup kwargs
stream = controller.stream(prompt="a cat", seed=42)
try:
for item in stream: # one element per yield, as it arrives (NOT gathered across ranks)
show(item)
finally:
stream.drain() # free the controller even on early exit
controller.shutdown()
```
### Passing tensors
Tensors are transparent:
- **Inputs.** Pass them as **top-level** kwargs (`stream(latent=t, steps=30)`) and
the relay (rank 0) broadcasts them to every rank over NCCL — `__call__` receives
them already on the local GPU. An input tensor nested inside a list/dict kwarg is
**not** broadcast; it falls back to the (slower) pickle path.
- **Outputs.** Yield tensors back (including nested inside a dict) and they return
via the result queue by shared memory / CUDA IPC — no pickling, regardless of
nesting.
Everything else must be picklable and small.
## Contract and limitations
- **Single machine only.** `MASTER_ADDR=localhost`, `RANK == LOCAL_RANK`, one rank per GPU.
- **One job at a time.** No job queue, no pipelining. Consume the `Stream` to the
end before the next `stream()`. Abandoning it is **not** cleaned up: the next
`stream()` raises `ControllerBusyError` until `stream.drain()` or `shutdown()` is
called. The recommended pattern is `try: ... finally: stream.drain()`.
- **SPMD lockstep.** Yields are forwarded individually (in result-queue order), not
gathered. Only per-rank terminals are collected to end the stream.
- **Thread ownership (baton-lock).** Any thread may call `stream()`. Each job
belongs to its **dispatching** thread — only that
thread may iterate or `drain()` its `Stream` (enforced in `Stream.__next__`). A
single lock guards only the in-flight check-and-set: among concurrent `stream()`
callers one proceeds and the rest raise `ControllerBusyError`.
## Error handling
| Situation | Outcome |
| --------- | ------- |
| Runner raises an **unexpected** exception | **Fatal** — a desynced NCCL collective cannot be unwound. The controller is poisoned and a new one must be constructed. |
| Runner raises `RunnerError` (or `ValueError`, auto-converted) **identically on every rank** | Recoverable. Iterating the `Stream` re-raises `SymmetricRunnerError`; the fleet survives — fix the input and retry. Raise it outside any collective (e.g. validating broadcast kwargs before the first one). |
| Some ranks raise `RunnerError`, others finish clean | `AsymmetricRunnerError` — surfaced prominently (a latent hang risk), but does not terminate the fleet. |
| Worker death / exceeded per-job `timeout` | Surfaced when the `Stream` is next iterated; the controller is poisoned. |
The public API (`from ltx_pipelines.multigpu import ...`) exports `MGPUController`,
`MGPURunner`, `Stream`, `RunnerError`, `SymmetricRunnerError`,
`AsymmetricRunnerError`, `ControllerBusyError`, and `NCCLGroups`.
packages/ltx-pipelines/docs/multigpu/distributed-decoder.md
+72
View File
@@ -0,0 +1,72 @@
# Distributed VAE Decoder
**Source**: [`ltx_core/multigpu/vae/distributed_decoder.py`](../../../ltx-core/src/ltx_core/multigpu/vae/distributed_decoder.py), [`multigpu/vae_builders.py`](../../src/ltx_pipelines/multigpu/vae_builders.py)
## What it is
VAE decode is expensive and embarrassingly parallel over space/time tiles. The
distributed decoder splits the latent into tiles, assigns them **round-robin** to
ranks (the tile count may exceed the GPU count), and every rank decodes its tiles in
parallel.
Workers ship their decoded tiles to the **driver rank** over an `mp.Queue` (CUDA
IPC — zero-copy handle sharing); the driver blends overlaps and yields the assembled
frames as temporal batches spread across the GPUs.
## Inter-GPU tiling vs intra-GPU tiling
Two independent tilings, commonly conflated:
| | Controls | Config | Set by |
| --- | --- | --- | --- |
| **Inter-GPU** (MGPU) | Which rank decodes which tile — **parallelism** | `vae_tiling: TileCountConfig` (at build time) | `DistributedDecoderBuilder` |
| **Intra-GPU** (SGPU) | Chunking *within* a rank's tile to bound **VRAM** | `tiling_config: TilingConfig` (per call) | the pipeline's usual tiling kwarg |
They compose — a rank can further chunk its assigned tile for VRAM — with one
guard: **multi-GPU temporal tiling and single-GPU temporal tiling cannot both be
on.** If `vae_tiling.frames.num_tiles > 1` and `tiling_config.temporal_config` is
set, `decode_video` raises, because two causal temporal splits would conflict.
## API
### `DistributedDecoderBuilder`
```python
from ltx_pipelines.multigpu.vae_builders import DistributedDecoderBuilder
from ltx_core.tiling import TileCountConfig, DimensionTilingConfig
DistributedDecoderBuilder(
inner: BuilderProtocol, # the block's single-GPU decoder builder
queue: Queue, # spawn-context mp.Queue, shared across ranks (CUDA IPC)
vae_group: dist.ProcessGroup, # self.groups.vae_group
vae_tiling: TileCountConfig, # inter-GPU split
driver_rank: int, # rank that collects + assembles (usually 0)
registry: Registry,
)
```
`build()` returns a `DistributedVideoDecoder`; its `decode_video(latent,
tiling_config=None, ...)` returns an iterator of temporal batches on the driver, and
an empty iterator on workers (they only `put` their tiles).
## Usage
```python
# The queue is created once, in the CLI __main__, and passed to controller.start(...)
# as a setup kwarg so every worker shares it:
vae_queue = torch.multiprocessing.get_context("spawn").SimpleQueue()
# inside runner.setup():
vae_tiling = TileCountConfig(height=DimensionTilingConfig(num_tiles=8, overlap=4))
pipeline.video_decoder._decoder_builder = DistributedDecoderBuilder(
inner=pipeline.video_decoder._decoder_builder,
queue=vae_queue,
vae_group=self.groups.vae_group,
vae_tiling=vae_tiling,
driver_rank=0,
registry=registry,
)
```
The shipped runners tile the VAE across **height** (8 tiles, overlap 4). Only the
driver ends up with the assembled video, which is why the runner's `__call__`
encodes the file on `driver_rank` and yields `None` on the others.
packages/ltx-pipelines/docs/multigpu/gemma.md
+80
View File
@@ -0,0 +1,80 @@
# Gemma Text Encoder (Multi-GPU)
**Source**: [`multigpu/gemma_builders.py`](../../src/ltx_pipelines/multigpu/gemma_builders.py), [`multigpu/bp_gemma_builder.py`](../../src/ltx_pipelines/multigpu/bp_gemma_builder.py)
Two ways to run the Gemma text encoder across the fleet. Both swap in for
`pipeline.prompt_encoder._text_encoder_builder` and broadcast the resulting
embeddings to every rank (so the transformer ranks all have them).
## `AccelerateGemmaBuilder` (Accelerate `device_map` — the default)
Loads Gemma **once, on the source rank**, with Accelerate `device_map="auto"`, which
**shards Gemma's layers across the available GPUs**. Non-source ranks receive a
lightweight `AccelerateGemmaWrapper` stub that receives the encoded embeddings over
NCCL. The source rank fuses all prompts into one Gemma call, then broadcasts each
output.
The first `build()` loads via HuggingFace `from_pretrained` and caches the full
state dict (including non-persistent buffers) in the registry; later builds recreate
the model from cache and reinstall the dispatch hooks — no disk I/O.
```python
from ltx_pipelines.multigpu.gemma_builders import AccelerateGemmaBuilder
AccelerateGemmaBuilder(
gemma_root_path: str,
gemma_group: dist.ProcessGroup | None, # self.groups.gemma_group
broadcast_group: dist.ProcessGroup | None, # self.groups.transformer_group
registry: Registry,
*,
src_rank: int, # rank that loads + encodes (usually 0)
dtype: torch.dtype = torch.bfloat16,
)
```
Usage (in `runner.setup()`):
```python
pipeline.prompt_encoder._text_encoder_builder = AccelerateGemmaBuilder(
gemma_root_path=gemma_root,
gemma_group=self.groups.gemma_group,
broadcast_group=self.groups.transformer_group,
registry=registry,
src_rank=0,
dtype=pipeline.dtype,
)
```
The shipped runners (`ti2vid_two_stages_mgpu`, `ti2vid_two_stages_hq_mgpu`, `distilled_mgpu`) use this builder.
## `BatchParallelGemmaBuilder` (replicated — data-parallel over prompts)
Every rank materialises a **full** `GemmaTextEncoder` on its own GPU via the standard
`SingleGPUModelBuilder` path (no Accelerate, no `device_map`, no per-layer dispatch
hooks). The wrapper (`BatchParallelGemmaWrapper`) then **partitions the prompt list
across ranks** in `encode` and broadcasts each prompt's output, so the forwards run
concurrently on different GPUs. Non-deterministic prompt enhancement
(`enhance_t2v` / `enhance_i2v`) is routed through a single `src_rank`.
```python
from ltx_pipelines.multigpu.bp_gemma_builder import BatchParallelGemmaBuilder
BatchParallelGemmaBuilder(
gemma_root_path: str,
broadcast_group: dist.ProcessGroup | None,
registry: Registry,
*,
src_rank: int,
dtype: torch.dtype = torch.bfloat16,
)
```
### Not for the distilled pipeline
Batch-parallel is beneficial only when there is **more than one prompt to encode** —
the typical CFG case, positive + negative (B=2 on 2 ranks = one prompt per rank, both
forwards concurrent). The **distilled** pipeline runs **without CFG**: its `__call__`
accepts a single `prompt` and no `negative_prompt`, so there is only one prompt to
encode and no work to partition; batch-parallel provides no speedup in that case. Use
`AccelerateGemmaBuilder` for the distilled pipeline (as the shipped `distilled` runner
does).
packages/ltx-pipelines/docs/multigpu/pipeline-setup.md
+111
View File
@@ -0,0 +1,111 @@
# Setting Up an MGPU Pipeline
**Source**: [`ti2vid_two_stages_mgpu.py`](../../src/ltx_pipelines/ti2vid_two_stages_mgpu.py), [`multigpu/weight_tracker.py`](../../src/ltx_pipelines/multigpu/weight_tracker.py)
An MGPU pipeline **is** a single-GPU pipeline with its per-block builders swapped
for MGPU builders. Build the standard pipeline, then replace each block's
`_transformer_builder` / `_text_encoder_builder` / `_decoder_builder`.
## The pattern
This is performed inside a runner's `setup()` (which runs on every rank — see
[Controller](controller.md)).
```python
from ltx_core.loader.registry import StateDictRegistry
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
from ltx_pipelines.multigpu.sp_builder import SequenceParallelBuilder
from ltx_pipelines.multigpu.tdp_builder import TiledDataParallelBuilder
from ltx_pipelines.multigpu.gemma_builders import AccelerateGemmaBuilder
from ltx_pipelines.multigpu.vae_builders import DistributedDecoderBuilder
from ltx_pipelines.multigpu.weight_tracker import TransformerWeightTracker
from ltx_core.multigpu.transformer.attention import AttentionManager
# 1. ONE shared registry for every builder in this process.
registry = StateDictRegistry()
# 2. Build the normal pipeline, handing it the registry.
pipeline = TI2VidTwoStagesPipeline(
checkpoint_path=..., distilled_lora=..., spatial_upsampler_path=...,
gemma_root=..., loras=[], registry=registry, quantization=...,
)
# 3. One weight tracker per transformer process group (shared by the stages).
tracker = TransformerWeightTracker(group=self.groups.transformer_group)
# 4. Swap each block's builder.
pipeline.stage_1._transformer_builder = SequenceParallelBuilder(
inner=pipeline.stage_1._transformer_builder, attn_mgr=attn_mgr,
registry=registry, tracker=tracker,
)
pipeline.stage_2._transformer_builder = TiledDataParallelBuilder(
inner=pipeline.stage_2._transformer_builder, group=self.groups.transformer_group,
tiling=tdp_tiling, registry=registry, tracker=tracker,
)
pipeline.prompt_encoder._text_encoder_builder = AccelerateGemmaBuilder(...)
pipeline.video_decoder._decoder_builder = DistributedDecoderBuilder(...)
```
Each MGPU builder **wraps** the block's existing single-GPU builder (`inner=...`),
so it inherits the checkpoint path, quantization, compilation, and LoRA config —
only the parallelism is added. See the per-technique pages for each builder's
constructor.
> **`with_builder` vs direct assignment.** `DiffusionStage.with_builder(builder)`
> returns a *new* stage with the builder swapped (functional, never mutates). The
> runners assign `stage._transformer_builder = ...` directly because they mutate the
> pipeline once, in place, during `setup()`. Both reach the same builder slot.
## The shared weights registry
`StateDictRegistry` is an in-process cache of loaded state dicts, keyed by
`(resolved paths, sd_ops name)`. Passing **one** registry to every builder means:
- The transformer checkpoint is read from disk **once per process**, even though
stage 1 (SP) and stage 2 (TDP) are separate builders on the same file.
- Gemma and the VAE cache their weights the same way (rebuild the module tree from
the cached tensors, skip disk I/O).
The registry is **per process** — it is not shared across ranks. Each worker loads
its own copy, so the full checkpoint is resident on every GPU (see the
[memory disclaimer](README.md)).
## `TransformerWeightTracker` — working copy + sharded clean weights
```python
TransformerWeightTracker(group: dist.ProcessGroup, bucket_mb=256, no_lora_swap=False)
```
The tracker is shared by the transformer stage builders that operate on the same
checkpoint. It does **not** own weights — it references the tensors in the registry
and receives a builder at `build()` time. Two copies of the weights exist per rank,
and they are **not** the same shape of memory:
- **Working copy** — the model the builder returns, backed by the registry's
tensors. This is a **full replica on every GPU**. LoRAs are fused into it
**in place**; `broadcast_sd` (a zero-copy `ShardedSD` view over these tensors)
broadcasts each owner rank's freshly fused shards — bucketed, `bucket_mb` at a
time — so all ranks converge on identical working weights.
- **Clean weights** (`stored_sd`) — an immutable, cloned backup of the original
(pre-LoRA) weights, held so the working copy can be reset before a different LoRA
set is applied. This copy is **sharded** across ranks (deterministic
`md5(key) % world_size` ownership): each rank stores only its ~1/world_size slice,
not a full clone.
So per-GPU transformer memory is one full working model **plus** a ~1/world_size
clean-weights shard — the clean backup is distributed, the working copy is not.
This allows a two-stage pipeline to apply the distilled LoRA to stage 2 and reset it
for stage 1 without reloading the checkpoint. Pass `no_lora_swap=True` when the
LoRA set is fixed (none, or one set for the whole run): the clean-weights clone is
skipped (`stored_sd` becomes a zero-copy view) and any swap/reset raises — saves the
~1/N shard, and guards against accidental swaps.
## Full example
The shipped runners are the reference: read
[`ti2vid_two_stages_mgpu.py`](../../src/ltx_pipelines/ti2vid_two_stages_mgpu.py)
(`setup()` lines ~54–132) and
[`distilled_mgpu.py`](../../src/ltx_pipelines/distilled_mgpu.py). Each ends with a
`__main__` block wiring the runner into an `MGPUController` behind the standard
two-stage CLI parser.
packages/ltx-pipelines/docs/multigpu/sequence-parallel.md
+100
View File
@@ -0,0 +1,100 @@
# Sequence Parallelism (SP)
**Source**: [`ltx_core/multigpu/transformer/sequence_parallel.py`](../../../ltx-core/src/ltx_core/multigpu/transformer/sequence_parallel.py), [`multigpu/sp_builder.py`](../../src/ltx_pipelines/multigpu/sp_builder.py)
## What it is
SP splits the **token (sequence) dimension** of the video across GPUs. Each rank
holds a slice of the tokens, runs the transformer forward on its slice, and the
outputs are gathered back to all ranks. Self-attention still needs every token to
see every other token, so the Q/K/V heads are exchanged across ranks with a custom
**all2all** kernel: each rank ends up with all tokens for a subset of heads, does
local attention, then the results are shuffled back.
**SP is faithful — numerically equivalent to single-GPU inference.** Attention stays
global (all2all preserves the full token interaction); only the floating-point
reduction order changes. The all2all kernels move bytes only — the round-trip
`gather(send(x)) == x` is byte-exact. SP is the appropriate choice whenever the
single-GPU result is required at lower latency.
This is the default for **stage 1** (`ti2vid_two_stages_mgpu`) and the **shared stage**
(`distilled_mgpu`, where one SP wrapping covers both the half-res and full-res calls).
## How the forward pass works
Per denoising step, [`SequenceParallelModelWrapper`](../../../ltx-core/src/ltx_core/multigpu/transformer/sequence_parallel.py):
1. Pads the video seq dim up to a multiple of `world_size` (padded keys are masked
out; padded rows sliced off after the gather) so every rank gets an equal shard.
2. Tiles latent/timesteps/positions to this rank's slice.
3. Runs the model — video self-attention (`attn1`) and video→audio cross-attention
are patched to route Q/K/V through the all2all kernel.
4. `all_gather`s the output tokens back to full length on every rank and unpads.
## The all2all kernels (`ltx-kernels`)
The custom op is `ltx_kernels.All2All` (from the `ltx-kernels` package); the CUDA
kernels use CUDA-IPC peer buffers to exchange tokens directly between ranks' GPUs.
**`ltx-kernels` must be installed** — the SP builder imports it.
## API
### `AttentionManager`
```python
from ltx_core.multigpu.transformer.attention import AttentionManager
attn_mgr = AttentionManager(
max_tokens: int, # upper bound on total video tokens (raises above it)
num_heads: int, # transformer.num_attention_heads
head_dim: int, # transformer.attention_head_dim
tensor_dtype: torch.dtype,
group: dist.ProcessGroup, # self.groups.transformer_group
copy_out_: bool = False,
)
```
Owns the all2all buffers (sized `ceil(max_tokens / world_size)` tokens per rank) and, per step,
`set_seqlen_all2all(...)` updates the per-rank token counts. `num_heads` must be
divisible by `world_size`.
### `SequenceParallelBuilder`
```python
from ltx_pipelines.multigpu.sp_builder import SequenceParallelBuilder
SequenceParallelBuilder(
inner: ModelBuilderProtocol, # the stage's single-GPU transformer builder
attn_mgr: AttentionManager,
registry: Registry,
tracker: TransformerWeightTracker,
)
```
Wraps a `SingleGPUModelBuilder` (raises otherwise), injects the all2all attention
module-ops, and `build()` returns a `SequenceParallelModelWrapper`.
## Usage
```python
# inside runner.setup(), per stage:
model_cfg = pipeline.stage_1._transformer_builder.model_config().get("transformer", {})
attn_mgr = AttentionManager(
max_tokens=32768,
num_heads=model_cfg["num_attention_heads"],
head_dim=model_cfg["attention_head_dim"],
tensor_dtype=pipeline.dtype,
group=self.groups.transformer_group,
)
pipeline.stage_1._transformer_builder = SequenceParallelBuilder(
inner=pipeline.stage_1._transformer_builder,
attn_mgr=attn_mgr,
registry=registry,
tracker=tracker,
)
```
`max_tokens` must cover the largest step. Reference: stage 1 at 512x768x121 is
~6k video tokens; the distilled shared stage's full-res call (1024x1536x121) is
~24k — both ship with `sp_max_tokens=32768`. Exceeding it raises with a clear
"use a smaller resolution or fewer frames" message.
packages/ltx-pipelines/docs/multigpu/tiled-data-parallel.md
+130
View File
@@ -0,0 +1,130 @@
# Tiled Data Parallelism (TDP)
**Source**: [`ltx_core/multigpu/transformer/tiled_data_parallel.py`](../../../ltx-core/src/ltx_core/multigpu/transformer/tiled_data_parallel.py), [`multigpu/tdp_builder.py`](../../src/ltx_pipelines/multigpu/tdp_builder.py)
## What it is
TDP splits the patchified `(frames, height, width)` latent into **tiles** and gives
each tile to a GPU. Every rank runs the full transformer on its own tile(s),
overlapping regions are blended with trapezoidal masks, and a single `all_reduce`
sums the blended tiles into the final result (masks sum to 1 globally). Tiles are
assigned **round-robin**, so the tile count may exceed the GPU count (16 tiles on
4 GPUs = 4 tiles/rank). Audio is processed untiled on every tile forward and averaged
across tiles.
Unlike [sequence parallelism](sequence-parallel.md), TDP is **not** bit-faithful to
single-GPU: each tile is denoised with only local context and blended, so it is an
approximation.
> ## ⚠️ Do not use the TDP stage's audio output
>
> Audio is **not** tiled. It is denoised on **every** tile's forward pass — each with
> a different, partial video context — and those results are **averaged across all
> tiles**. That average is not a meaningful audio latent. Take the final audio from
> the first (SP) stage and keep it frozen through the TDP upscale; treat the TDP
> stage's audio only as the video-conditioning context it needs internally, never as
> output.
## When to use it
TDP is an **upscaler**. It produces video at resolutions the model never saw during
training by running each tile at a resolution the model handles well and blending the
results. This is why the shipped two-stage runner uses TDP for **stage 2** (the
high-resolution upscale).
TDP can also be **faster** than running the whole frame on one GPU: self-attention is
quadratic in the token count, so splitting `N` tokens into `T` tiles drops per-tile
attention cost from `O(N^2)` to `O((N/T)^2)`.
> **Do not run TDP as the first stage.** Starting from pure noise (a high first
> sigma), each tile denoises independently and produces **unrelated content** — the
> tiles never converge on a single coherent video. Generate the first stage with
> [SP](sequence-parallel.md) (faithful, full-frame), then **upscale** that result
> with TDP.
>
> Even as the upscale stage, tiles can **drift** apart, and the drift grows with the
> **first sigma** of the TDP stage (more noise re-injected means more freedom per
> tile). For consistency, either condition on the stage-1 result with
> **negative-index image conditioning** (for i2v), or use a **smaller first sigma**.
## Position normalization
A tile's tokens must carry positions in the range the model was trained on — not the
global positions of a tile in the corner of a large frame, which the model was never
trained to handle. With `normalize_positions=True` (default), each tile's positions
are shifted so the tile's **generated** tokens start at zero in every dimension:
```python
offset = gen_pos[..., 0].amin(dim=2, keepdim=True)... # min start per (batch, dim)
positions = positions - offset # shift generated + conditioning tokens
```
Interval widths are preserved (only the origin moves), so RoPE sees a valid,
in-distribution position grid per tile.
## Shared negative (reference) positions
Conditioning tokens are appended after the generated tokens. A tile keeps a
conditioning token when its `[start, end)` interval overlaps the tile in all three
dimensions — **or** when it has a **negative time coordinate**. Negative-time tokens
are **reference tokens** (e.g. reference frames / audio references): they are kept by
**every** tile so all tiles share the same reference context. To avoid
double-counting a token kept by several tiles, its blend weight is `1 / (number of
tiles that kept it)`.
## API
### Tiling config (`ltx_core.tiling`)
```python
from ltx_core.tiling import TileCountConfig, DimensionTilingConfig
TileCountConfig(
frames: DimensionTilingConfig = DimensionTilingConfig(num_tiles=1, overlap=0),
height: DimensionTilingConfig = DimensionTilingConfig(num_tiles=1, overlap=0),
width: DimensionTilingConfig = DimensionTilingConfig(num_tiles=1, overlap=0),
)
DimensionTilingConfig(num_tiles: int, overlap: int = 0) # counts, not sizes; overlap in latent grid units
```
`TileCountConfig` specifies tile **counts** per dimension (contrast the single-GPU
VAE `TilingConfig`, which specifies tile **sizes**).
### `TiledDataParallelBuilder`
```python
from ltx_pipelines.multigpu.tdp_builder import TiledDataParallelBuilder
TiledDataParallelBuilder(
inner: ModelBuilderProtocol, # the stage's single-GPU transformer builder
group: dist.ProcessGroup, # self.groups.transformer_group
tiling: TileCountConfig,
registry: Registry,
tracker: TransformerWeightTracker,
normalize_positions: bool = True,
)
```
Wraps a `SingleGPUModelBuilder`. Its `build()` requires a `video_tools` kwarg (the
`VideoLatentTools` for the target shape) so the wrapper can compute tiles — the
pipeline passes this through automatically.
## Usage
```python
# inside runner.setup(), stage 2 -- balanced 2D spatial (height x width) grid over the group:
from ltx_core.tiling import TileCountConfig, DimensionTilingConfig, balanced_tile_split
h_tiles, w_tiles = balanced_tile_split(dist.get_world_size(self.groups.transformer_group))
tdp_tiling = TileCountConfig(
height=DimensionTilingConfig(num_tiles=h_tiles, overlap=5),
width=DimensionTilingConfig(num_tiles=w_tiles, overlap=5),
)
pipeline.stage_2._transformer_builder = TiledDataParallelBuilder(
inner=pipeline.stage_2._transformer_builder,
group=self.groups.transformer_group,
tiling=tdp_tiling,
registry=registry,
tracker=tracker,
)
```
packages/ltx-pipelines/docs/multimodal-guidance.md
+52
View File
@@ -0,0 +1,52 @@
# 🎛️ Multimodal Guidance
LTX-2 pipelines use **multimodal guidance** to steer the diffusion process for both video and audio modalities. Each modality (video, audio) has its own guider with independent parameters, allowing fine-grained control over generation quality and adherence to prompts.
## Guidance Parameters
The `MultiModalGuiderParams` dataclass controls guidance behavior:
| Parameter | Description |
| --------- | ----------- |
| `cfg_scale` | **Classifier-Free Guidance** scale. Higher values make the output adhere more strongly to the text prompt. Typical values: 2.0–5.0. Set to **1.0** to disable. |
| `stg_scale` | **Spatio-Temporal Guidance** scale. Controls perturbation-based guidance for improved temporal coherence. Typical values: 0.5–1.5. Set to **0.0** to disable. |
| `stg_blocks` | Which transformer blocks to perturb for STG (e.g., `[29]` for the last block). Set to **`[]`** to disable STG. |
| `rescale_scale` | Rescales the guided prediction to match the variance of the conditional prediction. Helps prevent over-saturation. Typical values: 0.5–0.7. Set to **0.0** to disable. |
| `modality_scale` | **Modality CFG** scale. Steers the model away from unsynced video and audio results, improving audio-visual coherence. Set to **1.0** to disable. |
| `skip_step` | Skip guidance every N steps. Can speed up inference with minimal quality loss. Set to **0** to disable (never skip). |
## How It Works
The multimodal guider combines three guidance signals during each denoising step:
1. **CFG (Text Guidance)**: Steers generation toward the text prompt by computing `(cond - uncond_text)`.
2. **STG (Perturbation Guidance)**: Improves structural coherence by perturbing specific transformer blocks and steering away from the perturbed prediction.
3. **Modality CFG**: For joint audio-video generation, steers the model away from unsynced video and audio results.
## Example Configuration
```python
from ltx_core.components.guiders import MultiModalGuiderParams
# Video guider: moderate CFG, STG enabled, modality isolation
video_guider_params = MultiModalGuiderParams(
cfg_scale=3.0,
stg_scale=1.0,
rescale_scale=0.7,
modality_scale=3.0,
stg_blocks=[29],
)
# Audio guider: higher CFG for stronger prompt adherence
audio_guider_params = MultiModalGuiderParams(
cfg_scale=7.0,
stg_scale=1.0,
rescale_scale=0.7,
modality_scale=3.0,
stg_blocks=[29],
)
```
> **Tip:** Start with the default values from [`constants.py`](../src/ltx_pipelines/utils/constants.py) and adjust based on your use case. Higher `cfg_scale` = stronger prompt adherence but potentially less natural motion; higher `stg_scale` = better temporal coherence but slower inference (requires extra forward passes).
>
> **Tip:** When generating video with audio, set `modality_scale` > 1.0 (e.g., 3.0) to improve audio-visual sync. If generating video-only, set it to 1.0 to disable.
packages/ltx-pipelines/docs/optimization.md
+149
View File
@@ -0,0 +1,149 @@
# ⚡ Optimization Tips
## Memory Optimization
### FP8 Quantization (Lower Memory Footprint)
For smaller GPU memory footprint, use the `--quantization` flag and set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`.
Two quantization policies are available:
| Policy | CLI Flag | Description |
| ------ | -------- | ----------- |
| **FP8 Cast** | `--quantization fp8-cast` | Downcasts transformer linear weights to FP8 during loading; upcasts on the fly during inference. No extra dependencies. |
| **FP8 Scaled MM** | `--quantization fp8-scaled-mm` | Uses FP8 scaled matrix multiplication via PyTorch's `torch._scaled_mm`. Best performance on Hopper+ GPUs with native FP8 support. |
**CLI:**
```bash
# FP8 Cast (works on any GPU with FP8 support)
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python -m ltx_pipelines.ti2vid_two_stages \
--quantization fp8-cast --checkpoint-path=...
# FP8 Scaled MM (no extra deps, best on Hopper+ GPUs)
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python -m ltx_pipelines.ti2vid_two_stages \
--quantization fp8-scaled-mm --checkpoint-path=...
```
**Programmatically:**
When authoring custom scripts, pass a `QuantizationPolicy` to pipeline classes:
```python
from ltx_core.quantization.fp8_cast import build_policy as build_fp8_cast_policy
# Alternative:
# from ltx_core.quantization.fp8_scaled_mm import build_policy as build_fp8_scaled_mm_policy
pipeline = TI2VidTwoStagesPipeline(
checkpoint_path=ltx_model_path,
distilled_lora=distilled_lora,
spatial_upsampler_path=upsampler_path,
gemma_root=gemma_root_path,
loras=[],
quantization=build_fp8_cast_policy(ltx_model_path),
)
pipeline(...)
```
You still need to use `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` when launching:
```bash
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python my_denoising_pipeline.py
```
### Memory Cleanup Between Stages
By default, pipelines clean GPU memory (especially transformer weights) between stages. If you have enough memory, you can skip this cleanup to reduce running time:
```python
# In pipeline implementations, memory cleanup happens automatically
# between stages. For custom pipelines, you can skip:
# utils.cleanup_memory() # Comment out if you have enough VRAM
```
## Compilation (`torch.compile`)
Compiling the transformer blocks with `torch.compile` speeds up inference. It is **opt-in and off by default**. The blocks are compiled shape-polymorphically (the sequence dimension is marked dynamic), so one compiled artifact serves any token count without recompiling.
**CLI** — the `--compile` flag maps directly to `CompilationConfig`:
| Form | Result |
| ---- | ------ |
| *(flag absent)* | eager, no compilation |
| `--compile` | compile with defaults |
| `--compile KEY=VALUE ...` | compile, overriding individual fields |
```bash
# Defaults
python -m ltx_pipelines.ti2vid_two_stages --compile --checkpoint-path=...
# reduce-overhead captures CUDA graphs -- the main latency lever for the denoising loop.
# Off by default because graph capture reserves static memory pools (extra VRAM), so it
# trades memory for speed; enable it when you have headroom.
python -m ltx_pipelines.ti2vid_two_stages --compile mode=reduce-overhead --checkpoint-path=...
# Several overrides at once
python -m ltx_pipelines.ti2vid_two_stages \
--compile mode=max-autotune fullgraph=true dynamic=true --checkpoint-path=...
```
| Field | Values | Default | Notes |
| ----- | ------ | ------- | ----- |
| `mode` | `none`, `reduce-overhead`, `max-autotune`, … | `none` | `reduce-overhead`/`max-autotune` enable CUDA graphs |
| `backend` | `inductor`, `eager`, … | `inductor` | |
| `fullgraph` | `true`/`false` | `false` | |
| `dynamic` | `auto`/`true`/`false` | `auto` | the seq dim is marked dynamic regardless |
| `inductor_config` | JSON object or path to a `.json` | `{}` | `torch._inductor.config` overrides |
| `dynamo_config` | JSON object or path to a `.json` | `{"inline_inbuilt_nn_modules": true, "cache_size_limit": 256}` | `torch._dynamo.config` overrides |
**Controlling inductor / dynamo configs.** `inductor_config` and `dynamo_config` take either an inline JSON object or a path to a `.json` file, applied via `torch._inductor.config.patch(...)` / `torch._dynamo.config.patch(...)` around the compiled forward. They **replace the defaults wholesale — they do not merge**, so when overriding `dynamo_config` re-include any defaults you want to keep:
```bash
python -m ltx_pipelines.ti2vid_two_stages \
--compile 'inductor_config={"max_autotune": true}' \
'dynamo_config={"inline_inbuilt_nn_modules": true, "cache_size_limit": 256, "recompile_limit": 32}' \
--checkpoint-path=...
```
**Programmatically**, pass a `CompilationConfig` to the pipeline:
```python
from ltx_core.model.transformer.compiling import CompilationConfig
pipeline = TI2VidTwoStagesPipeline(
...,
compilation_config=CompilationConfig(mode="reduce-overhead"),
)
```
**Faster cache loads: `unsafe_skip_cache_dynamic_shape_guards` (unsafe, opt-in).** Inductor's FX-graph cache re-checks the dynamic-shape guards stored with each entry on every lookup. Setting this flag skips that re-check (every entry is treated as a guard hit), which speeds up warm and cross-process cache loads. It is **not enabled by default** because it is a correctness hazard: a kernel first compiled at a small sequence length keeps int32 address arithmetic, and reusing it at a larger sequence length (roughly **>58k tokens/rank**) overflows int32 and reads out of bounds — surfacing as a CUDA illegal memory access or silently corrupted output. Only enable it when your token counts stay within the range the cached kernels were compiled for:
```bash
python -m ltx_pipelines.ti2vid_two_stages \
--compile 'inductor_config={"unsafe_skip_cache_dynamic_shape_guards": true}' \
--checkpoint-path=...
```
## Denoising Loop Optimization
**Gradient Estimation Denoising Loop:**
Instead of the standard Euler denoising loop, you can use gradient estimation for fewer steps (~20-30 instead of 40):
```python
from ltx_pipelines.utils import gradient_estimating_euler_denoising_loop
# Use gradient estimation denoising loop
def denoising_loop(sigmas, video_state, audio_state, stepper):
return gradient_estimating_euler_denoising_loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
transformer=transformer,
denoiser=denoiser,
ge_gamma=2.0, # Gradient estimation coefficient
)
```
This allows you to use **20-30 steps instead of 40** while maintaining quality. The gradient estimation function is defined in [`samplers.py`](../src/ltx_pipelines/utils/samplers.py).
packages/ltx-pipelines/docs/pipeline-selection.md
+48
View File
@@ -0,0 +1,48 @@
# Pipeline Selection Guide
## Quick Decision Tree
```text
Do you have an existing video to modify?
├─ YES → Use RetakePipeline (regenerate a specific time region)
│
Do you have an audio file to drive generation?
├─ YES → Use A2VidPipelineTwoStage (audio-to-video)
│
Do you need HDR output (linear float frames for EXR / tonemapping)?
├─ YES → Use HDRICLoraPipeline (video-to-video with LogC3 inverse decode)
│
Do you need to condition on existing images/videos?
├─ YES → Do you have reference videos for video-to-video?
│ ├─ YES → Use ICLoraPipeline
│ └─ NO → Do you have multiple keyframe images to interpolate?
│ ├─ YES → Use KeyframeInterpolationPipeline
│ └─ NO → Use TI2VidTwoStagesPipeline (image conditioning only)
│
└─ NO → Text-to-video only
├─ Do you need best quality?
│ └─ YES → Use TI2VidTwoStagesPipeline (recommended for production)
│
└─ Do you need fastest inference?
└─ YES → Use DistilledPipeline (with 8 predefined sigmas)
```
> **Note:** [`TI2VidOneStagePipeline`](../src/ltx_pipelines/ti2vid_one_stage.py) is primarily for educational purposes. For best quality, use two-stage pipelines ([`TI2VidTwoStagesPipeline`](../src/ltx_pipelines/ti2vid_two_stages.py), [`TI2VidTwoStagesHQPipeline`](../src/ltx_pipelines/ti2vid_two_stages_hq.py), [`ICLoraPipeline`](../src/ltx_pipelines/ic_lora.py), [`KeyframeInterpolationPipeline`](../src/ltx_pipelines/keyframe_interpolation.py), [`A2VidPipelineTwoStage`](../src/ltx_pipelines/a2vid_two_stage.py), or [`DistilledPipeline`](../src/ltx_pipelines/distilled.py)). For editing existing videos, use [`RetakePipeline`](../src/ltx_pipelines/retake.py).
## Features Comparison
| Pipeline | Stages | [Multimodal Guidance](multimodal-guidance.md) | Upsampling | Conditioning | Best For |
| -------- | ------ | --- | ---------- | ------------- | -------- |
| [**TI2VidTwoStagesPipeline**](pipelines.md#1-ti2vidtwostagespipeline) | 2 | ✅ | ✅ | Image | **Production quality** (recommended) |
| [**TI2VidTwoStagesHQPipeline**](pipelines.md#2-ti2vidtwostageshqpipeline) | 2 | ✅ | ✅ | Image | Same as above, res_2s sampler (higher quality) |
| [**TI2VidOneStagePipeline**](pipelines.md#3-ti2vidonestagepipeline) | 1 | ✅ | ❌ | Image | Educational, prototyping |
| [**DistilledPipeline**](pipelines.md#4-distilledpipeline) | 2 | ❌ | ✅ | Image | Fastest inference (8 sigmas) |
| [**ICLoraPipeline**](pipelines.md#5-iclorapipeline) | 2 | ✅ | ✅ | Image + Video | Video-to-video transformations |
| [**KeyframeInterpolationPipeline**](pipelines.md#6-keyframeinterpolationpipeline) | 2 | ✅ | ✅ | Keyframes | Animation, interpolation |
| [**A2VidPipelineTwoStage**](pipelines.md#7-a2vidpipelinetwostage) | 2 | ✅ | ✅ | Audio + Image | Audio-driven video generation |
| [**RetakePipeline**](pipelines.md#8-retakepipeline) | 1 | ✅ | ❌ | Source Video | Regenerating a time region of a video |
| [**HDRICLoraPipeline**](pipelines.md#9-hdriclorapipeline) | 2 | ❌ | ✅ | Video | HDR video-to-video (linear float output for EXR) |
| [**LipDubPipeline**](pipelines.md#10-lipdubpipeline) | 2 | ✅ | ✅ | Video + Audio | Lip dubbing with audio ref conditioning |
| [**T2AOneStagePipeline**](pipelines.md#11-t2aonestagepipeline) | 1 | Audio only | ❌ | None (text) | Text-to-audio (audio-only output, no video) |
See [Available Pipelines](pipelines.md) for a full description of each.
packages/ltx-pipelines/docs/pipelines.md
+151
View File
@@ -0,0 +1,151 @@
# Available Pipelines
Full reference for each pipeline. See the [Pipeline Selection Guide](pipeline-selection.md) to pick one.
---
## 1. TI2VidTwoStagesPipeline
**Best for:** High-quality text/image-to-video generation with upsampling. **Recommended for production use.**
**Source**: [`src/ltx_pipelines/ti2vid_two_stages.py`](../src/ltx_pipelines/ti2vid_two_stages.py)
Two-stage generation: Stage 1 generates low-resolution video with [multimodal guidance](multimodal-guidance.md), Stage 2 upsamples to 2x resolution with distilled LoRA refinement. Supports image conditioning. Highest quality output, slower than one-stage but significantly better quality.
**Use when:** Production-quality video generation, higher resolution needed, quality over speed, text-to-video with image conditioning.
---
## 2. TI2VidTwoStagesHQPipeline
**Best for:** Same two-stage text/image-to-video as TI2VidTwoStagesPipeline but with a different sampler and step count.
**Source**: [`src/ltx_pipelines/ti2vid_two_stages_hq.py`](../src/ltx_pipelines/ti2vid_two_stages_hq.py)
Uses the **res_2s** second-order sampler instead of Euler. Same stage structure (stage 1 at target resolution with CFG, stage 2 upsampling with distilled LoRA) and image conditioning support. Typically allows fewer steps for comparable quality; trade-offs differ from the default Euler-based pipeline.
**Use when:** You want the same two-stage workflow with fewer steps or prefer the res_2s sampling behavior.
---
## 3. TI2VidOneStagePipeline
**Best for:** Educational purposes and quick prototyping.
**Source**: [`src/ltx_pipelines/ti2vid_one_stage.py`](../src/ltx_pipelines/ti2vid_one_stage.py)
> **⚠️ Important:** This pipeline is primarily for educational purposes. For production-quality results, use `TI2VidTwoStagesPipeline` or other two-stage pipelines.
Single-stage generation (no upsampling) with [multimodal guidance](multimodal-guidance.md) and image conditioning support. Faster inference but lower resolution output (typically 512x768).
**Use when:** Learning how the pipeline works, quick prototyping, testing, or when high resolution is not needed.
---
## 4. DistilledPipeline
**Best for:** Fastest inference with good quality using a distilled model with predefined sigma schedule.
**Source**: [`src/ltx_pipelines/distilled.py`](../src/ltx_pipelines/distilled.py)
Two-stage generation with 8 predefined sigmas (8 steps in stage 1, 4 steps in stage 2). No guidance required. Fastest inference among all pipelines. Supports image conditioning. Requires spatial upsampler.
**Use when:** Fastest inference is critical, batch processing many videos, or when you have a distilled model checkpoint.
---
## 5. ICLoraPipeline
**Best for:** Video-to-video and image-to-video transformations using IC-LoRA.
**Source**: [`src/ltx_pipelines/ic_lora.py`](../src/ltx_pipelines/ic_lora.py)
Two-stage generation with IC-LoRA support. Can condition on reference videos (video-to-video) or images at specific frames. CFG guidance in stage 1, upsampling in stage 2. Requires IC-LoRA trained model.
**Note:** ICLoraPipeline can only be used with a distilled model.
**Use when:** Video-to-video transformations, image-to-video with strong control, or when you have reference videos to guide generation.
---
## 6. KeyframeInterpolationPipeline
**Best for:** Generating videos by interpolating between keyframe images.
**Source**: [`src/ltx_pipelines/keyframe_interpolation.py`](../src/ltx_pipelines/keyframe_interpolation.py)
Two-stage generation with keyframe interpolation. Uses guiding latents (additive conditioning) instead of replacing latents for smoother transitions. [Multimodal guidance](multimodal-guidance.md) in stage 1, upsampling in stage 2.
**Use when:** You have keyframe images and want to interpolate between them, creating smooth transitions, or animation/motion interpolation tasks.
---
## 7. A2VidPipelineTwoStage
**Best for:** Generating video driven by an input audio.
**Source**: [`src/ltx_pipelines/a2vid_two_stage.py`](../src/ltx_pipelines/a2vid_two_stage.py)
Two-stage audio-to-video generation. Stage 1 generates video at half resolution with audio conditioning (video-only denoising with the audio frozen), then Stage 2 upsamples by 2x and refines the video while keeping the audio fixed, using a distilled LoRA. The input audio is encoded via the audio VAE and used as the initial audio latent, but the original audio waveform is passed through and returned in the output to preserve fidelity. Supports image conditioning and prompt enhancement.
**Extra CLI arguments:** `--audio-path` (required), `--audio-start-time`, `--audio-max-duration`.
**Use when:** You have an audio clip and want to generate a matching video, audio-reactive video generation, or music visualization.
---
## 8. RetakePipeline
**Best for:** Regenerating a specific time region of an existing video while keeping the rest unchanged.
**Source**: [`src/ltx_pipelines/retake.py`](../src/ltx_pipelines/retake.py)
Single-stage generation that encodes the source video and audio into latents, applies a temporal region mask to mark `[start_time, end_time]` for regeneration, and denoises only the masked region from a text prompt. Content outside the time window is preserved. Supports independent control over video and audio regeneration (`regenerate_video`, `regenerate_audio` flags), and can use either the full model with CFG guidance or the distilled model with a fixed sigma schedule.
**Extra CLI arguments:** `--video-path` (required), `--start-time` (required), `--end-time` (required).
**Constraints:** Source video frame count must satisfy the 8k+1 format (e.g. 97, 193) and resolution must be multiples of 32.
**Use when:** You want to re-do a specific section of a generated video (e.g. fix a bad segment), selectively regenerate audio or video in a time window, or iterate on part of a result without re-generating the entire clip.
---
## 9. HDRICLoraPipeline
**Best for:** Video-to-video generation with HDR output for EXR export and offline tonemapping.
**Source**: [`src/ltx_pipelines/hdr_ic_lora.py`](../src/ltx_pipelines/hdr_ic_lora.py)
Two-stage video-to-video on the distilled model with an HDR IC-LoRA. Decoded latents pass through an HDR inverse transform (ARRI LogC3, auto-detected from LoRA metadata) to produce a **linear HDR float** tensor `[f, h, w, c]`. Video-only (audio skipped). Text embeddings are pre-computed externally and loaded from a `.safetensors` file. Tonemapping and EXR saving are the caller's responsibility. LoRA and embeddings: [`Lightricks/LTX-2.3-22b-IC-LoRA-HDR`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-HDR).
**Extra CLI arguments:** `--input` (mp4 or directory, required), `--output-dir` (required), `--hdr-lora` (required), `--text-embeddings` (pre-computed `.safetensors`, required), `--num-frames`, `--spatial-tile` (tiled VAE decode tile size; reduce on lower-VRAM GPUs), `--skip-mp4` (EXR only, no H.264 preview), `--exr-half` (float16 EXR), `--high-quality` (generates 2x frames internally for smoother output, ~2x slower), `--offload {none,cpu,disk}` (weight offloading; disables FP8 quantization when not `none`).
**Use when:** You need linear HDR float output for EXR export, color grading, or custom tonemapping workflows.
---
## 10. LipDubPipeline
**Best for:** Lip dubbing, rephrasing while keeping the same speaker identity and matching lip movements to new audio.
**Source**: [`src/ltx_pipelines/lipdub.py`](../src/ltx_pipelines/lipdub.py)
Uses IC-LoRA on a **distilled** checkpoint with a **single** lip-dub IC-LoRA applied in **both** stages. The reference clip provides video and audio reference tokens whose VAE latents are appended to the target audio sequence as frozen reference tokens. The frame count and frame rate are derived from the reference video (frame count is silently snapped to the nearest `8k+1`), so the CLI does not accept `--num-frames` or `--frame-rate`. Required: `--reference-video`. Optional: `--reference-strength`. LoRA: [`Lightricks/LTX-2.3-22b-IC-LoRA-LipDub`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-LipDub).
**Note:** Requires a distilled model checkpoint and one lip-dub IC-LoRA (`--lora` exactly once).
**Use when:** Dubbing, rephrasing with matched lips and speaker identity.
---
## 11. T2AOneStagePipeline
**Best for:** Text-to-audio — generating speech/audio only (no video) from a text prompt, e.g. driving an audio-style LoRA such as an accent LoRA.
**Source**: [`src/ltx_pipelines/t2a_one_stage.py`](../src/ltx_pipelines/t2a_one_stage.py)
Single-stage, **audio-only** generation: the video branch is absent (`video=None`), so only the audio modality is denoised and decoded through the audio VAE + vocoder, producing a wave file. Audio duration is derived from `--num-frames` / `--frame-rate` (the same `8k+1` frame convention as video). Audio guidance (CFG/STG) is optional — the `--audio-*` flags default to the model's values; the video→audio cross-modal guidance is disabled since there is no video modality.
**Extra CLI arguments (all optional, with sensible defaults):** `--num-frames`, `--frame-rate`, `--negative-prompt`, `--audio-cfg-guidance-scale`, `--audio-stg-guidance-scale`, `--audio-stg-blocks`, `--audio-rescale-scale`, `--audio-skip-step`. No `--height/--width/--image` (audio has no spatial dimensions).
**Use when:** You need speech/audio from text alone, or to evaluate an audio-only LoRA (accent, voice style) without generating video.
packages/ltx-pipelines/pyproject.toml
+2 -2
View File
@@ -1,10 +1,10 @@
[project]
name = "ltx-pipelines"
version = "1.1.6"
version = "1.1.7"
description = "Pipelines implementation for Lightricks' LTX-2 model"
readme = "README.md"
requires-python = ">=3.10"
dependencies = ["ltx-core", "av", "tqdm", "pillow", "openimageio"]
dependencies = ["ltx-core", "av", "tqdm", "pillow", "openimageio", "cloudpickle>=3.1"]
[build-system]
requires = ["uv_build>=0.9.8,<0.10.0"]
packages/ltx-pipelines/src/ltx_pipelines/__init__.py
+43 -9
View File
@@ -11,17 +11,38 @@ This package provides ready-to-use pipelines for video generation:
- RetakePipeline: Regenerate a time region (retake) of an existing video
For more detailed components and utilities, import from specific submodules
like `ltx_pipelines.utils.media_io` or `ltx_pipelines.utils.constants`.
Pipeline classes are imported lazily (PEP 562). Importing this package therefore
does not eagerly pull in every pipeline module, which keeps `import ltx_pipelines`
light and avoids the runpy double-import warning when a pipeline is run as a module
(e.g. `python -m ltx_pipelines.distilled`).
"""
from ltx_pipelines.a2vid_two_stage import A2VidPipelineTwoStage
from ltx_pipelines.distilled import DistilledPipeline
from ltx_pipelines.ic_lora import ICLoraPipeline
from ltx_pipelines.keyframe_interpolation import KeyframeInterpolationPipeline
from ltx_pipelines.lipdub import LipDubPipeline
from ltx_pipelines.retake import RetakePipeline
from ltx_pipelines.t2a_one_stage import T2AOneStagePipeline
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
import importlib
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from ltx_pipelines.a2vid_two_stage import A2VidPipelineTwoStage
from ltx_pipelines.distilled import DistilledPipeline
from ltx_pipelines.ic_lora import ICLoraPipeline
from ltx_pipelines.keyframe_interpolation import KeyframeInterpolationPipeline
from ltx_pipelines.lipdub import LipDubPipeline
from ltx_pipelines.retake import RetakePipeline
from ltx_pipelines.t2a_one_stage import T2AOneStagePipeline
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
# Public name -> module that defines it. Used for lazy resolution in __getattr__.
_EXPORTS = {
"A2VidPipelineTwoStage": "ltx_pipelines.a2vid_two_stage",
"DistilledPipeline": "ltx_pipelines.distilled",
"ICLoraPipeline": "ltx_pipelines.ic_lora",
"KeyframeInterpolationPipeline": "ltx_pipelines.keyframe_interpolation",
"LipDubPipeline": "ltx_pipelines.lipdub",
"RetakePipeline": "ltx_pipelines.retake",
"T2AOneStagePipeline": "ltx_pipelines.t2a_one_stage",
"TI2VidOneStagePipeline": "ltx_pipelines.ti2vid_one_stage",
"TI2VidTwoStagesPipeline": "ltx_pipelines.ti2vid_two_stages",
}
__all__ = [
"A2VidPipelineTwoStage",
@@ -34,3 +55,16 @@ __all__ = [
"TI2VidOneStagePipeline",
"TI2VidTwoStagesPipeline",
]
def __getattr__(name: str) -> object:
module_path = _EXPORTS.get(name)
if module_path is None:
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
value = getattr(importlib.import_module(module_path), name)
globals()[name] = value # cache so later lookups skip __getattr__
return value
def __dir__() -> list[str]:
return sorted({*globals(), *_EXPORTS})
packages/ltx-pipelines/src/ltx_pipelines/a2vid_two_stage.py
+29 -8
View File
@@ -13,6 +13,7 @@ from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, AudioLatentShape, VideoPixelShape
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.args import default_2_stage_arg_parser
from ltx_pipelines.utils.blocks import (
AudioConditioner,
@@ -43,7 +44,7 @@ class A2VidPipelineTwoStage:
both video and audio using a distilled LoRA for higher quality output.
"""
def __init__(
def __init__( # noqa: PLR0913
self,
checkpoint_path: str,
distilled_lora: list[LoraPathStrengthAndSDOps],
@@ -55,17 +56,28 @@ class A2VidPipelineTwoStage:
registry: Registry | None = None,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
):
self.device = device or get_device()
self.dtype = torch.bfloat16
self._scheduler = LTX2Scheduler()
self.prompt_encoder = PromptEncoder(
checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
checkpoint_path,
gemma_root,
self.dtype,
self.device,
registry=registry,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
self.audio_conditioner = AudioConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
self.stage_1 = DiffusionStage(
self.image_conditioner = ImageConditioner(
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
)
self.audio_conditioner = AudioConditioner(
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
)
self.stage_1 = DiffusionStage.from_checkpoint(
checkpoint_path,
self.dtype,
self.device,
@@ -74,9 +86,10 @@ class A2VidPipelineTwoStage:
registry=registry,
compilation_config=compilation_config,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
self.stage_2 = DiffusionStage(
self.stage_2 = DiffusionStage.from_checkpoint(
checkpoint_path,
self.dtype,
self.device,
@@ -85,11 +98,19 @@ class A2VidPipelineTwoStage:
registry=registry,
compilation_config=compilation_config,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.upsampler = VideoUpsampler(
checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
checkpoint_path,
spatial_upsampler_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.video_decoder = VideoDecoder(
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
)
self.video_decoder = VideoDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
def __call__( # noqa: PLR0913
self,
packages/ltx-pipelines/src/ltx_pipelines/distilled.py
+34 -9
View File
@@ -10,10 +10,11 @@ from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.args import (
ImageConditioningInput,
default_2_stage_distilled_arg_parser,
detect_checkpoint_path,
resolve_cli_params,
)
from ltx_pipelines.utils.blocks import (
AudioDecoder,
@@ -26,7 +27,6 @@ from ltx_pipelines.utils.blocks import (
from ltx_pipelines.utils.constants import (
DISTILLED_SIGMAS,
STAGE_2_DISTILLED_SIGMAS,
detect_params,
)
from ltx_pipelines.utils.denoisers import SimpleDenoiser
from ltx_pipelines.utils.helpers import (
@@ -56,6 +56,7 @@ class DistilledPipeline:
registry: Registry | None = None,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
):
self.device = device or get_device()
self.dtype = torch.bfloat16
@@ -67,9 +68,16 @@ class DistilledPipeline:
self.device,
registry=registry,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
self.stage = DiffusionStage(
self.image_conditioner = ImageConditioner(
distilled_checkpoint_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.stage = DiffusionStage.from_checkpoint(
distilled_checkpoint_path,
self.dtype,
self.device,
@@ -78,12 +86,30 @@ class DistilledPipeline:
registry=registry,
compilation_config=compilation_config,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.upsampler = VideoUpsampler(
distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
distilled_checkpoint_path,
spatial_upsampler_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.video_decoder = VideoDecoder(
distilled_checkpoint_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.audio_decoder = AudioDecoder(
distilled_checkpoint_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
self.audio_decoder = AudioDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
def __call__( # noqa: PLR0913
self,
@@ -182,8 +208,7 @@ class DistilledPipeline:
@torch.inference_mode()
def main() -> None:
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path(distilled=True)
params = detect_params(checkpoint_path)
params = resolve_cli_params(distilled=True)
parser = default_2_stage_distilled_arg_parser(params=params)
args = parser.parse_args()
pipeline = DistilledPipeline(
packages/ltx-pipelines/src/ltx_pipelines/distilled_mgpu.py
+193
View File
@@ -0,0 +1,193 @@
"""Multi-GPU distilled video runner.
Runs :class:`DistilledPipeline` across multiple GPUs with:
- **Shared stage** -- sequence parallelism (SP); the same DiffusionStage is
reused for both stage 1 (half-res) and stage 2 (full-res), so a single
SP wrapping covers both invocations.
- **Gemma** -- Accelerate-based parallelization
- **VAE** -- distributed decoding
Requires ``ltx-kernels`` to be installed (transitive via SP builder).
"""
from __future__ import annotations
import logging
from collections.abc import Callable, Iterator
from multiprocessing import SimpleQueue
from typing import Any
import torch
import torch.distributed as dist
from ltx_core.loader.registry import StateDictRegistry
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import get_video_chunks_number
from ltx_core.model.video_vae.tiling import TilingConfig
from ltx_core.multigpu.transformer.attention import AttentionManager
from ltx_core.quantization import QuantizationPolicy
from ltx_core.quantization.fp8_cast import build_policy as _build_fp8_cast_policy
from ltx_core.tiling import DimensionTilingConfig, TileCountConfig, balanced_tile_split
from ltx_pipelines.distilled import DistilledPipeline
from ltx_pipelines.multigpu.controller import MGPUController
from ltx_pipelines.multigpu.gemma_builders import AccelerateGemmaBuilder
from ltx_pipelines.multigpu.runner import MGPURunner
from ltx_pipelines.multigpu.sp_builder import SequenceParallelBuilder
from ltx_pipelines.multigpu.vae_builders import DistributedDecoderBuilder
from ltx_pipelines.multigpu.weight_tracker import TransformerWeightTracker
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.media_io import encode_video
logger = logging.getLogger(__name__)
# Stage 2 (full-res) dominates: 1024x1536, 121 frames ~= 24576 video tokens + audio tokens.
_DEFAULT_SP_MAX_TOKENS = 32768
# Rank that collects distributed-VAE tiles and encodes the assembled video.
_DRIVER_RANK = 0
class DistilledRunner(MGPURunner):
"""Distributed :class:`DistilledPipeline`: SP shared stage + Accelerate Gemma + distributed VAE."""
@torch.inference_mode()
def setup(
self,
*,
distilled_checkpoint_path: str,
gemma_root: str,
spatial_upsampler_path: str,
vae_queue: SimpleQueue,
compilation_config: CompilationConfig | None = None,
sp_max_tokens: int = _DEFAULT_SP_MAX_TOKENS,
quantization: Callable[[], QuantizationPolicy] | None = None,
) -> None:
# quantization is a picklable zero-arg builder (built per worker, post-spawn); default fp8-cast.
quantization_policy = (
quantization() if quantization is not None else _build_fp8_cast_policy(distilled_checkpoint_path)
)
registry = StateDictRegistry()
pipeline = DistilledPipeline(
distilled_checkpoint_path=distilled_checkpoint_path,
gemma_root=gemma_root,
spatial_upsampler_path=spatial_upsampler_path,
loras=[],
registry=registry,
quantization=quantization_policy,
compilation_config=compilation_config,
alloc_trim_strategy=AllocatorTrimStrategy.DEFER,
)
tracker = TransformerWeightTracker(group=self.groups.transformer_group)
# Shared stage: sequence parallelism (covers both stage 1 and stage 2 invocations).
model_cfg = pipeline.stage._transformer_builder.model_config().get("transformer", {})
attn_mgr = AttentionManager(
max_tokens=sp_max_tokens,
num_heads=model_cfg["num_attention_heads"],
head_dim=model_cfg["attention_head_dim"],
tensor_dtype=pipeline.dtype,
group=self.groups.transformer_group,
)
pipeline.stage._transformer_builder = SequenceParallelBuilder(
inner=pipeline.stage._transformer_builder,
attn_mgr=attn_mgr,
registry=registry,
tracker=tracker,
)
# Accelerate Gemma parallelization.
pipeline.prompt_encoder._text_encoder_builder = AccelerateGemmaBuilder(
gemma_root_path=gemma_root,
gemma_group=self.groups.gemma_group,
broadcast_group=self.groups.transformer_group,
registry=registry,
src_rank=_DRIVER_RANK,
dtype=pipeline.dtype,
)
# Distributed VAE decoding: balanced 2D spatial grid over the group (one tile/rank).
# height takes the smaller factor of world_size, width the larger; size-aware split is a follow-up.
vae_height_tiles, vae_width_tiles = balanced_tile_split(dist.get_world_size(self.groups.vae_group))
vae_tiling = TileCountConfig(
height=DimensionTilingConfig(num_tiles=vae_height_tiles, overlap=4),
width=DimensionTilingConfig(num_tiles=vae_width_tiles, overlap=4),
)
pipeline.video_decoder._decoder_builder = DistributedDecoderBuilder( # type: ignore[assignment]
inner=pipeline.video_decoder._decoder_builder,
queue=vae_queue,
vae_group=self.groups.vae_group,
vae_tiling=vae_tiling,
driver_rank=_DRIVER_RANK,
registry=registry,
)
self._pipeline = pipeline
@torch.inference_mode()
def __call__(
self,
*,
output_path: str,
prompt: str,
seed: int,
height: int,
width: int,
num_frames: int,
frame_rate: float,
images: list[Any] | None = None,
) -> Iterator[str | None]:
# The pipeline raises ValueError on invalid input (symmetric across ranks); the controller
# catches that and turns it into a recoverable RunnerError. Anything else is fatal.
video, audio = self._pipeline(
prompt=prompt,
seed=seed,
height=height,
width=width,
num_frames=num_frames,
frame_rate=frame_rate,
images=images or [],
tiling_config=None,
)
if dist.get_rank() != _DRIVER_RANK:
yield None # workers: nothing to encode
return
encode_video(
video=video,
fps=frame_rate,
audio=audio,
output_path=output_path,
video_chunks_number=get_video_chunks_number(num_frames, TilingConfig.default()),
)
yield output_path
if __name__ == "__main__":
from ltx_pipelines.utils.args import (
default_2_stage_distilled_arg_parser,
resolve_cli_params,
)
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
params = resolve_cli_params(distilled=True)
args = default_2_stage_distilled_arg_parser(params=params).parse_args()
vae_queue = torch.multiprocessing.get_context("spawn").SimpleQueue()
controller = MGPUController(DistilledRunner)
controller.start(
distilled_checkpoint_path=args.distilled_checkpoint_path,
gemma_root=args.gemma_root,
spatial_upsampler_path=args.spatial_upsampler_path,
vae_queue=vae_queue,
compilation_config=args.compile,
)
try:
for _ in controller.stream(
output_path=args.output_path,
prompt=args.prompt,
seed=args.seed,
height=args.height,
width=args.width,
num_frames=args.num_frames,
frame_rate=args.frame_rate,
images=args.images,
):
pass # drive the job to completion; the runner writes the file as a side effect
finally:
controller.shutdown()
packages/ltx-pipelines/src/ltx_pipelines/hdr_ic_lora.py
+32 -10
View File
@@ -39,6 +39,7 @@ from ltx_core.conditioning import (
ConditioningItem,
VideoConditionByReferenceLatent,
)
from ltx_core.devices import empty_device_cache
from ltx_core.hdr import apply_hdr_decode_postprocess
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import Registry
@@ -49,6 +50,7 @@ from ltx_core.quantization import QuantizationPolicy
from ltx_core.tiling import DimensionTilingConfig, TileCountConfig
from ltx_core.tools import VideoLatentTools
from ltx_core.types import VideoLatentShape, VideoPixelShape
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.blocks import (
DiffusionStage,
ImageConditioner,
@@ -199,7 +201,7 @@ class HDRICLoraPipeline:
Tonemapping and EXR saving are the caller's responsibility.
"""
def __init__(
def __init__( # noqa: PLR0913
self,
distilled_checkpoint_path: str,
spatial_upsampler_path: str,
@@ -211,6 +213,7 @@ class HDRICLoraPipeline:
hdr_lora_config: HdrLoraConfig | None = None,
tiled_vae_encode_pixel_threshold: int = TILED_VAE_ENCODE_PIXEL_THRESHOLD,
offload_mode: OffloadMode = OffloadMode.NONE,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
):
"""
Args:
@@ -252,17 +255,14 @@ class HDRICLoraPipeline:
f.get_tensor("audio_context"),
)
self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
self.stage_1 = DiffusionStage(
self.image_conditioner = ImageConditioner(
distilled_checkpoint_path,
self.dtype,
self.device,
loras=loras,
quantization=quantization,
registry=registry,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.stage_2 = DiffusionStage(
self.stage_1 = DiffusionStage.from_checkpoint(
distilled_checkpoint_path,
self.dtype,
self.device,
@@ -270,11 +270,33 @@ class HDRICLoraPipeline:
quantization=quantization,
registry=registry,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.stage_2 = DiffusionStage.from_checkpoint(
distilled_checkpoint_path,
self.dtype,
self.device,
loras=loras,
quantization=quantization,
registry=registry,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.upsampler = VideoUpsampler(
distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
distilled_checkpoint_path,
spatial_upsampler_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.video_decoder = VideoDecoder(
distilled_checkpoint_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
# HDR config: explicit override, or auto-detect from LoRA metadata.
if hdr_lora_config is not None:
@@ -721,7 +743,7 @@ def _process_single_video( # noqa: PLR0913
del hdr_video
gc.collect()
torch.cuda.empty_cache()
empty_device_cache()
if not skip_mp4:
# Wait for EXR saves to finish before encoding.
packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py
+36 -10
View File
@@ -16,12 +16,13 @@ from ltx_pipelines.iclora_utils import (
read_lora_reference_downscale_factor,
read_lora_reference_temporal_scale_factor,
)
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.args import (
ImageConditioningInput,
VideoConditioningAction,
VideoMaskConditioningAction,
default_2_stage_distilled_arg_parser,
detect_checkpoint_path,
resolve_cli_params,
)
from ltx_pipelines.utils.blocks import (
AudioDecoder,
@@ -34,7 +35,6 @@ from ltx_pipelines.utils.blocks import (
from ltx_pipelines.utils.constants import (
DISTILLED_SIGMAS,
STAGE_2_DISTILLED_SIGMAS,
detect_params,
)
from ltx_pipelines.utils.denoisers import SimpleDenoiser
from ltx_pipelines.utils.helpers import assert_resolution, combined_image_conditionings, get_device
@@ -64,6 +64,7 @@ class ICLoraPipeline:
registry: Registry | None = None,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
):
self.device = device or get_device()
self.dtype = torch.bfloat16
@@ -75,9 +76,16 @@ class ICLoraPipeline:
self.device,
registry=registry,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
self.stage_1 = DiffusionStage(
self.image_conditioner = ImageConditioner(
distilled_checkpoint_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.stage_1 = DiffusionStage.from_checkpoint(
distilled_checkpoint_path,
self.dtype,
self.device,
@@ -86,8 +94,9 @@ class ICLoraPipeline:
registry=registry,
compilation_config=compilation_config,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.stage_2 = DiffusionStage(
self.stage_2 = DiffusionStage.from_checkpoint(
distilled_checkpoint_path,
self.dtype,
self.device,
@@ -96,12 +105,30 @@ class ICLoraPipeline:
registry=registry,
compilation_config=compilation_config,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.upsampler = VideoUpsampler(
distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
distilled_checkpoint_path,
spatial_upsampler_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.video_decoder = VideoDecoder(
distilled_checkpoint_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.audio_decoder = AudioDecoder(
distilled_checkpoint_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
self.audio_decoder = AudioDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
# Read reference scale factors from LoRA metadata.
# IC-LoRAs trained with scaled reference videos store these factors
@@ -344,8 +371,7 @@ class ICLoraPipeline:
@torch.inference_mode()
def main() -> None:
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path(distilled=True)
params = detect_params(checkpoint_path)
params = resolve_cli_params(distilled=True)
parser = default_2_stage_distilled_arg_parser(params=params)
parser.add_argument(
"--video-conditioning",
packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py
+31 -12
View File
@@ -16,10 +16,11 @@ from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, VideoPixelShape
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.args import (
ImageConditioningInput,
default_2_stage_arg_parser,
detect_checkpoint_path,
resolve_cli_params,
)
from ltx_pipelines.utils.blocks import (
AudioDecoder,
@@ -31,7 +32,6 @@ from ltx_pipelines.utils.blocks import (
)
from ltx_pipelines.utils.constants import (
STAGE_2_DISTILLED_SIGMAS,
detect_params,
)
from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser, SimpleDenoiser
from ltx_pipelines.utils.helpers import (
@@ -53,7 +53,7 @@ class KeyframeInterpolationPipeline:
as the upsampled video already has good quality and just needs refinement.
"""
def __init__(
def __init__( # noqa: PLR0913
self,
checkpoint_path: str,
distilled_lora: list[LoraPathStrengthAndSDOps],
@@ -65,16 +65,25 @@ class KeyframeInterpolationPipeline:
registry: Registry | None = None,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
):
self.device = device or get_device()
self.dtype = torch.bfloat16
self._scheduler = LTX2Scheduler()
self.prompt_encoder = PromptEncoder(
checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
checkpoint_path,
gemma_root,
self.dtype,
self.device,
registry=registry,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
self.stage_1 = DiffusionStage(
self.image_conditioner = ImageConditioner(
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
)
self.stage_1 = DiffusionStage.from_checkpoint(
checkpoint_path,
self.dtype,
self.device,
@@ -83,9 +92,10 @@ class KeyframeInterpolationPipeline:
registry=registry,
compilation_config=compilation_config,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
self.stage_2 = DiffusionStage(
self.stage_2 = DiffusionStage.from_checkpoint(
checkpoint_path,
self.dtype,
self.device,
@@ -94,12 +104,22 @@ class KeyframeInterpolationPipeline:
registry=registry,
compilation_config=compilation_config,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.upsampler = VideoUpsampler(
checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
checkpoint_path,
spatial_upsampler_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.video_decoder = VideoDecoder(
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
)
self.audio_decoder = AudioDecoder(
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
)
self.video_decoder = VideoDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
self.audio_decoder = AudioDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
def __call__( # noqa: PLR0913
self,
@@ -235,8 +255,7 @@ class KeyframeInterpolationPipeline:
@torch.inference_mode()
def main() -> None:
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path()
params = detect_params(checkpoint_path)
params = resolve_cli_params()
parser = default_2_stage_arg_parser(params=params)
args = parser.parse_args()
pipeline = KeyframeInterpolationPipeline(
packages/ltx-pipelines/src/ltx_pipelines/lipdub.py
+36 -9
View File
@@ -21,10 +21,11 @@ from ltx_pipelines.iclora_utils import (
append_ic_lora_reference_video_conditionings,
read_lora_reference_downscale_factor,
)
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.args import (
ImageConditioningInput,
detect_checkpoint_path,
lipdub_arg_parser,
resolve_cli_params,
)
from ltx_pipelines.utils.blocks import (
AudioConditioner,
@@ -35,7 +36,7 @@ from ltx_pipelines.utils.blocks import (
VideoDecoder,
VideoUpsampler,
)
from ltx_pipelines.utils.constants import DISTILLED_SIGMAS, STAGE_2_DISTILLED_SIGMAS, detect_params
from ltx_pipelines.utils.constants import DISTILLED_SIGMAS, STAGE_2_DISTILLED_SIGMAS
from ltx_pipelines.utils.denoisers import SimpleDenoiser
from ltx_pipelines.utils.helpers import assert_resolution, combined_image_conditionings, get_device
from ltx_pipelines.utils.media_io import decode_audio_from_file, encode_video, get_videostream_metadata
@@ -62,6 +63,7 @@ class LipDubPipeline:
registry: Registry | None = None,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
) -> None:
self.device = device or get_device()
self.dtype = torch.bfloat16
@@ -75,15 +77,23 @@ class LipDubPipeline:
self.device,
registry=registry,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.image_conditioner = ImageConditioner(
distilled_checkpoint_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
self.audio_conditioner = AudioConditioner(
distilled_checkpoint_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.stage = DiffusionStage(
self.stage = DiffusionStage.from_checkpoint(
distilled_checkpoint_path,
self.dtype,
self.device,
@@ -92,12 +102,30 @@ class LipDubPipeline:
registry=registry,
compilation_config=compilation_config,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.upsampler = VideoUpsampler(
distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
distilled_checkpoint_path,
spatial_upsampler_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.video_decoder = VideoDecoder(
distilled_checkpoint_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.audio_decoder = AudioDecoder(
distilled_checkpoint_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
self.audio_decoder = AudioDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
self.reference_downscale_factor = read_lora_reference_downscale_factor(ic_lora.path)
def _create_stage_conditionings(
@@ -291,8 +319,7 @@ def patchify_lipdub_audio_reference_latent(
@torch.inference_mode()
def main() -> None:
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path(distilled=True)
params = detect_params(checkpoint_path)
params = resolve_cli_params(distilled=True)
parser = lipdub_arg_parser(params=params)
args = parser.parse_args()
packages/ltx-pipelines/src/ltx_pipelines/multigpu/__init__.py
+31
View File
@@ -0,0 +1,31 @@
"""Multi-GPU inference controller.
Public API:
- ``MGPUController``: controller-driven persistent multi-GPU fleet (start / stream / drain / shutdown)
- ``MGPURunner``: base class implemented per pipeline (setup + __call__; __call__ is a generator)
- ``Stream``: handle returned by ``controller.stream(...)``; iterate it for each rank's yields as they arrive
- ``RunnerError``: raised by a runner for a recoverable, symmetric failure
- ``SymmetricRunnerError`` / ``AsymmetricRunnerError``: caller-side exceptions raised from a job's result
- ``ControllerBusyError``: ``stream()`` called while a previous job is still uncollected
- ``NCCLGroups``: per-component NCCL process groups passed to a runner's ``setup``
"""
from ltx_pipelines.multigpu.controller import (
AsymmetricRunnerError,
ControllerBusyError,
MGPUController,
Stream,
SymmetricRunnerError,
)
from ltx_pipelines.multigpu.nccl_groups import NCCLGroups
from ltx_pipelines.multigpu.runner import MGPURunner, RunnerError
__all__ = [
"AsymmetricRunnerError",
"ControllerBusyError",
"MGPUController",
"MGPURunner",
"NCCLGroups",
"RunnerError",
"Stream",
"SymmetricRunnerError",
]
packages/ltx-pipelines/src/ltx_pipelines/multigpu/_broadcast.py
+68
View File
@@ -0,0 +1,68 @@
"""Store-based broadcast signaling between the relay rank and workers.
Workers poll an incrementing counter in the process group's store instead of
blocking directly on a collective, so an idle fleet doesn't trip the NCCL
watchdog timeout while waiting for the next job.
"""
from __future__ import annotations
import time
from collections.abc import Iterator
from contextlib import contextmanager
import torch.distributed as dist
class BroadcastCoordinator:
"""Coordinates broadcast signaling between the relay and workers via a distributed store.
Uses an incrementing counter to signal when a broadcast is ready. Workers
poll for counter changes to avoid NCCL timeout issues during idle periods.
The counter wraps at a large value to prevent overflow.
"""
def __init__(self, store: dist.Store, is_driver: bool) -> None:
"""Initialize the broadcast coordinator.
Args:
store: Distributed store (TCPStore, PrefixStore, etc.) for coordination across ranks.
is_driver: Whether this rank is the relay (signals broadcasts).
"""
self.store = store
self.is_driver = is_driver
self.broadcast_key = "dist_pipeline_broadcast_ready"
self.signal = 0
self.signal_max = 2**31 - 1 # ~2.1 billion calls before wraparound
if is_driver:
self.store.set(self.broadcast_key, str(self.signal))
self.last_seen_signal = self.get_current_signal() if not is_driver else 0
@contextmanager
def broadcast_context(self) -> Iterator[None]:
"""Signal a broadcast (driver only).
Increments the signal counter on entry to notify workers that a
broadcast is ready. The counter stays incremented for change detection.
Raises:
RuntimeError: If called on a non-driver rank.
"""
if not self.is_driver:
raise RuntimeError("broadcast_context can only be called on driver rank")
self.signal = (self.signal + 1) % self.signal_max
self.store.set(self.broadcast_key, str(self.signal))
yield
def wait_for_signal_change(self, poll_interval: float = 0.01) -> None:
"""Wait until the signal changes from the last seen value (worker only).
Args:
poll_interval: Seconds to sleep between polls (default: 0.01).
"""
while True:
current = self.get_current_signal()
if current != self.last_seen_signal:
self.last_seen_signal = current
return
time.sleep(poll_interval)
def get_current_signal(self) -> int:
"""Get the current signal value from the store."""
return int(self.store.get(self.broadcast_key).decode("utf-8"))
packages/ltx-pipelines/src/ltx_pipelines/multigpu/bp_gemma_builder.py
+91
View File
@@ -0,0 +1,91 @@
"""Batch-parallel Gemma text encoder builder.
Each rank materialises a full :class:`GemmaTextEncoder` on its own local
CUDA device via the standard :class:`SingleGPUModelBuilder` pipeline (the
same code path used by the non-MGPU pipelines). No Accelerate, no
``device_map``, no per-layer dispatch hooks.
Every ``build()`` reconstructs the encoder through
:class:`SingleGPUModelBuilder`: a fresh meta module is created, the
``GEMMA_MODEL_OPS`` chain re-runs (recomputing the rotary / position
buffers that live outside the safetensors file), and the trained weights
are bound from the provided :class:`Registry`. The registry caches the
loaded state dict so subsequent calls skip disk I/O while still rebuilding
the module tree -- mirroring the rebuild logic of
:class:`AccelerateGemmaBuilder` on this branch. Encoder-instance caching is
intentionally left out; it will arrive later as a global builder refactor.
The result is wrapped in :class:`BatchParallelGemmaWrapper`, which
partitions prompt lists across ranks in ``encode`` and routes
non-deterministic sampling (``enhance_t2v`` / ``enhance_i2v``) through a
single ``src_rank``.
"""
from __future__ import annotations
import logging
import torch
import torch.distributed as dist
from ltx_core.loader.primitives import BuilderProtocol
from ltx_core.loader.registry import Registry
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
from ltx_core.multigpu.gemma.batch_parallel_wrapper import BatchParallelGemmaWrapper
from ltx_core.text_encoders.gemma import (
GEMMA_LLM_KEY_OPS,
GEMMA_MODEL_OPS,
GemmaTextEncoderConfigurator,
module_ops_from_gemma_root,
)
from ltx_core.utils import find_matching_file
logger = logging.getLogger(__name__)
class BatchParallelGemmaBuilder(BuilderProtocol):
"""Per-rank Gemma replica builder for the batch-parallel encode path.
Mirrors the inline Gemma builder construction inside
:class:`PromptEncoder` (single-GPU path) and adds the MGPU wiring --
broadcast group + source rank for non-deterministic methods. Each
``build()`` reconstructs the encoder via :class:`SingleGPUModelBuilder`;
the registry caches the state dict so only disk I/O is skipped across
calls.
"""
def __init__(
self,
gemma_root_path: str,
broadcast_group: dist.ProcessGroup | None,
registry: Registry,
*,
src_rank: int,
dtype: torch.dtype = torch.bfloat16,
) -> None:
model_folder = find_matching_file(gemma_root_path, "model*.safetensors").parent
weight_paths = tuple(str(p) for p in model_folder.rglob("*.safetensors"))
self._inner = Builder(
model_path=weight_paths,
model_class_configurator=GemmaTextEncoderConfigurator,
model_sd_ops=GEMMA_LLM_KEY_OPS,
module_ops=(GEMMA_MODEL_OPS, *module_ops_from_gemma_root(gemma_root_path)),
registry=registry,
)
self._broadcast_group = broadcast_group
self._src_rank = src_rank
self._dtype = dtype
def model_config(self) -> dict:
return {}
def build(
self,
device: torch.device | None = None,
dtype: torch.dtype | None = None,
) -> BatchParallelGemmaWrapper:
dtype = dtype or self._dtype
encoder = self._inner.build(device=device, dtype=dtype).eval()
return BatchParallelGemmaWrapper(
encoder=encoder,
broadcast_group=self._broadcast_group,
src_rank=dist.get_group_rank(self._broadcast_group, self._src_rank),
dtype=dtype,
device=device,
)
packages/ltx-pipelines/src/ltx_pipelines/multigpu/controller.py
+382
View File
@@ -0,0 +1,382 @@
"""Synchronous multi-GPU controller.
Spawns one worker process per GPU. Rank 0 is the *relay*: the controller hands it a job over a
queue and it NCCL-broadcasts the job to every rank. All ranks run the user's runner -- a
*generator* -- in SPMD lockstep. `stream(**kwargs)` dispatches and returns a `Stream` you iterate;
each yielded value is forwarded straight to you, one element per yield, as it comes off the result
queue (no gathering across ranks). One job at a time -- no job queue, no pipelining.
Constraints and contract:
- Single machine only: MASTER_ADDR is localhost, RANK == LOCAL_RANK, one rank per GPU. By default
rank r runs on cuda:r; pass `devices=[...]` to place the fleet on a specific physical GPU subset
(rank r -> cuda:devices[r]), e.g. to run two controllers side by side on disjoint GPUs.
- The runner is a generator and runs in SPMD LOCKSTEP. Yields are forwarded individually, not
gathered, so each rank's yields appear as their own stream elements (in result-queue order). Only
the ranks' terminals are collected: once all have ended, the controller classifies and ends iteration.
- Job kwargs (non-tensor parts) cross the queue and are pickled again through the NCCL broadcast;
yielded values ride only the result queue (pickled once). All must be picklable and small.
Tensors are the exception: pass them as top-level kwargs and the relay broadcasts them over NCCL
instead of pickling (see `stream`); tensors inside a yielded value ride the result queue by
shared memory / CUDA IPC.
- Each job belongs to the thread that dispatched it: only that thread may iterate it, enforced in
`Stream`.
- One job in flight at a time: consume the `Stream` to the end before dispatching the next.
Abandoning it is NOT cleaned up -- the next `stream()` raises `ControllerBusyError` until you
`stream.drain()` or `shutdown()`.
- A runner that *raises* an unexpected exception kills the controller (a desynced NCCL collective
cannot be unwound; make a new one). For a recoverable failure the runner raises a `RunnerError`
(or a `ValueError`, which the controller turns into one) identically on every rank: the worker loop
catches it, the fleet survives, and iterating the `Stream` re-raises it as `SymmetricRunnerError`
-- or `AsymmetricRunnerError` if only some ranks raised.
"""
from __future__ import annotations
import contextlib
import logging
import threading
import time
from collections.abc import Iterator, Sequence
from datetime import timedelta
from typing import Any
import torch
import torch.multiprocessing as torch_mp
from torch.distributed.elastic.multiprocessing.api import DefaultLogsSpecs, LogsSpecs, Std
from ltx_pipelines.multigpu.fleet import _POLL_S, _Channels, _Job, _RunnersFleet
from ltx_pipelines.multigpu.runner import MGPURunner, RunnerError
logger = logging.getLogger(__name__)
_DEFAULT_INIT_TIMEOUT = timedelta(minutes=30)
# =============================================================================
# Errors raised to the caller: a busy dispatch, or a runner failure classified across ranks.
# =============================================================================
class SymmetricRunnerError(Exception):
"""Every rank raised a RunnerError -- the contract was honored. Recoverable: fix the
input and retry; the controller is still alive."""
def __init__(self, errors: list[RunnerError]) -> None:
self.errors = errors
super().__init__(errors[0].message)
class AsymmetricRunnerError(Exception):
"""Some ranks raised a RunnerError and some finished cleanly. The fleet is provably healthy
(a mix of terminals means every rank ran to completion and reported), but the runner's error
path is non-deterministic across ranks -- a latent hang risk. Loud by design; does NOT kill
the fleet."""
def __init__(self, terminals: list[Any]) -> None:
self.terminals = terminals # the per-rank terminals (RunnerErrors mixed with clean StopIterations)
super().__init__("runner raised RunnerError on some ranks but not all")
class ControllerBusyError(RuntimeError):
"""`stream()` was called while a job is still in flight (uncollected).
`stream` returns IMMEDIATELY, so it raises rather than silently draining the in-flight job
(which would block for the full job, possibly one owned by another thread). The in-flight job belongs
to the thread that dispatched it: consume it (`try: ... finally: stream.drain()`) so you never
wedge yourself, and let a concurrent caller that loses the dispatch race simply bounce.
"""
def __init__(self, job_id: int) -> None:
self.job_id = job_id # the in-flight job's id, for catch-site logging
super().__init__(f"MGPU job {job_id} is still in flight; consume it (stream.drain()) first.")
# =============================================================================
# The caller's streaming handle: a thin wrapper over the collector generator.
# Iterating it drives the job and forwards each rank's yields as they arrive.
# =============================================================================
class Stream:
"""The iterable returned by `stream()` and the handle to one in-flight job. Each element is one
rank's yielded value, forwarded as it comes off the result queue -- yields are NOT gathered
across ranks. When every rank's generator has finished, iteration ends: all clean -> the
`StopIteration` value is the per-rank list of returns (normally Nones, since results are yielded
-- `result = yield from stream` to read it); all `RunnerError` -> `SymmetricRunnerError`; a mix
-> `AsymmetricRunnerError`.
The job belongs to the thread that dispatched it: only that thread may iterate (or `drain`) this
Stream -- a second thread advancing the same ordered collection is the one hazard forbidden here.
Consume the Stream to completion before the next `stream()`; abandoning it leaves the job in
flight (the next `stream()` raises `ControllerBusyError` until it is drained or shut down). The
controller does not clean up after you -- the recommended pattern is `try: ... finally:
stream.drain()`.
"""
def __init__(self, pump: Iterator[Any], job_thread: int, job_id: int) -> None:
self._pump = pump # generator: yields each rank's values, returns the per-rank returns
self._job_thread = job_thread # the thread that dispatched this job; only it may iterate/drain
self.job_id = job_id # this job's id (for messages/debugging)
def __iter__(self) -> Stream:
return self
def __next__(self) -> Any: # noqa: ANN401
if threading.get_ident() != self._job_thread:
raise RuntimeError(
f"Stream for job {self.job_id} is single-threaded: dispatched on thread "
f"{self._job_thread}, iterated from thread {threading.get_ident()} -- a job belongs "
f"to its dispatching thread."
)
# The pump yields each rank's value directly; once all ranks have ended it raises
# StopIteration(per-rank returns), or Symmetric/Asymmetric if any rank raised.
return next(self._pump)
def drain(self) -> None:
"""Exhaust the Stream and free the controller, discarding any unconsumed yields. drain() must
be called from the same thread that called stream() (it iterates the Stream, so the owner
check applies); the recommended pattern is `try: ... finally: stream.drain()`. A recoverable
Symmetric/AsymmetricRunnerError is swallowed (cleanup); a dead/timed-out/desynced fleet still
surfaces -- you need to know.
"""
try:
for _ in self: # iterate via __next__, so the dispatching-thread check applies
pass
except (SymmetricRunnerError, AsymmetricRunnerError):
pass # cleanup: a recoverable runner error isn't worth surfacing when draining
# =============================================================================
# The controller.
# =============================================================================
class MGPUController:
"""Persistent one-job-at-a-time GPU fleet.
controller = MGPUController(MyRunner, num_gpus=8)
controller.start(**setup_kwargs)
stream = controller.stream(prompt="...")
try:
for item in stream: # one element per yield, as it arrives (not gathered)
show(item)
finally:
stream.drain() # free the controller even on early exit -- abandoning wedges it
controller.shutdown()
"""
def __init__(
self,
runner_cls: type[MGPURunner],
*,
num_gpus: int | None = None,
devices: Sequence[int] | None = None,
logs_specs: LogsSpecs | None = None,
) -> None:
"""`num_gpus` uses GPUs 0..num_gpus-1 (default: all). `devices` places the fleet on specific
physical GPUs (e.g. `[2, 3]` -> rank r on cuda:devices[r]), so several controllers can share a
box on disjoint GPU sets; the two are mutually exclusive. Indices are as the controller sees
them; every GPU stays visible to each worker, which simply binds to its assigned one.
"""
if devices is not None:
if num_gpus is not None:
raise ValueError("Pass either num_gpus or devices, not both.")
if len(devices) == 0 or len(set(devices)) != len(devices):
raise ValueError(f"devices must be non-empty and unique: {list(devices)}.")
self._runner_cls = runner_cls
self._devices = list(devices) if devices is not None else None
self._num_gpus = num_gpus
self._logs_specs = logs_specs
self._spawn_ctx = torch_mp.get_context("spawn")
self._fleet: _RunnersFleet | None = None
self._channels: _Channels | None = None
self._next_job_id = 0 # monotonic job-id generator; persists across jobs (powers desync detection)
self._inflight: Stream | None = None # the one in-flight job, as its handle; None between jobs
# baton-lock: guards ONLY the _inflight check-and-set; never held across dispatch / iteration / _collect.
self._lock = threading.Lock()
self._fatal_error: BaseException | None = None # one-way: set once, then every call raises
self._started = False
@property
def is_alive(self) -> bool:
"""True while the fleet is up and unpoisoned -- a health check that needs no try/except."""
return self._started and self._fatal_error is None
# ---------------------------------------------------------------- lifecycle
def start(
self,
*,
timeout: timedelta = _DEFAULT_INIT_TIMEOUT,
**setup_kwargs: Any, # noqa: ANN401
) -> None:
"""Spawn the fleet, run setup() on every rank, block until all report ready.
`timeout` bounds both the NCCL `init_process_group` and the controller's wait for
every rank to finish CUDA init, `create_local_nccl_groups`, and `setup()`, so a
worker wedged ALIVE turns into a clear error instead of an infinite hang (poll()
only sees a process *exit*, never a wedge). It must comfortably exceed your slowest
model load; pass a large value to effectively wait forever.
"""
if self._started:
raise RuntimeError("MGPUController.start() called twice")
if self._devices is not None:
num_gpus = len(self._devices)
device_ids: list[int] | None = list(self._devices)
else:
num_gpus = torch.cuda.device_count() if self._num_gpus is None else self._num_gpus
device_ids = None
if num_gpus <= 0:
raise ValueError(f"No GPUs available: num_gpus={num_gpus}.")
self._num_gpus = num_gpus # resolve "all GPUs" to the actual count == world size (one rank per GPU)
self._channels = _Channels(
jobs=self._spawn_ctx.Queue(),
results=self._spawn_ctx.Queue(),
ready=self._spawn_ctx.Queue(),
)
self._fleet = _RunnersFleet.spawn(
runner_cls=self._runner_cls,
setup_kwargs=setup_kwargs,
init_timeout=timeout,
channels=self._channels,
num_gpus=num_gpus,
logs_specs=self._logs_specs or DefaultLogsSpecs(tee=Std.ALL),
device_ids=device_ids,
)
try:
self._await_ready(num_gpus, timeout.total_seconds())
except BaseException:
self.shutdown() # tear the half-up fleet down so the caller need not
raise
self._started = True
logger.info("MGPU fleet ready (%d workers).", num_gpus)
def _await_ready(self, num_gpus: int, timeout: float) -> None:
assert self._channels is not None
assert self._fleet is not None
deadline = time.monotonic() + timeout
seen: set[int] = set() # which ranks have checked in (ready.put sends the rank)
while len(seen) < num_gpus:
died = self._fleet.poll()
if died is not None: # a worker exited -- catches crashes, not wedges
raise died
if time.monotonic() > deadline: # catches the wedges poll() can't see
missing = sorted(set(range(num_gpus)) - seen)
raise TimeoutError(
f"MGPU startup: {len(seen)}/{num_gpus} workers ready after {timeout:.0f}s; "
f"ranks {missing} never checked in -- stuck in init_process_group / "
f"create_local_nccl_groups / setup()? Check the worker logs."
)
try:
seen.add(self._channels.ready.get(timeout=_POLL_S))
except Exception:
continue # queue.Empty: re-check death + deadline, then retry
def shutdown(self, *, graceful_timeout: float = 60.0) -> None:
"""Tell the fleet to exit, give it a moment, then make sure it is gone.
Both teardown and kill switch -- safe to call from another thread while a job is in flight: it
force-kills the fleet, so a thread wedged on the Stream surfaces an error and unwedges (this is
how you recover a job you can't drain). A mid-job shutdown waits out `graceful_timeout` before
forcing; pass 0 to skip it.
"""
if self._fleet is None:
return
if self._fatal_error is None and self._channels is not None:
with contextlib.suppress(Exception):
self._channels.jobs.put(None) # relay broadcasts the sentinel to all ranks
self._fleet.drain(graceful_timeout)
self._fleet.terminate()
self._fleet = None
self._started = False
# ---------------------------------------------------------------- dispatch
def stream(self, *, timeout: float | None = None, **kwargs: Any) -> Stream: # noqa: ANN401
"""Dispatch one job and return IMMEDIATELY; collect later by iterating the returned Stream.
The workers run the job on their own while the controller does nothing; iterate the Stream
for the runner's yields -- each rank's yield is forwarded as its own element, as it arrives
(not gathered across ranks). Worker death or a blown timeout (measured from dispatch)
surfaces when you come back to iterate, not before.
TENSORS ARE TRANSPARENT. Pass them as top-level kwargs (`stream(latent=t, steps=30)`): the
tensor rides the queue to the relay by shared memory / CUDA IPC and the relay broadcasts it
to every rank over NCCL, so `__call__` receives it on the local GPU. (Only top-level kwargs
are broadcast this way; tensors nested in a list/dict ride the pickle path.) To send one
back, `yield` it from every rank; the yield rides the result queue, so its tensors come back
without pickling.
ONE JOB AT A TIME. Consume the Stream to the end before dispatching the next. Abandoning a
Stream is NOT cleaned up: the job stays in flight and the next `stream()` raises
`ControllerBusyError` until you `stream.drain()` or `shutdown()`. Use `try: ... finally:
stream.drain()`.
"""
if not self._started:
raise RuntimeError("MGPUController not started; call start() first")
if self._fatal_error is not None:
raise RuntimeError("MGPUController is dead; create a new one.") from self._fatal_error
assert self._channels is not None
job_thread = threading.get_ident() # this job belongs to the dispatching thread (checked in Stream)
# The baton-lock guards exactly this check-and-set -- nothing else (see the lock's comment).
with self._lock:
if self._inflight is not None:
raise ControllerBusyError(self._inflight.job_id)
job_id = self._next_job_id
self._next_job_id += 1
deadline = None if timeout is None else time.monotonic() + timeout
stream = Stream(self._collect(job_id, deadline, timeout), job_thread, job_id)
self._inflight = stream # the single source of "a job is in flight"
self._channels.jobs.put(_Job(job_id=job_id, kwargs=kwargs)) # dispatch OUTSIDE the lock; fleet starts NOW
return stream
@staticmethod
def _is_terminal(value: object) -> bool:
"""True if `value` is a rank's end marker on the wire: a clean StopIteration or a RunnerError."""
return isinstance(value, (StopIteration, RunnerError))
@staticmethod
def _classify_terminal(values: list[Any]) -> list[Any]:
"""Classify every rank's terminal (a StopIteration or a RunnerError), ordered by rank. All
RunnerError -> SymmetricRunnerError; a mix of RunnerError and clean StopIteration ->
AsymmetricRunnerError; all clean -> the per-rank return values (the caller's StopIteration.value).
Neither typed error kills the fleet."""
errs = [v for v in values if isinstance(v, RunnerError)]
if errs and len(errs) == len(values):
raise SymmetricRunnerError(errs)
if errs:
logger.error("asymmetric RunnerError across ranks: %r", values)
raise AsymmetricRunnerError(values)
return [v.value for v in values] # all StopIteration -> their return values
def _collect(self, job_id: int, deadline: float | None, timeout: float | None) -> Iterator[Any]:
"""Drain the result queue for the in-flight job: forward each yielded value to the caller as
soon as it arrives (no gathering), collecting each rank's terminal as it ends. Once every
rank has ended, classify them -- all clean -> return the per-rank returns (the caller's
StopIteration.value); all/some RunnerError -> Symmetric/Asymmetric. Watches for a dead worker
/ blown timeout meanwhile. Runs on the caller's thread when they come back to the Stream.
"""
assert self._channels is not None
assert self._fleet is not None
assert self._num_gpus is not None
channels, fleet = self._channels, self._fleet
n = self._num_gpus
terminals: dict[int, Any] = {} # rank -> its end marker (StopIteration | RunnerError)
while True:
try:
msg = channels.results.get(timeout=_POLL_S)
except Exception:
# Queue empty: nothing ready, so NOW (and only now) check for death / timeout.
died = fleet.poll()
if died is not None:
self._fatal_error = died
fleet.terminate()
raise died from None
if deadline is not None and time.monotonic() > deadline:
self._fatal_error = TimeoutError(f"MGPU job {job_id} exceeded {timeout}s")
fleet.terminate()
raise self._fatal_error from None
continue
if msg.job_id != job_id: # a rank ran a different job than we dispatched -> SPMD desync
self._fatal_error = RuntimeError(
f"MGPU fleet desync: rank {msg.rank} sent job {msg.job_id}, expected {job_id}."
)
fleet.terminate()
raise self._fatal_error from None
if not self._is_terminal(msg.value):
yield msg.value # forward this rank's yield straight to the caller -- no gathering
continue
terminals[msg.rank] = msg.value # a rank ended; hold its terminal for classification
if len(terminals) == n: # every rank has ended -> fleet free; classify and finish
self._inflight = None
return self._classify_terminal([terminals[r] for r in range(n)])
packages/ltx-pipelines/src/ltx_pipelines/multigpu/delegating_builder.py
+87
View File
@@ -0,0 +1,87 @@
"""Base class for multigpu transformer builders that delegate to an inner builder.
Shared boilerplate for SP and TDP builders — both wrap a
:class:`SingleGPUModelBuilder` and forward the ``ModelBuilderProtocol`` surface.
"""
from __future__ import annotations
import copy
from typing import TYPE_CHECKING, Generic, TypeVar
import torch
from ltx_core.loader.fuse_loras import FuseRule
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import Registry
from ltx_core.loader.sd_ops import SDOps
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
if TYPE_CHECKING:
from typing_extensions import Self
InnerModelT = TypeVar("InnerModelT", bound=torch.nn.Module)
class DelegatingBuilder(Generic[InnerModelT]):
"""Thin wrapper that delegates all ``ModelBuilderProtocol`` accessors to *inner*.
``InnerModelT`` is the type produced by the inner builder. Subclasses only
need to implement ``__init__`` and ``build`` (whose return type may differ).
"""
_inner: Builder[InnerModelT]
# -- delegated properties / with_* methods --------------------------------
@property
def checkpoint(self) -> str | tuple[str, ...]:
return self._inner.checkpoint
@property
def model_sd_ops(self) -> SDOps | None:
return self._inner.model_sd_ops
def with_sd_ops(self, sd_ops: SDOps | None) -> Self:
clone = copy.copy(self)
clone._inner = self._inner.with_sd_ops(sd_ops)
return clone
@property
def module_ops(self) -> tuple[ModuleOps, ...]:
return self._inner.module_ops
def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> Self:
clone = copy.copy(self)
clone._inner = self._inner.with_module_ops(module_ops)
return clone
@property
def loras(self) -> tuple[LoraPathStrengthAndSDOps, ...]:
return self._inner.loras
def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> Self:
clone = copy.copy(self)
clone._inner = self._inner.with_loras(loras)
return clone
@property
def registry(self) -> Registry:
return self._inner.registry
def with_registry(self, registry: Registry) -> Self:
clone = copy.copy(self)
clone._inner = self._inner.with_registry(registry)
return clone
def with_lora_load_device(self, device: torch.device) -> Self:
clone = copy.copy(self)
clone._inner = self._inner.with_lora_load_device(device)
return clone
def with_fuse_rule(self, fuse_rule: FuseRule) -> Self:
clone = copy.copy(self)
clone._inner = self._inner.with_fuse_rule(fuse_rule)
return clone
def model_config(self) -> dict:
return self._inner.model_config()
packages/ltx-pipelines/src/ltx_pipelines/multigpu/fleet.py
+375
View File
@@ -0,0 +1,375 @@
"""Worker fleet, wire protocol, and SPMD job execution driven by the MGPU controller.
Everything the controller (``MGPUController`` in ``controller.py``) drives lives here: the on-the-wire
payloads, the per-job NCCL input broadcast (``_RankLink`` / ``_run_job``), the worker entrypoint +
loops, the persistent worker fleet (``_RunnersFleet``), and ``_RunnerShipper`` (ships the runner
class to workers by value). The runner contract it executes (``MGPURunner`` / ``RunnerError``)
lives in ``runner.py``.
"""
from __future__ import annotations
import contextlib
import logging
import os
import socket
import sys
import time
from dataclasses import dataclass
from datetime import timedelta
from multiprocessing import Queue
from typing import Any
import cloudpickle
import torch
import torch.distributed as dist
from torch.distributed.elastic.multiprocessing import start_processes
from torch.distributed.elastic.multiprocessing.api import LogsSpecs
from torch.distributed.elastic.multiprocessing.errors import ProcessFailure, record
from ltx_pipelines.multigpu._broadcast import BroadcastCoordinator
from ltx_pipelines.multigpu.nccl_groups import create_local_nccl_groups
from ltx_pipelines.multigpu.runner import MGPURunner, RunnerError
logger = logging.getLogger(__name__)
_RELAY_RANK = 0
_POLL_S = 0.2 # how often the controller re-checks "did a worker die?" while waiting for a result
# =============================================================================
# The wire: what crosses the queues between controller and workers.
# =============================================================================
@dataclass
class _TensorPlaceholder:
"""Marker left in a kwargs/return dict where a tensor was lifted out for separate
transport. Carries the shape/dtype so receivers can preallocate before NCCL fills it.
Picklable and tiny -- this is what rides the object collective in the tensor's place.
"""
idx: int # position in the lifted-out tensor list
shape: tuple[int, ...]
dtype: Any # torch.dtype
def _replace_tensors_by_placeholders(d: dict[str, Any]) -> tuple[dict[str, Any], list[torch.Tensor]]:
"""Split a dict into (skeleton, tensors): top-level Tensor values become _TensorPlaceholders.
Top level only -- a tensor buried inside a list or nested dict is left alone and
will ride the pickle path. Keep tensors as direct kwargs/return values.
"""
skeleton: dict[str, Any] = {}
tensors: list[torch.Tensor] = []
for k, v in d.items():
if isinstance(v, torch.Tensor):
skeleton[k] = _TensorPlaceholder(len(tensors), tuple(v.shape), v.dtype)
tensors.append(v)
else:
skeleton[k] = v
return skeleton, tensors
def _fill_tensors_into_placeholders(skeleton: dict[str, Any], tensors: list[torch.Tensor]) -> dict[str, Any]:
"""Inverse of _replace_tensors_by_placeholders: put the tensors back where the _TensorPlaceholders are."""
return {k: (tensors[v.idx] if isinstance(v, _TensorPlaceholder) else v) for k, v in skeleton.items()}
@dataclass
class _Job:
"""One dispatched job: the controller queues it, the relay (rank 0) broadcasts it to every
rank, and all ranks run it. `kwargs` may hold top-level tensors -- broadcast over NCCL rather
than pickled (see `_RankLink`). A `None` on the job queue is the shutdown sentinel.
"""
job_id: int # every rank echoes this back; a mismatch means the fleet desynced
kwargs: dict[str, Any]
@dataclass
class _JobResult:
"""worker -> controller: one item from one rank -- a yielded value, a RunnerError, or the
terminating StopIteration itself (carrying `.value`, the generator's `return`). A yielded value
is forwarded to the caller as soon as it arrives. Terminals are collected per rank: once all
`world_size` ranks have ended, the controller classifies them -- all StopIteration -> end
iteration with `StopIteration([returns...])`; all/some RunnerError -> Symmetric/Asymmetric.
Tensors in `value` ride the queue by shared memory / CUDA IPC, the same for every rank.
"""
job_id: int # the job this is for; the controller rejects a mismatch as a desync
rank: int # which rank produced this -- used to collect one terminal per rank (and order returns)
value: Any # a yielded value, a RunnerError, or the StopIteration that ended this rank
@dataclass
class _Channels:
"""The queues bridging the controller and the workers."""
jobs: Queue # type: ignore[type-arg] # controller -> relay: _Job or None
results: Queue # type: ignore[type-arg] # all ranks -> controller: chunk items + each rank's _JobResult
ready: Queue # type: ignore[type-arg] # workers -> controller: rank, on setup-complete
# =============================================================================
# The NCCL core: the per-job input broadcast (relay -> all ranks).
# =============================================================================
class _RankLink:
"""One rank's handle to a job's *input* broadcast (relay -> all ranks, over NCCL). A job's data
crosses in two directions and this covers only the inbound one: inputs arrive here, while each
rank sends its results back out-of-band on the result queue -- never through a collective -- so
the controller assembles them without a second NCCL gather.
Between jobs the coordinator parks idle non-relay ranks on a cheap store signal instead of
leaving them blocked inside a pending broadcast (which would trip the NCCL watchdog); a rank
enters the collective only once the relay signals it has a job to send.
"""
def __init__(self, coordinator: BroadcastCoordinator, device: torch.device) -> None:
self._coordinator = coordinator
self._device = device # this rank's GPU; where received tensors land
# ---- relay side
def announce_job(self, job: _Job | None) -> _Job | None:
"""Broadcast the job to every rank, sending kwargs tensors over NCCL instead of
pickling them. Returns the job with its kwargs tensors now on THIS rank's GPU,
ready for the relay to run; None for the shutdown sentinel.
"""
with self._coordinator.broadcast_context():
if job is None:
dist.broadcast_object_list([None], src=_RELAY_RANK)
return None
skeleton, tensors = _replace_tensors_by_placeholders(job.kwargs)
gpu = [t.to(self._device).contiguous() for t in tensors] # NCCL needs CUDA + contiguous
dist.broadcast_object_list([_Job(job.job_id, skeleton)], src=_RELAY_RANK)
for t in gpu: # same order every rank, driven by the skeleton broadcast above
dist.broadcast(t, src=_RELAY_RANK)
job.kwargs = _fill_tensors_into_placeholders(skeleton, gpu)
return job
# ---- non-relay side
def await_job(self) -> _Job | None:
self._coordinator.wait_for_signal_change()
payload: list[Any] = [None]
dist.broadcast_object_list(payload, src=_RELAY_RANK)
shell: _Job | None = payload[0]
if shell is None:
return None
holes = sorted((v for v in shell.kwargs.values() if isinstance(v, _TensorPlaceholder)), key=lambda h: h.idx)
tensors: list[Any] = []
for h in holes: # receive in idx order -- matches the relay's send order
buf = torch.empty(h.shape, dtype=h.dtype, device=self._device)
dist.broadcast(buf, src=_RELAY_RANK)
tensors.append(buf)
shell.kwargs = _fill_tensors_into_placeholders(shell.kwargs, tensors)
return shell
def _run_job(runner: MGPURunner, kwargs: dict[str, Any], job_id: int, rank: int, result: Queue[_JobResult]) -> None:
"""Run one job on this rank and put each item on the result queue: every yield is a `_JobResult`
(tagged by rank), forwarded straight to the caller; the rank's end -- a clean StopIteration or a
raised RunnerError (recoverable) -- is itself a `_JobResult`, which the controller collects and
classifies. Runners must be generators; a non-generator return is not iterable, so `next` raises
(fatal).
"""
out = runner(**kwargs)
while True:
try:
value = next(out)
except StopIteration as stop: # clean finish: carries the (normally None) return value
result.put(_JobResult(job_id, rank, stop))
return
except RunnerError as err: # runner raised it explicitly -> recoverable; collect and classify
result.put(_JobResult(job_id, rank, err))
return
except ValueError as err: # input validation (symmetric) -> synthesize a recoverable RunnerError
result.put(_JobResult(job_id, rank, RunnerError(str(err))))
return
result.put(_JobResult(job_id, rank, value))
# =============================================================================
# The worker process: entrypoint + the two loops.
# =============================================================================
def _relay_loop(runner: MGPURunner, link: _RankLink, channels: _Channels) -> None:
"""Rank 0: take a job, broadcast it (tensors over NCCL), run it, stream every item on the queue."""
while True:
job: _Job | None = channels.jobs.get()
job = link.announce_job(job) # returns the job with kwargs tensors on this GPU
if job is None: # shutdown sentinel (already broadcast to the others)
return
_run_job(runner, job.kwargs, job.job_id, _RELAY_RANK, channels.results)
def _worker_loop(runner: MGPURunner, link: _RankLink, channels: _Channels, rank: int) -> None:
"""Non-relay ranks: wait for the broadcast, run in SPMD, stream every item on the queue."""
while True:
job = link.await_job()
if job is None: # shutdown sentinel
return
_run_job(runner, job.kwargs, job.job_id, rank, channels.results)
def _shutdown_distributed() -> None:
# Drop torch.compile state before tearing down NCCL: compiled artifacts hold CUDA
# pool/stream refs that ncclCommDestroy waits on, and destroy_process_group can
# deadlock without this.
torch._dynamo.reset()
if torch.cuda.is_available():
torch.cuda.synchronize()
if dist.is_initialized():
dist.destroy_process_group()
@record
def _worker_entrypoint(
runner_cls: type[MGPURunner],
setup_kwargs: dict[str, Any],
init_timeout: timedelta,
channels: _Channels,
device_ids: list[int] | None = None,
) -> None:
"""Per-worker entry. @record turns a crash into a readable failure on the controller."""
local_rank = int(os.environ["LOCAL_RANK"])
# All GPUs stay visible; we bind this rank to its physical GPU by index. (Setting
# CUDA_VISIBLE_DEVICES here is unreliable -- the spawn bootstrap can touch CUDA, freezing
# the device list, before a late env override would apply.)
device_index = local_rank if device_ids is None else device_ids[local_rank]
torch.cuda.set_device(device_index)
# Build the runner before NCCL init so its __init__ runs as the per-worker pre-init hook, for setup
# that must precede init_process_group (e.g. tests set torch.use_deterministic_algorithms there).
runner = runner_cls()
if not dist.is_initialized():
dist.init_process_group(
backend="nccl",
device_id=torch.device("cuda", device_index),
timeout=init_timeout,
)
is_relay = dist.get_rank() == _RELAY_RANK
device = torch.device("cuda", device_index)
groups = create_local_nccl_groups()
store = dist.distributed_c10d.PrefixStore("ltx_pipeline_broadcast/", dist.distributed_c10d._get_default_store())
link = _RankLink(BroadcastCoordinator(store=store, is_driver=is_relay), device)
runner._groups = groups
runner.setup(**setup_kwargs)
channels.ready.put(dist.get_rank())
try:
if is_relay:
_relay_loop(runner, link, channels)
else:
_worker_loop(runner, link, channels, dist.get_rank())
finally:
_shutdown_distributed()
# =============================================================================
# The fleet: spawn, poll for death, terminate. No job knowledge.
# =============================================================================
class _RunnerShipper:
"""Ships a runner CLASS to spawned workers by value (cloudpickle), not by reference.
A runner defined in ``__main__`` (a plain script or ``python -m``) or in a test module is
not importable under that name in a freshly spawned worker, so the stock by-reference pickle
the elastic launcher uses would raise ModuleNotFoundError on every rank. This proxy serializes
the class by value; its ``__reduce__`` targets ``cloudpickle.loads`` (importable everywhere),
so only the runner crosses by value -- the controller's own channel payloads stay on the stock
pickler. The worker unpickles it straight back to the runner class.
"""
def __init__(self, runner_cls: type[MGPURunner]) -> None:
module = sys.modules.get(runner_cls.__module__)
if module is None:
self._payload = cloudpickle.dumps(runner_cls)
return
# cloudpickle pickles a class by reference when its module looks importable; force
# by-value so a __main__/test-module runner survives the spawn.
cloudpickle.register_pickle_by_value(module)
try:
self._payload = cloudpickle.dumps(runner_cls)
finally:
cloudpickle.unregister_pickle_by_value(module)
def __reduce__(self) -> tuple[object, tuple[bytes]]:
return (cloudpickle.loads, (self._payload,))
def _format_failures(failures: dict[int, ProcessFailure]) -> RuntimeError:
lines = [f" rank {r} (pid {f.pid}) exit {f.exitcode}:\n{f.message}" for r, f in failures.items()]
return RuntimeError("MGPU worker(s) failed:\n" + "\n".join(lines))
def _find_free_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("", 0))
return s.getsockname()[1]
class _RunnersFleet:
"""The worker processes. All three methods call torch.elastic from the controller
thread only -- that is the whole reason there is no lock anywhere.
"""
def __init__(self, pcontext: Any) -> None: # noqa: ANN401
self._pcontext = pcontext
@classmethod
def spawn(
cls,
*,
runner_cls: type[MGPURunner],
setup_kwargs: dict[str, Any],
init_timeout: timedelta,
channels: _Channels,
num_gpus: int,
logs_specs: LogsSpecs,
device_ids: list[int] | None = None,
) -> _RunnersFleet:
port = _find_free_port()
envs = {
r: {
"RANK": str(r),
"LOCAL_RANK": str(r),
"WORLD_SIZE": str(num_gpus),
"MASTER_ADDR": "127.0.0.1",
"MASTER_PORT": str(port),
}
for r in range(num_gpus)
}
# device_ids maps rank -> physical GPU (None = identity); the worker binds to it. The runner
# class ships by value so a __main__/test-module runner survives the spawn (see _RunnerShipper).
packed = (_RunnerShipper(runner_cls), setup_kwargs, init_timeout, channels, device_ids)
args = dict.fromkeys(range(num_gpus), packed)
logger.info("Spawning %d MGPU workers...", num_gpus)
return cls(
start_processes(
name="ltx_mgpu_worker",
entrypoint=_worker_entrypoint,
args=args,
envs=envs,
logs_specs=logs_specs,
start_method="spawn",
)
)
def poll(self) -> RuntimeError | None:
"""None while all workers are alive. Once any has exited, an error describing it.
Used only mid-job, where ANY exit is unexpected (workers only exit on the
shutdown sentinel), so a clean exit is reported as an error too.
"""
result = self._pcontext.wait(timeout=0)
if result is None:
return None
if result.failures:
return _format_failures(result.failures)
return RuntimeError("MGPU workers exited unexpectedly")
def drain(self, timeout: float) -> bool:
"""Wait up to `timeout` for all workers to exit on their own. True if they did."""
end = time.monotonic() + timeout
while time.monotonic() < end:
if self._pcontext.wait(timeout=0) is not None:
return True
time.sleep(_POLL_S)
return False
def terminate(self) -> None:
"""SIGTERM -> SIGKILL. Never raises; safe to call more than once."""
with contextlib.suppress(Exception):
self._pcontext.close()
packages/ltx-pipelines/src/ltx_pipelines/multigpu/gemma_builders.py
+150
View File
@@ -0,0 +1,150 @@
"""Multi-GPU Gemma text encoder builder.
Replaces the text encoder builder on the ``PromptEncoder`` block with an
:class:`AccelerateGemmaBuilder` that uses ``device_map="auto"`` on the
source rank and a broadcast stub elsewhere.
On the source rank the first ``build()`` loads via HuggingFace
``from_pretrained`` and caches the full state dict (including non-persistent
buffers) in the provided :class:`Registry`. Subsequent calls recreate the
model from cache and reinstall accelerate dispatch hooks — no disk I/O.
"""
from __future__ import annotations
import copy
import logging
from typing import TYPE_CHECKING, Any
import torch
import torch.distributed as dist
from accelerate import dispatch_model
from transformers import Gemma3ForConditionalGeneration
from ltx_core.loader.primitives import BuilderProtocol, StateDict
from ltx_core.loader.registry import Registry
from ltx_core.multigpu.gemma.accelerate_wrapper import AccelerateGemmaWrapper
from ltx_core.multigpu.gemma.loader import load_gemma_with_device_map
from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder
if TYPE_CHECKING:
from typing_extensions import Self
logger = logging.getLogger(__name__)
class AccelerateGemmaBuilder(BuilderProtocol):
"""Builder that loads Gemma with ``device_map="auto"`` on the source rank.
Conforms to the builder interface expected by ``PromptEncoder``:
``build(device, dtype) -> model``. Non-source ranks get a lightweight
:class:`AccelerateGemmaWrapper` that receives embeddings via broadcast.
"""
def __init__(
self,
gemma_root_path: str,
gemma_group: dist.ProcessGroup | None,
broadcast_group: dist.ProcessGroup | None,
registry: Registry,
*,
src_rank: int,
dtype: torch.dtype = torch.bfloat16,
) -> None:
self._gemma_root_path = gemma_root_path
self._gemma_group = gemma_group
self._broadcast_group = broadcast_group
self._registry = registry
self._src_rank = src_rank
self._is_src = dist.get_rank() == src_rank
self._dtype = dtype
# Cached on the src rank after first build (non-tensor objects).
self._config: object | None = None
self._hf_device_map: dict[str, int | str] | None = None
self._tokenizer: object | None = None
self._processor: object | None = None
@property
def registry(self) -> Registry:
return self._registry
def with_registry(self, registry: Registry) -> Self:
clone = copy.copy(self)
clone._registry = registry
return clone
def model_config(self) -> dict:
return {}
def build(
self,
device: torch.device | None = None,
dtype: torch.dtype | None = None,
**_kwargs: Any, # noqa: ANN401
) -> AccelerateGemmaWrapper:
dtype = dtype or self._dtype
encoder = self._build_encoder(dtype) if self._is_src else None
return AccelerateGemmaWrapper(
encoder=encoder,
broadcast_group=self._broadcast_group,
src_rank=dist.get_group_rank(self._broadcast_group, self._src_rank),
dtype=dtype,
device=device,
)
# -- src-rank helpers ---------------------------------------------------
def _build_encoder(self, dtype: torch.dtype) -> GemmaTextEncoder:
cached = self._registry.get([self._gemma_root_path], None)
if cached is not None:
logger.info("Rebuilding Gemma from cached state dict (no disk I/O).")
return self._rebuild_from_cache(cached, dtype)
encoder = load_gemma_with_device_map(self._gemma_root_path, dtype)
# Cache non-tensor objects on the builder instance.
self._config = encoder.model.config
self._hf_device_map = encoder.model.hf_device_map
self._tokenizer = encoder.tokenizer
self._processor = encoder.processor
# Cache full state dict including non-persistent buffers.
sd = encoder.model.state_dict()
for name, buf in encoder.model.named_buffers():
if name not in sd:
sd[name] = buf
total_size = sum(t.nelement() * t.element_size() for t in sd.values())
dtypes = {t.dtype for t in sd.values()}
self._registry.add(
[self._gemma_root_path],
None,
StateDict(sd=sd, device=torch.device("meta"), size=total_size, dtype=dtypes),
)
logger.info("Cached Gemma state dict in registry (%d entries).", len(sd))
return encoder
def _rebuild_from_cache(self, cached: StateDict, dtype: torch.dtype) -> GemmaTextEncoder:
with torch.device("meta"):
model = Gemma3ForConditionalGeneration(self._config)
# Split into persistent (load_state_dict) and non-persistent (manual assign).
expected_keys = set(model.state_dict().keys())
persistent_sd = {k: v for k, v in cached.sd.items() if k in expected_keys}
non_persistent_sd = {k: v for k, v in cached.sd.items() if k not in expected_keys}
model.load_state_dict(persistent_sd, strict=True, assign=True)
for name, tensor in non_persistent_sd.items():
parent_path, attr = name.rsplit(".", 1)
module = model
for part in parent_path.split("."):
module = getattr(module, part)
setattr(module, attr, tensor)
dispatch_model(model, self._hf_device_map)
return GemmaTextEncoder(
model=model,
tokenizer=self._tokenizer,
processor=self._processor,
dtype=dtype,
)
packages/ltx-pipelines/src/ltx_pipelines/multigpu/nccl_groups.py
+31
View File
@@ -0,0 +1,31 @@
"""NCCL process group management."""
from __future__ import annotations
from dataclasses import dataclass
import torch.distributed as dist
@dataclass
class NCCLGroups:
"""Container for the NCCL process groups used by each pipeline component."""
gemma_group: dist.ProcessGroup
transformer_group: dist.ProcessGroup
vae_group: dist.ProcessGroup
def create_local_nccl_groups() -> NCCLGroups:
"""Create NCCL process groups for each pipeline component.
All ranks must call this collectively because ``dist.new_group`` is a
collective operation. All ranks participate in every group.
Returns:
NCCLGroups with one process group per component.
"""
all_ranks = list(range(dist.get_world_size()))
return NCCLGroups(
gemma_group=dist.new_group(ranks=all_ranks),
transformer_group=dist.new_group(ranks=all_ranks),
vae_group=dist.new_group(ranks=all_ranks),
)
packages/ltx-pipelines/src/ltx_pipelines/multigpu/runner.py
+62
View File
@@ -0,0 +1,62 @@
"""The runner contract: the base class a pipeline subclasses and the error it raises.
``MGPURunner`` is what a pipeline subclasses (``setup`` + a generator ``__call__``); ``RunnerError``
is the recoverable, symmetric failure a runner raises. The fleet (``fleet.py``) runs runners and
ships them to the spawned workers; the controller (``controller.py``) classifies their results.
"""
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Any
from ltx_pipelines.multigpu.nccl_groups import NCCLGroups
class RunnerError(Exception):
"""Raised by a runner (or synthesized by the controller from a ValueError) for a recoverable,
SYMMETRIC failure. The worker loop catches it and puts it on the result queue as that rank's
end; the controller collects it with the other terminals and classifies them -- every rank ->
SymmetricRunnerError, a mix with clean finishes -> AsymmetricRunnerError. Raise it IDENTICALLY on
every rank, outside any collective (e.g. validating the broadcast kwargs before the first one);
a RunnerError on only some ranks is the contract violation AsymmetricRunnerError flags.
"""
def __init__(self, message: str) -> None:
self.message = message
super().__init__(message)
class MGPURunner(ABC):
"""Subclass this. The controller builds one per worker, injects groups,
calls setup() once, then calls the instance per job. Define it anywhere -- the controller
ships the class to workers by value, so a runner in ``__main__`` or a test module works.
setup() and __call__() run on EVERY rank. __call__ MUST be a generator (use `yield`, even
once): each yield is forwarded to the Stream on its own, as it arrives -- yields are NOT gathered
across ranks. Results are *yielded*; a `return` value, if any, rides the terminal
StopIteration.value (the rare path), collected per rank once every rank has ended. The framework
does not wrap them in inference mode -- annotate your setup()/__call__ with @torch.inference_mode()
if you want it.
Tensors are transparent: pass them as kwargs (e.g. `stream(latent=t)`) and the relay
broadcasts them to every rank over NCCL, so `__call__` receives them already on the local GPU.
Yield tensors back the same way -- as values in a yielded dict -- and they come back to the
controller without being pickled. See the module docstring.
Raising an unexpected exception from __call__ is FATAL: it kills the worker, poisons the
controller, and needs a new one. For a RECOVERABLE failure raise a `RunnerError` (or a `ValueError`,
which the controller converts) -- identically on every rank, outside any collective. The worker
loop catches it and the fleet stays alive; iterating the Stream re-raises it as
SymmetricRunnerError.
"""
_groups: NCCLGroups
@property
def groups(self) -> NCCLGroups:
return self._groups
@abstractmethod
def setup(self, *args: Any, **kwargs: Any) -> None: # noqa: ANN401
...
@abstractmethod
def __call__(self, *args: Any, **kwargs: Any) -> Any: # noqa: ANN401
...
packages/ltx-pipelines/src/ltx_pipelines/multigpu/sp_builder.py
+57
View File
@@ -0,0 +1,57 @@
"""Sequence parallel transformer builder.
Wrapping builder that produces a transformer model with sequence parallelism applied.
Requires ``ltx-kernels`` to be installed.
"""
from __future__ import annotations
from typing import Generic
import torch
from ltx_core.loader.primitives import ModelBuilderProtocol
from ltx_core.loader.registry import Registry
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
from ltx_core.model.model_protocol import LTXModelProtocol
from ltx_core.multigpu.transformer.attention import AttentionManager
from ltx_core.multigpu.transformer.sequence_parallel import (
SequenceParallelModelWrapper,
create_video_self_attention_module_ops,
)
from ltx_pipelines.multigpu.delegating_builder import DelegatingBuilder, InnerModelT
from ltx_pipelines.multigpu.weight_tracker import TransformerWeightTracker
class SequenceParallelBuilder(DelegatingBuilder[InnerModelT], Generic[InnerModelT]):
"""Builder that injects SP module ops and wraps with :class:`SequenceParallelModelWrapper`."""
def __init__(
self,
inner: ModelBuilderProtocol[LTXModelProtocol],
attn_mgr: AttentionManager,
registry: Registry,
tracker: TransformerWeightTracker,
) -> None:
if not isinstance(inner, Builder):
raise TypeError(f"SequenceParallelBuilder wraps a SingleGPUModelBuilder, got {type(inner).__name__}")
cuda_device = torch.device(f"cuda:{torch.cuda.current_device()}")
inner = inner.with_registry(registry).with_lora_load_device(cuda_device)
sp_ops = create_video_self_attention_module_ops(attn_mgr)
self._inner = inner.with_module_ops((*inner.module_ops, sp_ops))
self._tracker = tracker
self._attn_mgr = attn_mgr
@property
def all2all_timeout_seconds(self) -> float:
"""The SP all2all barrier timeout (seconds); forwards to the AttentionManager that owns the refs."""
return self._attn_mgr.all2all_timeout_seconds
@all2all_timeout_seconds.setter
def all2all_timeout_seconds(self, seconds: float) -> None:
self._attn_mgr.all2all_timeout_seconds = seconds
def build(
self, device: torch.device | None = None, dtype: torch.dtype | None = None, **kwargs: object
) -> SequenceParallelModelWrapper:
model = self._tracker.build(self._inner, device=device, dtype=dtype, **kwargs)
return SequenceParallelModelWrapper(model, self._attn_mgr)
packages/ltx-pipelines/src/ltx_pipelines/multigpu/tdp_builder.py
+70
View File
@@ -0,0 +1,70 @@
"""Tiled data parallel transformer builder.
Wrapping builder that produces a transformer model with tiled data parallelism applied.
"""
from __future__ import annotations
from typing import Generic
import torch
import torch.distributed as dist
from ltx_core.loader.primitives import ModelBuilderProtocol
from ltx_core.loader.registry import Registry
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
from ltx_core.model.model_protocol import LTXModelProtocol
from ltx_core.multigpu.transformer.tiled_data_parallel import (
TiledDataParallelModelWrapper,
)
from ltx_core.tiling import TileCountConfig
from ltx_core.tools import VideoLatentTools
from ltx_pipelines.multigpu.delegating_builder import DelegatingBuilder, InnerModelT
from ltx_pipelines.multigpu.weight_tracker import TransformerWeightTracker
class TiledDataParallelBuilder(DelegatingBuilder[InnerModelT], Generic[InnerModelT]):
"""Builder conforming to :class:`ModelBuilderProtocol` that wraps with
:class:`TiledDataParallelModelWrapper`.
Requires ``video_tools`` as a keyword argument to :meth:`build` so the
wrapper can compute the tile for this rank.
The underlying model must accept ``(video, audio, perturbations)`` and return
``(denoised_video, denoised_audio)`` — i.e. conform to the ``X0Model`` forward
signature used by the LTX transformer.
"""
def __init__(
self,
inner: ModelBuilderProtocol[LTXModelProtocol],
group: dist.ProcessGroup,
tiling: TileCountConfig,
registry: Registry,
tracker: TransformerWeightTracker,
normalize_positions: bool = True,
) -> None:
if not isinstance(inner, Builder):
raise TypeError(f"TiledDataParallelBuilder wraps a SingleGPUModelBuilder, got {type(inner).__name__}")
cuda_device = torch.device(f"cuda:{torch.cuda.current_device()}")
self._inner = inner.with_registry(registry).with_lora_load_device(cuda_device)
self._tracker = tracker
self._group = group
self._tiling = tiling
self._normalize_positions = normalize_positions
def build(
self,
device: torch.device | None = None,
dtype: torch.dtype | None = None,
*,
video_tools: VideoLatentTools | None = None,
**_kwargs: object,
) -> TiledDataParallelModelWrapper:
if video_tools is None:
raise ValueError("TiledDataParallelBuilder.build() requires video_tools")
model = self._tracker.build(self._inner, device=device, dtype=dtype, **_kwargs)
return TiledDataParallelModelWrapper(
model,
video_tools=video_tools,
tiling=self._tiling,
group=self._group,
normalize_positions=self._normalize_positions,
)
packages/ltx-pipelines/src/ltx_pipelines/multigpu/vae_builders.py
+63
View File
@@ -0,0 +1,63 @@
"""Multi-GPU VAE decoder builder.
Wrapping builder that produces a :class:`DistributedVideoDecoder`.
"""
from __future__ import annotations
import copy
from typing import TYPE_CHECKING, Any
import torch
import torch.distributed as dist
from torch.multiprocessing import Queue
from ltx_core.loader.primitives import BuilderProtocol
from ltx_core.loader.registry import Registry
from ltx_core.multigpu.vae.distributed_decoder import DistributedVideoDecoder
from ltx_core.tiling import TileCountConfig
if TYPE_CHECKING:
from typing_extensions import Self
class DistributedDecoderBuilder(BuilderProtocol):
"""Builder that wraps a base decoder builder with distributed logic."""
def __init__(
self,
inner: BuilderProtocol,
queue: Queue, # type: ignore[type-arg]
vae_group: dist.ProcessGroup,
vae_tiling: TileCountConfig,
driver_rank: int,
registry: Registry,
) -> None:
self._inner = inner.with_registry(registry)
self._queue = queue
self._vae_group = vae_group
self._vae_tiling = vae_tiling
self._driver_rank = driver_rank
@property
def registry(self) -> Registry:
return self._inner.registry
def with_registry(self, registry: Registry) -> Self:
clone = copy.copy(self)
clone._inner = self._inner.with_registry(registry)
return clone
def build(
self,
device: torch.device | None = None,
dtype: torch.dtype | None = None,
**kwargs: Any, # noqa: ANN401
) -> DistributedVideoDecoder:
base_decoder = self._inner.build(device=device, dtype=dtype, **kwargs)
return DistributedVideoDecoder(
base_decoder,
queue=self._queue,
vae_group=self._vae_group,
vae_tiling=self._vae_tiling,
driver_rank=dist.get_group_rank(self._vae_group, self._driver_rank),
)
packages/ltx-pipelines/src/ltx_pipelines/multigpu/weight_tracker.py
+184
View File
@@ -0,0 +1,184 @@
"""Distributed transformer weight tracker with LoRA hot-swap.
Shared infrastructure used by both TDP and SP builders.
"""
from __future__ import annotations
import torch
import torch.distributed as dist
from ltx_core.loader.fuse_loras import fuse_lora_weights
from ltx_core.loader.primitives import LoraPathStrengthAndSDOps, LoraStateDictWithStrength, StateDict
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
from ltx_core.model.model_protocol import ModelType
from ltx_core.multigpu.sharded_sd import ShardedSD
def _apply_loras_inplace(
source: dict[str, torch.Tensor],
target: dict[str, torch.Tensor],
loras: tuple[LoraPathStrengthAndSDOps, ...],
builder: Builder, # type: ignore[type-arg]
lora_keys: frozenset[str],
) -> None:
"""Reset *target* to clean weights from *source*, then fuse all LoRAs in one pass."""
for key, clean_weight in source.items():
target[key].copy_(clean_weight)
lora_sds = [
LoraStateDictWithStrength(
builder.load_sd(
[lora.path],
sd_ops=lora.sd_ops.with_additional_allowed_keys(lora_keys),
registry=builder.registry,
device=builder.lora_load_device,
),
lora.strength,
)
for lora in loras
if lora.strength != 0
]
target_sd = StateDict(
sd=target, device=next(iter(target.values())).device, size=0, dtype={next(iter(target.values())).dtype}
)
for key, fused in fuse_lora_weights(target_sd, lora_sds, fuse_rule=builder.fuse_rule):
target[key].copy_(fused)
class TransformerWeightTracker:
"""Tracks cached transformer weights with distributed LoRA hot-swap.
Shared across stage builders that operate on the same checkpoint.
Does **not** own the model weights — it references tensors stored in a
:class:`Registry` and receives a builder at :meth:`build` time.
Uses two :class:`ShardedSD` instances (created on first :meth:`build` call):
- ``stored_sd`` — cloned backup of the original (pre-LoRA) weights.
Used to restore registry tensors before applying a different LoRA set.
- ``broadcast_sd`` — zero-copy view into the registry tensors.
After in-place LoRA fusion on the owning rank, this broadcasts the
fused results to all other ranks so every rank sees the same weights.
Both are created together and are always either both ``None`` or both set.
With ``no_lora_swap``, the LoRA set is assumed fixed (none, or one set):
the backup clone is skipped and any swap or reset raises.
"""
def __init__(self, group: dist.ProcessGroup, bucket_mb: int = 256, no_lora_swap: bool = False) -> None:
if bucket_mb <= 0:
raise ValueError("bucket_mb must be > 0")
self._group = group
self._bucket_mb = bucket_mb
self._no_lora_swap = no_lora_swap
self._staging: torch.Tensor | None = None
self.stored_sd: ShardedSD | None = None
self.broadcast_sd: ShardedSD | None = None
self.loras: list[tuple[str, float, str]] = []
@property
def staging(self) -> torch.Tensor:
"""The single broadcast scratch buffer, shared by both SDs, allocated on first use."""
if self._staging is None:
device = torch.device(f"cuda:{torch.cuda.current_device()}")
self._staging = torch.empty(self._bucket_mb * 1024 * 1024, dtype=torch.uint8, device=device)
return self._staging
def loras_match(self, lora_list: list[tuple[str, float, str]]) -> bool:
if len(lora_list) != len(self.loras):
return False
return sorted(lora_list) == sorted(self.loras)
def reset_loras(self, target_sd: dict[str, torch.Tensor]) -> None:
"""Restore *target_sd* to original (pre-LoRA) weights.
No-op if no LoRAs are currently applied. This is a cooperative
operation — all ranks must call it simultaneously.
"""
if not self.loras:
return
if self._no_lora_swap:
raise RuntimeError("no_lora_swap tracker has no backup to reset from")
if self.stored_sd is None:
raise RuntimeError("stored_sd must be initialised before reset_loras")
self.loras = []
self.stored_sd.broadcast_shards_into(target_sd, self.staging)
def _local_lora_keys(self) -> frozenset[str]:
"""Derive LoRA key names from the locally owned model keys."""
if self.stored_sd is None:
return frozenset()
keys: set[str] = set()
for k in self.stored_sd.local_shard:
if k.endswith(".weight"):
prefix = k[: -len(".weight")]
keys.add(f"{prefix}.lora_A.weight")
keys.add(f"{prefix}.lora_B.weight")
return frozenset(keys)
def apply_loras_(
self,
target_sd: dict[str, torch.Tensor],
loras: tuple[LoraPathStrengthAndSDOps, ...],
builder: Builder, # type: ignore[type-arg]
) -> None:
"""Fuse *loras* into *target_sd* in-place (trailing ``_`` denotes in-place).
Skips work when the requested LoRAs already match. Restores stored
weights before applying new LoRAs. This is a cooperative operation —
all ranks must call it simultaneously.
"""
new_loras = [(lora.path, lora.strength, lora.sd_ops.name) for lora in loras]
if self.loras_match(new_loras):
return
if self._no_lora_swap and self.loras:
raise RuntimeError(f"no_lora_swap tracker cannot change LoRAs: have {self.loras}, requested {new_loras}")
if all(lora.strength == 0 for lora in loras):
self.reset_loras(target_sd)
return
if self.stored_sd is None or self.broadcast_sd is None:
raise RuntimeError("ShardedSDs must be initialised before apply_loras_ (call build first)")
source = self.stored_sd.local_shard
target = {k: v for k, v in target_sd.items() if k in source}
lora_keys = self._local_lora_keys()
_apply_loras_inplace(source, target, loras, builder, lora_keys)
self.broadcast_sd.broadcast_shards_into(target_sd, self.staging)
self.loras = new_loras
def build(
self,
builder: Builder[ModelType],
device: torch.device | None = None,
dtype: torch.dtype | None = None,
**kwargs: object,
) -> ModelType:
"""Build the transformer model with distributed LoRA hot-swap.
Populates the registry with clean weights on first call, then applies
LoRAs in-place and broadcasts to all ranks. Assumes the builder carries
a non-dummy :class:`Registry` so that weights can be cached and reused
across calls.
"""
loras = builder.loras
clean_builder = builder.with_loras(())
model_paths = list(builder.model_path) if isinstance(builder.model_path, tuple) else [builder.model_path]
# First call: populate the registry with clean weights.
if clean_builder.registry.get(model_paths, clean_builder.model_sd_ops) is None:
clean_builder.build(device=device, dtype=dtype, **kwargs)
cached_sd = clean_builder.registry.get(model_paths, clean_builder.model_sd_ops)
if cached_sd is None:
raise RuntimeError("Expected model state dict in registry but found None")
if self.stored_sd is None:
self.stored_sd = ShardedSD.from_state_dict(cached_sd.sd, self._group, clone=not self._no_lora_swap)
self.broadcast_sd = ShardedSD.from_state_dict(cached_sd.sd, self._group, clone=False)
if loras:
self.apply_loras_(cached_sd.sd, loras, builder)
else:
self.reset_loras(cached_sd.sd)
return clean_builder.build(device=device, dtype=dtype, **kwargs)
packages/ltx-pipelines/src/ltx_pipelines/retake.py
+9 -1
View File
@@ -17,6 +17,7 @@ from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import (
SpatioTemporalScaleFactors,
)
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.args import video_editing_arg_parser
from ltx_pipelines.utils.blocks import (
AudioConditioner,
@@ -76,6 +77,7 @@ class RetakePipeline:
distilled: bool = True,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
):
self.device = device or get_device()
self.dtype = torch.bfloat16
@@ -89,20 +91,23 @@ class RetakePipeline:
device=self.device,
registry=registry,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.image_conditioner = ImageConditioner(
checkpoint_path=checkpoint_path,
dtype=self.dtype,
device=self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.audio_conditioner = AudioConditioner(
checkpoint_path=checkpoint_path,
dtype=self.dtype,
device=self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.stage = DiffusionStage(
self.stage = DiffusionStage.from_checkpoint(
checkpoint_path=checkpoint_path,
dtype=self.dtype,
device=self.device,
@@ -111,18 +116,21 @@ class RetakePipeline:
registry=registry,
compilation_config=compilation_config,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.video_decoder = VideoDecoder(
checkpoint_path=checkpoint_path,
dtype=self.dtype,
device=self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.audio_decoder = AudioDecoder(
checkpoint_path=checkpoint_path,
dtype=self.dtype,
device=self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
# --------------------------------------------------------------------- #
packages/ltx-pipelines/src/ltx_pipelines/t2a_one_stage.py
+8 -5
View File
@@ -16,16 +16,16 @@ from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio
from ltx_pipelines.utils import get_device
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.args import (
default_1_stage_t2a_arg_parser,
detect_checkpoint_path,
resolve_cli_params,
)
from ltx_pipelines.utils.blocks import (
AudioDecoder,
DiffusionStage,
PromptEncoder,
)
from ltx_pipelines.utils.constants import detect_params
from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser
from ltx_pipelines.utils.media_io import encode_audio
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
@@ -56,6 +56,7 @@ class T2AOneStagePipeline:
registry: Registry | None = None,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
):
self.dtype = torch.bfloat16
self.device = device or get_device()
@@ -67,12 +68,13 @@ class T2AOneStagePipeline:
device=self.device,
registry=registry,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
# Audio-only: build an audio-only transformer (model_configurator) so the video
# weights are never instantiated, plus a use-case-specific SDOps that restricts
# checkpoint reads to the audio model's keys, so the video weights are never even
# read from disk (the loader skips any key the SDOps maps to None).
self.stage = DiffusionStage(
self.stage = DiffusionStage.from_checkpoint(
checkpoint_path=checkpoint_path,
dtype=self.dtype,
device=self.device,
@@ -83,12 +85,14 @@ class T2AOneStagePipeline:
offload_mode=offload_mode,
model_configurator=LTXAudioOnlyModelConfigurator,
model_sd_ops=LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP,
alloc_trim_strategy=alloc_trim_strategy,
)
self.audio_decoder = AudioDecoder(
checkpoint_path=checkpoint_path,
dtype=self.dtype,
device=self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
def __call__(
@@ -153,8 +157,7 @@ class T2AOneStagePipeline:
@torch.inference_mode()
def main() -> None:
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path()
params = detect_params(checkpoint_path)
params = resolve_cli_params()
parser = default_1_stage_t2a_arg_parser(params=params)
args = parser.parse_args()
pipeline = T2AOneStagePipeline(
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py
+10 -5
View File
@@ -21,10 +21,11 @@ from ltx_pipelines.utils import (
combined_image_conditionings,
get_device,
)
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.args import (
ImageConditioningInput,
default_1_stage_arg_parser,
detect_checkpoint_path,
resolve_cli_params,
)
from ltx_pipelines.utils.blocks import (
AudioDecoder,
@@ -33,7 +34,6 @@ from ltx_pipelines.utils.blocks import (
PromptEncoder,
VideoDecoder,
)
from ltx_pipelines.utils.constants import detect_params
from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
@@ -58,6 +58,7 @@ class TI2VidOneStagePipeline:
registry: Registry | None = None,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
):
self.dtype = torch.bfloat16
self.device = device or get_device()
@@ -69,14 +70,16 @@ class TI2VidOneStagePipeline:
device=self.device,
registry=registry,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.image_conditioner = ImageConditioner(
checkpoint_path=checkpoint_path,
dtype=self.dtype,
device=self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.stage = DiffusionStage(
self.stage = DiffusionStage.from_checkpoint(
checkpoint_path=checkpoint_path,
dtype=self.dtype,
device=self.device,
@@ -85,18 +88,21 @@ class TI2VidOneStagePipeline:
registry=registry,
compilation_config=compilation_config,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.video_decoder = VideoDecoder(
checkpoint_path=checkpoint_path,
dtype=self.dtype,
device=self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.audio_decoder = AudioDecoder(
checkpoint_path=checkpoint_path,
dtype=self.dtype,
device=self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
def __call__( # noqa: PLR0913
@@ -187,8 +193,7 @@ class TI2VidOneStagePipeline:
@torch.inference_mode()
def main() -> None:
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path()
params = detect_params(checkpoint_path)
params = resolve_cli_params()
parser = default_1_stage_arg_parser(params=params)
args = parser.parse_args()
pipeline = TI2VidOneStagePipeline(
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py
+34 -13
View File
@@ -16,10 +16,11 @@ from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, VideoPixelShape
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.args import (
ImageConditioningInput,
default_2_stage_arg_parser,
detect_checkpoint_path,
resolve_cli_params,
)
from ltx_pipelines.utils.blocks import (
AudioDecoder,
@@ -31,7 +32,6 @@ from ltx_pipelines.utils.blocks import (
)
from ltx_pipelines.utils.constants import (
STAGE_2_DISTILLED_SIGMAS,
detect_params,
)
from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser, SimpleDenoiser
from ltx_pipelines.utils.helpers import (
@@ -52,7 +52,7 @@ class TI2VidTwoStagesPipeline:
images parameter.
"""
def __init__(
def __init__( # noqa: PLR0913
self,
checkpoint_path: str,
distilled_lora: list[LoraPathStrengthAndSDOps],
@@ -64,22 +64,40 @@ class TI2VidTwoStagesPipeline:
registry: Registry | None = None,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
):
self.device = device or get_device()
self.dtype = torch.bfloat16
self._scheduler = LTX2Scheduler()
self.prompt_encoder = PromptEncoder(
checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
checkpoint_path,
gemma_root,
self.dtype,
self.device,
registry=registry,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.image_conditioner = ImageConditioner(
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
)
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
self.upsampler = VideoUpsampler(
checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
checkpoint_path,
spatial_upsampler_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.video_decoder = VideoDecoder(
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
)
self.audio_decoder = AudioDecoder(
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
)
self.video_decoder = VideoDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
self.audio_decoder = AudioDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
self.stage_1 = DiffusionStage(
self.stage_1 = DiffusionStage.from_checkpoint(
checkpoint_path,
self.dtype,
self.device,
@@ -88,8 +106,9 @@ class TI2VidTwoStagesPipeline:
registry=registry,
compilation_config=compilation_config,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.stage_2 = DiffusionStage(
self.stage_2 = DiffusionStage.from_checkpoint(
checkpoint_path,
self.dtype,
self.device,
@@ -98,6 +117,7 @@ class TI2VidTwoStagesPipeline:
registry=registry,
compilation_config=compilation_config,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
def __call__( # noqa: PLR0913
@@ -196,7 +216,9 @@ class TI2VidTwoStagesPipeline:
)
)
video_state, audio_state = self.stage_2(
# Stage 2 refines video only; discard its audio. On the multi-GPU path stage-2 audio
# runs under partial tiled/TDP video context, so the full-context stage-1 audio is kept.
video_state, _ = self.stage_2(
denoiser=SimpleDenoiser(v_context=v_context_p, a_context=a_context_p),
sigmas=stage_2_sigmas,
noiser=noiser,
@@ -225,8 +247,7 @@ class TI2VidTwoStagesPipeline:
@torch.inference_mode()
def main() -> None:
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path()
params = detect_params(checkpoint_path)
params = resolve_cli_params()
parser = default_2_stage_arg_parser(params=params)
args = parser.parse_args()
pipeline = TI2VidTwoStagesPipeline(
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_hq.py
+31 -8
View File
@@ -13,6 +13,7 @@ from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, VideoLatentShape, VideoPixelShape
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.args import ImageConditioningInput, hq_2_stage_arg_parser
from ltx_pipelines.utils.blocks import (
AudioDecoder,
@@ -63,6 +64,7 @@ class TI2VidTwoStagesHQPipeline:
registry: Registry | None = None,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
):
self.device = device or get_device()
self.dtype = torch.bfloat16
@@ -80,16 +82,33 @@ class TI2VidTwoStagesHQPipeline:
)
self.prompt_encoder = PromptEncoder(
checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
checkpoint_path,
gemma_root,
self.dtype,
self.device,
registry=registry,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.image_conditioner = ImageConditioner(
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
)
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
self.upsampler = VideoUpsampler(
checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
checkpoint_path,
spatial_upsampler_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.video_decoder = VideoDecoder(
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
)
self.audio_decoder = AudioDecoder(
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
)
self.video_decoder = VideoDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
self.audio_decoder = AudioDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
self.stage_1 = DiffusionStage(
self.stage_1 = DiffusionStage.from_checkpoint(
checkpoint_path,
self.dtype,
self.device,
@@ -98,8 +117,9 @@ class TI2VidTwoStagesHQPipeline:
registry=registry,
compilation_config=compilation_config,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.stage_2 = DiffusionStage(
self.stage_2 = DiffusionStage.from_checkpoint(
checkpoint_path,
self.dtype,
self.device,
@@ -108,6 +128,7 @@ class TI2VidTwoStagesHQPipeline:
registry=registry,
compilation_config=compilation_config,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
@torch.inference_mode()
@@ -213,7 +234,9 @@ class TI2VidTwoStagesHQPipeline:
)
)
video_state, audio_state = self.stage_2(
# Stage 2 refines video only; discard its audio. On the multi-GPU path stage-2 audio
# runs under partial tiled/TDP video context, so the full-context stage-1 audio is kept.
video_state, _ = self.stage_2(
denoiser=SimpleDenoiser(v_context=v_context_p, a_context=a_context_p),
sigmas=stage_2_sigmas,
noiser=noiser,
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_hq_mgpu.py
+241
View File
@@ -0,0 +1,241 @@
"""Multi-GPU two-stage HQ text/image-to-video runner.
Runs :class:`TI2VidTwoStagesHQPipeline` across multiple GPUs with:
- **Stage 1** -- sequence parallelism (SP) at half resolution
- **Stage 2** -- tiled data parallelism (TDP) on height + width with overlap,
at full resolution
- **Gemma** -- Accelerate-based parallelization
- **VAE** -- distributed decoding
The HQ pipeline applies the distilled LoRA in both stages with separate
strengths and uses the res_2s second-order sampler.
Requires ``ltx-kernels`` to be installed (transitive via SP builder).
"""
from __future__ import annotations
import logging
from collections.abc import Iterator
from multiprocessing import SimpleQueue
import torch
import torch.distributed as dist
from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
from ltx_core.loader.registry import StateDictRegistry
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import get_video_chunks_number
from ltx_core.model.video_vae.tiling import TilingConfig
from ltx_core.multigpu.transformer.attention import AttentionManager
from ltx_core.quantization.fp8_cast import build_policy as _build_fp8_cast_policy
from ltx_core.tiling import DimensionTilingConfig, TileCountConfig, balanced_tile_split
from ltx_pipelines.multigpu.controller import MGPUController
from ltx_pipelines.multigpu.gemma_builders import AccelerateGemmaBuilder
from ltx_pipelines.multigpu.runner import MGPURunner
from ltx_pipelines.multigpu.sp_builder import SequenceParallelBuilder
from ltx_pipelines.multigpu.tdp_builder import TiledDataParallelBuilder
from ltx_pipelines.multigpu.vae_builders import DistributedDecoderBuilder
from ltx_pipelines.multigpu.weight_tracker import TransformerWeightTracker
from ltx_pipelines.ti2vid_two_stages_hq import TI2VidTwoStagesHQPipeline
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.constants import TDP_DISTILLED_SIGMAS
from ltx_pipelines.utils.media_io import encode_video
logger = logging.getLogger(__name__)
# Stage 1 at half-res, 121 frames = 6144 video tokens + audio tokens.
_DEFAULT_SP_MAX_TOKENS = 32768
# Rank that collects distributed-VAE tiles and encodes the assembled video.
_DRIVER_RANK = 0
class TI2VidTwoStagesHQRunner(MGPURunner):
"""Distributed HQ pipeline: SP stage 1 (half-res) + TDP stage 2 (full-res) + Gemma + distributed VAE."""
@torch.inference_mode()
def setup(
self,
*,
checkpoint_path: str,
gemma_root: str,
spatial_upsampler_path: str,
vae_queue: SimpleQueue,
distilled_lora_path: str,
distilled_lora_strength_stage_1: float,
distilled_lora_strength_stage_2: float,
compilation_config: CompilationConfig | None = None,
sp_max_tokens: int = _DEFAULT_SP_MAX_TOKENS,
) -> None:
distilled_lora = [LoraPathStrengthAndSDOps(distilled_lora_path, 1.0, LTXV_LORA_COMFY_RENAMING_MAP)]
registry = StateDictRegistry()
pipeline = TI2VidTwoStagesHQPipeline(
checkpoint_path=checkpoint_path,
distilled_lora=distilled_lora,
distilled_lora_strength_stage_1=distilled_lora_strength_stage_1,
distilled_lora_strength_stage_2=distilled_lora_strength_stage_2,
spatial_upsampler_path=spatial_upsampler_path,
gemma_root=gemma_root,
loras=(),
registry=registry,
quantization=_build_fp8_cast_policy(checkpoint_path),
compilation_config=compilation_config,
alloc_trim_strategy=AllocatorTrimStrategy.DEFER,
)
tracker = TransformerWeightTracker(group=self.groups.transformer_group)
# Stage 1: sequence parallelism.
model_cfg = pipeline.stage_1._transformer_builder.model_config().get("transformer", {})
attn_mgr = AttentionManager(
max_tokens=sp_max_tokens,
num_heads=model_cfg["num_attention_heads"],
head_dim=model_cfg["attention_head_dim"],
tensor_dtype=pipeline.dtype,
group=self.groups.transformer_group,
)
pipeline.stage_1._transformer_builder = SequenceParallelBuilder(
inner=pipeline.stage_1._transformer_builder,
attn_mgr=attn_mgr,
registry=registry,
tracker=tracker,
)
# Stage 2: tiled data parallelism -- balanced 2D spatial grid over the group (one tile/rank).
# height takes the smaller factor of world_size, width the larger; size-aware split is a follow-up.
tdp_height_tiles, tdp_width_tiles = balanced_tile_split(dist.get_world_size(self.groups.transformer_group))
tdp_tiling = TileCountConfig(
height=DimensionTilingConfig(num_tiles=tdp_height_tiles, overlap=5),
width=DimensionTilingConfig(num_tiles=tdp_width_tiles, overlap=5),
)
pipeline.stage_2._transformer_builder = TiledDataParallelBuilder(
inner=pipeline.stage_2._transformer_builder,
group=self.groups.transformer_group,
tiling=tdp_tiling,
registry=registry,
tracker=tracker,
)
# Accelerate Gemma parallelization.
pipeline.prompt_encoder._text_encoder_builder = AccelerateGemmaBuilder(
gemma_root_path=gemma_root,
gemma_group=self.groups.gemma_group,
broadcast_group=self.groups.transformer_group,
registry=registry,
src_rank=_DRIVER_RANK,
dtype=pipeline.dtype,
)
# Distributed VAE decoding: balanced 2D spatial grid over the group (one tile/rank).
# height takes the smaller factor of world_size, width the larger; size-aware split is a follow-up.
vae_height_tiles, vae_width_tiles = balanced_tile_split(dist.get_world_size(self.groups.vae_group))
vae_tiling = TileCountConfig(
height=DimensionTilingConfig(num_tiles=vae_height_tiles, overlap=4),
width=DimensionTilingConfig(num_tiles=vae_width_tiles, overlap=4),
)
pipeline.video_decoder._decoder_builder = DistributedDecoderBuilder( # type: ignore[assignment]
inner=pipeline.video_decoder._decoder_builder,
queue=vae_queue,
vae_group=self.groups.vae_group,
vae_tiling=vae_tiling,
driver_rank=_DRIVER_RANK,
registry=registry,
)
self._pipeline = pipeline
@torch.inference_mode()
def __call__( # noqa: PLR0913
self,
*,
output_path: str,
prompt: str,
negative_prompt: str,
seed: int,
height: int,
width: int,
num_frames: int,
frame_rate: float,
num_inference_steps: int,
video_guider_params: MultiModalGuiderParams,
audio_guider_params: MultiModalGuiderParams,
images: list | None = None,
) -> Iterator[str | None]:
# The pipeline raises ValueError on invalid input (symmetric across ranks); the controller
# catches that and turns it into a recoverable RunnerError. Anything else is fatal.
video, audio = self._pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
seed=seed,
height=height,
width=width,
num_frames=num_frames,
frame_rate=frame_rate,
num_inference_steps=num_inference_steps,
video_guider_params=video_guider_params,
audio_guider_params=audio_guider_params,
images=images or [],
tiling_config=None,
stage_2_sigmas=TDP_DISTILLED_SIGMAS,
)
if dist.get_rank() != _DRIVER_RANK:
yield None # workers: nothing to encode
return
encode_video(
video=video,
fps=frame_rate,
audio=audio,
output_path=output_path,
video_chunks_number=get_video_chunks_number(num_frames, TilingConfig.default()),
)
yield output_path
if __name__ == "__main__":
from ltx_pipelines.utils.args import hq_2_stage_arg_parser
from ltx_pipelines.utils.constants import LTX_2_3_HQ_PARAMS
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
args = hq_2_stage_arg_parser(params=LTX_2_3_HQ_PARAMS).parse_args()
vae_queue = torch.multiprocessing.get_context("spawn").SimpleQueue()
controller = MGPUController(TI2VidTwoStagesHQRunner)
controller.start(
checkpoint_path=args.checkpoint_path,
gemma_root=args.gemma_root,
spatial_upsampler_path=args.spatial_upsampler_path,
vae_queue=vae_queue,
distilled_lora_path=args.distilled_lora[0].path,
distilled_lora_strength_stage_1=args.distilled_lora_strength_stage_1,
distilled_lora_strength_stage_2=args.distilled_lora_strength_stage_2,
compilation_config=args.compile,
)
try:
for _ in controller.stream(
output_path=args.output_path,
prompt=args.prompt,
negative_prompt=args.negative_prompt,
seed=args.seed,
height=args.height,
width=args.width,
num_frames=args.num_frames,
frame_rate=args.frame_rate,
num_inference_steps=args.num_inference_steps,
video_guider_params=MultiModalGuiderParams(
cfg_scale=args.video_cfg_guidance_scale,
stg_scale=args.video_stg_guidance_scale,
rescale_scale=args.video_rescale_scale,
modality_scale=args.a2v_guidance_scale,
skip_step=args.video_skip_step,
stg_blocks=args.video_stg_blocks,
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=args.audio_cfg_guidance_scale,
stg_scale=args.audio_stg_guidance_scale,
rescale_scale=args.audio_rescale_scale,
modality_scale=args.v2a_guidance_scale,
skip_step=args.audio_skip_step,
stg_blocks=args.audio_stg_blocks,
),
images=args.images,
):
pass # drive the job to completion; the runner writes the file as a side effect
finally:
controller.shutdown()
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_mgpu.py
+239
View File
@@ -0,0 +1,239 @@
"""Multi-GPU two-stage text/image-to-video runner.
Runs :class:`TI2VidTwoStagesPipeline` across multiple GPUs with:
- **Stage 1** -- sequence parallelism (SP)
- **Stage 2** -- tiled data parallelism (TDP) on height + width with overlap
- **Gemma** -- Accelerate-based parallelization
- **VAE** -- distributed decoding
Requires ``ltx-kernels`` to be installed (transitive via SP builder).
"""
from __future__ import annotations
import logging
from collections.abc import Callable, Iterator
from multiprocessing import SimpleQueue
import torch
import torch.distributed as dist
from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
from ltx_core.loader.registry import StateDictRegistry
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import get_video_chunks_number
from ltx_core.model.video_vae.tiling import TilingConfig
from ltx_core.multigpu.transformer.attention import AttentionManager
from ltx_core.quantization import QuantizationPolicy
from ltx_core.quantization.fp8_cast import build_policy as _build_fp8_cast_policy
from ltx_core.tiling import DimensionTilingConfig, TileCountConfig, balanced_tile_split
from ltx_pipelines.multigpu.controller import MGPUController
from ltx_pipelines.multigpu.gemma_builders import AccelerateGemmaBuilder
from ltx_pipelines.multigpu.runner import MGPURunner
from ltx_pipelines.multigpu.sp_builder import SequenceParallelBuilder
from ltx_pipelines.multigpu.tdp_builder import TiledDataParallelBuilder
from ltx_pipelines.multigpu.vae_builders import DistributedDecoderBuilder
from ltx_pipelines.multigpu.weight_tracker import TransformerWeightTracker
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.constants import TDP_DISTILLED_SIGMAS
from ltx_pipelines.utils.media_io import encode_video
logger = logging.getLogger(__name__)
# Stage 1 at 512x768, 121 frames = 6144 video tokens + audio tokens.
_DEFAULT_SP_MAX_TOKENS = 32768
# Rank that collects distributed-VAE tiles and encodes the assembled video.
_DRIVER_RANK = 0
class TI2VidTwoStagesRunner(MGPURunner):
"""Distributed :class:`TI2VidTwoStagesPipeline`: SP stage 1 + TDP stage 2 + Gemma + distributed VAE."""
@torch.inference_mode()
def setup(
self,
*,
checkpoint_path: str,
gemma_root: str,
spatial_upsampler_path: str,
vae_queue: SimpleQueue,
distilled_lora_path: str,
compilation_config: CompilationConfig | None = None,
sp_max_tokens: int = _DEFAULT_SP_MAX_TOKENS,
quantization: Callable[[], QuantizationPolicy] | None = None,
) -> None:
# quantization is a picklable zero-arg builder (built per worker, post-spawn); default fp8-cast.
quantization_policy = quantization() if quantization is not None else _build_fp8_cast_policy(checkpoint_path)
distilled_lora = [LoraPathStrengthAndSDOps(distilled_lora_path, 1.0, LTXV_LORA_COMFY_RENAMING_MAP)]
registry = StateDictRegistry()
pipeline = TI2VidTwoStagesPipeline(
checkpoint_path=checkpoint_path,
distilled_lora=distilled_lora,
spatial_upsampler_path=spatial_upsampler_path,
gemma_root=gemma_root,
loras=[],
registry=registry,
quantization=quantization_policy,
compilation_config=compilation_config,
alloc_trim_strategy=AllocatorTrimStrategy.DEFER,
)
tracker = TransformerWeightTracker(group=self.groups.transformer_group)
# Stage 1: sequence parallelism.
model_cfg = pipeline.stage_1._transformer_builder.model_config().get("transformer", {})
attn_mgr = AttentionManager(
max_tokens=sp_max_tokens,
num_heads=model_cfg["num_attention_heads"],
head_dim=model_cfg["attention_head_dim"],
tensor_dtype=pipeline.dtype,
group=self.groups.transformer_group,
)
pipeline.stage_1._transformer_builder = SequenceParallelBuilder(
inner=pipeline.stage_1._transformer_builder,
attn_mgr=attn_mgr,
registry=registry,
tracker=tracker,
)
# Stage 2: tiled data parallelism -- balanced 2D spatial grid over the group (one tile/rank).
# height takes the smaller factor of world_size, width the larger; size-aware split is a follow-up.
tdp_height_tiles, tdp_width_tiles = balanced_tile_split(dist.get_world_size(self.groups.transformer_group))
tdp_tiling = TileCountConfig(
height=DimensionTilingConfig(num_tiles=tdp_height_tiles, overlap=5),
width=DimensionTilingConfig(num_tiles=tdp_width_tiles, overlap=5),
)
pipeline.stage_2._transformer_builder = TiledDataParallelBuilder(
inner=pipeline.stage_2._transformer_builder,
group=self.groups.transformer_group,
tiling=tdp_tiling,
registry=registry,
tracker=tracker,
)
# Accelerate Gemma parallelization.
pipeline.prompt_encoder._text_encoder_builder = AccelerateGemmaBuilder(
gemma_root_path=gemma_root,
gemma_group=self.groups.gemma_group,
broadcast_group=self.groups.transformer_group,
registry=registry,
src_rank=_DRIVER_RANK,
dtype=pipeline.dtype,
)
# Distributed VAE decoding: balanced 2D spatial grid over the group (one tile/rank).
# height takes the smaller factor of world_size, width the larger; size-aware split is a follow-up.
vae_height_tiles, vae_width_tiles = balanced_tile_split(dist.get_world_size(self.groups.vae_group))
vae_tiling = TileCountConfig(
height=DimensionTilingConfig(num_tiles=vae_height_tiles, overlap=4),
width=DimensionTilingConfig(num_tiles=vae_width_tiles, overlap=4),
)
pipeline.video_decoder._decoder_builder = DistributedDecoderBuilder(
inner=pipeline.video_decoder._decoder_builder,
queue=vae_queue,
vae_group=self.groups.vae_group,
vae_tiling=vae_tiling,
driver_rank=_DRIVER_RANK,
registry=registry,
)
self._pipeline = pipeline
@torch.inference_mode()
def __call__( # noqa: PLR0913
self,
*,
output_path: str,
prompt: str,
negative_prompt: str,
seed: int,
height: int,
width: int,
num_frames: int,
frame_rate: float,
num_inference_steps: int,
video_guider_params: MultiModalGuiderParams,
audio_guider_params: MultiModalGuiderParams,
images: list | None = None,
) -> Iterator[str | None]:
# The pipeline raises ValueError on invalid input (symmetric across ranks); the controller
# catches that and turns it into a recoverable RunnerError. Anything else is fatal.
video, audio = self._pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
seed=seed,
height=height,
width=width,
num_frames=num_frames,
frame_rate=frame_rate,
num_inference_steps=num_inference_steps,
video_guider_params=video_guider_params,
audio_guider_params=audio_guider_params,
images=images or [],
tiling_config=None,
stage_2_sigmas=TDP_DISTILLED_SIGMAS,
)
if dist.get_rank() != _DRIVER_RANK:
yield None # workers: nothing to encode
return
encode_video(
video=video,
fps=frame_rate,
audio=audio,
output_path=output_path,
video_chunks_number=get_video_chunks_number(num_frames, TilingConfig.default()),
)
yield output_path
if __name__ == "__main__":
from ltx_pipelines.utils.args import (
default_2_stage_arg_parser,
resolve_cli_params,
)
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
params = resolve_cli_params()
args = default_2_stage_arg_parser(params=params).parse_args()
vae_queue = torch.multiprocessing.get_context("spawn").SimpleQueue()
controller = MGPUController(TI2VidTwoStagesRunner)
controller.start(
checkpoint_path=args.checkpoint_path,
gemma_root=args.gemma_root,
spatial_upsampler_path=args.spatial_upsampler_path,
vae_queue=vae_queue,
distilled_lora_path=args.distilled_lora[0].path,
compilation_config=args.compile,
)
try:
for _ in controller.stream(
output_path=args.output_path,
prompt=args.prompt,
negative_prompt=args.negative_prompt,
seed=args.seed,
height=args.height,
width=args.width,
num_frames=args.num_frames,
frame_rate=args.frame_rate,
num_inference_steps=args.num_inference_steps,
video_guider_params=MultiModalGuiderParams(
cfg_scale=args.video_cfg_guidance_scale,
stg_scale=args.video_stg_guidance_scale,
rescale_scale=args.video_rescale_scale,
modality_scale=args.a2v_guidance_scale,
skip_step=args.video_skip_step,
stg_blocks=args.video_stg_blocks,
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=args.audio_cfg_guidance_scale,
stg_scale=args.audio_stg_guidance_scale,
rescale_scale=args.audio_rescale_scale,
modality_scale=args.v2a_guidance_scale,
skip_step=args.audio_skip_step,
stg_blocks=args.audio_stg_blocks,
),
images=args.images,
):
pass # drive the job to completion; the runner writes the file as a side effect
finally:
controller.shutdown()
packages/ltx-pipelines/src/ltx_pipelines/utils/allocator_trim_strategy.py
+8
View File
@@ -0,0 +1,8 @@
from enum import Enum
class AllocatorTrimStrategy(Enum):
"""How a block releases its model's memory when its scope exits."""
TRIM = "trim" # sync, release storage (to meta), and empty_cache() back to the OS
DEFER = "defer" # skip teardown; let GC reclaim it and keep the CUDA cache warm
packages/ltx-pipelines/src/ltx_pipelines/utils/args.py
+20
View File
@@ -1,5 +1,6 @@
import argparse
import json
import sys
from collections.abc import Sequence
from pathlib import Path
from typing import Any, NamedTuple
@@ -14,6 +15,7 @@ from ltx_pipelines.utils.constants import (
LTX_2_3_HQ_PARAMS,
LTX_2_3_PARAMS,
PipelineParams,
detect_params,
)
from ltx_pipelines.utils.quantization_factory import QuantizationKind
from ltx_pipelines.utils.types import OffloadMode
@@ -263,6 +265,24 @@ def detect_checkpoint_path(distilled: bool = False) -> str:
return known.distilled_checkpoint_path if distilled else known.checkpoint_path
def help_requested() -> bool:
"""Whether ``-h``/``--help`` appears on the command line."""
return "-h" in sys.argv or "--help" in sys.argv
def resolve_cli_params(distilled: bool = False) -> PipelineParams:
"""Return the model params a pipeline CLI uses to build its argument parser.
Reads the model version from the checkpoint named on the command line so the
parser's defaults match the target model.
Args:
distilled: Whether the pipeline takes a distilled checkpoint
(``--distilled-checkpoint-path``) rather than a full one (``--checkpoint-path``).
"""
if help_requested():
return LTX_2_3_PARAMS
return detect_params(detect_checkpoint_path(distilled=distilled))
def basic_arg_parser(
params: PipelineParams = LTX_2_3_PARAMS,
distilled: bool = False,
packages/ltx-pipelines/src/ltx_pipelines/utils/blocks.py
+224 -93
View File
@@ -1,7 +1,7 @@
"""Pipeline blocks — each block owns its model lifecycle.
Blocks build a model on each ``__call__``, use it, then free GPU memory.
This eliminates manual ``del model; cleanup_memory()`` in pipelines and
removes the need for :class:`ModelLedger`.
This eliminates manual ``del model; cleanup_memory()`` in pipelines: each
block is self-contained, so no central model-coordinator object is needed.
"""
from __future__ import annotations
@@ -77,6 +77,7 @@ from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcesso
from ltx_core.tools import AudioLatentTools, LatentTools, VideoLatentTools
from ltx_core.types import Audio, AudioLatentShape, LatentState, VideoLatentShape, VideoPixelShape
from ltx_core.utils import find_matching_file
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.gpu_model import gpu_model
from ltx_pipelines.utils.helpers import (
cleanup_memory,
@@ -136,23 +137,25 @@ def _apply_compile_ops(
@contextmanager
def _streaming_model(
builder: StreamingModelBuilder,
offload_mode: OffloadMode,
target_device: torch.device,
dtype: torch.dtype,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
) -> Iterator:
"""Build a streaming wrapper, yield it, then tear down and free memory."""
cpu_slots_count = DISK_CPU_SLOTS if offload_mode == OffloadMode.DISK else None
wrapped = builder.build(
device=target_device,
dtype=dtype,
cpu_slots_count=cpu_slots_count,
)
"""Build a streaming wrapper, yield it, then tear down and free memory.
The builder's own ``cpu_slots_count`` selects RAM vs disk streaming.
``teardown()`` always runs -- it releases non-memory resources (forward
hooks, the disk I/O worker thread, open file handles) that GC would not
reclaim promptly. ``alloc_trim_strategy=DEFER`` only skips the eager allocator
reclaim (``to("meta")`` + ``cleanup_memory()``), leaving param storage for GC.
"""
wrapped = builder.build(device=target_device, dtype=dtype)
try:
yield wrapped
finally:
wrapped.teardown()
wrapped.to("meta")
cleanup_memory()
if alloc_trim_strategy == AllocatorTrimStrategy.TRIM:
wrapped.to("meta")
cleanup_memory()
def _build_state(
@@ -177,9 +180,13 @@ def _build_state(
return state
def _cleanup_iter(it: Iterator[torch.Tensor], model: torch.nn.Module) -> Iterator[torch.Tensor]:
"""Wrap an iterator to clean up *model* memory once it is exhausted or abandoned."""
with gpu_model(model):
def _cleanup_iter(
it: Iterator[torch.Tensor],
model: torch.nn.Module,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
) -> Iterator[torch.Tensor]:
"""Wrap an iterator to release *model* memory (per ``alloc_trim_strategy``) once exhausted or abandoned."""
with gpu_model(model, alloc_trim_strategy=alloc_trim_strategy):
yield from it
@@ -190,12 +197,41 @@ def _cleanup_iter(it: Iterator[torch.Tensor], model: torch.nn.Module) -> Iterato
class DiffusionStage:
"""Owns transformer lifecycle. Builds on each call, frees on exit.
Replaces the manual ``model_ledger.transformer()`` / ``del transformer``
pattern in every pipeline.
Replaces the manual build-transformer / ``del transformer`` pattern that
every pipeline previously repeated.
"""
def __init__( # noqa: PLR0913
def __init__(
self,
transformer_builder: ModelBuilderProtocol[LTXModelProtocol],
dtype: torch.dtype,
device: torch.device,
*,
quantization: QuantizationPolicy | None = None,
compilation_config: CompilationConfig | None = None,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
) -> None:
"""Construct a stage from a single pre-built transformer ``builder``.
Holds only that builder plus build-time configuration (dtype, device,
quantization, compilation). Turning a checkpoint path + LoRA set into a
builder -- and choosing a :class:`StreamingModelBuilder` when offloading --
lives in :meth:`from_checkpoint`, which is how pipelines normally create a
stage. A :class:`StreamingModelBuilder` selects the block-streaming build
path; any other builder uses the standard (all-on-GPU) path.
``quantization`` and ``compilation_config`` are applied lazily on the
standard path; on the streaming path they are already baked into the
streaming builder by :meth:`from_checkpoint` and these fields are unused.
"""
self._transformer_builder = transformer_builder
self._dtype = dtype
self._device = device
self._quantization = quantization
self._compilation_config = compilation_config
self._alloc_trim_strategy = alloc_trim_strategy
@classmethod
def from_checkpoint( # noqa: PLR0913
cls,
checkpoint_path: str,
dtype: torch.dtype,
device: torch.device,
@@ -203,17 +239,24 @@ class DiffusionStage:
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
compilation_config: CompilationConfig | None = None,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
offload_mode: OffloadMode = OffloadMode.NONE,
transformer_builder: ModelBuilderProtocol[LTXModelProtocol] | None = None,
model_configurator: type[ModelConfigurator] = LTXModelConfigurator,
model_sd_ops: SDOps = LTXV_MODEL_COMFY_RENAMING_MAP,
) -> None:
self._checkpoint_path = checkpoint_path
self._dtype = dtype
self._device = device
self._quantization = quantization
self._compilation_config = compilation_config
self._offload_mode = offload_mode
) -> "DiffusionStage":
"""Build a stage from a checkpoint path and LoRA set.
Constructs a single transformer builder from ``checkpoint_path`` +
``loras`` + ``quantization`` and delegates to ``__init__``. When
``offload_mode != OffloadMode.NONE`` that builder is a
:class:`StreamingModelBuilder` (with quantization/compilation baked in and
its ``cpu_slots_count`` set for the requested mode); otherwise it is the
standard single-GPU builder. This is the high-level entry point used by
pipelines; ``__init__`` itself takes an already-built builder.
``model_configurator`` / ``model_sd_ops`` let callers (e.g. the audio-only
T2A pipeline) override the model class configurator and the state-dict key
mapping. A quantization policy that pins its own configurator takes
precedence over ``model_configurator``.
"""
# A quantization policy may pin its own configurator; otherwise use the one
# provided by the caller (defaults to the audio-video LTXModelConfigurator).
configurator = (
@@ -221,51 +264,75 @@ class DiffusionStage:
if quantization is not None and quantization.model_configurator is not None
else model_configurator
)
if transformer_builder is not None:
self._transformer_builder = transformer_builder
else:
self._transformer_builder = Builder(
transformer_builder: ModelBuilderProtocol[LTXModelProtocol]
if offload_mode == OffloadMode.NONE:
transformer_builder = Builder(
model_path=checkpoint_path,
model_class_configurator=configurator,
model_sd_ops=model_sd_ops,
loras=tuple(loras),
registry=registry or DummyRegistry(),
)
if offload_mode != OffloadMode.NONE:
# WeightsProvider currently only supports plain bf16 + fp8_cast LoRA fusion
# (no companion-key emission). Quantization policies that emit
# companion keys (e.g. ``.weight_scale``) cannot be streamed yet.
if quantization is not None and quantization.fuse_rule is not fp8_cast_fuse_rule:
raise ValueError(
"Block streaming is not supported with this quantization policy "
"(only bf16 and fp8_cast are currently supported)."
)
streaming_sd_ops: SDOps = model_sd_ops
streaming_module_ops: tuple[ModuleOps, ...] = ()
streaming_loras = tuple(loras)
if compilation_config:
number_of_layers = self._transformer_builder.model_config()["transformer"]["num_layers"]
streaming_sd_ops, streaming_module_ops, streaming_loras = _apply_compile_ops(
streaming_sd_ops, streaming_module_ops, streaming_loras, number_of_layers
)
if quantization is not None:
streaming_sd_ops, streaming_module_ops = _chain_quantization(
streaming_sd_ops, streaming_module_ops, quantization
)
self._streaming_builder = StreamingModelBuilder(
model_class_configurator=configurator,
model_path=checkpoint_path,
model_sd_ops=streaming_sd_ops,
module_ops=streaming_module_ops,
loras=streaming_loras,
else:
transformer_builder = cls._build_streaming_builder(
checkpoint_path=checkpoint_path,
configurator=configurator,
model_sd_ops=model_sd_ops,
loras=tuple(loras),
quantization=quantization,
registry=registry or DummyRegistry(),
fuse_rule=quantization.fuse_rule if quantization is not None else bf16_fuse_rule,
blocks_attr="transformer_blocks",
blocks_prefix="transformer_blocks",
offload_mode=offload_mode,
)
return cls(
transformer_builder,
dtype,
device,
quantization=quantization,
compilation_config=compilation_config,
alloc_trim_strategy=alloc_trim_strategy,
)
@staticmethod
def _build_streaming_builder(
*,
checkpoint_path: str,
configurator: type[ModelConfigurator],
model_sd_ops: SDOps,
loras: tuple[LoraPathStrengthAndSDOps, ...],
quantization: QuantizationPolicy | None,
registry: Registry,
offload_mode: OffloadMode,
) -> StreamingModelBuilder:
"""Construct the streaming transformer builder for an offloading stage.
Holds only raw config (``model_sd_ops`` / ``loras``); compilation and
quantization are applied at build time by :meth:`_prepared_builder`, exactly
as on the standard path -- so the builder's LoRA set stays raw and
:meth:`with_loras` swaps it consistently. ``cpu_slots_count`` is pinned for
the requested ``offload_mode`` (disk streaming uses a small slot count;
CPU/RAM streaming pins every block).
"""
# WeightsProvider currently only supports plain bf16 + fp8_cast LoRA fusion
# (no companion-key emission). Quantization policies that emit
# companion keys (e.g. ``.weight_scale``) cannot be streamed yet.
if quantization is not None and quantization.fuse_rule is not fp8_cast_fuse_rule:
raise ValueError(
"Block streaming is not supported with this quantization policy "
"(only bf16 and fp8_cast are currently supported)."
)
return StreamingModelBuilder(
model_class_configurator=configurator,
model_path=checkpoint_path,
model_sd_ops=model_sd_ops,
loras=loras,
registry=registry,
fuse_rule=quantization.fuse_rule if quantization is not None else bf16_fuse_rule,
blocks_attr="transformer_blocks",
blocks_prefix="transformer_blocks",
cpu_slots_count=DISK_CPU_SLOTS if offload_mode == OffloadMode.DISK else None,
)
def with_attention(self, attention: AttentionFunction | AttentionCallable | None) -> "DiffusionStage":
"""Return a new ``DiffusionStage`` that pins the transformer build to ``attention``.
Functional: never mutates ``self``. The returned stage shares all other
@@ -280,41 +347,64 @@ class DiffusionStage:
new._transformer_builder = self._transformer_builder.with_module_ops(
(*self._transformer_builder.module_ops, op),
)
if self._offload_mode != OffloadMode.NONE:
new._streaming_builder = self._streaming_builder.with_module_ops(
(*self._streaming_builder.module_ops, op),
)
return new
def _build_transformer(self, *, device: torch.device | None = None, **kwargs: object) -> X0Model:
target = device or self._device
sd_ops = self._transformer_builder.model_sd_ops
module_ops = self._transformer_builder.module_ops
loras = self._transformer_builder.loras
def with_builder(self, builder: ModelBuilderProtocol[LTXModelProtocol]) -> "DiffusionStage":
"""Return a new ``DiffusionStage`` that builds its transformer from ``builder``.
Functional: never mutates ``self``; shares all other configuration (dtype, device,
quantization, compilation). Affects the standard (non-offload) build path.
"""
new = copy.copy(self)
new._transformer_builder = builder
return new
def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "DiffusionStage":
"""Return a new ``DiffusionStage`` built with exactly ``loras`` (replacing the current set)."""
return self.with_builder(self._transformer_builder.with_loras(loras))
def _prepared_builder(self) -> ModelBuilderProtocol[LTXModelProtocol]:
"""Return the configured builder with the stage's build-time ops applied.
Compilation and quantization live on the stage (not on the builder) and are
applied here, lazily, for both the standard and streaming paths. This keeps
the builder holding only raw sd_ops/module_ops/LoRAs, so ``with_loras`` /
``with_builder`` swap them consistently regardless of the build path. The
returned copy preserves the builder's concrete type (e.g. a
``StreamingModelBuilder`` stays one).
"""
builder = self._transformer_builder
sd_ops = builder.model_sd_ops
module_ops = builder.module_ops
loras = builder.loras
if self._compilation_config is not None:
number_of_layers = self._transformer_builder.model_config()["transformer"]["num_layers"]
number_of_layers = builder.model_config()["transformer"]["num_layers"]
sd_ops, module_ops, loras = _apply_compile_ops(
sd_ops, module_ops, loras, number_of_layers, self._compilation_config
)
if self._quantization is not None:
sd_ops, module_ops = _chain_quantization(sd_ops, module_ops, self._quantization)
builder = self._transformer_builder.with_module_ops(module_ops).with_sd_ops(sd_ops).with_loras(loras)
if self._quantization is not None:
builder = builder.with_fuse_rule(self._quantization.fuse_rule)
return X0Model(builder.build(device=target, **kwargs)).to(target).eval()
return builder.with_module_ops(module_ops).with_sd_ops(sd_ops).with_loras(loras)
def _build_transformer(self, *, device: torch.device | None = None, **kwargs: object) -> X0Model:
target = device or self._device
return X0Model(self._prepared_builder().build(device=target, **kwargs)).to(target).eval()
@property
def _is_streaming(self) -> bool:
"""Whether the configured builder uses the block-streaming build path."""
return isinstance(self._transformer_builder, StreamingModelBuilder)
@contextmanager
def _streaming_transformer_ctx(self) -> Iterator[X0Model]:
with _streaming_model(
self._streaming_builder, self._offload_mode, self._device, self._dtype
) as streaming_wrapper:
builder = self._prepared_builder()
assert isinstance(builder, StreamingModelBuilder)
with _streaming_model(builder, self._device, self._dtype, self._alloc_trim_strategy) as streaming_wrapper:
yield X0Model(streaming_wrapper).eval()
def _transformer_ctx(self, **kwargs: object) -> AbstractContextManager:
if self._offload_mode != OffloadMode.NONE:
if self._is_streaming:
return self._streaming_transformer_ctx()
return gpu_model(self._build_transformer(**kwargs))
return gpu_model(self._build_transformer(**kwargs), alloc_trim_strategy=self._alloc_trim_strategy)
def model_context(self, **kwargs: object) -> AbstractContextManager:
"""Build the transformer, yield it, then free its memory on exit.
@@ -419,8 +509,8 @@ class DiffusionStage:
v_shape = VideoLatentShape.from_pixel_shape(pixel_shape)
video_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape, fps)
mode = "streaming" if self._offload_mode != OffloadMode.NONE else "standard"
logger.info("Building transformer (%s) from %s", mode, self._checkpoint_path)
mode = "streaming" if self._is_streaming else "standard"
logger.info("Building transformer (%s) from %s", mode, self._transformer_builder.checkpoint)
with self._transformer_ctx(video_tools=video_tools) as transformer:
logger.info(
"Running denoising loop (%d steps, %dx%d %d frames @ %.1f fps)",
@@ -467,12 +557,14 @@ class PromptEncoder:
registry: Registry | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
text_encoder_builder: BuilderProtocol | None = None,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
) -> None:
self._gemma_root = gemma_root
self._checkpoint_path = checkpoint_path
self._dtype = dtype
self._device = device
self._offload_mode = offload_mode
self._alloc_trim_strategy = alloc_trim_strategy
if text_encoder_builder is not None:
if offload_mode != OffloadMode.NONE:
@@ -501,6 +593,7 @@ class PromptEncoder:
registry=registry or DummyRegistry(),
blocks_attr="model.model.language_model.layers",
blocks_prefix="model.model.language_model.layers",
cpu_slots_count=DISK_CPU_SLOTS if offload_mode == OffloadMode.DISK else None,
)
self._embeddings_processor_builder = Builder(
model_path=checkpoint_path,
@@ -519,8 +612,10 @@ class PromptEncoder:
def _text_encoder_ctx(self) -> AbstractContextManager:
if self._offload_mode != OffloadMode.NONE:
return _streaming_model(self._streaming_text_encoder_builder, self._offload_mode, self._device, self._dtype)
return gpu_model(self._build_text_encoder())
return _streaming_model(
self._streaming_text_encoder_builder, self._device, self._dtype, self._alloc_trim_strategy
)
return gpu_model(self._build_text_encoder(), alloc_trim_strategy=self._alloc_trim_strategy)
def __call__(
self,
@@ -541,7 +636,9 @@ class PromptEncoder:
raw_outputs = text_encoder.encode(prompts)
logger.info("Text encoder done, building embeddings processor from %s", self._checkpoint_path)
with gpu_model(self._build_embeddings_processor()) as embeddings_processor:
with gpu_model(
self._build_embeddings_processor(), alloc_trim_strategy=self._alloc_trim_strategy
) as embeddings_processor:
result = [embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs]
logger.info("Prompt encoding complete")
return result
@@ -563,6 +660,7 @@ class ImageConditioner:
dtype: torch.dtype,
device: torch.device,
registry: Registry | None = None,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
) -> None:
self._dtype = dtype
self._device = device
@@ -572,13 +670,14 @@ class ImageConditioner:
model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
registry=registry or DummyRegistry(),
)
self._alloc_trim_strategy = alloc_trim_strategy
def _build_encoder(self) -> VideoEncoder:
return self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()
def __call__(self, fn: Callable[[VideoEncoder], T]) -> T:
"""Build video encoder → call *fn(encoder)* → free encoder."""
with gpu_model(self._build_encoder()) as encoder:
with gpu_model(self._build_encoder(), alloc_trim_strategy=self._alloc_trim_strategy) as encoder:
return fn(encoder)
@@ -597,6 +696,7 @@ class VideoUpsampler:
dtype: torch.dtype,
device: torch.device,
registry: Registry | None = None,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
) -> None:
self._upsampler_path = upsampler_path
self._dtype = dtype
@@ -612,13 +712,20 @@ class VideoUpsampler:
model_class_configurator=LatentUpsamplerConfigurator,
registry=registry or DummyRegistry(),
)
self._alloc_trim_strategy = alloc_trim_strategy
def __call__(self, latent: torch.Tensor) -> torch.Tensor:
"""Upsample *latent* using video encoder + spatial upsampler, then free both."""
logger.info("Building video encoder + spatial upsampler from %s", self._upsampler_path)
with (
gpu_model(self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()) as encoder,
gpu_model(self._upsampler_builder.build(device=self._device, dtype=self._dtype).eval()) as upsampler,
gpu_model(
self._encoder_builder.build(device=self._device, dtype=self._dtype).eval(),
alloc_trim_strategy=self._alloc_trim_strategy,
) as encoder,
gpu_model(
self._upsampler_builder.build(device=self._device, dtype=self._dtype).eval(),
alloc_trim_strategy=self._alloc_trim_strategy,
) as upsampler,
):
return upsample_video(latent=latent, video_encoder=encoder, upsampler=upsampler)
@@ -641,6 +748,7 @@ class VideoDecoder:
registry: Registry | None = None,
memory_efficient: bool = True,
decoder_builder: BuilderProtocol | None = None,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
) -> None:
self._checkpoint_path = checkpoint_path
self._dtype = dtype
@@ -655,6 +763,7 @@ class VideoDecoder:
registry=registry or DummyRegistry(),
module_ops=(MEMORY_EFFICIENT_DECODE,) if memory_efficient else (),
)
self._alloc_trim_strategy = alloc_trim_strategy
def __call__(
self,
@@ -665,7 +774,11 @@ class VideoDecoder:
"""Decode *latent* to pixel-space video chunks. Decoder freed after exhaustion."""
logger.info("Building video decoder from %s", self._checkpoint_path)
decoder = self._decoder_builder.build(device=self._device, dtype=self._dtype).eval()
return _cleanup_iter(decoder.decode_video(latent, tiling_config, generator), decoder)
return _cleanup_iter(
decoder.decode_video(latent, tiling_config, generator),
decoder,
alloc_trim_strategy=self._alloc_trim_strategy,
)
# ---------------------------------------------------------------------------
@@ -682,6 +795,7 @@ class AudioDecoder:
dtype: torch.dtype,
device: torch.device,
registry: Registry | None = None,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
) -> None:
self._checkpoint_path = checkpoint_path
self._dtype = dtype
@@ -698,13 +812,25 @@ class AudioDecoder:
model_sd_ops=VOCODER_COMFY_KEYS_FILTER,
registry=registry or DummyRegistry(),
)
self._alloc_trim_strategy = alloc_trim_strategy
def __call__(self, latent: torch.Tensor) -> Audio:
"""Decode audio *latent* through VAE decoder + vocoder, then free both."""
logger.info("Building audio decoder + vocoder from %s", self._checkpoint_path)
# The vocoder always runs in fp32 (bf16 accumulation degrades spectral
# metrics). On CUDA/CPU it is stored in bf16 and autocast upcasts per-op to
# save memory; MPS has no fp32 autocast, so store it in fp32 directly and
# avoid the per-call cast. Negligible footprint for this small model.
vocoder_dtype = torch.float32 if self._device.type == "mps" else self._dtype
with (
gpu_model(self._decoder_builder.build(device=self._device, dtype=self._dtype).eval()) as decoder,
gpu_model(self._vocoder_builder.build(device=self._device, dtype=self._dtype).eval()) as vocoder,
gpu_model(
self._decoder_builder.build(device=self._device, dtype=self._dtype).eval(),
alloc_trim_strategy=self._alloc_trim_strategy,
) as decoder,
gpu_model(
self._vocoder_builder.build(device=self._device, dtype=vocoder_dtype).eval(),
alloc_trim_strategy=self._alloc_trim_strategy,
) as vocoder,
):
return vae_decode_audio(latent, decoder, vocoder)
@@ -726,9 +852,11 @@ class AudioConditioner:
dtype: torch.dtype,
device: torch.device,
registry: Registry | None = None,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
) -> None:
self._dtype = dtype
self._device = device
self._alloc_trim_strategy = alloc_trim_strategy
self._encoder_builder = Builder(
model_path=checkpoint_path,
model_class_configurator=AudioEncoderConfigurator,
@@ -738,5 +866,8 @@ class AudioConditioner:
def __call__(self, fn: Callable[[torch.nn.Module], T]) -> T:
"""Build audio encoder → call *fn(encoder)* → free encoder."""
with gpu_model(self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()) as encoder:
with gpu_model(
self._encoder_builder.build(device=self._device, dtype=self._dtype).eval(),
alloc_trim_strategy=self._alloc_trim_strategy,
) as encoder:
return fn(encoder)
packages/ltx-pipelines/src/ltx_pipelines/utils/constants.py
+2
View File
@@ -20,6 +20,8 @@ STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0]
DISTILLED_SIGMAS = torch.tensor(DISTILLED_SIGMA_VALUES)
STAGE_2_DISTILLED_SIGMAS = torch.tensor(STAGE_2_DISTILLED_SIGMA_VALUES)
# Stage 2 schedule for the tiled-data-parallel multi-GPU runner.
TDP_DISTILLED_SIGMAS = torch.tensor([0.625, 0.4, 0.0])
# =============================================================================
packages/ltx-pipelines/src/ltx_pipelines/utils/denoisers.py
+15 -8
View File
@@ -164,17 +164,24 @@ def _guided_denoise( # noqa: PLR0913,PLR0915
enabled=not a_skip,
)
# Replicate each pass's PerturbationConfig to all `orig_b` samples it
# carries, so `BatchedPerturbationConfig.mask_like` returns a per-sample
# mask (length n*orig_b) instead of a per-pass mask (length n). Without
# this expansion the mask is broadcast against a (n*orig_b, T, D) tensor
# and the multiplication fails with a batch-dim mismatch whenever
# `orig_b > 1` (e.g. multi-prompt benchmark panels).
# Replicate each pass's PerturbationConfig to all `orig_b` samples it carries, so the keep-mask
# has one row per sample (length n*orig_b) instead of per-pass (length n). Without this
# expansion the mask broadcasts against a (n*orig_b, T, D) tensor and the multiplication fails
# with a batch-dim mismatch whenever `orig_b > 1` (e.g. multi-prompt benchmark panels).
batched_ptb_configs = [ptb for ptb in ptb_configs for _ in range(orig_b)]
all_v, all_a = transformer(
video=batched_video, audio=batched_audio, perturbations=BatchedPerturbationConfig(batched_ptb_configs)
# Build the config with num_blocks/device/dtype so it precomputes its per-block mask tensor
# on init (transformer.num_blocks delegates through the SP/TDP/BatchSplit/X0 wrappers). The
# compiled forward then reads perturbation as a runtime tensor, not by querying the config
# in-graph -- so it doesn't recompile per perturbation config.
ref_modality = batched_video if batched_video is not None else batched_audio
perturbations = BatchedPerturbationConfig(
batched_ptb_configs,
num_blocks=transformer.num_blocks,
device=ref_modality.latent.device,
dtype=ref_modality.latent.dtype,
)
all_v, all_a = transformer(video=batched_video, audio=batched_audio, perturbations=perturbations)
# Split results back and combine via guiders.
splits_v = list(all_v.chunk(n)) if all_v is not None else [0.0] * n
packages/ltx-pipelines/src/ltx_pipelines/utils/gpu_model.py
+15 -10
View File
@@ -4,27 +4,32 @@ from typing import TypeVar
import torch
from ltx_core.devices import synchronize_device
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.helpers import cleanup_memory
_M = TypeVar("_M", bound=torch.nn.Module)
@contextmanager
def gpu_model(model: _M) -> Iterator[_M]:
def gpu_model(model: _M, alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM) -> Iterator[_M]:
"""Context manager that yields a model and releases its memory on exit.
Moves all parameters and buffers to ``meta`` device on exit, which
immediately releases the underlying storage on **both** GPU and CPU,
then runs ``cleanup_memory()`` to reclaim fragmented CUDA memory.
On ``TRIM`` (default): synchronize, move parameters/buffers to the ``meta``
device (releasing GPU+CPU storage), then ``cleanup_memory()`` to return
cached blocks to the OS. ``DEFER`` skips this -- the model's storage is
reclaimed by normal GC and the CUDA caching allocator stays warm for the
next build (cheaper for back-to-back runs).
Usage::
with gpu_model(build_encoder()) as encoder:
... # use encoder — typed as the concrete class
... # use encoder -- typed as the concrete class
# GPU + CPU memory freed automatically
"""
try:
yield model
finally:
torch.cuda.synchronize()
# .to("meta") releases storage for all parameters/buffers regardless
# of their original device (CUDA or CPU).
model.to("meta")
cleanup_memory()
if alloc_trim_strategy == AllocatorTrimStrategy.TRIM:
synchronize_device()
# .to("meta") releases storage for all parameters/buffers regardless
# of their original device (CUDA or CPU).
model.to("meta")
cleanup_memory()
packages/ltx-pipelines/src/ltx_pipelines/utils/helpers.py
+3 -12
View File
@@ -1,4 +1,3 @@
import gc
import logging
import torch
@@ -9,6 +8,7 @@ from ltx_core.conditioning import (
VideoConditionByKeyframeIndex,
VideoConditionByLatentIndex,
)
from ltx_core.devices import cleanup_accelerator_memory, get_preferred_device
from ltx_core.model.audio_vae import encode_audio
from ltx_core.model.transformer import Modality
from ltx_core.model.video_vae import TilingConfig, VideoEncoder
@@ -28,20 +28,11 @@ from ltx_pipelines.utils.media_io import (
def get_device() -> torch.device:
if torch.cuda.is_available():
return torch.device("cuda", torch.cuda.current_device())
return torch.device("cpu")
return get_preferred_device()
def cleanup_memory() -> None:
gc.collect()
torch.cuda.empty_cache()
torch.cuda.synchronize()
try:
if hasattr(torch._C, "_host_emptyCache"):
torch._C._host_emptyCache()
except Exception:
logging.warning("Host empty cache cleanup failed; ignoring.", exc_info=True)
cleanup_accelerator_memory()
def _conform_latent_length(latent: torch.Tensor, expected_frames_count: int) -> torch.Tensor:
packages/ltx-pipelines/src/ltx_pipelines/utils/media_io.py
+37 -8
View File
@@ -258,16 +258,22 @@ def decode_image(image_path: str) -> np.ndarray:
return np_array
def _write_audio(container: av.container.Container, audio_stream: av.audio.AudioStream, audio: Audio) -> None:
def _validate_audio_waveform(audio: Audio) -> None:
"""Raise ValueError if the waveform is empty or not stereo ``(2, N)`` / ``(N, 2)``."""
samples = audio.waveform
if samples.ndim == 1:
samples = samples[:, None]
if samples.numel() == 0:
raise ValueError("audio.waveform is empty; pass audio=None for no audio.")
if samples.ndim != 2 or 2 not in samples.shape:
raise ValueError(f"audio.waveform must be stereo (2, N) or (N, 2); got shape {tuple(samples.shape)}.")
if samples.shape[1] != 2 and samples.shape[0] == 2:
samples = samples.T
if samples.shape[1] != 2:
raise ValueError(f"Expected samples with 2 channels; got shape {samples.shape}.")
def _normalize_audio_waveform(samples: torch.Tensor) -> torch.Tensor:
"""Transpose a validated stereo waveform to channel-last ``(N, 2)``."""
return samples.T if samples.shape[1] != 2 else samples
def _write_audio(container: av.container.Container, audio_stream: av.audio.AudioStream, audio: Audio) -> None:
samples = _normalize_audio_waveform(audio.waveform)
# Convert to int16 packed for ingestion; resampler converts to encoder fmt.
if samples.dtype != torch.int16:
@@ -335,13 +341,35 @@ def encode_video(
preset: str = "veryfast",
thread_count: int = 0,
) -> None:
"""Encode RGB frames to an H.264 file, optionally muxing an audio track.
Args:
video: RGB frames as a ``(F, H, W, C)`` float ``[0, 1]`` tensor, or an iterator of
such per-chunk tensors (e.g. the VAE decoder output). An empty iterator raises.
fps: Output frame rate.
audio: Audio track to mux, or None for a video-only file. Waveform must be stereo
``(2, N)`` or ``(N, 2)``.
output_path: Destination path. Partial output is removed if encoding fails.
video_chunks_number: Number of chunks yielded by ``video``, for the progress bar.
frame_converter: Float-to-pixel converter (default YUV420p BT.709).
crf: libx264 constant rate factor; lower is higher quality (0-51).
preset: libx264 speed/compression preset.
thread_count: libx264 thread count (0 = auto).
Raises:
ValueError: On an empty ``video`` or a non-stereo ``audio.waveform``.
"""
if audio is not None:
_validate_audio_waveform(audio)
if isinstance(video, torch.Tensor):
video = iter([video])
def convert(chunk: torch.Tensor) -> torch.Tensor:
return frame_converter(chunk.movedim(-1, -3))
first_chunk = convert(next(video))
first_raw_chunk = next(video, None)
if first_raw_chunk is None:
raise ValueError("video is empty; expected at least one frame chunk.")
first_chunk = convert(first_raw_chunk)
if frame_converter.pixel_format == PixelFormat.RGB24:
height, width = first_chunk.shape[-3], first_chunk.shape[-2]
@@ -397,6 +425,7 @@ def encode_audio(audio: Audio, output_path: str) -> None:
the only difference is a PCM (``pcm_s16le``) stream in a WAV container instead of
the AAC stream used for muxed video.
"""
_validate_audio_waveform(audio)
container = av.open(output_path, mode="w")
audio_stream = container.add_stream("pcm_s16le", rate=audio.sampling_rate)
audio_stream.codec_context.sample_rate = audio.sampling_rate
packages/ltx-pipelines/src/ltx_pipelines/utils/samplers.py
+23 -16
View File
@@ -8,6 +8,7 @@ from tqdm import tqdm
from ltx_core.components.diffusion_steps import EulerCfgPpDiffusionStep, Res2sDiffusionStep
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.devices import highest_precision_float
from ltx_core.model.transformer import X0Model
from ltx_core.utils import to_denoised, to_velocity
from ltx_pipelines.utils.helpers import post_process_latent, timesteps_from_mask
@@ -157,7 +158,10 @@ def _channelwise_normalize(x: torch.Tensor) -> torch.Tensor:
def _get_new_noise(x: torch.Tensor, generator: torch.Generator) -> torch.Tensor:
noise = torch.randn(x.shape, generator=generator, dtype=torch.float64, device=generator.device)
# float64 on CUDA/CPU for numerical stability; MPS has no float64, so degrade to float32.
noise = torch.randn(
x.shape, generator=generator, dtype=highest_precision_float(generator.device), device=generator.device
)
noise = (noise - noise.mean()) / noise.std()
return _channelwise_normalize(noise)
@@ -175,10 +179,11 @@ def _inject_sde_noise(
eta: float = 0.5,
) -> torch.Tensor:
sigmas_copy = sigmas.clone()
hp = highest_precision_float(state.denoise_mask.device)
new_noise = new_noise_fn(state.latent, step_noise_generator)
if not legacy_mode:
timesteps = timesteps_from_mask(state.denoise_mask.double(), sigmas_copy[step_idx].double())
next_timesteps = timesteps_from_mask(state.denoise_mask.double(), sigmas_copy[step_idx + 1].double())
timesteps = timesteps_from_mask(state.denoise_mask.to(hp), sigmas_copy[step_idx].to(hp))
next_timesteps = timesteps_from_mask(state.denoise_mask.to(hp), sigmas_copy[step_idx + 1].to(hp))
sigmas = torch.stack([timesteps, next_timesteps])
step_idx = 0
x_next = stepper.step(
@@ -249,6 +254,8 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915,PLR0912
if present_state is None:
raise ValueError("At least one of video_state or audio_state must be provided")
state_device = present_state.latent.device
# float64 on CUDA/CPU for ODE numerical stability; MPS has no float64, so degrade to float32.
hp = highest_precision_float(state_device)
# Initialize noise generators with different seeds
if noise_seed_substep is None:
@@ -270,19 +277,19 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915,PLR0912
if sigmas[-1] == 0:
sigmas = torch.cat([sigmas[:-1], torch.tensor([0.0011, 0.0], device=sigmas.device)], dim=0)
# Compute step sizes in hyperbolic space
hs = -torch.log(sigmas[1:].double().cpu() / (sigmas[:-1].double().cpu()))
hs = -torch.log(sigmas[1:].to(hp).cpu() / (sigmas[:-1].to(hp).cpu()))
# Initialize phi cache for reuse across loop iterations
phi_cache = {}
c2 = 0.5 # Midpoint for res_2s
for step_idx in tqdm(range(n_full_steps)):
sigma = sigmas[step_idx].double()
sigma_next = sigmas[step_idx + 1].double()
sigma = sigmas[step_idx].to(hp)
sigma_next = sigmas[step_idx + 1].to(hp)
# Initialize anchor point
x_anchor_video = video_state.latent.clone().double() if video_state is not None else None
x_anchor_audio = audio_state.latent.clone().double() if audio_state is not None else None
x_anchor_video = video_state.latent.clone().to(hp) if video_state is not None else None
x_anchor_audio = audio_state.latent.clone().to(hp) if audio_state is not None else None
# ====================================================================
# STAGE 1: Evaluate at current point
@@ -307,15 +314,15 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915,PLR0912
# Compute substep x using RK coefficient a21
# ====================================================================
if x_anchor_video is not None and denoised_video_1 is not None:
eps_1_video = denoised_video_1.double() - x_anchor_video
x_mid_video = x_anchor_video.double() + h * a21 * eps_1_video
eps_1_video = denoised_video_1.to(hp) - x_anchor_video
x_mid_video = x_anchor_video.to(hp) + h * a21 * eps_1_video
else:
eps_1_video = None
x_mid_video = None
if x_anchor_audio is not None and denoised_audio_1 is not None:
eps_1_audio = denoised_audio_1.double() - x_anchor_audio
x_mid_audio = x_anchor_audio.double() + h * a21 * eps_1_audio
eps_1_audio = denoised_audio_1.to(hp) - x_anchor_audio
x_mid_audio = x_anchor_audio.to(hp) + h * a21 * eps_1_audio
else:
eps_1_audio = None
x_mid_audio = None
@@ -347,10 +354,10 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915,PLR0912
for _ in range(bongmath_max_iter):
if x_mid_video is not None and eps_1_video is not None:
x_anchor_video = x_mid_video - h * a21 * eps_1_video
eps_1_video = denoised_video_1.double() - x_anchor_video
eps_1_video = denoised_video_1.to(hp) - x_anchor_video
if x_mid_audio is not None and eps_1_audio is not None:
x_anchor_audio = x_mid_audio - h * a21 * eps_1_audio
eps_1_audio = denoised_audio_1.double() - x_anchor_audio
eps_1_audio = denoised_audio_1.to(hp) - x_anchor_audio
# ====================================================================
# STAGE 2: Evaluate at substep point (WITH NOISE)
@@ -384,13 +391,13 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915,PLR0912
# FINAL COMBINATION: Compute x_next using RK coefficients
# ====================================================================
if x_anchor_video is not None and eps_1_video is not None and denoised_video_2 is not None:
eps_2_video = denoised_video_2.double() - x_anchor_video
eps_2_video = denoised_video_2.to(hp) - x_anchor_video
x_next_video = x_anchor_video + h * (b1 * eps_1_video + b2 * eps_2_video)
else:
x_next_video = None
if x_anchor_audio is not None and eps_1_audio is not None and denoised_audio_2 is not None:
eps_2_audio = denoised_audio_2.double() - x_anchor_audio
eps_2_audio = denoised_audio_2.to(hp) - x_anchor_audio
x_next_audio = x_anchor_audio + h * (b1 * eps_1_audio + b2 * eps_2_audio)
else:
x_next_audio = None