Automated PR - 2026-04-13

packages/ltx-trainer/docs/custom-training-strategies.md

@@ -307,15 +307,14 @@ class InpaintingStrategy(TrainingStrategy):
             audio_pred: Tensor | None,
             inputs: ModelInputs,
     ) -> Tensor:
-        """Compute training loss on inpaint regions only."""
+        """Compute training loss on inpaint regions only. Returns [B,]."""
         # MSE loss
         loss = (video_pred - inputs.video_targets).pow(2)
 
-        # Apply loss mask
+        # Apply loss mask and reduce to per-element [B,]
         loss_mask = inputs.video_loss_mask.unsqueeze(-1).float()
-        loss = loss.mul(loss_mask).div(loss_mask.mean() + 1e-8)
-
-        return loss.mean()
+        masked = loss.mul(loss_mask)
+        return masked.mean(dim=[-2, -1]) / loss_mask.mean(dim=[-2, -1]).clamp(min=1e-8)
 ```
 
 ### Step 5: Register the Strategy

packages/ltx-trainer/docs/utility-scripts.md

@@ -48,6 +48,10 @@ uv run python scripts/caption_videos.py videos_dir/ --output dataset.json --use-
 uv run python scripts/caption_videos.py videos_dir/ --output dataset.json \
     --captioner-type gemini_flash --api-key YOUR_API_KEY
 
+# Use Gemini Flash with parallel workers for faster throughput
+uv run python scripts/caption_videos.py videos_dir/ --output dataset.json \
+    --captioner-type gemini_flash --num-workers 5
+
 # Caption without audio processing (video-only)
 uv run python scripts/caption_videos.py videos_dir/ --output dataset.json --no-audio
 
@@ -61,9 +65,10 @@ uv run python scripts/caption_videos.py videos_dir/ --output dataset.json --over
 - **Multiple backends**:
   - `qwen_omni` (default): Local Qwen2.5-Omni model - processes video + audio locally
   - `gemini_flash`: Google Gemini Flash API - cloud-based, requires API key
+- **Parallel captioning** (Gemini Flash only): Use `--num-workers` to run multiple API calls concurrently for faster throughput on large datasets
 - **Structured output**: Captions include visual description, speech transcription, sounds, and on-screen text
 - **Memory optimization**: 8-bit quantization option for limited VRAM
-- **Incremental processing**: Skips already-captioned files by default
+- **Incremental processing**: Skips already-captioned files by default; progress is saved every 5 videos
 - **Multiple output formats**: JSON, JSONL, CSV, or TXT
 
 **Caption format:**
@@ -74,6 +79,26 @@ The captioner produces structured captions with four sections:
 - `[SOUNDS]`: Description of music, ambient sounds, sound effects
 - `[TEXT]`: Any on-screen text visible in the video
 
+**Parallel captioning with Gemini Flash:**
+
+When using `--captioner-type gemini_flash`, you can speed up large dataset captioning by running multiple API calls at the same time using `--num-workers` (accepts 1–10, default is 1):
+
+```bash
+export GEMINI_API_KEY="your-key-here"
+
+# Caption a large dataset with 5 workers running concurrently
+uv run python scripts/caption_videos.py videos_dir/ \
+    --output dataset.json \
+    --captioner-type gemini_flash \
+    --num-workers 5
+```
+
+> [!NOTE]
+> `--num-workers` is only supported with `gemini_flash`. Using it with `qwen_omni` or any other local model will raise an error, because local GPU models are not thread-safe.
+
+> [!TIP]
+> Keep `--num-workers` between 3–5 for most use cases. Very high values (8–10) may hit Gemini API rate limits depending on your quota tier.
+
 **Environment variables (for Gemini Flash):**
 
 Set one of these to use Gemini Flash without passing `--api-key`:

packages/ltx-trainer/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ltx-trainer"
-version = "1.1.0"
+version = "1.1.1"
 description = "LTX-2 training, democratized."
 readme = "README.md"
 authors = [
@@ -48,7 +48,7 @@ build-backend = "hatchling.build"
 
 
 [tool.ruff]
-target-version = "1.1.0"
+target-version = "1.1.1"
 line-length = 120
 
 [tool.ruff.lint]

packages/ltx-trainer/scripts/caption_videos.py

@@ -19,6 +19,8 @@ Basic usage:
 Advanced usage:
     # Use Gemini Flash API (requires GEMINI_API_KEY or GOOGLE_API_KEY env var)
     caption_videos.py videos_dir/ --captioner-type gemini_flash
+    # Use Gemini Flash with parallel workers (2-10 workers, cloud API only)
+    caption_videos.py videos_dir/ --captioner-type gemini_flash --num-workers 5
     # Disable audio processing (video-only captions)
     caption_videos.py videos_dir/ --no-audio
     # Process videos with specific extensions and save as JSON
@@ -27,6 +29,7 @@ Advanced usage:
 
 import csv
 import json
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from enum import Enum
 from pathlib import Path
 
@@ -70,7 +73,7 @@ class OutputFormat(str, Enum):
     JSONL = "jsonl"  # JSON Lines file with one JSON object per line
 
 
-def caption_media(
+def caption_media(  # noqa: PLR0913
     input_path: Path,
     output_path: Path,
     captioner: MediaCaptioningModel,
@@ -81,6 +84,7 @@ def caption_media(
     clean_caption: bool,
     output_format: OutputFormat,
     override: bool,
+    num_workers: int = 1,
 ) -> None:
     """Caption videos and images using the provided captioning model.
     Args:
@@ -94,6 +98,7 @@ def caption_media(
         clean_caption: Whether to clean up captions
         output_format: Format to save the captions in
         override: Whether to override existing captions
+        num_workers: Number of parallel workers (only for cloud-based captioners like Gemini)
     """
 
     # Get list of media files to process
@@ -121,9 +126,13 @@ def caption_media(
         console.print("[bold yellow]All media already have captions. Use --override to recaption.[/]")
         return
 
-    # Process media files
+    if num_workers > 1:
+        console.print(f"Running with [bold cyan]{num_workers}[/] parallel workers.")
+
     captions = existing_captions.copy()
     successfully_captioned = 0
+    completed_since_save = 0
+
     progress = Progress(
         SpinnerColumn(),
         TextColumn("{task.description}"),
@@ -135,36 +144,47 @@ def caption_media(
         console=console,
     )
 
+    def process_one(media_file: Path) -> tuple[str, str]:
+        """Caption a single media file and return (relative_path, caption)."""
+        caption = captioner.caption(
+            path=media_file,
+            fps=fps,
+            include_audio=include_audio,
+            clean_caption=clean_caption,
+        )
+        rel_path = str(media_file.resolve().relative_to(base_dir))
+        return rel_path, caption
+
     with progress:
-        task = progress.add_task("Captioning", total=len(media_to_process))
+        task = progress.add_task(
+            f"Captioning (workers: {num_workers})" if num_workers > 1 else "Captioning",
+            total=len(media_to_process),
+        )
 
-        for i, media_file in enumerate(media_to_process):
-            progress.update(task, description=f"Captioning [bold blue]{media_file.name}[/]")
+        with ThreadPoolExecutor(max_workers=num_workers) as executor:
+            futures = {executor.submit(process_one, f): f for f in media_to_process}
 
-            try:
-                # Generate caption for the media
-                caption = captioner.caption(
-                    path=media_file,
-                    fps=fps,
-                    include_audio=include_audio,
-                    clean_caption=clean_caption,
-                )
+            for future in as_completed(futures):
+                media_file = futures[future]
+                progress.update(task, description=f"Captioning [bold blue]{media_file.name}[/]")
 
-                # Convert absolute path to relative path (relative to the output file's directory)
-                rel_path = str(media_file.resolve().relative_to(base_dir))
-                # Store the caption with the relative path as key
-                captions[rel_path] = caption
-                successfully_captioned += 1
-            except Exception as e:
-                console.print(f"[bold red]Error captioning {media_file}: {e}[/]")
+                try:
+                    rel_path, caption = future.result()
 
-            if i % SAVE_INTERVAL == 0:
-                _save_captions(captions, output_path, output_format)
+                    captions[rel_path] = caption
+                    successfully_captioned += 1
+                    completed_since_save += 1
 
-            # Advance progress bar
-            progress.advance(task)
+                    if completed_since_save >= SAVE_INTERVAL:
+                        _save_captions(captions, output_path, output_format)
+                        completed_since_save = 0
 
-    # Save captions to file
+                except Exception as e:
+                    console.print(f"[bold red]Error captioning {media_file.name}: {e}[/]")
+
+                progress.advance(task)
+
+    # Final save with everything accumulated
     _save_captions(captions, output_path, output_format)
 
     # Print summary
@@ -407,6 +427,18 @@ def main(  # noqa: PLR0913
         envvar=["GOOGLE_API_KEY", "GEMINI_API_KEY"],
         help="API key for Gemini Flash (can also use GOOGLE_API_KEY or GEMINI_API_KEY env var)",
     ),
+    num_workers: int = typer.Option(
+        1,
+        "--num-workers",
+        "-w",
+        min=1,
+        max=10,
+        help=(
+            "Number of parallel workers for captioning (1-10). "
+            "Values above 1 are only supported for cloud-based captioners (gemini_flash). "
+            "Using multiple workers with a local model will raise an error."
+        ),
+    ),
 ) -> None:
     """Auto-caption videos with audio using multimodal models.
     This script supports audio-visual captioning using:
@@ -424,6 +456,17 @@ def main(  # noqa: PLR0913
         caption_videos.py video.mp4 -o captions.json -i "Describe this video in detail"
     """
 
+    # Parallel workers are only safe for cloud-based (stateless) captioners.
+    # Local models like Qwen-Omni hold GPU state and are not thread-safe.
+    if num_workers > 1 and captioner_type != CaptionerType.GEMINI_FLASH:
+        console.print(
+            "[bold red]Error:[/] --num-workers > 1 is only supported with [bold]--captioner-type gemini_flash[/].\n"
+            "Local models (e.g. qwen_omni) run on GPU and are not thread-safe — "
+            "parallel calls would cause memory corruption or incorrect results.\n"
+            "Either set [bold]--num-workers 1[/] (default) or switch to [bold]--captioner-type gemini_flash[/]."
+        )
+        raise typer.Exit(code=1)
+
     # Determine device for local models
     device_str = device or ("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -479,6 +522,7 @@ def main(  # noqa: PLR0913
         clean_caption=clean_caption,
         output_format=output_format,
         override=override,
+        num_workers=num_workers,
     )
 
 

packages/ltx-trainer/src/ltx_trainer/sigma_tracker.py

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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]: