| --- |
| 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. |
| |
