# SPDX-License-Identifier: Apache-2.0 """CPU-only unit tests for :mod:`fastvideo.train.callbacks.validation`. Covers the parts of ``ValidationCallback`` that don't need a real pipeline or distributed init: * constructor type coercions and defaults, * ``on_validation_begin`` gating logic (every_steps + modulo), * ``_find_ema_callback`` lookup via ``_callback_dict``, * ``state_dict`` / ``load_state_dict`` rng round-trip. The heavy ``_run_validation`` path needs a real diffusion pipeline plus distributed init and is exercised by Phase 2/3 tests. """ from __future__ import annotations from types import SimpleNamespace import numpy as np import pytest import torch from fastvideo.api.sampling_param import SamplingParam from fastvideo.train.callbacks.callback import CallbackDict from fastvideo.train.callbacks.ema import EMACallback from fastvideo.train.callbacks.validation import ( DEFAULT_VALIDATION_VBENCH_METRICS, SYNTHETIC_OPTICAL_FLOW_LOG_KEYS, SYNTHETIC_OPTICAL_FLOW_METRIC, ValidationCallback, _ValidationMetricStats, ) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- _PIPE_TARGET = "fastvideo.pipelines.basic.wan.wan_pipeline.WanPipeline" def _make_callback( *, every_steps: int = 100, sampling_steps: list[int] | None = None, guidance_scale: float | None = None, num_frames: int | None = None, num_videos_per_prompt: int = 1, use_validation_media_conditioning: bool = True, sampling_timesteps: list[int] | None = None, output_dir: str | None = None, overlay_actions: bool = False, run_at_start: bool = True, ) -> ValidationCallback: """Construct a callback with the shared schedule and media-test defaults.""" return ValidationCallback( pipeline_target=_PIPE_TARGET, dataset_file="/tmp/does_not_exist.json", every_steps=every_steps, sampling_steps=sampling_steps, guidance_scale=guidance_scale, num_frames=num_frames, num_videos_per_prompt=num_videos_per_prompt, use_validation_media_conditioning=use_validation_media_conditioning, sampling_timesteps=sampling_timesteps, output_dir=output_dir, overlay_actions=overlay_actions, run_at_start=run_at_start, ) # --------------------------------------------------------------------------- # A. Constructor coercions / defaults # --------------------------------------------------------------------------- class TestConstructor: def test_defaults(self) -> None: """Verify defaults preserve baseline validation and media conditioning.""" cb = _make_callback() assert cb.pipeline_target == _PIPE_TARGET assert cb.dataset_file == "/tmp/does_not_exist.json" assert cb.every_steps == 100 assert cb.sampling_steps == [40] assert cb.guidance_scale is None assert cb.num_frames is None assert cb.num_videos_per_prompt == 1 assert cb.use_validation_media_conditioning is True assert cb.run_at_start is True assert cb.sampling_timesteps is None assert cb.output_dir is None assert cb.overlay_actions is False assert cb.offload_training_state is False assert cb.unload_pipeline_after_validation is False assert cb.attn_qat_infer is False # Lazy fields not yet populated. assert cb._pipeline is None assert cb._sampling_param is None assert cb.validation_random_generator is None assert cb.metrics_config.enabled is False def test_string_inputs_are_coerced(self) -> None: """Verify YAML-compatible scalar strings become callback value types.""" # YAML often produces strings for numeric fields; the # constructor must coerce them. cb = ValidationCallback( pipeline_target=_PIPE_TARGET, dataset_file="x.json", every_steps="50", # type: ignore[arg-type] sampling_steps=["20", "40"], # type: ignore[arg-type] guidance_scale="4.5", # type: ignore[arg-type] num_frames="77", # type: ignore[arg-type] num_videos_per_prompt="2", # type: ignore[arg-type] use_validation_media_conditioning="false", # type: ignore[arg-type] sampling_timesteps=["1000", "500"], run_at_start="false", # type: ignore[arg-type] overlay_actions=1, # type: ignore[arg-type] offload_training_state="1", # type: ignore[arg-type] unload_pipeline_after_validation="false", # type: ignore[arg-type] attn_qat_infer="true", # type: ignore[arg-type] ) assert cb.every_steps == 50 assert cb.sampling_steps == [20, 40] assert cb.guidance_scale == 4.5 assert cb.num_frames == 77 assert cb.num_videos_per_prompt == 2 assert cb.use_validation_media_conditioning is False assert cb.run_at_start is False assert cb.sampling_timesteps == [1000, 500] assert cb.overlay_actions is True assert cb.offload_training_state is True assert cb.unload_pipeline_after_validation is False assert cb.attn_qat_infer is True def test_init_rejects_nonpositive_video_count(self) -> None: """Verify every prompt requests at least one generated video.""" with pytest.raises(ValueError, match="num_videos_per_prompt must be positive"): _make_callback(num_videos_per_prompt=0) def test_pipeline_kwargs_collected(self) -> None: cb = ValidationCallback( pipeline_target=_PIPE_TARGET, dataset_file="x.json", extra_arg=123, another="value", ) # Unknown kwargs are stashed for the pipeline factory. assert cb.pipeline_kwargs == { "extra_arg": 123, "another": "value", } def test_metrics_true_uses_default_vbench_subset(self) -> None: cb = ValidationCallback( pipeline_target=_PIPE_TARGET, dataset_file="x.json", metrics=True, ) assert cb.metrics_config.enabled is True assert cb.metrics_config.names == DEFAULT_VALIDATION_VBENCH_METRICS assert "metrics" not in cb.pipeline_kwargs def test_metrics_mapping_is_coerced(self) -> None: cb = ValidationCallback( pipeline_target=_PIPE_TARGET, dataset_file="x.json", metrics={ "names": "vbench.aesthetic_quality", "device": "cpu", "calibration_path": "/tmp/calibration.json", "mouse_pitch_sign": "-1", "skip_missing_deps": False, "strict": True, "unload_after_validation": False, "loader_threads": "2", "prefetch_factor": "3", "log_prefix": "custom/validation", }, ) assert cb.metrics_config.enabled is True assert cb.metrics_config.names == ["vbench.aesthetic_quality"] assert cb.metrics_config.device == "cpu" assert cb.metrics_config.calibration_path == "/tmp/calibration.json" assert cb.metrics_config.mouse_pitch_sign == -1 assert cb.metrics_config.skip_missing_deps is False assert cb.metrics_config.strict is True assert cb.metrics_config.unload_after_validation is False assert cb.metrics_config.loader_threads == 2 assert cb.metrics_config.prefetch_factor == 3 assert cb.metrics_config.log_prefix == "custom/validation" # --------------------------------------------------------------------------- # B. on_validation_begin gating # --------------------------------------------------------------------------- class _NoRunValidation(ValidationCallback): """Subclass that records ``_run_validation`` calls instead of actually running them — lets us assert the gating logic without a real pipeline.""" def __init__(self, **kwargs) -> None: super().__init__(**kwargs) self.run_calls: list[int] = [] def _run_validation(self, method, step: int) -> None: # type: ignore[override] self.run_calls.append(step) def _make_recording(**kwargs) -> _NoRunValidation: return _NoRunValidation( pipeline_target=_PIPE_TARGET, dataset_file="x.json", **kwargs, ) class TestOnValidationBegin: def test_skipped_when_every_steps_zero(self) -> None: cb = _make_recording(every_steps=0) cb.on_validation_begin(method=None, iteration=0) cb.on_validation_begin(method=None, iteration=1000) assert cb.run_calls == [] def test_skipped_on_off_iter(self) -> None: cb = _make_recording(every_steps=50) cb.on_validation_begin(method=None, iteration=49) cb.on_validation_begin(method=None, iteration=51) assert cb.run_calls == [] def test_runs_on_match(self) -> None: cb = _make_recording(every_steps=50) cb.on_validation_begin(method=None, iteration=50) cb.on_validation_begin(method=None, iteration=100) assert cb.run_calls == [50, 100] def test_on_validation_begin_runs_step_zero_by_default(self) -> None: """Verify the callback records a pre-training baseline by default.""" cb = _make_recording(every_steps=50) cb.on_validation_begin(method=None, iteration=0) assert cb.run_calls == [0] def test_on_validation_begin_respects_start_gate_and_twenty_step_cadence(self) -> None: """Verify the H3 schedule from baseline through optimizer step 400.""" cb = _make_recording(every_steps=20, run_at_start=True) for iteration in range(401): cb.on_validation_begin(method=None, iteration=iteration) assert cb.run_calls == [0, *range(20, 401, 20)] def test_on_validation_begin_skips_step_zero_when_disabled(self) -> None: """Verify disabling the baseline preserves later cadence events.""" cb = _make_recording(every_steps=20, run_at_start=False) cb.on_validation_begin(method=None, iteration=0) cb.on_validation_begin(method=None, iteration=20) assert cb.run_calls == [20] class TestH3ValidationContract: """Verify synchronized MiniMax H3 validation through Weights & Biases logging.""" def test_prepare_validation_batch_forwards_video_count( self, monkeypatch: pytest.MonkeyPatch, ) -> None: """Verify that each ForwardBatch receives the configured output count.""" cb = _make_callback(num_videos_per_prompt=3) cb.training_config = SimpleNamespace( data=SimpleNamespace( num_height=64, num_width=96, num_latent_t=2, ), pipeline_config=SimpleNamespace(vae_config=SimpleNamespace( arch_config=SimpleNamespace(temporal_compression_ratio=4), ), ), model_path="unused", vsa_sparsity=0.0, ) monkeypatch.setattr( "fastvideo.train.callbacks.validation.make_inference_args", lambda *args, **kwargs: SimpleNamespace(), ) batch = cb._prepare_validation_batch( SamplingParam(), {"prompt": "Generate synchronized media."}, num_inference_steps=5, ) assert batch.num_videos_per_prompt == 3 def test_prepare_validation_batch_ignores_media_for_text_only_generation( self, monkeypatch: pytest.MonkeyPatch, tmp_path, ) -> None: """Verify text-to-video-with-audio uses captions without source media.""" video_path = tmp_path / "reference.mp4" video_path.touch() cb = _make_callback(use_validation_media_conditioning=False) cb.training_config = SimpleNamespace( data=SimpleNamespace( num_height=64, num_width=96, num_latent_t=2, ), pipeline_config=SimpleNamespace(vae_config=SimpleNamespace( arch_config=SimpleNamespace(temporal_compression_ratio=4), ), ), model_path="unused", vsa_sparsity=0.0, ) monkeypatch.setattr( "fastvideo.train.callbacks.validation.make_inference_args", lambda *args, **kwargs: SimpleNamespace(), ) batch = cb._prepare_validation_batch( SamplingParam(), { "prompt": "Generate synchronized media.", "video_path": str(video_path), }, num_inference_steps=5, ) assert batch.image_path is None def test_run_validation_for_steps_aligns_h3_audio_with_captions( self, monkeypatch: pytest.MonkeyPatch, ) -> None: """Verify one waveform and sample rate stay aligned with each H3 video.""" validation_records = [{"caption": f"prompt-{index}"} for index in range(8)] class FakeH3Pipeline: """Return deterministic video and stereo audio for each prompt.""" fastvideo_args = SimpleNamespace(pipeline_config=SimpleNamespace()) def forward(self, batch, inference_args): """Produce media whose values identify the source prompt.""" prompt_index = int(batch.prompt.rsplit("-", 1)[1]) return SimpleNamespace( output=torch.full((1, 3, 2, 2, 2), prompt_index / 8), extra={ "audio": torch.full((32, 2), float(prompt_index)), "audio_sample_rate": 32_000, }, ) cb = _make_callback() cb.training_config = SimpleNamespace( model_path="unused", pipeline_config=SimpleNamespace(), ) cb.rank_in_sp_group = 0 pipeline = FakeH3Pipeline() monkeypatch.setattr( "fastvideo.train.callbacks.validation.ValidationDataset", lambda filename: validation_records, ) monkeypatch.setattr( "fastvideo.train.callbacks.validation.make_inference_args", lambda *args, **kwargs: SimpleNamespace( pipeline_config=SimpleNamespace(dmd_denoising_steps=None), dit_cpu_offload=True, ), ) monkeypatch.setattr(cb, "_get_pipeline", lambda *, transformer: pipeline) monkeypatch.setattr(cb, "_get_sampling_param", SamplingParam) monkeypatch.setattr( cb, "_prepare_validation_batch", lambda sampling_param, validation_batch, num_inference_steps: SimpleNamespace(prompt=validation_batch[ "caption"], ), ) result = cb._run_validation_for_steps(50, transformer=torch.nn.Identity()) assert result.captions == [f"prompt-{index}" for index in range(8)] assert len(result.videos) == 8 assert result.audio_sample_rates == [32_000] * 8 assert len(result.audio_waveforms) == 8 for prompt_index, waveform in enumerate(result.audio_waveforms): assert torch.is_tensor(waveform) torch.testing.assert_close(waveform, torch.full((32, 2), float(prompt_index))) def test_log_validation_video_artifacts_emits_eight_wandb_videos(self) -> None: """Verify a MiniMax H3 event contains eight media objects and scalar evidence.""" class FakeWandbTracker: """Record tracker calls without contacting W&B.""" def __init__(self) -> None: self.video_calls = [] self.artifact_calls = [] def video(self, filename, *, caption, fps): self.video_calls.append((filename, caption, fps)) return filename def log_artifacts(self, artifacts, step): self.artifact_calls.append((artifacts, step)) cb = _make_callback() cb.tracker = FakeWandbTracker() filenames = [f"validation-{index}.mp4" for index in range(8)] captions = [f"caption-{index}" for index in range(8)] scalar_metrics = { "validation/50_steps_video_count": 8.0, "validation/50_steps_duration_sec": 12.5, "validation/50_steps_audio_video_count": 8.0, } cb._log_validation_video_artifacts( filenames, captions, key="validation_videos_50_steps", step=20, fps=24, scalar_metrics=scalar_metrics, ) assert len(cb.tracker.video_calls) == 8 assert len(cb.tracker.artifact_calls) == 1 artifacts, step = cb.tracker.artifact_calls[0] assert step == 20 assert artifacts["validation_videos_50_steps"] == filenames assert {key: artifacts[key] for key in scalar_metrics} == scalar_metrics class TestAttnQatInferValidation: def test_context_swaps_and_restores_qat_attention( self, monkeypatch: pytest.MonkeyPatch, ) -> None: from fastvideo.attention.backends import attn_qat_infer from fastvideo.attention.backends.attn_qat_infer import ( AttnQatInferImpl, ) from fastvideo.attention.backends.attn_qat_train import ( AttnQatTrainImpl, ) from fastvideo.platforms import AttentionBackendEnum monkeypatch.setattr(attn_qat_infer, "is_attn_qat_infer_available", lambda: True) impl = AttnQatTrainImpl( num_heads=1, head_size=128, causal=False, softmax_scale=128**-0.5, ) holder = torch.nn.Module() holder.attn_impl = impl holder.num_heads = 1 holder.head_size = 128 holder.num_kv_heads = 1 holder.softmax_scale = 128**-0.5 holder.backend = AttentionBackendEnum.ATTN_QAT_TRAIN transformer = torch.nn.Module() transformer.add_module("attention", holder) cb = _make_callback() cb.attn_qat_infer = True with pytest.raises(RuntimeError, match="stop"): with cb._attn_qat_infer_context(transformer): assert isinstance(holder.attn_impl, AttnQatInferImpl) assert holder.backend is AttentionBackendEnum.ATTN_QAT_INFER raise RuntimeError("stop") assert holder.attn_impl is impl assert holder.backend is AttentionBackendEnum.ATTN_QAT_TRAIN def test_context_requires_qat_training_attention(self) -> None: cb = _make_callback() cb.attn_qat_infer = True with pytest.raises(RuntimeError, match="no ATTN_QAT_TRAIN layers"): with cb._attn_qat_infer_context(torch.nn.Linear(1, 1)): pass # --------------------------------------------------------------------------- # C. _find_ema_callback # --------------------------------------------------------------------------- class TestFindEmaCallback: def test_returns_none_without_callback_dict(self) -> None: cb = _make_callback() # _callback_dict is not set on bare instances. assert cb._find_ema_callback() is None def test_returns_none_when_no_ema_registered(self) -> None: cb = _make_callback() cb_dict = CallbackDict({}, training_config=object()) cb._callback_dict = cb_dict assert cb._find_ema_callback() is None def test_finds_ema_callback(self) -> None: cb = _make_callback() cb_dict = CallbackDict({}, training_config=object()) ema = EMACallback(decay=0.99) cb_dict._callbacks["ema"] = ema cb_dict._callbacks["validation"] = cb cb._callback_dict = cb_dict found = cb._find_ema_callback() assert found is ema # --------------------------------------------------------------------------- # D. state_dict / load_state_dict (rng round-trip) # --------------------------------------------------------------------------- class TestStateDict: def test_state_dict_empty_without_generator(self) -> None: cb = _make_callback() # validation_random_generator is None until on_train_start. assert cb.state_dict() == {} def test_round_trip_preserves_rng_state(self) -> None: cb = _make_callback() gen = torch.Generator(device="cpu").manual_seed(123) # Advance RNG so a default-init generator on the receiving # side is observably different. for _ in range(5): torch.randn(4, generator=gen) cb.validation_random_generator = gen state = cb.state_dict() assert "validation_rng" in state # Receiver: fresh generator with a different seed. fresh = _make_callback() fresh.validation_random_generator = (torch.Generator(device="cpu").manual_seed(999)) fresh.load_state_dict(state) # After load, both generators draw the same next sample. a = torch.randn(8, generator=cb.validation_random_generator) b = torch.randn(8, generator=fresh.validation_random_generator) assert torch.equal(a, b) def test_load_without_generator_is_noop(self) -> None: cb = _make_callback() # Generator is None: load must not raise even when state has # an rng entry. cb.load_state_dict({"validation_rng": torch.tensor([1, 2, 3], dtype=torch.uint8)}) assert cb.validation_random_generator is None # --------------------------------------------------------------------------- # E. metric result aggregation # --------------------------------------------------------------------------- class TestMetricAggregation: def test_accumulates_scores_and_scalar_details(self) -> None: stats = _ValidationMetricStats() row: dict = {"path": "sample.mp4"} ValidationCallback._accumulate_metric_results( stats, row, { "vbench.aesthetic_quality": SimpleNamespace( name="vbench.aesthetic_quality", score=0.5, details={"ignored": [1, 2, 3]}, ), SYNTHETIC_OPTICAL_FLOW_METRIC: SimpleNamespace( name=SYNTHETIC_OPTICAL_FLOW_METRIC, score=1.5, details={ **{ key: float(i) for i, key in enumerate(SYNTHETIC_OPTICAL_FLOW_LOG_KEYS) }, "pixel_epe_mean_std": 99.0, "pixel_epe_mean_max": 100.0, "pixel_epe_mean_auc": 101.0, }, ), }, ) assert stats.sums["vbench.aesthetic_quality"] == 0.5 assert stats.counts["vbench.aesthetic_quality"] == 1.0 for i, key in enumerate(SYNTHETIC_OPTICAL_FLOW_LOG_KEYS): metric_key = f"{SYNTHETIC_OPTICAL_FLOW_METRIC}.{key}" assert stats.sums[metric_key] == float(i) assert stats.counts[metric_key] == 1.0 assert SYNTHETIC_OPTICAL_FLOW_METRIC not in stats.sums assert f"{SYNTHETIC_OPTICAL_FLOW_METRIC}.pixel_epe_mean_std" not in stats.sums assert f"{SYNTHETIC_OPTICAL_FLOW_METRIC}.pixel_epe_mean_max" not in stats.sums assert f"{SYNTHETIC_OPTICAL_FLOW_METRIC}.pixel_epe_mean_auc" not in stats.sums assert "vbench.aesthetic_quality.ignored" not in stats.sums assert row["vbench.aesthetic_quality"] == 0.5 def test_merges_metric_stats(self) -> None: dst = _ValidationMetricStats( sums={"a": 1.0}, counts={"a": 1.0}, per_video=[{ "path": "a.mp4" }], errors=["first"], ) src = _ValidationMetricStats( sums={ "a": 2.0, "b": 3.0 }, counts={ "a": 2.0, "b": 1.0 }, per_video=[{ "path": "b.mp4" }], errors=["second"], ) ValidationCallback._merge_metric_stats(dst, src) assert dst.sums == {"a": 3.0, "b": 3.0} assert dst.counts == {"a": 3.0, "b": 1.0} assert dst.per_video == [{"path": "a.mp4"}, {"path": "b.mp4"}] assert dst.errors == ["first", "second"] def test_metric_log_name_replaces_dots(self) -> None: cb = ValidationCallback( pipeline_target=_PIPE_TARGET, dataset_file="x.json", metrics={"names": ["vbench.aesthetic_quality"]}, ) assert (cb._metric_log_name("vbench.aesthetic_quality") == "metrics/validation/vbench/aesthetic_quality") def test_metric_device_uses_local_rank( self, monkeypatch: pytest.MonkeyPatch, ) -> None: cb = ValidationCallback( pipeline_target=_PIPE_TARGET, dataset_file="x.json", metrics={"names": ["vbench.aesthetic_quality"]}, ) cb.world_group = SimpleNamespace(local_rank=5) monkeypatch.setattr( torch.cuda, "is_available", lambda: True, ) assert cb._metric_device() == "cuda:5" def test_metric_extras_include_actions_and_calibration(self) -> None: cb = ValidationCallback( pipeline_target=_PIPE_TARGET, dataset_file="x.json", metrics={ "names": ["optical_flow.synthetic_optical_flow"], "calibration_path": "/tmp/calibration.json", "mouse_pitch_sign": -1, }, ) actions = { "keyboard": np.zeros((3, 6)), "mouse": np.zeros((3, 2)), } extras = cb._validation_metric_extras( video_filenames=["a.mp4", "b.mp4"], actions=[actions, None], mouse_pitch_signs=[None, 1], ) assert extras[0]["actions"] is actions assert extras[0]["calibration"] == "/tmp/calibration.json" assert extras[0]["mouse_pitch_sign"] == -1 assert "actions" not in extras[1] assert extras[1]["mouse_pitch_sign"] == 1 def test_validation_actions_prefers_loaded_conditions(self) -> None: keyboard = np.zeros((4, 6)) mouse = np.zeros((4, 2)) actions = ValidationCallback._validation_actions({ "keyboard_cond": keyboard, "mouse_cond": mouse, }) assert actions is not None assert np.array_equal(actions["keyboard"], keyboard) assert np.array_equal(actions["mouse"], mouse) def test_validation_actions_loads_action_path( self, tmp_path, ) -> None: path = tmp_path / "actions.npy" np.save( path, { "keyboard": np.ones((4, 6)), "mouse": np.ones((4, 2)), }, allow_pickle=True, ) actions = ValidationCallback._validation_actions({"action_path": str(path)}) assert actions is not None assert actions["keyboard"].shape == (4, 6) assert actions["mouse"].shape == (4, 2) def test_available_paths_requires_all_paths(self) -> None: assert ValidationCallback._available_paths(["a.mp4", "b.mp4"]) == ["a.mp4", "b.mp4"] assert ValidationCallback._available_paths(["a.mp4", None]) is None # --------------------------------------------------------------------------- # F. action overlay plumbing # --------------------------------------------------------------------------- class TestActionOverlay: def test_save_validation_videos_supports_overlay_suffix( self, tmp_path, monkeypatch: pytest.MonkeyPatch, ) -> None: """Verify overlay artifacts use distinct names and the media writer.""" cb = _make_callback() cb.global_rank = 3 calls: list[tuple[str, int]] = [] def fake_write_validation_mp4( fname: str, video: list[np.ndarray], *, fps: int, audio=None, audio_sample_rate=None, ) -> None: """Record media arguments without invoking a system codec.""" del video, audio, audio_sample_rate calls.append((fname, fps)) monkeypatch.setattr( "fastvideo.train.callbacks.validation.write_validation_mp4", fake_write_validation_mp4, ) saved = cb._save_validation_videos( [[np.zeros((2, 2, 3), dtype=np.uint8)]], output_dir=str(tmp_path), step=7, num_inference_steps=4, fps=25, suffix="_overlay", ) assert saved.filenames == [str(tmp_path / "validation_step_7_inference_steps_4_rank_3_video_0_overlay.mp4")] assert saved.indices == [0] assert calls == [(saved.filenames[0], 25)] def test_save_validation_videos_skips_failed_artifact( self, tmp_path, monkeypatch: pytest.MonkeyPatch, ) -> None: """Verify a failed video-only encode does not hide later artifacts.""" cb = _make_callback() cb.global_rank = 0 calls: list[str] = [] def fake_write_validation_mp4( fname: str, video: list[np.ndarray], *, fps: int, audio=None, audio_sample_rate=None, ) -> None: """Fail the first encode to exercise video-only best-effort logging.""" del video, fps, audio, audio_sample_rate calls.append(fname) if len(calls) == 1: raise OSError("ffmpeg failed") monkeypatch.setattr( "fastvideo.train.callbacks.validation.write_validation_mp4", fake_write_validation_mp4, ) saved = cb._save_validation_videos( [ [np.zeros((2, 2, 3), dtype=np.uint8)], [np.ones((2, 2, 3), dtype=np.uint8)], ], output_dir=str(tmp_path), step=8, num_inference_steps=5, fps=24, ) assert saved.indices == [1] assert saved.filenames == [str(tmp_path / "validation_step_8_inference_steps_5_rank_0_video_1.mp4")] assert len(calls) == 2 def test_save_validation_videos_counts_preserved_audio( self, tmp_path, monkeypatch: pytest.MonkeyPatch, ) -> None: """Verify the callback counts only MP4 files with requested waveforms.""" cb = _make_callback() cb.global_rank = 0 encoded_audio = [] def fake_write_validation_mp4( fname: str, video: list[np.ndarray], *, fps: int, audio=None, audio_sample_rate=None, ) -> None: del fname, video, fps encoded_audio.append((audio, audio_sample_rate)) monkeypatch.setattr( "fastvideo.train.callbacks.validation.write_validation_mp4", fake_write_validation_mp4, ) waveforms = [torch.zeros(32, 2), torch.ones(32, 2)] saved = cb._save_validation_videos( [[np.zeros((2, 2, 3), dtype=np.uint8)]] * 2, output_dir=str(tmp_path), step=0, num_inference_steps=50, fps=24, audio_waveforms=waveforms, audio_sample_rates=[32_000, 32_000], ) assert saved.audio_video_count == 2 assert saved.indices == [0, 1] assert len(encoded_audio) == 2 assert encoded_audio[0][0] is waveforms[0] assert encoded_audio[1][0] is waveforms[1] assert [sample_rate for _, sample_rate in encoded_audio] == [32_000, 32_000] def test_save_validation_videos_skips_audio_encode_failure( self, tmp_path, monkeypatch: pytest.MonkeyPatch, ) -> None: """Verify an audio encoding failure is omitted without stopping later writes.""" cb = _make_callback() cb.global_rank = 0 calls: list[str] = [] def fail_first_audio_encode(fname: str, *args, **kwargs) -> None: """Leave one incomplete file before allowing the second write.""" del args, kwargs calls.append(fname) if len(calls) == 1: with open(fname, "wb") as incomplete_file: incomplete_file.write(b"incomplete") raise RuntimeError("AAC unavailable") monkeypatch.setattr( "fastvideo.train.callbacks.validation.write_validation_mp4", fail_first_audio_encode, ) saved = cb._save_validation_videos( [[np.zeros((2, 2, 3), dtype=np.uint8)]] * 2, output_dir=str(tmp_path), step=0, num_inference_steps=50, fps=24, audio_waveforms=[torch.zeros(32, 2)] * 2, audio_sample_rates=[32_000, 32_000], ) assert saved.indices == [1] assert saved.audio_video_count == 1 assert len(calls) == 2 assert not (tmp_path / "validation_step_0_inference_steps_50_rank_0_video_0.mp4").exists() def test_post_process_validation_frames_uses_student_hook(self) -> None: cb = _make_callback(overlay_actions=True) frames = [np.zeros((2, 2, 3), dtype=np.uint8)] action = { "keyboard": np.ones((1, 4), dtype=np.float32), "mouse": np.zeros((1, 2), dtype=np.float32), } class Student: def post_process_validation_frames( self, received_frames: list[np.ndarray], *, action: dict, ) -> list[np.ndarray]: assert received_frames is frames assert action["keyboard"].shape == (1, 4) return [np.full((2, 2, 3), 255, dtype=np.uint8)] cb.method = SimpleNamespace(student=Student()) overlay = cb._post_process_validation_frames( frames, action=action, ) assert overlay is not None assert int(overlay[0][0, 0, 0]) == 255 def test_post_process_validation_frames_without_hook_is_noop(self) -> None: cb = _make_callback(overlay_actions=True) cb.method = SimpleNamespace(student=object()) overlay = cb._post_process_validation_frames( [np.zeros((2, 2, 3), dtype=np.uint8)], action=None, ) assert overlay is None class TestKeepLoadedEncoderWidths: """``_keep_loaded_encoder_widths`` merges the two config sources. Validation runs the encoder weights the loader brought in, but reaches the stages through configs derived from the training side, which never see the checkpoint's ``config.json``. Each side therefore holds the real value for a different set of fields: the loader knows the widths, training knows the FastVideo-only fields HF configs never carry. """ @staticmethod def _encoder(hidden_size: int, text_len: int = 0) -> SimpleNamespace: return SimpleNamespace(arch_config=SimpleNamespace( hidden_size=hidden_size, text_len=text_len, )) def test_copies_loaded_width(self) -> None: # ByT5 at the generic T5 default vs the checkpoint's real width. validation = SimpleNamespace(text_encoder_configs=(self._encoder(512), )) loaded = SimpleNamespace(text_encoder_configs=(self._encoder(1472), )) ValidationCallback._keep_loaded_encoder_widths(validation, loaded) assert validation.text_encoder_configs[0].arch_config.hidden_size == 1472 def test_preserves_training_owned_fields(self) -> None: # ``text_len`` is a FastVideo-only field absent from HF configs, so # the loader never populates it and the training value must survive. validation = SimpleNamespace(text_encoder_configs=(self._encoder(512, text_len=1000), )) loaded = SimpleNamespace(text_encoder_configs=(self._encoder(1472, text_len=0), )) ValidationCallback._keep_loaded_encoder_widths(validation, loaded) assert validation.text_encoder_configs[0].arch_config.text_len == 1000 def test_keeps_config_objects_unshared(self) -> None: # The second call site writes into ``tc.pipeline_config`` itself, so # the merge must not alias the loaded encoder objects into it. validation = SimpleNamespace(text_encoder_configs=(self._encoder(512), )) original = validation.text_encoder_configs loaded = SimpleNamespace(text_encoder_configs=(self._encoder(1472), )) ValidationCallback._keep_loaded_encoder_widths(validation, loaded) assert validation.text_encoder_configs is original assert validation.text_encoder_configs[0] is not loaded.text_encoder_configs[0] def test_skips_unpopulated_loaded_width(self) -> None: # ``TextEncoderArchConfig.hidden_size`` defaults to 0; a loader that # never filled it must not clobber a real training-side width. validation = SimpleNamespace(text_encoder_configs=(self._encoder(3584), )) loaded = SimpleNamespace(text_encoder_configs=(self._encoder(0), )) ValidationCallback._keep_loaded_encoder_widths(validation, loaded) assert validation.text_encoder_configs[0].arch_config.hidden_size == 3584 def test_multiple_encoders_are_matched_by_index(self) -> None: validation = SimpleNamespace(text_encoder_configs=( self._encoder(512), self._encoder(512), )) loaded = SimpleNamespace(text_encoder_configs=( self._encoder(3584), self._encoder(1472), )) ValidationCallback._keep_loaded_encoder_widths(validation, loaded) widths = [e.arch_config.hidden_size for e in validation.text_encoder_configs] assert widths == [3584, 1472] @pytest.mark.parametrize( "validation, loaded", [ (SimpleNamespace(), SimpleNamespace(text_encoder_configs=())), (SimpleNamespace(text_encoder_configs=()), SimpleNamespace()), (SimpleNamespace(text_encoder_configs=()), SimpleNamespace(text_encoder_configs=())), ], ) def test_missing_or_empty_encoders_is_noop( self, validation: SimpleNamespace, loaded: SimpleNamespace, ) -> None: # Pipelines without text encoders must not raise here. ValidationCallback._keep_loaded_encoder_widths(validation, loaded)