From 6ef5b78020133257eba1c095056fd4afef93ac4f Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Thu, 8 Oct 2026 05:54:34 +0800 Subject: [PATCH] [bugfix] LTX-2: train temporal RoPE at the clip's fps (#1892) The legacy LTX-2 trainer did not give the clip fps to the DiT. The DiT then kept temporal RoPE positions in frames during training, but inference uses seconds (fps 24). The trainer now passes the fps of the stored latents (24 when the value is missing). The video transform stage resampled frames to train_fps but kept the source fps in the batch. The saved clips then had the source fps, and the audio was cut to the wrong length. The stage now sets the batch fps to train_fps. The new regression test runs in the CI unit lane. --- .buildkite/scripts/unit_test.sh | 1 + .../pipelines/preprocess/preprocess_stages.py | 2 + .../tests/training/test_ltx2_rope_fps.py | 134 ++++++++++++++++++ fastvideo/training/ltx2_training_pipeline.py | 26 +++- 4 files changed, 161 insertions(+), 2 deletions(-) create mode 100644 fastvideo/tests/training/test_ltx2_rope_fps.py diff --git a/.buildkite/scripts/unit_test.sh b/.buildkite/scripts/unit_test.sh index 77a11189a..b201fcb90 100644 --- a/.buildkite/scripts/unit_test.sh +++ b/.buildkite/scripts/unit_test.sh @@ -26,6 +26,7 @@ exec pytest \ ./fastvideo/tests/worker/ \ ./fastvideo/tests/training/test_runner.py \ ./fastvideo/tests/training/test_trackers.py \ + ./fastvideo/tests/training/test_ltx2_rope_fps.py \ ./fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py \ ./fastvideo/tests/inference/test_inference_regional_compile.py \ ./fastvideo/tests/attention/ \ diff --git a/fastvideo/pipelines/preprocess/preprocess_stages.py b/fastvideo/pipelines/preprocess/preprocess_stages.py index 319e94618..0c649e256 100644 --- a/fastvideo/pipelines/preprocess/preprocess_stages.py +++ b/fastvideo/pipelines/preprocess/preprocess_stages.py @@ -76,6 +76,8 @@ class VideoTransformStage(PipelineStage): video_pixel_values = video_pixel_values.float() / 255.0 batch.latents = video_pixel_values + # The frames were resampled to train_fps, so describe the clip rather than the source video. + batch.fps = [self.train_fps] * len(batch.video_loader) batch.num_frames = [video_pixel_values.shape[2]] * len(batch.video_loader) batch.height = [video_pixel_values.shape[3]] * len(batch.video_loader) batch.width = [video_pixel_values.shape[4]] * len(batch.video_loader) diff --git a/fastvideo/tests/training/test_ltx2_rope_fps.py b/fastvideo/tests/training/test_ltx2_rope_fps.py new file mode 100644 index 000000000..b48a2a77e --- /dev/null +++ b/fastvideo/tests/training/test_ltx2_rope_fps.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +"""The legacy LTX-2 trainer must hand the clip's fps to the DiT. + +The DiT divides temporal RoPE positions by ``forward_batch.fps`` from the +forward context. Inference always sets it, so a training forward without it +sees temporal positions in frames instead of seconds. +""" +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo.configs.configs import VideoLoaderType +from fastvideo.fastvideo_args import WorkloadType +from fastvideo.forward_context import get_forward_context +from fastvideo.pipelines.pipeline_batch_info import TrainingBatch +from fastvideo.pipelines.preprocess.ltx2.ltx2_preprocess_pipelines import LTX2AudioEncodingStage +from fastvideo.pipelines.preprocess.preprocess_stages import VideoTransformStage +from fastvideo.training import ltx2_training_pipeline as ltx2_module +from fastvideo.training.ltx2_training_pipeline import LTX2TrainingPipeline +from fastvideo.training.trackers import DummyTracker +from fastvideo.workflow.preprocess.preprocess_workflow_ltx2_t2v import LTX2PrecomputedSaver + + +class _FpsRecordingDiT(torch.nn.Module): + + def __init__(self) -> None: + super().__init__() + self.scale = torch.nn.Parameter(torch.ones(())) + self.seen_fps: list[float | None] = [] + + def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + forward_batch = get_forward_context().forward_batch + self.seen_fps.append(None if forward_batch is None else forward_batch.fps) + return hidden_states * self.scale + + +def _pipeline(monkeypatch: pytest.MonkeyPatch) -> LTX2TrainingPipeline: + monkeypatch.setattr(ltx2_module, "get_local_torch_device", lambda: torch.device("cpu")) + monkeypatch.setattr(ltx2_module, "get_world_group", + lambda: SimpleNamespace(all_reduce=lambda tensor, op=None: tensor)) + pipeline = object.__new__(LTX2TrainingPipeline) + pipeline.transformer = _FpsRecordingDiT() + pipeline.tracker = DummyTracker() + pipeline.with_audio = False + pipeline.training_args = SimpleNamespace(num_latent_t=2, + ltx2_first_frame_conditioning_p=0.0, + gradient_accumulation_steps=1) + pipeline.noise_random_generator = torch.Generator("cpu").manual_seed(0) + pipeline.noise_gen_cuda = torch.Generator("cpu").manual_seed(0) + return pipeline + + +def _pt_batch(fps: torch.Tensor | None) -> dict: + batch_size = 1 if fps is None else fps.numel() + latents = {"latents": torch.randn(batch_size, 4, 2, 2, 2)} + if fps is not None: + latents["fps"] = fps + return { + "latents": latents, + "conditions": { + "video_prompt_embeds": torch.randn(batch_size, 3, 8), + "audio_prompt_embeds": torch.randn(batch_size, 3, 8), + "prompt_attention_mask": torch.ones(batch_size, 3), + }, + } + + +def _train_step_fps(pipeline: LTX2TrainingPipeline, batch: dict) -> float | None: + training_batch = pipeline._get_next_batch_pt(batch, TrainingBatch()) + training_batch = pipeline._prepare_dit_inputs(training_batch) + training_batch = pipeline._build_input_kwargs(training_batch) + training_batch.total_loss = 0.0 + pipeline._transformer_forward_and_compute_loss(training_batch) + return pipeline.transformer.seen_fps[-1] + + +def test_training_forward_sees_the_clip_fps(monkeypatch: pytest.MonkeyPatch) -> None: + pipeline = _pipeline(monkeypatch) + assert _train_step_fps(pipeline, _pt_batch(torch.tensor([24.0]))) == 24.0 + assert _train_step_fps(pipeline, _pt_batch(torch.tensor([30.0]))) == 30.0 + + +def test_batch_without_fps_uses_the_preset_fps(monkeypatch: pytest.MonkeyPatch) -> None: + pipeline = _pipeline(monkeypatch) + assert _train_step_fps(pipeline, _pt_batch(None)) == 24.0 + + +def test_mixed_fps_batch_uses_the_first_sample(monkeypatch: pytest.MonkeyPatch) -> None: + pipeline = _pipeline(monkeypatch) + assert _train_step_fps(pipeline, _pt_batch(torch.tensor([25.0, 30.0]))) == 25.0 + + +class _Clip(str): + """A 60-frame video path whose frames read like a torchcodec decoder's.""" + + frames = torch.zeros(60, 3, 16, 16, dtype=torch.uint8) + + def get_frames_at(self, indices) -> SimpleNamespace: + return SimpleNamespace(data=self.frames[list(indices)]) + + +def test_resampled_clip_keeps_train_fps_for_audio_and_saving(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None: + requested_audio_seconds: list[float] = [] + + def extract_audio(video_path: str, target_duration: float) -> None: + requested_audio_seconds.append(target_duration) + + monkeypatch.setattr(LTX2AudioEncodingStage, "_extract_audio", staticmethod(extract_audio)) + batch = SimpleNamespace(data_type="video", + video_loader=[_Clip("clip.mp4")], + video_file_name=["clip.mp4"], + fps=[30.0], + num_frames=[60], + height=[16], + width=[16], + prompt_embeds=[torch.zeros(1, 3, 8)], + prompt_attention_mask=[torch.ones(1, 3)], + extra={}) + args = SimpleNamespace(preprocess_config=SimpleNamespace(video_loader_type=VideoLoaderType.TORCHCODEC), + workload_type=WorkloadType.T2V) + + VideoTransformStage(train_fps=24, num_frames=9, max_height=16, max_width=16, + do_temporal_sample=False).forward(batch, args) + LTX2AudioEncodingStage(torch.nn.Linear(1, 1), audio_processor=None, fallback_fps=24).forward(batch, args) + LTX2PrecomputedSaver(tmp_path).save_batch(batch) + + # 9 frames sampled at 24 fps from a 30 fps source cover 9 / 24 s. + assert requested_audio_seconds == [pytest.approx(9 / 24)] + assert torch.load(tmp_path / "latents" / "clip.pt")["fps"] == 24.0 + assert batch.num_frames == [9] + assert batch.fps == [24] diff --git a/fastvideo/training/ltx2_training_pipeline.py b/fastvideo/training/ltx2_training_pipeline.py index 18f18b276..c352bbff5 100644 --- a/fastvideo/training/ltx2_training_pipeline.py +++ b/fastvideo/training/ltx2_training_pipeline.py @@ -13,7 +13,7 @@ from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger from fastvideo.models.dits.ltx2 import VideoLatentShape from fastvideo.pipelines.basic.ltx2.ltx2_pipeline import LTX2Pipeline -from fastvideo.pipelines.pipeline_batch_info import TrainingBatch +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch from fastvideo.training.training_pipeline import TrainingPipeline from fastvideo.training.trackers import (DummyTracker, TrackerType, Trackers, initialize_trackers) from fastvideo.training.training_utils import (clip_grad_norm_while_handling_failing_dtensor_cases, get_scheduler) @@ -21,6 +21,9 @@ from fastvideo.utils import set_random_seed logger = init_logger(__name__) +# Preset fps of the LTX-2 checkpoints, used when a batch carries no fps. +_DEFAULT_ROPE_FPS = 24.0 + class LTX2TrainingPipeline(TrainingPipeline): """Training pipeline for LTX-2 text-to-video (optional audio).""" @@ -37,6 +40,8 @@ class LTX2TrainingPipeline(TrainingPipeline): text_encoder: torch.nn.Module with_audio: bool = False tracker: TrackerType + # fps of the current batch; the DiT divides temporal RoPE positions by it. + _rope_fps: float = _DEFAULT_ROPE_FPS def initialize_pipeline(self, fastvideo_args: FastVideoArgs): # TODO (David): Change to port LTX2 scheduler into self.modules["scheduler"] @@ -257,6 +262,7 @@ class LTX2TrainingPipeline(TrainingPipeline): training_batch.encoder_attention_mask = attention_mask.to(device, dtype=torch.bfloat16) training_batch.infos = [] training_batch.raw_latent_shape = latents.shape + self._rope_fps = _DEFAULT_ROPE_FPS return training_batch def _get_next_batch_pt( @@ -297,8 +303,20 @@ class LTX2TrainingPipeline(TrainingPipeline): else: training_batch.infos = [] training_batch.raw_latent_shape = latents.shape + self._rope_fps = self._batch_fps(batch["latents"].get("fps")) return training_batch + @staticmethod + def _batch_fps(fps: torch.Tensor | list[float] | float | None) -> float: + if fps is None: + return _DEFAULT_ROPE_FPS + values = [float(v) for v in torch.as_tensor(fps).flatten().tolist()] + if not values or values[0] <= 0: + return _DEFAULT_ROPE_FPS + if any(v != values[0] for v in values): + logger.warning("Batch mixes fps %s; using %s for RoPE", values, values[0]) + return values[0] + def _normalize_dit_input(self, training_batch: TrainingBatch) -> TrainingBatch: return training_batch @@ -432,8 +450,12 @@ class LTX2TrainingPipeline(TrainingPipeline): assert training_batch.noisy_model_input is not None assert training_batch.noise is not None + # Without fps in the forward context the DiT leaves temporal RoPE positions in + # frames, while inference divides them by the preset fps into seconds. with self.tracker.timed("timing/forward_backward"), set_forward_context( - current_timestep=training_batch.current_timestep, attn_metadata=training_batch.attn_metadata): + current_timestep=training_batch.current_timestep, + attn_metadata=training_batch.attn_metadata, + forward_batch=ForwardBatch(data_type="video", fps=self._rope_fps)): with torch.autocast("cuda", dtype=training_batch.latents.dtype), torch.autograd.set_detect_anomaly(True): outputs = self.transformer(**input_kwargs)