Automated PR - 2026-05-11
This commit is contained in:
github-actions[bot]
@@ -166,6 +166,11 @@ acceleration:
|
||||
# Useful when GPU memory is limited
|
||||
load_text_encoder_in_8bit: false
|
||||
|
||||
# Offload optimizer state to CPU during validation video sampling.
|
||||
# Helps avoid OOM when VAE decoder + transformer + optimizer state can't coexist
|
||||
# on the GPU (full fine-tune, high-rank LoRA). No effect for FSDP.
|
||||
offload_optimizer_during_validation: false
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Data Configuration
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
@@ -178,6 +178,11 @@ acceleration:
|
||||
# Useful when GPU memory is limited
|
||||
load_text_encoder_in_8bit: true
|
||||
|
||||
# Offload optimizer state to CPU during validation video sampling.
|
||||
# Helps avoid OOM when VAE decoder + transformer + optimizer state can't coexist
|
||||
# on the GPU (full fine-tune, high-rank LoRA). No effect for FSDP.
|
||||
offload_optimizer_during_validation: true
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Data Configuration
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
@@ -166,6 +166,11 @@ acceleration:
|
||||
# Useful when GPU memory is limited
|
||||
load_text_encoder_in_8bit: false
|
||||
|
||||
# Offload optimizer state to CPU during validation video sampling.
|
||||
# Helps avoid OOM when VAE decoder + transformer + optimizer state can't coexist
|
||||
# on the GPU (full fine-tune, high-rank LoRA). No effect for FSDP.
|
||||
offload_optimizer_during_validation: false
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Data Configuration
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
@@ -215,18 +215,20 @@ Hardware acceleration and compute optimization settings.
|
||||
|
||||
```yaml
|
||||
acceleration:
|
||||
mixed_precision_mode: "bf16" # "no", "fp16", or "bf16"
|
||||
quantization: null # Quantization options
|
||||
load_text_encoder_in_8bit: false # Load text encoder in 8-bit
|
||||
mixed_precision_mode: "bf16" # "no", "fp16", or "bf16"
|
||||
quantization: null # Quantization options
|
||||
load_text_encoder_in_8bit: false # Load text encoder in 8-bit
|
||||
offload_optimizer_during_validation: false # Offload optimizer state to CPU during validation
|
||||
```
|
||||
|
||||
**Key parameters:**
|
||||
|
||||
| Parameter | Description |
|
||||
|-----------------------------|------------------------------------------------------------------------------------|
|
||||
| `mixed_precision_mode` | Precision mode - `"bf16"` recommended for modern GPUs |
|
||||
| `quantization` | Model quantization: `null`, `"int8-quanto"`, `"int4-quanto"`, `"fp8-quanto"`, etc. |
|
||||
| `load_text_encoder_in_8bit` | Load the Gemma text encoder in 8-bit to save GPU memory |
|
||||
| Parameter | Description |
|
||||
|---------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| `mixed_precision_mode` | Precision mode - `"bf16"` recommended for modern GPUs |
|
||||
| `quantization` | Model quantization: `null`, `"int8-quanto"`, `"int4-quanto"`, `"fp8-quanto"`, etc. |
|
||||
| `load_text_encoder_in_8bit` | Load the Gemma text encoder in 8-bit to save GPU memory |
|
||||
| `offload_optimizer_during_validation` | Move optimizer state to CPU before validation video sampling and back afterwards. Useful when validation OOMs because VAE decoder + transformer + optimizer state can't coexist on the GPU (full fine-tune, high-rank LoRA). No effect for FSDP. |
|
||||
|
||||
### DataConfig
|
||||
|
||||
|
||||
@@ -50,17 +50,20 @@ This will create a `dataset.json` file containing video paths and their captions
|
||||
|
||||
**Captioning options:**
|
||||
|
||||
| Option | Description |
|
||||
|--------|-------------|
|
||||
| `--captioner-type` | `qwen_omni` (default, local) or `gemini_flash` (API) |
|
||||
| `--use-8bit` | Enable 8-bit quantization for lower VRAM usage |
|
||||
| `--no-audio` | Disable audio processing (video-only captions) |
|
||||
| `--override` | Re-caption files that already have captions |
|
||||
| `--api-key` | API key for Gemini Flash (or set `GOOGLE_API_KEY` env var) |
|
||||
|
||||
| Option | Description |
|
||||
| ------------------ | ---------------------------------------------------------- |
|
||||
| `--captioner-type` | `qwen_omni` (default, local) or `gemini_flash` (API) |
|
||||
| `--use-8bit` | Enable 8-bit quantization for lower VRAM usage |
|
||||
| `--no-audio` | Disable audio processing (video-only captions) |
|
||||
| `--override` | Re-caption files that already have captions |
|
||||
| `--api-key` | API key for Gemini Flash (or set `GOOGLE_API_KEY` env var) |
|
||||
|
||||
|
||||
**Caption format:**
|
||||
|
||||
The captioner produces structured captions with sections for:
|
||||
|
||||
- **Visual content**: People, objects, actions, settings, colors, movements
|
||||
- **Speech transcription**: Word-for-word transcription of spoken content
|
||||
- **Sounds**: Music, ambient sounds, sound effects
|
||||
@@ -106,15 +109,57 @@ uv run python scripts/process_dataset.py dataset.json \
|
||||
--with-audio
|
||||
```
|
||||
|
||||
### 🚀 Multi-GPU Preprocessing
|
||||
|
||||
Preprocessing large datasets can take a while. To run it across multiple GPUs in parallel, wrap the command with
|
||||
`accelerate launch` (for example `--num_processes 4`). Each process handles an interleaved slice of the dataset.
|
||||
The same approach applies to `process_videos.py` and `process_captions.py` when you run them standalone.
|
||||
|
||||
```bash
|
||||
uv run accelerate launch --num_processes 4 scripts/process_dataset.py dataset.json \
|
||||
--resolution-buckets "960x544x49" \
|
||||
--model-path /path/to/ltx-2-model.safetensors \
|
||||
--text-encoder-path /path/to/gemma-model
|
||||
```
|
||||
|
||||
Outputs are written atomically (via a per-process temporary file, then renamed), so an interrupted run leaves no
|
||||
corrupt files. By default a rerun **resumes** — items whose output `.pt` already exists are skipped.
|
||||
|
||||
> [!IMPORTANT]
|
||||
> Pass `**--overwrite`** when rerunning with changed parameters (different model checkpoint, resolution buckets,
|
||||
> text encoder, `--lora-trigger`, etc.). Without it the script keeps the stale outputs from the previous run.
|
||||
>
|
||||
> ```bash
|
||||
> uv run accelerate launch --num_processes 4 scripts/process_dataset.py dataset.json \
|
||||
> --resolution-buckets "960x544x49" \
|
||||
> --model-path /path/to/ltx-2.3-model.safetensors \
|
||||
> --text-encoder-path /path/to/gemma-model \
|
||||
> --overwrite
|
||||
> ```
|
||||
|
||||
### 📊 Dataset Format
|
||||
|
||||
The trainer supports either videos or single images.
|
||||
Note that your dataset must be homogeneous - either all videos or all images, mixing is not supported.
|
||||
The trainer supports videos, single images, or a mix of both in the same dataset.
|
||||
|
||||
> [!TIP]
|
||||
> **Image Datasets:** When using images, follow the same preprocessing steps and format requirements as with videos,
|
||||
> but use `1` for the frame count in the resolution bucket (e.g., `960x544x1`).
|
||||
|
||||
> [!NOTE]
|
||||
> **Mixed image + video datasets:** Mixing stills and videos in a single dataset is supported, but requires some care:
|
||||
>
|
||||
> - Preprocess with **multiple resolution buckets** covering both frame counts — e.g.
|
||||
> `--resolution-buckets "960x544x1;960x544x49"`. Images are automatically assigned to the `F=1` bucket and
|
||||
> videos to an `F>1` bucket.
|
||||
> - You **must** set `optimization.batch_size: 1` in your training config (see the warning under
|
||||
> [Resolution Buckets](#-resolution-buckets)), since samples with different shapes cannot be collated into a
|
||||
> single batch. Use `gradient_accumulation_steps` if you need a larger effective batch.
|
||||
> - Per-step cost differs substantially between a single-frame sample and a many-frame sample, which can lead to
|
||||
> uneven gradient magnitudes across steps. Consider weighting the two subsets or tuning the learning rate if
|
||||
> you observe instability.
|
||||
> - If you prefer a fully officially-supported path, train two separate LoRAs (one on stills, one on video) and
|
||||
> stack them at inference.
|
||||
|
||||
The dataset must be a CSV, JSON, or JSONL metadata file with columns for captions and video paths:
|
||||
|
||||
**JSON format example:**
|
||||
@@ -197,6 +242,7 @@ uv run python scripts/process_dataset.py dataset.json \
|
||||
> ```
|
||||
>
|
||||
> Where:
|
||||
>
|
||||
> - H = Height of video
|
||||
> - W = Width of video
|
||||
> - F = Number of frames
|
||||
@@ -204,6 +250,7 @@ uv run python scripts/process_dataset.py dataset.json \
|
||||
> - 8 = VAE's temporal downsampling factor
|
||||
>
|
||||
> For example, a 768×448×89 video would have sequence length:
|
||||
>
|
||||
> ```
|
||||
> (768/32) * (448/32) * ((89-1)/8 + 1) = 24 * 14 * 12 = 4,032
|
||||
> ```
|
||||
@@ -268,7 +315,6 @@ uv run python scripts/process_dataset.py dataset.json \
|
||||
|
||||
This will create an additional `reference_latents/` directory containing the preprocessed reference video latents.
|
||||
|
||||
|
||||
### Generating Reference Videos
|
||||
|
||||
**Dataset Requirements for IC-LoRA:**
|
||||
@@ -277,7 +323,7 @@ This will create an additional `reference_latents/` directory containing the pre
|
||||
- Reference and target videos must have *identical* resolution and length
|
||||
- Both reference and target videos should be preprocessed together using the same resolution buckets
|
||||
|
||||
We provide an example script, [`scripts/compute_reference.py`](../scripts/compute_reference.py), to generate reference
|
||||
We provide an example script, `[scripts/compute_reference.py](../scripts/compute_reference.py)`, to generate reference
|
||||
videos for a given dataset. The default implementation generates Canny edge reference videos.
|
||||
|
||||
```bash
|
||||
@@ -293,7 +339,6 @@ If you want to generate a different type of condition (depth maps, pose skeleton
|
||||
|
||||
For reference, see our **[Canny Control Dataset](https://huggingface.co/datasets/Lightricks/Canny-Control-Dataset)** which demonstrates proper IC-LoRA dataset structure with paired videos and Canny edge maps.
|
||||
|
||||
|
||||
## 🎯 LoRA Trigger Words
|
||||
|
||||
When training a LoRA, you can specify a trigger token that will be prepended to all captions:
|
||||
|
||||
@@ -84,6 +84,20 @@ optimization:
|
||||
optimizer_type: "adamw8bit"
|
||||
```
|
||||
|
||||
#### 7. Offload Optimizer State During Validation
|
||||
|
||||
If you OOM specifically during validation video sampling — typically in
|
||||
full fine-tunes or high-rank LoRA runs where AdamW state and the VAE decoder
|
||||
can't coexist on the GPU — offload optimizer state to CPU during sampling:
|
||||
|
||||
```yaml
|
||||
acceleration:
|
||||
offload_optimizer_during_validation: true
|
||||
```
|
||||
|
||||
The offload + reload happens once per validation interval, not per step.
|
||||
No effect for FSDP (sharded state).
|
||||
|
||||
---
|
||||
|
||||
## ⚠️ Common Usage Issues
|
||||
|
||||
@@ -143,6 +143,27 @@ uv run python scripts/process_dataset.py dataset.json \
|
||||
> [!NOTE]
|
||||
> When training with multiple resolution buckets, set `optimization.batch_size: 1`.
|
||||
|
||||
**Multi-GPU preprocessing.** Launch with `accelerate launch` to shard the dataset across processes. Reruns resume
|
||||
by default (existing `.pt` outputs are skipped); writes are atomic so interrupted runs are safe. Pass `--overwrite`
|
||||
when rerunning with changed parameters (different model, resolution buckets, text encoder, `--lora-trigger`, etc.)
|
||||
so stale outputs are replaced. Use the same `accelerate launch` pattern (and `--overwrite` when needed) with
|
||||
`process_videos.py` or `process_captions.py` when you run those scripts standalone.
|
||||
|
||||
```bash
|
||||
# Multi-GPU preprocessing
|
||||
uv run accelerate launch --num_processes 4 scripts/process_dataset.py dataset.json \
|
||||
--resolution-buckets "960x544x49" \
|
||||
--model-path /path/to/ltx-2-model.safetensors \
|
||||
--text-encoder-path /path/to/gemma-model
|
||||
|
||||
# Force re-encoding of all items (e.g. after switching model or resolution)
|
||||
uv run accelerate launch --num_processes 4 scripts/process_dataset.py dataset.json \
|
||||
--resolution-buckets "960x544x49" \
|
||||
--model-path /path/to/ltx-2.3-model.safetensors \
|
||||
--text-encoder-path /path/to/gemma-model \
|
||||
--overwrite
|
||||
```
|
||||
|
||||
For detailed usage, see the [Dataset Preparation Guide](dataset-preparation.md).
|
||||
|
||||
### Reference Video Generation
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "ltx-trainer"
|
||||
version = "1.1.2"
|
||||
version = "1.1.3"
|
||||
description = "LTX-2 training, democratized."
|
||||
readme = "README.md"
|
||||
authors = [
|
||||
@@ -48,7 +48,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
|
||||
[tool.ruff]
|
||||
target-version = "1.1.2"
|
||||
target-version = "1.1.3"
|
||||
line-length = 120
|
||||
|
||||
[tool.ruff.lint]
|
||||
|
||||
@@ -13,12 +13,14 @@ Can be used as a standalone script:
|
||||
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pandas as pd
|
||||
import torch
|
||||
import typer
|
||||
from accelerate import PartialState
|
||||
from rich.console import Console
|
||||
from rich.progress import (
|
||||
BarColumn,
|
||||
@@ -30,7 +32,7 @@ from rich.progress import (
|
||||
TimeElapsedColumn,
|
||||
TimeRemainingColumn,
|
||||
)
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from torch.utils.data import DataLoader, Dataset, Subset
|
||||
from transformers.utils.logging import disable_progress_bar
|
||||
|
||||
from ltx_trainer import logger
|
||||
@@ -232,9 +234,14 @@ def compute_captions_embeddings( # noqa: PLR0913
|
||||
batch_size: int = 8,
|
||||
device: str = "cuda",
|
||||
load_in_8bit: bool = False,
|
||||
overwrite: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Process captions and save text embeddings.
|
||||
Under ``accelerate launch``, each process handles an interleaved shard of
|
||||
the dataset (rank/world read from ``accelerate.PartialState``). Already-
|
||||
computed ``.pt`` outputs are skipped unless ``overwrite=True``; writes are
|
||||
atomic so an interrupted run is safe to resume.
|
||||
Args:
|
||||
dataset_file: Path to metadata file (CSV/JSON/JSONL) containing captions and media paths
|
||||
output_dir: Directory to save embeddings
|
||||
@@ -247,11 +254,12 @@ def compute_captions_embeddings( # noqa: PLR0913
|
||||
batch_size: Batch size for processing
|
||||
device: Device to use for computation
|
||||
load_in_8bit: Whether to load the Gemma text encoder in 8-bit precision
|
||||
overwrite: Re-encode every item even if its output exists. Use when rerunning with
|
||||
changed parameters (different text encoder, lora_trigger, etc.) so stale
|
||||
outputs are replaced.
|
||||
"""
|
||||
|
||||
console = Console()
|
||||
|
||||
# Create dataset
|
||||
dataset = CaptionsDataset(
|
||||
dataset_file=dataset_file,
|
||||
caption_column=caption_column,
|
||||
@@ -264,6 +272,24 @@ def compute_captions_embeddings( # noqa: PLR0913
|
||||
output_path = Path(output_dir)
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# TODO(batch-tokenization): The current Gemma tokenizer doesn't support batched tokenization.
|
||||
if batch_size > 1:
|
||||
logger.warning(
|
||||
"Batch size greater than 1 is not currently supported with the Gemma tokenizer. "
|
||||
"Overriding batch_size to 1. This will be fixed in a future update."
|
||||
)
|
||||
batch_size = 1
|
||||
|
||||
dataloader = _build_sharded_dataloader(
|
||||
dataset,
|
||||
batch_size=batch_size,
|
||||
num_workers=2,
|
||||
is_done=lambda idx: (output_path / dataset.output_paths[idx]).is_file(),
|
||||
overwrite=overwrite,
|
||||
)
|
||||
if dataloader is None:
|
||||
return
|
||||
|
||||
# Load text encoder and embeddings processor
|
||||
with console.status("[bold]Loading Gemma text encoder...", spinner="dots"):
|
||||
text_encoder = load_text_encoder(
|
||||
@@ -279,21 +305,7 @@ def compute_captions_embeddings( # noqa: PLR0913
|
||||
)
|
||||
|
||||
logger.info("Text encoder and embeddings processor loaded successfully")
|
||||
|
||||
# TODO(batch-tokenization): The current Gemma tokenizer doesn't support batched tokenization.
|
||||
if batch_size > 1:
|
||||
logger.warning(
|
||||
"Batch size greater than 1 is not currently supported with the Gemma tokenizer. "
|
||||
"Overriding batch_size to 1. This will be fixed in a future update."
|
||||
)
|
||||
batch_size = 1
|
||||
|
||||
# Create dataloader
|
||||
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=2)
|
||||
|
||||
# Process batches
|
||||
total_batches = len(dataloader)
|
||||
logger.info(f"Processing captions in {total_batches:,} batches...")
|
||||
logger.info(f"Processing captions in {len(dataloader):,} batches...")
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
@@ -333,11 +345,44 @@ def compute_captions_embeddings( # noqa: PLR0913
|
||||
embedding_data["audio_prompt_embeds"] = audio_prompt_embeds[0].cpu().contiguous()
|
||||
|
||||
output_file = output_path / output_rel_path
|
||||
torch.save(embedding_data, output_file)
|
||||
_atomic_save(embedding_data, output_file)
|
||||
|
||||
progress.advance(task)
|
||||
|
||||
logger.info(f"Processed {len(dataset):,} captions. Embeddings saved to {output_path}")
|
||||
logger.info(f"Processed {len(dataloader.dataset):,} captions -> {output_path}") # type: ignore[arg-type]
|
||||
|
||||
|
||||
def _atomic_save(data: Any, out: Path) -> None: # noqa: ANN401
|
||||
"""Save to ``out`` atomically via per-PID temp file + replace.
|
||||
Crash mid-write leaves an orphan ``.tmp.<pid>`` file that the skip logic
|
||||
ignores. The per-PID suffix makes concurrent writes from multiple ranks
|
||||
collision-free.
|
||||
"""
|
||||
tmp = out.with_suffix(f"{out.suffix}.tmp.{os.getpid()}")
|
||||
torch.save(data, tmp)
|
||||
tmp.replace(out)
|
||||
|
||||
|
||||
def _build_sharded_dataloader(
|
||||
dataset: Dataset,
|
||||
*,
|
||||
batch_size: int,
|
||||
num_workers: int,
|
||||
is_done: Callable[[int], bool],
|
||||
overwrite: bool,
|
||||
) -> DataLoader | None:
|
||||
"""Return a DataLoader over this rank's interleaved shard of ``dataset``.
|
||||
When ``overwrite`` is False, items whose outputs already exist (per
|
||||
``is_done``) are filtered out. Returns ``None`` if this rank has nothing
|
||||
to do, so the caller can early-return without loading any models.
|
||||
"""
|
||||
state = PartialState()
|
||||
todo = [i for i in range(state.process_index, len(dataset), state.num_processes) if overwrite or not is_done(i)]
|
||||
if not todo:
|
||||
logger.info(f"Rank {state.process_index}/{state.num_processes}: nothing to do")
|
||||
return None
|
||||
logger.info(f"Rank {state.process_index}/{state.num_processes}: processing {len(todo):,} of {len(dataset):,} items")
|
||||
return DataLoader(Subset(dataset, todo), batch_size=batch_size, shuffle=False, num_workers=num_workers)
|
||||
|
||||
|
||||
@app.command()
|
||||
@@ -387,8 +432,15 @@ def main( # noqa: PLR0913
|
||||
default=False,
|
||||
help="Load the Gemma text encoder in 8-bit precision to save GPU memory (requires bitsandbytes)",
|
||||
),
|
||||
overwrite: bool = typer.Option(
|
||||
default=False,
|
||||
help="Re-encode every caption even if its output exists. Use when rerunning with "
|
||||
"changed parameters (different text encoder, lora_trigger, etc.) so stale outputs are replaced.",
|
||||
),
|
||||
) -> None:
|
||||
"""Process text captions and save embeddings for video generation training.
|
||||
For multi-GPU preprocessing, invoke under ``accelerate launch`` - each process
|
||||
will handle an interleaved shard of the dataset.
|
||||
This script processes captions from metadata files and saves text embeddings
|
||||
that can be used for training video generation models. The output embeddings
|
||||
will maintain the same folder structure and naming as the corresponding media files.
|
||||
@@ -428,6 +480,7 @@ def main( # noqa: PLR0913
|
||||
batch_size=batch_size,
|
||||
device=device,
|
||||
load_in_8bit=load_text_encoder_in_8bit,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -50,6 +50,7 @@ def preprocess_dataset( # noqa: PLR0913
|
||||
reference_downscale_factor: int = 1,
|
||||
with_audio: bool = False,
|
||||
load_text_encoder_in_8bit: bool = False,
|
||||
overwrite: bool = False,
|
||||
) -> None:
|
||||
"""Run the preprocessing pipeline with the given arguments."""
|
||||
# Validate dataset file
|
||||
@@ -77,6 +78,7 @@ def preprocess_dataset( # noqa: PLR0913
|
||||
batch_size=batch_size,
|
||||
device=device,
|
||||
load_in_8bit=load_text_encoder_in_8bit,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
|
||||
# Process videos using the dedicated function
|
||||
@@ -97,6 +99,7 @@ def preprocess_dataset( # noqa: PLR0913
|
||||
vae_tiling=vae_tiling,
|
||||
with_audio=with_audio,
|
||||
audio_output_dir=str(audio_latents_dir) if audio_latents_dir else None,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
|
||||
# Process reference videos if reference_column is provided
|
||||
@@ -133,6 +136,7 @@ def preprocess_dataset( # noqa: PLR0913
|
||||
batch_size=batch_size,
|
||||
device=device,
|
||||
vae_tiling=vae_tiling,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
|
||||
# Handle decoding if requested (for verification)
|
||||
@@ -252,8 +256,15 @@ def main( # noqa: PLR0913
|
||||
help="Downscale factor for reference video resolution. When > 1, reference videos are processed at "
|
||||
"1/n resolution (e.g., 2 means half resolution). Used for efficient IC-LoRA training.",
|
||||
),
|
||||
overwrite: bool = typer.Option(
|
||||
default=False,
|
||||
help="Re-compute every item even if its output exists. Use when rerunning with "
|
||||
"changed parameters (different model, resolution, etc.) so stale outputs are replaced.",
|
||||
),
|
||||
) -> None:
|
||||
"""Preprocess a video dataset by computing and saving latents and text embeddings.
|
||||
For multi-GPU preprocessing, invoke under ``accelerate launch`` - each process
|
||||
will handle an interleaved shard of the dataset.
|
||||
The dataset must be a CSV, JSON, or JSONL file with columns for captions and video paths.
|
||||
This script is designed for LTX-2 models which use the Gemma text encoder.
|
||||
Examples:
|
||||
@@ -310,6 +321,7 @@ def main( # noqa: PLR0913
|
||||
reference_downscale_factor=reference_downscale_factor,
|
||||
with_audio=with_audio,
|
||||
load_text_encoder_in_8bit=load_text_encoder_in_8bit,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -14,6 +14,8 @@ Can be used as a standalone script:
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -22,6 +24,7 @@ import pandas as pd
|
||||
import torch
|
||||
import torchaudio
|
||||
import typer
|
||||
from accelerate import PartialState
|
||||
from pillow_heif import register_heif_opener
|
||||
from rich.console import Console
|
||||
from rich.progress import (
|
||||
@@ -34,7 +37,7 @@ from rich.progress import (
|
||||
TimeElapsedColumn,
|
||||
TimeRemainingColumn,
|
||||
)
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from torch.utils.data import DataLoader, Dataset, Subset
|
||||
from torchvision import transforms
|
||||
from torchvision.transforms import InterpolationMode
|
||||
from torchvision.transforms.functional import crop, resize, to_tensor
|
||||
@@ -444,9 +447,14 @@ def compute_latents( # noqa: PLR0913, PLR0915
|
||||
vae_tiling: bool = False,
|
||||
with_audio: bool = False,
|
||||
audio_output_dir: str | None = None,
|
||||
overwrite: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Process videos and save latent representations.
|
||||
Under ``accelerate launch``, each process handles an interleaved shard of
|
||||
the dataset (rank/world read from ``accelerate.PartialState``). Already-
|
||||
computed ``.pt`` outputs are skipped unless ``overwrite=True``; writes are
|
||||
atomic so an interrupted run is safe to resume.
|
||||
Args:
|
||||
dataset_file: Path to metadata file (CSV/JSON/JSONL) containing video paths
|
||||
video_column: Column name for video paths in the metadata file
|
||||
@@ -460,15 +468,15 @@ def compute_latents( # noqa: PLR0913, PLR0915
|
||||
vae_tiling: Whether to enable VAE tiling
|
||||
with_audio: Whether to extract and encode audio from videos
|
||||
audio_output_dir: Directory to save audio latents (required if with_audio=True)
|
||||
overwrite: Re-process every item even if its output exists. Use when rerunning with
|
||||
changed parameters (different model, resolution, etc.) so stale outputs are replaced.
|
||||
"""
|
||||
# Validate audio parameters
|
||||
if with_audio and audio_output_dir is None:
|
||||
raise ValueError("audio_output_dir must be provided when with_audio=True")
|
||||
|
||||
console = Console()
|
||||
torch_device = torch.device(device)
|
||||
|
||||
# Create dataset
|
||||
dataset = MediaDataset(
|
||||
dataset_file=dataset_file,
|
||||
main_media_column=main_media_column or video_column,
|
||||
@@ -481,13 +489,34 @@ def compute_latents( # noqa: PLR0913, PLR0915
|
||||
|
||||
output_path = Path(output_dir)
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Set up audio output directory if needed
|
||||
audio_output_path = None
|
||||
audio_output_path: Path | None = None
|
||||
if with_audio:
|
||||
audio_output_path = Path(audio_output_dir)
|
||||
audio_output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Audio processing requires batch_size=1; must be applied before the dataloader is built.
|
||||
if with_audio and batch_size > 1:
|
||||
logger.warning("Audio processing requires batch_size=1. Overriding batch_size to 1.")
|
||||
batch_size = 1
|
||||
|
||||
data_root = dataset.dataset_file.parent
|
||||
|
||||
def _is_done(idx: int) -> bool:
|
||||
rel = dataset.main_media_paths[idx].relative_to(data_root).with_suffix(".pt")
|
||||
if not (output_path / rel).is_file():
|
||||
return False
|
||||
return audio_output_path is None or (audio_output_path / rel).is_file()
|
||||
|
||||
dataloader = _build_sharded_dataloader(
|
||||
dataset,
|
||||
batch_size=batch_size,
|
||||
num_workers=4,
|
||||
is_done=_is_done,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
if dataloader is None:
|
||||
return
|
||||
|
||||
# Load video VAE encoder
|
||||
with console.status(f"[bold]Loading video VAE encoder from [cyan]{model_path}[/]...", spinner="dots"):
|
||||
vae = load_video_vae_encoder(model_path, device=torch_device, dtype=torch.bfloat16)
|
||||
@@ -510,14 +539,6 @@ def compute_latents( # noqa: PLR0913, PLR0915
|
||||
n_fft=audio_vae_encoder.n_fft,
|
||||
).to(torch_device)
|
||||
|
||||
# Create dataloader
|
||||
# Note: batch_size=1 required when with_audio because audio extraction can fail for some videos,
|
||||
# and the default collate function can't handle mixed None/dict values across a batch.
|
||||
if with_audio and batch_size > 1:
|
||||
logger.warning("Audio processing requires batch_size=1. Overriding batch_size to 1.")
|
||||
batch_size = 1
|
||||
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4)
|
||||
|
||||
# Track audio statistics
|
||||
audio_success_count = 0
|
||||
audio_skip_count = 0
|
||||
@@ -560,7 +581,7 @@ def compute_latents( # noqa: PLR0913, PLR0915
|
||||
"fps": batch["video_metadata"]["fps"][i].item(),
|
||||
}
|
||||
|
||||
torch.save(latent_data, output_file)
|
||||
_atomic_save(latent_data, output_file)
|
||||
|
||||
# Process audio if enabled (audio is already extracted by the dataset)
|
||||
if with_audio:
|
||||
@@ -588,7 +609,7 @@ def compute_latents( # noqa: PLR0913, PLR0915
|
||||
"duration": audio_latents["duration"],
|
||||
}
|
||||
|
||||
torch.save(audio_save_data, audio_output_file)
|
||||
_atomic_save(audio_save_data, audio_output_file)
|
||||
audio_success_count += 1
|
||||
else:
|
||||
# Video has no audio track
|
||||
@@ -596,8 +617,7 @@ def compute_latents( # noqa: PLR0913, PLR0915
|
||||
|
||||
progress.advance(task)
|
||||
|
||||
# Log summary
|
||||
logger.info(f"Processed {len(dataset)} videos. Latents saved to {output_path}")
|
||||
logger.info(f"Processed {len(dataloader.dataset)} videos -> {output_path}") # type: ignore[arg-type]
|
||||
if with_audio:
|
||||
logger.info(
|
||||
f"Audio processing: {audio_success_count} videos with audio, "
|
||||
@@ -935,6 +955,39 @@ def compute_scaled_resolution_buckets(
|
||||
return scaled_buckets
|
||||
|
||||
|
||||
def _atomic_save(data: Any, out: Path) -> None: # noqa: ANN401
|
||||
"""Save to ``out`` atomically via per-PID temp file + replace.
|
||||
Crash mid-write leaves an orphan ``.tmp.<pid>`` file that the skip logic
|
||||
ignores. The per-PID suffix makes concurrent writes from multiple ranks
|
||||
collision-free.
|
||||
"""
|
||||
tmp = out.with_suffix(f"{out.suffix}.tmp.{os.getpid()}")
|
||||
torch.save(data, tmp)
|
||||
tmp.replace(out)
|
||||
|
||||
|
||||
def _build_sharded_dataloader(
|
||||
dataset: Dataset,
|
||||
*,
|
||||
batch_size: int,
|
||||
num_workers: int,
|
||||
is_done: Callable[[int], bool],
|
||||
overwrite: bool,
|
||||
) -> DataLoader | None:
|
||||
"""Return a DataLoader over this rank's interleaved shard of ``dataset``.
|
||||
When ``overwrite`` is False, items whose outputs already exist (per
|
||||
``is_done``) are filtered out. Returns ``None`` if this rank has nothing
|
||||
to do, so the caller can early-return without loading any models.
|
||||
"""
|
||||
state = PartialState()
|
||||
todo = [i for i in range(state.process_index, len(dataset), state.num_processes) if overwrite or not is_done(i)]
|
||||
if not todo:
|
||||
logger.info(f"Rank {state.process_index}/{state.num_processes}: nothing to do")
|
||||
return None
|
||||
logger.info(f"Rank {state.process_index}/{state.num_processes}: processing {len(todo):,} of {len(dataset):,} items")
|
||||
return DataLoader(Subset(dataset, todo), batch_size=batch_size, shuffle=False, num_workers=num_workers)
|
||||
|
||||
|
||||
@app.command()
|
||||
def main( # noqa: PLR0913
|
||||
dataset_file: str = typer.Argument(
|
||||
@@ -981,8 +1034,15 @@ def main( # noqa: PLR0913
|
||||
default=None,
|
||||
help="Output directory for audio latents (required if --with-audio is set)",
|
||||
),
|
||||
overwrite: bool = typer.Option(
|
||||
default=False,
|
||||
help="Re-encode every item even if its output exists. Use when rerunning with "
|
||||
"changed parameters (different model, resolution, etc.) so stale outputs are replaced.",
|
||||
),
|
||||
) -> None:
|
||||
"""Process videos/images and save latent representations for video generation training.
|
||||
For multi-GPU preprocessing, invoke under ``accelerate launch`` - each process
|
||||
will handle an interleaved shard of the dataset.
|
||||
This script processes videos and images from metadata files and saves latent representations
|
||||
that can be used for training video generation models. The output latents will maintain
|
||||
the same folder structure and naming as the corresponding media files.
|
||||
@@ -1032,6 +1092,7 @@ def main( # noqa: PLR0913
|
||||
vae_tiling=vae_tiling,
|
||||
with_audio=with_audio,
|
||||
audio_output_dir=audio_output_dir,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -168,6 +168,15 @@ class AccelerationConfig(ConfigBaseModel):
|
||||
description="Whether to load the text encoder in 8-bit precision to save memory",
|
||||
)
|
||||
|
||||
offload_optimizer_during_validation: bool = Field(
|
||||
default=False,
|
||||
description="Offload optimizer state to CPU before validation video sampling and reload "
|
||||
"it afterwards, to free VRAM for inference. Useful when optimizer state is large "
|
||||
"(e.g. AdamW for full fine-tuning or high-rank LoRA) and validation OOMs because the "
|
||||
"VAE decoder + transformer + optimizer state cannot coexist on the GPU. Has no effect "
|
||||
"for FSDP (sharded state). Disabled by default.",
|
||||
)
|
||||
|
||||
|
||||
class DataConfig(ConfigBaseModel):
|
||||
"""Configuration for data loading and processing"""
|
||||
|
||||
@@ -85,6 +85,7 @@ def print_config(config: LtxTrainerConfig) -> None:
|
||||
("Mixed Precision", accel.mixed_precision_mode or "[dim]—[/]"),
|
||||
("Quantization", str(accel.quantization) if accel.quantization else "[dim]—[/]"),
|
||||
("Text Encoder 8bit", fmt(accel.load_text_encoder_in_8bit)),
|
||||
("Optimizer CPU Offload", fmt(accel.offload_optimizer_during_validation)),
|
||||
],
|
||||
),
|
||||
(
|
||||
|
||||
@@ -12,6 +12,7 @@ Example usage:
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
@@ -22,7 +23,11 @@ from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder
|
||||
from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
|
||||
|
||||
|
||||
def load_8bit_gemma(gemma_model_path: str | Path, dtype: torch.dtype = torch.bfloat16) -> GemmaTextEncoder:
|
||||
def load_8bit_gemma(
|
||||
gemma_model_path: str | Path,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
device: torch.device | str | int | None = None,
|
||||
) -> GemmaTextEncoder:
|
||||
"""Load the Gemma text encoder in 8-bit precision using bitsandbytes.
|
||||
Only the Gemma LLM backbone is loaded here. The embeddings processor
|
||||
(feature extractor + connectors) should be loaded separately via
|
||||
@@ -30,6 +35,10 @@ def load_8bit_gemma(gemma_model_path: str | Path, dtype: torch.dtype = torch.bfl
|
||||
Args:
|
||||
gemma_model_path: Path to Gemma model directory
|
||||
dtype: Data type for non-quantized model weights
|
||||
device: Device to place the quantized model on. When ``None`` (default),
|
||||
the device is inferred from ``LOCAL_RANK`` if CUDA is available, so
|
||||
multi-process launches put each rank's encoder on its own GPU
|
||||
instead of all colliding on ``cuda:0``.
|
||||
Returns:
|
||||
GemmaTextEncoder with 8-bit quantized Gemma backbone
|
||||
Raises:
|
||||
@@ -46,13 +55,23 @@ def load_8bit_gemma(gemma_model_path: str | Path, dtype: torch.dtype = torch.bfl
|
||||
gemma_path = _find_gemma_subpath(gemma_model_path, "model*.safetensors")
|
||||
tokenizer_path = _find_gemma_subpath(gemma_model_path, "tokenizer.model")
|
||||
|
||||
# Pin the entire model to a single device. `device_map="auto"` collides on cuda:0
|
||||
# in multi-process launches because every rank picks the same default device.
|
||||
device_map: str | dict[str, int | str | torch.device]
|
||||
if device is not None:
|
||||
device_map = {"": device}
|
||||
elif torch.cuda.is_available():
|
||||
device_map = {"": int(os.environ.get("LOCAL_RANK", "0"))}
|
||||
else:
|
||||
device_map = "auto"
|
||||
|
||||
quantization_config = BitsAndBytesConfig(load_in_8bit=True)
|
||||
with _suppress_accelerate_memory_warnings():
|
||||
gemma_model = Gemma3ForConditionalGeneration.from_pretrained(
|
||||
gemma_path,
|
||||
quantization_config=quantization_config,
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map="auto",
|
||||
device_map=device_map,
|
||||
local_files_only=True,
|
||||
)
|
||||
|
||||
|
||||
@@ -199,8 +199,6 @@ def load_text_encoder(
|
||||
device: Device to load model on
|
||||
dtype: Data type for model weights
|
||||
load_in_8bit: Whether to load the Gemma model in 8-bit precision using bitsandbytes.
|
||||
When True, the model is loaded with device_map="auto" and the device argument
|
||||
is ignored for the Gemma backbone.
|
||||
Returns:
|
||||
Loaded GemmaTextEncoder
|
||||
"""
|
||||
@@ -211,7 +209,7 @@ def load_text_encoder(
|
||||
if load_in_8bit:
|
||||
from ltx_trainer.gemma_8bit import load_8bit_gemma
|
||||
|
||||
return load_8bit_gemma(gemma_model_path, dtype)
|
||||
return load_8bit_gemma(gemma_model_path, dtype, device=device)
|
||||
|
||||
# Standard loading path
|
||||
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
import contextlib
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import warnings
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
@@ -9,8 +12,8 @@ from typing import Any, Callable
|
||||
import torch
|
||||
import wandb
|
||||
import yaml
|
||||
from accelerate import Accelerator, DistributedType
|
||||
from accelerate.utils import set_seed
|
||||
from accelerate import Accelerator, DistributedDataParallelKwargs, DistributedType
|
||||
from accelerate.utils import gather_object, set_seed
|
||||
from peft import LoraConfig, get_peft_model, get_peft_model_state_dict, set_peft_model_state_dict
|
||||
from peft.tuners.tuners_utils import BaseTunerLayer
|
||||
from peft.utils import ModulesToSaveWrapper
|
||||
@@ -63,7 +66,7 @@ if not IS_MAIN_PROCESS:
|
||||
|
||||
disable_progress_bar()
|
||||
|
||||
StepCallback = Callable[[int, int, list[Path]], None] # (step, total, list[sampled_video_path]) -> None
|
||||
StepCallback = Callable[[int, int, list[Path] | None], None] # (step, total, sampled paths or None) -> None
|
||||
|
||||
MEMORY_CHECK_INTERVAL = 200
|
||||
|
||||
@@ -186,9 +189,8 @@ class LtxvTrainer:
|
||||
|
||||
with progress:
|
||||
if cfg.validation.interval and not cfg.validation.skip_initial_validation:
|
||||
sampled_videos_paths = self._sample_videos(progress)
|
||||
if IS_MAIN_PROCESS and sampled_videos_paths and self._config.wandb.log_validation_videos:
|
||||
self._log_validation_samples(sampled_videos_paths, cfg.validation.prompts)
|
||||
with self._offloaded_optimizer_state():
|
||||
sampled_videos_paths = self._run_distributed_validation(progress)
|
||||
|
||||
self._accelerator.wait_for_everyone()
|
||||
|
||||
@@ -228,16 +230,8 @@ class LtxvTrainer:
|
||||
and self._global_step % cfg.validation.interval == 0
|
||||
and is_optimization_step
|
||||
):
|
||||
if self._accelerator.distributed_type == DistributedType.FSDP:
|
||||
# FSDP: All processes must participate in validation
|
||||
sampled_videos_paths = self._sample_videos(progress)
|
||||
if IS_MAIN_PROCESS and sampled_videos_paths and self._config.wandb.log_validation_videos:
|
||||
self._log_validation_samples(sampled_videos_paths, cfg.validation.prompts)
|
||||
# DDP: Only main process runs validation
|
||||
elif IS_MAIN_PROCESS:
|
||||
sampled_videos_paths = self._sample_videos(progress)
|
||||
if sampled_videos_paths and self._config.wandb.log_validation_videos:
|
||||
self._log_validation_samples(sampled_videos_paths, cfg.validation.prompts)
|
||||
with self._offloaded_optimizer_state():
|
||||
sampled_videos_paths = self._run_distributed_validation(progress)
|
||||
|
||||
# Save checkpoint if needed
|
||||
if (
|
||||
@@ -398,11 +392,14 @@ class LtxvTrainer:
|
||||
# 3. If validation prompts are configured, computes and caches their embeddings
|
||||
# 4. Unloads the Gemma model entirely, keeps the embeddings processor for training
|
||||
|
||||
# Load text encoder (pure Gemma LLM) on GPU
|
||||
# Load text encoder (pure Gemma LLM) on GPU — LOCAL_RANK before Accelerator exists
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
|
||||
init_device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
logger.debug("Loading text encoder...")
|
||||
text_encoder = load_text_encoder(
|
||||
gemma_model_path=self._config.model.text_encoder_path,
|
||||
device="cuda",
|
||||
device=init_device,
|
||||
dtype=torch.bfloat16,
|
||||
load_in_8bit=self._config.acceleration.load_text_encoder_in_8bit,
|
||||
)
|
||||
@@ -411,7 +408,7 @@ class LtxvTrainer:
|
||||
logger.debug("Loading embeddings processor...")
|
||||
self._embeddings_processor = load_embeddings_processor(
|
||||
checkpoint_path=self._config.model.model_path,
|
||||
device="cuda",
|
||||
device=init_device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
@@ -788,6 +785,41 @@ class LtxvTrainer:
|
||||
# noinspection PyTypeChecker
|
||||
self._optimizer, self._lr_scheduler = self._accelerator.prepare(optimizer, lr_scheduler)
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _offloaded_optimizer_state(self) -> Iterator[None]:
|
||||
"""Context manager that offloads optimizer state to CPU during validation.
|
||||
Opt-in via `acceleration.offload_optimizer_during_validation`. Frees VRAM for
|
||||
validation video generation when optimizer state is large (e.g. full fine-tune
|
||||
AdamW, high-rank LoRA). No-op for FSDP (sharded state -- manual `.cpu()` breaks
|
||||
metadata).
|
||||
"""
|
||||
enabled = (
|
||||
self._config.acceleration.offload_optimizer_during_validation
|
||||
and self._accelerator.distributed_type != DistributedType.FSDP
|
||||
)
|
||||
|
||||
# Track exactly which tensors we move so we don't promote ones that were
|
||||
# intentionally on CPU (e.g. AdamW's `step` scalar on recent PyTorch).
|
||||
offloaded: list[tuple[dict, str]] = []
|
||||
if enabled:
|
||||
offloaded_bytes = 0
|
||||
for state in self._optimizer.state.values():
|
||||
for k, v in state.items():
|
||||
if isinstance(v, torch.Tensor) and v.is_cuda:
|
||||
offloaded.append((state, k))
|
||||
offloaded_bytes += v.nbytes
|
||||
if offloaded:
|
||||
logger.info(f"Offloading optimizer state to CPU ({offloaded_bytes / 1e9:.1f} GB)")
|
||||
for state, k in offloaded:
|
||||
state[k] = state[k].cpu()
|
||||
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
device = self._accelerator.device
|
||||
for state, k in offloaded:
|
||||
state[k] = state[k].to(device)
|
||||
|
||||
def _create_scheduler(self, optimizer: torch.optim.Optimizer) -> LRScheduler | None:
|
||||
"""Create learning rate scheduler based on config."""
|
||||
scheduler_type = self._config.optimization.scheduler_type
|
||||
@@ -844,11 +876,18 @@ class LtxvTrainer:
|
||||
def _setup_accelerator(self) -> None:
|
||||
"""Initialize the Accelerator with the appropriate settings."""
|
||||
|
||||
# find_unused_parameters=True keeps DDP happy when LoRA targets a branch the forward
|
||||
# pass skips (e.g. audio LoRA with `with_audio: false`, or short module patterns like
|
||||
# "to_k" that match the audio branch unintentionally). It's a no-op for FSDP and
|
||||
# single-GPU runs. The probing cost is paid only on the first step.
|
||||
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
|
||||
|
||||
# All distributed setup (DDP/FSDP, number of processes, etc.) is controlled by
|
||||
# the user's Accelerate configuration (accelerate config / accelerate launch).
|
||||
self._accelerator = Accelerator(
|
||||
mixed_precision=self._config.acceleration.mixed_precision_mode,
|
||||
gradient_accumulation_steps=self._config.optimization.gradient_accumulation_steps,
|
||||
kwargs_handlers=[ddp_kwargs],
|
||||
)
|
||||
|
||||
if self._accelerator.num_processes > 1:
|
||||
@@ -881,11 +920,42 @@ class LtxvTrainer:
|
||||
"Monitor training stability and consider disabling quantization if issues arise."
|
||||
)
|
||||
|
||||
def _run_distributed_validation(self, progress: TrainingProgress) -> list[Path]:
|
||||
"""Run validation across all ranks and log gathered results on rank 0.
|
||||
Each rank generates only its assigned subset of prompts (see `_sample_videos`),
|
||||
so all GPUs stay busy and no rank idles long enough to trigger NCCL timeouts.
|
||||
Paths are gathered across ranks so rank 0 has the full list for W&B logging.
|
||||
Note: Multi-node training requires a shared filesystem so rank 0 can read
|
||||
videos written by other ranks.
|
||||
"""
|
||||
sampled = self._sample_videos(progress)
|
||||
|
||||
if self._accelerator.num_processes > 1:
|
||||
# gather_object returns a flat list from all ranks
|
||||
sampled = sorted(gather_object(sampled), key=lambda x: x[0])
|
||||
|
||||
paths = [p for _, p in sampled]
|
||||
|
||||
if self._accelerator.is_main_process and paths:
|
||||
self._log_validation_samples(paths, self._config.validation.prompts)
|
||||
|
||||
# Non-main ranks must not reach checkpoint collectives while main is still logging to W&B.
|
||||
self._accelerator.wait_for_everyone()
|
||||
|
||||
return paths
|
||||
|
||||
# Note: Use @torch.no_grad() instead of @torch.inference_mode() to avoid FSDP inplace update errors after validation
|
||||
@torch.no_grad()
|
||||
@free_gpu_memory_context(after=True)
|
||||
def _sample_videos(self, progress: TrainingProgress) -> list[Path] | None:
|
||||
"""Run validation by generating videos from validation prompts."""
|
||||
def _sample_videos(self, progress: TrainingProgress) -> list[tuple[int, Path]]:
|
||||
"""Run validation by generating videos from this rank's share of the validation prompts.
|
||||
Prompts are split round-robin across ranks via `process_index` / `num_processes`,
|
||||
which collapses to "all prompts" when running on a single GPU. Returns
|
||||
(prompt_idx, path) tuples so the caller can reconstruct global order without
|
||||
relying on filename conventions.
|
||||
Under FSDP with multiple processes, ranks pad with extra generate passes (same prompt,
|
||||
no disk write) so every rank runs the same number of forwards — avoids collective mismatch.
|
||||
"""
|
||||
use_images = self._config.validation.images is not None
|
||||
use_reference_videos = self._config.validation.reference_videos is not None
|
||||
generate_audio = self._config.validation.generate_audio
|
||||
@@ -895,13 +965,24 @@ class LtxvTrainer:
|
||||
self._optimizer.zero_grad(set_to_none=True)
|
||||
free_gpu_memory()
|
||||
|
||||
# Start sampling progress tracking
|
||||
prompts = self._config.validation.prompts
|
||||
rank = self._accelerator.process_index
|
||||
world_size = self._accelerator.num_processes
|
||||
rank_indices = list(range(rank, len(prompts), world_size))
|
||||
|
||||
# FSDP: every rank must run the same number of forwards; pad with duplicate generates (no save).
|
||||
work: list[tuple[int, bool]] = [(i, True) for i in rank_indices]
|
||||
if self._accelerator.distributed_type == DistributedType.FSDP and world_size > 1:
|
||||
max_per_rank = math.ceil(len(prompts) / world_size)
|
||||
pad_seed = rank_indices[-1] if rank_indices else 0
|
||||
work += [(pad_seed, False)] * (max_per_rank - len(work))
|
||||
|
||||
sampling_ctx = progress.start_sampling(
|
||||
num_prompts=len(self._config.validation.prompts),
|
||||
num_prompts=len(work),
|
||||
num_steps=inference_steps,
|
||||
)
|
||||
|
||||
# Create validation sampler with loaded models and progress tracking
|
||||
# Create a validation sampler with loaded models and progress tracking
|
||||
sampler = ValidationSampler(
|
||||
transformer=self._transformer,
|
||||
vae_decoder=self._vae_decoder,
|
||||
@@ -915,12 +996,12 @@ class LtxvTrainer:
|
||||
output_dir = Path(self._config.output_dir) / "samples"
|
||||
output_dir.mkdir(exist_ok=True, parents=True)
|
||||
|
||||
video_paths = []
|
||||
results: list[tuple[int, Path]] = []
|
||||
width, height, num_frames = self._config.validation.video_dims
|
||||
|
||||
for prompt_idx, prompt in enumerate(self._config.validation.prompts):
|
||||
# Update progress to show current video
|
||||
sampling_ctx.start_video(prompt_idx)
|
||||
for local_i, (prompt_idx, save_output) in enumerate(work):
|
||||
prompt = prompts[prompt_idx]
|
||||
sampling_ctx.start_video(local_i)
|
||||
|
||||
# Load conditioning image if provided
|
||||
condition_image = None
|
||||
@@ -972,28 +1053,30 @@ class LtxvTrainer:
|
||||
device=self._accelerator.device,
|
||||
)
|
||||
|
||||
if not save_output:
|
||||
continue
|
||||
|
||||
# Save output (image for single frame, video otherwise)
|
||||
if IS_MAIN_PROCESS:
|
||||
ext = "png" if num_frames == 1 else "mp4"
|
||||
output_path = output_dir / f"step_{self._global_step:06d}_{prompt_idx + 1}.{ext}"
|
||||
if num_frames == 1:
|
||||
save_image(video, output_path)
|
||||
else:
|
||||
save_video(
|
||||
video_tensor=video,
|
||||
output_path=output_path,
|
||||
fps=self._config.validation.frame_rate,
|
||||
audio=audio,
|
||||
audio_sample_rate=self._vocoder.output_sampling_rate if audio is not None else None,
|
||||
)
|
||||
video_paths.append(output_path)
|
||||
ext = "png" if num_frames == 1 else "mp4"
|
||||
output_path = output_dir / f"step_{self._global_step:06d}_{prompt_idx + 1:02d}.{ext}"
|
||||
if num_frames == 1:
|
||||
save_image(video, output_path)
|
||||
else:
|
||||
save_video(
|
||||
video_tensor=video,
|
||||
output_path=output_path,
|
||||
fps=self._config.validation.frame_rate,
|
||||
audio=audio,
|
||||
audio_sample_rate=self._vocoder.output_sampling_rate if audio is not None else None,
|
||||
)
|
||||
results.append((prompt_idx, output_path))
|
||||
|
||||
# Clean up progress tasks
|
||||
sampling_ctx.cleanup()
|
||||
|
||||
rel_outputs_path = output_dir.relative_to(self._config.output_dir)
|
||||
logger.info(f"🎥 Validation samples for step {self._global_step} saved in {rel_outputs_path}")
|
||||
return video_paths
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def _log_training_stats(stats: TrainingStats) -> None:
|
||||
|
||||
Reference in New Issue
Block a user