Automated PR - 2026-03-04
This commit is contained in:
sync-bot
+262
-88
@@ -4,19 +4,28 @@ This file provides guidance to AI coding assistants (Claude, Cursor, etc.) when
|
||||
|
||||
## Project Overview
|
||||
|
||||
**LTX-2 Trainer** is a training toolkit for fine-tuning the Lightricks LTX-2 audio-video generation model. It supports:
|
||||
**LTX Trainer** is a training toolkit for fine-tuning the Lightricks LTX audio-video generation models. It supports:
|
||||
|
||||
- **LoRA training** - Efficient fine-tuning with adapters
|
||||
- **Full fine-tuning** - Complete model training
|
||||
- **Audio-video training** - Joint audio and video generation
|
||||
- **IC-LoRA training** - In-context control adapters for video-to-video transformations
|
||||
|
||||
**Supported model versions:**
|
||||
|
||||
- **LTX-2** (19B, initial audio-video model)
|
||||
- **LTX-2.3** (20B, improved text conditioning and audio quality)
|
||||
|
||||
Version detection is fully automatic — ltx-core reads the checkpoint config and selects the correct architecture
|
||||
components. The trainer does not need version-specific code paths.
|
||||
|
||||
**Key Dependencies:**
|
||||
|
||||
- **[`ltx-core`](../ltx-core/)** - Core model implementations (transformer, VAE, text encoder)
|
||||
- **[`ltx-core`](../ltx-core/)** - Core model implementations (transformer, VAE, text encoder, scheduler)
|
||||
- **[`ltx-pipelines`](../ltx-pipelines/)** - Inference pipeline components
|
||||
|
||||
> **Important:** This trainer only supports **LTX-2** (the audio-video model). The older LTXV models are not supported.
|
||||
> **Important:** This trainer only supports **LTX-2 and later** (audio-video models). The older LTXV (video-only) models
|
||||
> are not supported.
|
||||
|
||||
## Architecture Overview
|
||||
|
||||
@@ -24,36 +33,45 @@ This file provides guidance to AI coding assistants (Claude, Cursor, etc.) when
|
||||
|
||||
```
|
||||
packages/ltx-trainer/
|
||||
├── src/ltx_trainer/ # Main training module
|
||||
│ ├── config.py # Pydantic configuration models
|
||||
│ ├── trainer.py # Main training orchestration with Accelerate
|
||||
│ ├── model_loader.py # Model loading using ltx-core
|
||||
│ ├── validation_sampler.py # Inference for validation samples
|
||||
│ ├── datasets.py # PrecomputedDataset for latent-based training
|
||||
│ ├── training_strategies/ # Strategy pattern for different training modes
|
||||
│ │ ├── __init__.py # Factory function: get_training_strategy()
|
||||
│ │ ├── base_strategy.py # TrainingStrategy ABC, ModelInputs, TrainingStrategyConfigBase
|
||||
│ │ ├── text_to_video.py # TextToVideoStrategy, TextToVideoConfig
|
||||
│ │ └── video_to_video.py # VideoToVideoStrategy, VideoToVideoConfig
|
||||
│ ├── timestep_samplers.py # Flow matching timestep sampling
|
||||
│ ├── captioning.py # Video captioning utilities
|
||||
│ ├── video_utils.py # Video processing utilities
|
||||
│ └── hf_hub_utils.py # HuggingFace Hub integration
|
||||
├── scripts/ # User-facing CLI tools
|
||||
│ ├── train.py # Main training script
|
||||
│ ├── process_dataset.py # Dataset preprocessing
|
||||
│ ├── process_videos.py # Video latent encoding
|
||||
│ ├── process_captions.py # Text embedding computation
|
||||
│ ├── caption_videos.py # Automatic video captioning
|
||||
│ ├── decode_latents.py # Latent decoding for debugging
|
||||
│ ├── inference.py # Inference with trained models
|
||||
│ ├── compute_reference.py # Generate IC-LoRA reference videos
|
||||
│ └── split_scenes.py # Scene detection and splitting
|
||||
├── configs/ # Example training configurations
|
||||
│ ├── ltx2_av_lora.yaml # Audio-video LoRA training
|
||||
│ ├── ltx2_v2v_ic_lora.yaml # IC-LoRA video-to-video
|
||||
│ └── accelerate/ # Accelerate configs for distributed training
|
||||
└── docs/ # Documentation
|
||||
├── src/ltx_trainer/ # Main training module
|
||||
│ ├── __init__.py # Logger setup, path config
|
||||
│ ├── config.py # Pydantic configuration models
|
||||
│ ├── config_display.py # Config pretty-printing
|
||||
│ ├── trainer.py # Main training orchestration with Accelerate
|
||||
│ ├── model_loader.py # Model loading using ltx-core
|
||||
│ ├── validation_sampler.py # Inference for validation samples
|
||||
│ ├── datasets.py # PrecomputedDataset, DummyDataset
|
||||
│ ├── training_strategies/ # Strategy pattern for different training modes
|
||||
│ │ ├── __init__.py # Factory function: get_training_strategy()
|
||||
│ │ ├── base_strategy.py # TrainingStrategy ABC, ModelInputs, TrainingStrategyConfigBase
|
||||
│ │ ├── text_to_video.py # TextToVideoStrategy, TextToVideoConfig
|
||||
│ │ └── video_to_video.py # VideoToVideoStrategy, VideoToVideoConfig
|
||||
│ ├── timestep_samplers.py # Flow matching timestep sampling
|
||||
│ ├── gemma_8bit.py # 8-bit Gemma text encoder loading (bitsandbytes)
|
||||
│ ├── quantization.py # Transformer INT8/INT4/FP8 quantization
|
||||
│ ├── captioning.py # Video captioning utilities
|
||||
│ ├── video_utils.py # Video I/O and processing
|
||||
│ ├── gpu_utils.py # GPU memory helpers
|
||||
│ ├── hf_hub_utils.py # HuggingFace Hub integration
|
||||
│ ├── progress.py # Training progress display
|
||||
│ └── utils.py # Image I/O helpers
|
||||
├── scripts/ # User-facing CLI tools
|
||||
│ ├── train.py # Main training script
|
||||
│ ├── process_dataset.py # Dataset preprocessing (latents + captions)
|
||||
│ ├── process_videos.py # Video latent encoding
|
||||
│ ├── process_captions.py # Text embedding computation
|
||||
│ ├── caption_videos.py # Automatic video captioning
|
||||
│ ├── decode_latents.py # Latent decoding for debugging
|
||||
│ ├── inference.py # Inference with trained models
|
||||
│ ├── compute_reference.py # Generate IC-LoRA reference videos
|
||||
│ └── split_scenes.py # Scene detection and splitting
|
||||
├── configs/ # Example training configurations
|
||||
│ ├── ltx2_av_lora.yaml # Audio-video LoRA training
|
||||
│ ├── ltx2_av_lora_low_vram.yaml
|
||||
│ ├── ltx2_v2v_ic_lora.yaml # IC-LoRA video-to-video
|
||||
│ └── accelerate/ # FSDP, DDP configs
|
||||
├── tests/ # Pytest tests
|
||||
└── docs/ # Documentation
|
||||
```
|
||||
|
||||
### Key Architectural Patterns
|
||||
@@ -61,35 +79,42 @@ packages/ltx-trainer/
|
||||
**Model Loading:**
|
||||
|
||||
- `ltx_trainer.model_loader` provides component loaders using `ltx-core`
|
||||
- Individual loaders: `load_transformer()`, `load_video_vae_encoder()`, `load_video_vae_decoder()`, `load_text_encoder()`, etc.
|
||||
- Individual loaders: `load_transformer()`, `load_video_vae_encoder()`, `load_video_vae_decoder()`,
|
||||
`load_text_encoder()`, etc.
|
||||
- Combined loader: `load_model()` returns `LtxModelComponents` dataclass
|
||||
- Uses `SingleGPUModelBuilder` from ltx-core internally
|
||||
- 8-bit text encoder loading via `gemma_8bit.py` (bitsandbytes)
|
||||
|
||||
**Training Flow:**
|
||||
|
||||
1. Configuration loaded via Pydantic models in `config.py`
|
||||
2. `Trainer` class orchestrates the training loop
|
||||
3. Training strategies (`TextToVideoStrategy`, `VideoToVideoStrategy`) prepare inputs and compute loss
|
||||
4. Accelerate handles distributed training and device placement
|
||||
5. Data flows as precomputed latents through `PrecomputedDataset`
|
||||
2. `LtxvTrainer` class orchestrates the training loop
|
||||
3. Text encoder loaded on GPU → validation embeddings cached → heavy components unloaded (only `embeddings_processor`
|
||||
kept)
|
||||
4. Each training step: embedding connectors applied → strategy prepares `ModelInputs` → transformer forward pass →
|
||||
strategy computes loss
|
||||
5. Training strategies (`TextToVideoStrategy`, `VideoToVideoStrategy`) handle mode-specific logic
|
||||
6. Accelerate handles distributed training, mixed precision, and device placement
|
||||
7. Data flows as precomputed latents through `PrecomputedDataset`
|
||||
|
||||
**Model Interface (Modality-based):**
|
||||
|
||||
```python
|
||||
from ltx_core.model.transformer.modality import Modality
|
||||
|
||||
# Create modality objects for video and audio
|
||||
video = Modality(
|
||||
enabled=True,
|
||||
latent=video_latents, # [B, seq_len, 128]
|
||||
timesteps=video_timesteps, # [B, seq_len] per-token
|
||||
positions=video_positions, # [B, 3, seq_len, 2]
|
||||
context=video_embeds,
|
||||
context_mask=None,
|
||||
latent=video_latents, # [B, seq_len, 128] patchified latent tokens
|
||||
sigma=sigma, # [B,] current noise level (per-batch)
|
||||
timesteps=video_timesteps, # [B, seq_len] per-token timestep embeddings
|
||||
positions=video_positions, # [B, 3, seq_len, 2] positional coordinates
|
||||
context=video_embeds, # text conditioning embeddings
|
||||
context_mask=None, # optional attention mask for text context
|
||||
)
|
||||
audio = Modality(
|
||||
enabled=True,
|
||||
latent=audio_latents,
|
||||
sigma=sigma,
|
||||
timesteps=audio_timesteps,
|
||||
positions=audio_positions, # [B, 1, seq_len, 2]
|
||||
context=audio_embeds,
|
||||
@@ -102,14 +127,91 @@ video_pred, audio_pred = model(video=video, audio=audio, perturbations=None)
|
||||
|
||||
> **Note:** `Modality` is immutable (frozen dataclass). Use `dataclasses.replace()` to modify.
|
||||
|
||||
**`sigma` vs `timesteps`:** These serve different roles. `timesteps` is per-token (e.g. `sigma * denoise_mask` —
|
||||
conditioning tokens get 0, noisy tokens get sigma). `sigma` is per-batch and is used for prompt AdaLN conditioning (
|
||||
LTX-2.3) and cross-modality (video↔audio) attention conditioning (both versions).
|
||||
|
||||
**Configuration System:**
|
||||
|
||||
- All config in `src/ltx_trainer/config.py`
|
||||
- Main class: `LtxTrainerConfig`
|
||||
- Training strategy configs: `TextToVideoConfig`, `VideoToVideoConfig`
|
||||
- Uses Pydantic field validators and model validators
|
||||
- Config uses `extra="forbid"` — unknown fields cause validation errors
|
||||
- Config files in `configs/` directory
|
||||
|
||||
## LTX-2 vs LTX-2.3: Differences
|
||||
|
||||
Both model versions share the same latent space interface (see [Latent Space Constants](#latent-space-constants)).
|
||||
The differences lie in how text conditioning and audio generation work. Version detection is automatic via checkpoint
|
||||
config — the trainer uses a unified API.
|
||||
|
||||
| Component | LTX-2 (19B) | LTX-2.3 (20B) |
|
||||
|-----------------------|---------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------|
|
||||
| Feature extractor | `FeatureExtractorV1`: single `aggregate_embed`, same output for video and audio | `FeatureExtractorV2`: separate `video_aggregate_embed` + `audio_aggregate_embed`, per-token RMSNorm |
|
||||
| Caption projection | Inside the transformer (`caption_projection`) | Inside the feature extractor (before connector) |
|
||||
| Embeddings connectors | Same dimensions for video and audio | Separate dimensions (`AudioEmbeddings1DConnectorConfigurator`) |
|
||||
| Prompt AdaLN | Not present (`cross_attention_adaln=False`) | Active — modulates cross-attention to text using `sigma` |
|
||||
| Vocoder | HiFi-GAN (`Vocoder`) | BigVGAN v2 + bandwidth extension (`VocoderWithBWE`) |
|
||||
|
||||
**How version detection works in ltx-core:**
|
||||
|
||||
- **Feature extractor:** `_create_feature_extractor()` checks for V2 config keys (`caption_proj_before_connector`,
|
||||
etc.). Present → V2; absent → V1.
|
||||
- **Vocoder:** `VocoderConfigurator` checks for `config["vocoder"]["bwe"]`. Present → `VocoderWithBWE`; absent →
|
||||
`Vocoder`.
|
||||
- **Transformer:** `_build_caption_projections()` checks `caption_proj_before_connector`. True (V2) → no caption
|
||||
projection in transformer; False (V1) → caption projection created in transformer.
|
||||
- **Embeddings connectors:** `AudioEmbeddings1DConnectorConfigurator` reads `audio_connector_*` keys, falling back to
|
||||
video connector keys for V1 backward compatibility.
|
||||
|
||||
## Text Encoder Pipeline
|
||||
|
||||
The `GemmaTextEncoder` implements a 3-block pipeline:
|
||||
|
||||
1. **Block 1 — Gemma LLM:** Tokenizes text → runs through Gemma → extracts hidden states
|
||||
2. **Block 2 — Feature extractor:** Hidden states → normalized features (V1: single stream duplicated for video/audio;
|
||||
V2: separate video and audio projections)
|
||||
3. **Block 3 — Embeddings processor:** Features → embeddings connectors → final context embeddings for the transformer
|
||||
|
||||
**Precomputed embeddings (offline):** `process_captions.py` runs Blocks 1+2 via `text_encoder.precompute()` and saves
|
||||
the results. Block 3 (connectors) is applied during training via
|
||||
`text_encoder.embeddings_processor.create_embeddings()`.
|
||||
|
||||
**Precomputed embeddings formats:**
|
||||
|
||||
- **New format** (from `precompute()`): saves `video_prompt_embeds`, `audio_prompt_embeds` (optional),
|
||||
`prompt_attention_mask`
|
||||
- **Legacy format** (from old `_preprocess_text()`): saves `prompt_embeds`, `prompt_attention_mask`
|
||||
|
||||
The trainer handles both formats in `_training_step()`: if `video_prompt_embeds` is present, it uses the new format;
|
||||
otherwise, it duplicates `prompt_embeds` for both modalities (mirroring V1 behavior).
|
||||
|
||||
**After caching validation embeddings**, the trainer unloads heavy components to free VRAM:
|
||||
|
||||
```python
|
||||
self._text_encoder.model = None
|
||||
self._text_encoder.tokenizer = None
|
||||
self._text_encoder.feature_extractor = None
|
||||
# Only embeddings_processor (connectors) remains — used during training
|
||||
```
|
||||
|
||||
## Latent Space Constants
|
||||
|
||||
These values are shared across all supported model versions:
|
||||
|
||||
| Constant | Value | Where used |
|
||||
|------------------------------|----------------------------------|-----------------------------------------------------------|
|
||||
| Video latent channels | 128 | VAE encoder/decoder, patchifier, `VideoLatentShape` |
|
||||
| Spatial compression | 32× (H and W) | `SpatioTemporalScaleFactors.default()`, config validators |
|
||||
| Temporal compression | 8× | `SpatioTemporalScaleFactors.default()`, config validators |
|
||||
| Frame constraint | `frames % 8 == 1` | Config validators, validation sampler |
|
||||
| Resolution constraint | Width and height divisible by 32 | Config validators, validation sampler |
|
||||
| Audio latent channels | 8 | `AudioLatentShape`, audio patchifier |
|
||||
| Audio mel bins | 16 | `AudioLatentShape`, audio patchifier |
|
||||
| Patchified token dim (video) | 128 (`128 × 1 × 1 × 1`) | Transformer `in_channels` |
|
||||
| Patchified token dim (audio) | 128 (`8 × 16`) | Transformer `audio_in_channels` |
|
||||
|
||||
## Development Commands
|
||||
|
||||
### Setup and Installation
|
||||
@@ -180,25 +282,34 @@ uv run accelerate launch scripts/train.py configs/ltx2_av_lora.yaml
|
||||
**`src/ltx_trainer/config.py`** - Master config definitions
|
||||
|
||||
Key classes:
|
||||
|
||||
- `LtxTrainerConfig` - Main configuration container
|
||||
- `ModelConfig` - Model paths and training mode
|
||||
- `TrainingStrategyConfig` - Union of `TextToVideoConfig` | `VideoToVideoConfig`
|
||||
- `LoraConfig` - LoRA hyperparameters
|
||||
- `OptimizationConfig` - Learning rate, batch size, etc.
|
||||
- `ValidationConfig` - Validation settings
|
||||
- `WandbConfig` - W&B logging settings
|
||||
- `ModelConfig` - Model paths, training mode (`lora` | `full`), checkpoint loading
|
||||
- `TrainingStrategyConfig` - Union of `TextToVideoConfig` | `VideoToVideoConfig` (discriminated by `name`)
|
||||
- `LoraConfig` - Rank, alpha, dropout, target modules
|
||||
- `OptimizationConfig` - Learning rate, batch size, gradient accumulation, scheduler, gradient checkpointing
|
||||
- `AccelerationConfig` - Mixed precision, quantization, 8-bit text encoder
|
||||
- `DataConfig` - Preprocessed data root, dataloader workers
|
||||
- `ValidationConfig` - Prompts, video dimensions, CFG/STG guidance, audio generation, inference steps
|
||||
- `CheckpointsConfig` - Save interval, retention, precision
|
||||
- `FlowMatchingConfig` - Timestep sampling mode and parameters
|
||||
- `HubConfig` - HuggingFace Hub push settings
|
||||
- `WandbConfig` - Weights & Biases logging
|
||||
|
||||
**⚠️ When modifying config.py:**
|
||||
|
||||
1. Update ALL config files in `configs/`
|
||||
2. Update `docs/configuration-reference.md`
|
||||
3. Test that all configs remain valid
|
||||
|
||||
### Training Core
|
||||
|
||||
**`src/ltx_trainer/trainer.py`** - Main training loop
|
||||
**`src/ltx_trainer/trainer.py`** - Main training loop (`LtxvTrainer`)
|
||||
|
||||
- Implements distributed training with Accelerate
|
||||
- Handles mixed precision, gradient accumulation, checkpointing
|
||||
- `_training_step()` applies embedding connectors then delegates to strategy
|
||||
- `_load_text_encoder_and_cache_embeddings()` caches validation embeddings and unloads heavy components
|
||||
- Uses training strategies for mode-specific logic
|
||||
|
||||
**`src/ltx_trainer/training_strategies/`** - Strategy pattern
|
||||
@@ -208,35 +319,55 @@ Key classes:
|
||||
- `video_to_video.py`: IC-LoRA video-to-video transformations
|
||||
|
||||
Key methods each strategy implements:
|
||||
|
||||
- `get_data_sources()` - Required data directories
|
||||
- `prepare_training_inputs()` - Convert batch to `ModelInputs`
|
||||
- `compute_loss()` - Calculate training loss
|
||||
- `prepare_training_inputs()` - Convert batch to `ModelInputs` with `Modality` objects
|
||||
- `compute_loss()` - Calculate training loss (velocity prediction, MSE with masking)
|
||||
- `requires_audio` property - Whether audio components needed
|
||||
|
||||
**`src/ltx_trainer/model_loader.py`** - Model loading
|
||||
|
||||
Component loaders:
|
||||
|
||||
- `load_transformer()` → `LTXModel`
|
||||
- `load_video_vae_encoder()` → `VideoVAEEncoder`
|
||||
- `load_video_vae_decoder()` → `VideoVAEDecoder`
|
||||
- `load_audio_vae_decoder()` → `AudioVAEDecoder`
|
||||
- `load_vocoder()` → `Vocoder`
|
||||
- `load_text_encoder()` → `AVGemmaTextEncoderModel`
|
||||
- `load_video_vae_encoder()` → `VideoEncoder`
|
||||
- `load_video_vae_decoder()` → `VideoDecoder`
|
||||
- `load_audio_vae_decoder()` → `AudioDecoder`
|
||||
- `load_vocoder()` → `Vocoder` or `VocoderWithBWE` (auto-detected)
|
||||
- `load_text_encoder()` → `GemmaTextEncoder` (unified, handles V1/V2 automatically)
|
||||
- `load_model()` → `LtxModelComponents` (convenience wrapper)
|
||||
|
||||
**`src/ltx_trainer/validation_sampler.py`** - Inference for validation
|
||||
|
||||
Uses ltx-core components for denoising:
|
||||
|
||||
- `LTX2Scheduler` for sigma scheduling
|
||||
- `EulerDiffusionStep` for diffusion steps
|
||||
- `CFGGuider` for classifier-free guidance
|
||||
- `STGGuider` for spatio-temporal guidance
|
||||
|
||||
**`src/ltx_trainer/timestep_samplers.py`** - Flow matching timestep sampling
|
||||
|
||||
- `UniformTimestepSampler` - Uniform sampling in `[min, max]`
|
||||
- `ShiftedLogitNormalTimestepSampler` - Stretched shifted logit-normal distribution with:
|
||||
- Shift determined by sequence length (more noise at higher token counts)
|
||||
- Percentile stretching for better `[0, 1]` coverage
|
||||
- Uniform fallback (10% of samples) to prevent distribution collapse
|
||||
- Reflection around `eps` for numerical stability near zero
|
||||
|
||||
**`src/ltx_trainer/gemma_8bit.py`** - 8-bit text encoder loading
|
||||
|
||||
Bypasses ltx-core's standard loading path to enable bitsandbytes 8-bit quantization of the Gemma backbone. Manually
|
||||
constructs the `GemmaTextEncoder` with quantized model, feature extractor, and embeddings processor.
|
||||
|
||||
### Data
|
||||
|
||||
**`src/ltx_trainer/datasets.py`** - Dataset handling
|
||||
|
||||
- `PrecomputedDataset` loads pre-computed VAE latents
|
||||
- Supports video latents, audio latents, text embeddings, reference latents
|
||||
- `PrecomputedDataset` loads pre-computed VAE latents and text embeddings
|
||||
- Supports video latents, audio latents, text embeddings, reference latents (for IC-LoRA)
|
||||
- Handles legacy patchified format `[seq_len, C]` → automatically unpatchifies to `[C, F, H, W]`
|
||||
- `DummyDataset` for benchmarking and minimal testing
|
||||
|
||||
## Common Development Tasks
|
||||
|
||||
@@ -263,23 +394,40 @@ Uses ltx-core components for denoising:
|
||||
from dataclasses import replace
|
||||
from ltx_core.model.transformer.modality import Modality
|
||||
|
||||
# Create modality
|
||||
# Create modality — all fields except enabled and masks are required
|
||||
video = Modality(
|
||||
enabled=True,
|
||||
latent=latents,
|
||||
timesteps=timesteps,
|
||||
positions=positions,
|
||||
context=context,
|
||||
latent=latents, # [B, seq_len, 128]
|
||||
sigma=sigma, # [B,] — the per-batch noise level
|
||||
timesteps=timesteps, # [B, seq_len] — per-token (sigma * denoise_mask)
|
||||
positions=positions, # [B, 3, seq_len, 2]
|
||||
context=context, # text embeddings from embeddings_processor
|
||||
context_mask=None,
|
||||
)
|
||||
|
||||
# Update (immutable - must use replace)
|
||||
video = replace(video, latent=new_latent, timesteps=new_timesteps)
|
||||
# Update (immutable — must use replace)
|
||||
video = replace(video, latent=new_latent, sigma=new_sigma, timesteps=new_timesteps)
|
||||
|
||||
# Disable a modality
|
||||
audio = replace(audio, enabled=False)
|
||||
```
|
||||
|
||||
### Working with the Text Encoder
|
||||
|
||||
```python
|
||||
# Full forward pass (used for validation — runs all 3 blocks)
|
||||
video_embeds, audio_embeds, attention_mask = text_encoder(prompt)
|
||||
|
||||
# Precompute features (used in process_captions.py — runs blocks 1+2 only)
|
||||
video_features, audio_features, attention_mask = text_encoder.precompute(prompt, padding_side="left")
|
||||
|
||||
# Apply connectors during training (block 3 only)
|
||||
additive_mask = text_encoder._convert_to_additive_mask(attention_mask, video_features.dtype)
|
||||
video_embeds, audio_embeds, binary_mask = text_encoder.embeddings_processor.create_embeddings(
|
||||
video_features, audio_features, additive_mask
|
||||
)
|
||||
```
|
||||
|
||||
## Debugging Tips
|
||||
|
||||
**Training Issues:**
|
||||
@@ -293,18 +441,27 @@ audio = replace(audio, enabled=False)
|
||||
- Ensure `model_path` points to a local `.safetensors` file
|
||||
- Ensure `text_encoder_path` points to a Gemma model directory
|
||||
- URLs are NOT supported for model paths
|
||||
- For 8-bit loading: ensure `bitsandbytes` is installed
|
||||
|
||||
**Configuration:**
|
||||
|
||||
- Validation errors: Check validators in `config.py`
|
||||
- Unknown fields: Config uses `extra="forbid"` - all fields must be defined
|
||||
- Unknown fields: Config uses `extra="forbid"` — all fields must be defined
|
||||
- Strategy validation: IC-LoRA requires `reference_videos` in validation config
|
||||
- Video-to-video strategy requires `training_mode: "lora"`
|
||||
|
||||
**Precomputed Data:**
|
||||
|
||||
- Legacy data (`prompt_embeds`) works via backward-compat in `_training_step()`
|
||||
- New data (`video_prompt_embeds` + `audio_prompt_embeds`) is the expected format
|
||||
- Latents must be in `[C, F, H, W]` format (legacy `[seq_len, C]` is auto-converted)
|
||||
|
||||
## Key Constraints
|
||||
|
||||
### LTX-2 Frame Requirements
|
||||
### Frame Requirements
|
||||
|
||||
Frames must satisfy `frames % 8 == 1`:
|
||||
|
||||
- ✅ Valid: 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, 97, 121
|
||||
- ❌ Invalid: 24, 32, 48, 64, 100
|
||||
|
||||
@@ -321,7 +478,7 @@ Width and height must be divisible by 32.
|
||||
### Platform Requirements
|
||||
|
||||
- Linux required (uses `triton` which is Linux-only)
|
||||
- CUDA GPU with 24GB+ VRAM recommended
|
||||
- CUDA GPU with 24GB+ VRAM recommended (80GB+ for full fine-tuning)
|
||||
|
||||
## Reference: ltx-core Key Components
|
||||
|
||||
@@ -329,24 +486,41 @@ Width and height must be divisible by 32.
|
||||
packages/ltx-core/src/ltx_core/
|
||||
├── model/
|
||||
│ ├── transformer/
|
||||
│ │ ├── model.py # LTXModel
|
||||
│ │ ├── modality.py # Modality dataclass
|
||||
│ │ └── transformer.py # BasicAVTransformerBlock
|
||||
│ │ ├── model.py # LTXModel (diffusion transformer)
|
||||
│ │ ├── modality.py # Modality dataclass
|
||||
│ │ ├── transformer.py # BasicAVTransformerBlock
|
||||
│ │ ├── transformer_args.py # TransformerArgsPreprocessor (sigma → prompt AdaLN)
|
||||
│ │ ├── model_configurator.py # LTXModelConfigurator (version-aware)
|
||||
│ │ └── timestep_embedding.py # Timestep/sigma embedding
|
||||
│ ├── video_vae/
|
||||
│ │ └── video_vae.py # Encoder, Decoder
|
||||
│ │ ├── video_vae.py # VideoEncoder, VideoDecoder
|
||||
│ │ └── model_configurator.py # VideoEncoderConfigurator, VideoDecoderConfigurator
|
||||
│ ├── audio_vae/
|
||||
│ │ ├── audio_vae.py # Decoder
|
||||
│ │ └── vocoder.py # Vocoder
|
||||
│ └── clip/gemma/
|
||||
│ └── encoders/av_encoder.py # AVGemmaTextEncoderModel
|
||||
├── pipeline/
|
||||
│ ├── components/
|
||||
│ │ ├── schedulers.py # LTX2Scheduler
|
||||
│ │ ├── diffusion_steps.py # EulerDiffusionStep
|
||||
│ │ ├── guiders.py # CFGGuider
|
||||
│ │ └── patchifiers.py # VideoLatentPatchifier, AudioPatchifier
|
||||
│ └── conditioning/ # VideoLatentTools, AudioLatentTools
|
||||
└── loader/
|
||||
├── single_gpu_model_builder.py # SingleGPUModelBuilder
|
||||
└── sd_ops.py # Key remapping (SDOps)
|
||||
│ │ ├── audio_vae.py # AudioEncoder, AudioDecoder
|
||||
│ │ └── vocoder.py # Vocoder, VocoderWithBWE (output_sampling_rate)
|
||||
│ └── common/ # Shared model components
|
||||
├── text_encoders/gemma/
|
||||
│ ├── __init__.py # Exports: GemmaTextEncoder, GemmaTextEncoderConfigurator,
|
||||
│ │ # AV_GEMMA_TEXT_ENCODER_KEY_OPS, GEMMA_MODEL_OPS,
|
||||
│ │ # module_ops_from_gemma_root
|
||||
│ ├── encoders/
|
||||
│ │ ├── base_encoder.py # GemmaTextEncoder (unified 3-block pipeline)
|
||||
│ │ └── encoder_configurator.py # GemmaTextEncoderConfigurator, _create_feature_extractor
|
||||
│ ├── feature_extractor.py # FeatureExtractorV1 (19B), FeatureExtractorV2 (20B)
|
||||
│ ├── embeddings_connector.py # Embeddings1DConnector, Embeddings1DConnectorConfigurator,
|
||||
│ │ # AudioEmbeddings1DConnectorConfigurator
|
||||
│ ├── embeddings_processor.py # EmbeddingsProcessor (wraps video + audio connectors)
|
||||
│ └── tokenizer.py # LTXVGemmaTokenizer
|
||||
├── components/
|
||||
│ ├── schedulers.py # LTX2Scheduler
|
||||
│ ├── diffusion_steps.py # EulerDiffusionStep
|
||||
│ ├── guiders.py # CFGGuider, STGGuider
|
||||
│ └── patchifiers.py # VideoLatentPatchifier, AudioPatchifier
|
||||
├── conditioning/ # ConditioningItem, mask_utils, types
|
||||
├── tools.py # VideoLatentTools, AudioLatentTools
|
||||
├── loader/
|
||||
│ ├── single_gpu_model_builder.py # SingleGPUModelBuilder
|
||||
│ ├── sft_loader.py # SafetensorsModelStateDictLoader
|
||||
│ └── sd_ops.py # Key remapping (SDOps)
|
||||
└── types.py # SpatioTemporalScaleFactors, VideoLatentShape, AudioLatentShape
|
||||
```
|
||||
|
||||
@@ -282,6 +282,7 @@ class InpaintingStrategy(TrainingStrategy):
|
||||
video_modality = Modality(
|
||||
enabled=True,
|
||||
latent=noisy_latents,
|
||||
sigma=sigmas,
|
||||
timesteps=timesteps,
|
||||
positions=positions,
|
||||
context=video_prompt_embeds,
|
||||
|
||||
@@ -287,7 +287,7 @@ class LatentsDecoder:
|
||||
|
||||
# Save as WAV
|
||||
output_path = output_dir / f"{latent_file.stem}.wav"
|
||||
sample_rate = self.vocoder.output_sample_rate
|
||||
sample_rate = self.vocoder.output_sampling_rate
|
||||
torchaudio.save(str(output_path), waveform[0].cpu(), sample_rate)
|
||||
|
||||
|
||||
|
||||
@@ -410,7 +410,7 @@ def main() -> None: # noqa: PLR0912, PLR0915
|
||||
# Get audio sample rate from vocoder if audio was generated
|
||||
audio_sample_rate = None
|
||||
if audio is not None and components.vocoder is not None:
|
||||
audio_sample_rate = components.vocoder.output_sample_rate
|
||||
audio_sample_rate = components.vocoder.output_sampling_rate
|
||||
|
||||
save_video(
|
||||
video_tensor=video,
|
||||
|
||||
@@ -303,14 +303,13 @@ def compute_captions_embeddings( # noqa: PLR0913
|
||||
) as progress:
|
||||
task = progress.add_task("Processing captions", total=len(dataloader))
|
||||
for batch in dataloader:
|
||||
# Encode prompts using _preprocess_text (returns embeddings before connector)
|
||||
# This is what we want to save - the connector is applied during training
|
||||
# Encode prompts using precompute() (returns video/audio features before connector)
|
||||
# The connector is applied during training via embeddings_processor
|
||||
with torch.inference_mode():
|
||||
# TODO(batch-tokenization): When tokenizer supports batching, encode all prompts at once:
|
||||
# prompt_embeds, prompt_attention_mask = text_encoder._preprocess_text(batch["prompt"]) # noqa: ERA001
|
||||
# TODO(batch-tokenization): When tokenizer supports batching, encode all prompts at once.
|
||||
# For now, process one at a time:
|
||||
for i in range(len(batch["prompt"])):
|
||||
prompt_embeds, prompt_attention_mask = text_encoder._preprocess_text(
|
||||
video_prompt_embeds, audio_prompt_embeds, prompt_attention_mask = text_encoder.precompute(
|
||||
batch["prompt"][i], padding_side="left"
|
||||
)
|
||||
|
||||
@@ -321,9 +320,11 @@ def compute_captions_embeddings( # noqa: PLR0913
|
||||
output_dir_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
embedding_data = {
|
||||
"prompt_embeds": prompt_embeds[0].cpu().contiguous(),
|
||||
"video_prompt_embeds": video_prompt_embeds[0].cpu().contiguous(),
|
||||
"prompt_attention_mask": prompt_attention_mask[0].cpu().contiguous(),
|
||||
}
|
||||
if audio_prompt_embeds is not None:
|
||||
embedding_data["audio_prompt_embeds"] = audio_prompt_embeds[0].cpu().contiguous()
|
||||
|
||||
output_file = output_path / output_rel_path
|
||||
torch.save(embedding_data, output_file)
|
||||
|
||||
@@ -41,6 +41,7 @@ from torchvision.transforms.functional import crop, resize, to_tensor
|
||||
from transformers.utils.logging import disable_progress_bar
|
||||
|
||||
from ltx_core.model.audio_vae import AudioProcessor
|
||||
from ltx_core.types import Audio
|
||||
from ltx_trainer import logger
|
||||
from ltx_trainer.model_loader import load_audio_vae_encoder, load_video_vae_encoder
|
||||
from ltx_trainer.utils import open_image_as_srgb
|
||||
@@ -503,7 +504,7 @@ def compute_latents( # noqa: PLR0913, PLR0915
|
||||
)
|
||||
# Create audio processor for waveform-to-spectrogram conversion
|
||||
audio_processor = AudioProcessor(
|
||||
sample_rate=audio_vae_encoder.sample_rate,
|
||||
target_sample_rate=audio_vae_encoder.sample_rate,
|
||||
mel_bins=audio_vae_encoder.mel_bins,
|
||||
mel_hop_length=audio_vae_encoder.mel_hop_length,
|
||||
n_fft=audio_vae_encoder.n_fft,
|
||||
@@ -567,10 +568,10 @@ def compute_latents( # noqa: PLR0913, PLR0915
|
||||
if audio_batch is not None:
|
||||
# Extract the i-th item from batched audio data
|
||||
# DataLoader collates [channels, samples] -> [batch, channels, samples]
|
||||
audio_data = {
|
||||
"waveform": audio_batch["waveform"][i],
|
||||
"sample_rate": audio_batch["sample_rate"][i].item(),
|
||||
}
|
||||
audio_data = Audio(
|
||||
waveform=audio_batch["waveform"][i],
|
||||
sampling_rate=audio_batch["sample_rate"][i].item(),
|
||||
)
|
||||
|
||||
# Encode audio
|
||||
with torch.inference_mode():
|
||||
@@ -822,13 +823,13 @@ def tiled_encode_video( # noqa: PLR0912, PLR0915
|
||||
def encode_audio(
|
||||
audio_vae_encoder: torch.nn.Module,
|
||||
audio_processor: torch.nn.Module,
|
||||
audio_data: dict[str, torch.Tensor | int],
|
||||
audio: Audio,
|
||||
) -> dict[str, torch.Tensor | int | float]:
|
||||
"""Encode audio waveform into latent representation.
|
||||
Args:
|
||||
audio_vae_encoder: Audio VAE encoder model from ltx-core
|
||||
audio_processor: AudioProcessor for waveform-to-spectrogram conversion
|
||||
audio_data: Dict with {"waveform": Tensor[channels, samples], "sample_rate": int}
|
||||
audio: Audio container with waveform tensor and sampling rate.
|
||||
Returns:
|
||||
Dict containing audio latents and shape information:
|
||||
{
|
||||
@@ -841,18 +842,17 @@ def encode_audio(
|
||||
device = next(audio_vae_encoder.parameters()).device
|
||||
dtype = next(audio_vae_encoder.parameters()).dtype
|
||||
|
||||
waveform = audio_data["waveform"].to(device=device, dtype=dtype)
|
||||
sample_rate = audio_data["sample_rate"]
|
||||
waveform = audio.waveform.to(device=device, dtype=dtype)
|
||||
|
||||
# Add batch dimension if needed: [channels, samples] -> [batch, channels, samples]
|
||||
if waveform.dim() == 2:
|
||||
waveform = waveform.unsqueeze(0)
|
||||
|
||||
# Calculate duration
|
||||
duration = waveform.shape[-1] / sample_rate
|
||||
duration = waveform.shape[-1] / audio.sampling_rate
|
||||
|
||||
# Convert waveform to mel spectrogram using AudioProcessor
|
||||
mel_spectrogram = audio_processor.waveform_to_mel(waveform, waveform_sample_rate=sample_rate)
|
||||
mel_spectrogram = audio_processor.waveform_to_mel(Audio(waveform=waveform, sampling_rate=audio.sampling_rate))
|
||||
mel_spectrogram = mel_spectrogram.to(dtype=dtype)
|
||||
|
||||
# Encode mel spectrogram to latents
|
||||
|
||||
@@ -67,14 +67,18 @@ class DummyDataset(Dataset):
|
||||
"fps": self.fps,
|
||||
},
|
||||
"text_conditions": {
|
||||
"prompt_embeds": torch.randn(
|
||||
"video_prompt_embeds": torch.randn(
|
||||
self.prompt_sequence_length,
|
||||
self.prompt_embed_dim,
|
||||
), # random text embeddings
|
||||
),
|
||||
"audio_prompt_embeds": torch.randn(
|
||||
self.prompt_sequence_length,
|
||||
self.prompt_embed_dim,
|
||||
),
|
||||
"prompt_attention_mask": torch.ones(
|
||||
self.prompt_sequence_length,
|
||||
dtype=torch.bool,
|
||||
), # random attention mask
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -18,28 +18,26 @@ import logging
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
|
||||
from ltx_core.text_encoders.gemma.embeddings_connector import Embeddings1DConnectorConfigurator
|
||||
from ltx_core.text_encoders.gemma.encoders.av_encoder import (
|
||||
AV_GEMMA_TEXT_ENCODER_KEY_OPS,
|
||||
AVGemmaTextEncoderModel,
|
||||
from ltx_core.text_encoders.gemma import AV_GEMMA_TEXT_ENCODER_KEY_OPS
|
||||
from ltx_core.text_encoders.gemma.embeddings_connector import (
|
||||
AudioEmbeddings1DConnectorConfigurator,
|
||||
Embeddings1DConnectorConfigurator,
|
||||
)
|
||||
from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
|
||||
from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessor
|
||||
from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder
|
||||
from ltx_core.text_encoders.gemma.encoders.encoder_configurator import _create_feature_extractor
|
||||
from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ltx_core.text_encoders.gemma import AVGemmaTextEncoderModel
|
||||
|
||||
|
||||
def load_8bit_gemma(
|
||||
checkpoint_path: str | Path,
|
||||
gemma_model_path: str | Path,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
) -> "AVGemmaTextEncoderModel":
|
||||
) -> GemmaTextEncoder:
|
||||
"""Load the Gemma text encoder in 8-bit precision using bitsandbytes.
|
||||
This function bypasses ltx-core's standard loading path to enable 8-bit quantization
|
||||
via the bitsandbytes library. The Gemma model is loaded with load_in_8bit=True and
|
||||
@@ -50,7 +48,7 @@ def load_8bit_gemma(
|
||||
gemma_model_path: Path to Gemma model directory
|
||||
dtype: Data type for non-quantized model weights (feature extractor, connectors)
|
||||
Returns:
|
||||
Loaded AVGemmaTextEncoderModel with 8-bit quantized Gemma backbone
|
||||
Loaded GemmaTextEncoder with 8-bit quantized Gemma backbone
|
||||
Raises:
|
||||
ImportError: If bitsandbytes is not installed
|
||||
FileNotFoundError: If required model files are not found
|
||||
@@ -88,28 +86,35 @@ def load_8bit_gemma(
|
||||
def extract_state_dict(prefix: str) -> dict[str, torch.Tensor]:
|
||||
return {k.replace(prefix, ""): v for k, v in sd.sd.items() if k.startswith(prefix)}
|
||||
|
||||
# Create and load feature extractor
|
||||
feature_extractor = GemmaFeaturesExtractorProjLinear()
|
||||
feature_extractor.load_state_dict(extract_state_dict("feature_extractor_linear."))
|
||||
feature_extractor = feature_extractor.to(device=gemma_model.device, dtype=dtype)
|
||||
|
||||
# Create and load video embeddings connector
|
||||
embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
|
||||
embeddings_connector.load_state_dict(extract_state_dict("embeddings_connector."))
|
||||
embeddings_connector.load_state_dict(extract_state_dict("embeddings_processor.video_connector."))
|
||||
embeddings_connector = embeddings_connector.to(device=gemma_model.device, dtype=dtype)
|
||||
|
||||
# Create and load audio embeddings connector
|
||||
audio_embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
|
||||
audio_embeddings_connector.load_state_dict(extract_state_dict("audio_embeddings_connector."))
|
||||
audio_embeddings_connector = AudioEmbeddings1DConnectorConfigurator.from_config(config)
|
||||
audio_embeddings_connector.load_state_dict(extract_state_dict("embeddings_processor.audio_connector."))
|
||||
audio_embeddings_connector = audio_embeddings_connector.to(device=gemma_model.device, dtype=dtype)
|
||||
|
||||
# Construct the text encoder
|
||||
text_encoder = AVGemmaTextEncoderModel(
|
||||
feature_extractor_linear=feature_extractor,
|
||||
embeddings_connector=embeddings_connector,
|
||||
audio_embeddings_connector=audio_embeddings_connector,
|
||||
# Create embeddings processor
|
||||
embeddings_processor = EmbeddingsProcessor(
|
||||
video_connector=embeddings_connector,
|
||||
audio_connector=audio_embeddings_connector,
|
||||
)
|
||||
|
||||
transformer_config = config.get("transformer", {})
|
||||
feature_extractor = _create_feature_extractor(transformer_config)
|
||||
feature_extractor.load_state_dict(
|
||||
{k.removeprefix("feature_extractor."): v for k, v in sd.sd.items() if k.startswith("feature_extractor.")},
|
||||
)
|
||||
feature_extractor = feature_extractor.to(device=gemma_model.device, dtype=dtype)
|
||||
|
||||
text_encoder = GemmaTextEncoder(
|
||||
feature_extractor=feature_extractor,
|
||||
embeddings_processor=embeddings_processor,
|
||||
tokenizer=tokenizer,
|
||||
model=gemma_model,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
return text_encoder
|
||||
|
||||
@@ -32,7 +32,7 @@ if TYPE_CHECKING:
|
||||
from ltx_core.model.audio_vae import AudioDecoder, AudioEncoder, Vocoder
|
||||
from ltx_core.model.transformer import LTXModel
|
||||
from ltx_core.model.video_vae import VideoDecoder, VideoEncoder
|
||||
from ltx_core.text_encoders.gemma import AVGemmaTextEncoderModel
|
||||
from ltx_core.text_encoders.gemma import GemmaTextEncoder
|
||||
|
||||
|
||||
def _to_torch_device(device: Device) -> torch.device:
|
||||
@@ -192,7 +192,7 @@ def load_text_encoder(
|
||||
device: Device = "cpu",
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
load_in_8bit: bool = False,
|
||||
) -> "AVGemmaTextEncoderModel":
|
||||
) -> "GemmaTextEncoder":
|
||||
"""Load the Gemma text encoder.
|
||||
Args:
|
||||
checkpoint_path: Path to the LTX-2 safetensors checkpoint file
|
||||
@@ -203,7 +203,7 @@ def load_text_encoder(
|
||||
When True, the model is loaded with device_map="auto" and the device argument
|
||||
is ignored for the Gemma backbone (feature extractor still uses dtype).
|
||||
Returns:
|
||||
Loaded AVGemmaTextEncoderModel
|
||||
Loaded GemmaTextEncoder (unified encoder handling V1/V2/V3)
|
||||
"""
|
||||
if not Path(gemma_model_path).is_dir():
|
||||
raise ValueError(f"Gemma model path is not a directory: {gemma_model_path}")
|
||||
@@ -216,12 +216,12 @@ def load_text_encoder(
|
||||
|
||||
# Standard loading path
|
||||
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
|
||||
from ltx_core.text_encoders.gemma.encoders.av_encoder import (
|
||||
from ltx_core.text_encoders.gemma import (
|
||||
AV_GEMMA_TEXT_ENCODER_KEY_OPS,
|
||||
GEMMA_MODEL_OPS,
|
||||
AVGemmaTextEncoderModelConfigurator,
|
||||
GemmaTextEncoderConfigurator,
|
||||
module_ops_from_gemma_root,
|
||||
)
|
||||
from ltx_core.text_encoders.gemma.encoders.base_encoder import module_ops_from_gemma_root
|
||||
from ltx_core.utils import find_matching_file
|
||||
|
||||
torch_device = _to_torch_device(device)
|
||||
@@ -231,7 +231,7 @@ def load_text_encoder(
|
||||
|
||||
text_encoder = SingleGPUModelBuilder(
|
||||
model_path=(str(checkpoint_path), *gemma_weight_paths),
|
||||
model_class_configurator=AVGemmaTextEncoderModelConfigurator,
|
||||
model_class_configurator=GemmaTextEncoderConfigurator,
|
||||
model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS,
|
||||
module_ops=(GEMMA_MODEL_OPS, *module_ops_from_gemma_root(str(gemma_model_path))),
|
||||
).build(device=torch_device, dtype=dtype)
|
||||
@@ -253,7 +253,7 @@ class LtxModelComponents:
|
||||
video_vae_decoder: "VideoDecoder | None" = None
|
||||
audio_vae_decoder: "AudioDecoder | None" = None
|
||||
vocoder: "Vocoder | None" = None
|
||||
text_encoder: "AVGemmaTextEncoderModel | None" = None
|
||||
text_encoder: "GemmaTextEncoder | None" = None
|
||||
scheduler: "LTX2Scheduler | None" = None
|
||||
|
||||
|
||||
|
||||
@@ -45,29 +45,61 @@ class UniformTimestepSampler(TimestepSampler):
|
||||
return self.sample(batch.shape[0], device=batch.device)
|
||||
|
||||
|
||||
class ShiftedLogitNormalTimestepSampler:
|
||||
class ShiftedLogitNormalTimestepSampler(TimestepSampler):
|
||||
"""
|
||||
Samples timesteps from a shifted logit-normal distribution,
|
||||
Samples timesteps from a stretched shifted logit-normal distribution,
|
||||
where the shift is determined by the sequence length.
|
||||
The stretching normalizes samples between percentile bounds to ensure
|
||||
the distribution covers [0, 1] more evenly. A uniform fallback prevents
|
||||
collapse at high token counts.
|
||||
"""
|
||||
|
||||
def __init__(self, std: float = 1.0):
|
||||
def __init__(self, std: float = 1.0, eps: float = 1e-3, uniform_prob: float = 0.1):
|
||||
self.std = std
|
||||
self.eps = eps
|
||||
self.uniform_prob = uniform_prob
|
||||
# Percentile values for stretching (scaled by std)
|
||||
# 99.9th percentile of standard normal ≈ 3.0902
|
||||
# 0.5th percentile of standard normal ≈ -2.5758
|
||||
self.normal_999_percentile = 3.0902 * std
|
||||
self.normal_005_percentile = -2.5758 * std
|
||||
|
||||
def sample(self, batch_size: int, seq_length: int, device: torch.device = None) -> torch.Tensor:
|
||||
"""Sample timesteps for a batch from a shifted logit-normal distribution.
|
||||
"""Sample timesteps for a batch from a stretched shifted logit-normal distribution.
|
||||
Args:
|
||||
batch_size: Number of timesteps to sample
|
||||
seq_length: Length of the sequence being processed, used to determine the shift
|
||||
device: Device to place the samples on
|
||||
Returns:
|
||||
Tensor of shape (batch_size,) containing timesteps sampled from a shifted
|
||||
logit-normal distribution, where the shift is determined by seq_length
|
||||
Tensor of shape (batch_size,) containing timesteps sampled from a stretched
|
||||
shifted logit-normal distribution, where the shift is determined by seq_length
|
||||
"""
|
||||
shift = self._get_shift_for_sequence_length(seq_length)
|
||||
normal_samples = torch.randn((batch_size,), device=device) * self.std + shift
|
||||
timesteps = torch.sigmoid(normal_samples)
|
||||
return timesteps
|
||||
mu = self._get_shift_for_sequence_length(seq_length)
|
||||
|
||||
# Sample from shifted logit-normal
|
||||
normal_samples = torch.randn((batch_size,), device=device) * self.std + mu
|
||||
logitnormal_samples = torch.sigmoid(normal_samples)
|
||||
|
||||
# Compute percentile bounds for stretching
|
||||
percentile_999 = torch.sigmoid(torch.tensor(mu + self.normal_999_percentile, device=device))
|
||||
percentile_005 = torch.sigmoid(torch.tensor(mu + self.normal_005_percentile, device=device))
|
||||
|
||||
# Stretch to [0, 1] range by normalizing between percentiles
|
||||
zero_terminal_raw = (logitnormal_samples - percentile_005) / (percentile_999 - percentile_005)
|
||||
|
||||
# Reflect small values around eps for numerical stability
|
||||
stretched_logit = torch.where(
|
||||
zero_terminal_raw >= self.eps,
|
||||
zero_terminal_raw,
|
||||
2 * self.eps - zero_terminal_raw,
|
||||
)
|
||||
stretched_logit = torch.clamp(stretched_logit, 0, 1)
|
||||
|
||||
# Mix with uniform samples (uniform_prob of the time)
|
||||
uniform = (1 - self.eps) * torch.rand((batch_size,), device=device) + self.eps
|
||||
prob = torch.rand((batch_size,), device=device)
|
||||
|
||||
return torch.where(prob > self.uniform_prob, stretched_logit, uniform)
|
||||
|
||||
def sample_for(self, batch: torch.Tensor) -> torch.Tensor:
|
||||
"""Sample timesteps for a specific batch tensor.
|
||||
|
||||
@@ -309,9 +309,22 @@ class LtxvTrainer:
|
||||
"""Perform a single training step using the configured strategy."""
|
||||
# Apply embedding connectors to transform pre-computed text embeddings
|
||||
conditions = batch["conditions"]
|
||||
video_embeds, audio_embeds, attention_mask = self._text_encoder._run_connectors(
|
||||
conditions["prompt_embeds"], conditions["prompt_attention_mask"]
|
||||
|
||||
if "video_prompt_embeds" in conditions:
|
||||
# New format: separate video/audio features from precompute()
|
||||
video_features = conditions["video_prompt_embeds"]
|
||||
audio_features = conditions.get("audio_prompt_embeds")
|
||||
else:
|
||||
# Legacy format: single prompt_embeds tensor — duplicate for both modalities
|
||||
video_features = conditions["prompt_embeds"]
|
||||
audio_features = conditions["prompt_embeds"]
|
||||
|
||||
mask = conditions["prompt_attention_mask"]
|
||||
additive_mask = self._text_encoder._convert_to_additive_mask(mask, video_features.dtype)
|
||||
video_embeds, audio_embeds, attention_mask = self._text_encoder.embeddings_processor.create_embeddings(
|
||||
video_features, audio_features, additive_mask
|
||||
)
|
||||
|
||||
conditions["video_prompt_embeds"] = video_embeds
|
||||
conditions["audio_prompt_embeds"] = audio_embeds
|
||||
conditions["prompt_attention_mask"] = attention_mask
|
||||
@@ -375,7 +388,7 @@ class LtxvTrainer:
|
||||
# Unload heavy components to free VRAM, keeping only the embedding connectors
|
||||
self._text_encoder.model = None
|
||||
self._text_encoder.tokenizer = None
|
||||
self._text_encoder.feature_extractor_linear = None
|
||||
self._text_encoder.feature_extractor = None
|
||||
|
||||
logger.debug("Validation prompt embeddings cached. Gemma model unloaded")
|
||||
return cached_embeddings
|
||||
@@ -822,7 +835,7 @@ class LtxvTrainer:
|
||||
output_path=output_path,
|
||||
fps=self._config.validation.frame_rate,
|
||||
audio=audio,
|
||||
audio_sample_rate=self._vocoder.output_sample_rate if audio is not None else None,
|
||||
audio_sample_rate=self._vocoder.output_sampling_rate if audio is not None else None,
|
||||
)
|
||||
video_paths.append(output_path)
|
||||
|
||||
|
||||
@@ -163,6 +163,7 @@ class TextToVideoStrategy(TrainingStrategy):
|
||||
# Create video Modality
|
||||
video_modality = Modality(
|
||||
enabled=True,
|
||||
sigma=sigmas,
|
||||
latent=noisy_video,
|
||||
timesteps=video_timesteps,
|
||||
positions=video_positions,
|
||||
@@ -254,6 +255,7 @@ class TextToVideoStrategy(TrainingStrategy):
|
||||
audio_modality = Modality(
|
||||
enabled=True,
|
||||
latent=noisy_audio,
|
||||
sigma=sigmas,
|
||||
timesteps=audio_timesteps,
|
||||
positions=audio_positions,
|
||||
context=audio_prompt_embeds,
|
||||
|
||||
@@ -210,6 +210,7 @@ class VideoToVideoStrategy(TrainingStrategy):
|
||||
video_modality = Modality(
|
||||
enabled=True,
|
||||
latent=combined_latents,
|
||||
sigma=sigmas,
|
||||
timesteps=timesteps,
|
||||
positions=positions,
|
||||
context=prompt_embeds,
|
||||
|
||||
@@ -36,7 +36,7 @@ if TYPE_CHECKING:
|
||||
from ltx_core.model.audio_vae import AudioDecoder, Vocoder
|
||||
from ltx_core.model.transformer import LTXModel
|
||||
from ltx_core.model.video_vae import VideoDecoder, VideoEncoder
|
||||
from ltx_core.text_encoders.gemma import AVGemmaTextEncoderModel
|
||||
from ltx_core.text_encoders.gemma import GemmaTextEncoder
|
||||
|
||||
VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
|
||||
|
||||
@@ -124,7 +124,7 @@ class ValidationSampler:
|
||||
transformer: "LTXModel",
|
||||
vae_decoder: "VideoDecoder",
|
||||
vae_encoder: "VideoEncoder | None",
|
||||
text_encoder: "AVGemmaTextEncoderModel | None" = None,
|
||||
text_encoder: "GemmaTextEncoder | None" = None,
|
||||
audio_decoder: "AudioDecoder | None" = None,
|
||||
vocoder: "Vocoder | None" = None,
|
||||
sampling_context: SamplingContext | None = None,
|
||||
@@ -497,6 +497,7 @@ class ValidationSampler:
|
||||
video = Modality(
|
||||
enabled=True,
|
||||
latent=video_state.latent,
|
||||
sigma=sigmas[0].repeat(video_state.latent.shape[0]),
|
||||
timesteps=video_state.denoise_mask,
|
||||
positions=video_state.positions,
|
||||
context=v_ctx_pos,
|
||||
@@ -509,6 +510,7 @@ class ValidationSampler:
|
||||
audio = Modality(
|
||||
enabled=True,
|
||||
latent=audio_state.latent,
|
||||
sigma=sigmas[0].repeat(audio_state.latent.shape[0]),
|
||||
timesteps=audio_state.denoise_mask,
|
||||
positions=audio_state.positions,
|
||||
context=a_ctx_pos,
|
||||
@@ -525,6 +527,7 @@ class ValidationSampler:
|
||||
video = replace(
|
||||
video,
|
||||
latent=video_state.latent,
|
||||
sigma=sigma.repeat(video_state.latent.shape[0]),
|
||||
timesteps=sigma * video_state.denoise_mask,
|
||||
positions=video_state.positions,
|
||||
)
|
||||
@@ -533,6 +536,7 @@ class ValidationSampler:
|
||||
audio = replace(
|
||||
audio,
|
||||
latent=audio_state.latent,
|
||||
sigma=sigma.repeat(audio_state.latent.shape[0]),
|
||||
timesteps=sigma * audio_state.denoise_mask,
|
||||
positions=audio_state.positions,
|
||||
)
|
||||
@@ -703,7 +707,8 @@ class ValidationSampler:
|
||||
# Move the base Gemma model to CPU but keep embeddings connectors on GPU
|
||||
# as this module is also used during training
|
||||
self._text_encoder.model.to("cpu")
|
||||
self._text_encoder.feature_extractor_linear.to("cpu")
|
||||
if self._text_encoder.feature_extractor is not None:
|
||||
self._text_encoder.feature_extractor.to("cpu")
|
||||
|
||||
return v_ctx_pos, a_ctx_pos, v_ctx_neg, a_ctx_neg
|
||||
|
||||
|
||||
Reference in New Issue
Block a user