LTX-2/packages/ltx-trainer/AGENTS.md

AGENTS.md

This file provides guidance to AI coding assistants (Claude, Cursor, etc.) when working with code in this repository.

Project Overview

LTX-2 Trainer is a training toolkit for fine-tuning the Lightricks LTX-2 audio-video generation model. 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

Key Dependencies:

• ltx-core - Core model implementations (transformer, VAE, text encoder)
• ltx-pipelines - Inference pipeline components

Important: This trainer only supports LTX-2 (the audio-video model). The older LTXV models are not supported.

Architecture Overview

Package Structure

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

Key Architectural Patterns

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.
• Combined loader: load_model() returns LtxModelComponents dataclass
• Uses SingleGPUModelBuilder from ltx-core internally

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

Model Interface (Modality-based):

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,
)
audio = Modality(
    enabled=True,
    latent=audio_latents,
    timesteps=audio_timesteps,
    positions=audio_positions,  # [B, 1, seq_len, 2]
    context=audio_embeds,
    context_mask=None,
)

# Forward pass returns predictions for both modalities
video_pred, audio_pred = model(video=video, audio=audio, perturbations=None)

Note: Modality is immutable (frozen dataclass). Use dataclasses.replace() to modify.

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 files in configs/ directory

Development Commands

Setup and Installation

# From the repository root
uv sync
cd packages/ltx-trainer

Code Quality

# Run ruff linting and formatting
uv run ruff check .
uv run ruff format .

# Run pre-commit checks
uv run pre-commit run --all-files

Running Tests

cd packages/ltx-trainer
uv run pytest

Running Training

# Single GPU
uv run python scripts/train.py configs/ltx2_av_lora.yaml

# Multi-GPU with Accelerate
uv run accelerate launch scripts/train.py configs/ltx2_av_lora.yaml

Code Standards

Type Hints

• Always use type hints for all function arguments and return values
• Use Python 3.10+ syntax: list[str] not List[str], str | Path not Union[str, Path]
• Use pathlib.Path for file operations

Class Methods

• Mark methods as @staticmethod if they don't access instance or class state
• Use @classmethod for alternative constructors

AI/ML Specific

• Use @torch.inference_mode() for inference (prefer over @torch.no_grad())
• Use accelerator.device for distributed compatibility
• Support mixed precision (bfloat16 via dtype parameters)
• Use gradient checkpointing for memory-intensive training

Logging

• Use from ltx_trainer import logger for all messages
• Avoid print statements in production code

Important Files & Modules

Configuration (CRITICAL)

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

⚠️ 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

• Implements distributed training with Accelerate
• Handles mixed precision, gradient accumulation, checkpointing
• Uses training strategies for mode-specific logic

src/ltx_trainer/training_strategies/ - Strategy pattern

• base_strategy.py: TrainingStrategy ABC, ModelInputs dataclass
• text_to_video.py: Standard text-to-video (with optional audio)
• 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
• 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_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

Data

src/ltx_trainer/datasets.py - Dataset handling

• PrecomputedDataset loads pre-computed VAE latents
• Supports video latents, audio latents, text embeddings, reference latents

Common Development Tasks

Adding a New Configuration Parameter

1. Add field to appropriate config class in src/ltx_trainer/config.py
2. Add validator if needed
3. Update ALL config files in configs/
4. Update docs/configuration-reference.md

Implementing a New Training Strategy

1. Create new file in src/ltx_trainer/training_strategies/
2. Create config class inheriting TrainingStrategyConfigBase
3. Create strategy class inheriting TrainingStrategy
4. Implement: get_data_sources(), prepare_training_inputs(), compute_loss()
5. Add to __init__.py: import, add to TrainingStrategyConfig union, update factory
6. Add discriminator tag to config.py's TrainingStrategyConfig
7. Create example config file in configs/

Working with Modalities

from dataclasses import replace
from ltx_core.model.transformer.modality import Modality

# Create modality
video = Modality(
    enabled=True,
    latent=latents,
    timesteps=timesteps,
    positions=positions,
    context=context,
    context_mask=None,
)

# Update (immutable - must use replace)
video = replace(video, latent=new_latent, timesteps=new_timesteps)

# Disable a modality
audio = replace(audio, enabled=False)

Debugging Tips

Training Issues:

• Check logs first (rich logger provides context)
• GPU memory: Look for OOM errors, enable enable_gradient_checkpointing: true
• Distributed training: Check accelerator.state and device placement

Model Loading:

• 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

Configuration:

• Validation errors: Check validators in config.py
• Unknown fields: Config uses extra="forbid" - all fields must be defined
• Strategy validation: IC-LoRA requires reference_videos in validation config

Key Constraints

LTX-2 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

Resolution Requirements

Width and height must be divisible by 32.

Model Paths

• Must be local paths (URLs not supported)
• model_path: Path to .safetensors checkpoint
• text_encoder_path: Path to Gemma model directory

Platform Requirements

• Linux required (uses triton which is Linux-only)
• CUDA GPU with 24GB+ VRAM recommended

Reference: ltx-core Key Components

packages/ltx-core/src/ltx_core/
├── model/
│   ├── transformer/
│   │   ├── model.py              # LTXModel
│   │   ├── modality.py           # Modality dataclass
│   │   └── transformer.py        # BasicAVTransformerBlock
│   ├── video_vae/
│   │   └── video_vae.py          # Encoder, Decoder
│   ├── 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)