From 608657a2f5a530c2eecb9b85881b05f0f0c7d3a9 Mon Sep 17 00:00:00 2001 From: shaoxiongduan Date: Mon, 11 May 2026 04:48:45 +0000 Subject: [PATCH 01/10] [feat] eval: async VideoPool + GPU-side common metrics + safe optical-flow chunks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pipelines path → tensor decode behind GPU metric compute via a new VideoPool, so multi-sample eval runs no longer serialize disk I/O and metric work. The Evaluator owns one pool per evaluate(samples=...) call; each EvalWorker is a single-GPU consumer that grabs decoded samples from the shared queue (work-stealing across replicas when num_gpus > 1). Worker pre-uploads video/reference to its device once per sample so every metric in the loop consumes the same GPU-resident tensor (no per-metric .to(device) traffic). SSIM and PSNR move to the GPU — at 1080p × 121 frames the CPU path was both slow (5–10 s/pair) and contended with the loader thread for DDR bandwidth. LPIPS gains a chunk_size knob (default 8) that caps peak from ~60 GB to ~5 GB with bit-identical output. Optical-flow metrics drop chunk_size to 1 because DPFlow's cost volume is ~4 GB per frame pair at 1080p (matches mhuo/ptlflow upstream). physics_iq and a handful of vbench metrics drop their list-batch shims — the per-sample contract is now uniform across the suite. --- fastvideo/eval/__init__.py | 3 +- fastvideo/eval/evaluator.py | 113 ++++++++++++---- fastvideo/eval/metrics/common/lpips/metric.py | 10 +- fastvideo/eval/metrics/common/psnr/metric.py | 21 ++- fastvideo/eval/metrics/common/ssim/metric.py | 12 +- .../eval/metrics/optical_flow/_shared.py | 15 ++- .../optical_flow/gt_optical_flow/metric.py | 4 +- .../synthetic_optical_flow/metric.py | 7 +- fastvideo/eval/metrics/physics_iq/metric.py | 20 +-- .../eval/metrics/physics_iq/mse/metric.py | 4 +- .../metrics/physics_iq/spatial_iou/metric.py | 4 +- .../physics_iq/spatiotemporal_iou/metric.py | 4 +- fastvideo/eval/metrics/physics_iq/utils.py | 63 +++------ .../physics_iq/weighted_spatial_iou/metric.py | 4 +- .../metrics/vbench/appearance_style/metric.py | 2 - fastvideo/eval/metrics/vbench/color/metric.py | 12 +- .../metrics/vbench/human_action/metric.py | 2 - .../metrics/vbench/multiple_objects/metric.py | 6 +- .../metrics/vbench/object_class/metric.py | 6 +- .../vbench/overall_consistency/metric.py | 2 - fastvideo/eval/metrics/vbench/scene/metric.py | 6 +- .../vbench/spatial_relationship/metric.py | 6 +- fastvideo/eval/pool.py | 127 ++++++++++++++++++ fastvideo/eval/types.py | 25 ++++ fastvideo/eval/worker.py | 46 ++++--- 25 files changed, 351 insertions(+), 173 deletions(-) create mode 100644 fastvideo/eval/pool.py diff --git a/fastvideo/eval/__init__.py b/fastvideo/eval/__init__.py index e4db83f45a..0fe56ed5a8 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 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 @@ -35,6 +35,7 @@ def _redirect_third_party_caches() -> None: "Evaluator", "create_evaluator", "MetricResult", + "Video", "BaseMetric", "register", "list_metrics", diff --git a/fastvideo/eval/evaluator.py b/fastvideo/eval/evaluator.py index 5ab74430c0..3c1adba086 100644 --- a/fastvideo/eval/evaluator.py +++ b/fastvideo/eval/evaluator.py @@ -3,19 +3,21 @@ 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 @@ -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: @@ -71,17 +96,16 @@ def evaluate( ``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. + Paths in the list form are decoded asynchronously by a + :class:`VideoPool` that runs alongside metric compute, hiding + decode latency behind GPU work. 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:: + Many samples — pipelined decode + work-stealing across replicas:: ev.evaluate(samples=[ {"video": "a.mp4", "reference": "ref_a.mp4"}, @@ -89,23 +113,59 @@ 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. + Multi-GPU dispatch fires automatically when ``num_gpus > 1`` and + the list form is used: every worker runs a consumer thread, + pulling decoded samples from the shared pool as it frees up. + The kwargs form always runs on worker 0 with no pool overhead. """ if samples is None: return self._workers[0].evaluate(**kwargs) samples = list(samples) - if self._pool is None or len(samples) <= 1: - return [self._workers[0].evaluate(**s) for s in samples] - - 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] + if not samples: + return [] + return self._evaluate_with_pool(samples) + + def _evaluate_with_pool(self, samples: list[dict]) -> list[dict[str, MetricResult]]: + """Run samples through a :class:`VideoPool`, writing results in input order.""" + from fastvideo.eval.pool import VideoPool + + n_workers = len(self._workers) + max_size = self._prefetch_factor * n_workers + results: 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 + results[idx] = self._workers[0].evaluate(**decoded) + else: + # Multi-GPU: every worker drains the shared pool (work-stealing). + threads: list[threading.Thread] = [] + for w in self._workers: + t = threading.Thread( + target=self._consumer_loop, + args=(w, pool, results), + daemon=True, + ) + t.start() + threads.append(t) + for t in threads: + t.join() + + return results + + @staticmethod + def _consumer_loop(worker: EvalWorker, pool: Any, results: list) -> None: + while True: + item = pool.get() + if item is None: + return + idx, decoded = item + results[idx] = worker.evaluate(**decoded) def release_cuda_memory(self) -> None: """Free CUDA caches on every replica without dropping models.""" @@ -123,10 +183,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/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..b8ea52934d 100644 --- a/fastvideo/eval/metrics/physics_iq/utils.py +++ b/fastvideo/eval/metrics/physics_iq/utils.py @@ -86,63 +86,30 @@ 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, - ) - ] + return prepare_pair_inputs( + sample["video"], + sample["reference"], + generated_mask=sample.get("video_mask"), + reference_mask=sample.get("reference_mask"), + **(prep_kwargs or {}), + ) 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..1ceec8d56e --- /dev/null +++ b/fastvideo/eval/pool.py @@ -0,0 +1,127 @@ +"""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 + +from fastvideo.eval.types import Video + +_SENTINEL = object() + + +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. + 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 + 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 + decoded = self._decode(sample) + try: + self._ready_q.put((idx, decoded), timeout=10.0) + except queue.Full: + return + + def _decode(self, sample: dict) -> dict: + """Materialize any path-shaped video values in *sample*. + + Recognises ``Video`` instances (populates ``.frames``) and bare + path strings under ``video`` / ``reference``. Other entries pass + 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") and isinstance(val, str | Path): + out[key] = load_video(str(val)) + return out diff --git a/fastvideo/eval/types.py b/fastvideo/eval/types.py index b25f7e2c03..a776115d4e 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,27 @@ class MetricResult: name: str score: float | None details: dict[str, Any] = field(default_factory=dict) + + +@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..208cf225cc 100644 --- a/fastvideo/eval/worker.py +++ b/fastvideo/eval/worker.py @@ -13,6 +13,7 @@ """ from __future__ import annotations +import contextlib from pathlib import Path from typing import Any @@ -21,16 +22,21 @@ from fastvideo.eval.memory import clear_cache 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. + + Metrics receive one sample per ``compute(sample)`` call: scalar + values, not list-wrapped or batch-dim-prefixed. See + :class:`Evaluator` for ``pre_upload`` semantics. + """ - def __init__(self, metric_names: list[str], device: str, *, compile: bool = False) -> None: + 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._pre_upload = pre_upload self._metrics: dict = {} self._unloaded = False self._load() @@ -57,13 +63,10 @@ def _load(self) -> None: 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. - - A ``(1, T, C, H, W)`` tensor is also accepted for back-compat - and gets unwrapped to ``(T, C, H, W)`` before reaching metrics. + ``video`` / ``reference`` may be a ``(T, C, H, W)`` tensor or a + path-like (``str`` / ``Path``). Paths are decoded here so the + dispatcher can queue cheap strings. A ``(1, T, C, H, W)`` tensor + is unwrapped to ``(T, C, H, W)`` for back-compat. """ if self._unloaded: raise RuntimeError("EvalWorker was unloaded; call reload() before evaluating.") @@ -72,6 +75,10 @@ def evaluate(self, **kwargs) -> dict[str, MetricResult]: 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(): @@ -97,16 +104,21 @@ def reload(self) -> None: self._load() +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) + + 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. + Paths → decoded ``(T, C, H, W)`` tensor; ``(1, T, C, H, W)`` → + squeezed to ``(T, C, H, W)``; everything else returned untouched. """ if value is None: return None From eab418d4fc0fb09e3874f479867c0ccacab4d9de Mon Sep 17 00:00:00 2001 From: shaoxiongduan Date: Mon, 11 May 2026 21:16:30 +0000 Subject: [PATCH 02/10] [bugfix] eval: cache prepare_pair result; unbounded pool put --- fastvideo/eval/metrics/physics_iq/utils.py | 4 +++- fastvideo/eval/pool.py | 8 ++++---- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/fastvideo/eval/metrics/physics_iq/utils.py b/fastvideo/eval/metrics/physics_iq/utils.py index b8ea52934d..453ef3b535 100644 --- a/fastvideo/eval/metrics/physics_iq/utils.py +++ b/fastvideo/eval/metrics/physics_iq/utils.py @@ -103,13 +103,15 @@ def prepare_pair( if "reference" not in sample: raise KeyError("Physics-IQ pair metrics require sample['reference'].") - return prepare_pair_inputs( + 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/pool.py b/fastvideo/eval/pool.py index 1ceec8d56e..23747a7d72 100644 --- a/fastvideo/eval/pool.py +++ b/fastvideo/eval/pool.py @@ -102,10 +102,10 @@ def _loader_loop(self) -> None: return idx, sample = item decoded = self._decode(sample) - try: - self._ready_q.put((idx, decoded), timeout=10.0) - except queue.Full: - return + # 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 any path-shaped video values in *sample*. From 5cc88e1e994de2249664613339265289c6f29ca8 Mon Sep 17 00:00:00 2001 From: shaoxiongduan Date: Mon, 11 May 2026 23:52:51 +0000 Subject: [PATCH 03/10] [feat] eval: set-vs-set metric protocol; unify single/list dispatch through the pool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds is_set_metric/reset/accumulate/finalize/merge_from to BaseMetric so corpus-level metrics (FAD, IS, KL on distributions, …) ride the same Evaluator pipeline as per-sample metrics. The pool delivers each sample once; per-sample metrics call compute(), set metrics call accumulate(); finalize() runs once per set metric after the pool drains and after worker-local accumulators are folded into worker 0. Collapses the kwargs (single-sample) and samples=[...] (list) paths onto one Evaluator._run pipeline. The kwargs case wraps as [kwargs], runs through the pool, and unwraps the single dict on return — so the public return shape is unchanged for both forms. Removes the worker's _resolve_video_input duplicate of the pool's _decode (Video instances, path strings, and (1,T,C,H,W) tensors are now handled in one place). VideoPool now forwards loader exceptions to the consumer thread instead of hanging — fixes a real bug surfaced by test_missing_path... once the single-sample form started using the pool. samples=[...] form now returns EvalResults (a list subclass) carrying corpus-level scores under .corpus; per-sample iteration is unchanged. All eval tests pass (31/31). --- fastvideo/eval/__init__.py | 3 +- fastvideo/eval/evaluator.py | 88 +++++++++++++++++++++------------- fastvideo/eval/metrics/base.py | 54 +++++++++++++++------ fastvideo/eval/pool.py | 48 +++++++++++++++---- fastvideo/eval/types.py | 19 ++++++++ fastvideo/eval/worker.py | 56 ++++++++++------------ 6 files changed, 180 insertions(+), 88 deletions(-) diff --git a/fastvideo/eval/__init__.py b/fastvideo/eval/__init__.py index 0fe56ed5a8..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, Video # 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,6 +34,7 @@ def _redirect_third_party_caches() -> None: "evaluate", "Evaluator", "create_evaluator", + "EvalResults", "MetricResult", "Video", "BaseMetric", diff --git a/fastvideo/eval/evaluator.py b/fastvideo/eval/evaluator.py index 3c1adba086..3fccf8745d 100644 --- a/fastvideo/eval/evaluator.py +++ b/fastvideo/eval/evaluator.py @@ -20,7 +20,7 @@ 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 @@ -91,20 +91,20 @@ 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 in the list form are decoded asynchronously by a - :class:`VideoPool` that runs alongside metric compute, hiding - decode latency behind GPU work. + 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) + Returns a ``dict[str, MetricResult]``. + Many samples — pipelined decode + work-stealing across replicas:: ev.evaluate(samples=[ @@ -113,26 +113,34 @@ def evaluate( ... ]) - Multi-GPU dispatch fires automatically when ``num_gpus > 1`` and - the list form is used: every worker runs a consumer thread, - pulling decoded samples from the shared pool as it frees up. - The kwargs form always runs on worker 0 with no pool overhead. + 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) - samples = list(samples) - if not samples: - return [] - return self._evaluate_with_pool(samples) + if single: + return per_sample[0] + return EvalResults(samples=per_sample, corpus=corpus) - def _evaluate_with_pool(self, samples: list[dict]) -> list[dict[str, MetricResult]]: - """Run samples through a :class:`VideoPool`, writing results in input order.""" + def _run(self, samples: list[dict]) -> tuple[list[dict[str, MetricResult]], dict[str, MetricResult]]: + """Pool-driven sample pipeline + set-metric finalize. + + 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 - results: list[Any] = [None] * len(samples) + 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: @@ -141,31 +149,43 @@ def _evaluate_with_pool(self, samples: list[dict]) -> list[dict[str, MetricResul if item is None: break idx, decoded = item - results[idx] = self._workers[0].evaluate(**decoded) + 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, results), - daemon=True, - ) + 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 results + return per_sample, corpus @staticmethod - def _consumer_loop(worker: EvalWorker, pool: Any, results: list) -> None: - while True: - item = pool.get() - if item is None: - return - idx, decoded = item - results[idx] = worker.evaluate(**decoded) + 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.""" diff --git a/fastvideo/eval/metrics/base.py b/fastvideo/eval/metrics/base.py index 14d79f8828..9cffc641c1 100644 --- a/fastvideo/eval/metrics/base.py +++ b/fastvideo/eval/metrics/base.py @@ -1,6 +1,6 @@ from __future__ import annotations -from abc import ABC, abstractmethod +from abc import ABC import torch @@ -10,14 +10,24 @@ class BaseMetric(ABC): """Abstract base class for all eval metrics. - Subclasses must implement :meth:`compute`. Optionally override - :meth:`setup` to eagerly load models. + Two execution shapes: - 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. + * **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. + + 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 +36,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 +67,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/pool.py b/fastvideo/eval/pool.py index 23747a7d72..cb5dd795cc 100644 --- a/fastvideo/eval/pool.py +++ b/fastvideo/eval/pool.py @@ -16,11 +16,27 @@ 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. @@ -82,6 +98,8 @@ 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: @@ -93,6 +111,9 @@ def get(self, timeout: float | None = None) -> tuple[int, dict] | None: 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: @@ -101,18 +122,26 @@ def _loader_loop(self) -> None: if item is _SENTINEL: return idx, sample = item - decoded = self._decode(sample) + 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 any path-shaped video values in *sample*. - - Recognises ``Video`` instances (populates ``.frames``) and bare - path strings under ``video`` / ``reference``. Other entries pass - through unchanged. + """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 @@ -122,6 +151,9 @@ def _decode(self, sample: dict) -> dict: 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") and isinstance(val, str | Path): - out[key] = load_video(str(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 a776115d4e..63d89be3b8 100644 --- a/fastvideo/eval/types.py +++ b/fastvideo/eval/types.py @@ -17,6 +17,25 @@ class MetricResult: 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 diff --git a/fastvideo/eval/worker.py b/fastvideo/eval/worker.py index 208cf225cc..7d512d0fed 100644 --- a/fastvideo/eval/worker.py +++ b/fastvideo/eval/worker.py @@ -14,12 +14,12 @@ from __future__ import annotations import contextlib -from pathlib import Path 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 @@ -27,9 +27,11 @@ class EvalWorker: """Owns metric replicas on one device. Single-GPU, single-sample. - Metrics receive one sample per ``compute(sample)`` call: scalar - values, not list-wrapped or batch-dim-prefixed. See - :class:`Evaluator` for ``pre_upload`` semantics. + 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: @@ -37,7 +39,7 @@ def __init__(self, metric_names: list[str], device: str, *, compile: bool = Fals self._device = device self._compile = compile self._pre_upload = pre_upload - self._metrics: dict = {} + self._metrics: dict[str, BaseMetric] = {} self._unloaded = False self._load() @@ -61,20 +63,17 @@ def _load(self) -> None: self._unloaded = False def evaluate(self, **kwargs) -> dict[str, MetricResult]: - """Score one sample. + """Score one already-decoded sample. - ``video`` / ``reference`` may be a ``(T, C, H, W)`` tensor or a - path-like (``str`` / ``Path``). Paths are decoded here so the - dispatcher can queue cheap strings. A ``(1, T, C, H, W)`` tensor - is unwrapped to ``(T, C, H, W)`` for back-compat. + 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: @@ -82,9 +81,22 @@ def evaluate(self, **kwargs) -> dict[str, MetricResult]: 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() @@ -112,19 +124,3 @@ def _to_device(value: Any, device: str | torch.device) -> Any: if value.device == target: return value return value.to(target, non_blocking=True) - - -def _resolve_video_input(value: Any) -> Any: - """Normalize a sample's ``video`` / ``reference`` field for metrics. - - Paths → decoded ``(T, C, H, W)`` tensor; ``(1, T, C, H, W)`` → - squeezed to ``(T, C, H, W)``; everything 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 From b9e4f5fab5400c64861351c6bc1972b8deeedd41 Mon Sep 17 00:00:00 2001 From: shaoxiongduan Date: Mon, 11 May 2026 23:58:29 +0000 Subject: [PATCH 04/10] [lint] eval: drop ABC inheritance from BaseMetric (B024) --- fastvideo/eval/metrics/base.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/fastvideo/eval/metrics/base.py b/fastvideo/eval/metrics/base.py index 9cffc641c1..984acb8ffc 100644 --- a/fastvideo/eval/metrics/base.py +++ b/fastvideo/eval/metrics/base.py @@ -1,13 +1,11 @@ from __future__ import annotations -from abc import ABC - import torch from fastvideo.eval.types import MetricResult -class BaseMetric(ABC): +class BaseMetric: """Abstract base class for all eval metrics. Two execution shapes: From 874ffd2e37afe95930323b9612c6ab4e652bd8b2 Mon Sep 17 00:00:00 2001 From: abaghyangor Date: Tue, 12 May 2026 01:19:03 -0700 Subject: [PATCH 05/10] =?UTF-8?q?[feat]:=20Add=20common.fvd=20=E2=80=94=20?= =?UTF-8?q?Fr=C3=A9chet=20Video=20Distance=20metric?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements FVD using I3D features (Kinetics-400) as a set-vs-set metric following the accumulate/finalize protocol introduced in shao/eval-pool. - accumulate(): extracts I3D features per video, buffers them, builds real feature cache from sample["reference"] on first encounter - finalize(): computes Fréchet distance between generated and real feature distributions - merge_from(): folds multi-GPU worker state for parallel evaluation - reset(): clears generated feature buffer between runs I3D model downloaded automatically from HuggingFace (flateon/FVD-I3D-torchscript). Reference features cached to ~/.cache/fastvideo/eval/fvd/real_features.pt after first extraction. Warns when fewer than 256 videos are used (standard protocol: 2048). --- fastvideo/eval/metrics/common/fvd/__init__.py | 0 fastvideo/eval/metrics/common/fvd/metric.py | 326 ++++++++++++++++++ 2 files changed, 326 insertions(+) create mode 100644 fastvideo/eval/metrics/common/fvd/__init__.py create mode 100644 fastvideo/eval/metrics/common/fvd/metric.py 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..becaef5201 --- /dev/null +++ b/fastvideo/eval/metrics/common/fvd/metric.py @@ -0,0 +1,326 @@ +"""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.""" + parts = [] + 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) + 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).""" + mu = features.mean(axis=0) + sigma = np.cov(features, rowvar=False) + 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) + return data.numpy() + 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) From 71a2b3dc5c1586cb5999116807777dbf6c638d5e Mon Sep 17 00:00:00 2001 From: abaghyangor Date: Tue, 12 May 2026 01:52:34 -0700 Subject: [PATCH 06/10] [bugfix] disable NVRTC fusion in FVD I3D forward pass for CUDA compat --- fastvideo/eval/metrics/common/fvd/metric.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/fastvideo/eval/metrics/common/fvd/metric.py b/fastvideo/eval/metrics/common/fvd/metric.py index becaef5201..0c2d3c2646 100644 --- a/fastvideo/eval/metrics/common/fvd/metric.py +++ b/fastvideo/eval/metrics/common/fvd/metric.py @@ -102,12 +102,21 @@ def _extract_features( chunk: int, device: torch.device, ) -> np.ndarray: - """Extract I3D features → (B, 400) numpy array, chunked to fit VRAM.""" + """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 = [] - 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) - parts.append(feats.cpu().numpy()) + 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) + parts.append(feats.cpu().numpy()) return np.concatenate(parts, axis=0) From 2de3bad84aea07a7d99e6fdad5237114645f6117 Mon Sep 17 00:00:00 2001 From: abaghyangor Date: Tue, 12 May 2026 01:59:58 -0700 Subject: [PATCH 07/10] [bugfix]: fix I3D feature shape in FVD metric --- fastvideo/eval/metrics/common/fvd/metric.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/fastvideo/eval/metrics/common/fvd/metric.py b/fastvideo/eval/metrics/common/fvd/metric.py index 0c2d3c2646..d28b69b8da 100644 --- a/fastvideo/eval/metrics/common/fvd/metric.py +++ b/fastvideo/eval/metrics/common/fvd/metric.py @@ -116,6 +116,8 @@ def _extract_features( 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) From 9adc6909c6a7a99f14b144f45d3d475ca184eaba Mon Sep 17 00:00:00 2001 From: abaghyangor Date: Tue, 12 May 2026 02:12:01 -0700 Subject: [PATCH 08/10] [bugfix] guard against 1-D features in gaussian_params and load_cache --- fastvideo/eval/metrics/common/fvd/metric.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/fastvideo/eval/metrics/common/fvd/metric.py b/fastvideo/eval/metrics/common/fvd/metric.py index d28b69b8da..a896d12de6 100644 --- a/fastvideo/eval/metrics/common/fvd/metric.py +++ b/fastvideo/eval/metrics/common/fvd/metric.py @@ -127,9 +127,17 @@ def _extract_features( # --------------------------------------------------------------------------- def _gaussian_params(features: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - """Compute mean and covariance of a feature matrix (N, D).""" + """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 @@ -329,7 +337,8 @@ def merge_from(self, other: BaseMetric) -> None: 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) - return data.numpy() + arr = np.atleast_2d(data.numpy()) # guard against stale 1-D cache + return arr return None def _save_cache(self, features: np.ndarray) -> None: From 8fa0688a59fb50f20a1d95f4a6a3478b09ea4880 Mon Sep 17 00:00:00 2001 From: abaghyangor Date: Tue, 12 May 2026 02:29:27 -0700 Subject: [PATCH 09/10] [misc] pre-commit fixes for common.fvd --- fastvideo/eval/metrics/common/fvd/metric.py | 70 ++++++++++----------- 1 file changed, 33 insertions(+), 37 deletions(-) diff --git a/fastvideo/eval/metrics/common/fvd/metric.py b/fastvideo/eval/metrics/common/fvd/metric.py index a896d12de6..a4ac5603dd 100644 --- a/fastvideo/eval/metrics/common/fvd/metric.py +++ b/fastvideo/eval/metrics/common/fvd/metric.py @@ -51,17 +51,17 @@ # 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 +_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. @@ -82,10 +82,8 @@ def _preprocess(video: torch.Tensor) -> torch.Tensor: """ 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." - ) + 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) @@ -98,7 +96,7 @@ def _preprocess(video: torch.Tensor) -> torch.Tensor: @torch.no_grad() def _extract_features( model: torch.nn.Module, - video: torch.Tensor, # (B, T, C, H, W) float [0, 1] + video: torch.Tensor, # (B, T, C, H, W) float [0, 1] chunk: int, device: torch.device, ) -> np.ndarray: @@ -114,7 +112,7 @@ def _extract_features( parts = [] with torch.jit.fuser("none"): for i in range(0, video.shape[0], chunk): - batch = _preprocess(video[i : i + chunk].to(device)) + 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 @@ -126,6 +124,7 @@ def _extract_features( # 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). @@ -134,9 +133,9 @@ def _gaussian_params(features: np.ndarray) -> tuple[np.ndarray, np.ndarray]: cause ``np.cov`` to return a 0-d scalar and break ``sigma.shape[0]``. """ features = np.atleast_2d(features) - mu = features.mean(axis=0) + mu = features.mean(axis=0) sigma = np.cov(features, rowvar=False) - if sigma.ndim == 0: # n==1 edge case: variance scalar + if sigma.ndim == 0: # n==1 edge case: variance scalar sigma = sigma.reshape(1, 1) return mu, sigma @@ -152,7 +151,7 @@ def _frechet_distance( sigma1 = sigma1 + eps * np.eye(sigma1.shape[0]) sigma2 = sigma2 + eps * np.eye(sigma2.shape[0]) - diff = mu1 - mu2 + diff = mu1 - mu2 covmean = scipy.linalg.sqrtm(sigma1 @ sigma2) if np.iscomplexobj(covmean): @@ -164,13 +163,14 @@ def _frechet_distance( ) covmean = covmean.real - return float(np.sum(diff ** 2) + np.trace(sigma1 + sigma2 - 2.0 * covmean)) + 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. @@ -194,12 +194,12 @@ class FVDMetric(BaseMetric): 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"] + 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, @@ -210,11 +210,11 @@ def __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._chunk = chunk_size self._i3d: torch.nn.Module | None = None # Accumulated feature buffers — cleared by reset() - self._gen_features: list[np.ndarray] = [] + self._gen_features: list[np.ndarray] = [] self._real_features: np.ndarray | None = None def to(self, device): @@ -253,9 +253,9 @@ def accumulate(self, sample: dict) -> None: if self._i3d is None: self.setup() - video = sample["video"] # (T, C, H, W) from evaluator + video = sample["video"] # (T, C, H, W) from evaluator if video.dim() == 4: - video = video.unsqueeze(0) # → (1, T, C, H, W) + video = video.unsqueeze(0) # → (1, T, C, H, W) # Buffer generated features feats = _extract_features(self._i3d, video, self._chunk, self.device) @@ -270,9 +270,7 @@ def accumulate(self, sample: dict) -> None: 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._real_features = _extract_features(self._i3d, ref, self._chunk, self.device) self._save_cache(self._real_features) def finalize(self) -> MetricResult: @@ -289,17 +287,15 @@ def finalize(self) -> 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}" - ) + "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) + 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( @@ -309,7 +305,7 @@ def finalize(self) -> MetricResult: stacklevel=2, ) - mu_gen, sigma_gen = _gaussian_params(all_gen) + 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) From ab369e34b1ef50e525fe43c24e9959d07b5ca1b2 Mon Sep 17 00:00:00 2001 From: abaghyangor Date: Tue, 12 May 2026 02:33:25 -0700 Subject: [PATCH 10/10] [docs] update eval README to reflect common.fvd addition --- fastvideo/eval/README.md | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) 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.