diffeval 0.1.0__py3-none-any.whl

diffeval/metrics/cmmd/score.py

@@ -0,0 +1,24 @@
+"""Reference folder or .npy against generated PNGs."""
+
+from diffeval.metrics.cmmd.embed_batch import load_vision
+from diffeval.metrics.cmmd.embed_paths import paired
+from diffeval.metrics.cmmd.mmd import mmd
+from diffeval.metrics.cmmd.protocol import PROTOCOL
+from diffeval.metrics.count_warning import count_warning
+from diffeval.metrics.ok import ok
+from diffeval.metrics.skipped import skipped
+
+
+def score_cmmd(context):
+    if context.reference_dir is None and context.cmmd_ref_embed is None:
+        return skipped("cmmd", "reference set is missing", False, PROTOCOL)
+    loaded = load_vision(PROTOCOL["model"])
+    if loaded is None:
+        return skipped("cmmd", "metrics extra is not installed", False, PROTOCOL)
+    model, processor = loaded
+    model.to(context.device)
+    generated, reference = paired(context, model, processor)
+    if generated is None or reference is None:
+        return skipped("cmmd", "reference set is missing", False, PROTOCOL)
+    warning = count_warning(len(generated))
+    return ok("cmmd", float(mmd(reference, generated)), False, PROTOCOL, {} if warning is None else {"warning": warning})

diffeval/metrics/count_warning.py

@@ -0,0 +1,7 @@
+"""Warn when the image count is below the published setup."""
+
+
+def count_warning(count: int, usual: int = 30000) -> str | None:
+    if count >= usual:
+        return None
+    return f"{count} images is below the published COCO setup of {usual}"

diffeval/metrics/distribution/__init__.py

@@ -0,0 +1 @@
+"""FID and KID."""

diffeval/metrics/distribution/call.py

@@ -0,0 +1,24 @@
+"""Call clean-fid for FID or KID."""
+
+from diffeval.metrics.count_warning import count_warning
+from diffeval.metrics.distribution.kwargs_folder import distribution_kwargs
+from diffeval.metrics.distribution.seed_numpy import seeded
+from diffeval.metrics.fid.guard import fid_ready
+from diffeval.metrics.images import list_images
+from diffeval.metrics.ok import ok
+from diffeval.metrics.skipped import skipped
+
+
+def score_distribution(kind, context, protocol, model_name=None):
+    reason = fid_ready(context)
+    if reason:
+        return skipped(kind, reason, False, protocol)
+    try:
+        from cleanfid import fid as clean_fid
+    except ImportError:
+        return skipped(kind, "metrics extra is not installed", False, protocol)
+    fn = clean_fid.compute_fid if kind == "fid" else clean_fid.compute_kid
+    kwargs = distribution_kwargs(context, model_name)
+    value = seeded(context.seed, lambda: fn(**kwargs))
+    warning = count_warning(len(list_images(context.generated_dir)))
+    return ok(kind, float(value), False, protocol, {} if warning is None else {"warning": warning})

diffeval/metrics/distribution/kwargs_folder.py

@@ -0,0 +1,16 @@
+"""Folder-mode arguments for clean-fid."""
+
+from diffeval.metrics.distribution.kwargs_stats import stats_kwargs
+
+
+def folder_kwargs(reference, generated, device, batch_size, model_name=None) -> dict:
+    body = {"fdir1": str(reference), "fdir2": str(generated), "mode": "clean", "num_workers": 0, "batch_size": int(batch_size), "device": device, "verbose": False, "use_dataparallel": device == "cuda"}
+    if model_name:
+        body["model_name"] = model_name
+    return body
+
+
+def distribution_kwargs(context, model_name=None) -> dict:
+    if context.clean_fid_stats:
+        return stats_kwargs(context.generated_dir, context.clean_fid_stats, context.device, context.batch_size, model_name)
+    return folder_kwargs(context.reference_dir, context.generated_dir, context.device, context.batch_size, model_name)

diffeval/metrics/distribution/kwargs_stats.py

@@ -0,0 +1,8 @@
+"""Stats-mode arguments for clean-fid."""
+
+
+def stats_kwargs(generated, stats, device, batch_size, model_name=None) -> dict:
+    body = {"fdir1": str(generated), "dataset_name": stats["dataset_name"], "dataset_res": stats["dataset_res"], "dataset_split": stats["dataset_split"], "mode": "clean", "num_workers": 0, "batch_size": int(batch_size), "device": device, "verbose": False, "use_dataparallel": device == "cuda"}
+    if model_name:
+        body["model_name"] = model_name
+    return body

diffeval/metrics/distribution/protocol.py

@@ -0,0 +1,12 @@
+"""Shared FID and KID protocol."""
+
+PROTOCOL = {
+    "library": "clean-fid",
+    "mode": "clean",
+    "model": "inception_v3",
+    "resize": "clean-fid resizes each RGB image to 299x299 with PIL bicubic per channel in float and does not quantize back to uint8",
+    "pre_resize": "diffeval does not pre-resize",
+    "num_workers": 0,
+    "dataparallel": "only when the metric device is cuda",
+    "kid": "degree-3 polynomial, 100 subsets, subset size at most 1000",
+}

diffeval/metrics/distribution/seed_numpy.py

@@ -0,0 +1,12 @@
+"""Seed NumPy around a call and restore the previous state."""
+
+import numpy as np
+
+
+def seeded(seed, fn):
+    state = np.random.get_state()
+    np.random.seed(seed)
+    try:
+        return fn()
+    finally:
+        np.random.set_state(state)

diffeval/metrics/fid/__init__.py

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

diffeval/metrics/fid/guard.py

@@ -0,0 +1,13 @@
+"""Require two images and a reference or stats."""
+
+from diffeval.metrics.images import list_images
+
+
+def fid_ready(context) -> str | None:
+    if context.reference_dir is None and not context.clean_fid_stats:
+        return "reference set is missing"
+    if len(list_images(context.generated_dir)) < 2:
+        return "fewer than two images"
+    if context.reference_dir is not None and not context.clean_fid_stats and len(list_images(context.reference_dir)) < 2:
+        return "fewer than two images"
+    return None

diffeval/metrics/fid/metric.py

@@ -0,0 +1,20 @@
+"""Register fid."""
+
+from diffeval.metrics.distribution.call import score_distribution
+from diffeval.metrics.distribution.protocol import PROTOCOL
+from diffeval.registry.register import register
+from diffeval.registry.tables import METRICS
+
+
+class Fid:
+    name = "fid"
+    description = "FID via clean-fid, mode clean, Inception-v3."
+    higher_is_better = False
+    requires_references = True
+    protocol = PROTOCOL
+
+    def compute(self, samples, context):
+        return score_distribution("fid", context, self.protocol, "inception_v3")
+
+
+register(METRICS, "fid")(Fid)

diffeval/metrics/geneval/__init__.py

@@ -0,0 +1 @@
+"""GenEval stub."""

diffeval/metrics/geneval/metric.py

@@ -0,0 +1,22 @@
+"""Register geneval and return skipped."""
+
+from diffeval.metrics.geneval.protocol import PROTOCOL
+from diffeval.metrics.geneval.tags import TAGS
+from diffeval.metrics.skipped import skipped
+from diffeval.registry.register import register
+from diffeval.registry.tables import METRICS
+
+
+class Geneval:
+    name = "geneval"
+    description = "GenEval v1 stub. The Mask2Former detector is not imported."
+    higher_is_better = True
+    requires_references = False
+    protocol = PROTOCOL
+
+    def compute(self, samples, context):
+        details = {"by_tag": {tag: None for tag in TAGS}}
+        return skipped(self.name, "GenEval detector is not in v1", True, self.protocol, details)
+
+
+register(METRICS, "geneval")(Geneval)

diffeval/metrics/geneval/protocol.py

@@ -0,0 +1,5 @@
+"""Stub protocol and tag schema."""
+
+from diffeval.metrics.geneval.tags import TAGS
+
+PROTOCOL = {"detector": "Mask2Former is not imported", "tags": list(TAGS)}

diffeval/metrics/geneval/tags.py

@@ -0,0 +1,3 @@
+"""The six GenEval tag names."""
+
+TAGS = ("single_object", "two_object", "counting", "colors", "position", "color_attr")

diffeval/metrics/images.py

@@ -0,0 +1,12 @@
+"""List image files in a directory."""
+
+from pathlib import Path
+
+_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
+
+
+def list_images(folder) -> list[Path]:
+    root = Path(folder)
+    if not root.is_dir():
+        return []
+    return sorted(path for path in root.iterdir() if path.suffix.lower() in _SUFFIXES)

diffeval/metrics/kid/__init__.py

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

diffeval/metrics/kid/guard.py

@@ -0,0 +1,3 @@
+"""Same image and reference guard as FID."""
+
+from diffeval.metrics.fid.guard import fid_ready as kid_ready

diffeval/metrics/kid/metric.py

@@ -0,0 +1,20 @@
+"""Register kid."""
+
+from diffeval.metrics.distribution.call import score_distribution
+from diffeval.metrics.distribution.protocol import PROTOCOL
+from diffeval.registry.register import register
+from diffeval.registry.tables import METRICS
+
+
+class Kid:
+    name = "kid"
+    description = "KID via clean-fid, mode clean, polynomial degree 3."
+    higher_is_better = False
+    requires_references = True
+    protocol = PROTOCOL
+
+    def compute(self, samples, context):
+        return score_distribution("kid", context, self.protocol)
+
+
+register(METRICS, "kid")(Kid)

diffeval/metrics/ok.py

@@ -0,0 +1,9 @@
+"""Build an ok MetricResult."""
+
+from diffeval.types.result import MetricResult
+from diffeval.types.result_rules import check
+
+
+def ok(name: str, value: float, higher_is_better, protocol: dict, details: dict | None = None):
+    check("ok", value, None)
+    return MetricResult(name, value, "ok", None, higher_is_better, details or {}, dict(protocol))

diffeval/metrics/protocol.py

@@ -0,0 +1,17 @@
+"""Metric contract."""
+
+from typing import Protocol
+
+from diffeval.types.context import EvalContext
+from diffeval.types.result import MetricResult
+from diffeval.types.sample import Sample
+
+
+class Metric(Protocol):
+    name: str
+    description: str
+    higher_is_better: bool | None
+    requires_references: bool
+    protocol: dict
+
+    def compute(self, samples: list[Sample], context: EvalContext) -> MetricResult: ...

diffeval/metrics/skipped.py

@@ -0,0 +1,10 @@
+"""Build a skipped MetricResult."""
+
+from diffeval.types.result import MetricResult
+from diffeval.types.result_rules import check
+
+
+def skipped(name: str, error: str, higher_is_better, protocol: dict, details: dict | None = None):
+    check("skipped", None, error)
+    body = {} if details is None else details
+    return MetricResult(name, None, "skipped", error, higher_is_better, body, dict(protocol))

diffeval/plugins.py

@@ -0,0 +1,7 @@
+"""Import concrete tasks, metrics, and backends."""
+
+
+def load() -> None:
+    import diffeval.backends
+    import diffeval.metrics
+    import diffeval.tasks

diffeval/registry/__init__.py

@@ -0,0 +1 @@
+"""Task, metric, and backend registries."""

diffeval/registry/lookup.py

@@ -0,0 +1,9 @@
+"""Look up a registered object by name."""
+
+
+def get(table: dict, name: str, kind: str):
+    try:
+        return table[name]
+    except KeyError:
+        known = ", ".join(sorted(table)) or "(none)"
+        raise KeyError(f"Unknown {kind} {name!r}. Known: {known}") from None

diffeval/registry/names.py

@@ -0,0 +1,5 @@
+"""Sorted registry names."""
+
+
+def names(table: dict) -> list[str]:
+    return sorted(table)

diffeval/registry/register.py

@@ -0,0 +1,11 @@
+"""Register one object under one name."""
+
+
+def register(table: dict, name: str):
+    def decorator(obj):
+        if name in table and table[name] is not obj:
+            raise ValueError(f"{name} is already registered")
+        table[name] = obj
+        return obj
+
+    return decorator

diffeval/registry/tables.py

@@ -0,0 +1,5 @@
+"""Registries for tasks, metrics, and backends."""
+
+TASKS: dict = {}
+METRICS: dict = {}
+BACKENDS: dict = {}

diffeval/registry/unregister.py

@@ -0,0 +1,5 @@
+"""Drop one registry entry."""
+
+
+def unregister(table: dict, name: str) -> None:
+    table.pop(name, None)

diffeval/report/__init__.py

@@ -0,0 +1 @@
+"""Console table, Markdown report, and HTML report."""

diffeval/report/console.py

@@ -0,0 +1,18 @@
+"""Plain text table of aggregate scores."""
+
+from diffeval.report.console_widths import widths
+from diffeval.report.direction import direction
+from diffeval.report.number import number
+
+
+def _cells(name, row) -> list[str]:
+    seeds = ",".join(number(item) for item in row["per_seed"])
+    return [name, row["status"], number(row["mean"]), number(row["std"]), seeds, direction(row.get("higher_is_better"))]
+
+
+def console(body: dict) -> str:
+    header = ["metric", "status", "mean", "std", "per-seed", "direction"]
+    rows = [header, *[_cells(name, row) for name, row in body["aggregate"].items()]]
+    size = widths(rows)
+    lines = ["  ".join(cell.ljust(size[i]) for i, cell in enumerate(row)) for row in rows]
+    return "\n".join(lines) + "\n"

diffeval/report/console_widths.py

@@ -0,0 +1,5 @@
+"""Column widths for the text table."""
+
+
+def widths(rows: list[list[str]]) -> list[int]:
+    return [max(len(row[index]) for row in rows) for index in range(len(rows[0]))]

diffeval/report/direction.py

@@ -0,0 +1,7 @@
+"""higher or lower from a metric flag."""
+
+
+def direction(higher_is_better) -> str:
+    if higher_is_better is None:
+        return "n/a"
+    return "higher" if higher_is_better else "lower"

diffeval/report/html_figure.py

@@ -0,0 +1,9 @@
+"""One escaped figure for one sample."""
+
+from html import escape
+
+
+def html_figure(sample: dict) -> str:
+    src = escape(sample["path"], quote=True)
+    label = escape(f"seed {sample['seed']}")
+    return f'<figure><img src="{src}" alt="{label}"><figcaption>{label}</figcaption></figure>'

diffeval/report/html_grid.py

@@ -0,0 +1,19 @@
+"""Image grids grouped by prompt."""
+
+from html import escape
+
+from diffeval.report.html_figure import html_figure
+
+
+def _samples(body, prompt_id):
+    return [sample for seed in body["per_seed"] for sample in seed["samples"] if sample["prompt_id"] == prompt_id]
+
+
+def html_grid(body: dict) -> str:
+    blocks = []
+    for prompt in body["suite"]["prompts"]:
+        figures = "".join(html_figure(sample) for sample in _samples(body, prompt["id"]))
+        heading = escape(prompt["id"])
+        text = escape(prompt["prompt"])
+        blocks.append(f'<section><h2>{heading}</h2><p>{text}</p>{figures}</section>')
+    return "\n".join(blocks)

diffeval/report/html_page.py

@@ -0,0 +1,15 @@
+"""Fill the static HTML template."""
+
+from html import escape
+from pathlib import Path
+
+from diffeval.report.html_grid import html_grid
+from diffeval.report.html_row import html_row
+
+
+def html_page(body: dict) -> str:
+    rows = "\n".join(html_row(name, row) for name, row in body["aggregate"].items())
+    template = (Path(__file__).parent / "templates" / "page.html").read_text(encoding="utf-8")
+    text = template.replace("{{TITLE}}", escape(str(body["suite"]["name"])))
+    text = text.replace("{{ROWS}}", rows)
+    return text.replace("{{GRID}}", html_grid(body))

diffeval/report/html_row.py

@@ -0,0 +1,12 @@
+"""One escaped HTML table row."""
+
+from html import escape
+
+from diffeval.report.direction import direction
+from diffeval.report.number import number
+
+
+def html_row(name: str, row: dict) -> str:
+    seeds = ", ".join(number(item) for item in row["per_seed"])
+    cells = (name, row["status"], number(row["mean"]), number(row["std"]), seeds, direction(row.get("higher_is_better")))
+    return "<tr>" + "".join(f"<td>{escape(cell)}</td>" for cell in cells) + "</tr>"

diffeval/report/markdown.py

@@ -0,0 +1,11 @@
+"""Markdown report for one run."""
+
+from diffeval.report.markdown_row import markdown_row
+
+
+def markdown(body: dict) -> str:
+    header = "| metric | status | mean | std | per-seed | direction |"
+    rule = "| --- | --- | --- | --- | --- | --- |"
+    rows = [markdown_row(name, row) for name, row in body["aggregate"].items()]
+    title = f"# {body['suite']['name']}"
+    return "\n".join([title, "", header, rule, *rows, ""]) + "\n"

diffeval/report/markdown_row.py

@@ -0,0 +1,10 @@
+"""One Markdown table row."""
+
+from diffeval.report.direction import direction
+from diffeval.report.number import number
+
+
+def markdown_row(name: str, row: dict) -> str:
+    seeds = ", ".join(number(item) for item in row["per_seed"])
+    cells = (name, row["status"], number(row["mean"]), number(row["std"]), seeds, direction(row.get("higher_is_better")))
+    return "| " + " | ".join(cells) + " |"

diffeval/report/number.py

@@ -0,0 +1,7 @@
+"""Format a float, or n/a when the value is missing."""
+
+
+def number(value) -> str:
+    if value is None:
+        return "n/a"
+    return f"{float(value):.6g}"

diffeval/report/templates/page.html

@@ -0,0 +1,35 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<title>{{TITLE}}</title>
+<style>
+body { font-family: sans-serif; margin: 1.5rem; color: #222; }
+.note { max-width: 40rem; }
+.layout { display: flex; gap: 2rem; align-items: flex-start; flex-wrap: wrap; }
+table { border-collapse: collapse; }
+th, td { border: 1px solid #ccc; padding: 0.3rem 0.6rem; text-align: left; }
+section { margin-bottom: 1.5rem; }
+figure { display: inline-block; margin: 0.4rem; vertical-align: top; }
+img { display: block; max-width: 12rem; height: auto; background: #eee; }
+figcaption { font-size: 0.85rem; }
+</style>
+</head>
+<body>
+<h1>{{TITLE}}</h1>
+<p class="note">Metric numbers have to be checked against the published reference before they are compared to a paper. Scores have to be checked against the published reference implementation and dataset. The core suite is a smoke suite. The bundled drawings are not COCO.</p>
+<div class="layout">
+<table>
+<thead>
+<tr><th>metric</th><th>status</th><th>mean</th><th>std</th><th>per-seed</th><th>direction</th></tr>
+</thead>
+<tbody>
+{{ROWS}}
+</tbody>
+</table>
+<div class="grids">
+{{GRID}}
+</div>
+</div>
+</body>
+</html>

diffeval/runner/__init__.py

@@ -0,0 +1 @@
+"""Generate images, score them, and write results.json."""

diffeval/runner/generate_save.py

@@ -0,0 +1,19 @@
+"""Generate in chunks, save each PNG, and drop the pixels from memory."""
+
+from pathlib import Path
+
+from diffeval.runner.relpath import relpath
+from diffeval.runner.save_png import save_png
+
+CHUNK = 64
+
+
+def generate_save(backend, requests, folder: Path, output_dir) -> list:
+    samples = []
+    for start in range(0, len(requests), CHUNK):
+        for sample in backend.generate(requests[start:start + CHUNK]):
+            path = folder / f"{sample.prompt_id}.png"
+            save_png(sample.image, path)
+            sample.path, sample.image = relpath(path, Path(output_dir)), None
+            samples.append(sample)
+    return samples

diffeval/runner/metric_device.py

@@ -0,0 +1,16 @@
+"""Pick the device metrics use."""
+
+
+def metric_device(requested: str | None) -> str:
+    if requested:
+        return requested
+    try:
+        import torch
+    except ImportError:
+        return "cpu"
+    if torch.cuda.is_available():
+        return "cuda"
+    mps = getattr(torch.backends, "mps", None)
+    if mps is not None and mps.is_available():
+        return "mps"
+    return "cpu"

diffeval/runner/one_seed.py

@@ -0,0 +1,20 @@
+"""Generate, save, and score one seed."""
+
+from pathlib import Path
+
+from diffeval.runner.generate_save import generate_save
+from diffeval.runner.requests import requests_for
+from diffeval.runner.safe_metric import safe_metric
+from diffeval.types.context import EvalContext
+
+
+def _public(sample) -> dict:
+    return {"prompt_id": sample.prompt_id, "prompt": sample.prompt, "seed": sample.seed, "width": sample.width, "height": sample.height, "path": sample.path, "negative_prompt": sample.negative_prompt}
+
+
+def one_seed(backend, suite, seed, picked, metrics, output_dir, device):
+    folder = Path(output_dir) / "images" / f"seed_{seed}"
+    samples = generate_save(backend, requests_for(suite, seed, picked["width"], picked["height"]), folder, output_dir)
+    context = EvalContext(folder, suite.reference_dir, device, seed, suite.clean_fid_stats, suite.cmmd_ref_embed, backend.describe()["batch_size"])
+    scored = [safe_metric(metrics[name], samples, context) for name in picked["metrics"]]
+    return {"seed": seed, "samples": [_public(item) for item in samples], "metrics": scored}

diffeval/runner/options.py

@@ -0,0 +1,23 @@
+"""Options for one evaluation run."""
+
+from dataclasses import dataclass
+from pathlib import Path
+
+from diffeval.types.suite import Suite
+
+
+@dataclass(frozen=True)
+class RunOptions:
+    suite: Suite
+    backend: str
+    backend_config: dict
+    output_dir: Path
+    seeds: tuple[int, ...] | None = None
+    metrics: tuple[str, ...] | None = None
+    width: int | None = None
+    height: int | None = None
+
+
+def chosen(options) -> dict:
+    suite = options.suite
+    return {"seeds": options.seeds or suite.seeds, "metrics": options.metrics or suite.metrics, "width": options.width or suite.width, "height": options.height or suite.height}

diffeval/runner/payload.py

@@ -0,0 +1,26 @@
+"""Assemble the results object."""
+
+from datetime import datetime, timezone
+
+from diffeval.system.describe import describe as system_info
+from diffeval.system.versions import versions
+
+
+def _text(path):
+    return None if path is None else str(path)
+
+
+def _protocols(per_seed):
+    found = {}
+    for seed in per_seed:
+        for metric in seed["metrics"]:
+            found.setdefault(metric["name"], metric.pop("protocol", {}))
+    return found
+
+
+def payload(options, backend, device, seeds, per_seed, aggregate) -> dict:
+    suite = options.suite
+    prompts = [{"id": item.id, "prompt": item.prompt, "negative_prompt": item.negative_prompt} for item in suite.prompts]
+    info = {"name": suite.name, "description": suite.description, "source": suite.source, "width": options.width or suite.width, "height": options.height or suite.height, "prompts": prompts, "reference_dir": _text(suite.reference_dir), "clean_fid_stats": suite.clean_fid_stats, "cmmd_ref_embed": _text(suite.cmmd_ref_embed)}
+    stamp = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
+    return {"schema_version": 1, "created_at": stamp, "suite": info, "backend": backend.describe(), "metric_device": device, "seeds": list(seeds), "versions": versions(), "system": system_info(), "metric_protocols": _protocols(per_seed), "per_seed": per_seed, "aggregate": aggregate, "outputs": ["results.json", "report.md", "report.html"]}

diffeval/runner/relpath.py

@@ -0,0 +1,7 @@
+"""Path relative to the output directory."""
+
+from pathlib import Path
+
+
+def relpath(path: Path, output_dir: Path) -> str:
+    return Path(path).relative_to(output_dir).as_posix()

diffeval/runner/requests.py

@@ -0,0 +1,7 @@
+"""Generation requests for one seed."""
+
+from diffeval.types.request import GenerationRequest
+
+
+def requests_for(suite, seed: int, width: int, height: int) -> list:
+    return [GenerationRequest(item.id, item.prompt, seed, width, height, item.negative_prompt) for item in suite.prompts]