diffeval 0.1.0__py3-none-any.whl

diffeval/__init__.py

@@ -0,0 +1,7 @@
+"""Evaluation harness for image and video diffusion models."""
+
+__version__ = "0.1.0"
+
+from diffeval.plugins import load as load_plugins
+
+load_plugins()

diffeval/__main__.py

@@ -0,0 +1,6 @@
+"""python -m diffeval"""
+
+from diffeval.cli.main import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())

diffeval/aggregate/__init__.py

@@ -0,0 +1 @@
+"""Mean and sample standard deviation across seeds."""

diffeval/aggregate/mean.py

@@ -0,0 +1,7 @@
+"""Arithmetic mean of successful values."""
+
+
+def mean(values: list[float]):
+    if not values:
+        return None
+    return sum(values) / len(values)

diffeval/aggregate/names.py

@@ -0,0 +1,11 @@
+"""Metric names in first-seen order."""
+
+
+def names(per_seed: list[dict]) -> list[str]:
+    found = []
+    for seed in per_seed:
+        for metric in seed.get("metrics", []):
+            name = metric["name"]
+            if name not in found:
+                found.append(name)
+    return found

diffeval/aggregate/row.py

@@ -0,0 +1,26 @@
+"""One aggregate row."""
+
+from diffeval.aggregate.mean import mean
+from diffeval.aggregate.sample_std import sample_std
+from diffeval.aggregate.series import series
+
+
+def _status(statuses: list) -> str:
+    ok_n = sum(item == "ok" for item in statuses)
+    if statuses and ok_n == len(statuses):
+        return "ok"
+    if ok_n:
+        return "partial"
+    if statuses and not any(item == "error" for item in statuses):
+        return "skipped"
+    return "error"
+
+
+def row(per_seed: list[dict], name: str) -> dict:
+    data = series(per_seed, name)
+    values = [item for item in data["per_seed"] if item is not None]
+    data.update(mean=mean(values), std=sample_std(values), n=len(values))
+    data["status"] = _status(data.pop("statuses"))
+    data["name"] = name
+    data["std_ddof"] = 1
+    return data

diffeval/aggregate/sample_std.py

@@ -0,0 +1,9 @@
+"""Sample standard deviation, null below two values."""
+
+import statistics
+
+
+def sample_std(values: list[float]):
+    if len(values) < 2:
+        return None
+    return statistics.stdev(values)

diffeval/aggregate/series.py

@@ -0,0 +1,22 @@
+"""Per-seed values and statuses for one metric."""
+
+
+def series(per_seed: list[dict], name: str) -> dict:
+    values = []
+    statuses = []
+    higher = None
+    for seed in per_seed:
+        metrics = seed.get("metrics", [])
+        match = next((item for item in metrics if item["name"] == name), None)
+        if match is None:
+            values.append(None)
+            statuses.append("skipped")
+        elif match["status"] != "ok":
+            values.append(None)
+            statuses.append(match["status"])
+        else:
+            values.append(match["value"])
+            statuses.append("ok")
+        if match is not None and higher is None:
+            higher = match.get("higher_is_better")
+    return {"per_seed": values, "statuses": statuses, "higher_is_better": higher}

diffeval/aggregate/table.py

@@ -0,0 +1,8 @@
+"""Aggregate mean and sample std across seeds."""
+
+from diffeval.aggregate.names import names
+from diffeval.aggregate.row import row
+
+
+def aggregate(per_seed: list[dict]) -> dict:
+    return {name: row(per_seed, name) for name in names(per_seed)}

diffeval/backends/__init__.py

@@ -0,0 +1,4 @@
+"""Import backends so they register."""
+
+import diffeval.backends.diffusers.backend
+import diffeval.backends.dummy.backend

diffeval/backends/diffusers/__init__.py

@@ -0,0 +1 @@
+"""Diffusers text-to-image backend."""

diffeval/backends/diffusers/attention_apply.py

@@ -0,0 +1,16 @@
+"""Dispatch sdpa, xformers, eager, or detect."""
+
+from diffeval.backends.diffusers.attention_eager import apply_eager
+from diffeval.backends.diffusers.attention_name import attention_name
+from diffeval.backends.diffusers.attention_sdpa import apply_sdpa
+from diffeval.backends.diffusers.attention_xformers import apply_xformers
+
+
+def apply_attention(pipe, choice: str | None) -> str:
+    if choice == "sdpa":
+        apply_sdpa(pipe)
+    elif choice == "eager":
+        apply_eager(pipe)
+    elif choice == "xformers":
+        apply_xformers(pipe)
+    return attention_name(pipe)

diffeval/backends/diffusers/attention_eager.py

@@ -0,0 +1,6 @@
+"""Attach the eager AttnProcessor."""
+
+
+def apply_eager(pipe) -> None:
+    from diffusers.models.attention_processor import AttnProcessor
+    pipe.set_attn_processor(AttnProcessor())

diffeval/backends/diffusers/attention_name.py

@@ -0,0 +1,15 @@
+"""Classify the current attention processor."""
+
+
+def attention_name(pipe) -> str:
+    processors = getattr(pipe, "attn_processors", None) or {}
+    if not processors:
+        return "none"
+    cls = type(next(iter(processors.values()))).__name__
+    if "xformers" in cls.lower():
+        return "xformers"
+    if cls == "AttnProcessor2_0":
+        return "sdpa"
+    if cls == "AttnProcessor":
+        return "eager"
+    return cls

diffeval/backends/diffusers/attention_sdpa.py

@@ -0,0 +1,6 @@
+"""Attach AttnProcessor2_0."""
+
+
+def apply_sdpa(pipe) -> None:
+    from diffusers.models.attention_processor import AttnProcessor2_0
+    pipe.set_attn_processor(AttnProcessor2_0())

diffeval/backends/diffusers/attention_xformers.py

@@ -0,0 +1,5 @@
+"""Enable xformers attention."""
+
+
+def apply_xformers(pipe) -> None:
+    pipe.enable_xformers_memory_efficient_attention()

diffeval/backends/diffusers/backend.py

@@ -0,0 +1,29 @@
+"""Bind and register the diffusers backend."""
+
+from diffeval.backends.diffusers.describe import describe as describe_state
+from diffeval.backends.diffusers.from_config import from_config as read_config
+from diffeval.backends.diffusers.generate import generate as render
+from diffeval.registry.register import register
+from diffeval.registry.tables import BACKENDS
+
+
+@register(BACKENDS, "diffusers")
+class DiffusersBackend:
+    description = "Diffusers AutoPipelineForText2Image."
+
+    def __init__(self, config: dict):
+        self.config, self.state = config, {}
+
+    def prepare(self) -> None:
+        from diffeval.backends.diffusers.load_pipe import prepare_pipe
+        self.state = prepare_pipe(self.config)
+
+    def generate(self, requests):
+        return render(self.state["pipe"], requests, self.config["batch_size"], self.state["sampling"]["steps"], self.state["sampling"]["guidance"])
+
+    def describe(self) -> dict:
+        return describe_state(self.config, self.state)
+
+    @staticmethod
+    def from_config(config: dict):
+        return DiffusersBackend(read_config(config))

diffeval/backends/diffusers/call_kwargs.py

@@ -0,0 +1,15 @@
+"""Compose pipeline call kwargs."""
+
+from diffeval.backends.diffusers.kwargs_generator import generators
+from diffeval.backends.diffusers.kwargs_prompt import prompt_kwargs
+from diffeval.backends.diffusers.kwargs_sampling import sampling_kwargs
+from diffeval.backends.diffusers.kwargs_size import size_kwargs
+
+
+def call_kwargs(pipe, requests, steps, guidance) -> dict:
+    body = prompt_kwargs(requests)
+    body.update(sampling_kwargs(pipe, steps, guidance))
+    body.update(size_kwargs(pipe, requests[0].width, requests[0].height))
+    gens = generators([item.seed for item in requests])
+    body["generator"] = gens[0] if len(requests) == 1 else gens
+    return body

diffeval/backends/diffusers/coerce_dtype.py

@@ -0,0 +1,10 @@
+"""Load float16 and bfloat16 as float32 on CPU."""
+
+from diffeval.backends.diffusers.dtype_name import parse_dtype, torch_dtype
+
+
+def coerce_dtype(requested: str | None, device: str):
+    name = parse_dtype(requested)
+    if device == "cpu" and name in ("float16", "bfloat16"):
+        name = "float32"
+    return torch_dtype(name), name

diffeval/backends/diffusers/describe.py

@@ -0,0 +1,6 @@
+"""describe() after the pipeline loads."""
+
+
+def describe(config: dict, state: dict) -> dict:
+    sampling = state.get("sampling") or {"steps": None, "guidance": None, "steps_source": "pipeline-default", "guidance_source": "pipeline-default"}
+    return {"name": "diffusers", "pipeline": config.get("pipeline"), "dtype": state.get("dtype_text"), "requested_dtype": config.get("dtype"), "steps": sampling["steps"], "guidance": sampling["guidance"], "steps_source": sampling["steps_source"], "guidance_source": sampling["guidance_source"], "batch_size": config.get("batch_size", 1), "device": state.get("device"), "attention_backend": state.get("attention"), "variant": config.get("variant")}

diffeval/backends/diffusers/dtype_name.py

@@ -0,0 +1,23 @@
+"""Map dtype strings to torch dtypes."""
+
+_NAMES = {"float32": "float32", "fp32": "float32", "float16": "float16", "fp16": "float16", "bfloat16": "bfloat16", "bf16": "bfloat16"}
+
+
+def parse_dtype(text: str | None) -> str | None:
+    if text is None:
+        return None
+    if text not in _NAMES:
+        raise ValueError(f"unknown dtype {text}")
+    return _NAMES[text]
+
+
+def torch_dtype(name: str | None):
+    import torch
+    table = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16}
+    return None if name is None else table[name]
+
+
+def label(dtype) -> str | None:
+    if dtype is None:
+        return None
+    return str(dtype).split(".")[-1]

diffeval/backends/diffusers/from_config.py

@@ -0,0 +1,17 @@
+"""Require --pipeline and store backend settings."""
+
+
+def from_config(config: dict) -> dict:
+    pipeline = config.get("pipeline")
+    if not isinstance(pipeline, str) or not pipeline:
+        raise ValueError("diffusers backend requires --pipeline")
+    raw = config.get("batch_size", 1)
+    if type(raw) is not int or raw < 1:
+        raise ValueError("batch_size must be a positive int")
+    attention = config.get("attention")
+    if attention not in (None, "sdpa", "xformers", "eager"):
+        raise ValueError("attention must be sdpa, xformers, or eager")
+    keys = ("dtype", "steps", "guidance", "variant", "device")
+    stored = {key: config.get(key) for key in keys}
+    stored.update(pipeline=pipeline, batch_size=raw, attention=attention)
+    return stored

diffeval/backends/diffusers/generate.py

@@ -0,0 +1,18 @@
+"""Walk requests in batches."""
+
+from diffeval.backends.diffusers.one_batch import one_batch
+from diffeval.types.sample import Sample
+
+
+def _sample(request, image) -> Sample:
+    return Sample(request.prompt_id, request.prompt, request.seed, request.width, request.height, image, None, request.negative_prompt)
+
+
+def generate(pipe, requests, batch_size, steps, guidance):
+    samples = []
+    size = max(1, int(batch_size))
+    for start in range(0, len(requests), size):
+        chunk = list(requests[start:start + size])
+        images = one_batch(pipe, chunk, steps, guidance)
+        samples.extend(_sample(request, image) for request, image in zip(chunk, images))
+    return samples

diffeval/backends/diffusers/kwargs_generator.py

@@ -0,0 +1,6 @@
+"""One CPU generator per seed."""
+
+
+def generators(seeds):
+    import torch
+    return [torch.Generator(device="cpu").manual_seed(int(seed)) for seed in seeds]

diffeval/backends/diffusers/kwargs_prompt.py

@@ -0,0 +1,11 @@
+"""Prompt and negative prompt arguments."""
+
+
+def prompt_kwargs(requests) -> dict:
+    prompts = [item.prompt for item in requests]
+    body = {"prompt": prompts[0] if len(requests) == 1 else prompts}
+    negatives = [item.negative_prompt for item in requests]
+    if any(negatives):
+        filled = [text or "" for text in negatives]
+        body["negative_prompt"] = filled[0] if len(requests) == 1 else filled
+    return body

diffeval/backends/diffusers/kwargs_sampling.py

@@ -0,0 +1,18 @@
+"""Steps and guidance when the pipeline accepts them."""
+
+import inspect
+
+
+def _accepts(pipe, name: str) -> bool:
+    params = inspect.signature(pipe.__call__).parameters
+    varied = any(item.kind == inspect.Parameter.VAR_KEYWORD for item in params.values())
+    return name in params or varied
+
+
+def sampling_kwargs(pipe, steps, guidance) -> dict:
+    body = {}
+    if steps is not None and _accepts(pipe, "num_inference_steps"):
+        body["num_inference_steps"] = steps
+    if guidance is not None and _accepts(pipe, "guidance_scale"):
+        body["guidance_scale"] = guidance
+    return body

diffeval/backends/diffusers/kwargs_size.py

@@ -0,0 +1,18 @@
+"""Width and height when the pipeline accepts them."""
+
+import inspect
+
+
+def _accepts(pipe, name: str) -> bool:
+    params = inspect.signature(pipe.__call__).parameters
+    varied = any(item.kind == inspect.Parameter.VAR_KEYWORD for item in params.values())
+    return name in params or varied
+
+
+def size_kwargs(pipe, width, height) -> dict:
+    body = {}
+    if _accepts(pipe, "width"):
+        body["width"] = width
+    if _accepts(pipe, "height"):
+        body["height"] = height
+    return body

diffeval/backends/diffusers/load_kwargs.py

@@ -0,0 +1,10 @@
+"""from_pretrained keyword arguments."""
+
+
+def load_kwargs(config: dict, dtype) -> dict:
+    body = {}
+    if dtype is not None:
+        body["torch_dtype"] = dtype
+    if config.get("variant"):
+        body["variant"] = config["variant"]
+    return body

diffeval/backends/diffusers/load_pipe.py

@@ -0,0 +1,23 @@
+"""Load AutoPipelineForText2Image and move it to the device."""
+
+
+def load_pipe(pipeline_id, kwargs, device):
+    from diffusers import AutoPipelineForText2Image
+    return AutoPipelineForText2Image.from_pretrained(pipeline_id, **kwargs).to(device)
+
+
+def prepare_pipe(config: dict) -> dict:
+    from diffeval.backends.diffusers.attention_apply import apply_attention
+    from diffeval.backends.diffusers.coerce_dtype import coerce_dtype
+    from diffeval.backends.diffusers.dtype_name import label
+    from diffeval.backends.diffusers.load_kwargs import load_kwargs
+    from diffeval.backends.diffusers.missing_dep import require_diffusers
+    from diffeval.backends.diffusers.pick_device import pick_device
+    from diffeval.backends.diffusers.sampling import sampling
+    require_diffusers()
+    device = pick_device(config.get("device"))
+    dtype, name = coerce_dtype(config.get("dtype"), device)
+    pipe = load_pipe(config["pipeline"], load_kwargs(config, dtype), device)
+    seen = label(getattr(pipe, "dtype", None)) or name
+    attention = apply_attention(pipe, config.get("attention"))
+    return {"pipe": pipe, "device": device, "dtype_text": seen, "sampling": sampling(config), "attention": attention}

diffeval/backends/diffusers/missing_dep.py

@@ -0,0 +1,10 @@
+"""Error telling the user to install the extra."""
+
+
+def require_diffusers() -> None:
+    try:
+        import diffusers
+        import torch
+    except ImportError as exc:
+        raise RuntimeError("install the diffusers extra: pip install 'diffeval[diffusers]'") from exc
+    return diffusers, torch

diffeval/backends/diffusers/one_batch.py

@@ -0,0 +1,11 @@
+"""One pipeline call for a same-size batch."""
+
+from diffeval.backends.diffusers.call_kwargs import call_kwargs
+
+
+def one_batch(pipe, requests, steps, guidance):
+    if len({(item.width, item.height) for item in requests}) != 1:
+        raise ValueError("a batch must share width and height")
+    result = pipe(**call_kwargs(pipe, requests, steps, guidance))
+    images = result.images if hasattr(result, "images") else result[0]
+    return list(images)

diffeval/backends/diffusers/pick_device.py

@@ -0,0 +1,13 @@
+"""Pick cuda, mps, or cpu."""
+
+
+def pick_device(requested: str | None) -> str:
+    if requested:
+        return requested
+    import torch
+    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/backends/diffusers/presets.py

@@ -0,0 +1,8 @@
+"""Published steps and guidance by model id."""
+
+SCHNELL = "black-forest-labs/FLUX.1-schnell"
+PRESETS = {SCHNELL: {"steps": 4, "guidance": 0.0}}
+
+
+def preset_for(pipeline: str) -> dict:
+    return PRESETS.get(pipeline, {})

diffeval/backends/diffusers/sampling.py

@@ -0,0 +1,18 @@
+"""Resolve steps and guidance to cli, preset, or pipeline-default."""
+
+from diffeval.backends.diffusers.presets import preset_for
+
+
+def _one(cli_value, preset_value):
+    if cli_value is not None:
+        return cli_value, "cli"
+    if preset_value is not None:
+        return preset_value, "preset"
+    return None, "pipeline-default"
+
+
+def sampling(config: dict) -> dict:
+    found = preset_for(config["pipeline"])
+    steps, steps_source = _one(config.get("steps"), found.get("steps"))
+    guidance, guidance_source = _one(config.get("guidance"), found.get("guidance"))
+    return {"steps": steps, "guidance": guidance, "steps_source": steps_source, "guidance_source": guidance_source}

diffeval/backends/dummy/__init__.py

@@ -0,0 +1 @@
+"""Deterministic RGB images for CPU tests."""

diffeval/backends/dummy/backend.py

@@ -0,0 +1,29 @@
+"""Bind and register the dummy backend."""
+
+from diffeval.backends.dummy.describe import describe as describe_config
+from diffeval.backends.dummy.from_config import from_config as read_config
+from diffeval.backends.dummy.generate import generate as render_batch
+from diffeval.backends.noop_prepare import prepare as noop
+from diffeval.registry.register import register
+from diffeval.registry.tables import BACKENDS
+
+
+@register(BACKENDS, "dummy")
+class DummyBackend:
+    description = "Deterministic RGB images for CPU tests."
+
+    def __init__(self, config: dict):
+        self.config = config
+
+    def prepare(self) -> None:
+        noop()
+
+    def generate(self, requests):
+        return render_batch(requests, self.config["batch_size"])
+
+    def describe(self) -> dict:
+        return describe_config(self.config)
+
+    @staticmethod
+    def from_config(config: dict):
+        return DummyBackend(read_config(config))

diffeval/backends/dummy/box.py

@@ -0,0 +1,9 @@
+"""Seed-shifted box inside the image."""
+
+
+def box(seed: int, width: int, height: int) -> tuple[int, int, int, int]:
+    span_w = max(1, width // 4)
+    span_h = max(1, height // 4)
+    left = abs(seed) % max(1, width - span_w + 1)
+    top = abs(seed // 3) % max(1, height - span_h + 1)
+    return left, top, left + span_w, top + span_h

diffeval/backends/dummy/colors.py

@@ -0,0 +1,10 @@
+"""Background and foreground from a prompt hash."""
+
+import hashlib
+
+
+def colors(prompt: str) -> tuple[tuple[int, int, int], tuple[int, int, int]]:
+    digest = hashlib.sha256(prompt.encode("utf-8")).digest()
+    background = (digest[0], digest[1], digest[2])
+    foreground = ((digest[3] + 128) % 256, digest[4], digest[5])
+    return background, foreground

diffeval/backends/dummy/describe.py

@@ -0,0 +1,18 @@
+"""describe() for the dummy backend."""
+
+
+def describe(config: dict) -> dict:
+    return {
+        "name": "dummy",
+        "pipeline": config.get("pipeline"),
+        "dtype": config.get("dtype"),
+        "requested_dtype": config.get("dtype"),
+        "steps": config.get("steps"),
+        "guidance": config.get("guidance"),
+        "steps_source": "pipeline-default",
+        "guidance_source": "pipeline-default",
+        "batch_size": config["batch_size"],
+        "device": "cpu",
+        "attention_backend": "none",
+        "variant": config.get("variant"),
+    }

diffeval/backends/dummy/draw.py

@@ -0,0 +1,16 @@
+"""Draw a rectangle, ellipse, or triangle."""
+
+from PIL import ImageDraw
+
+
+def draw(image, shape: str, box, color):
+    pen = ImageDraw.Draw(image)
+    if shape == "ellipse":
+        pen.ellipse(box, fill=color)
+    elif shape == "triangle":
+        left, top, right, bottom = box
+        mid = (left + right) // 2
+        pen.polygon([(left, bottom), (mid, top), (right, bottom)], fill=color)
+    else:
+        pen.rectangle(box, fill=color)
+    return image

diffeval/backends/dummy/from_config.py

@@ -0,0 +1,14 @@
+"""Batch size and an unused pipeline id."""
+
+
+def from_config(config: dict) -> dict:
+    raw = config.get("batch_size", 1)
+    if type(raw) is not int or raw < 1:
+        raise ValueError("batch_size must be a positive int")
+    pipeline = config.get("pipeline")
+    if pipeline is not None and not isinstance(pipeline, str):
+        raise ValueError("pipeline must be a string")
+    kept = ("dtype", "steps", "guidance", "variant")
+    stored = {key: config.get(key) for key in kept}
+    stored.update(batch_size=raw, pipeline=pipeline)
+    return stored

diffeval/backends/dummy/generate.py

@@ -0,0 +1,17 @@
+"""Generate samples in batch-size slices."""
+
+from diffeval.backends.dummy.render import render
+from diffeval.types.sample import Sample
+
+
+def _one(request) -> Sample:
+    image = render(request.prompt, request.seed, request.width, request.height)
+    return Sample(request.prompt_id, request.prompt, request.seed, request.width, request.height, image, None, request.negative_prompt)
+
+
+def generate(requests, batch_size: int) -> list:
+    size = max(1, batch_size)
+    samples = []
+    for start in range(0, len(requests), size):
+        samples.extend(_one(item) for item in requests[start:start + size])
+    return samples

diffeval/backends/dummy/render.py

@@ -0,0 +1,18 @@
+"""One RGB image for a prompt and a seed."""
+
+import hashlib
+
+from PIL import Image
+
+from diffeval.backends.dummy.box import box
+from diffeval.backends.dummy.colors import colors
+from diffeval.backends.dummy.draw import draw
+
+_SHAPES = ("rectangle", "ellipse", "triangle")
+
+
+def render(prompt: str, seed: int, width: int, height: int) -> Image.Image:
+    background, foreground = colors(prompt)
+    image = Image.new("RGB", (width, height), background)
+    kind = _SHAPES[hashlib.sha256(prompt.encode("utf-8")).digest()[6] % 3]
+    return draw(image, kind, box(seed, width, height), foreground)

diffeval/backends/noop_prepare.py

@@ -0,0 +1,5 @@
+"""prepare that does nothing."""
+
+
+def prepare() -> None:
+    return None

diffeval/backends/protocol.py

@@ -0,0 +1,15 @@
+"""Backend contract."""
+
+from typing import Protocol
+
+from diffeval.types.request import GenerationRequest
+from diffeval.types.sample import Sample
+
+
+class Backend(Protocol):
+    description: str
+
+    def from_config(self, config: dict) -> "Backend": ...
+    def prepare(self) -> None: ...
+    def generate(self, requests: list[GenerationRequest]) -> list[Sample]: ...
+    def describe(self) -> dict: ...

diffeval/bench/__init__.py

@@ -0,0 +1 @@
+"""Benchmark helpers: the CMMD check, the COCO-30k suite, and the results table."""