Add ImageNet LPIPS benchmark workflow

2026-01-24 10:19:36 +08:00
parent 060fcc4475
commit fb35a1d032
3 changed files with 357 additions and 0 deletions
@@ -0,0 +1,66 @@
+name: LanPaintBench LPIPS (ImageNet subset)
+
+on:
+  pull_request:
+    branches: [ master, main ]
+    paths:
+      - 'src/**'
+      - 'tests/benchmark_lpips.py'
+      - 'tests/compare_benchmarks.py'
+      - '.github/workflows/benchmark.yml'
+  workflow_dispatch:
+
+jobs:
+  benchmark:
+    name: Compare Base vs Head
+    runs-on: ubuntu-latest
+    env:
+      HF_HOME: ~/.cache/huggingface
+      BASE_REF: ${{ github.event.pull_request.base.sha || github.event.repository.default_branch }}
+    steps:
+      - name: Checkout Head (Current PR)
+        uses: actions/checkout@v4
+        with:
+          path: head
+          fetch-depth: 1
+
+      - name: Checkout Base (Baseline)
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ env.BASE_REF }}
+          path: base
+          fetch-depth: 1
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+          cache: 'pip'
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+          pip install lpips datasets numpy pillow
+          # We don't need to install LanPaint as a package because we use sys.path in the script
+
+      - name: Cache HF + torch
+        uses: actions/cache@v4
+        with:
+          path: |
+            ~/.cache/huggingface
+            ~/.cache/torch
+          key: ${{ runner.os }}-lpips-imagenet-${{ hashFiles('head/tests/benchmark_lpips.py', 'head/tests/compare_benchmarks.py') }}
+
+      - name: Run Baseline Benchmark
+        run: |
+          # Use the benchmark script from the HEAD branch to ensure test consistency
+          python head/tests/benchmark_lpips.py --src base/src --output baseline.json --subset 8 --seed 42 --mask-type box --n-steps 5 --sigma 1.0 --device cpu
+
+      - name: Run Current Benchmark
+        run: |
+          python head/tests/benchmark_lpips.py --src head/src --output current.json --subset 8 --seed 42 --mask-type box --n-steps 5 --sigma 1.0 --device cpu
+
+      - name: Compare Results
+        run: |
+          python head/tests/compare_benchmarks.py --baseline baseline.json --current current.json --max-rel-diff 0.01 --max-abs-diff 0.0
@@ -0,0 +1,228 @@
+"""Lightweight CI benchmark inspired by LanPaintBench.
+
+Runs LPIPS on a tiny deterministic ImageNet-1k-256x256 subset and compares base vs head.
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+
+def _make_mask(mask_type: str, height: int, width: int) -> torch.Tensor:
+    """Return a (1, 3, H, W) mask where 1=known region, 0=to-inpaint."""
+    if height <= 0 or width <= 0:
+        raise ValueError("height/width must be positive")
+
+    mask = torch.ones((1, 3, height, width), dtype=torch.float32)
+
+    if mask_type == "half":
+        mask[:, :, :, width // 2 :] = 0.0
+        return mask
+
+    if mask_type == "box":
+        h0, h1 = height // 4, (3 * height) // 4
+        w0, w1 = width // 4, (3 * width) // 4
+        mask[:, :, h0:h1, w0:w1] = 0.0
+        return mask
+
+    if mask_type == "outpaint":
+        mask = torch.zeros((1, 3, height, width), dtype=torch.float32)
+        h0, h1 = height // 4, (3 * height) // 4
+        w0, w1 = width // 4, (3 * width) // 4
+        mask[:, :, h0:h1, w0:w1] = 1.0
+        return mask
+
+    if mask_type == "checkerboard":
+        grid = max(1, min(height, width) // 16)
+        mask_2d = torch.zeros((height, width), dtype=torch.float32)
+        for row in range(height):
+            row_parity = (row // grid) % 2
+            start = 0 if row_parity == 0 else 1
+            for block in range(start, (width + grid - 1) // grid, 2):
+                c0 = block * grid
+                c1 = min(width, (block + 1) * grid)
+                mask_2d[row, c0:c1] = 1.0
+        return mask_2d.unsqueeze(0).unsqueeze(0).repeat(1, 3, 1, 1)
+
+    raise ValueError(f"unknown mask type: {mask_type}")
+
+
+class _DummySampling:
+    def noise_scaling(self, sigma, noise, latent_image):  # type: ignore[no-untyped-def]
+        # Variance exploding convention (ComfyUI/k-diffusion style).
+        return latent_image + noise * sigma
+
+
+class _DummyModel:
+    def __init__(self) -> None:
+        self.inner_model = self
+        self.model_sampling = _DummySampling()
+
+    def __call__(self, x, sigma, model_options=None, seed=None):  # type: ignore[no-untyped-def]
+        # Cheap "denoiser" so the benchmark exercises LanPaint math deterministically on CPU.
+        sigma_ = sigma.reshape((sigma.shape[0],) + (1,) * (x.ndim - 1))
+        x0 = x / (1.0 + sigma_**2)
+        return x0, x0
+
+
+def _load_images(*, dataset: str, split: str, subset: int, seed: int, shuffle_buffer: int) -> list[torch.Tensor]:
+    from datasets import load_dataset
+
+    ds = load_dataset(dataset, split=split, streaming=True)
+    try:
+        ds = ds.shuffle(seed=seed, buffer_size=shuffle_buffer)
+    except Exception:
+        pass
+
+    images: list[torch.Tensor] = []
+    for idx, item in enumerate(ds):
+        if idx >= subset:
+            break
+        image = item["image"].convert("RGB")
+        arr = np.asarray(image, dtype=np.uint8)
+        img = torch.from_numpy(arr).permute(2, 0, 1).float() / 127.5 - 1.0
+        images.append(img)
+    return images
+
+
+def _run_bench(  # type: ignore[no-untyped-def]
+    LanPaintClass,
+    images: list[torch.Tensor],
+    device: torch.device,
+    *,
+    n_steps: int,
+    sigma_value: float,
+    mask_type: str,
+    seed: int,
+) -> float:
+    import lpips
+
+    lpips_fn = lpips.LPIPS(net="alex").to(device)
+
+    sigma = torch.tensor([sigma_value], device=device)
+    abt = 1.0 / (1.0 + sigma**2)
+    flow_t = torch.zeros_like(sigma)
+    current_times = (sigma, abt, flow_t)
+
+    engine = LanPaintClass(
+        Model=_DummyModel(),
+        NSteps=n_steps,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.1,
+    )
+
+    scores: list[float] = []
+    for idx, img in enumerate(images):
+        img = img.unsqueeze(0).to(device)
+        latent_image = img.clone()
+
+        height, width = img.shape[-2:]
+        latent_mask = _make_mask(mask_type, height, width).to(device)
+
+        gen = torch.Generator(device=device.type).manual_seed(seed + idx)
+        noise = torch.randn(img.shape, generator=gen, device=device)
+        x = noise * sigma.reshape((1,) + (1,) * (img.ndim - 1))
+
+        output = engine(
+            x=x,
+            latent_image=latent_image,
+            noise=noise,
+            sigma=sigma,
+            latent_mask=latent_mask,
+            current_times=current_times,
+            model_options={},
+            seed=seed,
+            n_steps=n_steps,
+        ).clamp(-1.0, 1.0)
+
+        dist = lpips_fn(output, img)
+        scores.append(float(dist.item()))
+
+    return float(np.mean(scores))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="LanPaint LPIPS benchmark (ImageNet-1k-256x256 subset)")
+    parser.add_argument("--src", type=Path, required=True, help="Path to LanPaint src directory (contains LanPaint/)")
+    parser.add_argument("--output", type=Path, required=True, help="Output JSON path")
+
+    parser.add_argument("--dataset", type=str, default="benjamin-paine/imagenet-1k-256x256")
+    parser.add_argument("--split", type=str, default="train")
+    parser.add_argument("--subset", type=int, default=8)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--shuffle-buffer", type=int, default=1024)
+
+    parser.add_argument("--mask-type", type=str, default="box", choices=["half", "box", "outpaint", "checkerboard"])
+    parser.add_argument("--n-steps", type=int, default=5, help="LanPaint inner iterations per call")
+    parser.add_argument("--sigma", type=float, default=1.0, help="Variance-exploding sigma value")
+    parser.add_argument("--device", type=str, default="cpu", choices=["cpu", "cuda", "auto"])
+    args = parser.parse_args()
+
+    if args.subset <= 0:
+        raise SystemExit("--subset must be > 0")
+    if args.n_steps <= 0:
+        raise SystemExit("--n-steps must be > 0")
+    if args.shuffle_buffer <= 0:
+        raise SystemExit("--shuffle-buffer must be > 0")
+
+    device_str = args.device
+    if device_str == "auto":
+        device_str = "cuda" if torch.cuda.is_available() else "cpu"
+    device = torch.device(device_str)
+
+    sys.path.insert(0, str(args.src.resolve()))
+    from LanPaint.lanpaint import LanPaint as LanPaintEngine
+
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    images = _load_images(
+        dataset=args.dataset,
+        split=args.split,
+        subset=args.subset,
+        seed=args.seed,
+        shuffle_buffer=args.shuffle_buffer,
+    )
+
+    score = _run_bench(
+        LanPaintEngine,
+        images,
+        device,
+        n_steps=args.n_steps,
+        sigma_value=args.sigma,
+        mask_type=args.mask_type,
+        seed=args.seed,
+    )
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    with args.output.open("w", encoding="utf-8") as f:
+        json.dump(
+            {
+                "lpips": score,
+                "config": {
+                    "dataset": args.dataset,
+                    "split": args.split,
+                    "subset": args.subset,
+                    "seed": args.seed,
+                    "shuffle_buffer": args.shuffle_buffer,
+                    "mask_type": args.mask_type,
+                    "n_steps": args.n_steps,
+                    "sigma": args.sigma,
+                    "device": device_str,
+                },
+            },
+            f,
+            indent=2,
+        )
+
+    print(f"LPIPS: {score:.6f}")
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,63 @@
+"""Compare two LPIPS benchmark JSON outputs (baseline vs current)."""
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+
+def _append_step_summary(markdown: str) -> None:
+    summary_path = os.environ.get("GITHUB_STEP_SUMMARY")
+    if not summary_path:
+        return
+    with Path(summary_path).open("a", encoding="utf-8", errors="ignore") as f:
+        f.write(markdown)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--baseline", type=Path, required=True)
+    parser.add_argument("--current", type=Path, required=True)
+    parser.add_argument("--max-rel-diff", type=float, default=0.01, help="Fail if (current-baseline)/baseline exceeds this")
+    parser.add_argument("--max-abs-diff", type=float, default=0.0, help="Fail if current-baseline exceeds this")
+    args = parser.parse_args()
+
+    baseline = json.loads(args.baseline.read_text(encoding="utf-8"))
+    current = json.loads(args.current.read_text(encoding="utf-8"))
+
+    b_score = float(baseline["lpips"])
+    c_score = float(current["lpips"])
+    diff = c_score - b_score
+    rel_diff = diff / b_score if b_score != 0.0 else float("inf")
+
+    lines = [
+        f"Baseline LPIPS: {b_score:.6f}",
+        f"Current LPIPS:  {c_score:.6f}",
+        f"Difference:     {diff:+.6f}",
+    ]
+    if b_score != 0.0:
+        lines.append(f"Relative diff:   {rel_diff:+.2%}")
+    print("\n".join(lines))
+
+    summary = "\n".join(
+        [
+            "### LPIPS benchmark",
+            "",
+            "| Metric | Baseline | Current | Diff |",
+            "|---|---:|---:|---:|",
+            f"| LPIPS | {b_score:.6f} | {c_score:.6f} | {diff:+.6f} |",
+            "",
+        ]
+    )
+    _append_step_summary(summary)
+
+    if diff > args.max_abs_diff:
+        if b_score == 0.0 or rel_diff > args.max_rel_diff:
+            print("FAIL: LPIPS regression exceeds configured threshold(s)")
+            raise SystemExit(1)
+
+    print("PASS: LPIPS within configured threshold(s)")
+
+
+if __name__ == "__main__":
+    main()