eval_harness.py

#!/usr/bin/env python3
"""
eval_harness.py — ImageForge Eval Harness
==========================================
Drives the live FastAPI /generate endpoint, collects images, and computes
clean-FID + HPSv2 + GenEval2 (Soft-TIFA) for two model keys (base vs candidate).

Usage
-----
  # Base vs base (noise floor, different seeds):
  python eval_harness.py --base sd-turbo --candidate sd-turbo --noise-floor

  # Base vs LoRA candidate:
  python eval_harness.py --base sd-turbo --candidate sdxl-turbo

  # Skip GenEval2 if Qwen3-VL is not available locally:
  python eval_harness.py --base sd-turbo --candidate sdxl-turbo --skip-geneval

  # Run FID+HPS only (fastest):
  python eval_harness.py --base sd-turbo --candidate sdxl-turbo --skip-geneval

Requirements
------------
Install on the Mac (MPS box) — the serving environment:
  pip install clean-fid hpsv2 httpx Pillow tqdm

GenEval2 (runs on CUDA box — Qwen3-VL-8B needs ~19 GB VRAM at FP16, ~5 GB at INT4):
  git clone https://github.com/facebookresearch/GenEval2
  cd GenEval2
  pip install torch transformers==4.57.0 pillow tqdm scipy
  # For INT4 on 4090: pip install bitsandbytes
  # Copy geneval2_data.jsonl from the repo into this directory or pass --geneval2-data.

API
---
Drives http://127.0.0.1:8765/generate with GenerateRequest schema:
  {
    "prompt": str, "model": str|null, "steps": int|null, "guidance": float|null,
    "width": int|null, "height": int|null, "seed": int,
    "return_mode": "base64"
  }
Response: GenerationResponse.image_base64 (PNG, base64).
Seeds are pinned per-prompt for reproducibility; noise-floor run uses seed+1 offsets.
The engine uses torch.Generator(device="cpu").manual_seed(seed) — deterministic across runs.
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import io
import json
import logging
import os
import shutil
import sys
import time
from pathlib import Path
from typing import Optional

# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)-8s  %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("eval")

# ---------------------------------------------------------------------------
# Constants / thresholds
# ---------------------------------------------------------------------------
API_URL = "http://127.0.0.1:8765"
GENERATE_ENDPOINT = f"{API_URL}/generate"

# Pass/fail gate thresholds (see emit_report for the gate model).
#
# Gate model (redesigned):
#   * HPSv2 is the PRIMARY gate — higher is better, always measurable, so the
#     gate always rests on it. Delta ≥ +0.20 AND > 2×noise_std is pass.
#   * clean-FID is a CORROBORATING signal computed ONLY against a real-image
#     reference set (--reference-dir): fid_base = FID(ref, base),
#     fid_cand = FID(ref, cand); delta = fid_cand - fid_base must be ≤ -5 (and
#     beyond the noise floor) — a LOWER candidate FID is the improvement.
#     With no reference set FID is N/A and does not gate (never fabricated).
#   * GenEval2 Soft-TIFA_AM: higher is better. Delta ≥ +2.0 pp AND > 2×noise_std.
FID_PASS_DELTA = -5.0          # candidate FID(ref,cand) lower than base FID(ref,base) by ≥5 → corroborates
HPS_PASS_DELTA = 0.20          # candidate HPS higher by ≥0.20 points (100-scale) → pass
GENEVAL_PASS_DELTA = 2.0       # candidate Soft-TIFA_AM higher by ≥2.0 pp → pass
NOISE_SIGMA_MULTIPLIER = 2.0   # delta must also exceed 2×noise_std (per metric)

# Number of images per model run. 200 images gives ±0.3 FID CI vs ±0.2 at N=1000.
# For a LoRA vs base comparison on a held-out prompt set, N=200 is adequate.
DEFAULT_N_IMAGES = 200  # use --n-images to override

# Default resolution — matches the sd-turbo/sdxl-turbo defaults in the registry.
DEFAULT_WIDTH = 512
DEFAULT_HEIGHT = 512

# ---------------------------------------------------------------------------
# Prompt set loader
# ---------------------------------------------------------------------------
PROMPTS_FILE = Path(__file__).parent / "prompts_heldout.jsonl"


def load_prompts(path: Path) -> list[dict]:
    """Load JSONL prompt set. Each line: {"prompt": str, "seed": int, ...}"""
    if not path.exists():
        raise FileNotFoundError(
            f"Prompt file not found: {path}\n"
            "Create prompts_heldout.jsonl next to eval_harness.py, or pass --prompts."
        )
    records = []
    with path.open() as fh:
        for line in fh:
            line = line.strip()
            if line and not line.startswith("//"):
                records.append(json.loads(line))
    log.info("Loaded %d prompts from %s", len(records), path)
    return records


# ---------------------------------------------------------------------------
# Image generation driver
# ---------------------------------------------------------------------------
def generate_images(
    prompts: list[dict],
    model: Optional[str],
    out_dir: Path,
    *,
    seed_offset: int = 0,
    steps: Optional[int] = None,
    guidance: Optional[float] = None,
    width: int = DEFAULT_WIDTH,
    height: int = DEFAULT_HEIGHT,
    n_images: int = DEFAULT_N_IMAGES,
) -> list[Path]:
    """
    POST each prompt to /generate, save returned PNG to out_dir.

    Respects the engine's single-inference-lock: sends one request at a time
    (no concurrency) with a 300s per-image timeout.

    Seeds are pinned per-prompt (prompt["seed"] + seed_offset) for reproducibility.
    The engine uses torch.Generator(device="cpu").manual_seed(seed) which is
    deterministic across MPS restarts (pipeline.py:573).

    Returns list of saved image paths.
    """
    try:
        import httpx
    except ImportError:
        sys.exit("httpx not installed. Run: pip install httpx")

    out_dir.mkdir(parents=True, exist_ok=True)
    saved: list[Path] = []

    # Tile prompts up to n_images if fewer prompts than requested
    tiled: list[dict] = []
    while len(tiled) < n_images:
        tiled.extend(prompts)
    tiled = tiled[:n_images]

    label = model or "default"
    log.info("Generating %d images with model=%s → %s", len(tiled), label, out_dir)

    with httpx.Client(base_url=API_URL, timeout=300.0) as client:
        # Health check first — fail fast rather than timeout N times
        try:
            health = client.get("/health").json()
            log.info("Engine health: status=%s device=%s", health.get("status"), health.get("device"))
        except Exception as e:
            sys.exit(f"Cannot reach ImageForge at {API_URL}. Start with: ./run_api.sh\nError: {e}")

        for idx, rec in enumerate(tiled):
            seed = (rec.get("seed", idx * 1000) + seed_offset) & 0xFFFFFFFF
            payload: dict = {
                "prompt": rec["prompt"],
                "seed": seed,
                "width": width,
                "height": height,
                "return_mode": "base64",
                "use_cache": False,  # always generate fresh for eval
            }
            if model:
                payload["model"] = model
            if steps is not None:
                payload["steps"] = steps
            if guidance is not None:
                payload["guidance"] = guidance

            out_file = out_dir / f"{idx:04d}_seed{seed}.png"
            if out_file.exists():
                log.debug("[%d/%d] cached %s", idx + 1, len(tiled), out_file.name)
                saved.append(out_file)
                continue

            t0 = time.perf_counter()
            try:
                resp = client.post("/generate", json=payload)
                resp.raise_for_status()
            except httpx.HTTPStatusError as e:
                log.error("[%d/%d] HTTP %s: %s", idx + 1, len(tiled), e.response.status_code, e.response.text[:200])
                continue
            except Exception as e:
                log.error("[%d/%d] request failed: %s", idx + 1, len(tiled), e)
                continue

            data = resp.json()
            b64 = data.get("image_base64")
            if not b64:
                log.warning("[%d/%d] no image_base64 in response: %s", idx + 1, len(tiled), list(data.keys()))
                continue

            img_bytes = base64.b64decode(b64)
            out_file.write_bytes(img_bytes)
            elapsed = time.perf_counter() - t0
            log.info("[%d/%d] seed=%d %.1fs → %s", idx + 1, len(tiled), seed, elapsed, out_file.name)
            saved.append(out_file)

    log.info("Generated %d/%d images", len(saved), len(tiled))
    return saved


# ---------------------------------------------------------------------------
# clean-FID scoring (on-Mac, CPU or MPS via torch device override)
# ---------------------------------------------------------------------------
def score_fid(dir_base: Path, dir_candidate: Path, *, device_str: str = "cpu") -> float:
    """
    Compute clean-FID between two image directories.

    Uses mode='clean' (PIL-bicubic resizing) to avoid the Parmar 2022 resize bug
    where OpenCV/PyTorch bicubic produced FID ≥6 offset vs correct PIL-bicubic.

    device_str: 'cpu' for Mac MPS box (Inception runs fine on CPU, ~2min for N=200).
                'cuda' for the CUDA training box.
    """
    try:
        from cleanfid import fid
    except ImportError:
        sys.exit("clean-fid not installed. Run: pip install clean-fid")

    import torch

    device = torch.device(device_str)
    log.info("Computing FID: %s vs %s (device=%s) …", dir_base.name, dir_candidate.name, device_str)
    score = fid.compute_fid(
        str(dir_base),
        str(dir_candidate),
        mode="clean",          # correct PIL-bicubic resize — never legacy_pytorch/legacy_tensorflow
        model_name="inception_v3",
        num_workers=0,         # 0 = no multiprocessing (avoids fork issues on macOS)
        batch_size=32,
        device=device,
        verbose=True,
    )
    return float(score)


# ---------------------------------------------------------------------------
# HPSv2 scoring (on-Mac, CPU or MPS)
# ---------------------------------------------------------------------------
def score_hps(images: list[Path], prompts: list[dict], *, hps_version: str = "v2.1") -> float:
    """
    Compute mean HPSv2 score for a list of images against their prompts.

    hpsv2.score(img, prompt) accepts: PIL.Image, file path str, or list.
    Returns a logit score; typical range for SDXL-class models: 26–30 on 100-scale.
    We use v2.1 (released Sep 2024, trained on higher-quality HPD v2.1 dataset).

    Runs on CPU perfectly; MPS also works via standard PyTorch dispatch.
    The reward model is a CLIP-ViT-H variant (~900MB download on first run).
    """
    try:
        import hpsv2
    except ImportError:
        sys.exit("hpsv2 not installed. Run: pip install hpsv2")

    if len(images) != len(prompts):
        # Tile prompts to match images (same tiling as generate_images)
        tiled_prompts: list[dict] = []
        while len(tiled_prompts) < len(images):
            tiled_prompts.extend(prompts)
        tiled_prompts = tiled_prompts[:len(images)]
    else:
        tiled_prompts = list(prompts)

    scores: list[float] = []
    log.info("Scoring %d images with HPSv2.%s …", len(images), hps_version)
    for img_path, rec in zip(images, tiled_prompts):
        try:
            result = hpsv2.score(str(img_path), rec["prompt"], hps_version=hps_version)
            # hpsv2.score returns a list when given a list of images; a scalar for a single path
            val = result[0] if isinstance(result, (list, tuple)) else float(result)
            scores.append(float(val))
        except Exception as e:
            log.warning("HPSv2 error on %s: %s", img_path.name, e)

    if not scores:
        return float("nan")
    mean_score = sum(scores) / len(scores)
    log.info("HPSv2 mean=%.4f (N=%d)", mean_score, len(scores))
    return mean_score


# ---------------------------------------------------------------------------
# GenEval2 Soft-TIFA scoring (run on CUDA box via SSH or local if 4090 present)
# ---------------------------------------------------------------------------
def build_geneval_image_map(images: list[Path], prompts: list[dict], run_dir: Path) -> Path:
    """
    Build the image_filepath_data.json that GenEval2's evaluation.py expects.
    Format: {"<prompt_text>": "<absolute_path_to_image>"}
    """
    if len(images) != len(prompts):
        tiled_prompts: list[dict] = []
        while len(tiled_prompts) < len(images):
            tiled_prompts.extend(prompts)
        tiled_prompts = tiled_prompts[:len(images)]
    else:
        tiled_prompts = list(prompts)

    mapping: dict[str, str] = {}
    for img_path, rec in zip(images, tiled_prompts):
        prompt_text = rec["prompt"]
        # GenEval2 uses prompt as key; use last image if duplicated prompts
        mapping[prompt_text] = str(img_path.resolve())

    out_path = run_dir / "geneval_image_map.json"
    out_path.write_text(json.dumps(mapping, indent=2))
    log.info("Wrote GenEval2 image map → %s (%d entries)", out_path, len(mapping))
    return out_path


def score_geneval2(
    image_map_path: Path,
    geneval2_repo: str,
    *,
    method: str = "soft_tifa_am",
    geneval2_data: Optional[str] = None,
    output_file: Optional[Path] = None,
) -> float:
    """
    Run GenEval2 evaluation.py as a subprocess.

    geneval2_repo: path to the cloned facebookresearch/GenEval2 directory.
    method: 'soft_tifa_am' (atom-level arithmetic mean, recommended for LoRA delta).
            'soft_tifa_gm' is more conservative (geometric mean, sensitive to zero atoms).
    Returns: aggregate score (0–100).

    NOTE: GenEval2 requires Qwen3-VL-8B-Instruct (~19GB VRAM at FP16, ~5GB at INT4).
    Run on the RTX 4090 (CUDA box). Pass --skip-geneval on the Mac if not available.
    """
    import subprocess

    repo = Path(geneval2_repo).resolve()
    eval_script = repo / "evaluation.py"
    if not eval_script.exists():
        raise FileNotFoundError(f"GenEval2 evaluation.py not found at {eval_script}")

    data_path = geneval2_data or str(repo / "geneval2_data.jsonl")
    out_path = output_file or image_map_path.parent / "geneval2_scores.json"

    cmd = [
        sys.executable, str(eval_script),
        "--benchmark_data", data_path,
        "--image_filepath_data", str(image_map_path),
        "--method", method,
        "--output_file", str(out_path),
    ]
    log.info("Running GenEval2: %s", " ".join(cmd))
    result = subprocess.run(cmd, capture_output=True, text=True, cwd=str(repo))
    if result.returncode != 0:
        log.error("GenEval2 stderr:\n%s", result.stderr[-2000:])
        raise RuntimeError(f"GenEval2 evaluation.py exited {result.returncode}")

    log.info("GenEval2 stdout:\n%s", result.stdout[-2000:])

    if not out_path.exists():
        raise FileNotFoundError(f"GenEval2 output not written: {out_path}")

    scores_data = json.loads(out_path.read_text())
    # Expected structure: list of per-atom scores or a dict with "score" key.
    # Aggregate: arithmetic mean of all atom scores (0–1), then ×100.
    if isinstance(scores_data, list):
        flat = [v for entry in scores_data for v in (entry if isinstance(entry, list) else [entry]) if isinstance(v, (int, float))]
    elif isinstance(scores_data, dict) and "score" in scores_data:
        return float(scores_data["score"])
    else:
        flat = list(scores_data.values()) if isinstance(scores_data, dict) else []

    if not flat:
        return float("nan")
    return 100.0 * sum(flat) / len(flat)


# ---------------------------------------------------------------------------
# Noise floor estimation (base vs base, different seeds)
# ---------------------------------------------------------------------------
def estimate_noise_floor(
    prompts: list[dict],
    model: Optional[str],
    run_dir: Path,
    *,
    width: int,
    height: int,
    n_images: int,
    hps_version: str,
    seed_offset_a: int = 0,
    seed_offset_b: int = 100_000,
) -> dict[str, float]:
    """
    Generate two batches from the same model with offset seeds.
    Returns FID and HPS std-dev between the two runs as noise floor estimates.
    """
    log.info("Estimating noise floor: %s run A vs run B …", model or "default")
    dir_a = run_dir / "noise_a"
    dir_b = run_dir / "noise_b"

    imgs_a = generate_images(prompts, model, dir_a, seed_offset=seed_offset_a, width=width, height=height, n_images=n_images)
    imgs_b = generate_images(prompts, model, dir_b, seed_offset=seed_offset_b, width=width, height=height, n_images=n_images)

    fid_noise = score_fid(dir_a, dir_b)
    hps_a = score_hps(imgs_a, prompts, hps_version=hps_version)
    hps_b = score_hps(imgs_b, prompts, hps_version=hps_version)

    # std of two-point estimate = |a - b| / sqrt(2) (conservative)
    import math
    hps_noise_std = abs(hps_a - hps_b) / math.sqrt(2)

    return {
        "fid_a": fid_noise,     # FID between same-model different-seed runs
        "hps_a": hps_a,
        "hps_b": hps_b,
        "hps_noise_std": hps_noise_std,
        "fid_noise_floor": fid_noise,  # treat same-model FID as noise floor for FID
    }


# ---------------------------------------------------------------------------
# Report emitter
# ---------------------------------------------------------------------------
def emit_report(
    run_dir: Path,
    *,
    base_model: str,
    cand_model: str,
    fid_base: Optional[float],
    fid_cand: Optional[float],
    fid_noise: float,
    hps_base: float,
    hps_cand: float,
    hps_noise_std: float,
    geneval_base: Optional[float],
    geneval_cand: Optional[float],
    geneval_noise_std: float,
    n_images: int,
) -> None:
    """Print and write the markdown report table.

    Gate semantics (redesigned — see module docstring "Gate model"):

    * **HPSv2 is the PRIMARY gate.** It is a higher-is-better learned human
      preference score; the candidate must beat the base by ≥ HPS_PASS_DELTA
      AND by ≥ 2×noise_std. HPSv2 is always measurable (no reference set
      needed), so the gate can ALWAYS rest on it.

    * **clean-FID is a SECONDARY / corroborating signal, and only when a real
      reference set was provided.** FID is a distribution distance to a real
      reference: ``fid_base = FID(reference, base)``, ``fid_cand =
      FID(reference, candidate)`` against the SAME reference. A LOWER candidate
      FID is the improvement, so ``fid_delta = fid_cand - fid_base`` must be
      ≤ FID_PASS_DELTA (and beyond the noise floor) to corroborate.
      When no reference set is available, ``fid_base``/``fid_cand`` are ``None``
      → FID is reported as **N/A** and the gate rests on HPSv2 alone (with a
      logged note). FID is never fabricated to 0.0.

    The overall GATE passes iff every *available* metric passes. FID rows that
    are N/A do not count against the gate (they are corroborating, not required).
    """
    import math

    def fid_moved(delta: float, noise: float) -> str:
        return "YES" if delta <= FID_PASS_DELTA and delta <= -NOISE_SIGMA_MULTIPLIER * noise else "no"

    def hps_moved(delta: float, noise_std: float) -> str:
        return "YES" if delta >= HPS_PASS_DELTA and delta >= NOISE_SIGMA_MULTIPLIER * noise_std else "no"

    def ge_moved(delta: float, noise_std: float) -> str:
        return "YES" if delta >= GENEVAL_PASS_DELTA and delta >= NOISE_SIGMA_MULTIPLIER * noise_std else "no"

    hps_delta = hps_cand - hps_base

    # HPSv2 — PRIMARY gate, always present.
    rows = [
        ("HPSv2.1 (PRIMARY)", f"{hps_base:.4f}", f"{hps_cand:.4f}", f"{hps_delta:+.4f}", f"±{hps_noise_std:.4f}", hps_moved(hps_delta, hps_noise_std)),
    ]

    # clean-FID — corroborating signal, only when a real reference set produced
    # both scores. N/A (None) means no reference set → does not gate.
    fid_available = fid_base is not None and fid_cand is not None and not (
        math.isnan(fid_base) or math.isnan(fid_cand)
    )
    if fid_available:
        fid_delta = fid_cand - fid_base
        rows.append(
            ("clean-FID (corroborating)", f"{fid_base:.2f}", f"{fid_cand:.2f}", f"{fid_delta:+.2f}", f"{fid_noise:.2f}", fid_moved(fid_delta, fid_noise))
        )
    else:
        rows.append(("clean-FID (corroborating)", "N/A", "N/A", "—", "—", "—"))

    if geneval_base is not None and geneval_cand is not None and not math.isnan(geneval_base):
        ge_delta = geneval_cand - geneval_base
        rows.append(
            ("GenEval2 Soft-TIFA_AM", f"{geneval_base:.1f}", f"{geneval_cand:.1f}", f"{ge_delta:+.1f}", f"±{geneval_noise_std:.1f}", ge_moved(ge_delta, geneval_noise_std))
        )
    else:
        rows.append(("GenEval2 Soft-TIFA_AM", "SKIPPED", "SKIPPED", "—", "—", "—"))

    # Gate logic: every AVAILABLE metric (Moved? != "—") must pass. The PRIMARY
    # HPSv2 metric is always available, so the gate is never vacuous.
    gate_results = [r[5] for r in rows if r[5] != "—"]
    gate = "PASS" if gate_results and all(g == "YES" for g in gate_results) else "FAIL"

    fid_note = (
        ""
        if fid_available
        else "\n> **Note:** No real-image reference set provided (`--reference-dir`); "
             "clean-FID is N/A and the gate rests on HPSv2 alone. FID is a "
             "distribution distance to a real reference, not a base-vs-candidate "
             "quantity — supply `--reference-dir` for the corroborating FID signal.\n"
    )

    header = f"""
# ImageForge Eval Report
Base: `{base_model}` | Candidate: `{cand_model}` | N={n_images} images
{fid_note}
| Metric | Base | Candidate | Delta | Noise Floor | Moved? |
|--------|------|-----------|-------|-------------|--------|
"""
    table_rows = "\n".join(f"| {r[0]} | {r[1]} | {r[2]} | {r[3]} | {r[4]} | {r[5]} |" for r in rows)
    thresholds = f"""
## Pass/Fail Thresholds
- HPSv2.1 (PRIMARY): delta ≥ **+{HPS_PASS_DELTA}** AND delta ≥ **{NOISE_SIGMA_MULTIPLIER}×noise_std**
- clean-FID (corroborating, only with `--reference-dir`): delta ≤ **{FID_PASS_DELTA}** AND delta ≤ **-{NOISE_SIGMA_MULTIPLIER}×noise_floor**, where delta = FID(ref,cand) − FID(ref,base)
- GenEval2: delta ≥ **+{GENEVAL_PASS_DELTA} pp** AND delta ≥ **{NOISE_SIGMA_MULTIPLIER}×noise_std**
"""
    report = header + table_rows + thresholds + f"\n## GATE: {gate}\n"

    print(report)
    report_path = run_dir / "eval_report.md"
    report_path.write_text(report)
    log.info("Report written → %s", report_path)


# ---------------------------------------------------------------------------
# Dry-run validator (checks config + API reachability; no generation)
# ---------------------------------------------------------------------------
def _dry_run(args: "argparse.Namespace") -> None:
    """
    Validate configuration and probe API reachability.

    Prints a structured summary and exits:
      0 — config valid AND API reachable
      1 — config valid BUT API unreachable (clear message printed)
    The caller has already validated numeric/string constraints via argparse
    before reaching here, so those checks are not repeated.
    """
    print("=== eval_harness --dry-run ===")
    print(f"  base model       : {args.base}")
    print(f"  candidate model  : {args.candidate}")
    print(f"  n-images         : {args.n_images}")
    print(f"  resolution       : {args.width}×{args.height}")
    print(f"  prompts file     : {args.prompts}")
    print(f"  api-url          : {args.api_url}")
    print(f"  noise-floor      : {args.noise_floor}")
    print(f"  skip-geneval     : {args.skip_geneval}")
    print(f"  hps-version      : {args.hps_version}")
    print()

    # Optional-dep availability report (informational, not fatal in dry-run)
    dep_status: list[tuple[str, bool, str]] = []
    for pkg, import_name, install_hint in [
        ("httpx",     "httpx",    "pip install httpx"),
        ("clean-fid", "cleanfid", "pip install clean-fid"),
        ("hpsv2",     "hpsv2",    "pip install hpsv2"),
    ]:
        try:
            __import__(import_name)
            dep_status.append((pkg, True, ""))
        except ImportError:
            dep_status.append((pkg, False, install_hint))

    print("  Optional dependencies:")
    all_deps_ok = True
    for pkg, ok, hint in dep_status:
        status_str = "OK" if ok else f"MISSING  →  {hint}"
        print(f"    {pkg:<12} {status_str}")
        if not ok:
            all_deps_ok = False
    print()

    # Prompts file check (not fatal; generation would fail later if missing)
    prompts_path = Path(args.prompts)
    if prompts_path.exists():
        print(f"  Prompts file     : FOUND ({prompts_path})")
    else:
        print(f"  Prompts file     : NOT FOUND — create {prompts_path} before running eval")
    print()

    # API reachability check
    try:
        import httpx as _httpx
    except ImportError:
        print("  API reachability : SKIPPED (httpx not installed — run: pip install httpx)")
        print()
        print("Config valid. Install missing dependencies before running eval.")
        sys.exit(0)

    api_url = args.api_url
    print(f"  Checking API reachability at {api_url}/health …")
    try:
        with _httpx.Client(timeout=5.0) as client:
            resp = client.get(f"{api_url}/health")
            resp.raise_for_status()
            health = resp.json()
        status = health.get("status", "?")
        device = health.get("device", "?")
        model  = health.get("model", "?")
        print(f"  API status       : REACHABLE  (status={status} device={device} model={model})")
        print()
        print("Config valid and API reachable. Ready to run eval.")
        sys.exit(0)
    except Exception as exc:
        print(f"  API status       : UNREACHABLE")
        print(f"  Error            : {exc}")
        print()
        print(
            f"Cannot reach ImageForge at {api_url}.\n"
            "Start the server with:  ./run_api.sh\n"
            "Then re-run eval_harness.py (without --dry-run) to begin evaluation."
        )
        sys.exit(1)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def main() -> None:
    global API_URL, GENERATE_ENDPOINT
    parser = argparse.ArgumentParser(description="ImageForge eval harness — clean-FID + HPSv2 + GenEval2")
    parser.add_argument("--base", default="sd-turbo", help="Base model key (e.g. sd-turbo)")
    parser.add_argument("--candidate", default="sdxl-turbo", help="Candidate model key (e.g. sdxl-turbo)")
    parser.add_argument("--prompts", default=str(PROMPTS_FILE), help="Path to prompts_heldout.jsonl")
    parser.add_argument("--run-dir", default=None, help="Directory for this eval run (auto-named if omitted)")
    parser.add_argument("--n-images", type=int, default=DEFAULT_N_IMAGES, help=f"Images per model (default {DEFAULT_N_IMAGES})")
    parser.add_argument("--width", type=int, default=DEFAULT_WIDTH)
    parser.add_argument("--height", type=int, default=DEFAULT_HEIGHT)
    parser.add_argument("--steps", type=int, default=None, help="Inference steps override")
    parser.add_argument("--guidance", type=float, default=None, help="CFG guidance override")
    parser.add_argument("--hps-version", default="v2.1", choices=["v2", "v2.1"], help="HPSv2 model version")
    parser.add_argument("--fid-device", default="cpu", help="Device for FID Inception forward pass (cpu or cuda)")
    parser.add_argument(
        "--reference-dir",
        default=None,
        help=(
            "Directory of REAL held-out images used as the clean-FID reference "
            "distribution. FID is computed as FID(reference, base) and "
            "FID(reference, candidate) against this SAME set; a lower candidate "
            "FID is the improvement. If omitted, FID is reported as N/A and the "
            "gate rests on HPSv2 alone (FID is never fabricated)."
        ),
    )
    parser.add_argument("--skip-geneval", action="store_true", help="Skip GenEval2 scoring (e.g. Qwen3-VL not available)")
    parser.add_argument("--geneval2-repo", default=None, help="Path to cloned facebookresearch/GenEval2 repo")
    parser.add_argument("--geneval2-data", default=None, help="Path to geneval2_data.jsonl (defaults to <repo>/geneval2_data.jsonl)")
    parser.add_argument("--geneval-method", default="soft_tifa_am", choices=["soft_tifa_am", "soft_tifa_gm", "tifa", "vqascore"])
    parser.add_argument("--noise-floor", action="store_true", help="Estimate noise floor (base vs base, different seeds) before comparison")
    parser.add_argument("--api-url", default=API_URL, help=f"ImageForge API base URL (default {API_URL})")
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help=(
            "Validate config and check API reachability without generating images. "
            "Exits 0 if the API is reachable and config is valid; exits 1 with a "
            "clear message if the API is unreachable or config is invalid."
        ),
    )
    args = parser.parse_args()

    # ---------------------------------------------------------------------------
    # Argparse validation — reject bad values early so errors surface before any
    # expensive work (image gen, scoring) begins.
    # ---------------------------------------------------------------------------

    # --n-images must be ≥ 1
    if args.n_images < 1:
        parser.error(f"--n-images must be ≥ 1 (got {args.n_images})")

    # --width / --height must be multiples of 8 (diffusion latent-space requirement)
    if args.width % 8 != 0:
        parser.error(f"--width must be a multiple of 8 (got {args.width})")
    if args.height % 8 != 0:
        parser.error(f"--height must be a multiple of 8 (got {args.height})")

    # --base / --candidate must be non-empty strings
    if not args.base or not args.base.strip():
        parser.error("--base must be a non-empty model key string")
    if not args.candidate or not args.candidate.strip():
        parser.error("--candidate must be a non-empty model key string")

    # Override global API URL if provided
    API_URL = args.api_url
    GENERATE_ENDPOINT = f"{API_URL}/generate"

    # ---------------------------------------------------------------------------
    # --dry-run: validate config + check API reachability, then exit without gen.
    # ---------------------------------------------------------------------------
    if args.dry_run:
        _dry_run(args)

    # Load prompts
    prompts = load_prompts(Path(args.prompts))

    # Run directory
    run_name = f"eval_{args.base}_vs_{args.candidate}_{int(time.time())}"
    run_dir = Path(args.run_dir) if args.run_dir else Path(__file__).parent / "eval_runs" / run_name
    run_dir.mkdir(parents=True, exist_ok=True)
    log.info("Run dir: %s", run_dir)

    # Noise floor estimation
    fid_noise_floor = 10.0    # conservative default if not measured
    hps_noise_std = 0.05      # conservative default
    geneval_noise_std = 1.5   # conservative default (1.5 pp)

    if args.noise_floor:
        nf = estimate_noise_floor(
            prompts, args.base, run_dir / "noise_floor",
            width=args.width, height=args.height, n_images=min(50, args.n_images),
            hps_version=args.hps_version,
        )
        fid_noise_floor = nf["fid_noise_floor"]
        hps_noise_std = nf["hps_noise_std"]
        log.info("Noise floor — FID: %.2f  HPS std: %.4f", fid_noise_floor, hps_noise_std)

    # Generate base images
    dir_base = run_dir / f"imgs_{args.base}"
    imgs_base = generate_images(
        prompts, args.base, dir_base,
        seed_offset=0, width=args.width, height=args.height, n_images=args.n_images,
        steps=args.steps, guidance=args.guidance,
    )

    # Generate candidate images
    dir_cand = run_dir / f"imgs_{args.candidate}"
    imgs_cand = generate_images(
        prompts, args.candidate, dir_cand,
        seed_offset=0, width=args.width, height=args.height, n_images=args.n_images,
        steps=args.steps, guidance=args.guidance,
    )

    if not imgs_base or not imgs_cand:
        sys.exit("No images generated. Check that ImageForge is running: ./run_api.sh")

    # ---------------------------------------------------------------------------
    # clean-FID — coherent before/after ONLY when a real reference set exists.
    #
    # FID is a distance between a REAL-image reference distribution and a
    # GENERATED distribution. The only meaningful before/after is:
    #     fid_base = FID(reference, base)      (lower = base closer to real)
    #     fid_cand = FID(reference, candidate) (lower = candidate closer to real)
    #     delta    = fid_cand - fid_base       (negative = candidate improved)
    # against the SAME reference set. Without a real reference there is NO
    # coherent FID-as-quality signal, so FID is reported as N/A (None) and the
    # gate rests on HPSv2 alone — we never fabricate fid_base = 0.0.
    # ---------------------------------------------------------------------------
    fid_base_score: Optional[float] = None
    fid_cand_score: Optional[float] = None

    ref_dir_arg = getattr(args, "reference_dir", None)
    if ref_dir_arg:
        ref_dir = Path(ref_dir_arg)
        if not ref_dir.is_dir():
            sys.exit(
                f"--reference-dir {ref_dir} is not a directory. Provide a folder of "
                "real held-out images for the FID reference, or omit it (FID will be N/A)."
            )
        log.info("Computing clean-FID against real reference set %s", ref_dir)
        fid_base_score = score_fid(ref_dir, dir_base, device_str=args.fid_device)
        fid_cand_score = score_fid(ref_dir, dir_cand, device_str=args.fid_device)
    else:
        log.warning(
            "No --reference-dir provided: clean-FID is N/A (a coherent FID needs a "
            "real-image reference distribution). The gate will rest on HPSv2 alone."
        )

    # HPSv2
    hps_base_score = score_hps(imgs_base, prompts, hps_version=args.hps_version)
    hps_cand_score = score_hps(imgs_cand, prompts, hps_version=args.hps_version)

    # GenEval2
    geneval_base_score: Optional[float] = None
    geneval_cand_score: Optional[float] = None

    if not args.skip_geneval:
        if not args.geneval2_repo:
            log.warning("--geneval2-repo not set; skipping GenEval2. Pass path to cloned repo.")
            args.skip_geneval = True
        else:
            try:
                map_base = build_geneval_image_map(imgs_base, prompts, run_dir / "geneval_base")
                geneval_base_score = score_geneval2(
                    map_base, args.geneval2_repo,
                    method=args.geneval_method, geneval2_data=args.geneval2_data,
                    output_file=run_dir / "geneval_base" / "scores.json",
                )
                map_cand = build_geneval_image_map(imgs_cand, prompts, run_dir / "geneval_cand")
                geneval_cand_score = score_geneval2(
                    map_cand, args.geneval2_repo,
                    method=args.geneval_method, geneval2_data=args.geneval2_data,
                    output_file=run_dir / "geneval_cand" / "scores.json",
                )
            except Exception as e:
                log.error("GenEval2 scoring failed: %s", e)
                log.info("Re-run with --skip-geneval to skip this metric.")

    # Report
    emit_report(
        run_dir,
        base_model=args.base,
        cand_model=args.candidate,
        fid_base=fid_base_score,
        fid_cand=fid_cand_score,
        fid_noise=fid_noise_floor,
        hps_base=hps_base_score,
        hps_cand=hps_cand_score,
        hps_noise_std=hps_noise_std,
        geneval_base=geneval_base_score,
        geneval_cand=geneval_cand_score,
        geneval_noise_std=geneval_noise_std,
        n_images=args.n_images,
    )


if __name__ == "__main__":
    main()