imagery.py

"""Optional: generate a photo of the finished dish with FLUX.2 klein.

A visual payoff after the plan — "here's what it'll look like." FLUX.2 klein (4B)
runs via Hugging Face Inference (fal-ai provider), billed to the hackathon org,
so it's a hosted API call — no GPU to host. 4B keeps us under the 32B total cap
(Mellum 12B + MiniCPM ~8B + FLUX 4B ≈ 24B).

Needs HF_TOKEN in the environment; without it the feature is a no-op with a hint.
"""

import os

IMAGE_MODEL = os.environ.get("IMAGE_MODEL", "black-forest-labs/FLUX.2-klein-4B")
IMAGE_PROVIDER = os.environ.get("IMAGE_PROVIDER", "fal-ai")
# Empty = bill the personal account (your own credits). Set to an org slug to
# bill that org instead (only works if the org has an inference-credit pool).
HF_BILL_TO = os.environ.get("HF_BILL_TO", "").strip()
HF_TOKEN = os.environ.get("HF_TOKEN")


def dish_prompt(dish: str, ingredients: list[str]) -> str:
    ing = ", ".join(ingredients[:6])
    return (
        f"Transform these raw ingredients into a beautifully plated, finished {dish}"
        + (f" made with {ing}" if ing else "")
        + ". Professional, appetizing food photograph, freshly plated on a simple "
          "ceramic dish, soft natural light, shallow depth of field, high detail, no text"
    )


def generate_dish_image(dish: str, ingredients: list[str], input_image=None):
    """Return (PIL.Image | None, note). FLUX.2 klein on fal-ai is image-to-image,
    so it transforms the cook's ingredients photo into the plated dish. Never
    raises — degrades to a hint."""
    if not (dish or "").strip():
        return None, "Plan a dish first, then I can picture it."
    if input_image is None:
        return None, "📷 Upload an ingredients photo (step 1) — FLUX renders the dish from it."
    if not HF_TOKEN:
        return None, "🔌 Set HF_TOKEN to generate the dish image (FLUX.2 klein via HF/fal-ai)."
    try:
        import io

        from huggingface_hub import InferenceClient

        kwargs = {"provider": IMAGE_PROVIDER, "api_key": HF_TOKEN}
        if HF_BILL_TO:  # omitted → bills the personal account's credits
            kwargs["bill_to"] = HF_BILL_TO
        client = InferenceClient(**kwargs)
        buf = io.BytesIO()
        input_image.convert("RGB").save(buf, format="PNG")
        image = client.image_to_image(
            buf.getvalue(), prompt=dish_prompt(dish, ingredients), model=IMAGE_MODEL)
        return image, f"🎨 {IMAGE_MODEL} · {IMAGE_PROVIDER}"
    except Exception as exc:
        return None, f"Image unavailable ({type(exc).__name__}: {exc})."