diff --git a/fastvideo/eval/README.md b/fastvideo/eval/README.md index 0f0fd64a18..ae0a9dbb13 100644 --- a/fastvideo/eval/README.md +++ b/fastvideo/eval/README.md @@ -1,9 +1,9 @@ # `fastvideo.eval` In-process evaluation suite for video generations. Includes pixel -metrics (SSIM, PSNR, LPIPS), optical-flow comparisons, the full VBench -suite, Physics-IQ, and a VLM scorer (VideoScore-2) behind a single -registry-driven API. +metrics (SSIM, PSNR, LPIPS), FVD (Fréchet Video Distance), optical-flow +comparisons, the full VBench suite, Physics-IQ, and a VLM scorer +(VideoScore-2) behind a single registry-driven API. ## Install @@ -78,7 +78,7 @@ fastvideo/ │ ├── datasets/ # prompt corpora (vbench, physics_iq) │ └── metrics/ │ ├── base.py # BaseMetric + @register contract -│ ├── common/ # SSIM, PSNR, LPIPS +│ ├── common/ # SSIM, PSNR, LPIPS, FVD │ ├── optical_flow/ # gt_optical_flow, synthetic_optical_flow │ ├── videoscore2/ # VideoScore-2 (Qwen2.5-VL) │ ├── physics_iq/ # PhysicsIQ + sub-metrics @@ -170,6 +170,7 @@ ${FASTVIDEO_CACHE_ROOT}/eval/ ├── models/ # URL-fetched checkpoints (LAION head, AMT, GRiT) ├── torch/ # redirected TORCH_HOME (DINO via torch.hub, lpips) ├── clip/ # passed as download_root= to clip.load callsites +├── fvd/ # common.fvd: real_features.pt (cached I3D reference features) └── datasets/ # auto-fetched dataset assets, one subdir per benchmark # (e.g. datasets/physics_iq/{split-videos,switch-frames,...}) ``` @@ -206,9 +207,8 @@ the metric's docstring if it matters. - **MIND** metrics. Depend on a separate `vipe` upstream submodule. - **VBench-2.0**. Sibling vbench2 package; needs its own port. -- **FVD as a registered metric**. Currently still at `benchmarks/fvd/`. - FVD is a set-vs-set distribution distance and does not fit the - per-sample `BaseMetric.compute` API without a stateful accumulator; - conversion is a designed follow-up. +- **FVD** is now available as `common.fvd` (see `metrics/common/fvd/`). + Uses the `is_set_metric` accumulate/finalize interface; requires + ≥ 256 videos for statistically reliable scores (standard: 2048). - **Training-time eval callback** (`EvalCallback`) and the `RolloutEvaluator` helper. diff --git a/fastvideo/eval/__init__.py b/fastvideo/eval/__init__.py index e4db83f45a..a1fd8ae20c 100644 --- a/fastvideo/eval/__init__.py +++ b/fastvideo/eval/__init__.py @@ -21,7 +21,7 @@ def _redirect_third_party_caches() -> None: _redirect_third_party_caches() -from fastvideo.eval.types import MetricResult # noqa: E402 +from fastvideo.eval.types import EvalResults, MetricResult, Video # noqa: E402 from fastvideo.eval.metrics.base import BaseMetric # noqa: E402 from fastvideo.eval.registry import register, list_metrics, get_metric # noqa: E402 from fastvideo.eval.api import evaluate # noqa: E402 @@ -34,7 +34,9 @@ def _redirect_third_party_caches() -> None: "evaluate", "Evaluator", "create_evaluator", + "EvalResults", "MetricResult", + "Video", "BaseMetric", "register", "list_metrics", diff --git a/fastvideo/eval/evaluator.py b/fastvideo/eval/evaluator.py index 5ab74430c0..3fccf8745d 100644 --- a/fastvideo/eval/evaluator.py +++ b/fastvideo/eval/evaluator.py @@ -3,22 +3,24 @@ Layering (mirrors FastVideo's VideoGenerator → Worker pattern, but in-process):: - Evaluator ← user-facing; round-robins samples across workers + Evaluator ← user-facing └── EvalWorker × N ← single-GPU; owns metric replicas + └── VideoPool ← async path-→-tensor prefetch (per evaluate call) The constructor builds one :class:`EvalWorker` per GPU and loads every metric on every worker eagerly. :meth:`evaluate` is the single entry point: pass kwargs for one sample, or pass a list of sample dicts to -fan-out across GPU replicas — same method, return type follows the -input shape. +fan-out across GPU replicas with pipelined decoding — same method, +return type follows the input shape. """ from __future__ import annotations -from concurrent.futures import ThreadPoolExecutor +import threading from collections.abc import Iterable +from typing import Any from fastvideo.eval.registry import (list_metrics, missing_dependencies, resolve_group) -from fastvideo.eval.types import MetricResult +from fastvideo.eval.types import EvalResults, MetricResult from fastvideo.eval.worker import EvalWorker from fastvideo.logger import init_logger @@ -38,6 +40,22 @@ class Evaluator: Number of GPU replicas. Each gets its own :class:`EvalWorker`. compile : bool Apply :func:`torch.compile` to each metric's ``_model``. + loader_threads : int + Background decode threads in the :class:`VideoPool`. Default 1 + (hide decode behind compute). Bump for I/O-heavy benchmark sets + where one loader can't keep up with the workers. + prefetch_factor : int + ``pool max_size = prefetch_factor * num_workers``. Default 2 — + one sample being consumed, one prefetched per worker. + pre_upload : bool + When ``True`` (default), the worker performs a single + host→device upload of ``video`` / ``reference`` per sample + before the metric loop, and every metric reads from that + shared GPU-resident tensor. Without it, each metric pays its + own ``.to(self.device)`` — N transfers of the same clip for N + metrics, which dominates at high resolution. Set ``False`` for + training-time eval, where keeping a clip resident on GPU + across the metric loop would fight the training step for VRAM. """ def __init__( @@ -46,13 +64,20 @@ def __init__( device: str = "cuda:0", num_gpus: int = 1, compile: bool = False, + *, + loader_threads: int = 1, + prefetch_factor: int = 2, + pre_upload: bool = True, ) -> None: names = _resolve_metric_names(metrics) if num_gpus > 1: - self._workers = [EvalWorker(names, f"cuda:{i}", compile=compile) for i in range(num_gpus)] + self._workers = [ + EvalWorker(names, f"cuda:{i}", compile=compile, pre_upload=pre_upload) for i in range(num_gpus) + ] else: - self._workers = [EvalWorker(names, device, compile=compile)] - self._pool = (ThreadPoolExecutor(max_workers=num_gpus) if num_gpus > 1 else None) + self._workers = [EvalWorker(names, device, compile=compile, pre_upload=pre_upload)] + self._loader_threads = max(1, loader_threads) + self._prefetch_factor = max(1, prefetch_factor) @property def num_gpus(self) -> int: @@ -66,22 +91,21 @@ def evaluate( self, samples: Iterable[dict] | None = None, **kwargs, - ) -> dict[str, MetricResult] | list[dict[str, MetricResult]]: + ) -> dict[str, MetricResult] | EvalResults: """Score one sample (kwargs form) or many samples (list form). - ``video`` and ``reference`` may be either a pre-loaded - ``(T, C, H, W)`` tensor or a path-like (``str`` / ``Path``). - Paths are decoded inside the worker thread that picks up the - sample, so memory stays bounded by ``num_gpus`` even when - thousands of paths are queued — see ``score_folder.py`` for - the canonical pattern. + Both forms go through the same :class:`VideoPool` pipeline; + ``video`` / ``reference`` paths are decoded asynchronously, + ``(1, T, C, H, W)`` tensors are squeezed. One sample:: ev.evaluate(video=tensor, text_prompt="...", fps=24.0) ev.evaluate(video="path/to/clip.mp4", fps=24.0) - Many samples — fan out across GPU replicas, results in input order:: + Returns a ``dict[str, MetricResult]``. + + Many samples — pipelined decode + work-stealing across replicas:: ev.evaluate(samples=[ {"video": "a.mp4", "reference": "ref_a.mp4"}, @@ -89,23 +113,79 @@ def evaluate( ... ]) - Multi-GPU dispatch fires automatically iff ``num_gpus > 1`` *and* - the list form is used. The kwargs form always runs on worker 0; - if you have a single sample but multiple GPUs, wrap it in a - one-element list to use the pool, or just accept that a single - call uses one GPU — that's fine. + Returns an :class:`EvalResults` (list-of-dict subclass): per-sample + dicts in input order, with set-metric scores under ``.corpus``. """ - if samples is None: - return self._workers[0].evaluate(**kwargs) + single = samples is None + sample_list: list[dict] = [kwargs] if samples is None else list(samples) + if not sample_list: + return EvalResults(samples=[], corpus={}) + + per_sample, corpus = self._run(sample_list) + + if single: + return per_sample[0] + return EvalResults(samples=per_sample, corpus=corpus) - samples = list(samples) - if self._pool is None or len(samples) <= 1: - return [self._workers[0].evaluate(**s) for s in samples] + def _run(self, samples: list[dict]) -> tuple[list[dict[str, MetricResult]], dict[str, MetricResult]]: + """Pool-driven sample pipeline + set-metric finalize. - n = len(self._workers) - # Round-robin: worker i handles samples i, i+n, i+2n, ... - futures = [self._pool.submit(self._workers[idx % n].evaluate, **sample) for idx, sample in enumerate(samples)] - return [f.result() for f in futures] + Returns ``(per_sample_results, corpus_results)``. + """ + from fastvideo.eval.pool import VideoPool + + # Reset every worker's set-metric buffers — per-call isolation. + for w in self._workers: + w.reset_set_metrics() + + n_workers = len(self._workers) + max_size = self._prefetch_factor * n_workers + per_sample: list[Any] = [None] * len(samples) + + with VideoPool(samples, loader_threads=self._loader_threads, max_size=max_size) as pool: + if n_workers == 1: + while True: + item = pool.get() + if item is None: + break + idx, decoded = item + per_sample[idx] = self._workers[0].evaluate(**decoded) + else: + # Multi-GPU: every worker drains the shared pool (work-stealing). + errors: list[BaseException] = [] + threads: list[threading.Thread] = [] + for w in self._workers: + t = threading.Thread(target=self._consumer_loop, args=(w, pool, per_sample, errors), daemon=True) + t.start() + threads.append(t) + for t in threads: + t.join() + if errors: + raise errors[0] + + # Finalize set metrics. With multiple workers, fold per-worker + # accumulator state into worker 0 first, then finalize once. + corpus: dict[str, MetricResult] = {} + base_set = self._workers[0].set_metrics() + if base_set: + for w in self._workers[1:]: + for name, m in w.set_metrics().items(): + base_set[name].merge_from(m) + corpus = {name: m.finalize() for name, m in base_set.items()} + + return per_sample, corpus + + @staticmethod + def _consumer_loop(worker: EvalWorker, pool: Any, results: list, errors: list) -> None: + try: + while True: + item = pool.get() + if item is None: + return + idx, decoded = item + results[idx] = worker.evaluate(**decoded) + except BaseException as e: # noqa: BLE001 — surface to parent thread via shared list + errors.append(e) def release_cuda_memory(self) -> None: """Free CUDA caches on every replica without dropping models.""" @@ -123,10 +203,7 @@ def reload(self) -> None: w.reload() def shutdown(self) -> None: - """Tear down the worker thread pool. Idempotent.""" - if self._pool is not None: - self._pool.shutdown(wait=True) - self._pool = None + """No-op; kept for API compatibility with older callers.""" def create_evaluator( diff --git a/fastvideo/eval/metrics/base.py b/fastvideo/eval/metrics/base.py index 14d79f8828..984acb8ffc 100644 --- a/fastvideo/eval/metrics/base.py +++ b/fastvideo/eval/metrics/base.py @@ -1,23 +1,31 @@ from __future__ import annotations -from abc import ABC, abstractmethod - import torch from fastvideo.eval.types import MetricResult -class BaseMetric(ABC): +class BaseMetric: """Abstract base class for all eval metrics. - Subclasses must implement :meth:`compute`. Optionally override - :meth:`setup` to eagerly load models. + Two execution shapes: + + * **Per-sample** (``is_set_metric=False``, default) — implement + :meth:`compute`. The Evaluator calls it once per input sample and + returns one :class:`MetricResult` per sample. + + * **Set-vs-set** (``is_set_metric=True``) — implement + :meth:`accumulate` (called once per sample to buffer features) + and :meth:`finalize` (called once after all samples to compute + the corpus-level result). Use :meth:`reset` to clear buffers and + :meth:`merge_from` to fold multi-GPU per-worker state together. - Metrics that need to chunk along the time dimension (frames or frame - pairs) for memory reasons should hardcode their own chunk size in - ``__init__`` (see ``optical_flow`` for the canonical example). Eval - always processes one video per :meth:`Evaluator.evaluate` call; - ``compute`` therefore receives a single sample, not a batch. + Optionally override :meth:`setup` to eagerly load models. Metrics + that chunk along the time dim for memory hardcode their own chunk + size in ``__init__`` (see ``optical_flow`` for the canonical + example). Eval always processes one video per + :meth:`Evaluator.evaluate` call; ``compute`` / ``accumulate`` + receive a single sample, not a batch. """ name: str = "" @@ -26,6 +34,7 @@ class BaseMetric(ABC): dependencies: list[str] = [] needs_gpu: bool = False backbone: str | None = None + is_set_metric: bool = False # Default time-dim chunk size for metrics that batch internally over # frames or frame-pairs. Override in subclass __init__ if needed @@ -56,14 +65,27 @@ def _skip(self, sample: dict, reason: str) -> MetricResult: """Return a skipped result (``score=None`` + reason in details).""" return MetricResult(name=self.name, score=None, details={"skipped": reason}) - @abstractmethod def compute(self, sample: dict) -> MetricResult: - """Compute the metric on a single sample. + """Per-sample metrics: compute the score for one sample. ``sample["video"]`` is ``(T, C, H, W)`` float in ``[0, 1]``. - ``sample["reference"]`` (if used) has the same shape. - - If required inputs are missing, return ``self._skip(sample, reason)`` - instead of raising. + ``sample["reference"]`` (if used) has the same shape. Return + ``self._skip(sample, reason)`` for missing inputs. """ - ... + raise NotImplementedError(f"{type(self).__name__}.compute is not implemented") + + # --- set-vs-set protocol (only invoked when is_set_metric=True) --- + + def reset(self) -> None: # noqa: B027 - intentionally optional override + """Clear accumulator state at the start of each evaluate() call.""" + + def accumulate(self, sample: dict) -> None: + """Buffer per-sample features for a corpus-level metric.""" + raise NotImplementedError(f"{type(self).__name__}.accumulate is not implemented") + + def finalize(self) -> MetricResult: + """Compute the corpus-level result from buffered state.""" + raise NotImplementedError(f"{type(self).__name__}.finalize is not implemented") + + def merge_from(self, other: BaseMetric) -> None: # noqa: B027 - intentionally optional override + """Multi-GPU: fold another worker's accumulator state into this one.""" diff --git a/fastvideo/eval/metrics/common/fvd/__init__.py b/fastvideo/eval/metrics/common/fvd/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/fastvideo/eval/metrics/common/fvd/metric.py b/fastvideo/eval/metrics/common/fvd/metric.py new file mode 100644 index 0000000000..a4ac5603dd --- /dev/null +++ b/fastvideo/eval/metrics/common/fvd/metric.py @@ -0,0 +1,342 @@ +"""common.fvd — Fréchet Video Distance. + +Measures distributional similarity between generated and reference videos +using I3D features (trained on Kinetics-400). Lower score → generated +videos are closer to the real video distribution. + +FVD is a dataset-level metric — it compares the +distribution of a collection of videos rather than scoring each video +individually. For statistically reliable results at least 256 videos are +recommended; the standard protocol uses 2048. + +This metric follows the set-vs-set protocol (is_set_metric=True): + - :meth:`accumulate` is called once per video to buffer I3D features. + - :meth:`finalize` is called once after all videos to compute FVD. + - :meth:`reset` clears buffers between evaluation runs. + +Reference features +------------------ +On the first call to :meth:`accumulate` that includes +``sample["reference"]``, I3D features are extracted from those reference +videos and saved to *cache_path*. Every subsequent run loads that cache +automatically — no need to pass ``sample["reference"]`` again. + +Cache default: ``${FASTVIDEO_EVAL_CACHE}/fvd/real_features.pt`` +(override via ``FASTVIDEO_EVAL_CACHE`` env-var; see +:func:`fastvideo.eval.models.get_cache_dir`). + +I3D model +--------- +Downloaded automatically from HuggingFace on first use via +:func:`fastvideo.eval.models.ensure_checkpoint` (filelock-safe): + flateon/FVD-I3D-torchscript (i3d_torchscript.pt) +Requires ``huggingface_hub`` (included in fastvideo deps). +""" + +from __future__ import annotations + +import os +import warnings + +import numpy as np +import scipy.linalg +import torch +import torch.nn.functional as F + +from fastvideo.eval.metrics.base import BaseMetric +from fastvideo.eval.registry import register +from fastvideo.eval.types import MetricResult + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +_I3D_REPO_ID = "flateon/FVD-I3D-torchscript" +_I3D_FILENAME = "i3d_torchscript.pt" +_I3D_MIN_FRAMES = 10 # I3D hard minimum +_MIN_VIDEOS_WARN = 256 # below this FVD is unreliable +_REAL_FEAT_RELPATH = "fvd/real_features.pt" # relative to get_cache_dir() + +# --------------------------------------------------------------------------- +# I3D helpers (adapted from FastVideo benchmarks/fvd/i3d_model.py) +# --------------------------------------------------------------------------- + + +def _load_i3d(device: torch.device) -> torch.nn.Module: + """Download I3D TorchScript from HuggingFace Hub and load it. + + Uses :func:`fastvideo.eval.models.ensure_checkpoint` so the download + is filelock-safe across threads, processes, and SLURM ranks. + """ + from fastvideo.eval.models import ensure_checkpoint + path = ensure_checkpoint(_I3D_FILENAME, source=_I3D_REPO_ID, filename=_I3D_FILENAME) + model = torch.jit.load(path, map_location=device) + model.eval() + return model + + +def _preprocess(video: torch.Tensor) -> torch.Tensor: + """(B, T, C, H, W) float [0, 1] → (B, C, T, 224, 224) float [-1, 1]. + + Matches I3D preprocessing from FastVideo benchmarks/fvd/feature_extractors.py. + """ + B, T, C, H, W = video.shape + if T < _I3D_MIN_FRAMES: + raise ValueError(f"I3D requires at least {_I3D_MIN_FRAMES} frames, got {T}. " + "Increase num_frames or use a longer video.") + if H != 224 or W != 224: + video = video.reshape(B * T, C, H, W) + video = F.interpolate(video, size=(224, 224), mode="bilinear", align_corners=False) + video = video.reshape(B, T, C, 224, 224) + # Scale [0, 1] → [-1, 1] and permute to (B, C, T, H, W) + video = video * 2.0 - 1.0 + return video.permute(0, 2, 1, 3, 4).contiguous() + + +@torch.no_grad() +def _extract_features( + model: torch.nn.Module, + video: torch.Tensor, # (B, T, C, H, W) float [0, 1] + chunk: int, + device: torch.device, +) -> np.ndarray: + """Extract I3D features → (B, 400) numpy array, chunked to fit VRAM. + + ``torch.jit.fuser("none")`` disables NVRTC kernel fusion for the I3D + TorchScript forward pass. Without it, PyTorch tries to JIT-compile fused + kernels via ``libnvrtc-builtins``, which is only available on the exact + CUDA version the binary was built against (e.g. fails on Colab CUDA 12 + when the lib expects CUDA 13). Disabling fusion has no effect on + numerical correctness — the I3D model still runs in full precision on GPU. + """ + parts = [] + with torch.jit.fuser("none"): + for i in range(0, video.shape[0], chunk): + batch = _preprocess(video[i:i + chunk].to(device)) + feats = model(batch, rescale=False, resize=False, return_features=True) + if feats.dim() == 1: + feats = feats.unsqueeze(0) # (D,) → (1, D) when I3D squeezes B=1 + parts.append(feats.cpu().numpy()) + return np.concatenate(parts, axis=0) + + +# --------------------------------------------------------------------------- +# Fréchet distance (adapted from FastVideo benchmarks/fvd/fvd.py) +# --------------------------------------------------------------------------- + + +def _gaussian_params(features: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Compute mean and covariance of a feature matrix (N, D). + + ``np.atleast_2d`` guards against a 1-D array (e.g. a single feature + vector squeezed by I3D or loaded from a stale cache), which would + cause ``np.cov`` to return a 0-d scalar and break ``sigma.shape[0]``. + """ + features = np.atleast_2d(features) + mu = features.mean(axis=0) + sigma = np.cov(features, rowvar=False) + if sigma.ndim == 0: # n==1 edge case: variance scalar + sigma = sigma.reshape(1, 1) + return mu, sigma + + +def _frechet_distance( + mu1: np.ndarray, + sigma1: np.ndarray, + mu2: np.ndarray, + sigma2: np.ndarray, + eps: float = 1e-6, +) -> float: + """Compute Fréchet distance between two Gaussians N(mu1,sigma1) and N(mu2,sigma2).""" + sigma1 = sigma1 + eps * np.eye(sigma1.shape[0]) + sigma2 = sigma2 + eps * np.eye(sigma2.shape[0]) + + diff = mu1 - mu2 + covmean = scipy.linalg.sqrtm(sigma1 @ sigma2) + + if np.iscomplexobj(covmean): + if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3): + warnings.warn( + f"FVD: large imaginary component in sqrtm " + f"({np.max(np.abs(covmean.imag)):.4f}). Result may be inaccurate.", + stacklevel=3, + ) + covmean = covmean.real + + return float(np.sum(diff**2) + np.trace(sigma1 + sigma2 - 2.0 * covmean)) + + +# --------------------------------------------------------------------------- +# Metric +# --------------------------------------------------------------------------- + + +@register("common.fvd") +class FVDMetric(BaseMetric): + """Fréchet Video Distance (FVD) using I3D features. + + Set-vs-set metric (``is_set_metric=True``). The Evaluator calls + :meth:`accumulate` once per video and :meth:`finalize` once after + all videos have been processed. + + For meaningful scores, evaluate over ≥ 256 videos (2048 is the + standard protocol used in the literature). + + Parameters + ---------- + cache_path : str, optional + Where extracted reference (real video) features are cached. + Built from ``sample["reference"]`` on the first run, reused + automatically on every subsequent run. Defaults to + ``${FASTVIDEO_EVAL_CACHE}/fvd/real_features.pt`` (resolved at + ``setup()`` time so the env-var can be set after import). + chunk_size : int + Videos per I3D forward pass. Reduce if GPU runs OOM. + """ + + name = "common.fvd" + is_set_metric = True + requires_reference = False # uses cached real features, not per-sample ref + higher_is_better = False # lower FVD = better + needs_gpu = True + dependencies = ["huggingface_hub", "scipy"] + + def __init__( + self, + cache_path: str | None = None, + chunk_size: int = 32, + ) -> None: + super().__init__() + # None → resolved lazily from get_cache_dir() in setup() so that + # FASTVIDEO_EVAL_CACHE is honoured even if set after import time. + self._cache_path_arg = cache_path + self._chunk = chunk_size + self._i3d: torch.nn.Module | None = None + + # Accumulated feature buffers — cleared by reset() + self._gen_features: list[np.ndarray] = [] + self._real_features: np.ndarray | None = None + + def to(self, device): + super().to(device) + if self._i3d is not None: + self._i3d = self._i3d.to(self.device) + return self + + def setup(self) -> None: + if self._i3d is not None: + return + from fastvideo.eval.models import get_cache_dir + # Resolve cache_path now so FASTVIDEO_EVAL_CACHE is read at run-time. + if self._cache_path_arg is None: + self.cache_path = str(get_cache_dir() / _REAL_FEAT_RELPATH) + else: + self.cache_path = os.path.expanduser(self._cache_path_arg) + self._i3d = _load_i3d(self.device) + # Pre-load cached real features if available + self._real_features = self._load_cache() + + # ------------------------------------------------------------------ + # Set-vs-set protocol + # ------------------------------------------------------------------ + + def reset(self) -> None: + """Clear generated feature buffer. Called before each evaluation run.""" + self._gen_features = [] + + def accumulate(self, sample: dict) -> None: + """Extract I3D features from one generated video and buffer them. + + If ``sample["reference"]`` is provided and no cache exists yet, + reference features are extracted and saved to *cache_path*. + """ + if self._i3d is None: + self.setup() + + video = sample["video"] # (T, C, H, W) from evaluator + if video.dim() == 4: + video = video.unsqueeze(0) # → (1, T, C, H, W) + + # Buffer generated features + feats = _extract_features(self._i3d, video, self._chunk, self.device) + self._gen_features.append(feats) + + # Build real feature cache on first encounter + if self._real_features is None: + self._real_features = self._load_cache() + + if self._real_features is None: + ref = sample.get("reference") + if ref is not None: + if ref.dim() == 4: + ref = ref.unsqueeze(0) + self._real_features = _extract_features(self._i3d, ref, self._chunk, self.device) + self._save_cache(self._real_features) + + def finalize(self) -> MetricResult: + """Compute FVD from all accumulated generated features vs. real features.""" + if not self._gen_features: + return MetricResult( + name=self.name, + score=None, + details={"skipped": "No generated videos accumulated before finalize()."}, + ) + + if self._real_features is None: + return MetricResult( + name=self.name, + score=None, + details={ + "skipped": ("No reference features available. Pass sample['reference'] " + "in at least one accumulate() call to build the cache at: " + f"{self.cache_path}") + }, + ) + + all_gen = np.concatenate(self._gen_features, axis=0) + n_gen = len(all_gen) + n_real = len(self._real_features) + + if n_gen < _MIN_VIDEOS_WARN or n_real < _MIN_VIDEOS_WARN: + warnings.warn( + f"FVD computed with only {n_gen} generated and {n_real} real videos. " + f"At least {_MIN_VIDEOS_WARN} recommended (standard protocol: 2048). " + "Score may not be statistically reliable.", + stacklevel=2, + ) + + mu_gen, sigma_gen = _gaussian_params(all_gen) + mu_real, sigma_real = _gaussian_params(self._real_features) + fvd = _frechet_distance(mu_gen, sigma_gen, mu_real, sigma_real) + + return MetricResult( + name=self.name, + score=fvd, + details={ + "n_generated": n_gen, + "n_reference": n_real, + }, + ) + + def merge_from(self, other: BaseMetric) -> None: + """Fold another worker's accumulated features into this one (multi-GPU).""" + assert isinstance(other, FVDMetric) + self._gen_features.extend(other._gen_features) + # Real features are identical across workers (same cache); keep ours. + if self._real_features is None and other._real_features is not None: + self._real_features = other._real_features + + # ------------------------------------------------------------------ + # Cache helpers + # ------------------------------------------------------------------ + + def _load_cache(self) -> np.ndarray | None: + if os.path.exists(self.cache_path): + data = torch.load(self.cache_path, map_location="cpu", weights_only=True) + arr = np.atleast_2d(data.numpy()) # guard against stale 1-D cache + return arr + return None + + def _save_cache(self, features: np.ndarray) -> None: + os.makedirs(os.path.dirname(self.cache_path), exist_ok=True) + torch.save(torch.from_numpy(features), self.cache_path) diff --git a/fastvideo/eval/metrics/common/lpips/metric.py b/fastvideo/eval/metrics/common/lpips/metric.py index 32747895e8..fec858088c 100644 --- a/fastvideo/eval/metrics/common/lpips/metric.py +++ b/fastvideo/eval/metrics/common/lpips/metric.py @@ -17,9 +17,12 @@ class LPIPSMetric(BaseMetric): needs_gpu = True dependencies = ["lpips"] - def __init__(self, net: str = "alex") -> None: + def __init__(self, net: str = "alex", chunk_size: int = 8) -> None: super().__init__() self.net = net + # Chunk the per-frame forward: AlexNet feature maps at 1080p + # peak ~60 GB on a 121-frame clip; chunk=8 caps it at ~5 GB. + self._chunk_size = chunk_size self._model: Any = None def to(self, device: str | torch.device) -> LPIPSMetric: @@ -39,8 +42,9 @@ def compute(self, sample: dict) -> MetricResult: if self._model is None: self.setup() - gen = sample["video"].float().to(self.device) # (T, C, H, W) - ref = sample["reference"].float().to(self.device) + gen = sample["video"].float().to(self.device, non_blocking=True) + ref = sample["reference"].float().to(self.device, non_blocking=True) + n = min(gen.shape[0], ref.shape[0]) gen, ref = gen[:n] * 2.0 - 1.0, ref[:n] * 2.0 - 1.0 diff --git a/fastvideo/eval/metrics/common/psnr/metric.py b/fastvideo/eval/metrics/common/psnr/metric.py index b6fa60b442..e96fe9dbeb 100644 --- a/fastvideo/eval/metrics/common/psnr/metric.py +++ b/fastvideo/eval/metrics/common/psnr/metric.py @@ -12,20 +12,29 @@ class PSNRMetric(BaseMetric): name = "common.psnr" requires_reference = True higher_is_better = True - needs_gpu = False + # On CPU the squared-diff is memory-bandwidth-bound at 1080p and + # fights the loader for the host bus; trivial on GPU. + needs_gpu = True - def __init__(self, max_val: float = 1.0) -> None: + def __init__(self, max_val: float = 1.0, chunk_size: int = 32) -> None: super().__init__() self.max_val = max_val + # Keeps the (gen - ref)**2 intermediate under ~800 MB at 1080p. + self._chunk_size = chunk_size def compute(self, sample: dict) -> MetricResult: - gen = sample["video"].float() # (T, C, H, W) - ref = sample["reference"].float() + gen = sample["video"].float().to(self.device) # (T, C, H, W) + ref = sample["reference"].float().to(self.device) n = min(gen.shape[0], ref.shape[0]) gen, ref = gen[:n], ref[:n] - # Per-frame MSE → PSNR. - mse = ((gen - ref)**2).mean(dim=(1, 2, 3)) # (T,) + chunk = self._chunk_size or n + mse_parts = [] + for i in range(0, n, chunk): + g = gen[i:i + chunk] + r = ref[i:i + chunk] + mse_parts.append(((g - r)**2).mean(dim=(1, 2, 3))) + mse = torch.cat(mse_parts) # (T,) psnr = 10.0 * torch.log10(self.max_val**2 / mse.clamp(min=1e-10)) return MetricResult( diff --git a/fastvideo/eval/metrics/common/ssim/metric.py b/fastvideo/eval/metrics/common/ssim/metric.py index ef4f9c9fe1..3a1c352735 100644 --- a/fastvideo/eval/metrics/common/ssim/metric.py +++ b/fastvideo/eval/metrics/common/ssim/metric.py @@ -51,15 +51,19 @@ class SSIMMetric(BaseMetric): name = "common.ssim" requires_reference = True higher_is_better = True - needs_gpu = False + # Five depthwise conv2d's per chunk — GPU is ~100× the CPU path at + # 1080p, and it frees the host bus for the loader thread. + needs_gpu = True - def __init__(self, window_size: int = 11) -> None: + def __init__(self, window_size: int = 11, chunk_size: int = 16) -> None: super().__init__() self.window_size = window_size + # Caps the ~5-7 conv intermediates at ~4 GB on a 1080p clip. + self._chunk_size = chunk_size def compute(self, sample: dict) -> MetricResult: - gen = sample["video"].float() # (T, C, H, W) - ref = sample["reference"].float() + gen = sample["video"].float().to(self.device) # (T, C, H, W) + ref = sample["reference"].float().to(self.device) n = min(gen.shape[0], ref.shape[0]) gen, ref = gen[:n], ref[:n] diff --git a/fastvideo/eval/metrics/optical_flow/_shared.py b/fastvideo/eval/metrics/optical_flow/_shared.py index a2830e4bcc..f9506dc79d 100644 --- a/fastvideo/eval/metrics/optical_flow/_shared.py +++ b/fastvideo/eval/metrics/optical_flow/_shared.py @@ -232,13 +232,14 @@ def aggregate_temporal(per_frame: list[dict[str, float]], ) -> dict[str, float | def tensor_to_bgr_list(video: torch.Tensor) -> list[np.ndarray]: - """Convert ``(T, C, H, W)`` float [0,1] to a list of HWC BGR uint8 frames.""" - frames = [] - for t in range(video.shape[0]): - rgb = (video[t].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8) - bgr = rgb[:, :, ::-1].copy() - frames.append(bgr) - return frames + """Convert ``(T, C, H, W)`` float [0,1] to a list of HWC BGR uint8 frames. + + Casts + permutes + BGR-swaps on-device and transfers once, instead + of T per-frame ``.cpu()`` round-trips. + """ + bgr_u8 = (video.float() * 255.0).clamp(0, 255).to(torch.uint8).permute(0, 2, 3, 1).flip(-1).contiguous() + arr = bgr_u8.cpu().numpy() + return [arr[t] for t in range(arr.shape[0])] def load_ptlflow_model(model_name: str, ckpt: str, device: torch.device): diff --git a/fastvideo/eval/metrics/optical_flow/gt_optical_flow/metric.py b/fastvideo/eval/metrics/optical_flow/gt_optical_flow/metric.py index 664f95fb4a..0bc67cb78c 100644 --- a/fastvideo/eval/metrics/optical_flow/gt_optical_flow/metric.py +++ b/fastvideo/eval/metrics/optical_flow/gt_optical_flow/metric.py @@ -53,7 +53,9 @@ def __init__( self.max_mag_pct = max_mag_pct self.grid_size = grid_size self._model = None - self._chunk_size = 16 + # One frame-pair per DPFlow forward: the cost volume is ~4 GB + # per pair at 1080p, so batching OOMs. Bump for low-res inputs. + self._chunk_size = 1 def to(self, device: str | torch.device) -> GtOpticalFlowMetric: super().to(device) diff --git a/fastvideo/eval/metrics/optical_flow/synthetic_optical_flow/metric.py b/fastvideo/eval/metrics/optical_flow/synthetic_optical_flow/metric.py index 0e4841a34c..389275a14f 100644 --- a/fastvideo/eval/metrics/optical_flow/synthetic_optical_flow/metric.py +++ b/fastvideo/eval/metrics/optical_flow/synthetic_optical_flow/metric.py @@ -94,7 +94,8 @@ def __init__( self._calibration: ThirdPersonCalibration | None = (_resolve_calibration(calibration_path) if calibration_path else None) self._model = None - self._chunk_size = 16 + # One frame-pair per DPFlow forward; see ``gt_optical_flow``. + self._chunk_size = 1 def to(self, device: str | torch.device) -> SyntheticOpticalFlowMetric: super().to(device) @@ -114,10 +115,6 @@ def compute(self, sample: dict) -> MetricResult: actions = sample.get("actions") if actions is None: return self._skip(sample, "missing 'actions' (keyboard + mouse)") - if isinstance(actions, list): - actions = actions[0] if actions else None - if actions is None: - return self._skip(sample, "empty 'actions' list") cal_obj = sample.get("calibration") cal = self._calibration if cal_obj is None else _resolve_calibration(cal_obj) diff --git a/fastvideo/eval/metrics/physics_iq/metric.py b/fastvideo/eval/metrics/physics_iq/metric.py index af8828c6e1..7a73bb3593 100644 --- a/fastvideo/eval/metrics/physics_iq/metric.py +++ b/fastvideo/eval/metrics/physics_iq/metric.py @@ -16,7 +16,6 @@ mean, prepare_pair_inputs, prepare_triplet_inputs, - unpack_batch_value, ) from fastvideo.eval.registry import register from fastvideo.eval.types import MetricResult @@ -172,22 +171,15 @@ def compute(self, sample: dict) -> MetricResult: if take2_key is None: raise KeyError("PhysicsIQMetric requires sample['reference_take2'] or an alias.") - # Polymorphic input handling: tensors / ndarrays / file paths / list-of-one. - # ``unpack_batch_value`` always yields a list; with B=1 we take element 0. - [video] = unpack_batch_value(sample["video"]) - [reference] = unpack_batch_value(sample["reference"]) - [reference_take2] = unpack_batch_value(sample[take2_key]) - generated_mask = unpack_batch_value(sample["video_mask"])[0] if "video_mask" in sample else None - reference_mask = unpack_batch_value(sample["reference_mask"])[0] if "reference_mask" in sample else None - reference_take2_mask = (unpack_batch_value(sample["reference_take2_mask"])[0] - if "reference_take2_mask" in sample else None) + video = sample["video"] + reference = sample["reference"] + reference_take2 = sample[take2_key] + generated_mask = sample.get("video_mask") + reference_mask = sample.get("reference_mask") + reference_take2_mask = sample.get("reference_take2_mask") scenario = sample.get("scenario") - if isinstance(scenario, list): - scenario = scenario[0] if scenario else None view = sample.get("view") - if isinstance(view, list): - view = view[0] if view else None details = self.compute_single( video, diff --git a/fastvideo/eval/metrics/physics_iq/mse/metric.py b/fastvideo/eval/metrics/physics_iq/mse/metric.py index 6fd58484ca..d6c3e36fd5 100644 --- a/fastvideo/eval/metrics/physics_iq/mse/metric.py +++ b/fastvideo/eval/metrics/physics_iq/mse/metric.py @@ -5,7 +5,7 @@ from fastvideo.eval.metrics.base import BaseMetric from fastvideo.eval.registry import register from fastvideo.eval.types import MetricResult -from fastvideo.eval.metrics.physics_iq.utils import compute_mse, prepare_pair_batch +from fastvideo.eval.metrics.physics_iq.utils import compute_mse, prepare_pair @register("physics_iq.mse") @@ -19,7 +19,7 @@ def __init__(self, **kwargs: Any) -> None: self._kwargs = kwargs def compute(self, sample: dict) -> MetricResult: - [prepared] = prepare_pair_batch(sample, prep_kwargs=self._kwargs) + prepared = prepare_pair(sample, prep_kwargs=self._kwargs) per_frame = compute_mse(prepared.reference_quarter, prepared.generated_quarter) score = sum(per_frame) / len(per_frame) return MetricResult(name=self.name, score=score, details={"per_frame": per_frame}) diff --git a/fastvideo/eval/metrics/physics_iq/spatial_iou/metric.py b/fastvideo/eval/metrics/physics_iq/spatial_iou/metric.py index 506c65f498..69a328743c 100644 --- a/fastvideo/eval/metrics/physics_iq/spatial_iou/metric.py +++ b/fastvideo/eval/metrics/physics_iq/spatial_iou/metric.py @@ -5,7 +5,7 @@ from fastvideo.eval.metrics.base import BaseMetric from fastvideo.eval.registry import register from fastvideo.eval.types import MetricResult -from fastvideo.eval.metrics.physics_iq.utils import compute_spatial_iou, prepare_pair_batch +from fastvideo.eval.metrics.physics_iq.utils import compute_spatial_iou, prepare_pair @register("physics_iq.spatial_iou") @@ -19,6 +19,6 @@ def __init__(self, **kwargs: Any) -> None: self._kwargs = kwargs def compute(self, sample: dict) -> MetricResult: - [prepared] = prepare_pair_batch(sample, prep_kwargs=self._kwargs) + prepared = prepare_pair(sample, prep_kwargs=self._kwargs) score = compute_spatial_iou(prepared.reference_masks, prepared.generated_masks) return MetricResult(name=self.name, score=score, details={}) diff --git a/fastvideo/eval/metrics/physics_iq/spatiotemporal_iou/metric.py b/fastvideo/eval/metrics/physics_iq/spatiotemporal_iou/metric.py index 440599cd95..e8e0a62b49 100644 --- a/fastvideo/eval/metrics/physics_iq/spatiotemporal_iou/metric.py +++ b/fastvideo/eval/metrics/physics_iq/spatiotemporal_iou/metric.py @@ -5,7 +5,7 @@ from fastvideo.eval.metrics.base import BaseMetric from fastvideo.eval.registry import register from fastvideo.eval.types import MetricResult -from fastvideo.eval.metrics.physics_iq.utils import compute_spatiotemporal_iou, prepare_pair_batch +from fastvideo.eval.metrics.physics_iq.utils import compute_spatiotemporal_iou, prepare_pair @register("physics_iq.spatiotemporal_iou") @@ -19,7 +19,7 @@ def __init__(self, **kwargs: Any) -> None: self._kwargs = kwargs def compute(self, sample: dict) -> MetricResult: - [prepared] = prepare_pair_batch(sample, prep_kwargs=self._kwargs) + prepared = prepare_pair(sample, prep_kwargs=self._kwargs) per_frame = compute_spatiotemporal_iou(prepared.reference_masks, prepared.generated_masks) score = sum(per_frame) / len(per_frame) return MetricResult(name=self.name, score=score, details={"per_frame": per_frame}) diff --git a/fastvideo/eval/metrics/physics_iq/utils.py b/fastvideo/eval/metrics/physics_iq/utils.py index a1e4268ab4..453ef3b535 100644 --- a/fastvideo/eval/metrics/physics_iq/utils.py +++ b/fastvideo/eval/metrics/physics_iq/utils.py @@ -86,63 +86,32 @@ def as_numpy_video(source: Any) -> tuple[np.ndarray, str]: raise TypeError(f"Unsupported Physics-IQ video source type: {type(source)!r}") -def unpack_batch_value(value: Any) -> list[Any]: - if isinstance(value, torch.Tensor): - if value.ndim == 4: - return [value] - if value.ndim == 5: - return [value[idx] for idx in range(value.shape[0])] - raise ValueError(f"Unsupported tensor rank for Physics-IQ metric: {value.ndim}") - if isinstance(value, np.ndarray): - if value.ndim == 4: - return [value] - if value.ndim == 5: - return [value[idx] for idx in range(value.shape[0])] - raise ValueError(f"Unsupported ndarray rank for Physics-IQ metric: {value.ndim}") - if isinstance(value, str | Path): - return [value] - if isinstance(value, list): - return value - raise TypeError(f"Unsupported batched Physics-IQ value type: {type(value)!r}") - - -def prepare_pair_batch( +def prepare_pair( sample: dict[str, Any], *, prep_kwargs: dict[str, Any] | None = None, -) -> list[PreparedPhysicsIQPair]: +) -> PreparedPhysicsIQPair: + """Resolve a sample into a prepared (gen, ref) pair. + + Caches the result on ``sample['_physics_iq_pair']`` so other physics_iq + sub-metrics on the same sample reuse it instead of re-decoding. + """ prepared = sample.get("_physics_iq_pair") if prepared is not None: - return [prepared] + return prepared if "reference" not in sample: raise KeyError("Physics-IQ pair metrics require sample['reference'].") - videos = unpack_batch_value(sample["video"]) - references = unpack_batch_value(sample["reference"]) - generated_masks = unpack_batch_value(sample["video_mask"]) if "video_mask" in sample else [None] * len(videos) - reference_masks = unpack_batch_value( - sample["reference_mask"]) if "reference_mask" in sample else [None] * len(videos) - - if not (len(videos) == len(references) == len(generated_masks) == len(reference_masks)): - raise ValueError("Physics-IQ pair metric inputs must have the same batch size.") - - prep_kwargs = prep_kwargs or {} - return [ - prepare_pair_inputs( - video, - reference, - generated_mask=generated_mask, - reference_mask=reference_mask, - **prep_kwargs, - ) for video, reference, generated_mask, reference_mask in zip( - videos, - references, - generated_masks, - reference_masks, - strict=False, - ) - ] + prepared = prepare_pair_inputs( + sample["video"], + sample["reference"], + generated_mask=sample.get("video_mask"), + reference_mask=sample.get("reference_mask"), + **(prep_kwargs or {}), + ) + sample["_physics_iq_pair"] = prepared + return prepared def select_window(frames: np.ndarray, *, target_frames: int, selection: str = "first") -> np.ndarray: diff --git a/fastvideo/eval/metrics/physics_iq/weighted_spatial_iou/metric.py b/fastvideo/eval/metrics/physics_iq/weighted_spatial_iou/metric.py index 55b60793eb..3b50336109 100644 --- a/fastvideo/eval/metrics/physics_iq/weighted_spatial_iou/metric.py +++ b/fastvideo/eval/metrics/physics_iq/weighted_spatial_iou/metric.py @@ -5,7 +5,7 @@ from fastvideo.eval.metrics.base import BaseMetric from fastvideo.eval.registry import register from fastvideo.eval.types import MetricResult -from fastvideo.eval.metrics.physics_iq.utils import compute_weighted_spatial_iou, prepare_pair_batch +from fastvideo.eval.metrics.physics_iq.utils import compute_weighted_spatial_iou, prepare_pair @register("physics_iq.weighted_spatial_iou") @@ -19,6 +19,6 @@ def __init__(self, **kwargs: Any) -> None: self._kwargs = kwargs def compute(self, sample: dict) -> MetricResult: - [prepared] = prepare_pair_batch(sample, prep_kwargs=self._kwargs) + prepared = prepare_pair(sample, prep_kwargs=self._kwargs) score = compute_weighted_spatial_iou(prepared.reference_masks, prepared.generated_masks) return MetricResult(name=self.name, score=score, details={}) diff --git a/fastvideo/eval/metrics/vbench/appearance_style/metric.py b/fastvideo/eval/metrics/vbench/appearance_style/metric.py index 89b09b7466..2f765cdf35 100644 --- a/fastvideo/eval/metrics/vbench/appearance_style/metric.py +++ b/fastvideo/eval/metrics/vbench/appearance_style/metric.py @@ -69,8 +69,6 @@ def compute(self, sample: dict) -> MetricResult: text_prompt = sample.get("text_prompt") if text_prompt is None: return self._skip(sample, "missing text_prompt") - if isinstance(text_prompt, list): - text_prompt = text_prompt[0] frames = _clip_transform(video.to(self.device)) diff --git a/fastvideo/eval/metrics/vbench/color/metric.py b/fastvideo/eval/metrics/vbench/color/metric.py index 854084ebc1..250f5ddceb 100644 --- a/fastvideo/eval/metrics/vbench/color/metric.py +++ b/fastvideo/eval/metrics/vbench/color/metric.py @@ -55,17 +55,11 @@ def compute(self, sample: dict) -> MetricResult: from fastvideo.eval.metrics.vbench._grit_helper import prepare_frames video = sample["video"] # (T, C, H, W) - aux = sample.get("auxiliary_info") - if isinstance(aux, list): - aux = aux[0] if aux else None - if not aux or "color" not in aux: + aux = sample.get("auxiliary_info") or {} + if "color" not in aux: return self._skip(sample, "missing 'color' in auxiliary_info") - text_prompt = sample.get("text_prompt") - if isinstance(text_prompt, list): - text_prompt = text_prompt[0] if text_prompt else "" - prompt = text_prompt or "" - + prompt = sample.get("text_prompt") or "" color_key = aux["color"] # Parse object name: remove "a ", "an ", and the color word object_key = prompt.replace("a ", "").replace("an ", "").replace(color_key, "").strip() diff --git a/fastvideo/eval/metrics/vbench/human_action/metric.py b/fastvideo/eval/metrics/vbench/human_action/metric.py index de60ad6655..af62b6cfbd 100644 --- a/fastvideo/eval/metrics/vbench/human_action/metric.py +++ b/fastvideo/eval/metrics/vbench/human_action/metric.py @@ -101,8 +101,6 @@ def compute(self, sample: dict) -> MetricResult: text_prompt = sample.get("text_prompt") if text_prompt is None: return self._skip(sample, "missing text_prompt with action labels") - if isinstance(text_prompt, list): - text_prompt = text_prompt[0] cat_dict = _load_cat_dict() diff --git a/fastvideo/eval/metrics/vbench/multiple_objects/metric.py b/fastvideo/eval/metrics/vbench/multiple_objects/metric.py index 2fce255828..31a19eb486 100644 --- a/fastvideo/eval/metrics/vbench/multiple_objects/metric.py +++ b/fastvideo/eval/metrics/vbench/multiple_objects/metric.py @@ -38,10 +38,8 @@ def compute(self, sample: dict) -> MetricResult: from fastvideo.eval.metrics.vbench._grit_helper import prepare_frames, detect_frames video = sample["video"] # (T, C, H, W) - aux = sample.get("auxiliary_info") - if isinstance(aux, list): - aux = aux[0] if aux else None - if not aux or "object" not in aux: + aux = sample.get("auxiliary_info") or {} + if "object" not in aux: return self._skip(sample, "missing 'object' in auxiliary_info") object_info = aux["object"] diff --git a/fastvideo/eval/metrics/vbench/object_class/metric.py b/fastvideo/eval/metrics/vbench/object_class/metric.py index 9cee4e8474..e6e360b018 100644 --- a/fastvideo/eval/metrics/vbench/object_class/metric.py +++ b/fastvideo/eval/metrics/vbench/object_class/metric.py @@ -38,10 +38,8 @@ def compute(self, sample: dict) -> MetricResult: from fastvideo.eval.metrics.vbench._grit_helper import prepare_frames, detect_frames video = sample["video"] # (T, C, H, W) - aux = sample.get("auxiliary_info") - if isinstance(aux, list): - aux = aux[0] if aux else None - if not aux or "object" not in aux: + aux = sample.get("auxiliary_info") or {} + if "object" not in aux: return self._skip(sample, "missing 'object' in auxiliary_info") object_key = aux["object"] diff --git a/fastvideo/eval/metrics/vbench/overall_consistency/metric.py b/fastvideo/eval/metrics/vbench/overall_consistency/metric.py index 1cc29487af..ff944db16a 100644 --- a/fastvideo/eval/metrics/vbench/overall_consistency/metric.py +++ b/fastvideo/eval/metrics/vbench/overall_consistency/metric.py @@ -80,8 +80,6 @@ def compute(self, sample: dict) -> MetricResult: text_prompt = sample.get("text_prompt") if text_prompt is None: return self._skip(sample, "missing text_prompt") - if isinstance(text_prompt, list): - text_prompt = text_prompt[0] frames = _clip_transform(extract_frames(video, 8)) # (8, C, H, W) clip_in = frames.unsqueeze(0).to(self.device) # (1, 8, C, H, W) diff --git a/fastvideo/eval/metrics/vbench/scene/metric.py b/fastvideo/eval/metrics/vbench/scene/metric.py index bd76f545e4..d8dc70354a 100644 --- a/fastvideo/eval/metrics/vbench/scene/metric.py +++ b/fastvideo/eval/metrics/vbench/scene/metric.py @@ -142,10 +142,8 @@ def _generate_caption(self, video_path: str) -> str: @torch.no_grad() def compute(self, sample: dict) -> MetricResult: video = sample["video"] # (T, C, H, W) - aux = sample.get("auxiliary_info") - if isinstance(aux, list): - aux = aux[0] if aux else None - if aux is None or "scene" not in aux: + aux = sample.get("auxiliary_info") or {} + if "scene" not in aux: return self._skip(sample, "missing 'scene' in auxiliary_info") scene_keywords = aux["scene"] diff --git a/fastvideo/eval/metrics/vbench/spatial_relationship/metric.py b/fastvideo/eval/metrics/vbench/spatial_relationship/metric.py index 497718aa03..b18858c29a 100644 --- a/fastvideo/eval/metrics/vbench/spatial_relationship/metric.py +++ b/fastvideo/eval/metrics/vbench/spatial_relationship/metric.py @@ -77,10 +77,8 @@ def compute(self, sample: dict) -> MetricResult: from fastvideo.eval.metrics.vbench._grit_helper import prepare_frames video = sample["video"] # (T, C, H, W) - aux = sample.get("auxiliary_info") - if isinstance(aux, list): - aux = aux[0] if aux else None - if not aux or "spatial_relationship" not in aux: + aux = sample.get("auxiliary_info") or {} + if "spatial_relationship" not in aux: return self._skip(sample, "missing 'spatial_relationship' in auxiliary_info") sp_info = aux["spatial_relationship"] diff --git a/fastvideo/eval/pool.py b/fastvideo/eval/pool.py new file mode 100644 index 0000000000..cb5dd795cc --- /dev/null +++ b/fastvideo/eval/pool.py @@ -0,0 +1,159 @@ +"""Async path-→-tensor prefetcher for the Evaluator. + +Hides video-decode latency behind metric compute by running a small +thread pool of decoders that fill a bounded queue. One :class:`VideoPool` +is owned by the Evaluator per ``evaluate(samples=...)`` call; workers +consume via :meth:`VideoPool.get`. Decode order is non-deterministic; +each yielded item carries its original input index so consumers can +write back into a result list in input order. + +Pool sizing: ``max_size = prefetch_factor * num_workers``. +""" +from __future__ import annotations + +import queue +import threading +from pathlib import Path +from typing import Any + +import torch + +from fastvideo.eval.types import Video + +_SENTINEL = object() + + +class _DecodeError: + """Marker pushed onto the ready queue when a loader thread raises. + + The consumer re-raises in its own thread, surfacing the error to + the caller of ``Evaluator.evaluate`` instead of hanging on + ``_ready_q.get()`` forever. + """ + + __slots__ = ("exc", ) + + def __init__(self, exc: BaseException) -> None: + self.exc = exc + + +class VideoPool: + """Bounded prefetch queue feeding decoded samples to consumers. + + Use as a context manager so loader threads are always cleaned up:: + + with VideoPool(samples, loader_threads=1, max_size=4) as pool: + while True: + item = pool.get() + if item is None: + break + idx, decoded = item + results[idx] = worker.evaluate(**decoded) + """ + + def __init__( + self, + samples: list[dict], + *, + loader_threads: int = 1, + max_size: int = 4, + ) -> None: + if loader_threads < 1: + raise ValueError("loader_threads must be >= 1") + self._samples = samples + self._loader_threads_n = loader_threads + self._max_size = max(max_size, 1) + + self._task_q: queue.Queue = queue.Queue() + self._ready_q: queue.Queue = queue.Queue(maxsize=self._max_size) + self._loaders: list[threading.Thread] = [] + self._stop = threading.Event() + + self._consumed = 0 + self._consume_lock = threading.Lock() + + def __enter__(self) -> VideoPool: + for idx, sample in enumerate(self._samples): + self._task_q.put((idx, sample)) + for _ in range(self._loader_threads_n): + self._task_q.put(_SENTINEL) + for _ in range(self._loader_threads_n): + t = threading.Thread(target=self._loader_loop, daemon=True) + t.start() + self._loaders.append(t) + return self + + def __exit__(self, *_exc: Any) -> None: + self._stop.set() + # Drain ready queue so any blocked-on-put loader unblocks. + while True: + try: + self._ready_q.get_nowait() + except queue.Empty: + break + for t in self._loaders: + t.join(timeout=5.0) + + def get(self, timeout: float | None = None) -> tuple[int, dict] | None: + """Pop the next decoded ``(idx, sample)``. + + Returns ``None`` when all input samples have been consumed. + Re-raises any exception caught in a loader thread on the + consumer's stack so callers don't hang on a dead loader. + Thread-safe: multiple consumer threads may share one pool. + """ + with self._consume_lock: + if self._consumed >= len(self._samples): + return None + try: + item = self._ready_q.get(timeout=timeout) + except queue.Empty: + return None + with self._consume_lock: + self._consumed += 1 + idx, payload = item + if isinstance(payload, _DecodeError): + raise payload.exc + return item + + def _loader_loop(self) -> None: + while not self._stop.is_set(): + item = self._task_q.get() + if item is _SENTINEL: + return + idx, sample = item + try: + decoded: Any = self._decode(sample) + except BaseException as exc: # noqa: BLE001 — forward to consumer + self._ready_q.put((idx, _DecodeError(exc))) + continue + # Blocking put: under normal flow the consumer drains the + # queue; under shutdown ``__exit__`` drains it for us. A + # timeout would silently drop samples and hang the consumer. + self._ready_q.put((idx, decoded)) + + def _decode(self, sample: dict) -> dict: + """Materialize and normalize ``video`` / ``reference`` entries. + + * ``Video`` instance → populate ``.frames`` via decode. + * ``str`` / ``Path`` under ``video`` / ``reference`` → decoded + ``(T, C, H, W)`` tensor. + * ``(1, T, C, H, W)`` tensor under ``video`` / ``reference`` → + squeezed to ``(T, C, H, W)`` (back-compat with callers that + still pass a leading batch dim). + * Everything else passes through unchanged. + """ + from fastvideo.eval.io.video import load_video + + out = dict(sample) + for key, val in sample.items(): + if isinstance(val, Video): + if val.frames is None and val.source is not None: + val.frames = load_video(val.source) + out[key] = val + elif key in ("video", "reference"): + if isinstance(val, str | Path): + out[key] = load_video(str(val)) + elif isinstance(val, torch.Tensor) and val.dim() == 5 and val.shape[0] == 1: + out[key] = val.squeeze(0) + return out diff --git a/fastvideo/eval/types.py b/fastvideo/eval/types.py index b25f7e2c03..63d89be3b8 100644 --- a/fastvideo/eval/types.py +++ b/fastvideo/eval/types.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass, field +from pathlib import Path from typing import Any @@ -14,3 +15,46 @@ class MetricResult: name: str score: float | None details: dict[str, Any] = field(default_factory=dict) + + +class EvalResults(list): + """Return type for :meth:`Evaluator.evaluate` with ``samples=...``. + + Behaves like a ``list[dict[str, MetricResult]]`` — one dict per + input sample, in input order — so existing iteration and indexing + keeps working. The ``corpus`` attribute carries set-metric results + (FAD, IS, …) that are properties of the whole input set, not of + any individual sample. Empty dict when no set metric ran. + """ + + def __init__( + self, + samples: list[dict[str, MetricResult]] | None = None, + corpus: dict[str, MetricResult] | None = None, + ) -> None: + super().__init__(samples or []) + self.corpus: dict[str, MetricResult] = corpus or {} + + +@dataclass +class Video: + """Path-backed media handle. The :class:`VideoPool` populates + ``frames`` (and optionally ``audio``) before the metric loop sees + the sample. + """ + + source: Any + fps: float | None = None + frames: Any = None + audio: Any = None + audio_sr: int | None = None + + def has_frames(self) -> bool: + return self.frames is not None + + def has_audio(self) -> bool: + return self.audio is not None + + def __post_init__(self) -> None: + if isinstance(self.source, Path): + self.source = str(self.source) diff --git a/fastvideo/eval/worker.py b/fastvideo/eval/worker.py index d4a7b638fa..7d512d0fed 100644 --- a/fastvideo/eval/worker.py +++ b/fastvideo/eval/worker.py @@ -13,25 +13,33 @@ """ from __future__ import annotations -from pathlib import Path +import contextlib from typing import Any import torch from fastvideo.eval.memory import clear_cache +from fastvideo.eval.metrics.base import BaseMetric from fastvideo.eval.registry import get_metric from fastvideo.eval.types import MetricResult -import contextlib class EvalWorker: - """Owns metric replicas on one device. Single-GPU, single-sample.""" + """Owns metric replicas on one device. Single-GPU, single-sample. - def __init__(self, metric_names: list[str], device: str, *, compile: bool = False) -> None: + Per-sample metrics (``is_set_metric=False``) return one + :class:`MetricResult` per ``evaluate`` call. Set metrics + (``is_set_metric=True``) accumulate state on the worker's own + instance and contribute nothing to the per-sample return — the + Evaluator finalizes them after the pool drains. + """ + + def __init__(self, metric_names: list[str], device: str, *, compile: bool = False, pre_upload: bool = True) -> None: self._names = list(metric_names) self._device = device self._compile = compile - self._metrics: dict = {} + self._pre_upload = pre_upload + self._metrics: dict[str, BaseMetric] = {} self._unloaded = False self._load() @@ -55,29 +63,40 @@ def _load(self) -> None: self._unloaded = False def evaluate(self, **kwargs) -> dict[str, MetricResult]: - """Score one sample. - - ``video`` may be a ``(T, C, H, W)`` tensor or a path-like - (``str`` / ``Path``) — paths are loaded inside this method so - the dispatcher can hold a queue of cheap path strings instead - of fully-decoded tensors. ``reference`` follows the same rule. + """Score one already-decoded sample. - A ``(1, T, C, H, W)`` tensor is also accepted for back-compat - and gets unwrapped to ``(T, C, H, W)`` before reaching metrics. + Inputs (``video`` / ``reference`` paths, 5-D tensors) are + decoded and normalized upstream by the :class:`VideoPool`. The + worker only handles the optional pre-upload to its device and + dispatches to each metric's ``compute`` or ``accumulate``. """ if self._unloaded: raise RuntimeError("EvalWorker was unloaded; call reload() before evaluating.") sample = dict(kwargs) - sample["video"] = _resolve_video_input(sample.get("video")) - if "reference" in sample: - sample["reference"] = _resolve_video_input(sample["reference"]) + if self._pre_upload: + sample["video"] = _to_device(sample.get("video"), self._device) + if "reference" in sample: + sample["reference"] = _to_device(sample["reference"], self._device) results: dict[str, MetricResult] = {} for name, m in self._metrics.items(): - results[name] = m.compute(sample) + if m.is_set_metric: + m.accumulate(sample) + else: + results[name] = m.compute(sample) return results + def set_metrics(self) -> dict[str, BaseMetric]: + """Return ``{name: instance}`` for set metrics on this worker.""" + return {n: m for n, m in self._metrics.items() if m.is_set_metric} + + def reset_set_metrics(self) -> None: + """Clear accumulator state on every set metric.""" + for m in self._metrics.values(): + if m.is_set_metric: + m.reset() + def release_cuda_memory(self) -> None: """Free CUDA caches without dropping models.""" clear_cache() @@ -97,22 +116,11 @@ def reload(self) -> None: self._load() -def _resolve_video_input(value: Any) -> Any: - """Normalize a sample's ``video`` / ``reference`` field for metrics. - - * ``str`` / ``Path`` → decoded ``(T, C, H, W)`` tensor via - :func:`fastvideo.eval.io.video.load_video`. Decoding happens in - the worker thread so the dispatcher can keep paths queued - instead of full tensors. - * ``(1, T, C, H, W)`` tensor → squeezed to ``(T, C, H, W)`` - (back-compat with callers that still pass the leading batch dim). - * anything else → returned untouched. - """ - if value is None: - return None - if isinstance(value, str | Path): - from fastvideo.eval.io.video import load_video - return load_video(str(value)) - if isinstance(value, torch.Tensor) and value.dim() == 5 and value.shape[0] == 1: - return value.squeeze(0) - return value +def _to_device(value: Any, device: str | torch.device) -> Any: + """Move *value* to *device* if it's a tensor not already there.""" + if value is None or not isinstance(value, torch.Tensor): + return value + target = torch.device(device) + if value.device == target: + return value + return value.to(target, non_blocking=True)