sophialan/pg-map-sdxl

PG-MAP for Stable Diffusion XL

Custom diffusers pipeline for PG-MAP (Preference-Guided Adaptive MAP) on SDXL. Per-step refinement of the token-level conditioning $c$ and the latent $z_t$ at every DDIM step. Pooled embeddings and time_ids are kept frozen (paper §3.5).

Preprint — arXiv:2606.22958; see github.com/sophialanlan/PG-MAP for the paper, full configs, and reproduction scripts.

Install

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

Usage

from diffusers import DiffusionPipeline
from pgmap import sdxl_defaults, FrozenRewardModel
import torch

pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    custom_pipeline="sophialan/pg-map-sdxl",
    torch_dtype=torch.float16,
    variant="fp16",
).to("cuda")

cfg = sdxl_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]

Compute

Setting	Per-image wall-clock (RTX PRO 6000 Blackwell, fp16)
Vanilla pass-through, 50 steps, 1024²	~6 s
PG-MAP default (K=2, λ=0.1), 50 steps	~30 s
Reward-free MAP-cz (λ=0)	~12 s

The reward backward dominates per-image cost; the reward-free MAP-cz path tracks PG-MAP within 0.3 pp PickScore at ~2.6× lower wall-clock and is the recommended deployment when latency matters.

Paper headline (SDXL, PartiPrompts $n=1632$, seed 123)

Method	PickScore	HPS	Aesthetic	CLIP
PG-MAP (default)	56.4%	47.1%	56.2%	48.1%
Tuned-CFG + PG-MAP (w*=7.5)	51.3%	64.6%	56.5%	52.8%

Win-rate vs. same-seed static SDXL baseline (w=5.0).

Citation

@misc{sun2026pgmap,
  title={{PG-MAP}: Joint {MAP} Optimization for Inference-Time Alignment of Diffusion and Flow-Matching Models},
  author={Sun, Ruolan and Polak, Pawel},
  year={2026},
  eprint={2606.22958},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2606.22958}
}

License

MIT (see LICENSE). SDXL weights are under the CreativeML Open RAIL++-M license.