imagepipeline/imagepipeline/modules/openrouter_edit.py

from __future__ import annotations

import base64
import io
import json
import os
import urllib.error
import urllib.request
from pathlib import Path

from PIL import Image

from imagepipeline.core.context import ModuleContext
from imagepipeline.core.exceptions import DependencyError
from imagepipeline.core.params import Param
from imagepipeline.modules.ai_base import AIModule
from imagepipeline.modules.registry import register

OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
TEMPLATE_PROMPT_PREFIX = (
    "You are given two images. The FIRST image is a style reference from an existing "
    "gallery. The SECOND image is the photo to edit. "
)
# Cap reference uploads so multi-image requests stay within API limits.
TEMPLATE_API_MAX_EDGE = 1536


@register
class OpenRouterEditModule(AIModule):
    name = "openrouter_edit"
    description = "Generative image editing via OpenRouter (cloud)"

    @classmethod
    def parameters(cls) -> dict[str, Param]:
        params = dict(super().parameters())
        params.update(
            {
                "prompt": Param(
                    "string",
                    required=True,
                    help="Edit instruction for the model",
                ),
                "model": Param(
                    "string",
                    default="black-forest-labs/flux.2-klein-4b",
                    help="OpenRouter model id with image output",
                ),
                "strength": Param(
                    "float",
                    default=0.3,
                    help="Edit strength (image_config.strength where supported)",
                ),
                "template_image": Param(
                    "path",
                    default=None,
                    help=(
                        "Optional style-reference image (e.g. existing gallery player). "
                        "Sent as the first image when set."
                    ),
                ),
                "api_key_env": Param(
                    "string",
                    default="OPENROUTER_API_KEY",
                    help="Environment variable containing the API key",
                ),
            }
        )
        return params

    @classmethod
    def check_dependencies(cls) -> None:
        if not os.environ.get("OPENROUTER_API_KEY"):
            raise DependencyError("OPENROUTER_API_KEY environment variable is not set")

    def run(self, ctx: ModuleContext) -> None:
        api_key_env = ctx.params["api_key_env"]
        api_key = os.environ.get(api_key_env)
        if not api_key:
            raise DependencyError(f"Environment variable {api_key_env!r} is not set")

        prompt = ctx.params["prompt"]
        model = ctx.params["model"]
        strength = ctx.params["strength"]
        template_path = ctx.params["template_image"]
        template_data_url: str | None = None
        if template_path is not None:
            template_path = Path(template_path)
            if not template_path.is_file():
                raise FileNotFoundError(f"Template image not found: {template_path}")
            template_data_url = self._path_to_data_url(
                template_path, max_edge=TEMPLATE_API_MAX_EDGE
            )

        def process(src: Path, dst: Path, index: int, total: int) -> None:
            with Image.open(src) as image:
                megapixels = (image.size[0] * image.size[1]) / 1_000_000
            if ctx.logger is not None:
                ctx.logger.info(
                    f"  OpenRouter request [{index}/{total}]: model={model!r}, "
                    f"~{megapixels:.1f} MP (cost varies by model)"
                )
            source_data_url = self._path_to_data_url(src, max_edge=0)
            payload = self._build_payload(
                source_data_url,
                prompt,
                model,
                strength,
                template_data_url=template_data_url,
            )
            response = self._post(api_key, payload)
            result_bytes = self._extract_image_bytes(response)
            dst.parent.mkdir(parents=True, exist_ok=True)
            self._save_result_matching_source(src, result_bytes, dst)

        self.iter_input_images(ctx, process)

    @classmethod
    def _modalities_for_model(cls, model: str) -> list[str]:
        if "gemini" in model.lower():
            return ["image", "text"]
        return ["image"]

    @classmethod
    def _strength_supported(cls, model: str) -> bool:
        lowered = model.lower()
        return "recraft" in lowered or "flux" in lowered

    @classmethod
    def _path_to_data_url(cls, path: Path, *, max_edge: int) -> str:
        with Image.open(path) as image:
            if max_edge > 0:
                width, height = image.size
                long_edge = max(width, height)
                if long_edge > max_edge:
                    scale = max_edge / long_edge
                    image = image.resize(
                        (max(1, int(width * scale)), max(1, int(height * scale))),
                        Image.Resampling.LANCZOS,
                    )
            rgb = image.convert("RGB")
            buffer = io.BytesIO()
            rgb.save(buffer, format="JPEG", quality=90)
            encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
        return f"data:image/jpeg;base64,{encoded}"

    @classmethod
    def _build_payload(
        cls,
        source_data_url: str,
        prompt: str,
        model: str,
        strength: float,
        *,
        template_data_url: str | None = None,
    ) -> dict:
        full_prompt = f"{TEMPLATE_PROMPT_PREFIX}{prompt}" if template_data_url else prompt
        content: list[dict] = [{"type": "text", "text": full_prompt}]
        if template_data_url is not None:
            content.append({"type": "image_url", "image_url": {"url": template_data_url}})
        content.append({"type": "image_url", "image_url": {"url": source_data_url}})
        payload: dict = {
            "model": model,
            "modalities": cls._modalities_for_model(model),
            "messages": [{"role": "user", "content": content}],
        }
        if strength is not None and cls._strength_supported(model):
            payload["image_config"] = {"strength": strength}
        return payload

    @classmethod
    def _save_result_matching_source(cls, source: Path, result_bytes: bytes, dest: Path) -> None:
        with Image.open(source) as original:
            orig_format = original.format
            orig_size = original.size
            orig_mode = original.mode
            orig_alpha = original.getchannel("A") if "A" in original.getbands() else None

        with Image.open(io.BytesIO(result_bytes)) as edited:
            if edited.size != orig_size:
                edited = edited.resize(orig_size, Image.Resampling.LANCZOS)
            if orig_alpha is not None:
                edited = edited.convert("RGB").convert("RGBA")
                edited.putalpha(orig_alpha)
            elif orig_mode not in ("RGB", "RGBA"):
                edited = edited.convert(orig_mode)

            save_format = orig_format
            if not save_format:
                suffix = dest.suffix.lower().lstrip(".")
                save_format = {"jpg": "JPEG", "jpeg": "JPEG"}.get(suffix, suffix.upper())

            save_kwargs: dict = {}
            if save_format == "JPEG":
                if edited.mode == "RGBA":
                    edited = edited.convert("RGB")
                save_kwargs["quality"] = 95
            elif save_format == "PNG" and edited.mode not in ("RGBA", "RGB", "P"):
                edited = edited.convert("RGBA")

            dest.parent.mkdir(parents=True, exist_ok=True)
            edited.save(dest, format=save_format, **save_kwargs)

    @staticmethod
    def _post(api_key: str, payload: dict) -> dict:
        body = json.dumps(payload).encode("utf-8")
        request = urllib.request.Request(
            OPENROUTER_URL,
            data=body,
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
                "HTTP-Referer": "https://github.com/froxxxy/imagepipeline",
                "X-Title": "imagepipeline",
            },
            method="POST",
        )
        try:
            with urllib.request.urlopen(request, timeout=600) as response:
                return json.loads(response.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            detail = exc.read().decode("utf-8", errors="replace")
            raise RuntimeError(f"OpenRouter API error ({exc.code}): {detail}") from exc

    @staticmethod
    def _extract_image_bytes(response: dict) -> bytes:
        choices = response.get("choices") or []
        if not choices:
            raise RuntimeError("OpenRouter response contained no choices")
        message = choices[0].get("message") or {}
        images = message.get("images") or []
        if not images:
            raise RuntimeError("OpenRouter response contained no images")
        url = images[0].get("image_url", {}).get("url", "")
        if url.startswith("data:"):
            _, encoded = url.split(",", 1)
            return base64.b64decode(encoded)
        if url.startswith("http://") or url.startswith("https://"):
            with urllib.request.urlopen(url, timeout=120) as image_response:
                return image_response.read()
        raise RuntimeError(f"Unsupported image URL in OpenRouter response: {url!r}")