src/aamcq/prompt_render.py

"""Prompt rendering for the FLUX.2-klein backend."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

from aamcq.profile import VisualProfile

# Template structure follows the BFL klein prompting guide: scene front-loaded,
# natural-language prose, trailing "Key: value." markers. Lighting is embedded
# mid-sentence as an environmental modifier — klein responds to it more
# strongly there than as a trailing key.
#   docs.bfl.ml/guides/prompting_guide_flux2_klein
PROMPT_TEMPLATES: dict[str, str] = {
    "flux2_klein": (
        "{scene}, lit by {lighting_phrase}, rendered as {medium_phrase} "
        "in {style_phrase}. Color palette: {color_lc}."
    ),
}

# Medium-to-phrase map: bare vocab words like "Ink Drawing" or "Pixel Art"
# aren't idiomatic article-bearing noun phrases, so we standardize the wording
# here to keep the rendered prompt grammatical.
ART_MEDIUM_PHRASES: dict[str, str] = {
    "Oil Painting": "an oil painting",
    "Watercolor": "a watercolor painting",
    "Ink Drawing": "an ink drawing",
    "Digital Painting": "a digital painting",
    "Pixel Art": "a pixel-art illustration",
    "Pencil Sketch": "a pencil sketch",
}

# Style-to-phrase map. Most styles are rendered as "{X} style"; Minimalism and
# Art Deco use expanded phrases because the bare word doesn't activate klein's
# flat-design / decorative-geometric signals.
ART_STYLE_PHRASES: dict[str, str] = {
    "Impressionism": "Impressionism style",
    "Anime": "Anime style",
    "Photorealism": "Photorealism style",
    "Cubism": "Cubism style",
    "Minimalism": "flat minimalist style with sparse composition",
    "Art Deco": "Art Deco style with bold geometric shapes, clean symmetry, and ornamental lines",
}

# Lighting-to-phrase map, spanning a (temperature × hardness × key × source)
# grid so the 6 values are maximally distinguishable in the rendered image.
LIGHTING_PHRASES: dict[str, str] = {
    "Golden Hour": "golden hour sunset light",
    "Moody Low-Key": "moody low-key dramatic light",
    "Soft Overcast": "soft diffused overcast daylight",
    "Harsh Noon": "harsh direct noon sunlight",
    "Neon Glow": "pink and cyan neon glow",
    "Candlelit": "warm dim amber glow",
}

NEGATIVE_PROMPTS: dict[str, str | None] = {
    "flux2_klein": None,
}

MODEL_SPECS: dict[str, dict[str, object]] = {
    "flux2_klein": {
        "model_id": "black-forest-labs/FLUX.2-klein-9B",
        "num_inference_steps": 4,
        "guidance_scale": 1.0,
    },
}

Backend = Literal["flux2_klein"]


@dataclass(frozen=True)
class RenderedPrompt:
    backend: str
    prompt: str
    negative_prompt: str | None


def _format_kwargs(profile: VisualProfile, base_prompt: str) -> dict[str, str]:
    if profile.art_medium not in ART_MEDIUM_PHRASES:
        raise ValueError(f"no ART_MEDIUM_PHRASES entry for {profile.art_medium!r}")
    if profile.lighting not in LIGHTING_PHRASES:
        raise ValueError(f"no LIGHTING_PHRASES entry for {profile.lighting!r}")
    if profile.art_style not in ART_STYLE_PHRASES:
        raise ValueError(f"no ART_STYLE_PHRASES entry for {profile.art_style!r}")
    return {
        "scene": base_prompt,
        "medium_phrase": ART_MEDIUM_PHRASES[profile.art_medium],
        "lighting_phrase": LIGHTING_PHRASES[profile.lighting],
        "style_phrase": ART_STYLE_PHRASES[profile.art_style],
        "color_lc": profile.color.lower(),
    }


def render(profile: VisualProfile, base_prompt: str, backend: Backend = "flux2_klein") -> RenderedPrompt:
    if backend not in PROMPT_TEMPLATES:
        raise ValueError(f"unknown backend {backend!r}; expected one of {list(PROMPT_TEMPLATES)}")
    prompt = PROMPT_TEMPLATES[backend].format(**_format_kwargs(profile, base_prompt))
    prompt = _post_process(prompt)
    return RenderedPrompt(
        backend=backend,
        prompt=prompt,
        negative_prompt=NEGATIVE_PROMPTS[backend],
    )


def _post_process(prompt: str) -> str:
    return (
        " ".join(prompt.split())
        .replace(" ,", ",")
        .replace(",,", ",")
        .rstrip(", ")
    )