pipeline/images.py

"""Step 6: turn each [IMAGE: ...] marker into a FLUX prompt and render it.

Images are produced entirely by **remote HF Inference Provider calls** (FLUX.1-schnell)
billed to the user's token — this Space has no GPU, so nothing is generated locally.
The render tries auto provider routing first, then falls back across the providers that
serve the model, so a single provider being unavailable for the user's token doesn't
break the run.
"""
from __future__ import annotations

from pathlib import Path
from typing import List, Optional, Tuple

from huggingface_hub import InferenceClient

from . import config, llm

_PROMPT_SYSTEM = (
    "You are a prompt engineer for the FLUX text-to-image model. Given a short scene "
    "description for a blog illustration and the article topic, write ONE vivid, concrete "
    "image prompt (single line, <60 words). Describe subject, setting, composition, "
    "lighting and style. Prefer clean, editorial, photographic or tasteful illustrative "
    "styles suitable for a professional blog. No text/words in the image. Return only the prompt."
)


def _flux_prompt(client: InferenceClient, topic: str, scene: str) -> str:
    try:
        p = llm.chat(
            client,
            config.MODEL_REASONING,
            _PROMPT_SYSTEM,
            f"Article topic: {topic}\nScene: {scene}\nWrite the FLUX prompt.",
            max_tokens=150,
            temperature=0.8,
        )
        p = p.strip().strip('"')
        return p or scene
    except Exception:
        return f"{scene}, editorial photography, clean composition, natural lighting"


def _render(hf_token: str, prompt: str) -> Tuple[Optional[object], Optional[str]]:
    """Generate one image via Inference Providers, trying auto then explicit providers.

    Returns (PIL image, None) on success or (None, error message) if every provider fails.
    """
    # None => let the router auto-select; then try each known provider explicitly.
    attempts: List[Optional[str]] = [None] + config.IMAGE_PROVIDERS
    errors: List[str] = []
    for provider in attempts:
        try:
            client = (
                InferenceClient(token=hf_token, provider=provider)
                if provider
                else InferenceClient(token=hf_token)
            )
            image = client.text_to_image(prompt=prompt, model=config.MODEL_IMAGE)
            return image, None
        except Exception as e:  # noqa: BLE001 - try the next provider
            errors.append(f"{provider or 'auto'}: {e}")
            continue
    return None, " | ".join(errors[-3:])


def generate_images(
    client: InferenceClient,
    hf_token: str,
    topic: str,
    scenes: List[str],
    run_dir: Path,
) -> List[dict]:
    """Render one image per scene via inference calls. Returns [{scene, prompt, path|None, error?}]."""
    run_dir.mkdir(parents=True, exist_ok=True)
    out: List[dict] = []
    for i, scene in enumerate(scenes):
        prompt = _flux_prompt(client, topic, scene)
        item = {"scene": scene, "prompt": prompt, "path": None}
        image, err = _render(hf_token, prompt)
        if image is not None:
            try:
                path = run_dir / f"image_{i + 1}.png"
                image.save(path)
                item["path"] = str(path)
            except Exception as e:  # noqa: BLE001
                item["error"] = f"save failed: {e}"
        else:
            item["error"] = err or "image generation failed"
        out.append(item)
    return out