| 1 | from __future__ import annotations |
| 2 | |
| 3 | import asyncio |
| 4 | import base64 |
| 5 | import json |
| 6 | import os |
| 7 | from io import BytesIO |
| 8 | from typing import Any, List |
| 9 | |
| 10 | import aiohttp |
| 11 | from PIL import Image |
| 12 | from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential |
| 13 | |
| 14 | from interfaces.image_output import ImageOutput |
| 15 | from tools.image_orientation import ensure_not_portrait, landscape_guard_requested |
| 16 | from utils.image import image_path_to_b64 |
| 17 | from utils.rate_limiter import RateLimiter |
| 18 | from utils.retry import after_func |
| 19 | |
| 20 | |
| 21 | class OpenRouterImageAPIError(RuntimeError): |
| 22 | def __init__(self, status_code: int, payload: Any) -> None: |
| 23 | self.status_code = status_code |
| 24 | super().__init__(f"OpenRouter image generation failed with HTTP {status_code}: {payload}") |
| 25 | |
| 26 | |
| 27 | def _request_timeout_seconds() -> float: |
| 28 | raw = os.environ.get("VIMAX_IMAGE_REQUEST_TIMEOUT_SECONDS", "300") |
| 29 | try: |
| 30 | return max(1.0, float(raw)) |
| 31 | except ValueError: |
| 32 | return 300.0 |
| 33 | |
| 34 | |
| 35 | def _is_retryable_image_error(exc: BaseException) -> bool: |
| 36 | if isinstance(exc, OpenRouterImageAPIError): |
| 37 | return exc.status_code in {408, 409, 425, 429} or exc.status_code >= 500 |
| 38 | if isinstance(exc, (aiohttp.ClientError, asyncio.TimeoutError)): |
| 39 | return True |
| 40 | return isinstance(exc, ValueError) and "portrait-oriented" in str(exc) |
| 41 | |
| 42 | |
| 43 | class ImageGeneratorOpenRouterAPI: |
| 44 | """Generate images through OpenRouter's dedicated Images API.""" |
| 45 | |
| 46 | def __init__( |
| 47 | self, |
| 48 | api_key: str, |
| 49 | model: str = "openai/gpt-image-2", |
| 50 | base_url: str = "https://openrouter.ai/api/v1", |
| 51 | quality: str = "auto", |
| 52 | background: str = "auto", |
| 53 | output_compression: int | None = None, |
| 54 | rate_limiter: RateLimiter | None = None, |
| 55 | http_referer: str = "", |
| 56 | app_title: str = "ViMax", |
| 57 | ) -> None: |
| 58 | self.api_key = api_key |
| 59 | self.model = model |
| 60 | self.base_url = base_url.rstrip("/") |
| 61 | self.quality = quality |
| 62 | self.background = background |
| 63 | self.output_compression = output_compression |
| 64 | self.rate_limiter = rate_limiter |
| 65 | self.http_referer = http_referer |
| 66 | self.app_title = app_title |
| 67 | |
| 68 | @retry( |
| 69 | stop=stop_after_attempt(3), |
| 70 | wait=wait_exponential(multiplier=1, min=1, max=10), |
| 71 | retry=retry_if_exception(_is_retryable_image_error), |
| 72 | after=after_func, |
| 73 | reraise=True, |
| 74 | ) |
| 75 | async def generate_single_image( |
| 76 | self, |
| 77 | prompt: str, |
| 78 | reference_image_paths: List[str] | None = None, |
| 79 | aspect_ratio: str | None = "16:9", |
| 80 | **kwargs: Any, |
| 81 | ) -> ImageOutput: |
| 82 | references = list(reference_image_paths or []) |
| 83 | if len(references) > 16: |
| 84 | raise ValueError("OpenRouter GPT Image supports at most 16 reference images") |
| 85 | if self.rate_limiter is not None: |
| 86 | await self.rate_limiter.acquire() |
| 87 | |
| 88 | enforce_landscape = landscape_guard_requested( |
| 89 | size=kwargs.get("size"), |
| 90 | aspect_ratio=aspect_ratio, |
| 91 | enforce_landscape=kwargs.get("enforce_landscape", True), |
| 92 | allow_portrait=kwargs.get("allow_portrait", False), |
| 93 | ) |
| 94 | request_prompt = _prompt_with_landscape_requirement(prompt, aspect_ratio) if enforce_landscape else prompt |
| 95 | payload: dict[str, Any] = { |
| 96 | "model": self.model, |
| 97 | "prompt": request_prompt, |
| 98 | "n": 1, |
| 99 | "quality": kwargs.get("quality", self.quality), |
| 100 | "background": kwargs.get("background", self.background), |
| 101 | } |
| 102 | compression = kwargs.get("output_compression", self.output_compression) |
| 103 | if compression is not None: |
| 104 | payload["output_compression"] = compression |
| 105 | if references: |
| 106 | payload["input_references"] = [ |
| 107 | {"type": "image_url", "image_url": {"url": image_path_to_b64(path, mime=True)}} |
| 108 | for path in references |
| 109 | ] |
| 110 | |
| 111 | progress = kwargs.get("progress") |
| 112 | _emit_progress( |
| 113 | progress, |
| 114 | "image_generation", |
| 115 | f"Generating image with {self.model}", |
| 116 | {"model": self.model, "reference_count": len(references)}, |
| 117 | ) |
| 118 | timeout = aiohttp.ClientTimeout(total=_request_timeout_seconds()) |
| 119 | status, response = await _post_json( |
| 120 | f"{self.base_url}/images", |
| 121 | headers=self._headers(), |
| 122 | payload=payload, |
| 123 | timeout=timeout, |
| 124 | ) |
| 125 | if status >= 400: |
| 126 | raise OpenRouterImageAPIError(status, response) |
| 127 | |
| 128 | image, extension = _decode_image_response(response) |
| 129 | if enforce_landscape: |
| 130 | ensure_not_portrait(image) |
| 131 | _emit_progress( |
| 132 | progress, |
| 133 | "image_completed", |
| 134 | "OpenRouter image generation completed", |
| 135 | {"model": self.model, "width": image.width, "height": image.height}, |
| 136 | ) |
| 137 | return ImageOutput(fmt="pil", ext=extension, data=image) |
| 138 | |
| 139 | def _headers(self) -> dict[str, str]: |
| 140 | headers = { |
| 141 | "Authorization": f"Bearer {self.api_key}", |
| 142 | "Content-Type": "application/json", |
| 143 | } |
| 144 | if self.http_referer: |
| 145 | headers["HTTP-Referer"] = self.http_referer |
| 146 | if self.app_title: |
| 147 | headers["X-OpenRouter-Title"] = self.app_title |
| 148 | return headers |
| 149 | |
| 150 | |
| 151 | def _prompt_with_landscape_requirement(prompt: str, aspect_ratio: str | None) -> str: |
| 152 | ratio = aspect_ratio or "16:9" |
| 153 | return f"{prompt}\n\nComposition requirement: create a landscape image with an approximate {ratio} aspect ratio; the width must be greater than the height." |
| 154 | |
| 155 | |
| 156 | def _decode_image_response(payload: Any) -> tuple[Image.Image, str]: |
| 157 | data = payload.get("data") if isinstance(payload, dict) else None |
| 158 | item = data[0] if isinstance(data, list) and data and isinstance(data[0], dict) else None |
| 159 | encoded = item.get("b64_json") if item else None |
| 160 | if not isinstance(encoded, str) or not encoded: |
| 161 | raise ValueError(f"OpenRouter image response missing data[0].b64_json: {payload}") |
| 162 | if encoded.startswith("data:"): |
| 163 | encoded = encoded.split(",", 1)[-1] |
| 164 | try: |
| 165 | raw = base64.b64decode(encoded, validate=True) |
| 166 | with Image.open(BytesIO(raw)) as opened: |
| 167 | opened.load() |
| 168 | image = opened.copy() |
| 169 | except Exception as exc: |
| 170 | raise ValueError("OpenRouter image response contained invalid image data") from exc |
| 171 | media_type = item.get("media_type", "image/png") |
| 172 | extension = {"image/jpeg": "jpg", "image/webp": "webp"}.get(media_type, "png") |
| 173 | return image, extension |
| 174 | |
| 175 | |
| 176 | def _emit_progress(progress: Any, stage: str, message: str, metadata: dict[str, Any]) -> None: |
| 177 | if progress is not None: |
| 178 | progress(stage, message, metadata) |
| 179 | |
| 180 | |
| 181 | async def _post_json( |
| 182 | url: str, |
| 183 | *, |
| 184 | headers: dict[str, str], |
| 185 | payload: dict[str, Any], |
| 186 | timeout: aiohttp.ClientTimeout, |
| 187 | ) -> tuple[int, Any]: |
| 188 | async with aiohttp.ClientSession(timeout=timeout) as session: |
| 189 | async with session.post(url, headers=headers, json=payload) as response: |
| 190 | text = await response.text() |
| 191 | try: |
| 192 | body = json.loads(text) |
| 193 | except json.JSONDecodeError: |
| 194 | body = {"message": text} |
| 195 | return response.status, body |
| 196 |