image_client.py

"""Image generation client via the Hugging Face Inference API (text_to_image)."""
from typing import Optional
from PIL import Image

try:
    from huggingface_hub import InferenceClient
except ImportError:
    InferenceClient = None

import config


class ImageGenerationClient:
    """Generate artifact images for alternate timelines via the HF Inference API."""

    def __init__(self, hf_token: Optional[str] = None):
        self.hf_token = hf_token or config.HF_TOKEN
        self.model = config.MODEL_NAME_IMAGE
        self.fallback_model = config.FALLBACK_IMAGE_MODEL
        self.client = None
        self.available = False

        if not config.ENABLE_IMAGE_GEN:
            return
        if not InferenceClient:
            print("[warn]huggingface_hub not installed; image generation disabled")
            return
        if not self.hf_token:
            print("[warn]HF_TOKEN not set; image generation disabled")
            return

        try:
            self.client = InferenceClient(token=self.hf_token, timeout=config.TIMEOUT_IMAGE)
            self.available = True
        except Exception as e:
            print(f"[warn]Image client initialization failed: {e}")
            self.available = False

    def generate(self, prompt: str) -> Optional[Image.Image]:
        """Generate an image from a prompt, with a fallback model."""
        if not self.available:
            return None

        try:
            return self.client.text_to_image(prompt, model=self.model)
        except Exception as e:
            print(f"[warn]Primary image generation failed ({self.model}): {e}")

        try:
            return self.client.text_to_image(prompt, model=self.fallback_model)
        except Exception as e:
            print(f"[warn]Fallback image generation failed ({self.fallback_model}): {e}")

        return None

    def generate_artifact_newspaper(self, headline: str, subheading: str) -> Optional[Image.Image]:
        """Generate a vintage newspaper artifact."""
        prompt = (
            f'Vintage newspaper front page with headline "{headline}" and subheading '
            f'"{subheading}". Aged paper, period typography, realistic newspaper layout.'
        )
        return self.generate(prompt)

    def generate_artifact_product(self, product: str, context: str) -> Optional[Image.Image]:
        """Generate a product advertisement or packaging."""
        prompt = (
            f'Vintage product advertisement for "{product}" in a world where {context}. '
            f"Period-appropriate packaging design, realistic product mockup, authentic styling."
        )
        return self.generate(prompt)

    def generate_artifact_document(self, doc_type: str, context: str) -> Optional[Image.Image]:
        """Generate a historical document artifact."""
        prompt = (
            f"Historical {doc_type} from an alternate timeline where {context}. "
            f"Aged paper, period authentic, realistic museum quality."
        )
        return self.generate(prompt)


# Global client instance
_image_client = None


def get_image_client() -> ImageGenerationClient:
    """Get or create the global image generation client."""
    global _image_client
    if _image_client is None:
        _image_client = ImageGenerationClient()
    return _image_client


def generate_artifact(artifact_type: str, **kwargs) -> Optional[Image.Image]:
    """Generate an artifact image by type."""
    client = get_image_client()

    if artifact_type == "newspaper":
        return client.generate_artifact_newspaper(
            kwargs.get("headline", "Alternate History"),
            kwargs.get("subheading", "A world that never was"),
        )
    elif artifact_type == "product":
        return client.generate_artifact_product(
            kwargs.get("product", "Product"),
            kwargs.get("context", "divergence occurred"),
        )
    elif artifact_type == "document":
        return client.generate_artifact_document(
            kwargs.get("doc_type", "document"),
            kwargs.get("context", "divergence occurred"),
        )

    return None