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

Automated PR - 2026-04-23

Browse Source
This commit is contained in:
github-actions[bot]
2026-04-23 12:43:54 +00:00
parent a2c3f24078
commit b604d3fab3
49 changed files with 2664 additions and 568 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-trainer/src/ltx_trainer/config.py
-6
View File
@@ -260,12 +260,6 @@ class ValidationConfig(ConfigBaseModel):
gt=0,
)
videos_per_prompt: int = Field(
default=1,
description="Number of videos to generate per validation prompt",
gt=0,
)
guidance_scale: float = Field(
default=4.0,
description="CFG guidance scale to use during validation",
packages/ltx-trainer/src/ltx_trainer/datasets.py
+64 -23
View File
@@ -1,3 +1,4 @@
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import torch
@@ -155,40 +156,77 @@ class PrecomputedDataset(Dataset):
return source_paths
def _discover_samples(self) -> dict[str, list[Path]]:
"""Discover all valid sample files across all data sources."""
# Use first data source as the reference to discover samples
"""Discover all valid sample files across all data sources.
Uses a fast two-pass approach: first globs all sources in parallel to build
full-path sets in memory, then checks expected paths via set membership.
This avoids O(N * num_sources) stat calls on networked filesystems while
correctly handling path remapping (e.g. latent_X.pt -> condition_X.pt).
"""
if not self.data_sources:
raise ValueError("No data sources configured")
data_key = "latents" if "latents" in self.data_sources else next(iter(self.data_sources.keys()))
data_path = self.source_paths[data_key]
data_files = list(data_path.glob("**/*.pt"))
# Pass 1: Glob all sources in parallel, build full-path sets
def _glob_source(dir_name: str) -> tuple[list[Path], set[str]]:
source_path = self.source_paths[dir_name]
paths = list(source_path.glob("**/*.pt"))
path_set = {str(p) for p in paths}
return paths, path_set
with ThreadPoolExecutor(max_workers=len(self.data_sources)) as executor:
glob_results = dict(
zip(
self.data_sources.keys(),
executor.map(_glob_source, self.data_sources.keys()),
strict=True,
)
)
# Get primary source files (cached from glob, no second scan)
data_files, _ = glob_results[data_key]
if not data_files:
raise ValueError(f"No data files found in {data_path}")
data_files.sort()
# Initialize sample files dict
sample_files = {output_key: [] for output_key in self.data_sources.values()}
# Log source sizes
for dir_name, (paths, _) in glob_results.items():
logger.debug(f"Source {dir_name}: {len(paths)} files")
# Build path sets for non-primary sources
other_path_sets = {
dir_name: path_set for dir_name, (_, path_set) in glob_results.items() if dir_name != data_key
}
# Pass 2: For each primary file, check if expected paths exist in other sources' sets
sample_files: dict[str, list[Path]] = {output_key: [] for output_key in self.data_sources.values()}
valid_count = 0
# For each data file, find corresponding files in other sources
for data_file in data_files:
rel_path = data_file.relative_to(data_path)
# Check if corresponding files exist in ALL sources
if self._all_source_files_exist(data_file, rel_path):
# Check all other sources via set lookup (O(1) per source, no stat calls)
all_exist = True
for dir_name, path_set in other_path_sets.items():
expected = self._get_expected_file_path(dir_name, data_file, rel_path)
if str(expected) not in path_set:
logger.debug(f"Skipping {data_file.name}: no matching {dir_name} file at {expected}")
all_exist = False
break
if all_exist:
self._fill_sample_data_files(data_file, rel_path, sample_files)
valid_count += 1
skipped = len(data_files) - valid_count
if skipped > 0:
logger.info(f"Fast index: {valid_count} valid samples from {len(data_files)} total ({skipped} skipped)")
else:
logger.debug(f"Fast index: {valid_count} valid samples from {len(data_files)} total")
return sample_files
def _all_source_files_exist(self, data_file: Path, rel_path: Path) -> bool:
"""Check if corresponding files exist in all data sources."""
for dir_name in self.data_sources:
expected_path = self._get_expected_file_path(dir_name, data_file, rel_path)
if not expected_path.exists():
logger.warning(
f"No matching {dir_name} file found for: {data_file.name} (expected in: {expected_path})"
)
return False
return True
def _get_expected_file_path(self, dir_name: str, data_file: Path, rel_path: Path) -> Path:
"""Get the expected file path for a given data source."""
source_path = self.source_paths[dir_name]
@@ -207,11 +245,14 @@ class PrecomputedDataset(Dataset):
def _validate_setup(self) -> None:
"""Validate that the dataset setup is correct."""
if not self.sample_files:
raise ValueError("No valid samples found - all data sources must have matching files")
sample_counts = {key: len(files) for key, files in self.sample_files.items()}
if not sample_counts or all(count == 0 for count in sample_counts.values()):
raise ValueError(
f"No valid samples found in {self.data_root} - all configured data sources "
f"({list(self.data_sources)}) must have matching files (per-source counts: {sample_counts})"
)
# Verify all output keys have the same number of samples
sample_counts = {key: len(files) for key, files in self.sample_files.items()}
if len(set(sample_counts.values())) > 1:
raise ValueError(f"Mismatched sample counts across sources: {sample_counts}")
packages/ltx-trainer/src/ltx_trainer/hf_hub_utils.py
+7 -3
View File
@@ -1,7 +1,7 @@
import shutil
import tempfile
from pathlib import Path
from typing import List, Union
from typing import List, Optional, Union
import imageio
from huggingface_hub import HfApi, create_repo
@@ -12,7 +12,11 @@ from ltx_trainer import logger
from ltx_trainer.config import LtxTrainerConfig
def push_to_hub(weights_path: Path, sampled_videos_paths: List[Path], config: LtxTrainerConfig) -> None:
def push_to_hub(
weights_path: Path,
sampled_videos_paths: Optional[List[Path]],
config: LtxTrainerConfig,
) -> None:
"""Push the trained LoRA weights to HuggingFace Hub."""
if not config.hub.hub_model_id:
logger.warning("⚠️ HuggingFace hub_model_id not specified, skipping push to hub")
@@ -108,7 +112,7 @@ def convert_video_to_gif(video_path: Path, output_path: Path) -> None:
def _create_model_card(
output_dir: Union[str, Path],
videos: List[Path],
videos: Optional[List[Path]],
config: LtxTrainerConfig,
) -> Path:
"""Generate and save a model card for the trained model."""
packages/ltx-trainer/src/ltx_trainer/training_strategies/base_strategy.py
+7 -4
View File
@@ -3,7 +3,6 @@ This module defines the abstract base class that all training strategies must im
along with the base configuration class.
"""
import random
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, Literal
@@ -251,13 +250,17 @@ class TrainingStrategy(ABC):
device: Target device
first_frame_conditioning_p: Probability of conditioning on the first frame
Returns:
Boolean mask where True indicates first frame tokens (if conditioning is enabled)
Boolean mask where True indicates first frame tokens (if conditioning is enabled).
The conditioning decision is drawn independently per batch element so the training
signal across samples in a batch is i.i.d.
"""
conditioning_mask = torch.zeros(batch_size, sequence_length, dtype=torch.bool, device=device)
if first_frame_conditioning_p > 0 and random.random() < first_frame_conditioning_p:
if first_frame_conditioning_p > 0:
first_frame_end_idx = height * width
if first_frame_end_idx < sequence_length:
conditioning_mask[:, :first_frame_end_idx] = True
# Per-sample Bernoulli draw so each batch element is independently conditioned.
per_sample_condition = torch.rand(batch_size, device=device) < first_frame_conditioning_p
conditioning_mask[per_sample_condition, :first_frame_end_idx] = True
return conditioning_mask
packages/ltx-trainer/src/ltx_trainer/video_utils.py
+23 -5
View File
@@ -5,12 +5,15 @@ with optional audio support.
from fractions import Fraction
from pathlib import Path
from typing import Literal
import av
import numpy as np
import torch
from torch import Tensor
VideoFormat = Literal["CFHW", "FCHW"]
def get_video_frame_count(video_path: str | Path) -> int:
"""Get the number of frames in a video file.
@@ -68,6 +71,7 @@ def save_video(
fps: float = 24.0,
audio: torch.Tensor | None = None,
audio_sample_rate: int | None = None,
video_format: VideoFormat | None = None,
) -> None:
"""Save a video tensor to a file using PyAV, optionally with audio.
Args:
@@ -76,12 +80,16 @@ def save_video(
fps: Frames per second for the output video
audio: Optional audio tensor of shape [C, samples] or [samples, C] in range [-1, 1]
audio_sample_rate: Sample rate for the audio (required if audio is provided)
video_format: Explicit layout of ``video_tensor``, either ``"CFHW"`` or ``"FCHW"``.
When ``None`` (default), the layout is auto-detected using a heuristic that only
works when ``shape[1] > 3`` — the ambiguous ``[C=3, F=3, H, W]`` / ``[F=3, C=3, H, W]``
case requires passing this argument explicitly.
"""
output_path = Path(output_path)
output_path.parent.mkdir(parents=True, exist_ok=True)
# Normalize to [F, H, W, C] uint8 numpy array
video_np = _prepare_video_array(video_tensor)
video_np = _prepare_video_array(video_tensor, video_format=video_format)
_, height, width, _ = video_np.shape
with av.open(str(output_path), mode="w") as container:
@@ -113,11 +121,21 @@ def save_video(
_write_audio(container, audio_stream, audio, audio_sample_rate)
def _prepare_video_array(video_tensor: torch.Tensor) -> np.ndarray:
"""Convert video tensor to [F, H, W, C] uint8 numpy array."""
# Handle [C, F, H, W] vs [F, C, H, W] format
if video_tensor.shape[0] == 3 and video_tensor.shape[1] > 3:
def _prepare_video_array(
video_tensor: torch.Tensor,
video_format: VideoFormat | None = None,
) -> np.ndarray:
"""Convert video tensor to [F, H, W, C] uint8 numpy array.
If ``video_format`` is provided, it is trusted. Otherwise, the layout is auto-detected
using a heuristic that only fires when ``shape[0] == 3 and shape[1] > 3`` (CFHW). The
ambiguous ``[C=3, F=3, H, W]`` / ``[F=3, C=3, H, W]`` case cannot be disambiguated and
defaults to the FCHW interpretation — callers must pass ``video_format`` explicitly for
3-frame CFHW tensors.
"""
if video_format == "CFHW":
video_tensor = video_tensor.permute(1, 0, 2, 3) # [C, F, H, W] -> [F, C, H, W]
elif video_format is None and video_tensor.shape[0] == 3 and video_tensor.shape[1] > 3:
video_tensor = video_tensor.permute(1, 0, 2, 3)
# Normalize to [0, 255] uint8
if video_tensor.max() <= 1.0: