[feat] MiniMax-H3 Ref2VA SFT: train each clip at its own canvas and frame count (#2003)

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

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`,

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,
     })

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)

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 == []

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)]

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]