metric

@register("common.fvd")
class FVDMetric(BaseMetric):
    """Fréchet Video Distance (FVD) over a pluggable feature backbone.
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
----------
extractor : str
    Feature backbone name. One of ``"i3d"`` (default), ``"clip"``,
    ``"videomae"``. See module docstring for guidance.
cache_path : str, optional
    Where extracted reference features are cached. Resolution order:
    constructor kwarg → ``$FASTVIDEO_FVD_REF_FEATURES`` env-var →
    ``${FASTVIDEO_EVAL_CACHE}/fvd/real_features_{extractor}.pt``
    (default, resolved at :meth:`setup` time so the env-var can be set
    after import).
cache_mode : str
    ``"read_write"`` (default) — read at setup, write at finalize on
    cache miss. ``"read"`` — read at setup, never write.  ``"off"`` —
    ignore the cache entirely.  Auto-write happens only when the cache
    was missing AND reference features streamed in this run.
chunk_size : int
    Videos per forward pass. Reduce if GPU runs OOM.
"""

name = "common.fvd"
is_set_metric = True
requires_reference = False  # reference is supplied via role="reference" samples or cache
higher_is_better = False  # lower FVD = better
needs_gpu = True
dependencies = ["huggingface_hub", "scipy", "transformers"]

def __init__(
    self,
    extractor: str = "i3d",
    cache_path: str | None = None,
    cache_mode: str = "read_write",
    chunk_size: int = 32,
) -> None:
    super().__init__()
    if extractor not in available_extractors():
        raise ValueError(f"Unknown FVD extractor '{extractor}'. "
                         f"Available: {available_extractors()}")
    if cache_mode not in {"off", "read", "read_write"}:
        raise ValueError(f"cache_mode must be one of 'off', 'read', 'read_write'; got {cache_mode!r}")
    self._extractor_name = extractor
    self._cache_path_arg = cache_path
    self._cache_mode = cache_mode
    self._chunk = chunk_size
    self._extractor: _BaseExtractor | None = None

    # Accumulated feature buffers — cleared by reset()
    self._gen_buf: list[np.ndarray] = []
    self._real_buf: list[np.ndarray] = []
    # Cache loaded at setup() time; survives reset() (the cached corpus
    # is part of the metric's configured state, not its run state).
    self._cached_real: np.ndarray | None = None

def to(self, device: str | torch.device) -> FVDMetric:
    super().to(device)
    if self._extractor is not None:
        self._extractor.to(self.device)
    return self

def setup(self) -> None:
    if self._extractor is not None:
        return
    # Resolve cache_path at runtime so $FASTVIDEO_EVAL_CACHE /
    # $FASTVIDEO_FVD_REF_FEATURES are both honoured even when set after
    # import. Precedence: kwarg > env-var > default.
    if self._cache_path_arg is not None:
        self.cache_path = os.path.expanduser(self._cache_path_arg)
    elif env_path := envs.FASTVIDEO_FVD_REF_FEATURES.get():
        self.cache_path = os.path.expanduser(env_path)
    else:
        self.cache_path = _default_cache_path(self._extractor_name)
    self._extractor = load_extractor(self._extractor_name, self.device)
    if self._cache_mode != "off":
        self._cached_real = self._load_cache()

# ------------------------------------------------------------------
# Set-vs-set protocol
# ------------------------------------------------------------------

def reset(self) -> None:
    """Clear generated and streamed-reference buffers. Cache survives."""
    self._gen_buf = []
    self._real_buf = []

def accumulate(self, sample: dict) -> None:
    """Extract features from one video and route by ``sample["role"]``.

    Routing rules:

    * ``role="reference"`` → real buffer; the reference key is ignored.
    * Otherwise → generated buffer.  If ``sample["reference"]`` is
      also present, its features also go into the real buffer — this
      lets a paired samples list (the shape per-sample metrics like
      LPIPS need) feed FVD in the same pass without a second decode.

    Accepts either a raw tensor or a populated :class:`Video` instance
    under either key.
    """
    if self._extractor is None:
        self.setup()
    assert self._extractor is not None  # for type narrowing

    if sample.get("role") == "reference":
        self._real_buf.append(self._extract(sample["video"]))
        return

    self._gen_buf.append(self._extract(sample["video"]))
    if "reference" in sample:
        self._real_buf.append(self._extract(sample["reference"]))

def _extract(self, video: torch.Tensor) -> np.ndarray:
    """Batchify (4-D → 5-D) and run the chunked extractor forward.

    Pre-condition: ``video`` is already a tensor — :class:`EvalWorker`
    unwraps :class:`Video` instances before any metric sees the sample.
    """
    if video.dim() == 4:
        video = video.unsqueeze(0)  # (T, C, H, W) → (1, T, C, H, W)
    assert self._extractor is not None
    return _extract_chunked(self._extractor, video, self._chunk)

def finalize(self) -> MetricResult:
    """Compute FVD from buffered generated features vs. reference features.

    Reference source priority: streamed (``_real_buf``) > disk cache
    (``_cached_real``).  On a fresh cache + streamed refs, write the
    accumulated real features back to disk if ``cache_mode="read_write"``.
    """
    if not self._gen_buf:
        return MetricResult(
            name=self.name,
            score=None,
            details={"skipped": "No generated videos accumulated before finalize()."},
        )
    all_gen = np.concatenate(self._gen_buf, axis=0)

    if self._real_buf:
        all_real = np.concatenate(self._real_buf, axis=0)
        ref_source = "streamed"
        # Persist to cache only when we built real features fresh and
        # the cache slot was empty — never silently overwrite an
        # existing cache.
        if self._cache_mode == "read_write" and self._cached_real is None:
            self._save_cache(all_real)
    elif self._cached_real is not None:
        all_real = self._cached_real
        ref_source = "cached"
    else:
        return MetricResult(
            name=self.name,
            score=None,
            details={
                "skipped": ("No reference features available. Pass role='reference' samples "
                            "(or a paired samples list with 'reference' key) to "
                            f"Evaluator.evaluate(), or pre-build the cache at: {self.cache_path}")
            },
        )

    n_gen = len(all_gen)
    n_real = len(all_real)

    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(all_real)
    fvd = _frechet_distance(mu_gen, sigma_gen, mu_real, sigma_real)

    return MetricResult(
        name=self.name,
        score=fvd,
        details={
            "extractor": self._extractor_name,
            "n_generated": n_gen,
            "n_reference": n_real,
            "ref_source": ref_source,
        },
    )

def merge_from(self, other: BaseMetric) -> None:
    """Fold another worker's accumulated features into this one (multi-GPU).

    Both gen and ref buffers concatenate.  The disk cache is the same
    across workers, so we only adopt *other*'s ``_cached_real`` if our
    own slot is empty (defensive — should always already match).
    """
    assert isinstance(other, FVDMetric)
    assert other._extractor_name == self._extractor_name, (
        f"merge_from extractor mismatch: {other._extractor_name!r} vs {self._extractor_name!r}")
    self._gen_buf.extend(other._gen_buf)
    self._real_buf.extend(other._real_buf)
    if self._cached_real is None and other._cached_real is not None:
        self._cached_real = other._cached_real

# ------------------------------------------------------------------
# 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 np.atleast_2d(data.numpy())  # guard against stale 1-D cache
    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)</code></pre></details>
fastvideo.eval.metrics.common.fvd.metric.FVDMetric.to
to(device: str | torch.device) -> FVDMetric
Source
def to(self, device: str | torch.device) -> FVDMetric:
    super().to(device)
    if self._extractor is not None:
        self._extractor.to(self.device)
    return self
fastvideo.eval.metrics.common.fvd.metric.FVDMetric.setup
setup() -> None
Source
def setup(self) -> None:
    if self._extractor is not None:
        return
    # Resolve cache_path at runtime so $FASTVIDEO_EVAL_CACHE /
    # $FASTVIDEO_FVD_REF_FEATURES are both honoured even when set after
    # import. Precedence: kwarg > env-var > default.
    if self._cache_path_arg is not None:
        self.cache_path = os.path.expanduser(self._cache_path_arg)
    elif env_path := envs.FASTVIDEO_FVD_REF_FEATURES.get():
        self.cache_path = os.path.expanduser(env_path)
    else:
        self.cache_path = _default_cache_path(self._extractor_name)
    self._extractor = load_extractor(self._extractor_name, self.device)
    if self._cache_mode != "off":
        self._cached_real = self._load_cache()
fastvideo.eval.metrics.common.fvd.metric.FVDMetric.reset
reset() -> None
Clear generated and streamed-reference buffers. Cache survives.
Source
def reset(self) -> None:
    """Clear generated and streamed-reference buffers. Cache survives."""
    self._gen_buf = []
    self._real_buf = []
fastvideo.eval.metrics.common.fvd.metric.FVDMetric.accumulate
accumulate(sample: dict) -> None
Extract features from one video and route by sample["role"].
Routing rules:

role="reference" → real buffer; the reference key is ignored.
Otherwise → generated buffer.  If sample["reference"] is
  also present, its features also go into the real buffer — this
  lets a paired samples list (the shape per-sample metrics like
  LPIPS need) feed FVD in the same pass without a second decode.

Accepts either a raw tensor or a populated :class:Video instance
under either key.
Source
def accumulate(self, sample: dict) -> None:
    """Extract features from one video and route by ``sample["role"]``.
Routing rules:

* ``role=&quot;reference&quot;`` → real buffer; the reference key is ignored.
* Otherwise → generated buffer.  If ``sample&#91;&quot;reference&quot;&#93;`` is
  also present, its features also go into the real buffer — this
  lets a paired samples list (the shape per-sample metrics like
  LPIPS need) feed FVD in the same pass without a second decode.

Accepts either a raw tensor or a populated :class:`Video` instance
under either key.
&quot;&quot;&quot;
if self._extractor is None:
    self.setup()
assert self._extractor is not None  # for type narrowing

if sample.get(&quot;role&quot;) == &quot;reference&quot;:
    self._real_buf.append(self._extract(sample&#91;&quot;video&quot;&#93;))
    return

self._gen_buf.append(self._extract(sample&#91;&quot;video&quot;&#93;))
if &quot;reference&quot; in sample:
    self._real_buf.append(self._extract(sample&#91;&quot;reference&quot;&#93;))</code></pre></details>
fastvideo.eval.metrics.common.fvd.metric.FVDMetric.finalize
finalize() -> MetricResult
Compute FVD from buffered generated features vs. reference features.
Reference source priority: streamed (_real_buf) > disk cache
(_cached_real).  On a fresh cache + streamed refs, write the
accumulated real features back to disk if cache_mode="read_write".
Source
def finalize(self) -> MetricResult:
    """Compute FVD from buffered generated features vs. reference features.
Reference source priority: streamed (``_real_buf``) &gt; disk cache
(``_cached_real``).  On a fresh cache + streamed refs, write the
accumulated real features back to disk if ``cache_mode=&quot;read_write&quot;``.
&quot;&quot;&quot;
if not self._gen_buf:
    return MetricResult(
        name=self.name,
        score=None,
        details=&#123;&quot;skipped&quot;: &quot;No generated videos accumulated before finalize().&quot;&#125;,
    )
all_gen = np.concatenate(self._gen_buf, axis=0)

if self._real_buf:
    all_real = np.concatenate(self._real_buf, axis=0)
    ref_source = &quot;streamed&quot;
    # Persist to cache only when we built real features fresh and
    # the cache slot was empty — never silently overwrite an
    # existing cache.
    if self._cache_mode == &quot;read_write&quot; and self._cached_real is None:
        self._save_cache(all_real)
elif self._cached_real is not None:
    all_real = self._cached_real
    ref_source = &quot;cached&quot;
else:
    return MetricResult(
        name=self.name,
        score=None,
        details=&#123;
            &quot;skipped&quot;: (&quot;No reference features available. Pass role=&#x27;reference&#x27; samples &quot;
                        &quot;(or a paired samples list with &#x27;reference&#x27; key) to &quot;
                        f&quot;Evaluator.evaluate(), or pre-build the cache at: &#123;self.cache_path&#125;&quot;)
        &#125;,
    )

n_gen = len(all_gen)
n_real = len(all_real)

if n_gen &lt; _MIN_VIDEOS_WARN or n_real &lt; _MIN_VIDEOS_WARN:
    warnings.warn(
        f&quot;FVD computed with only &#123;n_gen&#125; generated and &#123;n_real&#125; real videos. &quot;
        f&quot;At least &#123;_MIN_VIDEOS_WARN&#125; recommended (standard protocol: 2048). &quot;
        &quot;Score may not be statistically reliable.&quot;,
        stacklevel=2,
    )

mu_gen, sigma_gen = _gaussian_params(all_gen)
mu_real, sigma_real = _gaussian_params(all_real)
fvd = _frechet_distance(mu_gen, sigma_gen, mu_real, sigma_real)

return MetricResult(
    name=self.name,
    score=fvd,
    details=&#123;
        &quot;extractor&quot;: self._extractor_name,
        &quot;n_generated&quot;: n_gen,
        &quot;n_reference&quot;: n_real,
        &quot;ref_source&quot;: ref_source,
    &#125;,
)</code></pre></details>
fastvideo.eval.metrics.common.fvd.metric.FVDMetric.merge_from
merge_from(other: BaseMetric) -> None
Fold another worker's accumulated features into this one (multi-GPU).
Both gen and ref buffers concatenate.  The disk cache is the same
across workers, so we only adopt other's _cached_real if our
own slot is empty (defensive — should always already match).
Source
def merge_from(self, other: BaseMetric) -> None:
    """Fold another worker's accumulated features into this one (multi-GPU).
Both gen and ref buffers concatenate.  The disk cache is the same
across workers, so we only adopt *other*&#x27;s ``_cached_real`` if our
own slot is empty (defensive — should always already match).
&quot;&quot;&quot;
assert isinstance(other, FVDMetric)
assert other._extractor_name == self._extractor_name, (
    f&quot;merge_from extractor mismatch: &#123;other._extractor_name!r&#125; vs &#123;self._extractor_name!r&#125;&quot;)
self._gen_buf.extend(other._gen_buf)
self._real_buf.extend(other._real_buf)
if self._cached_real is None and other._cached_real is not None:
    self._cached_real = other._cached_real</code></pre></details>