diff --git a/docs/training/minimax_h3_ref2va_preprocess.md b/docs/training/minimax_h3_ref2va_preprocess.md index aa5033fd4..f0fbdcc28 100644 --- a/docs/training/minimax_h3_ref2va_preprocess.md +++ b/docs/training/minimax_h3_ref2va_preprocess.md @@ -12,6 +12,17 @@ format, which stores already augmented `ref_visual_anchor` tensors. Keep using `preprocess_minimax_h3_ref2va` for that format. Do not rename or directly combine shards from these two schemas. +The SFT preprocessor keeps each target clip's own geometry: the canvas follows +the clip's aspect ratio at the inference short edge (`resolve_canvas_size`), and +the frame count is the longest `17 * n + 5` that fits the clip, capped at 362 +(about 15 seconds; a longer clip trains on its start). Clips under 124 frames at +24 fps fail. Target audio is cut to the samples covered by that frame count's +whole audio latents, so the record matches what inference generates. To train +on those lengths, leave `training.data.num_latent_t` and `num_frames` unset in +the training config. When they are set, every row is cropped to that length; +for example, the 124-frame overfit values crop a 10-second clip to its first +5.17 seconds while its caption still describes all 10. + ## Input Provide your own local MiniMax-H3 checkpoint with `vae`, `audio_vae`, `tokenizer`, diff --git a/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py b/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py index 9ef03db80..cd786f313 100644 --- a/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py @@ -25,7 +25,15 @@ from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2va from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.models.loader.component_loader import PipelineComponentLoader from fastvideo.pipelines import ForwardBatch -from fastvideo.pipelines.basic.minimax_h3.packing import MINIMAX_H3_FPS +from fastvideo.pipelines.basic.minimax_h3.packing import ( + MINIMAX_H3_AUDIO_LATENTS_PER_SECOND, + MINIMAX_H3_FPS, + MINIMAX_H3_FRAMES_PER_CHUNK, + MINIMAX_H3_LATENTS_PER_CHUNK, + MINIMAX_H3_MAX_ALIGNED_FRAMES, + MINIMAX_H3_MIN_ALIGNED_FRAMES, + audio_latent_num_frames, +) from fastvideo.pipelines.basic.minimax_h3.reference import ( decode_reference_video, prepare_reference_frames, @@ -60,20 +68,41 @@ def _init_single_process_distributed() -> None: maybe_init_distributed_environment_and_model_parallel(1, 1) -def load_training_media(path: Path) -> tuple[np.ndarray, torch.Tensor]: +def native_num_frames(available: int) -> int: + """Return the longest H3-native ``17 * n + 5`` frame count within ``available``. + + Counts are capped at the largest accepted bucket, so a long clip trains on + its first 15 seconds instead of failing. + """ + if available < MINIMAX_H3_MIN_ALIGNED_FRAMES: + raise ValueError(f"Training video has {available} frames at {MINIMAX_H3_FPS} FPS; " + f"at least {MINIMAX_H3_MIN_ALIGNED_FRAMES} are required") + num_frames = min(available, MINIMAX_H3_MAX_ALIGNED_FRAMES) + while num_frames % MINIMAX_H3_FRAMES_PER_CHUNK != MINIMAX_H3_LATENTS_PER_CHUNK: + num_frames -= 1 + return num_frames + + +def load_training_media(path: Path, num_frames: int | None = NUM_FRAMES) -> tuple[np.ndarray, torch.Tensor]: """Decode paired targets at the geometry and rates consumed by H3. - The video uses exactly 124 frames, and the audio is capped at the duration - of those frames before variational autoencoder (VAE) encoding. + With the default ``num_frames`` the video must resolve to the fixed + 124-frame 768x1344 overfit geometry. ``num_frames=None`` keeps the clip's + own canvas (``resolve_canvas_size``) and its longest native frame count. + The audio is capped at the duration of the selected frames, and at the + samples their ``audio_latent_num_frames`` cover, before VAE encoding. """ decoded_frames, source_fps, soundtrack = decode_reference_video(path) resampled_frames = resample_reference_frames(decoded_frames, source_fps) - if resampled_frames.shape[0] < NUM_FRAMES: + fixed_geometry = num_frames is not None + if num_frames is None: + num_frames = native_num_frames(resampled_frames.shape[0]) + if resampled_frames.shape[0] < num_frames: raise ValueError(f"Training video {path} has {resampled_frames.shape[0]} frames at {MINIMAX_H3_FPS} FPS; " - f"{NUM_FRAMES} are required") - frames = prepare_reference_frames(resampled_frames, NUM_FRAMES) - if frames.shape != (NUM_FRAMES, VIDEO_HEIGHT, VIDEO_WIDTH, 3): - raise ValueError(f"Training video must resolve to {(NUM_FRAMES, VIDEO_HEIGHT, VIDEO_WIDTH, 3)}, " + f"{num_frames} are required") + frames = prepare_reference_frames(resampled_frames, num_frames) + if fixed_geometry and frames.shape != (num_frames, VIDEO_HEIGHT, VIDEO_WIDTH, 3): + raise ValueError(f"Training video must resolve to {(num_frames, VIDEO_HEIGHT, VIDEO_WIDTH, 3)}, " f"got {tuple(frames.shape)}") if soundtrack is None: raise ValueError(f"Training video {path} requires a soundtrack") @@ -82,8 +111,14 @@ def load_training_media(path: Path) -> tuple[np.ndarray, torch.Tensor]: waveform, source_sample_rate, AUDIO_SAMPLE_RATE, - max_duration=NUM_FRAMES / MINIMAX_H3_FPS, + max_duration=num_frames / MINIMAX_H3_FPS, ) + # The audio VAE pads to whole latents, so audio past the last whole latent + # would add one more than inference generates for these frames (e.g. 264 + # instead of 263 at 158 frames). The packed Ref2VA preprocessor applies the + # same cap. + samples_per_latent = AUDIO_SAMPLE_RATE // MINIMAX_H3_AUDIO_LATENTS_PER_SECOND + waveform = waveform[..., :audio_latent_num_frames(num_frames) * samples_per_latent].contiguous() return frames, waveform @@ -234,6 +269,9 @@ def build_parquet_record( video_latents: torch.Tensor, audio_latents: torch.Tensor, text_embedding: torch.Tensor, + num_frames: int = NUM_FRAMES, + height: int = VIDEO_HEIGHT, + width: int = VIDEO_WIDTH, ) -> dict[str, Any]: """Serialize one synchronized H3 sample for the Parquet schema collator. @@ -255,10 +293,10 @@ def build_parquet_record( "file_name": file_name, "caption": caption, "media_type": "video_with_audio", - "width": VIDEO_WIDTH, - "height": VIDEO_HEIGHT, - "num_frames": NUM_FRAMES, - "duration_sec": NUM_FRAMES / MINIMAX_H3_FPS, + "width": width, + "height": height, + "num_frames": num_frames, + "duration_sec": num_frames / MINIMAX_H3_FPS, "fps": float(MINIMAX_H3_FPS), "audio_sample_rate": AUDIO_SAMPLE_RATE, }) diff --git a/fastvideo/pipelines/preprocess/preprocess_minimax_h3_ref2va.py b/fastvideo/pipelines/preprocess/preprocess_minimax_h3_ref2va.py index c45032d4f..9bab4a203 100644 --- a/fastvideo/pipelines/preprocess/preprocess_minimax_h3_ref2va.py +++ b/fastvideo/pipelines/preprocess/preprocess_minimax_h3_ref2va.py @@ -305,6 +305,9 @@ def build_ref2va_parquet_record( ref_visual_anchor: torch.Tensor, ref_audio_anchor: torch.Tensor, references: list[MiniMaxH3PreparedReference], + num_frames: int = NUM_FRAMES, + height: int = VIDEO_HEIGHT, + width: int = VIDEO_WIDTH, ) -> dict[str, Any]: record = build_parquet_record( file_name=file_name, @@ -312,6 +315,9 @@ def build_ref2va_parquet_record( video_latents=video_latents, audio_latents=audio_latents, text_embedding=text_embedding, + num_frames=num_frames, + height=height, + width=width, ) record["schema_version"] = MINIMAX_H3_REF2VA_SCHEMA_VERSION record["text_token_tags"] = text_token_tags.to(dtype=torch.long, device="cpu").tolist() @@ -326,8 +332,13 @@ def validate_record_contract(record: dict[str, Any]) -> None: if missing: raise ValueError(f"Ref2VA record is missing schema fields: {missing}") expected_target_shapes = { - "vae_latent_shape": [24, video_latent_num_frames(NUM_FRAMES), VIDEO_HEIGHT // 16, VIDEO_WIDTH // 16], - "audio_latent_shape": [2, 32, audio_latent_num_frames(NUM_FRAMES)], + "vae_latent_shape": [ + 24, + video_latent_num_frames(record["num_frames"]), + record["height"] // 16, + record["width"] // 16, + ], + "audio_latent_shape": [2, 32, audio_latent_num_frames(record["num_frames"])], } for name, expected in expected_target_shapes.items(): if record[name] != expected: @@ -337,9 +348,12 @@ def validate_record_contract(record: dict[str, Any]) -> None: collate_minimax_h3_ref2va_rows([record]) -def _prepare_references(references: tuple[MiniMaxH3RawReference, ...], ) -> list[MiniMaxH3PreparedReference]: +def _prepare_references( + references: tuple[MiniMaxH3RawReference, ...], + num_frames: int = NUM_FRAMES, +) -> list[MiniMaxH3PreparedReference]: raw_references = build_minimax_h3_references(references) - prepared = [prepare_reference(reference, NUM_FRAMES, AUDIO_SAMPLE_RATE) for reference in raw_references] + prepared = [prepare_reference(reference, num_frames, AUDIO_SAMPLE_RATE) for reference in raw_references] for index, (raw, resolved) in enumerate(zip(references, prepared, strict=True)): expected_audio = raw.media_type in {"audio", "video_audio"} if resolved.has_audio != expected_audio: @@ -377,8 +391,9 @@ def _build_record( fastvideo_args: FastVideoArgs, patch_size: tuple[int, int, int], ) -> dict[str, Any]: - target_frames, target_waveform = load_training_media(sample.target_video_path) - references = _prepare_references(sample.references) + target_frames, target_waveform = load_training_media(sample.target_video_path, num_frames=None) + num_frames, height, width = (int(value) for value in target_frames.shape[:3]) + references = _prepare_references(sample.references, num_frames) # The established single-sample helpers deliberately release each large # component before the next one is loaded, keeping this multi-sample path @@ -410,6 +425,9 @@ def _build_record( ref_visual_anchor=ref_visual_anchor, ref_audio_anchor=ref_audio_anchor, references=references, + num_frames=num_frames, + height=height, + width=width, ) record["id"] = sample.sample_id validate_record_contract(record) diff --git a/fastvideo/tests/train/models/test_minimax_h3_row_length_latents.py b/fastvideo/tests/train/models/test_minimax_h3_row_length_latents.py new file mode 100644 index 000000000..eba9b3b5d --- /dev/null +++ b/fastvideo/tests/train/models/test_minimax_h3_row_length_latents.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Without a fixed training shape, H3 rows train at their own length instead of being cropped.""" +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo.pipelines.basic.minimax_h3.packing import audio_latent_num_frames, video_latent_num_frames +from fastvideo.train.models.minimax_h3.minimax_h3 import MiniMaxH3Model + + +def _model(num_latent_t: int, num_frames: int) -> MiniMaxH3Model: + model = MiniMaxH3Model.__new__(MiniMaxH3Model) + model.training_config = SimpleNamespace(data=SimpleNamespace( + native_shape_bucketing=False, num_latent_t=num_latent_t, num_frames=num_frames)) + return model + + +def _row(num_frames: int, audio_frames: int | None = None) -> dict[str, torch.Tensor]: + if audio_frames is None: + audio_frames = audio_latent_num_frames(num_frames) + return { + "vae_latent": torch.zeros(1, 24, video_latent_num_frames(num_frames), 4, 2), + "audio_latent": torch.zeros(1, 2, 32, audio_frames), + } + + +def _resolve(model: MiniMaxH3Model, raw_batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + return model._resolve_clean_latents(raw_batch, "data", torch.float32, torch.device("cpu")) + + +@pytest.mark.parametrize("num_frames", [124, 158, 243, 362]) +def test_unset_shape_keeps_each_row_length(num_frames: int) -> None: + video, audio = _resolve(_model(0, 0), _row(num_frames)) + + assert video.shape[2] == video_latent_num_frames(num_frames) + assert audio.shape[-1] == audio_latent_num_frames(num_frames) + + +def test_unset_shape_rejects_mismatched_latent_clocks() -> None: + with pytest.raises(ValueError, match="158 video frames need 263"): + _resolve(_model(0, 0), _row(158, audio_frames=264)) + + +def test_fixed_shape_still_crops_and_warns_once(monkeypatch: pytest.MonkeyPatch) -> None: + import fastvideo.train.models.minimax_h3.minimax_h3 as module + + warnings = [] + monkeypatch.setattr(module.logger, "warning_once", warnings.append) + + video, audio = _resolve(_model(37, 124), _row(243)) + + assert video.shape[2] == 37 + assert audio.shape[-1] == audio_latent_num_frames(124) + assert len(warnings) == 1 and "Cropping vae_latent from 72 to" in warnings[0] + + +def test_fixed_shape_at_its_own_length_does_not_warn(monkeypatch: pytest.MonkeyPatch) -> None: + import fastvideo.train.models.minimax_h3.minimax_h3 as module + + warnings = [] + monkeypatch.setattr(module.logger, "warning_once", warnings.append) + + _resolve(_model(37, 124), _row(124)) + + assert warnings == [] diff --git a/fastvideo/tests/workflow/test_minimax_h3_native_geometry_preprocess.py b/fastvideo/tests/workflow/test_minimax_h3_native_geometry_preprocess.py new file mode 100644 index 000000000..4e9210337 --- /dev/null +++ b/fastvideo/tests/workflow/test_minimax_h3_native_geometry_preprocess.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Ref2VA SFT preprocessing keeps each target clip's own canvas and native frame count.""" + +from pathlib import Path + +import numpy as np +import pytest +import torch + +from fastvideo.pipelines.basic.minimax_h3.packing import ( + MINIMAX_H3_TEXT_TAG, + audio_latent_num_frames, + resolve_canvas_size, + video_latent_num_frames, +) +from fastvideo.pipelines.basic.minimax_h3.ref2va_manifest import MiniMaxH3Ref2VARawSample +from fastvideo.pipelines.preprocess import preprocess_minimax_h3_overfit as overfit +from fastvideo.pipelines.preprocess import preprocess_minimax_h3_ref2va as ref2va + +_SAMPLES_PER_AUDIO_LATENT = 800 # 32 kHz / 40 latents per second + + +@pytest.mark.parametrize( + ("available", "expected"), + [(124, 124), (140, 124), (141, 141), (157, 141), (209, 209), (362, 362), (1000, 362)], +) +def test_native_num_frames_rounds_down_to_17n_plus_5(available: int, expected: int) -> None: + assert overfit.native_num_frames(available) == expected + + +def test_native_num_frames_rejects_clips_under_five_seconds() -> None: + with pytest.raises(ValueError, match="at least 124"): + overfit.native_num_frames(123) + + +def _stub_media(monkeypatch: pytest.MonkeyPatch, num_source_frames: int, seconds: float) -> None: + """Decode to tiny 30-fps frames and a 48 kHz stereo soundtrack; skip the canvas resize.""" + frames = np.zeros((num_source_frames, 8, 4, 3), dtype=np.uint8) + waveform = torch.zeros(2, int(seconds * 48000)) + monkeypatch.setattr(overfit, "decode_reference_video", lambda _path: (frames, 30.0, (waveform, 48000))) + monkeypatch.setattr(overfit, "prepare_reference_frames", lambda decoded, num_frames: decoded[:num_frames]) + + +def test_native_media_keeps_native_length_and_whole_audio_latents(monkeypatch: pytest.MonkeyPatch) -> None: + # 210 frames at 30 fps resample to 168 at 24 fps; the longest native count is 158. + _stub_media(monkeypatch, num_source_frames=210, seconds=7.0) + + frames, waveform = overfit.load_training_media(Path("clip.mp4"), num_frames=None) + + assert frames.shape[0] == 158 + # 158 frames last 210666 samples, past the 263rd whole latent; the VAE would pad a 264th. + assert audio_latent_num_frames(158) == 263 + assert waveform.shape == (2, 263 * _SAMPLES_PER_AUDIO_LATENT) + + +def test_fixed_overfit_media_audio_length_is_unchanged(monkeypatch: pytest.MonkeyPatch) -> None: + _stub_media(monkeypatch, num_source_frames=200, seconds=7.0) + monkeypatch.setattr( + overfit, + "prepare_reference_frames", + lambda decoded, num_frames: np.zeros((num_frames, overfit.VIDEO_HEIGHT, overfit.VIDEO_WIDTH, 3), np.uint8), + ) + + frames, waveform = overfit.load_training_media(Path("clip.mp4")) + + # 124 frames end inside the 207th audio latent, so the whole-latent cap cuts nothing. + uncapped = overfit.prepare_reference_waveform(torch.zeros(2, 7 * 48000), + 48000, + overfit.AUDIO_SAMPLE_RATE, + max_duration=overfit.NUM_FRAMES / 24) + assert frames.shape[0] == overfit.NUM_FRAMES + assert uncapped.shape[-1] < audio_latent_num_frames(overfit.NUM_FRAMES) * _SAMPLES_PER_AUDIO_LATENT + assert waveform.shape == uncapped.shape + + +def test_build_record_uses_the_clip_geometry(monkeypatch: pytest.MonkeyPatch) -> None: + """A portrait 158-frame clip yields a record the training collator accepts.""" + num_frames = 158 + height, width = resolve_canvas_size(9, 16) + target = np.zeros((num_frames, height, width, 3), dtype=np.uint8) + waveform = torch.zeros(2, audio_latent_num_frames(num_frames) * _SAMPLES_PER_AUDIO_LATENT) + monkeypatch.setattr(ref2va, "load_training_media", lambda _path, num_frames: (target, waveform)) + monkeypatch.setattr(ref2va, "_prepare_references", lambda _refs, _num_frames: []) + monkeypatch.setattr( + ref2va, + "encode_video_latents", + lambda frames, *_args: torch.zeros(24, video_latent_num_frames(frames.shape[0]), frames.shape[1] // 16, + frames.shape[2] // 16), + ) + monkeypatch.setattr( + ref2va, + "encode_audio_latents", + lambda audio, *_args: torch.zeros(2, 32, audio.shape[-1] // _SAMPLES_PER_AUDIO_LATENT), + ) + monkeypatch.setattr(ref2va, "encode_ref_visual_anchor", lambda *_args: torch.empty(0, 96)) + monkeypatch.setattr(ref2va, "encode_ref_audio_anchor", lambda *_args: torch.empty(0, 32)) + monkeypatch.setattr( + ref2va, + "encode_ref2va_conditioning", + lambda *_args: (torch.zeros(4, 5120), torch.full((4, ), MINIMAX_H3_TEXT_TAG)), + ) + sample = MiniMaxH3Ref2VARawSample( + sample_id="portrait", + target_file="portrait.mp4", + target_video_path=Path("portrait.mp4"), + caption="A person talks to the camera.", + references=(), + ) + + record = ref2va._build_record( + sample, + model_path=Path("model"), + model_index={}, + fastvideo_args=None, + patch_size=(1, 2, 2), + ) + + assert (record["num_frames"], record["height"], record["width"]) == (num_frames, height, width) + assert record["vae_latent_shape"] == [24, video_latent_num_frames(num_frames), height // 16, width // 16] + assert record["audio_latent_shape"] == [2, 32, audio_latent_num_frames(num_frames)] diff --git a/fastvideo/train/models/minimax_h3/minimax_h3.py b/fastvideo/train/models/minimax_h3/minimax_h3.py index 313dec6ac..a8c022202 100644 --- a/fastvideo/train/models/minimax_h3/minimax_h3.py +++ b/fastvideo/train/models/minimax_h3/minimax_h3.py @@ -19,6 +19,7 @@ from fastvideo.dataset.sequence_capacity import select_sequence_capacity, valida from fastvideo.distributed import get_sp_group from fastvideo.forward_context import set_forward_context from fastvideo.layers.pdd import shifted_noise_amount +from fastvideo.logger import init_logger from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler from fastvideo.pipelines import TrainingBatch from fastvideo.pipelines.basic.minimax_h3.packing import ( @@ -49,6 +50,8 @@ from fastvideo.train.utils.module_state import apply_trainable from fastvideo.train.utils.moduleloader import load_module_from_path from fastvideo.train.utils.lora import finalize_lora_training, LoraConfig +logger = init_logger(__name__) + if TYPE_CHECKING: from fastvideo.train.utils.training_config import TrainingConfig @@ -370,10 +373,26 @@ class MiniMaxH3Model(ModelBase): if latents_source == "data" and native_shapes: self._validate_native_latents(raw_batch, video_latents, audio_latents) + elif latents_source == "data" and data_config.num_latent_t <= 0 and data_config.num_frames <= 0: + # With no fixed shape configured, each row trains at its own + # length, as the Ref2VA SFT preprocessor records it. Its video and + # audio latents must still describe the same frame count. + num_video_latents = int(video_latents.shape[2]) + if (num_video_latents - 2) % 5: + raise ValueError(f"vae_latent has {num_video_latents} frames; MiniMax H3 encodes 5 * n + 2") + num_frames = (num_video_latents - 2) // 5 * 17 + 5 + if audio_latents.shape[-1] != audio_latent_num_frames(num_frames): + raise ValueError(f"audio_latent has {audio_latents.shape[-1]} frames; {num_frames} video frames " + f"need {audio_latent_num_frames(num_frames)}") elif not native_shapes: # Preserve the legacy fixed-shape contract for configs that have # not opted into exact-shape bucketing. if data_config.num_latent_t > 0: + if video_latents.shape[2] > data_config.num_latent_t: + logger.warning_once( + f"Cropping vae_latent from {video_latents.shape[2]} to training.data.num_latent_t=" + f"{data_config.num_latent_t} frames, and audio to match. Leave num_latent_t and " + "num_frames unset to train each row at its own length.") video_latents = video_latents[:, :, :data_config.num_latent_t] expected_audio_frames = audio_latent_num_frames(data_config.num_frames) audio_latents = audio_latents[:, :, :, :expected_audio_frames]