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

Automated PR - 2026-04-13

Browse Source
This commit is contained in:
github-actions[bot]
2026-04-13 14:29:35 +00:00
parent 59ca828d5a
commit d887bbd1e0
29 changed files with 463 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-trainer/src/ltx_trainer/sigma_tracker.py
+59
View File
@@ -0,0 +1,59 @@
"""Sigma-bucketed loss tracking.
Maps each training step's per-element sigmas and losses to buckets.
Smoothing is left to wandb's UI.
"""
import bisect
from collections import defaultdict
class SigmaBucketTracker:
"""Map per-element sigma values to named buckets for per-bucket loss logging.
By default, partitions [0, 1] into four equal-width buckets.
Custom boundaries can be provided for non-uniform bucketing.
Each call to update() receives per-element sigmas and losses (both [B,]),
buckets each element, and computes the mean loss per bucket. This gives
accurate per-sigma loss tracking even for batch_size > 1.
"""
def __init__(
self,
bucket_boundaries: list[float] | None = None,
) -> None:
if bucket_boundaries is None:
bucket_boundaries = [0.0, 0.25, 0.5, 0.75, 1.0]
if len(bucket_boundaries) < 2:
raise ValueError("bucket_boundaries must have at least 2 elements")
if any(bucket_boundaries[i] >= bucket_boundaries[i + 1] for i in range(len(bucket_boundaries) - 1)):
raise ValueError("bucket_boundaries must be strictly increasing")
self._boundaries = list(bucket_boundaries)
self._num_buckets = len(bucket_boundaries) - 1
self._bucket_labels = [
f"{bucket_boundaries[i]:.2f}-{bucket_boundaries[i + 1]:.2f}" for i in range(self._num_buckets)
]
self._last_metrics: dict[str, float] = {}
def _get_bucket_index(self, sigma: float) -> int:
"""Map sigma value to bucket index."""
idx = bisect.bisect_right(self._boundaries, sigma) - 1
return max(0, min(idx, self._num_buckets - 1))
def update(self, sigmas: list[float], losses: list[float]) -> None:
"""Record per-element losses into their sigma buckets.
Args:
sigmas: Per-element sigma values, one per batch element.
losses: Per-element losses, one per batch element.
"""
if not sigmas:
self._last_metrics = {}
return
bucket_losses: dict[int, list[float]] = defaultdict(list)
for sigma, loss in zip(sigmas, losses, strict=True):
bucket_losses[self._get_bucket_index(sigma)].append(loss)
self._last_metrics = {self._bucket_labels[b]: sum(vals) / len(vals) for b, vals in bucket_losses.items()}
def get_metrics(self, prefix: str = "train") -> dict[str, float]:
"""Return the mean loss for each bucket hit on the last update.
Wandb handles smoothing in the UI.
"""
return {f"{prefix}/loss_sigma_{label}": loss for label, loss in self._last_metrics.items()}
packages/ltx-trainer/src/ltx_trainer/trainer.py
+50 -26
View File
@@ -2,8 +2,9 @@ import os
import re
import time
import warnings
from dataclasses import dataclass
from pathlib import Path
from typing import Callable
from typing import Any, Callable
import torch
import wandb
@@ -39,6 +40,7 @@ from ltx_trainer.model_loader import load_embeddings_processor, load_text_encode
from ltx_trainer.model_loader import load_model as load_ltx_model
from ltx_trainer.progress import TrainingProgress
from ltx_trainer.quantization import quantize_model
from ltx_trainer.sigma_tracker import SigmaBucketTracker
from ltx_trainer.timestep_samplers import SAMPLERS
from ltx_trainer.training_state import ConfigFingerprint, RngStates, TrainingState
from ltx_trainer.training_strategies import get_training_strategy
@@ -77,6 +79,14 @@ class TrainingStats(BaseModel):
num_processes: int
@dataclass(frozen=True)
class TrainingStepOutput:
"""Output from a single training step."""
loss: Tensor # [B,] per-element loss (unreduced)
sigma: Tensor # [B,] sampled sigma, detached from computational graph
class LtxvTrainer:
def __init__(self, trainer_config: LtxTrainerConfig) -> None:
self._config = trainer_config
@@ -95,7 +105,8 @@ class LtxvTrainer:
self._checkpoint_paths: list[Path] = []
self._training_state_paths: list[Path] = []
self._training_state_size_warned = False
self._init_wandb()
self._wandb_run = None
self._sigma_tracker = SigmaBucketTracker()
def train( # noqa: PLR0912, PLR0915
self,
@@ -128,6 +139,10 @@ class LtxvTrainer:
initial_step = 0
resuming = False
# Initialize W&B after restore so we only resume the run when state restore succeeds.
resume_run_id = training_state.wandb_run_id if resuming and training_state is not None else None
self._init_wandb(resume_run_id=resume_run_id)
self._init_dataloader()
data_iter = iter(self._dataloader)
self._init_timestep_sampler()
@@ -191,8 +206,8 @@ class LtxvTrainer:
if is_optimization_step:
self._global_step += 1
loss = self._training_step(batch)
self._accelerator.backward(loss)
output = self._training_step(batch)
self._accelerator.backward(output.loss.mean())
if self._accelerator.sync_gradients and cfg.optimization.max_grad_norm > 0:
self._accelerator.clip_grad_norm_(
@@ -244,9 +259,10 @@ class LtxvTrainer:
# Update progress and log metrics
current_lr = self._optimizer.param_groups[0]["lr"]
step_time = (time.time() - step_start_time) * cfg.optimization.gradient_accumulation_steps
step_loss = output.loss.detach().mean().item()
progress.update_training(
loss=loss.item(),
loss=step_loss,
lr=current_lr,
step_time=step_time,
advance=is_optimization_step,
@@ -254,14 +270,16 @@ class LtxvTrainer:
# Log metrics to W&B (only on main process and optimization steps)
if IS_MAIN_PROCESS and is_optimization_step:
self._log_metrics(
{
"train/loss": loss.item(),
"train/learning_rate": current_lr,
"train/step_time": step_time,
"train/global_step": self._global_step,
}
)
# Track per-element loss by sigma bucket
self._sigma_tracker.update(output.sigma.cpu().tolist(), output.loss.detach().cpu().tolist())
metrics = {
"train/loss": step_loss,
"train/learning_rate": current_lr,
"train/step_time": step_time,
"train/global_step": self._global_step,
}
metrics.update(self._sigma_tracker.get_metrics())
self._log_metrics(metrics)
# Fallback logging when progress bars are disabled
if disable_progress_bars and IS_MAIN_PROCESS and self._global_step % 20 == 0:
@@ -274,7 +292,7 @@ class LtxvTrainer:
total_time = "calculating..."
logger.info(
f"Step {self._global_step}/{cfg.optimization.steps} - "
f"Loss: {loss.item():.4f}, LR: {current_lr:.2e}, "
f"Loss: {step_loss:.4f}, LR: {current_lr:.2e}, "
f"Time/Step: {step_time:.2f}s, Total Time: {total_time}",
)
@@ -330,7 +348,7 @@ class LtxvTrainer:
return saved_path, stats
def _training_step(self, batch: dict[str, dict[str, Tensor]]) -> Tensor:
def _training_step(self, batch: dict[str, dict[str, Tensor]]) -> TrainingStepOutput:
"""Perform a single training step using the configured strategy."""
# Apply embedding connectors to transform pre-computed text embeddings
conditions = batch["conditions"]
@@ -366,8 +384,9 @@ class LtxvTrainer:
# Use strategy to compute loss
loss = self._training_strategy.compute_loss(video_pred, audio_pred, model_inputs)
sigma = model_inputs.video.sigma.detach() if model_inputs.video.enabled else model_inputs.audio.sigma.detach()
return loss
return TrainingStepOutput(loss=loss, sigma=sigma)
@free_gpu_memory_context(after=True)
def _load_text_encoder_and_cache_embeddings(self) -> list[CachedPromptEmbeddings] | None:
@@ -1064,8 +1083,8 @@ class LtxvTrainer:
def _save_training_state(self, save_dir: Path) -> None:
"""Save training state alongside checkpoint for resume.
Respects checkpoints.save_training_state config:
- "full": optimizer + scheduler + RNG + step
- "minimal": scheduler + RNG + step only
- "full": optimizer + scheduler + RNG + step + wandb_run_id
- "minimal": scheduler + RNG + step + wandb_run_id
- "off": skip entirely
"""
if not IS_MAIN_PROCESS:
@@ -1101,6 +1120,7 @@ class LtxvTrainer:
),
lr_scheduler_state_dict=self._lr_scheduler.state_dict() if self._lr_scheduler is not None else None,
optimizer_state_dict=optimizer_state,
wandb_run_id=self._wandb_run.id if self._wandb_run is not None else None,
)
state_path = save_dir / f"training_state_step_{self._global_step:05d}.pt"
@@ -1166,20 +1186,24 @@ class LtxvTrainer:
logger.info(f"💾 Training configuration saved to: {config_path.relative_to(self._config.output_dir)}")
def _init_wandb(self) -> None:
def _init_wandb(self, resume_run_id: str | None = None) -> None:
"""Initialize Weights & Biases run."""
if not self._config.wandb.enabled or not IS_MAIN_PROCESS:
self._wandb_run = None
return
wandb_config = self._config.wandb
run = wandb.init(
project=wandb_config.project,
entity=wandb_config.entity,
name=Path(self._config.output_dir).name,
tags=wandb_config.tags,
config=self._config.model_dump(),
)
init_kwargs: dict[str, Any] = {
"project": wandb_config.project,
"entity": wandb_config.entity,
"name": Path(self._config.output_dir).name,
"tags": wandb_config.tags,
"config": self._config.model_dump(),
}
if resume_run_id is not None:
init_kwargs["id"] = resume_run_id
init_kwargs["resume"] = "allow"
run = wandb.init(**init_kwargs)
self._wandb_run = run
def _log_metrics(self, metrics: dict[str, float]) -> None:
packages/ltx-trainer/src/ltx_trainer/training_state.py
+2
View File
@@ -28,6 +28,7 @@ class TrainingState(BaseModel):
rng_states: RngStates
lr_scheduler_state_dict: dict[str, Any] | None = None
optimizer_state_dict: dict[str, Any] | None = None
wandb_run_id: str | None = None
def to_save_dict(self) -> dict[str, Any]:
"""Build dict suitable for torch.save -- recurses BaseModel sub-models, passes tensors/dicts through."""
@@ -48,4 +49,5 @@ class TrainingState(BaseModel):
rng_states=RngStates(**data["rng_states"]),
lr_scheduler_state_dict=data.get("lr_scheduler_state_dict"),
optimizer_state_dict=data.get("optimizer_state_dict"),
wandb_run_id=data.get("wandb_run_id"),
)
packages/ltx-trainer/src/ltx_trainer/training_strategies/base_strategy.py
+2 -1
View File
@@ -125,7 +125,8 @@ class TrainingStrategy(ABC):
audio_pred: Audio prediction from the transformer model (None for video-only)
inputs: The prepared model inputs containing targets and masks
Returns:
Scalar loss tensor
Per-element loss tensor of shape [B,]. The trainer reduces to a scalar
before backward(). Returning unreduced loss enables per-sigma-bucket tracking.
"""
def get_checkpoint_metadata(self) -> dict[str, Any]:
packages/ltx-trainer/src/ltx_trainer/training_strategies/text_to_video.py
+7 -7
View File
@@ -273,19 +273,19 @@ class TextToVideoStrategy(TrainingStrategy):
audio_pred: Tensor | None,
inputs: ModelInputs,
) -> Tensor:
"""Compute masked MSE loss for video and optionally audio."""
# Video loss
"""Compute masked MSE loss for video and optionally audio. Returns [B,]."""
# Video loss: per-element mean over (seq, channels), [B,]
video_loss = (video_pred - inputs.video_targets).pow(2)
video_loss_mask = inputs.video_loss_mask.unsqueeze(-1).float()
video_loss = video_loss.mul(video_loss_mask).div(video_loss_mask.mean())
video_loss = video_loss.mean()
masked = video_loss.mul(video_loss_mask)
video_loss = masked.mean(dim=[-2, -1]) / video_loss_mask.mean(dim=[-2, -1]).clamp(min=1e-8)
# If no audio, return video loss only
if not self.config.with_audio or audio_pred is None or inputs.audio_targets is None:
return video_loss
# Audio loss (no conditioning mask)
audio_loss = (audio_pred - inputs.audio_targets).pow(2).mean()
# Audio loss: per-element mean over (seq, channels), [B,]
audio_loss = (audio_pred - inputs.audio_targets).pow(2).mean(dim=[-2, -1])
# Combined loss
# Combined loss [B,]
return video_loss + audio_loss
packages/ltx-trainer/src/ltx_trainer/training_strategies/video_to_video.py
+4 -7
View File
@@ -240,7 +240,7 @@ class VideoToVideoStrategy(TrainingStrategy):
_audio_pred: Tensor | None,
inputs: ModelInputs,
) -> Tensor:
"""Compute masked loss only on target portion."""
"""Compute masked loss only on target portion. Returns [B,]."""
# Extract target portion of prediction
ref_seq_len = inputs.ref_seq_len
target_pred = video_pred[:, ref_seq_len:, :]
@@ -248,14 +248,11 @@ class VideoToVideoStrategy(TrainingStrategy):
# Get target portion of loss mask
target_loss_mask = inputs.video_loss_mask[:, ref_seq_len:]
# Compute loss
# Compute per-element loss [B,]
loss = (target_pred - inputs.video_targets).pow(2)
# Apply loss mask
loss_mask = target_loss_mask.unsqueeze(-1).float()
loss = loss.mul(loss_mask).div(loss_mask.mean())
return loss.mean()
masked = loss.mul(loss_mask)
return masked.mean(dim=[-2, -1]) / loss_mask.mean(dim=[-2, -1]).clamp(min=1e-8)
def get_checkpoint_metadata(self) -> dict[str, Any]:
"""Get metadata for checkpoint files."""
packages/ltx-trainer/src/ltx_trainer/video_utils.py
+16 -5
View File
@@ -14,6 +14,9 @@ from torch import Tensor
def get_video_frame_count(video_path: str | Path) -> int:
"""Get the number of frames in a video file.
Tries three approaches in order: stream metadata, duration*fps estimate,
full decode. The estimate may be off by a few frames for VFR videos or
containers with edit lists — exact for the min_frames filtering use case.
Args:
video_path: Path to the video file
Returns:
@@ -21,11 +24,19 @@ def get_video_frame_count(video_path: str | Path) -> int:
"""
with av.open(str(video_path)) as container:
video_stream = container.streams.video[0]
frame_count = video_stream.frames
if frame_count == 0:
# Fallback: count frames by decoding
frame_count = sum(1 for _ in container.decode(video=0))
return frame_count
if video_stream.frames > 0:
return video_stream.frames
# Fast estimate from container metadata (avoids full decode).
# Uses Fraction arithmetic to prevent float precision loss.
rate = video_stream.average_rate or video_stream.base_rate
if video_stream.duration and video_stream.time_base and rate:
duration = Fraction(video_stream.duration) * Fraction(video_stream.time_base)
return round(duration * Fraction(rate))
# Last resort: full decode (very slow for 4K)
return sum(1 for _ in container.decode(video=0))
def read_video(video_path: str | Path, max_frames: int | None = None) -> tuple[Tensor, float]: