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)