itsulu-blog-publisher/image_router.py

# -*- coding: utf-8 -*-
"""
Image Router — itsulu_blog_publisher
======================================
Independent from the text LLM router. Generates a cover image for the blog
post and returns it as base64-encoded bytes for storage as an ir.attachment.

Supported image providers:
- openai_dalle  : DALL·E 3 (via OpenAI /v1/images/generations)
- google_imagen : Google Imagen (via Gemini API — requires Gemini API key)
- stable_diff   : Stable Diffusion (via Ollama / local Automatic1111 / SDXL)
- none          : Skip image generation (default if not configured)

The image prompt is built from the blog post title and meta keywords so that
the cover image is always contextually relevant.
"""
import base64
import logging
import requests
from odoo.exceptions import UserError

_logger = logging.getLogger(__name__)

SUPPORTED_IMAGE_PROVIDERS = ('openai_dalle', 'google_imagen', 'stable_diff', 'none')


def build_image_prompt(title: str, keywords: str) -> str:
    """
    Construct a cover image prompt from title + keywords.
    Professional, editorial style — no faces, no text, no brand logos.
    """
    kw_fragment = f" Related concepts: {keywords}." if keywords else ''
    return (
        f"Professional editorial cover image for a technology blog post titled: '{title}'."
        f"{kw_fragment} "
        "Modern flat design illustration, clean composition, vibrant but professional colors. "
        "No human faces. No text or words in the image. No brand logos."
    )


class ImageRouter:
    """
    Usage::

        router = ImageRouter(env, provider='openai_dalle')
        result = router.generate_cover(title='AI Governance 2026', keywords='AI, enterprise')
        # result is None (if provider='none') or ImageResult(b64_data, mime_type)
    """

    def __init__(self, env, provider: str):
        self.env = env
        self.provider = (provider or 'none').lower().strip()
        if self.provider not in SUPPORTED_IMAGE_PROVIDERS:
            raise UserError(
                f"Image provider '{self.provider}' is not supported. "
                f"Choose one of: {', '.join(SUPPORTED_IMAGE_PROVIDERS)}."
            )

    def _get_param(self, key: str) -> str:
        return self.env['ir.config_parameter'].sudo().get_param(
            f'itsulu_blog_publisher.{key}', default=''
        )

    def generate_cover(self, title: str, keywords: str = '') -> 'ImageResult | None':
        """
        Generate a cover image. Returns None when provider='none'.
        :raises UserError: on API failure.
        """
        if self.provider == 'none':
            return None

        prompt = build_image_prompt(title=title, keywords=keywords)
        _logger.info("ImageRouter.generate_cover: provider=%s", self.provider)

        if self.provider == 'openai_dalle':
            return self._dalle(prompt)
        elif self.provider == 'google_imagen':
            return self._imagen(prompt)
        elif self.provider == 'stable_diff':
            return self._stable_diff(prompt)

        return None

    # ------------------------------------------------------------------ #
    # DALL·E 3                                                             #
    # ------------------------------------------------------------------ #

    def _dalle(self, prompt: str) -> 'ImageResult':
        api_key = self._get_param('openai_api_key')
        if not api_key:
            raise UserError(
                "OpenAI API key is required for DALL·E image generation. "
                "Set it in Settings → Blog Publisher."
            )

        payload = {
            'model': 'dall-e-3',
            'prompt': prompt,
            'n': 1,
            'size': '1792x1024',       # landscape for blog cover
            'response_format': 'b64_json',
            'quality': 'standard',
        }

        try:
            resp = requests.post(
                'https://api.openai.com/v1/images/generations',
                headers={'Authorization': f'Bearer {api_key}'},
                json=payload,
                timeout=120,
            )
        except requests.RequestException as exc:
            raise UserError(f"DALL·E 3 network error: {exc}") from exc

        if not resp.ok:
            try:
                err = resp.json().get('error', {}).get('message', resp.text[:300])
            except Exception:
                err = resp.text[:300]
            raise UserError(f"DALL·E 3 error {resp.status_code}: {err}")

        data = resp.json()
        b64 = data['data'][0]['b64_json']
        return ImageResult(b64_data=b64, mime_type='image/png')

    # ------------------------------------------------------------------ #
    # Google Imagen (via Gemini API)                                       #
    # ------------------------------------------------------------------ #

    def _imagen(self, prompt: str) -> 'ImageResult':
        api_key = self._get_param('gemini_api_key')
        if not api_key:
            raise UserError(
                "Google Gemini API key is required for Imagen. "
                "Set it in Settings → Blog Publisher."
            )

        url = (
            f"https://generativelanguage.googleapis.com/v1beta/models/"
            f"imagen-3.0-generate-002:predict?key={api_key}"
        )
        payload = {
            'instances': [{'prompt': prompt}],
            'parameters': {
                'sampleCount': 1,
                'aspectRatio': '16:9',
            },
        }

        try:
            resp = requests.post(url, json=payload, timeout=120)
        except requests.RequestException as exc:
            raise UserError(f"Google Imagen network error: {exc}") from exc

        if not resp.ok:
            try:
                err = resp.json().get('error', {}).get('message', resp.text[:300])
            except Exception:
                err = resp.text[:300]
            raise UserError(f"Google Imagen error {resp.status_code}: {err}")

        predictions = resp.json().get('predictions', [])
        if not predictions:
            raise UserError("Google Imagen returned no image predictions.")

        b64 = predictions[0].get('bytesBase64Encoded', '')
        if not b64:
            raise UserError("Google Imagen returned an empty image.")

        return ImageResult(b64_data=b64, mime_type='image/png')

    # ------------------------------------------------------------------ #
    # Stable Diffusion (Ollama / Automatic1111 / ComfyUI)                 #
    # ------------------------------------------------------------------ #

    def _stable_diff(self, prompt: str) -> 'ImageResult':
        base_url = self._get_param('ollama_base_url')
        if not base_url:
            raise UserError(
                "Ollama / Stable Diffusion base URL is required for image generation. "
                "Set it in Settings → Blog Publisher."
            )

        # Try Ollama's native image generation endpoint (Ollama >= 0.5 supports it
        # for vision/image models like llava, bakllava, and SD models via ollama pull)
        url = f"{base_url.rstrip('/')}/api/generate"
        payload = {
            'model': self._get_param('image_ollama_model') or 'stable-diffusion',
            'prompt': prompt,
            'stream': False,
        }

        try:
            resp = requests.post(url, json=payload, timeout=300)
        except requests.RequestException as exc:
            raise UserError(f"Stable Diffusion (Ollama) network error: {exc}") from exc

        if not resp.ok:
            raise UserError(
                f"Stable Diffusion returned HTTP {resp.status_code}: {resp.text[:300]}"
            )

        data = resp.json()
        images = data.get('images', [])
        if not images:
            raise UserError(
                "Stable Diffusion returned no images. "
                "Check that a compatible image model is loaded in Ollama."
            )

        return ImageResult(b64_data=images[0], mime_type='image/png')


class ImageResult:
    """Holds the generated image as base64 + mime type."""

    def __init__(self, b64_data: str, mime_type: str = 'image/png'):
        self.b64_data = b64_data
        self.mime_type = mime_type

    def to_bytes(self) -> bytes:
        return base64.b64decode(self.b64_data)

    def to_odoo_attachment_vals(self, name: str = 'cover_image.png') -> dict:
        """Return vals dict for ir.attachment.create()."""
        return {
            'name': name,
            'datas': self.b64_data,
            'mimetype': self.mime_type,
        }