diffeval 0.1.0__py3-none-any.whl

diffeval/bench/cmmd_check.py

@@ -0,0 +1,20 @@
+"""diffeval-check-cmmd: score the cmmd-pytorch samples and expect 7.696."""
+
+import tempfile
+from pathlib import Path
+
+from diffeval.bench.cmmd_samples import cmmd_samples
+
+EXPECTED = 7.696
+
+
+def main() -> int:
+    from diffeval.metrics.cmmd.embed_batch import load_vision
+    from diffeval.metrics.cmmd.embed_paths import embed_dir
+    from diffeval.metrics.cmmd.mmd import mmd
+    with tempfile.TemporaryDirectory() as tmp:
+        ref_dir, gen_dir = cmmd_samples(Path(tmp))
+        model, processor = load_vision("openai/clip-vit-large-patch14-336")
+        value = round(mmd(embed_dir(model, processor, ref_dir, "cpu"), embed_dir(model, processor, gen_dir, "cpu")), 3)
+    print(f"CMMD {value} (expected {EXPECTED})")
+    return 0 if abs(value - EXPECTED) < 0.001 else 1

diffeval/bench/cmmd_samples.py

@@ -0,0 +1,15 @@
+"""The six cmmd-pytorch sample PNGs (Apache-2.0) behind its 7.696 check."""
+
+from pathlib import Path
+
+from diffeval.bench.fetch import fetch
+
+BASE = "https://raw.githubusercontent.com/sayakpaul/cmmd-pytorch/main"
+FILES = {"reference_images": "glowing_512_1", "generated_images": "glowing_512_2"}
+
+
+def cmmd_samples(root: Path) -> tuple[Path, Path]:
+    for folder, stem in FILES.items():
+        for suffix in ("", "_copy", "_copy_2"):
+            fetch(f"{BASE}/{folder}/{stem}{suffix}.png", root / folder / f"{stem}{suffix}.png")
+    return root / "reference_images", root / "generated_images"

diffeval/bench/coco30k_cli.py

@@ -0,0 +1,23 @@
+"""diffeval-coco30k: write <out>/suite.yaml for the cmmd-pytorch COCO-30k protocol."""
+
+import argparse
+from pathlib import Path
+
+import yaml
+
+from diffeval.bench.coco30k_files import coco30k_captions
+from diffeval.bench.coco30k_suite import coco30k_suite
+
+
+def main(argv=None) -> int:
+    parser = argparse.ArgumentParser(prog="diffeval-coco30k")
+    parser.add_argument("--out", required=True, type=Path)
+    parser.add_argument("--limit", type=int, default=30000, help="first N captions (default all 30000)")
+    parser.add_argument("--seed", type=int, default=2024, help="cmmd-pytorch uses 2024")
+    parser.add_argument("--reference-images", type=Path, default=None, help="COCO image folder; adds fid and kid")
+    args = parser.parse_args(argv)
+    refs = None if args.reference_images is None else args.reference_images.resolve()
+    suite = coco30k_suite(coco30k_captions(args.out, args.limit), args.seed, refs)
+    (args.out / "suite.yaml").write_text(yaml.safe_dump(suite, sort_keys=False, allow_unicode=True), encoding="utf-8")
+    print(f"wrote {args.out / 'suite.yaml'} with {len(suite['prompts'])} prompts and metrics {suite['metrics']}")
+    return 0

diffeval/bench/coco30k_export.py

@@ -0,0 +1,19 @@
+"""Write the COCO-30k reference images from the dataset's parquet shards."""
+
+import io
+from pathlib import Path
+
+from PIL import Image
+
+
+def coco30k_export(out: Path) -> int:
+    import pyarrow.parquet as pq
+    from huggingface_hub import snapshot_download
+    out.mkdir(parents=True, exist_ok=True)
+    root = Path(snapshot_download("sayakpaul/coco-30-val-2014", repo_type="dataset", allow_patterns=["data/*.parquet"]))
+    count = 0
+    for shard in sorted((root / "data").glob("*.parquet")):
+        for row in pq.read_table(shard, columns=["image"]).column("image").to_pylist():
+            Image.open(io.BytesIO(row["bytes"])).convert("RGB").save(out / f"ref_{count:05d}.png")
+            count += 1
+    return count

diffeval/bench/coco30k_files.py

@@ -0,0 +1,16 @@
+"""The captions and reference embeddings cmmd-pytorch publishes for COCO-30k."""
+
+import csv
+from pathlib import Path
+
+from diffeval.bench.fetch import fetch
+
+CSV_URL = "https://huggingface.co/datasets/sayakpaul/sample-datasets/raw/main/coco_30k_randomly_sampled_2014_val.csv"
+EMB_URL = "https://huggingface.co/datasets/sayakpaul/coco-30-val-2014/resolve/main/ref_embs_coco_30k.npy"
+EMB_NAME = "ref_embs_coco_30k.npy"
+
+
+def coco30k_captions(out: Path, limit: int) -> list[str]:
+    fetch(EMB_URL, out / EMB_NAME)
+    with fetch(CSV_URL, out / "captions.csv").open(encoding="utf-8") as handle:
+        return [row["caption"].strip() for row in csv.DictReader(handle)][:limit]

diffeval/bench/coco30k_images.py

@@ -0,0 +1,21 @@
+"""diffeval-coco30k-images: write the 30k COCO reference PNGs for fid and kid (about 5 GB).
+
+Needs the benchmarks extra: pip install "diffeval[benchmarks]".
+"""
+
+import argparse
+import importlib.util
+from pathlib import Path
+
+from diffeval.bench.coco30k_export import coco30k_export
+
+
+def main(argv=None) -> int:
+    parser = argparse.ArgumentParser(prog="diffeval-coco30k-images")
+    parser.add_argument("--out", required=True, type=Path)
+    out = parser.parse_args(argv).out
+    if any(importlib.util.find_spec(name) is None for name in ("pyarrow", "huggingface_hub")):
+        print('install the benchmarks extra: pip install "diffeval[benchmarks]"')
+        return 2
+    print(f"wrote {coco30k_export(out)} images to {out}")
+    return 0

diffeval/bench/coco30k_suite.py

@@ -0,0 +1,16 @@
+"""The suite mapping for the first N COCO-30k captions."""
+
+from diffeval.bench.coco30k_files import EMB_NAME
+
+
+def coco30k_suite(captions: list[str], seed: int, reference_images=None) -> dict:
+    name = "coco30k" if len(captions) >= 30000 else f"coco30k-{len(captions)}"
+    suite = {"name": name, "description": f"cmmd-pytorch COCO-30k captions (COCO 2014 val), first {len(captions)}"}
+    suite.update(seeds=[seed], width=512, height=512, cmmd_ref_embed=EMB_NAME)
+    metrics = ["clip_score", "cmmd"]
+    if reference_images is not None:
+        suite["reference_dir"] = str(reference_images)
+        metrics += ["fid", "kid"]
+    suite["metrics"] = metrics
+    suite["prompts"] = [{"id": f"c{index:05d}", "prompt": text} for index, text in enumerate(captions)]
+    return suite

diffeval/bench/fetch.py

@@ -0,0 +1,12 @@
+"""Download a URL to a path once."""
+
+import urllib.request
+from pathlib import Path
+
+
+def fetch(url: str, path: Path) -> Path:
+    if not path.exists():
+        path.parent.mkdir(parents=True, exist_ok=True)
+        print(f"downloading {path.name}")
+        urllib.request.urlretrieve(url, path)
+    return path

diffeval/bench/published.py

@@ -0,0 +1,10 @@
+"""cmmd-pytorch's published COCO-30k CMMD by model id (github.com/sayakpaul/cmmd-pytorch)."""
+
+PUBLISHED_CMMD = {
+    "runwayml/stable-diffusion-v1-5": 0.582,
+    "stable-diffusion-v1-5/stable-diffusion-v1-5": 0.582,
+    "stabilityai/sdxl-turbo": 0.548,
+    "stabilityai/stable-diffusion-xl-base-1.0": 0.696,
+    "segmind/SSD-1B": 0.669,
+    "PixArt-alpha/PixArt-XL-2-1024-MS": 1.140,
+}

diffeval/bench/table_cli.py

@@ -0,0 +1,16 @@
+"""diffeval-table: print a Markdown table from results.json files."""
+
+import argparse
+
+from diffeval.bench.table_row import HEADER, table_row
+
+
+def main(argv=None) -> int:
+    parser = argparse.ArgumentParser(prog="diffeval-table")
+    parser.add_argument("results", nargs="+")
+    args = parser.parse_args(argv)
+    print("| " + " | ".join(HEADER) + " |")
+    print("|" + " --- |" * len(HEADER))
+    for path in args.results:
+        print(table_row(path))
+    return 0

diffeval/bench/table_row.py

@@ -0,0 +1,24 @@
+"""One Markdown row from one results.json."""
+
+import json
+from pathlib import Path
+
+from diffeval.bench.published import PUBLISHED_CMMD
+
+METRICS = ("clip_score", "cmmd", "fid", "kid")
+HEADER = ["model", "steps", "guidance", "size", "device", "images", *METRICS, "published CMMD (30k)"]
+
+
+def _cell(row) -> str:
+    if row is None or row.get("mean") is None:
+        return "-"
+    return f"{row['mean']:.4g}" if row.get("std") is None else f"{row['mean']:.4g} +/- {row['std']:.2g}"
+
+
+def table_row(path) -> str:
+    body = json.loads(Path(path).read_text(encoding="utf-8"))
+    backend, suite, agg = body["backend"], body["suite"], body["aggregate"]
+    model = backend.get("pipeline") or backend["name"]
+    cells = [model, str(backend.get("steps")), str(backend.get("guidance")), f"{suite['width']}x{suite['height']}", str(backend.get("device")), str(len(suite["prompts"]))]
+    published = PUBLISHED_CMMD.get(model)
+    return "| " + " | ".join(cells + [_cell(agg.get(name)) for name in METRICS] + ["-" if published is None else str(published)]) + " |"

diffeval/cli/__init__.py

@@ -0,0 +1 @@
+"""Command line interface."""

diffeval/cli/command_list.py

@@ -0,0 +1,13 @@
+"""List tasks, metrics, or backends."""
+
+import sys
+
+from diffeval.cli.print_registered import print_registered
+from diffeval.registry.tables import BACKENDS, METRICS, TASKS
+
+_TABLES = {"tasks": TASKS, "metrics": METRICS, "backends": BACKENDS}
+
+
+def command_list(args, stream=None) -> int:
+    print_registered(_TABLES[args.kind], stream or sys.stdout)
+    return 0

diffeval/cli/command_run.py

@@ -0,0 +1,25 @@
+"""Resolve a suite, run it, and print output paths."""
+
+import sys
+
+from diffeval.cli.config import backend_config, backend_name
+from diffeval.cli.output_dir import output_dir
+from diffeval.cli.print_paths import print_paths
+from diffeval.runner.options import RunOptions
+from diffeval.runner.run import run_evaluation
+from diffeval.tasks.resolve import resolve
+
+
+def command_run(args, stream=None) -> int:
+    out = stream or sys.stdout
+    try:
+        suite = resolve(args.suite)
+        name = backend_name(args)
+        folder = output_dir(args.output, suite.name, name)
+        options = RunOptions(suite, name, backend_config(args), folder, args.seeds, args.metrics, args.width, args.height)
+        body = run_evaluation(options)
+    except Exception as exc:
+        print(exc, file=out)
+        return 2
+    print_paths(folder, out, body)
+    return 2 if args.strict and any(row["status"] != "ok" for row in body["aggregate"].values()) else 0

diffeval/cli/config.py

@@ -0,0 +1,14 @@
+"""Backend name and config dict from CLI args."""
+
+
+def backend_name(args) -> str:
+    if args.backend:
+        return args.backend
+    if args.pipeline:
+        return "diffusers"
+    return "dummy"
+
+
+def backend_config(args) -> dict:
+    raw = {"pipeline": args.pipeline, "dtype": args.dtype, "steps": args.steps, "guidance": args.guidance, "device": args.device, "attention": getattr(args, "attention", None), "variant": getattr(args, "variant", None), "batch_size": getattr(args, "batch_size", None)}
+    return {key: value for key, value in raw.items() if value is not None}

diffeval/cli/csv_ints.py

@@ -0,0 +1,8 @@
+"""Parse a comma-separated list of seeds."""
+
+
+def csv_ints(text: str) -> tuple[int, ...]:
+    parts = [item.strip() for item in text.split(",") if item.strip()]
+    if not parts or any(not item.lstrip("-").isdigit() for item in parts):
+        raise ValueError(f"expected integers, got {text!r}")
+    return tuple(int(item) for item in parts)

diffeval/cli/csv_names.py

@@ -0,0 +1,8 @@
+"""Parse a comma-separated list of names."""
+
+
+def csv_names(text: str) -> tuple[str, ...]:
+    parts = tuple(item.strip() for item in text.split(",") if item.strip())
+    if not parts:
+        raise ValueError("expected at least one name")
+    return parts

diffeval/cli/main.py

@@ -0,0 +1,12 @@
+"""Dispatch diffeval commands and return an exit code."""
+
+from diffeval.cli.command_list import command_list
+from diffeval.cli.command_run import command_run
+from diffeval.cli.parser import build_parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    if args.command == "list":
+        return command_list(args)
+    return command_run(args)

diffeval/cli/output_dir.py

@@ -0,0 +1,12 @@
+"""Default directory outputs/<utc>-<suite>-<backend>."""
+
+from datetime import datetime, timezone
+from pathlib import Path
+
+
+def output_dir(explicit, suite_name: str, backend: str) -> Path:
+    if explicit:
+        return Path(explicit)
+    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+    label = f"{stamp}-{suite_name}-{backend}".replace(" ", "_").replace("/", "_")
+    return Path("outputs") / label

diffeval/cli/parser.py

@@ -0,0 +1,17 @@
+"""argparse for run and list."""
+
+import argparse
+
+from diffeval.cli.run_parser_eval import add_eval_flags
+from diffeval.cli.run_parser_model import add_model_flags
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="diffeval")
+    commands = parser.add_subparsers(dest="command", required=True)
+    run = commands.add_parser("run")
+    add_model_flags(run)
+    add_eval_flags(run)
+    listed = commands.add_parser("list")
+    listed.add_argument("kind", choices=("tasks", "metrics", "backends"))
+    return parser

diffeval/cli/print_paths.py

@@ -0,0 +1,12 @@
+"""Print the table and the written output paths."""
+
+from diffeval.report.console import console
+
+
+def print_paths(folder, stream, body=None) -> None:
+    if body is not None:
+        print(console(body), end="", file=stream)
+    print(folder / "results.json", file=stream)
+    print(folder / "report.md", file=stream)
+    print(folder / "report.html", file=stream)
+    print(folder / "images", file=stream)

diffeval/cli/print_registered.py

@@ -0,0 +1,6 @@
+"""Print registry names and descriptions."""
+
+
+def print_registered(table: dict, stream) -> None:
+    for name in sorted(table):
+        print(f"{name}\t{getattr(table[name], 'description', '')}", file=stream)

diffeval/cli/run_parser_eval.py

@@ -0,0 +1,14 @@
+"""Evaluation flags for diffeval run."""
+
+from diffeval.cli.csv_ints import csv_ints
+from diffeval.cli.csv_names import csv_names
+
+
+def add_eval_flags(parser) -> None:
+    parser.add_argument("--suite", required=True)
+    parser.add_argument("--seeds", type=csv_ints, default=None)
+    parser.add_argument("--metrics", type=csv_names, default=None)
+    parser.add_argument("--width", type=int, default=None)
+    parser.add_argument("--height", type=int, default=None)
+    parser.add_argument("--output", default=None)
+    parser.add_argument("--strict", action="store_true")

diffeval/cli/run_parser_model.py

@@ -0,0 +1,13 @@
+"""Model flags for diffeval run."""
+
+
+def add_model_flags(parser) -> None:
+    parser.add_argument("--backend", default=None)
+    parser.add_argument("--pipeline", default=None)
+    parser.add_argument("--dtype", default=None)
+    parser.add_argument("--steps", type=int, default=None)
+    parser.add_argument("--guidance", type=float, default=None)
+    parser.add_argument("--device", default=None)
+    parser.add_argument("--attention", choices=("sdpa", "xformers", "eager"), default=None)
+    parser.add_argument("--variant", default=None)
+    parser.add_argument("--batch-size", dest="batch_size", type=int, default=None)

diffeval/metrics/__init__.py

@@ -0,0 +1,7 @@
+"""Import metrics so they register."""
+
+import diffeval.metrics.clip.metric
+import diffeval.metrics.cmmd.metric
+import diffeval.metrics.fid.metric
+import diffeval.metrics.geneval.metric
+import diffeval.metrics.kid.metric

diffeval/metrics/clip/__init__.py

@@ -0,0 +1 @@
+"""CLIPScore metric."""

diffeval/metrics/clip/features.py

@@ -0,0 +1,14 @@
+"""L2-normalized CLIP image and text features."""
+
+
+def _norm(tensor):
+    return tensor / tensor.norm(dim=-1, keepdim=True)
+
+
+def features(model, processor, images, texts, device):
+    import torch
+    batch = processor(text=texts, images=images, return_tensors="pt", padding=True, truncation=True, max_length=77)
+    batch = batch.to(device)
+    with torch.no_grad():
+        output = model(**batch)
+    return _norm(output.image_embeds), _norm(output.text_embeds)

diffeval/metrics/clip/formula.py

@@ -0,0 +1,6 @@
+"""Mean of weight * max(cosine, 0)."""
+
+
+def clip_score(cosines, weight: float) -> float:
+    scores = [weight * max(float(value), 0.0) for value in cosines]
+    return sum(scores) / len(scores)

diffeval/metrics/clip/load.py

@@ -0,0 +1,17 @@
+"""Load CLIPModel and CLIPProcessor."""
+
+_CACHE = {}
+
+
+def load_clip(model_name: str):
+    if model_name in _CACHE:
+        return _CACHE[model_name]
+    try:
+        from transformers import CLIPModel, CLIPProcessor
+    except ImportError:
+        return None
+    model = CLIPModel.from_pretrained(model_name)
+    model.eval()
+    processor = CLIPProcessor.from_pretrained(model_name)
+    _CACHE[model_name] = (model, processor)
+    return _CACHE[model_name]

diffeval/metrics/clip/metric.py

@@ -0,0 +1,20 @@
+"""Register clip_score."""
+
+from diffeval.metrics.clip.protocol import PROTOCOL
+from diffeval.metrics.clip.score import score_seed
+from diffeval.registry.register import register
+from diffeval.registry.tables import METRICS
+
+
+class ClipScore:
+    name = "clip_score"
+    description = "CLIPScore (Hessel et al. 2021), ViT-B/32, weight 2.5."
+    higher_is_better = True
+    requires_references = False
+    protocol = PROTOCOL
+
+    def compute(self, samples, context):
+        return score_seed(samples, context, self.protocol)
+
+
+register(METRICS, "clip_score")(ClipScore)

diffeval/metrics/clip/protocol.py

@@ -0,0 +1,10 @@
+"""CLIPScore definition text."""
+
+PROTOCOL = {
+    "paper": "Hessel et al. 2021",
+    "model": "openai/clip-vit-base-patch32",
+    "weight": 2.5,
+    "crop": "shorter-side resize, center crop 224",
+    "text": "raw prompt, truncated to 77 tokens, no prefix",
+    "formula": "mean of 2.5 * max(cosine, 0)",
+}

diffeval/metrics/clip/sample_image.py

@@ -0,0 +1,10 @@
+"""A sample's pixels, read back from its saved PNG when the runner dropped them."""
+
+from PIL import Image
+
+
+def sample_image(sample, folder):
+    if sample.image is not None:
+        return sample.image.convert("RGB")
+    with Image.open(folder / f"{sample.prompt_id}.png") as image:
+        return image.convert("RGB")

diffeval/metrics/clip/score.py

@@ -0,0 +1,25 @@
+"""Per-image CLIP scores for one seed."""
+
+from diffeval.metrics.clip.features import features
+from diffeval.metrics.clip.formula import clip_score
+from diffeval.metrics.clip.load import load_clip
+from diffeval.metrics.clip.sample_image import sample_image
+from diffeval.metrics.ok import ok
+from diffeval.metrics.skipped import skipped
+
+CHUNK = 32
+
+
+def score_seed(samples, context, protocol):
+    loaded = load_clip(protocol["model"])
+    if loaded is None:
+        return skipped("clip_score", "metrics extra is not installed", True, protocol)
+    model, processor = loaded
+    model.to(context.device)
+    cosines = []
+    for start in range(0, len(samples), CHUNK):
+        part = samples[start:start + CHUNK]
+        images = [sample_image(sample, context.generated_dir) for sample in part]
+        image_f, text_f = features(model, processor, images, [sample.prompt for sample in part], context.device)
+        cosines.extend((image_f * text_f).sum(dim=-1).tolist())
+    return ok("clip_score", float(clip_score(cosines, protocol["weight"])), True, protocol)

diffeval/metrics/cmmd/__init__.py

@@ -0,0 +1 @@
+"""CMMD metric."""

diffeval/metrics/cmmd/embed_batch.py

@@ -0,0 +1,26 @@
+"""One CLIP vision batch, L2-normalized like the official CMMD."""
+
+_CACHE = {}
+
+
+def load_vision(model_name: str):
+    if model_name in _CACHE:
+        return _CACHE[model_name]
+    try:
+        from transformers import CLIPImageProcessor, CLIPModel
+    except ImportError:
+        return None
+    model = CLIPModel.from_pretrained(model_name)
+    model.eval()
+    processor = CLIPImageProcessor.from_pretrained(model_name)
+    processor.do_resize = processor.do_center_crop = processor.do_rescale = False
+    _CACHE[model_name] = (model, processor)
+    return _CACHE[model_name]
+
+
+def embed_batch(model, processor, images, device):
+    import torch
+    pixels = processor(images=images, return_tensors="pt")["pixel_values"].to(device)
+    with torch.no_grad():
+        embeds = model.visual_projection(model.vision_model(pixel_values=pixels).pooler_output)
+        return (embeds / embeds.norm(dim=-1, keepdim=True)).detach().cpu().numpy()

diffeval/metrics/cmmd/embed_paths.py

@@ -0,0 +1,30 @@
+"""Embed image files and the reference bank."""
+
+from PIL import Image
+
+import numpy as np
+
+from diffeval.metrics.cmmd.embed_batch import embed_batch
+from diffeval.metrics.cmmd.resize import resize_image
+from diffeval.metrics.images import list_images
+
+CHUNK = 32
+
+
+def embed_paths(model, processor, paths, device):
+    chunks = []
+    for start in range(0, len(paths), CHUNK):
+        images = [resize_image(Image.open(path)) for path in paths[start:start + CHUNK]]
+        chunks.append(embed_batch(model, processor, images, device))
+    return np.concatenate(chunks) if chunks else None
+
+
+def embed_dir(model, processor, folder, device):
+    return embed_paths(model, processor, list_images(folder), device)
+
+
+def paired(context, model, processor):
+    generated = embed_dir(model, processor, context.generated_dir, context.device)
+    if context.cmmd_ref_embed is not None:
+        return generated, np.load(context.cmmd_ref_embed)
+    return generated, embed_dir(model, processor, context.reference_dir, context.device)

diffeval/metrics/cmmd/metric.py

@@ -0,0 +1,20 @@
+"""Register cmmd."""
+
+from diffeval.metrics.cmmd.protocol import PROTOCOL
+from diffeval.metrics.cmmd.score import score_cmmd
+from diffeval.registry.register import register
+from diffeval.registry.tables import METRICS
+
+
+class Cmmd:
+    name = "cmmd"
+    description = "CMMD (Jayasumana et al. 2024), ViT-L/14 at 336, L2-normalized embeddings."
+    higher_is_better = False
+    requires_references = True
+    protocol = PROTOCOL
+
+    def compute(self, samples, context):
+        return score_cmmd(context)
+
+
+register(METRICS, "cmmd")(Cmmd)

diffeval/metrics/cmmd/mmd.py

@@ -0,0 +1,23 @@
+"""Minimum-variance RBF MMD, sigma 10, scale 1000, in row blocks."""
+
+import numpy as np
+
+_SIGMA = 10
+_SCALE = 1000
+BLOCK = 2048
+
+
+def _kernel(x, y, x_sq, y_sq, gamma):
+    total = 0.0
+    for start in range(0, len(x), BLOCK):
+        dist = -2 * (x[start:start + BLOCK] @ y.T) + x_sq[start:start + BLOCK, None] + y_sq[None, :]
+        total += np.exp(-gamma * dist).sum()
+    return total / (len(x) * len(y))
+
+
+def mmd(x, y) -> float:
+    x, y = np.asarray(x, dtype=np.float64), np.asarray(y, dtype=np.float64)
+    gamma = 1 / (2 * _SIGMA**2)
+    x_sq, y_sq = np.sum(x * x, axis=1), np.sum(y * y, axis=1)
+    cross = _kernel(x, y, x_sq, y_sq, gamma)
+    return float(_SCALE * (_kernel(x, x, x_sq, x_sq, gamma) + _kernel(y, y, y_sq, y_sq, gamma) - 2 * cross))

diffeval/metrics/cmmd/protocol.py

@@ -0,0 +1,11 @@
+"""CMMD definition text."""
+
+PROTOCOL = {
+    "paper": "Jayasumana et al. 2024",
+    "model": "openai/clip-vit-large-patch14-336",
+    "resize": "bicubic resize of the 8-bit image to 336 (PIL), then divided by 255",
+    "crop": "center square crop",
+    "embeddings": "vision projection, L2-normalized",
+    "check": "cmmd-pytorch sample folders give 7.696",
+    "distance": "minimum-variance RBF MMD, sigma 10, scale 1000, diagonal included",
+}

diffeval/metrics/cmmd/resize.py

@@ -0,0 +1,17 @@
+"""Center square crop, 8-bit bicubic resize to 336, then divide by 255."""
+
+from PIL import Image
+
+import numpy as np
+
+SIZE = 336
+
+
+def resize_image(image: Image.Image) -> np.ndarray:
+    image = image.convert("RGB")
+    width, height = image.size
+    side = min(width, height)
+    top, left = (height - side) // 2, (width - side) // 2
+    square = image.crop((left, top, left + side, top + side))
+    resized = square.resize((SIZE, SIZE), Image.Resampling.BICUBIC)
+    return np.asarray(resized, dtype=np.float32) / 255.0