Merge pull request #36 from Lightricks/pr-2026-01-08
Open a PR to sync - 2026-01-08
This commit is contained in:
@@ -3,17 +3,15 @@
|
||||
[](https://ltx.io)
|
||||
[](https://huggingface.co/Lightricks/LTX-2)
|
||||
[](https://app.ltx.studio/ltx-2-playground/i2v)
|
||||
[](https://videos.ltx.io/LTX-2/grants/LTX_2_Technical_Report_compressed.pdf)
|
||||
[](https://videos.ltx.io/LTX-2/grants/LTX_2_Technical_Report_compressed.pdf)
|
||||
[](https://discord.gg/ltxplatform)
|
||||
|
||||
**LTX-2** is the first DiT-based audio-video foundation model that contains all core capabilities of modern video generation in one model: synchronized audio and video, high fidelity, multiple performance modes, production-ready outputs, API access, and open access.
|
||||
|
||||
|
||||
<div align="center">
|
||||
<video src="https://github.com/user-attachments/assets/4414adc0-086c-43de-b367-9362eeb20228" width="70%" poster=""> </video>
|
||||
</div>
|
||||
|
||||
|
||||
## 🚀 Quick Start
|
||||
|
||||
```bash
|
||||
@@ -40,7 +38,7 @@ Download the following models from the [LTX-2 HuggingFace repository](https://hu
|
||||
**Spatial Upscaler** - Required for current two-stage pipeline implementations in this repository
|
||||
* [`ltx-2-spatial-upscaler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-spatial-upscaler-x2-1.0.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-spatial-upscaler-x2-1.0.safetensors)
|
||||
|
||||
**Temporal Upscaler** - Required for current two-stage pipeline implementations in this repository
|
||||
**Temporal Upscaler** - Supported by the model and will be required for future pipeline implementations
|
||||
* [`ltx-2-temporal-upscaler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-temporal-upscaler-x2-1.0.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-temporal-upscaler-x2-1.0.safetensors)
|
||||
|
||||
**Distilled LoRA** - Required for current two-stage pipeline implementations in this repository (except DistilledPipeline and ICLoraPipeline)
|
||||
@@ -64,7 +62,7 @@ Download the following models from the [LTX-2 HuggingFace repository](https://hu
|
||||
|
||||
### Available Pipelines
|
||||
|
||||
* **[TI2VidTwoStagesPipeline](packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py)** - Production-quality text-to-video with 2x upsampling (recommended)
|
||||
* **[TI2VidTwoStagesPipeline](packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py)** - Production-quality text/image-to-video with 2x upsampling (recommended)
|
||||
* **[TI2VidOneStagePipeline](packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py)** - Single-stage generation for quick prototyping
|
||||
* **[DistilledPipeline](packages/ltx-pipelines/src/ltx_pipelines/distilled.py)** - Fastest inference with 8 predefined sigmas
|
||||
* **[ICLoraPipeline](packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py)** - Video-to-video and image-to-video transformations
|
||||
|
||||
+66
-47
@@ -31,11 +31,11 @@ pip install -e packages/ltx-core
|
||||
|
||||
### Core Models
|
||||
|
||||
- **Transformer** ([`model/transformer/`](src/ltx_core/model/transformer/)): The 48-layer LTX-2 transformer with cross-modal attention for joint audio-video processing. Expects inputs in [`Modality`](src/ltx_core/model/transformer/modality.py) format
|
||||
- **Transformer** ([`model/transformer/`](src/ltx_core/model/transformer/)): The asymmetric dual-stream LTX-2 transformer (14B-parameter video stream, 5B-parameter audio stream) with bidirectional cross-modal attention for joint audio-video processing. Expects inputs in [`Modality`](src/ltx_core/model/transformer/modality.py) format
|
||||
- **Video VAE** ([`model/video_vae/`](src/ltx_core/model/video_vae/)): Encodes/decodes video pixels to/from latent space with temporal and spatial compression
|
||||
- **Audio VAE** ([`model/audio_vae/`](src/ltx_core/model/audio_vae/)): Encodes/decodes audio spectrograms to/from latent space
|
||||
- **Vocoder** ([`model/audio_vae/`](src/ltx_core/model/audio_vae/)): Neural vocoder that converts mel spectrograms to audio waveforms
|
||||
- **Text Encoder** ([`text_encoders/`](src/ltx_core/text_encoders/)): Gemma-based encoder that produces separate embeddings for video and audio conditioning
|
||||
- **Text Encoder** ([`text_encoders/`](src/ltx_core/text_encoders/)): Gemma 3-based multilingual encoder with multi-layer feature extraction and thinking tokens that produces separate embeddings for video and audio conditioning
|
||||
- **Spatial Upscaler** ([`model/upsampler/`](src/ltx_core/model/upsampler/)): Upsamples latent representations for higher-resolution generation
|
||||
|
||||
### Diffusion Components
|
||||
@@ -76,7 +76,14 @@ This section provides a deep dive into the internal architecture of the LTX-2 Au
|
||||
|
||||
## High-Level Architecture
|
||||
|
||||
LTX-2 is a **joint Audio-Video diffusion transformer** that processes both modalities simultaneously in a unified architecture. Unlike traditional models that handle video and audio separately, LTX-2 uses cross-modal attention to enable natural synchronization.
|
||||
LTX-2 is an **asymmetric dual-stream diffusion transformer** that jointly models the text-conditioned distribution of video and audio signals, capturing true joint dependencies (unlike sequential T2V→V2A pipelines).
|
||||
|
||||
### Key Design Principles
|
||||
|
||||
- **Decoupled Latent Representations**: Separate modality-specific VAEs enable 3D RoPE (video) vs 1D RoPE (audio), independent compression optimization, and native V2A/A2V editing workflows
|
||||
- **Asymmetric Dual-Stream**: 14B-parameter video stream (spatiotemporal dynamics) + 5B-parameter audio stream (1D temporal), sharing 48 transformer blocks but differing in width
|
||||
- **Bidirectional Cross-Modal Attention**: 1D temporal RoPE enables sub-frame alignment, mapping visual cues to auditory events (lip-sync, foley, environmental acoustics)
|
||||
- **Cross-Modality AdaLN**: Scaling/shift parameters conditioned on the other modality's hidden states for synchronization across differing diffusion timesteps/temporal resolutions
|
||||
|
||||
```text
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
@@ -84,21 +91,26 @@ LTX-2 is a **joint Audio-Video diffusion transformer** that processes both modal
|
||||
│ │
|
||||
│ Video Pixels → Video VAE Encoder → Video Latents │
|
||||
│ Audio Waveform → Audio VAE Encoder → Audio Latents │
|
||||
│ Text Prompt → Gemma Encoder → Text Embeddings │
|
||||
│ Text Prompt → Gemma 3 Encoder → Text Embeddings │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
↓
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ LTX-2 TRANSFORMER (48 Blocks) │
|
||||
│ LTX-2 ASYMMETRIC DUAL-STREAM TRANSFORMER (48 Blocks) │
|
||||
│ │
|
||||
│ ┌──────────────┐ ┌──────────────┐ │
|
||||
│ │ Video Stream │ │ Audio Stream │ │
|
||||
│ │ │ │ │ │
|
||||
│ │ Self-Attn │ │ Self-Attn │ │
|
||||
│ │ Cross-Attn │ │ Cross-Attn │ │
|
||||
│ │ │◄────────────►│ │ │
|
||||
│ │ A↔V Cross │ │ A↔V Cross │ │
|
||||
│ │ Feed-Forward │ │ Feed-Forward │ │
|
||||
│ └──────────────┘ └──────────────┘ │
|
||||
│ ┌──────────────────────┐ ┌──────────────────────┐ │
|
||||
│ │ Video Stream (14B) │ │ Audio Stream (5B) │ │
|
||||
│ │ │ │ │ │
|
||||
│ │ 3D RoPE (x,y,t) │ │ 1D RoPE (temporal) │ │
|
||||
│ │ │ │ │ │
|
||||
│ │ Self-Attn │ │ Self-Attn │ │
|
||||
│ │ Text Cross-Attn │ │ Text Cross-Attn │ │
|
||||
│ │ │◄────►│ │ │
|
||||
│ │ A↔V Cross-Attn │ │ A↔V Cross-Attn │ │
|
||||
│ │ (1D temporal RoPE) │ │ (1D temporal RoPE) │ │
|
||||
│ │ Cross-modality │ │ Cross-modality │ │
|
||||
│ │ AdaLN │ │ AdaLN │ │
|
||||
│ │ Feed-Forward │ │ Feed-Forward │ │
|
||||
│ └──────────────────────┘ └──────────────────────┘ │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
↓
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
@@ -106,7 +118,7 @@ LTX-2 is a **joint Audio-Video diffusion transformer** that processes both modal
|
||||
│ │
|
||||
│ Video Latents → Video VAE Decoder → Video Pixels │
|
||||
│ Audio Latents → Audio VAE Decoder → Mel Spectrogram │
|
||||
│ Mel Spectrogram → Vocoder → Audio Waveform │
|
||||
│ Mel Spectrogram → Vocoder → Audio Waveform (24 kHz) │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
@@ -114,7 +126,7 @@ LTX-2 is a **joint Audio-Video diffusion transformer** that processes both modal
|
||||
|
||||
## The Transformer
|
||||
|
||||
The core of LTX-2 is a 48-layer transformer that processes both video and audio tokens simultaneously.
|
||||
The core of LTX-2 is an **asymmetric dual-stream diffusion transformer** with 48 layers that processes both video and audio tokens simultaneously. The architecture allocates 14B parameters to the video stream and 5B parameters to the audio stream, reflecting the different information densities of the two modalities.
|
||||
|
||||
### Model Structure
|
||||
|
||||
@@ -126,29 +138,36 @@ The `LTXModel` class implements the transformer. It supports both video-only and
|
||||
|
||||
**Source**: [`src/ltx_core/model/transformer/transformer.py`](src/ltx_core/model/transformer/transformer.py)
|
||||
|
||||
Each dual-stream block performs four operations sequentially:
|
||||
|
||||
1. **Self-Attention**: Within-modality attention for each stream
|
||||
2. **Text Cross-Attention**: Textual prompt conditioning for both streams
|
||||
3. **Audio-Visual Cross-Attention**: Bidirectional inter-modal exchange
|
||||
4. **Feed-Forward Network (FFN)**: Feature refinement
|
||||
|
||||
```text
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ TRANSFORMER BLOCK │
|
||||
│ │
|
||||
│ VIDEO PATH: │
|
||||
│ Input → RMSNorm → AdaLN → Self-Attn (attn1) │
|
||||
│ → RMSNorm → Cross-Attn (attn2, text) │
|
||||
│ → RMSNorm → AdaLN → A↔V Cross-Attn │
|
||||
│ → RMSNorm → AdaLN → Feed-Forward (ff) → Output │
|
||||
│ VIDEO (14B): Input → RMSNorm → AdaLN → Self-Attn → │
|
||||
│ RMSNorm → Text Cross-Attn → │
|
||||
│ RMSNorm → AdaLN → A↔V Cross-Attn (1D RoPE) → │
|
||||
│ RMSNorm → AdaLN → FFN → Output │
|
||||
│ │
|
||||
│ AUDIO PATH: │
|
||||
│ Input → RMSNorm → AdaLN → Self-Attn (audio_attn1) │
|
||||
│ → RMSNorm → Cross-Attn (audio_attn2, text) │
|
||||
│ → RMSNorm → AdaLN → A↔V Cross-Attn │
|
||||
│ → RMSNorm → AdaLN → Feed-Forward (audio_ff) │
|
||||
│ AUDIO (5B): Input → RMSNorm → AdaLN → Self-Attn → │
|
||||
│ RMSNorm → Text Cross-Attn → │
|
||||
│ RMSNorm → AdaLN → A↔V Cross-Attn (1D RoPE) → │
|
||||
│ RMSNorm → AdaLN → FFN → Output │
|
||||
│ │
|
||||
│ AdaLN (Adaptive Layer Normalization): │
|
||||
│ - Uses scale_shift_table (6 params) for video/audio │
|
||||
│ - Uses scale_shift_table_a2v_ca (5 params) for A↔V CA │
|
||||
│ - Conditioned on per-token timestep embeddings │
|
||||
│ RoPE: Video=3D (x,y,t), Audio=1D (t), Cross-Attn=1D (t) │
|
||||
│ AdaLN: Timestep-conditioned, cross-modality for A↔V CA │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Audio-Visual Cross-Attention Details
|
||||
|
||||
Bidirectional cross-attention enables tight temporal alignment: video and audio streams exchange information bidirectionally using 1D temporal RoPE (synchronization only, no spatial alignment). AdaLN gates condition on each modality's timestep for cross-modal synchronization.
|
||||
|
||||
### Perturbations
|
||||
|
||||
The transformer supports [**perturbations**](src/ltx_core/guidance/perturbations.py) that selectively skip attention operations.
|
||||
@@ -189,12 +208,11 @@ The Audio VAE ([`src/ltx_core/model/audio_vae/`](src/ltx_core/model/audio_vae/))
|
||||
|
||||
### Audio VAE Architecture
|
||||
|
||||
- **Encoder**: Compresses mel spectrogram `[B, mel_bins, T]` → `[B, 8, T/4, 16]` latents
|
||||
- Temporal downsampling: 4× (`LATENT_DOWNSAMPLE_FACTOR = 4`)
|
||||
- Frequency bins: Fixed 16 mel bins in latent space
|
||||
- Latent channels: 8
|
||||
- **Decoder**: Expands `[B, 8, T, 16]` latents → mel spectrogram `[B, mel_bins, T*4]`
|
||||
- **Vocoder**: Converts mel spectrogram → audio waveform
|
||||
Compact neural audio representation optimized for diffusion-based training. Natively supports stereo: processes two-channel mel-spectrograms (16 kHz input) with channel concatenation before encoding.
|
||||
|
||||
- **Encoder**: `[B, mel_bins, T]` → `[B, 8, T/4, 16]` latents (4× temporal downsampling, 8 channels, 16 mel bins in latent space, ~1/25s per token, 128-dim feature vector)
|
||||
- **Decoder**: `[B, 8, T, 16]` → `[B, mel_bins, T*4]` mel spectrogram
|
||||
- **Vocoder**: HiFi-GAN-based, modified for stereo synthesis and upsampling (16 kHz mel → 24 kHz waveform, doubled generator capacity for stereo)
|
||||
|
||||
**Downsampling**:
|
||||
|
||||
@@ -207,19 +225,20 @@ The Audio VAE is used internally by pipelines for encoding mel spectrograms to l
|
||||
|
||||
## Text Encoding (Gemma)
|
||||
|
||||
LTX-2 uses **Gemma** (Google's open LLM) as the text encoder, located in [`src/ltx_core/text_encoders/gemma/`](src/ltx_core/text_encoders/gemma/).
|
||||
LTX-2 uses **Gemma 3** (Gemma 3-12B) as the multilingual text encoder backbone, located in [`src/ltx_core/text_encoders/gemma/`](src/ltx_core/text_encoders/gemma/). Advanced text understanding is critical not only for global language support but for the phonetic and semantic accuracy of generated speech.
|
||||
|
||||
### Text Encoder Architecture
|
||||
|
||||
- **Tokenizer**: Converts text → token IDs
|
||||
- **Gemma Model**: Processes tokens → embeddings
|
||||
- **Text Projection**: Uses `PixArtAlphaTextProjection` to project caption embeddings
|
||||
- Two-layer MLP with GELU (tanh approximation) or SiLU activation
|
||||
- Projects from caption channels (3840) to model dimensions
|
||||
- **Feature Extractor**: Extracts video/audio-specific embeddings
|
||||
- **Separate Encoders**:
|
||||
- `AVEncoder`: For audio-video generation (outputs separate video and audio contexts)
|
||||
- `VideoOnlyEncoder`: For video-only generation
|
||||
The text conditioning pipeline consists of three stages:
|
||||
|
||||
1. **Gemma 3 Backbone**: Decoder-only LLM processes text tokens → embeddings across all layers `[B, T, D, L]`
|
||||
2. **Multi-Layer Feature Extractor**: Aggregates features from all decoder layers (not just final layer), applies mean-centered scaling, flattens to `[B, T, D×L]`, and projects via learnable matrix W (jointly optimized with LTX-2, LLM weights frozen)
|
||||
3. **Text Connector**: Bidirectional transformer blocks with learnable registers (replacing padded positions, also referred to as "thinking tokens" in the paper) for contextual mixing. Separate connectors for video and audio streams (`Embeddings1DConnector`)
|
||||
|
||||
**Encoders**:
|
||||
|
||||
- `AVGemmaTextEncoderModel`: Audio-video generation (two connectors → `AVGemmaEncoderOutput` with separate video/audio contexts)
|
||||
- `VideoGemmaTextEncoderModel`: Video-only generation (single connector → `VideoGemmaEncoderOutput`)
|
||||
|
||||
### System Prompts
|
||||
|
||||
@@ -228,7 +247,7 @@ System prompts are also used to enhance user's prompts.
|
||||
- **Text-to-Video**: [`gemma_t2v_system_prompt.txt`](src/ltx_core/text_encoders/gemma/encoders/prompts/gemma_t2v_system_prompt.txt)
|
||||
- **Image-to-Video**: [`gemma_i2v_system_prompt.txt`](src/ltx_core/text_encoders/gemma/encoders/prompts/gemma_i2v_system_prompt.txt)
|
||||
|
||||
**Important**: Video and audio receive **different** context embeddings, even from the same prompt. This allows better modality-specific conditioning.
|
||||
**Important**: Video and audio receive **different** context embeddings, even from the same prompt. This allows better modality-specific conditioning and enables the model to synthesize speech that is synchronized with visual lip movement while being natural in cadence, accent, and emotional tone.
|
||||
|
||||
**Output Format**:
|
||||
|
||||
|
||||
@@ -150,8 +150,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
|
||||
run_vx = video is not None and video.enabled and vx.numel() > 0
|
||||
run_ax = audio is not None and audio.enabled and ax.numel() > 0
|
||||
|
||||
run_a2v = run_vx and (audio is not None and audio.enabled and ax.numel() > 0)
|
||||
run_v2a = run_ax and (video is not None and video.enabled and vx.numel() > 0)
|
||||
run_a2v = run_vx and (audio is not None and ax.numel() > 0)
|
||||
run_v2a = run_ax and (video is not None and vx.numel() > 0)
|
||||
|
||||
if run_vx:
|
||||
vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(
|
||||
|
||||
@@ -43,7 +43,7 @@ All pipelines can be run directly from the command line. Each pipeline module is
|
||||
# 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 path/to/distilled_lora.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" \
|
||||
@@ -55,11 +55,11 @@ python -m ltx_pipelines.ti2vid_two_stages --help
|
||||
|
||||
Available pipeline modules:
|
||||
|
||||
- `ltx_pipelines.ti2vid_two_stages` - Two-stage text-to-video (recommended)
|
||||
- `ltx_pipelines.ti2vid_one_stage` - Single-stage text-to-video
|
||||
- `ltx_pipelines.distilled` - Fast distilled pipeline
|
||||
- `ltx_pipelines.ic_lora` - Video-to-video with IC-LoRA
|
||||
- `ltx_pipelines.keyframe_interpolation` - Keyframe interpolation
|
||||
- `ltx_pipelines.ti2vid_two_stages` - Two-stage text/image-to-video (recommended).
|
||||
- `ltx_pipelines.ti2vid_one_stage` - Single-stage text/image-to-video.
|
||||
- `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.
|
||||
|
||||
Use `--help` with any pipeline module to see all available options and parameters.
|
||||
|
||||
@@ -73,9 +73,9 @@ Use `--help` with any pipeline module to see all available options and parameter
|
||||
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 keyframe images to interpolate?
|
||||
│ └─ NO → Do you have multiple keyframe images to interpolate?
|
||||
│ ├─ YES → Use KeyframeInterpolationPipeline
|
||||
│ └─ NO → Use ICLoraPipeline (image conditioning only)
|
||||
│ └─ NO → Use TI2VidTwoStagesPipeline (image conditioning only)
|
||||
│
|
||||
└─ NO → Text-to-video only
|
||||
├─ Do you need best quality?
|
||||
@@ -103,7 +103,7 @@ Do you need to condition on existing images/videos?
|
||||
|
||||
### 1. TI2VidTwoStagesPipeline
|
||||
|
||||
**Best for:** High-quality text-to-video generation with upsampling. **Recommended for production use.**
|
||||
**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)
|
||||
|
||||
@@ -212,8 +212,7 @@ When authoring custom scripts, pass the `fp8transformer` flag to pipeline classe
|
||||
```python
|
||||
pipeline = TI2VidTwoStagesPipeline(
|
||||
checkpoint_path=ltx_model_path,
|
||||
distilled_lora_path=distilled_lora_path,
|
||||
distilled_lora_strength=0.6,
|
||||
distilled_lora=distilled_lora,
|
||||
spatial_upsampler_path=upsampler_path,
|
||||
gemma_root=gemma_root_path,
|
||||
loras=[],
|
||||
@@ -275,11 +274,20 @@ This allows you to use **20-30 steps instead of 40** while maintaining quality.
|
||||
## 📖 Example: Image-to-Video
|
||||
|
||||
```python
|
||||
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
|
||||
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
|
||||
|
||||
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_path="/path/to/distilled_lora.safetensors",
|
||||
distilled_lora=distilled_lora,
|
||||
spatial_upsampler_path="/path/to/upsampler.safetensors",
|
||||
gemma_root="/path/to/gemma",
|
||||
loras=[],
|
||||
|
||||
@@ -68,7 +68,10 @@ class LatentsDecoder:
|
||||
self.vae = load_video_vae_decoder(model_path, device=self.device, dtype=torch.bfloat16)
|
||||
|
||||
if vae_tiling:
|
||||
self.vae.enable_tiling()
|
||||
logger.warning(
|
||||
"VAE tiling is not yet implemented in this script. "
|
||||
"Continuing without tiling - this may cause OOM errors for large resolutions."
|
||||
)
|
||||
|
||||
if with_audio:
|
||||
with console.status(f"[bold]Loading audio VAE decoder from {model_path}...", spinner="dots"):
|
||||
|
||||
@@ -48,6 +48,13 @@ def preprocess_dataset( # noqa: PLR0913
|
||||
with_audio: bool = False,
|
||||
) -> None:
|
||||
"""Run the preprocessing pipeline with the given arguments."""
|
||||
# VAE tiling is not yet implemented
|
||||
if vae_tiling:
|
||||
logger.warning(
|
||||
"VAE tiling is not yet implemented in this script. "
|
||||
"Continuing without tiling - this may cause OOM errors for large resolutions."
|
||||
)
|
||||
|
||||
# Validate dataset file
|
||||
_validate_dataset_file(dataset_file)
|
||||
|
||||
|
||||
@@ -489,7 +489,10 @@ def compute_latents( # noqa: PLR0913, PLR0915
|
||||
vae = load_video_vae_encoder(model_path, device=torch_device, dtype=torch.bfloat16)
|
||||
|
||||
if vae_tiling:
|
||||
vae.enable_tiling()
|
||||
logger.warning(
|
||||
"VAE tiling is not yet implemented in this script. "
|
||||
"Continuing without tiling - this may cause OOM errors for large resolutions."
|
||||
)
|
||||
|
||||
# Load audio VAE encoder and audio processor if needed
|
||||
audio_vae_encoder = None
|
||||
|
||||
Reference in New Issue
Block a user