# 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 os from dataclasses import dataclass from typing import Any, TYPE_CHECKING import imageio 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.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: videos: list[list[np.ndarray]] captions: list[str] 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, sampling_steps: list[int] | None = None, guidance_scale: float | None = None, num_frames: int | None = None, output_dir: str | None = None, sampling_timesteps: list[int] | None = None, offload_training_state: bool = False, unload_pipeline_after_validation: bool = False, **pipeline_kwargs: Any, ) -> None: self.pipeline_target = str(pipeline_target) self.dataset_file = str(dataset_file) self.every_steps = int(every_steps) 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.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.offload_training_state = self._coerce_bool(offload_training_state) self.unload_pipeline_after_validation = self._coerce_bool(unload_pipeline_after_validation) 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.tracker: Any = DummyTracker() self.validation_random_generator: (torch.Generator | None) = None self.seed: int = 0 @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: if self.every_steps <= 0: 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._run_validation_inner( method, step, t, ) finally: if self.unload_pipeline_after_validation: self._clear_pipeline_cache() @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: 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) for num_inference_steps in self.sampling_steps: result = self._run_validation_for_steps( num_inference_steps, transformer=transformer, ) if self.rank_in_sp_group != 0: continue if self.global_rank == 0: all_videos = list(result.videos) all_captions = list(result.captions) for sp_idx in range(1, num_sp_groups): 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)) all_videos.extend(recv_v) all_captions.extend(recv_c) os.makedirs( output_dir, exist_ok=True, ) video_filenames: list[str] = [] sp = self._get_sampling_param() for i, video in enumerate(all_videos): fname = os.path.join( output_dir, f"validation_step_{step}" f"_inference_steps_" f"{num_inference_steps}" f"_video_{i}.mp4", ) imageio.mimsave( fname, video, fps=sp.fps, ) video_filenames.append(fname) video_logs = [] for fname, cap in zip( video_filenames, all_captions, strict=True, ): art = self.tracker.video( fname, caption=cap, ) if art is not None: video_logs.append(art) if video_logs: logs = { f"validation_videos_" f"{num_inference_steps}" f"_steps": video_logs } self.tracker.log_artifacts( logs, step, ) else: self.world_group.send_object( result.videos, dst=0, ) self.world_group.send_object( result.captions, dst=0, ) finally: if was_training: transformer.train() # ---------------------------------------------------------- # Pipeline management # ---------------------------------------------------------- def _get_sampling_param(self) -> SamplingParam: if self._sampling_param is None: self._sampling_param = (SamplingParam.from_pretrained(self.training_config.model_path)) return self._sampling_param 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, ) kwargs: dict[str, Any] = { "inference_mode": True, "loaded_modules": { "transformer": transformer, }, "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) self._pipeline = PipelineCls.from_pretrained( tc.model_path, **kwargs, ) scheduler = self._pipeline.get_module("scheduler") if (scheduler is not None and type(scheduler).__name__ == "SelfForcingFlowMatchScheduler"): scheduler.sigma_min = 0.0 scheduler.extra_one_step = True scheduler.set_timesteps(num_inference_steps=1000, training=True) 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: 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 # image_path for I2V pipelines. 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, ) 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"] return batch # ---------------------------------------------------------- # Validation loop # ---------------------------------------------------------- def _run_validation_for_steps( self, num_inference_steps: int, *, transformer: torch.nn.Module, ) -> _ValidationStepResult: 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, ) # 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]] = [] captions: list[str] = [] 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)) captions.append(batch.prompt) with torch.no_grad(): output_batch = pipeline.forward( batch, inference_args, ) samples = output_batch.output.cpu() if self.rank_in_sp_group != 0: continue 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) return _ValidationStepResult( videos=videos, captions=captions, ) # ---------------------------------------------------------- # 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)