sophialan/pg-map-sd15

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 for the paper, full configs, and reproduction scripts.

Install

pip install pg-map
# or
pip install git+https://github.com/sophialanlan/PG-MAP

Usage

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)=-\frac{1}{2{\beta }_{t\mathrm{\mid }s}}\mathrm{\parallel }{r}_t(c,z_t){\mathrm{\parallel }}^2-\frac{1}{2{\sigma }_c^2}\mathrm{\parallel }c-{\mu }_t{\mathrm{\parallel }}^2-\frac{1}{2{\sigma }_z(t{)}^2}\mathrm{\parallel }{z}_t-z_t^{\text{ddim}}{\mathrm{\parallel }}^2+\lambda Q({\widehat{x}}_0(z_t,c),y)\backslash mathcal\{J\}\_t(c,\; z\_t)\; =\; -\backslash tfrac\{1\}\{2\backslash beta\_\{t|s\}\}\backslash |r\_t(c,z\_t)\backslash |^2\; -\; \backslash tfrac\{1\}\{2\backslash sigma\_c^2\}\backslash |c-\backslash mu\_t\backslash |^2\; -\; \backslash tfrac\{1\}\{2\backslash sigma\_z(t)^2\}\backslash |z\_t-z\_t^\{\backslash text\{ddim\}\}\backslash |^2\; +\; \backslash lambda\backslash ,Q(\backslash 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

@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). Pretrained weights remain under their original licenses.