app/models/generator.py

"""
AI template generation using Stable Diffusion / Flux API.
"""

import httpx
import base64
import asyncio
from io import BytesIO

import cv2
import numpy as np
from PIL import Image


class TemplateGenerator:
    """Generate coloring page templates from text prompts."""

    def __init__(
        self,
        api_key: str = "",
        provider: str = "replicate",
        base_url: str = "",
    ):
        self.api_key = api_key
        self.provider = provider
        self.base_url = base_url or self._default_url()

    def _default_url(self) -> str:
        urls = {
            "replicate": "https://api.replicate.com/v1/predictions",
            "together": "https://api.together.xyz/v1/images/generations",
            "stability": "https://api.stability.ai/v2beta/stable-image/generate/core",
        }
        return urls.get(self.provider, "")

    @property
    def available(self) -> bool:
        return bool(self.api_key)

    async def generate_from_prompt(
        self,
        prompt: str,
        style: str = "coloring_book",
        width: int = 1024,
        height: int = 1024,
    ) -> np.ndarray | None:
        style_suffixes = {
            "coloring_book": (
                "clean line art, coloring book page, black outlines on white "
                "background, no shading, no color fill, simple shapes, "
                "clear boundaries between regions, vector art style"
            ),
            "realistic": (
                "detailed illustration, clear edges, distinct color regions, "
                "flat colors, no gradients, poster style"
            ),
            "cartoon": (
                "cartoon style, bold black outlines, flat colors, "
                "simple shapes, children's coloring book"
            ),
            "mandala": (
                "mandala pattern, symmetrical, intricate line art, "
                "black outlines on white, no fill colors"
            ),
            "pixel": (
                "pixel art style, clear grid, distinct color blocks, "
                "retro game aesthetic, flat colors"
            ),
        }

        enhanced_prompt = f"{prompt}, {style_suffixes.get(style, style_suffixes['coloring_book'])}"

        try:
            if self.provider == "replicate":
                return await self._generate_replicate(enhanced_prompt, width, height)
            elif self.provider == "together":
                return await self._generate_together(enhanced_prompt, width, height)
            elif self.provider == "stability":
                return await self._generate_stability(enhanced_prompt, width, height)
        except Exception as e:
            print(f"Generation failed: {e}")
            return None

    async def _generate_replicate(
        self, prompt: str, width: int, height: int
    ) -> np.ndarray | None:
        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.post(
                self.base_url,
                headers={
                    "Authorization": f"Bearer {self.api_key}",
                    "Content-Type": "application/json",
                },
                json={
                    "version": "39ed52f2a78e934b3ba6e2a89f5b1c712de7dfea535525255b1aa35c5565e08b",
                    "input": {
                        "prompt": prompt,
                        "width": width,
                        "height": height,
                        "num_inference_steps": 28,
                        "guidance_scale": 7.5,
                    },
                },
            )
            prediction = resp.json()
            poll_url = prediction.get("urls", {}).get("get", "")
            for _ in range(60):
                await asyncio.sleep(2)
                status_resp = await client.get(
                    poll_url,
                    headers={"Authorization": f"Bearer {self.api_key}"},
                )
                status = status_resp.json()
                if status["status"] == "succeeded":
                    image_url = status["output"][0]
                    img_resp = await client.get(image_url)
                    return self._bytes_to_cv2(img_resp.content)
                elif status["status"] == "failed":
                    return None
        return None

    async def _generate_together(
        self, prompt: str, width: int, height: int
    ) -> np.ndarray | None:
        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.post(
                self.base_url,
                headers={
                    "Authorization": f"Bearer {self.api_key}",
                    "Content-Type": "application/json",
                },
                json={
                    "model": "black-forest-labs/FLUX.1-schnell",
                    "prompt": prompt,
                    "width": width,
                    "height": height,
                    "n": 1,
                    "response_format": "b64_json",
                },
            )
            data = resp.json()
            b64 = data["data"][0]["b64_json"]
            img_bytes = base64.b64decode(b64)
            return self._bytes_to_cv2(img_bytes)

    async def _generate_stability(
        self, prompt: str, width: int, height: int
    ) -> np.ndarray | None:
        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.post(
                self.base_url,
                headers={
                    "Authorization": f"Bearer {self.api_key}",
                    "Accept": "image/png",
                },
                files={"none": ""},
                data={
                    "prompt": prompt,
                    "output_format": "png",
                    "aspect_ratio": f"{width}:{height}",
                },
            )
            if resp.status_code == 200:
                return self._bytes_to_cv2(resp.content)
        return None

    def _bytes_to_cv2(self, img_bytes: bytes) -> np.ndarray:
        nparr = np.frombuffer(img_bytes, np.uint8)
        return cv2.imdecode(nparr, cv2.IMREAD_COLOR)