File size: 2,678 Bytes
54b17f2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | ---
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.
|