# SPDX-License-Identifier: Apache-2.0 """Validation callback. All configuration is read from the YAML ``callbacks.validation`` section. The pipeline class is resolved from ``pipeline_target``. """ from __future__ import annotations import contextlib import gc import json import os import time from copy import deepcopy from dataclasses import dataclass, field from typing import Any, TYPE_CHECKING import numpy as np import torch import torchvision from einops import rearrange from torch.utils.data import DataLoader from fastvideo.api.sampling_param import SamplingParam from fastvideo.dataset.validation_dataset import ( ValidationDataset, ) from fastvideo.distributed import ( get_sp_group, get_world_group, ) from fastvideo.logger import init_logger from fastvideo.pipelines import ForwardBatch from fastvideo.train.callbacks.callback import Callback from fastvideo.train.utils.instantiate import resolve_target from fastvideo.train.utils.moduleloader import ( make_inference_args, ) from fastvideo.train.utils.validation_media import write_validation_mp4 from fastvideo.training.trackers import DummyTracker from fastvideo.utils import shallow_asdict if TYPE_CHECKING: from fastvideo.train.methods.base import TrainingMethod logger = init_logger(__name__) @dataclass(slots=True) class _ValidationStepResult: """Keep decoded media aligned with the prompt and conditions that produced it.""" videos: list[list[np.ndarray]] captions: list[str] audio_waveforms: list[torch.Tensor | np.ndarray | None] = field(default_factory=list) audio_sample_rates: list[int | None] = field(default_factory=list) overlay_videos: list[list[np.ndarray]] = field(default_factory=list) overlay_captions: list[str] = field(default_factory=list) ref_videos: list[str | None] = field(default_factory=list) actions: list[dict[str, Any] | None] = field(default_factory=list) mouse_pitch_signs: list[int | None] = field(default_factory=list) @dataclass(slots=True) class _ValidationMetricStats: sums: dict[str, float] = field(default_factory=dict) counts: dict[str, float] = field(default_factory=dict) per_video: list[dict[str, Any]] = field(default_factory=list) errors: list[str] = field(default_factory=list) @dataclass(slots=True) class _SavedValidationVideos: """Identify verified MP4 files and their positions in the generated batch.""" filenames: list[str] = field(default_factory=list) indices: list[int] = field(default_factory=list) audio_video_count: int = 0 @dataclass(slots=True) class _ValidationMetricsConfig: enabled: bool = False names: list[str] = field(default_factory=list) device: str = "cuda" calibration_path: str | None = None mouse_pitch_sign: int = 1 skip_missing_deps: bool = True strict: bool = False unload_after_validation: bool = True loader_threads: int = 1 prefetch_factor: int = 2 log_prefix: str = "metrics/validation" DEFAULT_VALIDATION_VBENCH_METRICS = [ "vbench.imaging_quality", "vbench.aesthetic_quality", "vbench.temporal_flickering", "vbench.motion_smoothness", "vbench.subject_consistency", "vbench.background_consistency", "vbench.dynamic_degree", ] SYNTHETIC_OPTICAL_FLOW_METRIC = "optical_flow.synthetic_optical_flow" SYNTHETIC_OPTICAL_FLOW_LOG_KEYS = ( "mf_epe_mean", "mf_angle_err_mean", "mf_cosine_mean", "mf_mag_ratio_mean", "pixel_epe_mean_mean", "px_angle_rmse_mean", "fl_all_mean", "foe_dist_mean", "flow_kl_2d_mean", ) class ValidationCallback(Callback): """Generic validation callback driven entirely by YAML config. Works with any pipeline that follows the ``PipelineCls.from_pretrained(...)`` + ``pipeline.forward()`` contract. """ def __init__( self, *, pipeline_target: str, dataset_file: str, every_steps: int = 100, run_at_start: bool = True, 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, output_dir: str | None = None, sampling_timesteps: list[int] | None = None, overlay_actions: bool = False, keyboard_value_scale: float = 1.0, offload_training_state: bool = False, unload_pipeline_after_validation: bool = False, attn_qat_infer: bool = False, **pipeline_kwargs: Any, ) -> None: """Configure validation cadence, generation parameters, and pipeline loading. ``run_at_start`` controls the pre-training baseline event. ``use_validation_media_conditioning`` lets text-to-video recipes use captions from a dataset that also contains source-media paths. """ self.pipeline_target = str(pipeline_target) self.dataset_file = str(dataset_file) self.every_steps = int(every_steps) self.run_at_start = self._coerce_bool(run_at_start) self.sampling_steps = ([int(s) for s in sampling_steps] if sampling_steps else [40]) self.guidance_scale = (float(guidance_scale) if guidance_scale is not None else None) self.num_frames = (int(num_frames) if num_frames is not None else None) self.num_videos_per_prompt = int(num_videos_per_prompt) if self.num_videos_per_prompt <= 0: raise ValueError("callbacks.validation.num_videos_per_prompt must be positive") self.use_validation_media_conditioning = self._coerce_bool(use_validation_media_conditioning) self.output_dir = (str(output_dir) if output_dir is not None else None) self.sampling_timesteps = ([int(s) for s in sampling_timesteps] if sampling_timesteps is not None else None) self.overlay_actions = self._coerce_bool(overlay_actions) # Validation-only action amplification for world model; training keeps raw action values. self.keyboard_value_scale = float(keyboard_value_scale) metrics_config = pipeline_kwargs.pop("metrics", None) self.metrics_config = self._parse_metrics_config(metrics_config) self.offload_training_state = self._coerce_bool(offload_training_state) self.unload_pipeline_after_validation = self._coerce_bool(unload_pipeline_after_validation) self.attn_qat_infer = self._coerce_bool(attn_qat_infer) self.pipeline_kwargs = dict(pipeline_kwargs) # Set after on_train_start. self._pipeline: Any | None = None self._pipeline_key: tuple[Any, ...] | None = None self._sampling_param: SamplingParam | None = None self._metric_evaluator: Any | None = None self.tracker: Any = DummyTracker() self.validation_random_generator: (torch.Generator | None) = None self.seed: int = 0 @staticmethod def _parse_metrics_config(config: Any) -> _ValidationMetricsConfig: if config is None or config is False: return _ValidationMetricsConfig(enabled=False) if config is True: return _ValidationMetricsConfig( enabled=True, names=list(DEFAULT_VALIDATION_VBENCH_METRICS), ) if isinstance(config, str): return _ValidationMetricsConfig( enabled=True, names=[config], ) if isinstance(config, list | tuple): return _ValidationMetricsConfig( enabled=bool(config), names=[str(name) for name in config], ) if not isinstance(config, dict): raise TypeError("callbacks.validation.metrics must be a bool, string, list, or mapping") enabled = bool(config.get("enabled", True)) raw_names = config.get("names", config.get("metrics", None)) if raw_names is None: names = list(DEFAULT_VALIDATION_VBENCH_METRICS) elif isinstance(raw_names, str): names = [raw_names] else: names = [str(name) for name in raw_names] return _ValidationMetricsConfig( enabled=enabled and bool(names), names=names, device=str(config.get("device", "cuda")), calibration_path=(str(config["calibration_path"]) if config.get("calibration_path") is not None else None), mouse_pitch_sign=int(config.get("mouse_pitch_sign", 1)), skip_missing_deps=bool(config.get("skip_missing_deps", True)), strict=bool(config.get("strict", False)), unload_after_validation=bool(config.get("unload_after_validation", True)), loader_threads=int(config.get("loader_threads", 1)), prefetch_factor=int(config.get("prefetch_factor", 2)), log_prefix=str(config.get("log_prefix", "metrics/validation")), ) @staticmethod def _coerce_bool(value: Any) -> bool: if isinstance(value, str): return value.strip().lower() in {"1", "true", "yes", "on"} return bool(value) # ---------------------------------------------------------- # Callback hooks # ---------------------------------------------------------- def on_train_start( self, method: TrainingMethod, iteration: int = 0, ) -> None: self.method = method tc = self.training_config self.world_group = get_world_group() self.sp_group = get_sp_group() self.global_rank = self.world_group.rank self.rank_in_sp_group = (self.sp_group.rank_in_group) self.sp_world_size = self.sp_group.world_size seed = tc.data.seed if seed is None: raise ValueError("training.data.seed must be set " "for validation") self.seed = int(seed) self.validation_random_generator = (torch.Generator(device="cpu").manual_seed(self.seed)) tracker = getattr(method, "tracker", None) if tracker is not None: self.tracker = tracker def on_validation_begin( self, method: TrainingMethod, iteration: int = 0, ) -> None: """Run the optional step-zero baseline and each scheduled validation event.""" if self.every_steps <= 0: return # Step zero measures the checkpoint before the first optimizer update. if iteration == 0 and not self.run_at_start: return if iteration % self.every_steps != 0: return self._run_validation(method, iteration) # ---------------------------------------------------------- # Core validation logic # ---------------------------------------------------------- def _run_validation( self, method: TrainingMethod, step: int, ) -> None: transformer = method.student.transformer try: with self._validation_memory_context( method, validation_transformer=transformer, ): # Look for an EMA callback to temporarily swap # EMA weights during validation. ema_cb = self._find_ema_callback() ctx = ema_cb.ema_context(transformer) if ema_cb is not None else contextlib.nullcontext(transformer) with ctx as t, self._attn_qat_infer_context(t): self._run_validation_inner( method, step, t, ) finally: if self.unload_pipeline_after_validation: self._clear_pipeline_cache() @contextlib.contextmanager def _attn_qat_infer_context(self, transformer: torch.nn.Module): if not self.attn_qat_infer: yield return from fastvideo.attention.backends.attn_qat_infer import (AttnQatInferImpl, is_attn_qat_infer_available) from fastvideo.attention.backends.attn_qat_train import ( AttnQatTrainImpl, ) from fastvideo.platforms import AttentionBackendEnum layers = [ module for module in transformer.modules() if isinstance(getattr(module, "attn_impl", None), AttnQatTrainImpl) ] if not layers: raise RuntimeError("attn_qat_infer validation requested, but the transformer has no ATTN_QAT_TRAIN layers") if not is_attn_qat_infer_available(): from fastvideo.attention.backends.attn_qat_infer import ( attn_qat_infer_receipt, ) raise RuntimeError("attn_qat_infer validation requested but no ATTN_QAT_INFER kernel serves " f"this device ({attn_qat_infer_receipt()}). Set " "callbacks.validation.attn_qat_infer=false to validate with ATTN_QAT_TRAIN.") previous = [(layer, layer.attn_impl, layer.backend) for layer in layers] try: for layer, impl, _ in previous: layer.attn_impl = AttnQatInferImpl( num_heads=layer.num_heads, head_size=layer.head_size, num_kv_heads=layer.num_kv_heads, causal=impl.causal, softmax_scale=layer.softmax_scale, ) layer.backend = AttentionBackendEnum.ATTN_QAT_INFER logger.info("Enabled ATTN_QAT_INFER for %d validation attention layers.", len(layers)) yield finally: for layer, impl, backend in previous: layer.attn_impl = impl layer.backend = backend @contextlib.contextmanager def _validation_memory_context( self, method: TrainingMethod, *, validation_transformer: torch.nn.Module, ): if not self.offload_training_state: yield return optimizer_tensor_records: list[tuple[Any, Any, torch.device]] = [] module_records: list[tuple[str, torch.nn.Module, torch.device]] = [] try: self._offload_optimizer_states_to_cpu( method, optimizer_tensor_records, ) self._offload_inactive_role_modules_to_cpu( method, validation_transformer=validation_transformer, module_records=module_records, ) self._empty_cuda_cache() yield finally: self._restore_inactive_role_modules(module_records) self._restore_optimizer_states(optimizer_tensor_records) self._empty_cuda_cache() def _offload_optimizer_states_to_cpu( self, method: TrainingMethod, records: list[tuple[Any, Any, torch.device]], ) -> None: optimizers = getattr(method, "_optimizer_dict", {}) if not optimizers: return moved = 0 for optimizer in optimizers.values(): state = getattr(optimizer, "state", None) if not isinstance(state, dict): continue for param_state in state.values(): moved += self._offload_tensor_container_to_cpu( param_state, records, ) if moved: logger.info( "Offloaded %d optimizer state tensors to CPU for validation.", moved, ) def _offload_tensor_container_to_cpu( self, obj: Any, records: list[tuple[Any, Any, torch.device]], ) -> int: moved = 0 if isinstance(obj, dict): for key, value in list(obj.items()): if torch.is_tensor(value) and value.device.type == "cuda": records.append((obj, key, value.device)) obj[key] = value.detach().cpu() moved += 1 else: moved += self._offload_tensor_container_to_cpu(value, records) return moved if isinstance(obj, list): for idx, value in enumerate(list(obj)): if torch.is_tensor(value) and value.device.type == "cuda": records.append((obj, idx, value.device)) obj[idx] = value.detach().cpu() moved += 1 else: moved += self._offload_tensor_container_to_cpu(value, records) return moved def _restore_optimizer_states( self, records: list[tuple[Any, Any, torch.device]], ) -> None: for container, key, device in reversed(records): value = container[key] if torch.is_tensor(value): container[key] = value.to(device=device) if records: logger.info( "Restored %d optimizer state tensors after validation.", len(records), ) def _offload_inactive_role_modules_to_cpu( self, method: TrainingMethod, *, validation_transformer: torch.nn.Module, module_records: list[tuple[str, torch.nn.Module, torch.device]], ) -> None: role_models = getattr(method, "_role_models", {}) if not isinstance(role_models, dict): return for role, model in role_models.items(): module = getattr(model, "transformer", None) if not isinstance(module, torch.nn.Module): continue if module is validation_transformer: continue device = self._first_cuda_tensor_device(module) if device is None: continue try: module.to("cpu") except Exception as exc: logger.warning( "Could not offload role %r transformer to CPU before validation: %s", role, exc, ) continue module_records.append((str(role), module, device)) logger.info( "Offloaded role %r transformer from %s to CPU for validation.", role, device, ) def _restore_inactive_role_modules( self, module_records: list[tuple[str, torch.nn.Module, torch.device]], ) -> None: for role, module, device in reversed(module_records): module.to(device) logger.info( "Restored role %r transformer to %s after validation.", role, device, ) @staticmethod def _first_cuda_tensor_device(module: torch.nn.Module) -> torch.device | None: for tensor in list(module.parameters(recurse=True)) + list(module.buffers(recurse=True)): device = getattr(tensor, "device", None) if isinstance(device, torch.device) and device.type == "cuda": return device return None def _clear_pipeline_cache(self) -> None: self._pipeline = None self._pipeline_key = None self._empty_cuda_cache() @staticmethod def _empty_cuda_cache() -> None: if torch.cuda.is_available(): torch.cuda.empty_cache() def _find_ema_callback(self) -> Any | None: """Find the EMA callback in the callback dict.""" from fastvideo.train.callbacks.ema import ( EMACallback, ) cb_dict = getattr(self, "_callback_dict", None) if cb_dict is not None: for cb in cb_dict._callbacks.values(): if isinstance(cb, EMACallback): return cb return None def _run_validation_inner( self, method: TrainingMethod, step: int, transformer: torch.nn.Module, ) -> None: """Generate one event and gather sequence-parallel group outputs for logging. Each sequence-parallel group leader saves its local media. Global rank zero gathers the leaders' paths and statistics into one tracker event. """ tc = self.training_config was_training = bool(getattr(transformer, "training", False)) output_dir = (self.output_dir or tc.checkpoint.output_dir) try: transformer.eval() num_sp_groups = (self.world_group.world_size // self.sp_group.world_size) sp = self._get_sampling_param() for num_inference_steps in self.sampling_steps: validation_started_at = time.perf_counter() result = self._run_validation_for_steps( num_inference_steps, transformer=transformer, ) # Every rank participates in sequence-parallel inference, but # only the group leader retains decoded media for saving. if self.rank_in_sp_group != 0: continue os.makedirs( output_dir, exist_ok=True, ) local_videos = self._save_validation_videos( result.videos, output_dir=output_dir, step=step, num_inference_steps=num_inference_steps, fps=sp.fps, audio_waveforms=result.audio_waveforms, audio_sample_rates=result.audio_sample_rates, ) local_overlay_videos = self._save_validation_videos( result.overlay_videos, output_dir=output_dir, step=step, num_inference_steps=num_inference_steps, fps=sp.fps, suffix="_overlay", ) local_video_filenames = local_videos.filenames local_overlay_video_filenames = local_overlay_videos.filenames local_captions = self._select_by_indices( result.captions, local_videos.indices, ) local_overlay_captions = self._select_by_indices( result.overlay_captions, local_overlay_videos.indices, ) local_ref_videos = self._select_by_indices( result.ref_videos, local_videos.indices, ) local_actions = self._select_by_indices( result.actions, local_videos.indices, ) local_mouse_pitch_signs = self._select_by_indices( result.mouse_pitch_signs, local_videos.indices, ) local_metric_stats = self._evaluate_validation_metrics( video_filenames=local_video_filenames, captions=local_captions, ref_videos=local_ref_videos, actions=local_actions, mouse_pitch_signs=local_mouse_pitch_signs, fps=sp.fps, output_dir=output_dir, step=step, num_inference_steps=num_inference_steps, ) if self.global_rank == 0: all_video_filenames = list(local_video_filenames) all_overlay_video_filenames = list(local_overlay_video_filenames) all_captions = list(local_captions) all_overlay_captions = list(local_overlay_captions) all_audio_video_count = local_videos.audio_video_count all_metric_stats = local_metric_stats for sp_idx in range(1, num_sp_groups): # Sequence-parallel group leaders occupy the first # global rank in each contiguous group. src = (sp_idx * self.sp_world_size) recv_v = (self.world_group.recv_object(src=src)) recv_c = (self.world_group.recv_object(src=src)) recv_ov = (self.world_group.recv_object(src=src)) recv_oc = (self.world_group.recv_object(src=src)) recv_m = (self.world_group.recv_object(src=src)) recv_audio_video_count = (self.world_group.recv_object(src=src)) all_video_filenames.extend(recv_v) all_overlay_video_filenames.extend(recv_ov) all_captions.extend(recv_c) all_overlay_captions.extend(recv_oc) all_audio_video_count += int(recv_audio_video_count) self._merge_metric_stats( all_metric_stats, recv_m, ) self._log_validation_metrics( all_metric_stats, step=step, ) # Media and completion counts share one tracker event so # artifacts and verification data remain aligned. self._log_validation_video_artifacts( all_video_filenames, all_captions, key=f"validation_videos_{num_inference_steps}_steps", step=step, fps=sp.fps, scalar_metrics={ f"validation/{num_inference_steps}_steps_video_count": float(len(all_video_filenames)), f"validation/{num_inference_steps}_steps_duration_sec": (time.perf_counter() - validation_started_at), f"validation/{num_inference_steps}_steps_audio_video_count": float(all_audio_video_count), }, ) if all_overlay_video_filenames: self._log_validation_video_artifacts( all_overlay_video_filenames, all_overlay_captions, key=(f"validation_videos_{num_inference_steps}" f"_steps_overlay"), step=step, fps=sp.fps, ) else: self.world_group.send_object( local_video_filenames, dst=0, ) self.world_group.send_object( local_captions, dst=0, ) self.world_group.send_object( local_overlay_video_filenames, dst=0, ) self.world_group.send_object( local_overlay_captions, dst=0, ) self.world_group.send_object( local_metric_stats, dst=0, ) self.world_group.send_object( local_videos.audio_video_count, dst=0, ) finally: if was_training: transformer.train() self._release_metric_evaluator() def _save_validation_videos( self, videos: list[list[np.ndarray]], *, output_dir: str, step: int, num_inference_steps: int, fps: int, suffix: str = "", audio_waveforms: list[torch.Tensor | np.ndarray | None] | None = None, audio_sample_rates: list[int | None] | None = None, ) -> _SavedValidationVideos: """Save aligned validation media and skip artifacts that fail encoding. Audio and sample-rate metadata must pass the alignment checks before encoding. Media writer failures are logged and omitted so validation artifact creation does not interrupt the training loop. """ if (audio_waveforms is None) != (audio_sample_rates is None): raise ValueError("Validation audio waveforms and sample rates must be provided together.") if audio_waveforms is not None and len(audio_waveforms) != len(videos): raise ValueError("Validation audio waveform count must match the video count.") if audio_sample_rates is not None and len(audio_sample_rates) != len(videos): raise ValueError("Validation audio sample-rate count must match the video count.") saved = _SavedValidationVideos() for i, video in enumerate(videos): audio = audio_waveforms[i] if audio_waveforms is not None else None audio_sample_rate = audio_sample_rates[i] if audio_sample_rates is not None else None if (audio is None) != (audio_sample_rate is None): raise ValueError("Each validation waveform must have one aligned sample rate.") fname = os.path.join( output_dir, f"validation_step_{step}" f"_inference_steps_{num_inference_steps}" f"_rank_{self.global_rank}" f"_video_{i}{suffix}.mp4", ) try: write_validation_mp4( fname, video, fps=fps, audio=audio, audio_sample_rate=audio_sample_rate, ) except Exception as exc: # Validation media is diagnostic output, so one failed write # must not terminate training or prevent later artifact writes. logger.exception( "Failed to save validation media %s on rank %s; skipping artifact: %s", fname, self.global_rank, exc, ) with contextlib.suppress(OSError): os.remove(fname) continue saved.filenames.append(fname) saved.indices.append(i) if audio is not None: # The media writer verifies requested streams before returning. saved.audio_video_count += 1 return saved @staticmethod def _select_by_indices( values: list[Any], indices: list[int], ) -> list[Any]: return [values[i] for i in indices if i < len(values)] def _log_validation_video_artifacts( self, video_filenames: list[str], captions: list[str], *, key: str, step: int, fps: int, scalar_metrics: dict[str, float] | None = None, ) -> None: """Log validation media and its scalar verification data at one step.""" video_logs = [] for fname, cap in zip( video_filenames, captions, strict=True, ): art = self.tracker.video( fname, caption=cap, fps=fps, ) if art is not None: video_logs.append(art) if video_logs: artifacts: dict[str, Any] = {key: video_logs} if scalar_metrics: artifacts.update(scalar_metrics) self.tracker.log_artifacts( artifacts, step, ) # ---------------------------------------------------------- # Metric evaluation # ---------------------------------------------------------- def _metric_device(self) -> str: device = self.metrics_config.device if device == "cuda" and torch.cuda.is_available(): local_rank = int(getattr(self.world_group, "local_rank", 0)) return f"cuda:{local_rank}" return device def _get_metric_evaluator(self) -> Any: if self._metric_evaluator is not None: return self._metric_evaluator from fastvideo.eval import Evaluator cfg = self.metrics_config self._metric_evaluator = Evaluator( metrics=cfg.names, device=self._metric_device(), num_gpus=1, loader_threads=cfg.loader_threads, prefetch_factor=cfg.prefetch_factor, pre_upload=False, skip_missing_deps=cfg.skip_missing_deps, ) logger.info( "Initialized validation metrics: %s", ", ".join(self._metric_evaluator.metric_names), ) return self._metric_evaluator def _release_metric_evaluator(self) -> None: if self._metric_evaluator is None: return if self.metrics_config.unload_after_validation: self._metric_evaluator.unload() self._metric_evaluator = None gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() def _evaluate_validation_metrics( self, *, video_filenames: list[str], captions: list[str], ref_videos: list[str | None], actions: list[dict[str, Any] | None], mouse_pitch_signs: list[int | None], fps: int, output_dir: str, step: int, num_inference_steps: int, ) -> _ValidationMetricStats: stats = _ValidationMetricStats() if not self.metrics_config.enabled or not video_filenames: return stats try: evaluator = self._get_metric_evaluator() from fastvideo.eval import samples_from samples = samples_from( video=video_filenames, reference=self._available_paths(ref_videos), text_prompts=captions, fps=float(fps), extras=self._validation_metric_extras( video_filenames=video_filenames, actions=actions, mouse_pitch_signs=mouse_pitch_signs, ), ) results = evaluator.evaluate(samples=samples) for filename, metric_results in zip( video_filenames, results, strict=True, ): row: dict[str, Any] = {"path": filename} self._accumulate_metric_results( stats, row, metric_results, ) stats.per_video.append(row) for metric_name, metric_result in results.corpus.items(): row = {"path": ""} self._accumulate_metric_results( stats, row, {metric_name: metric_result}, ) stats.per_video.append(row) except Exception as exc: message = ("Validation metric evaluation failed on rank " f"{self.global_rank}: {exc}") logger.exception(message) stats.errors.append(message) if self.metrics_config.strict: raise finally: if self._metric_evaluator is not None: self._metric_evaluator.release_cuda_memory() self._write_metric_summary( stats, output_dir=output_dir, step=step, num_inference_steps=num_inference_steps, ) return stats @staticmethod def _accumulate_metric_results( stats: _ValidationMetricStats, row: dict[str, Any], metric_results: dict[str, Any], ) -> None: for metric_name, metric_result in metric_results.items(): details = getattr(metric_result, "details", {}) or {} if metric_name == SYNTHETIC_OPTICAL_FLOW_METRIC: if not isinstance(details, dict): continue for detail_name in SYNTHETIC_OPTICAL_FLOW_LOG_KEYS: key = f"{metric_name}.{detail_name}" ValidationCallback._accumulate_scalar( stats, row, key, details.get(detail_name), ) continue score = getattr(metric_result, "score", None) ValidationCallback._accumulate_scalar( stats, row, metric_name, score, ) if not isinstance(details, dict): continue for detail_name, value in details.items(): key = f"{metric_name}.{detail_name}" ValidationCallback._accumulate_scalar( stats, row, key, value, ) @staticmethod def _accumulate_scalar( stats: _ValidationMetricStats, row: dict[str, Any], key: str, value: Any, ) -> None: if not isinstance(value, float | int | np.floating | np.integer): return value_float = float(value) if not np.isfinite(value_float): return row[key] = value_float stats.sums[key] = stats.sums.get(key, 0.0) + value_float stats.counts[key] = stats.counts.get(key, 0.0) + 1.0 @staticmethod def _merge_metric_stats( dst: _ValidationMetricStats, src: _ValidationMetricStats, ) -> None: for key, value in src.sums.items(): dst.sums[key] = dst.sums.get(key, 0.0) + float(value) for key, value in src.counts.items(): dst.counts[key] = dst.counts.get(key, 0.0) + float(value) dst.per_video.extend(src.per_video) dst.errors.extend(src.errors) def _log_validation_metrics( self, stats: _ValidationMetricStats, *, step: int, ) -> None: if stats.errors: message = "Validation metric evaluation failed:\n" + "\n".join(stats.errors) if self.metrics_config.strict: raise RuntimeError(message) logger.warning(message) logs: dict[str, float] = {} for key, metric_sum in stats.sums.items(): metric_count = stats.counts.get(key, 0.0) if metric_count <= 0: continue value = float(metric_sum / metric_count) if not np.isfinite(value): continue logs[self._metric_log_name(key)] = value if logs: self.tracker.log( logs, step, ) def _metric_log_name( self, key: str, ) -> str: return f"{self.metrics_config.log_prefix}/{key.replace('.', '/')}" def _write_metric_summary( self, stats: _ValidationMetricStats, *, output_dir: str, step: int, num_inference_steps: int, ) -> None: if not self.metrics_config.enabled: return summary_dir = os.path.join( output_dir, "eval", f"step_{step}", ) os.makedirs( summary_dir, exist_ok=True, ) metrics = { key: float(stats.sums[key] / stats.counts[key]) for key in sorted(stats.sums) if stats.counts.get(key, 0.0) > 0 } payload = { "step": int(step), "num_inference_steps": int(num_inference_steps), "rank": int(self.global_rank), "metrics": metrics, "per_video": stats.per_video, "errors": stats.errors, } path = os.path.join( summary_dir, f"inference_steps_{num_inference_steps}_rank_{self.global_rank}.json", ) with open( path, "w", encoding="utf-8", ) as f: json.dump( payload, f, indent=2, ) def _validation_metric_extras( self, *, video_filenames: list[str], actions: list[dict[str, Any] | None], mouse_pitch_signs: list[int | None], ) -> list[dict[str, Any]]: extras: list[dict[str, Any]] = [] for i, _filename in enumerate(video_filenames): extra: dict[str, Any] = {} action = actions[i] if i < len(actions) else None if action is not None: extra["actions"] = action if self.metrics_config.calibration_path is not None: extra["calibration"] = self.metrics_config.calibration_path mouse_pitch_sign = mouse_pitch_signs[i] if i < len(mouse_pitch_signs) else None extra["mouse_pitch_sign"] = int(mouse_pitch_sign or self.metrics_config.mouse_pitch_sign) extras.append(extra) return extras @staticmethod def _available_paths(paths: list[str | None]) -> list[str] | None: if not paths or any(path is None for path in paths): return None return [str(path) for path in paths] @staticmethod def _validation_actions(validation_batch: dict[str, Any]) -> dict[str, Any] | None: keyboard = validation_batch.get("keyboard_cond") mouse = validation_batch.get("mouse_cond") if keyboard is not None and mouse is not None: return { "keyboard": np.asarray(keyboard), "mouse": np.asarray(mouse), } action_path = validation_batch.get("action_path") if not isinstance(action_path, str) or not os.path.isfile(action_path): return None try: actions_obj = np.load(action_path, allow_pickle=True) if isinstance(actions_obj, np.ndarray) and actions_obj.dtype == object: actions_obj = actions_obj.item() except Exception as exc: logger.warning( "Failed to load validation action file %s for metrics: %s", action_path, exc, ) return None if not isinstance(actions_obj, dict): return None keyboard = actions_obj.get("keyboard") mouse = actions_obj.get("mouse") if keyboard is None or mouse is None: return None return { "keyboard": np.asarray(keyboard), "mouse": np.asarray(mouse), } def _validation_mouse_pitch_sign( self, validation_batch: dict[str, Any], ) -> int: if validation_batch.get("mouse_pitch_sign") is not None: return int(validation_batch["mouse_pitch_sign"]) flipped = validation_batch.get("mouse_pitch_flipped") if isinstance(flipped, str): flipped = flipped.lower() in {"1", "true", "yes", "y"} if bool(flipped): return -1 return int(self.metrics_config.mouse_pitch_sign) # ---------------------------------------------------------- # Pipeline management # ---------------------------------------------------------- def _get_sampling_param(self) -> SamplingParam: if self._sampling_param is None: self._sampling_param = (SamplingParam.from_pretrained(self._pipeline_model_path())) return self._sampling_param def _pipeline_model_path(self) -> str: tc = self.training_config pipeline_config = getattr(tc, "pipeline_config", None) model_path = getattr(pipeline_config, "_fastvideo_train_model_path", None) return str(model_path or tc.model_path) def _validation_pipeline_config(self, transformer: torch.nn.Module) -> Any: tc = self.training_config pipeline_config = deepcopy(tc.pipeline_config) self._sync_runtime_dit_arch_config( pipeline_config, transformer, ) return pipeline_config @staticmethod def _keep_loaded_encoder_widths( validation_config: Any, loaded_config: Any, ) -> None: """Carry the loader-populated encoder widths onto ``validation_config``. Validation reaches the stages through two pipeline configs that never pass through ``ModelConfig.update_model_arch``: the deep copy ``_validation_pipeline_config`` makes of the training-side config, and ``tc.pipeline_config`` itself, which ``make_inference_args`` hands to ``pipeline.forward`` by reference. Both still hold the encoder dataclass defaults for whatever the checkpoint would have supplied. Stages read ``hidden_size`` only when they have to synthesise an embedding instead of measuring one: HunyuanVideo 1.5 sizes its zero-length ByT5 placeholder from it, so the generic ``T5ArchConfig`` default of 512 collides with the checkpoint's real width of 1472. Only ``hidden_size`` is copied. Training owns the rest: model plugins set ``text_len`` from ``text_encoder_max_lengths`` to size the parquet text padding, and ``tokenizer_kwargs`` carry run-specific settings, so both keep their training values. A falsy width means the loader never populated that encoder, so there is nothing to carry over. """ loaded_encoders = getattr(loaded_config, "text_encoder_configs", None) validation_encoders = getattr( validation_config, "text_encoder_configs", None, ) if not loaded_encoders or not validation_encoders: return for validation_encoder, loaded_encoder in zip( validation_encoders, loaded_encoders, strict=False, ): hidden_size = getattr( getattr(loaded_encoder, "arch_config", None), "hidden_size", None, ) arch_config = getattr(validation_encoder, "arch_config", None) if hidden_size and arch_config is not None: arch_config.hidden_size = hidden_size @staticmethod def _sync_runtime_dit_arch_config( pipeline_config: Any, transformer: torch.nn.Module, ) -> None: dit_config = getattr( pipeline_config, "dit_config", None, ) arch_config = getattr( dit_config, "arch_config", None, ) if arch_config is None: return for name in ("local_attn_size", "sink_size"): if not hasattr(arch_config, name) or not hasattr(transformer, name): continue setattr( arch_config, name, getattr(transformer, name), ) def _get_pipeline( self, *, transformer: torch.nn.Module, ) -> Any: key = (id(transformer), ) if (self._pipeline is not None and self._pipeline_key == key): return self._pipeline tc = self.training_config PipelineCls = resolve_target(self.pipeline_target) flow_shift = getattr( tc.pipeline_config, "flow_shift", None, ) loaded_modules: dict[str, Any] = {"transformer": transformer} # Distillation methods build the flow-match scheduler their few-step DMD # sampler needs; inject it so the pipeline doesn't fall back to UniPC. method_scheduler = getattr(self.method, "_sf_scheduler", None) if method_scheduler is not None: loaded_modules["scheduler"] = method_scheduler kwargs: dict[str, Any] = { "inference_mode": True, "loaded_modules": loaded_modules, "tp_size": tc.distributed.tp_size, "sp_size": tc.distributed.sp_size, "num_gpus": tc.distributed.num_gpus, "pin_cpu_memory": (tc.distributed.pin_cpu_memory), "dit_cpu_offload": False, "dit_layerwise_offload": False, } if flow_shift is not None: kwargs["flow_shift"] = float(flow_shift) kwargs.update(self.pipeline_kwargs) # The pipeline class comes from a YAML target, so static analysis cannot # infer the dynamically resolved ``from_pretrained`` class method. self._pipeline = PipelineCls.from_pretrained( # type: ignore[attr-defined] self._pipeline_model_path(), **kwargs, ) if tc.pipeline_config is not None: loaded_config = self._pipeline.fastvideo_args.pipeline_config validation_config = self._validation_pipeline_config(transformer) self._keep_loaded_encoder_widths( validation_config, loaded_config, ) self._pipeline.fastvideo_args.pipeline_config = validation_config arch_config = self._pipeline.fastvideo_args.pipeline_config.dit_config.arch_config logger.info( "Validation pipeline runtime config: local_attn_size=%s sink_size=%s boundary_ratio=%s", getattr(arch_config, "local_attn_size", None), getattr(arch_config, "sink_size", None), getattr(self._pipeline.fastvideo_args.pipeline_config.dit_config, "boundary_ratio", None), ) self._pipeline_key = key return self._pipeline # ---------------------------------------------------------- # Batch preparation # ---------------------------------------------------------- def _prepare_validation_batch( self, sampling_param: SamplingParam, validation_batch: dict[str, Any], num_inference_steps: int, ) -> ForwardBatch: """Build a pipeline batch with the recipe's output and conditioning policy.""" tc = self.training_config sampling_param.prompt = validation_batch["prompt"] sampling_param.height = tc.data.num_height sampling_param.width = tc.data.num_width sampling_param.num_inference_steps = int(num_inference_steps) sampling_param.data_type = "video" if self.guidance_scale is not None: sampling_param.guidance_scale = float(self.guidance_scale) sampling_param.seed = self.seed # Output multiplicity belongs in SamplingParam so pipeline stages # allocate the same batch dimension that validation expects to log. sampling_param.num_videos_per_prompt = self.num_videos_per_prompt # SamplingParam is cached across records; clearing the path prevents a # prior record's image or video from conditioning a later prompt. sampling_param.image_path = None if self.use_validation_media_conditioning: # Image-to-video pipelines use an image or the first frame of a validation video. img_path = (validation_batch.get("image_path") or validation_batch.get("video_path")) if img_path is not None and (img_path.startswith("http") or os.path.isfile(img_path)): sampling_param.image_path = img_path temporal_compression_factor = int( tc.pipeline_config.vae_config.arch_config.temporal_compression_ratio # type: ignore[union-attr] ) default_num_frames = ((tc.data.num_latent_t - 1) * temporal_compression_factor + 1) if self.num_frames is not None: sampling_param.num_frames = int(self.num_frames) else: sampling_param.num_frames = int(default_num_frames) latents_size = [ (sampling_param.num_frames - 1) // 4 + 1, sampling_param.height // 8, sampling_param.width // 8, ] n_tokens = (latents_size[0] * latents_size[1] * latents_size[2]) sampling_timesteps_tensor = (torch.tensor( [int(s) for s in self.sampling_timesteps], dtype=torch.long, ) if self.sampling_timesteps is not None else None) inference_args = make_inference_args( tc, model_path=tc.model_path, ) batch = ForwardBatch( **shallow_asdict(sampling_param), generator=self.validation_random_generator, n_tokens=n_tokens, eta=0.0, VSA_sparsity=tc.vsa_sparsity, timesteps=sampling_timesteps_tensor, ) # shallow_asdict(sampling_param) copies list-typed fields by # reference. sampling_param is cached and reused across every # validation sample/step, so without this reset every ForwardBatch # would share (and keep appending to) the same # prompt_attention_mask/negative_attention_mask list forever -- # index [0] would then hold the *first-ever* validation sample's # mask instead of the current one, mismatching prompt_embeds. batch.prompt_attention_mask = [] batch.negative_attention_mask = [] batch._inference_args = inference_args # type: ignore[attr-defined] # Conditionally set I2V fields. if ("image" in validation_batch and validation_batch["image"] is not None): batch.pil_image = validation_batch["image"] self._attach_action_conditions( batch, validation_batch, sampling_param.num_frames, ) return batch def _attach_action_conditions( self, batch: ForwardBatch, validation_batch: dict[str, Any], num_frames: int, ) -> None: for name in ("keyboard_cond", "mouse_cond"): value = validation_batch.get(name) if value is None: continue array = np.asarray(value) if array.size == 0: continue array = array[:num_frames] tensor = torch.as_tensor( array, dtype=torch.bfloat16, ) student = getattr(self.method, "student", None) prepare_action = getattr(student, "prepare_validation_action_condition", None) if prepare_action is not None: tensor = prepare_action( tensor, name=name, keyboard_value_scale=self.keyboard_value_scale, ) tensor = tensor.unsqueeze(0) setattr( batch, name, tensor, ) # ---------------------------------------------------------- # Validation loop # ---------------------------------------------------------- def _run_validation_for_steps( self, num_inference_steps: int, *, transformer: torch.nn.Module, ) -> _ValidationStepResult: """Generate the prompts assigned to one sequence-parallel group. ``ValidationDataset`` pads the prompt count to the data-parallel degree and assigns an equal prompt shard to each sequence-parallel group. All ranks in a group execute each forward pass; the group leader retains decoded media. """ tc = self.training_config pipeline = self._get_pipeline(transformer=transformer, ) sampling_param = self._get_sampling_param() dataset = ValidationDataset(self.dataset_file) dataloader = DataLoader( dataset, batch_size=None, num_workers=0, ) inference_args = make_inference_args( tc, model_path=tc.model_path, ) self._sync_runtime_dit_arch_config( inference_args.pipeline_config, transformer, ) self._keep_loaded_encoder_widths( inference_args.pipeline_config, pipeline.fastvideo_args.pipeline_config, ) # Propagate sampling_timesteps to pipeline_config so # causal/DMD denoising stages can read them. if (self.sampling_timesteps is not None and inference_args.pipeline_config.dmd_denoising_steps is None): inference_args.pipeline_config.dmd_denoising_steps = ([int(s) for s in self.sampling_timesteps]) videos: list[list[np.ndarray]] = [] audio_waveforms: list[torch.Tensor | np.ndarray | None] = [] audio_sample_rates: list[int | None] = [] overlay_videos: list[list[np.ndarray]] = [] captions: list[str] = [] overlay_captions: list[str] = [] ref_videos: list[str | None] = [] actions: list[dict[str, Any] | None] = [] mouse_pitch_signs: list[int | None] = [] for validation_batch in dataloader: batch = self._prepare_validation_batch( sampling_param, validation_batch, num_inference_steps, ) assert (batch.prompt is not None and isinstance(batch.prompt, str)) ref_video = validation_batch.get("ref_video") action = self._validation_actions(validation_batch) with torch.no_grad(): output_batch = pipeline.forward( batch, inference_args, ) samples = output_batch.output.cpu() if self.rank_in_sp_group != 0: continue # Append metadata only on the group leader so every list position # describes the same decoded video throughout save and logging. output_audio = output_batch.extra.get("audio") output_audio_sample_rate = output_batch.extra.get("audio_sample_rate") if (output_audio is None) != (output_audio_sample_rate is None): raise ValueError("Validation pipeline outputs must provide audio and its sample rate together.") if output_audio is not None and not isinstance(output_audio, np.ndarray) and not torch.is_tensor(output_audio): raise TypeError("Validation pipeline audio must be a torch.Tensor or numpy.ndarray; " f"got {type(output_audio).__name__}.") if torch.is_tensor(output_audio): # The returned validation result stays on CPU so MP4 encoding # does not retain a generation tensor on the GPU. output_audio = output_audio.detach().cpu() video = rearrange( samples, "b c t h w -> t b c h w", ) frames: list[np.ndarray] = [] for x in video: x = torchvision.utils.make_grid( x, nrow=6, ) x = (x.transpose(0, 1).transpose(1, 2).squeeze(-1)) frames.append((x * 255).numpy().astype(np.uint8)) videos.append(frames) captions.append(batch.prompt) audio_waveforms.append(output_audio) audio_sample_rates.append(int(output_audio_sample_rate) if output_audio_sample_rate is not None else None) ref_videos.append(ref_video if isinstance(ref_video, str) else None) actions.append(action) mouse_pitch_signs.append(self._validation_mouse_pitch_sign(validation_batch)) if self.overlay_actions: overlay_frames = self._post_process_validation_frames( frames, action=action, ) if overlay_frames is not None: overlay_videos.append(overlay_frames) overlay_captions.append(batch.prompt) return _ValidationStepResult( videos=videos, captions=captions, audio_waveforms=audio_waveforms, audio_sample_rates=audio_sample_rates, overlay_videos=overlay_videos, overlay_captions=overlay_captions, ref_videos=ref_videos, actions=actions, mouse_pitch_signs=mouse_pitch_signs, ) def _post_process_validation_frames( self, frames: list[np.ndarray], *, action: dict[str, Any] | None, ) -> list[np.ndarray] | None: student = getattr(self.method, "student", None) post_process = getattr(student, "post_process_validation_frames", None) if post_process is None: return None return post_process( frames, action=action, ) # ---------------------------------------------------------- # State management # ---------------------------------------------------------- def state_dict(self) -> dict[str, Any]: state: dict[str, Any] = {} if self.validation_random_generator is not None: state["validation_rng"] = (self.validation_random_generator.get_state()) return state def load_state_dict( self, state_dict: dict[str, Any], ) -> None: rng_state = state_dict.get("validation_rng") if (rng_state is not None and self.validation_random_generator is not None): self.validation_random_generator.set_state(rng_state)