pipeline.py

"""sophialan/pg-map-sd3 — PG-MAP custom pipeline for Stable Diffusion 3.5-medium.

Loaded by the HuggingFace community-pipeline registry::

    from diffusers import DiffusionPipeline
    pipe = DiffusionPipeline.from_pretrained(
        "stabilityai/stable-diffusion-3.5-medium",
        custom_pipeline="sophialan/pg-map-sd3",
        torch_dtype=torch.float16,
    ).to("cuda")

This file is a thin re-export shim. Install the ``pg-map`` PyPI package
first::

    pip install pg-map

Defaults to **UG-FM** (the paper's 91.9 % PickScore / 75.7 % HPS row on
PartiPrompts $n{=}1632$): data-side gate, K_UG=4, eta_z=0.1, full
backprop through the velocity prediction. To run the more expensive
full PG-MAP-FM (joint c + z_t with flow consistency + Gaussian priors
+ reward), pass ``pg_map_config`` with ``optimize_c=True``::

    from pgmap import sdxl_defaults
    from pgmap_config import PGMAPConfig
    from dataclasses import replace

    cfg = sdxl_defaults()                # use as starting point
    cfg = replace(cfg, optimize_c=True, optimize_z=True)
    # ... cfg.refinement.K, cfg.refinement.eta_c, etc. as desired

    image = pipe("a phoenix rising from ashes", pg_map_config=cfg).images[0]

Note: SD3.5-medium requires acceptance of the Stability AI Community
License on huggingface.co before first load.

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.sd3 import PGMAPStableDiffusion3Pipeline
except ImportError as e:
    raise ImportError(
        "Custom pipeline `sophialan/pg-map-sd3` 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__ = ["PGMAPStableDiffusion3Pipeline"]