PG-MAP NeurIPS 2026 — v1.0 custom-pipeline release

README.md

@@ -0,0 +1,86 @@
+---
+language:
+- en
+license: mit
+library_name: diffusers
+tags:
+- text-to-image
+- stable-diffusion
+- inference-time-alignment
+- preference-optimization
+- pg-map
+- neurips-2026
+pipeline_tag: text-to-image
+---
+
+# PG-MAP for Stable Diffusion 1.5
+
+Custom diffusers pipeline for **PG-MAP** (Preference-Guided Adaptive MAP) on SD 1.5. Per-step joint optimization of conditioning $c$ and latent $z_t$ via a trajectory-level Gibbs-MAP / proximal energy objective, optionally guided by a frozen preference reward (PickScore by default).
+
+NeurIPS 2026 — see [github.com/sophialanlan/PG-MAP](https://github.com/sophialanlan/PG-MAP) for the paper, full configs, and reproduction scripts.
+
+## Install
+
+```bash
+pip install pg-map
+# or
+pip install git+https://github.com/sophialanlan/PG-MAP
+```
+
+## Usage
+
+```python
+from diffusers import DiffusionPipeline
+from pgmap import sd15_defaults, FrozenRewardModel
+import torch
+
+pipe = DiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    custom_pipeline="sophialan/pg-map-sd15",
+    torch_dtype=torch.float16,
+    safety_checker=None,
+).to("cuda")
+
+cfg = sd15_defaults()                                       # paper defaults
+reward = FrozenRewardModel("pickscore", device="cuda")
+
+image = pipe(
+    "a phoenix rising from ashes, vivid orange and red feathers",
+    pg_map_config=cfg,
+    reward_model=reward,
+).images[0]
+```
+
+Passing `pg_map_config=None` falls through to the vanilla `StableDiffusionPipeline`, so the class is a strict superset of the parent.
+
+## Method overview
+
+Per denoising step $t$, PG-MAP solves the proximal MAP problem:
+
+$$\mathcal{J}_t(c, z_t) = -\tfrac{1}{2\beta_{t|s}}\|r_t(c,z_t)\|^2 - \tfrac{1}{2\sigma_c^2}\|c-\mu_t\|^2 - \tfrac{1}{2\sigma_z(t)^2}\|z_t-z_t^{\text{ddim}}\|^2 + \lambda\,Q(\hat x_0(z_t,c), y)$$
+
+with $K$ inner gradient-ascent steps and a schedule-adaptive trust region $\sigma_z(t)=\gamma\sqrt{1-\bar\alpha_t}$.
+
+## Paper headline (SD 1.5, PartiPrompts $n=1632$, seed 123)
+
+| Method | PickScore | HPS | Aesthetic | CLIP |
+|---|---|---|---|---|
+| PG-MAP (default)           | **56.8%** | 52.8% | 54.0% | 50.6% |
+| Tuned-CFG + PG-MAP         | 53.6% | **66.0%** | **60.2%** | **56.0%** |
+
+Win-rate vs. same-seed static baseline.
+
+## Citation
+
+```bibtex
+@inproceedings{sun2026pgmap,
+  title={{PG-MAP}: Joint {MAP} Optimization for Inference-Time Alignment of Diffusion and Flow-Matching Models},
+  author={Sun, Ruolan and Polak, Pawel},
+  booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
+  year={2026}
+}
+```
+
+## License
+
+MIT (see [LICENSE](https://github.com/sophialanlan/PG-MAP/blob/main/LICENSE)). Pretrained weights remain under their original licenses.

pipeline.py

@@ -0,0 +1,57 @@
+"""sophialan/pg-map-sd15 — PG-MAP custom pipeline for Stable Diffusion 1.5.
+
+Loaded by the HuggingFace community-pipeline registry::
+
+    from diffusers import DiffusionPipeline
+    pipe = DiffusionPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        custom_pipeline="sophialan/pg-map-sd15",
+        torch_dtype=torch.float16,
+    ).to("cuda")
+
+This file is a thin re-export shim. The actual pipeline class plus all
+PG-MAP machinery lives in the ``pg-map`` PyPI package; install it first::
+
+    pip install pg-map
+
+Usage after loading::
+
+    from pgmap import sd15_defaults, FrozenRewardModel
+
+    cfg = sd15_defaults()                                    # paper defaults
+    reward = FrozenRewardModel("pickscore", device="cuda")
+
+    image = pipe(
+        "a phoenix rising from ashes",
+        pg_map_config=cfg,
+        reward_model=reward,
+    ).images[0]
+
+Passing ``pg_map_config=None`` falls through to the vanilla
+``StableDiffusionPipeline``, so this class is a strict superset of the
+parent.
+
+Citation::
+
+    @inproceedings{sun2026pgmap,
+      title={PG-MAP: Joint MAP Optimization for Inference-Time Alignment
+             of Diffusion and Flow-Matching Models},
+      author={Sun, Ruolan and Polak, Pawel},
+      booktitle={NeurIPS},
+      year={2026},
+    }
+"""
+from __future__ import annotations
+
+try:
+    from pgmap.pipelines.sd15 import PGMAPStableDiffusionPipeline
+except ImportError as e:
+    raise ImportError(
+        "Custom pipeline `sophialan/pg-map-sd15` requires the `pg-map` "
+        "package.\n"
+        "  Install from PyPI:   pip install pg-map\n"
+        "  Install from source: pip install git+https://github.com/sophialanlan/PG-MAP\n"
+        "Repo: https://github.com/sophialanlan/PG-MAP"
+    ) from e
+
+__all__ = ["PGMAPStableDiffusionPipeline"]