argus-proof 0.2.0__py3-none-any.whl

argus_proof/__init__.py

@@ -0,0 +1,16 @@
+"""argus-proof — Post-training LoRA evaluation and optimisation: generate samples from a trained LoRA and score them against the curated dataset it was trained from"""
+
+from __future__ import annotations
+
+try:
+    # Written by hatch-vcs at build time (see pyproject [tool.hatch.build.hooks.vcs]).
+    from argus_proof._version import __version__
+except ImportError:  # running from a source checkout that hasn't been built
+    from importlib.metadata import PackageNotFoundError, version
+
+    try:
+        __version__ = version("argus-proof")
+    except PackageNotFoundError:
+        __version__ = "0.0.0+unknown"
+
+__all__ = ["__version__"]

argus_proof/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.2.0'
+__version_tuple__ = version_tuple = (0, 2, 0)
+
+__commit_id__ = commit_id = None

argus_proof/acceptance.py

@@ -0,0 +1,118 @@
+"""The CI acceptance gate (#12): a machine-readable pass/fail on an EvalReport.
+
+Closes the loop for CI — "was this dataset/LoRA good enough?" becomes an
+automatable yes/no against declared :class:`~argus_proof.models.AcceptanceThresholds`
+(a min pass-rate, a statistically-defensible pass-rate lower bound, an identity
+floor, a safety ceiling). Only the configured checks run; the run is accepted
+iff every check passes.
+"""
+
+from __future__ import annotations
+
+from argus_proof.models import (
+    AcceptanceThresholds,
+    EvalReport,
+    GateCheck,
+    GateResult,
+)
+from argus_proof.stats import wilson_lower_bound
+
+
+def _unsafe_rate(report: EvalReport, safety_floor: float) -> float:
+    """Fraction of images that are unsafe.
+
+    An image counts as unsafe if its ``safety`` metric is below *safety_floor*
+    OR it carries an ``unsafe`` reject reason. Keying off the metric (not only
+    the reject reason) matters: the ``unsafe`` reason is only produced when the
+    gate has a ``safety`` hard gate, so a run scored with the default gate could
+    otherwise report a 0% unsafe rate while containing unsafe imagery.
+    """
+    n = len(report.images)
+    if n == 0:
+        return 0.0
+    unsafe = sum(
+        1
+        for img in report.images
+        if (img.metrics.safety is not None and img.metrics.safety < safety_floor)
+        or any(r.code == "unsafe" for r in img.reject_reasons)
+    )
+    return unsafe / n
+
+
+def evaluate_acceptance(report: EvalReport, thresholds: AcceptanceThresholds) -> GateResult:
+    """Evaluate *report* against *thresholds* → a :class:`GateResult`.
+
+    A configured metric that wasn't measured (e.g. ``min_identity_mean`` with no
+    identity scorer) fails its check rather than passing silently — you can't
+    accept on evidence you don't have.
+    """
+    checks: list[GateCheck] = []
+    agg = report.aggregate
+    # The pass-rate is over near-dup groups when deduped, else over images.
+    n = agg.n_groups if agg.n_groups is not None else agg.n_images
+    n_passed = min(agg.n_passed, n)  # defensive: never let successes exceed the denominator
+
+    if thresholds.min_pass_rate is not None:
+        ok = agg.pass_rate >= thresholds.min_pass_rate
+        checks.append(
+            GateCheck(
+                name="pass_rate",
+                passed=ok,
+                actual=agg.pass_rate,
+                threshold=thresholds.min_pass_rate,
+                detail=f"pass_rate {agg.pass_rate:.3f} {'>=' if ok else '<'} {thresholds.min_pass_rate:.3f}",
+            )
+        )
+
+    if thresholds.min_pass_rate_ci_lower is not None:
+        lb = wilson_lower_bound(n_passed, n, thresholds.confidence)
+        ok = lb >= thresholds.min_pass_rate_ci_lower
+        checks.append(
+            GateCheck(
+                name="pass_rate_ci_lower",
+                passed=ok,
+                actual=lb,
+                threshold=thresholds.min_pass_rate_ci_lower,
+                detail=(
+                    f"{thresholds.confidence:.0%} Wilson lower bound {lb:.3f} "
+                    f"({n_passed}/{n}) {'>=' if ok else '<'} {thresholds.min_pass_rate_ci_lower:.3f}"
+                ),
+            )
+        )
+
+    if thresholds.min_identity_mean is not None:
+        val = agg.means.identity
+        ok = val is not None and val >= thresholds.min_identity_mean
+        detail = (
+            "identity not measured"
+            if val is None
+            else f"identity mean {val:.3f} {'>=' if ok else '<'} {thresholds.min_identity_mean:.3f}"
+        )
+        checks.append(
+            GateCheck(
+                name="identity_mean", passed=ok, actual=val, threshold=thresholds.min_identity_mean, detail=detail
+            )
+        )
+
+    if thresholds.max_unsafe_rate is not None:
+        rate = _unsafe_rate(report, thresholds.unsafe_safety_floor)
+        ok = rate <= thresholds.max_unsafe_rate
+        checks.append(
+            GateCheck(
+                name="unsafe_rate",
+                passed=ok,
+                actual=rate,
+                threshold=thresholds.max_unsafe_rate,
+                detail=f"unsafe rate {rate:.3f} {'<=' if ok else '>'} {thresholds.max_unsafe_rate:.3f}",
+            )
+        )
+
+    if not checks:
+        # No thresholds configured -> nothing was verified. Refuse rather than
+        # accept on zero evidence (all([]) would otherwise be True).
+        return GateResult(
+            passed=False,
+            checks=[GateCheck(name="no_checks", passed=False, detail="no acceptance thresholds configured")],
+        )
+
+    return GateResult(passed=all(c.passed for c in checks), checks=checks)

argus_proof/archive.py

@@ -0,0 +1,168 @@
+"""Metadata-only archival of a reviewed run (#9).
+
+Keep the analytical value of a run without hoarding imagery: after review, a run
+directory is reduced to
+
+* ``manifest.json`` + ``eval_report.json`` — the run's params and scores;
+* ``passing_images.zip`` — the images that passed, retained zipped;
+* ``reject_archive.json`` — a :class:`~argus_proof.models.RejectArchive`: params
+  (the embedded :class:`~argus_proof.models.RunManifest`) + scores + structured
+  reject reasons + rater id for every rejected image, and **no image or
+  thumbnail reference** (explicit product decision — the metadata is the useful
+  signal for diagnosing failure modes and could later train an auto-classifier).
+
+With ``blank_slate=True`` the loose per-image files that were archived (passing →
+zipped, rejected → recorded) are deleted, leaving only the artifacts above.
+Images still awaiting review (``passed is None``) are left untouched. Every
+record joins back to its run by ``run_id``, so a reject is traceable to its exact
+config without the pixels.
+"""
+
+from __future__ import annotations
+
+import zipfile
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import structlog
+
+from argus_proof.models import (
+    EvalReport,
+    GeneratedImage,
+    RejectArchive,
+    RejectRecord,
+    RunManifest,
+)
+
+logger = structlog.get_logger()
+
+# Version of the on-disk archive layout (the set/naming of artifacts below).
+# Bump if the directory contract changes so consumers can refuse an old layout.
+ARCHIVE_VERSION = "1.0"
+
+MANIFEST_NAME = "manifest.json"
+REPORT_NAME = "eval_report.json"
+PASSING_ZIP_NAME = "passing_images.zip"
+REJECT_ARCHIVE_NAME = "reject_archive.json"
+
+
+@dataclass
+class ArchiveResult:
+    """What :func:`archive_run` produced: counts + the artifact paths."""
+
+    run_id: str
+    n_passing: int
+    n_rejected: int
+    n_pending: int  # images left for review (passed is None)
+    passing_zip: Path | None
+    reject_archive_path: Path
+    deleted_images: list[str] = field(default_factory=list)
+
+
+def build_reject_archive(report: EvalReport, manifest: RunManifest, *, rater_id: str | None = None) -> RejectArchive:
+    """A :class:`RejectArchive` of the report's rejected images (``passed is False``).
+
+    Embeds *manifest* (keyed by ``run_id``) so the archive is self-describing, and
+    stamps ``rater_id`` on each record when a human reviewed the run. Carries zero
+    image references by construction.
+    """
+    records = [
+        RejectRecord(
+            run_id=report.run_id,
+            seed=img.seed,
+            metrics=img.metrics,
+            hitl_rating=img.hitl_rating,
+            reasons=img.reject_reasons,
+            rater_id=rater_id,
+        )
+        for img in report.images
+        if img.passed is False
+    ]
+    return RejectArchive(manifests={manifest.run_id: manifest}, records=records)
+
+
+def archive_run(
+    run_dir: Path,
+    images: list[GeneratedImage],
+    report: EvalReport,
+    manifest: RunManifest,
+    *,
+    rater_id: str | None = None,
+    blank_slate: bool = True,
+) -> ArchiveResult:
+    """Reduce a reviewed *run_dir* to metadata + a zip of the passing images.
+
+    *images* supplies the on-disk paths (joined to *report* by ``image_id``).
+    Writes ``manifest.json``, ``eval_report.json``, ``passing_images.zip``, and
+    ``reject_archive.json``; with ``blank_slate`` deletes the loose image files
+    that were archived (passing + rejected), leaving pending images in place.
+    """
+    run_dir.mkdir(parents=True, exist_ok=True)
+    verdict_by_id = {img.image_id: img.passed for img in report.images}
+    by_id = {img.image_id: img for img in images}
+
+    (run_dir / MANIFEST_NAME).write_text(manifest.model_dump_json(indent=2), encoding="utf-8")
+    (run_dir / REPORT_NAME).write_text(report.model_dump_json(indent=2), encoding="utf-8")
+
+    reject_archive = build_reject_archive(report, manifest, rater_id=rater_id)
+    reject_path = run_dir / REJECT_ARCHIVE_NAME
+    reject_path.write_text(reject_archive.model_dump_json(indent=2), encoding="utf-8")
+
+    # Warn about report images with no matching file — they can't be archived.
+    missing = [i.image_id for i in report.images if i.passed is not None and i.image_id not in by_id]
+    if missing:
+        logger.warning("archive.missing_images", run_id=report.run_id, image_ids=missing)
+
+    # Zip the passing images (that still exist on disk). Arcname is the unique
+    # image_id (+ suffix) so two passing images that share a basename in
+    # different dirs don't collide and silently drop one.
+    passing = [img for img in images if verdict_by_id.get(img.image_id) is True]
+    passing_zip: Path | None = None
+    if passing:
+        passing_zip = run_dir / PASSING_ZIP_NAME
+        with zipfile.ZipFile(passing_zip, "w", zipfile.ZIP_DEFLATED) as zf:
+            for img in passing:
+                src = Path(img.path)
+                if src.is_file():
+                    zf.write(src, arcname=f"{img.image_id}{src.suffix}")
+
+    run_resolved = run_dir.resolve()
+    n_pending = sum(1 for img in report.images if img.passed is None)
+    deleted: list[str] = []
+    if blank_slate:
+        # Delete the loose files that are now archived: passing (in the zip) and
+        # rejected (recorded as metadata). Pending images are left for review.
+        # Only touch files inside run_dir — never delete an image the run doesn't own.
+        for image_id, passed in verdict_by_id.items():
+            if passed is None:
+                continue  # not yet reviewed
+            img = by_id.get(image_id)
+            if img is None:
+                continue
+            src = Path(img.path)
+            if not src.is_file():
+                continue
+            if not src.resolve().is_relative_to(run_resolved):
+                logger.warning("archive.skip_delete_outside_run_dir", run_id=report.run_id, path=str(src))
+                continue
+            src.unlink()
+            deleted.append(src.name)
+
+    result = ArchiveResult(
+        run_id=report.run_id,
+        n_passing=len(passing),
+        n_rejected=len(reject_archive.records),
+        n_pending=n_pending,
+        passing_zip=passing_zip,
+        reject_archive_path=reject_path,
+        deleted_images=deleted,
+    )
+    logger.info(
+        "archive.run",
+        run_id=result.run_id,
+        passing=result.n_passing,
+        rejected=result.n_rejected,
+        pending=result.n_pending,
+        blank_slate=blank_slate,
+    )
+    return result

argus_proof/backends/__init__.py

@@ -0,0 +1,67 @@
+"""Generation backends — one module per engine, selected by name (config).
+
+Swapping the engine that generates the eval grid is a config change, not a code
+change: a caller names a backend and constructs it through :func:`get_backend`,
+and scoring/report code is unchanged regardless of which one produced the run.
+
+* ``comfyui`` — a running ComfyUI instance (the first adapter).
+* ``diffusers`` — in-process diffusers SDXL pipeline (deterministic, no service).
+* ``a1111`` — an AUTOMATIC1111 / SD.Next ``/sdapi`` server.
+* ``remote`` — a hosted/cloud endpoint that speaks the proof wire.
+"""
+
+from __future__ import annotations
+
+from argus_proof.backends.a1111 import A1111Backend
+from argus_proof.backends.base import (
+    BackendError,
+    GenBackend,
+    GenResult,
+    ModelResolver,
+    ProgressSink,
+    build_local_manifest,
+    hash_model,
+)
+from argus_proof.backends.comfyui import ComfyUIBackend
+from argus_proof.backends.diffusers import DiffusersBackend
+from argus_proof.backends.remote import RemoteBackend
+
+# The backends this build knows how to construct, by name.
+_BACKENDS: dict[str, type] = {
+    "comfyui": ComfyUIBackend,
+    "diffusers": DiffusersBackend,
+    "a1111": A1111Backend,
+    "remote": RemoteBackend,
+}
+KNOWN_BACKENDS: tuple[str, ...] = tuple(_BACKENDS)
+
+
+def get_backend(name: str, **kwargs: object) -> GenBackend:
+    """Construct the backend called *name*, forwarding ``kwargs`` to it.
+
+    Raises :class:`BackendError` for an unknown name so a bad config fails
+    loudly instead of silently doing nothing.
+    """
+    try:
+        backend_cls = _BACKENDS[name]
+    except KeyError:
+        known = ", ".join(KNOWN_BACKENDS)
+        raise BackendError(f"unknown generation backend {name!r} (known: {known})") from None
+    return backend_cls(**kwargs)  # type: ignore[arg-type]
+
+
+__all__ = [
+    "KNOWN_BACKENDS",
+    "A1111Backend",
+    "BackendError",
+    "ComfyUIBackend",
+    "DiffusersBackend",
+    "GenBackend",
+    "GenResult",
+    "ModelResolver",
+    "ProgressSink",
+    "RemoteBackend",
+    "build_local_manifest",
+    "get_backend",
+    "hash_model",
+]

argus_proof/backends/a1111.py

@@ -0,0 +1,150 @@
+"""AUTOMATIC1111 / SD.Next generation backend — drives their ``/sdapi`` HTTP API.
+
+Submits a ``txt2img`` request per seed, decodes the base64 PNG(s) it returns, and
+emits a reproducible :class:`~argus_proof.models.RunManifest`. The base checkpoint
+is selected per-request via ``override_settings`` and LoRA(s) via the A1111
+``<lora:name:weight>`` prompt syntax. Model files are SHA256-pinned from disk (via
+``resolve_model``), so this targets an A1111/SD.Next whose models are locally
+resolvable — the manifest reconstructs the run exactly.
+
+The HTTP layer is the shared injectable :class:`~argus_proof.backends.http.Transport`,
+so the adapter is unit-testable without a live A1111.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import structlog
+
+from argus_proof.backends.base import (
+    BackendError,
+    GenResult,
+    ModelResolver,
+    ProgressSink,
+    build_local_manifest,
+    write_manifest,
+)
+from argus_proof.backends.http import Transport, UrllibTransport, decode_base64_image
+from argus_proof.backends.pnginfo import read_dimensions
+from argus_proof.models import BackendCapabilities, GeneratedImage, ProgressEvent, RunSpec
+
+logger = structlog.get_logger()
+
+BACKEND_NAME = "a1111"
+DEFAULT_BASE_URL = "http://127.0.0.1:7860"
+
+
+class A1111Backend:
+    """Generate images through an AUTOMATIC1111 / SD.Next ``/sdapi`` server.
+
+    ``resolve_model`` hashes the base checkpoint + LoRA files into the manifest;
+    ``transport`` defaults to :class:`UrllibTransport` against ``base_url`` and is
+    injectable for testing. ``engine_version`` is recorded on the manifest (A1111
+    exposes no portable version endpoint); pass it to pin the server build.
+    """
+
+    def __init__(
+        self,
+        resolve_model: ModelResolver,
+        base_url: str = DEFAULT_BASE_URL,
+        transport: Transport | None = None,
+        engine_version: str | None = None,
+        timeout: float = 600.0,
+    ) -> None:
+        self.resolve_model = resolve_model
+        self.transport = transport or UrllibTransport(base_url, timeout=timeout, label="A1111")
+        self._engine_version = engine_version or "unknown"
+
+    def capabilities(self) -> BackendCapabilities:
+        return BackendCapabilities(
+            name=BACKEND_NAME,
+            supports_seed_set=True,
+            max_loras=None,
+            reads_pnginfo=False,
+            streams_progress=True,
+        )
+
+    def generate(self, spec: RunSpec, out_dir: Path, progress: ProgressSink | None = None) -> GenResult:
+        """Run *spec*, writing one image per seed (+ manifest.json) under *out_dir*."""
+        logger.debug("a1111.generate", run_id=spec.run_id, seeds=len(spec.seeds), out_dir=str(out_dir))
+        out_dir.mkdir(parents=True, exist_ok=True)
+
+        def emit(event: ProgressEvent) -> None:
+            if progress is not None:
+                progress(event)
+
+        total = len(spec.seeds)
+        emit(ProgressEvent(run_id=spec.run_id, type="start", total=total))
+
+        images: list[GeneratedImage] = []
+        try:
+            manifest = build_local_manifest(
+                spec, resolve_model=self.resolve_model, engine=BACKEND_NAME, engine_version=self._engine_version
+            )
+            for done, seed in enumerate(spec.seeds, start=1):
+                images.extend(self._generate_seed(spec, seed, out_dir, emit))
+                emit(ProgressEvent(run_id=spec.run_id, type="progress", completed=done, total=total))
+        except BackendError as exc:
+            emit(ProgressEvent(run_id=spec.run_id, type="error", message=str(exc)))
+            raise
+
+        write_manifest(out_dir, manifest)
+        emit(ProgressEvent(run_id=spec.run_id, type="done", completed=total, total=total))
+        return GenResult(manifest=manifest, images=images)
+
+    def _generate_seed(self, spec: RunSpec, seed: int, out_dir: Path, emit: ProgressSink) -> list[GeneratedImage]:
+        resp = self.transport.post_json("/sdapi/v1/txt2img", self._payload(spec, seed))
+        encoded = resp.get("images") or []
+        if not encoded:
+            raise BackendError(f"A1111 produced no images for seed {seed}")
+        produced: list[GeneratedImage] = []
+        for index, b64 in enumerate(encoded):
+            data = decode_base64_image(b64, context=f"A1111 seed {seed}")
+            suffix = "" if index == 0 else f"-{index}"
+            image_id = f"{spec.run_id}-{seed}{suffix}"
+            path = out_dir / f"{image_id}.png"
+            path.write_bytes(data)
+            dims = read_dimensions(data) or (spec.sampling.width, spec.sampling.height)
+            produced.append(
+                GeneratedImage(
+                    image_id=image_id,
+                    run_id=spec.run_id,
+                    seed=seed,
+                    path=str(path),
+                    width=dims[0],
+                    height=dims[1],
+                )
+            )
+            emit(ProgressEvent(run_id=spec.run_id, type="image", seed=seed, image_id=image_id))
+        return produced
+
+    def _payload(self, spec: RunSpec, seed: int) -> dict:
+        # A1111 applies LoRAs via the prompt: "<lora:name:weight>" — keep any
+        # subdirectory (with_suffix drops only the extension) so a nested LoRA
+        # (loras/char/subject.safetensors) resolves as "<lora:char/subject:...>".
+        lora_tags = "".join(f" <lora:{Path(lo.name).with_suffix('')}:{lo.weight}>" for lo in spec.loras)
+        # Pin the checkpoint, VAE, and clip-skip via override_settings so the run
+        # matches the manifest; restore_afterwards=False keeps the checkpoint loaded
+        # across the run's seeds instead of reloading it per request.
+        override: dict = {
+            "sd_model_checkpoint": spec.base_checkpoint,
+            "CLIP_stop_at_last_layers": spec.sampling.clip_skip,
+        }
+        if spec.vae:
+            override["sd_vae"] = spec.vae
+        return {
+            "prompt": spec.prompt + lora_tags,
+            "negative_prompt": spec.negative_prompt,
+            "steps": spec.sampling.steps,
+            "cfg_scale": spec.sampling.cfg,
+            "width": spec.sampling.width,
+            "height": spec.sampling.height,
+            "seed": seed,
+            "sampler_name": spec.sampling.sampler,
+            "scheduler": spec.sampling.scheduler,
+            "batch_size": 1,
+            "n_iter": 1,
+            "override_settings": override,
+            "override_settings_restore_afterwards": False,
+        }