concept_configs.py

"""Concept Guidance presets for inference.

The canonical ``text``, ``hands``, and ``aesthetics`` entries reproduce the
layer configurations in Table 6 and the selected guidance scales in Table 11
of the paper. Additional research presets are retained for convenience but are
not claimed as Table 6 configurations.

The module is intentionally dependency-free so scripts and notebooks can
inspect presets without loading a diffusion model.
"""

from dataclasses import dataclass
from typing import Dict, Tuple


@dataclass(frozen=True)
class ConceptPreset:
    """A concept-specific set of layers, weights, and guidance scale."""

    layers: Tuple[int, ...]
    weights: Tuple[float, ...]
    scale: float

    def __post_init__(self) -> None:
        if not self.layers:
            raise ValueError("A concept preset must contain at least one layer.")
        if len(self.layers) != len(self.weights):
            raise ValueError("Each concept-relevant layer needs exactly one weight.")
        if sum(self.weights) <= 0:
            raise ValueError("The sum of the layer weights must be positive.")
        if self.scale < 1.0:
            raise ValueError("Concept Guidance scales must be greater than or equal to 1.")


MODEL_IDS = {
    "pixart": "PixArt-alpha/PixArt-XL-2-1024-MS",
    "sd3": "stabilityai/stable-diffusion-3-medium-diffusers",
    "sd35": "stabilityai/stable-diffusion-3.5-large",
    "flux": "black-forest-labs/FLUX.1-dev",
}


CONCEPT_PRESETS: Dict[str, Dict[str, ConceptPreset]] = {
    "pixart": {
        "text": ConceptPreset((5, 25, 13, 14, 9), (0.8, 0.7, 0.5, 0.2, 0.1), 2.0),
        "hands": ConceptPreset((3, 23, 8, 18), (1.0, 0.88, 0.73, 0.63), 2.0),
        "aesthetics": ConceptPreset((9, 20, 5), (1.0, 0.28, 0.12), 2.5),
        "counting": ConceptPreset((15, 14, 21, 7), (0.5, 0.4, 0.2, 0.2), 2.0),
        "position": ConceptPreset((10, 21, 22, 4, 5), (0.4, 0.2, 0.2, 0.2, 0.2), 2.5),
        "hands_and_aesthetics": ConceptPreset((20, 9, 15), (3.0, 1.5, 0.4), 2.0),
    },
    "sd3": {
        "text": ConceptPreset((7, 11, 4), (1.0, 0.55, 0.45), 2.25),
        "hands": ConceptPreset((9, 6, 5, 10), (1.0, 0.48, 0.45, 0.21), 2.0),
        "aesthetics": ConceptPreset((6, 7, 8), (1.0, 0.71, 0.70), 2.5),
        "complex_text": ConceptPreset((7, 11, 4), (1.0, 0.55, 0.45), 2.25),
        "counting": ConceptPreset((8, 6, 9), (0.2, 0.1, 0.1), 2.5),
        "position": ConceptPreset((8, 9, 7, 6, 22), (0.6, 0.6, 0.3, 0.3, 0.3), 2.5),
        "complex_text_and_hands": ConceptPreset((7, 9), (3.0, 2.0), 2.5),
        "complex_text_and_aesthetics": ConceptPreset(
            (7, 9, 11, 5), (1.4, 1.2, 1.0, 0.8), 2.5
        ),
        "hands_and_aesthetics": ConceptPreset((6, 9, 7, 10), (1.9, 1.5, 1.3, 0.9), 2.5),
    },
    "sd35": {
        "text": ConceptPreset((15, 10), (1.0, 0.49), 2.5),
        "hands": ConceptPreset((7, 18), (1.0, 0.98), 2.25),
        "aesthetics": ConceptPreset((4,), (1.0,), 2.0),
        "counting": ConceptPreset((18, 5, 7, 10, 11), (0.12, 0.7, 0.7, 0.3, 0.3), 2.0),
        "position": ConceptPreset((5, 11, 14, 22, 2), (0.13, 0.13, 0.13, 0.13, 0.8), 2.0),
        "text_and_hands": ConceptPreset((2, 18, 7, 11), (0.9, 0.8, 0.7, 0.4), 2.25),
        "text_and_aesthetics": ConceptPreset((11, 18, 3, 2), (0.7, 0.5, 0.5, 0.3), 2.25),
        "hands_and_aesthetics": ConceptPreset((6, 9, 18, 2), (1.7, 1.4, 1.1, 1.0), 2.25),
        "ukiyo_e": ConceptPreset((13, 10, 9), (5.0, 3.0, 2.0), 2.0),
        "background_separation": ConceptPreset((12, 10, 13), (6.0, 3.0, 1.0), 2.0),
        "symmetry": ConceptPreset((3, 12, 16), (4.0, 3.0, 3.0), 2.0),
    },
    "flux": {
        "text": ConceptPreset((11, 3, 9), (1.0, 0.85, 0.74), 2.25),
        "hands": ConceptPreset((9, 12, 15), (1.0, 0.81, 0.78), 1.75),
        "aesthetics": ConceptPreset((15, 14, 13), (1.0, 0.18, 0.09), 3.0),
        "complex_text": ConceptPreset((11, 3, 9), (1.0, 0.85, 0.74), 2.25),
        "counting": ConceptPreset((10, 17, 3), (3.4, 3.0, 1.5), 2.0),
        "position": ConceptPreset((13, 17, 9, 11, 12), (0.24, 0.19, 0.14, 0.14, 0.14), 3.5),
        "complex_text_and_hands": ConceptPreset(
            (6, 4, 15, 12), (7.0, 5.7, 5.3, 5.0), 2.5
        ),
        "complex_text_and_aesthetics": ConceptPreset((15, 14, 13), (5.9, 2.6, 1.0), 2.5),
        "hands_and_aesthetics": ConceptPreset((15,), (6.3,), 2.5),
    },
}

# Canonical presets reported in Table 6. The associated ``scale`` fields are
# the bold, selected settings in the paper's Table 11.
PAPER_PRESET_KEYS = {
    "pixart": ("hands", "aesthetics"),
    "sd3": ("text", "hands", "aesthetics"),
    "sd35": ("text", "hands", "aesthetics"),
    "flux": ("text", "hands", "aesthetics"),
}

PAPER_PRESETS: Dict[str, Dict[str, ConceptPreset]] = {
    model: {concept: CONCEPT_PRESETS[model][concept] for concept in concepts}
    for model, concepts in PAPER_PRESET_KEYS.items()
}


def get_preset(model: str, concept: str) -> ConceptPreset:
    """Return a concept preset with a helpful error for unknown names."""

    model_key = model.lower().replace("-", "").replace(".", "")
    model_key = {"sd30": "sd3", "sd350": "sd35", "pixartalpha": "pixart"}.get(
        model_key, model_key
    )
    concept_key = concept.lower().replace("-", "_").replace(" ", "_")

    if model_key not in CONCEPT_PRESETS:
        available = ", ".join(CONCEPT_PRESETS)
        raise KeyError(f"Unknown model '{model}'. Available models: {available}.")
    if concept_key not in CONCEPT_PRESETS[model_key]:
        available = ", ".join(CONCEPT_PRESETS[model_key])
        raise KeyError(
            f"Unknown concept '{concept}' for {model_key}. Available concepts: {available}."
        )
    return CONCEPT_PRESETS[model_key][concept_key]