| 1 | #!/usr/bin/env python3 |
| 2 | """ |
| 3 | OpenAI Compatible Image Generation Backend |
| 4 | |
| 5 | Generates images via OpenAI-compatible APIs (OpenAI, local models like Qwen-Image, etc.). |
| 6 | Used by image_gen.py as a backend module. |
| 7 | |
| 8 | Configuration keys: |
| 9 | OPENAI_API_KEY (required) API key |
| 10 | OPENAI_BASE_URL (optional) Custom API endpoint (e.g. http://127.0.0.1:3000/v1) |
| 11 | OPENAI_MODEL (optional) Model name (default: gpt-image-2) |
| 12 | OPENAI_SIZE_PRESET (optional) auto, legacy, gpt-image, gpt-image-2, or dall-e-2 |
| 13 | OPENAI_RESPONSE_FORMAT (optional) auto, b64_json, url, or omit |
| 14 | OPENAI_QUALITY (optional) auto, omit, low, medium, high, standard, or hd |
| 15 | OPENAI_OUTPUT_FORMAT (optional) png, jpeg, or webp for GPT image models |
| 16 | OPENAI_OUTPUT_COMPRESSION (optional) 0-100, only for jpeg/webp GPT image output |
| 17 | OPENAI_BACKGROUND (optional) auto or opaque for gpt-image-2 |
| 18 | OPENAI_MODERATION (optional) auto or low for GPT image models |
| 19 | |
| 20 | Image editing (image-to-image): |
| 21 | When image_gen.py passes reference_image=<path> (single-image CLI only, |
| 22 | via --reference-image), this backend calls /v1/images/edits with the source |
| 23 | image + the prompt as an edit instruction, instead of /v1/images/generations. |
| 24 | |
| 25 | Dependencies: |
| 26 | pip install requests Pillow |
| 27 | """ |
| 28 | |
| 29 | import sys |
| 30 | from pathlib import Path |
| 31 | |
| 32 | _SCRIPTS_DIR = Path(__file__).resolve().parents[1] |
| 33 | if str(_SCRIPTS_DIR) not in sys.path: |
| 34 | sys.path.insert(0, str(_SCRIPTS_DIR)) |
| 35 | |
| 36 | from console_encoding import configure_utf8_stdio # noqa: E402 |
| 37 | |
| 38 | configure_utf8_stdio() |
| 39 | |
| 40 | if __name__ == "__main__": |
| 41 | print(__doc__) |
| 42 | print("Use via: python3 skills/ppt-master/scripts/image_gen.py \"prompt\" --backend openai") |
| 43 | raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1) |
| 44 | |
| 45 | import base64 |
| 46 | import mimetypes |
| 47 | import os |
| 48 | import time |
| 49 | import threading |
| 50 | from collections.abc import Mapping |
| 51 | |
| 52 | import requests |
| 53 | from image_backends.backend_common import ( |
| 54 | MAX_RETRIES, |
| 55 | download_image, |
| 56 | http_error, |
| 57 | is_rate_limit_error, |
| 58 | normalize_image_size, |
| 59 | resolve_output_path, |
| 60 | retry_delay, |
| 61 | save_image_bytes, |
| 62 | ) |
| 63 | |
| 64 | |
| 65 | # ╔══════════════════════════════════════════════════════════════════╗ |
| 66 | # ║ Constants ║ |
| 67 | # ╚══════════════════════════════════════════════════════════════════╝ |
| 68 | |
| 69 | # Aspect ratio -> DALL-E 3 / legacy compatible size mapping. |
| 70 | # Unknown OpenAI-compatible models use this table to preserve old behavior. |
| 71 | LEGACY_COMPAT_ASPECT_RATIO_TO_SIZE = { |
| 72 | "1:1": "1024x1024", |
| 73 | "16:9": "1792x1024", |
| 74 | "9:16": "1024x1792", |
| 75 | "3:2": "1536x1024", |
| 76 | "2:3": "1024x1536", |
| 77 | "4:3": "1536x1024", # closest available |
| 78 | "3:4": "1024x1536", # closest available |
| 79 | "4:5": "1024x1024", # fallback to square |
| 80 | "5:4": "1024x1024", # fallback to square |
| 81 | "21:9": "1792x1024", # closest wide format |
| 82 | } |
| 83 | |
| 84 | # GPT Image 1/1.5/mini officially support only square, landscape, portrait, or auto. |
| 85 | GPT_IMAGE_LEGACY_ASPECT_RATIO_TO_SIZE = { |
| 86 | "1:1": "1024x1024", |
| 87 | "16:9": "1536x1024", |
| 88 | "9:16": "1024x1536", |
| 89 | "3:2": "1536x1024", |
| 90 | "2:3": "1024x1536", |
| 91 | "4:3": "1536x1024", |
| 92 | "3:4": "1024x1536", |
| 93 | "4:5": "1024x1536", |
| 94 | "5:4": "1536x1024", |
| 95 | "21:9": "1536x1024", |
| 96 | } |
| 97 | |
| 98 | # GPT Image 2 supports flexible sizes when both edges are multiples of 16, |
| 99 | # the edge ratio is <= 3:1, and the total pixels are within model limits. |
| 100 | GPT_IMAGE_2_SIZES = { |
| 101 | "512px": { |
| 102 | "1:1": "1024x1024", "16:9": "1280x720", "9:16": "720x1280", |
| 103 | "3:2": "1248x832", "2:3": "832x1248", "4:3": "1024x768", |
| 104 | "3:4": "768x1024", "4:5": "896x1120", "5:4": "1120x896", |
| 105 | "21:9": "1280x544", |
| 106 | }, |
| 107 | "1K": { |
| 108 | "1:1": "1024x1024", "16:9": "1280x720", "9:16": "720x1280", |
| 109 | "3:2": "1248x832", "2:3": "832x1248", "4:3": "1024x768", |
| 110 | "3:4": "768x1024", "4:5": "896x1120", "5:4": "1120x896", |
| 111 | "21:9": "1280x544", |
| 112 | }, |
| 113 | "2K": { |
| 114 | "1:1": "2048x2048", "16:9": "2048x1152", "9:16": "1152x2048", |
| 115 | "3:2": "2016x1344", "2:3": "1344x2016", "4:3": "1920x1440", |
| 116 | "3:4": "1440x1920", "4:5": "1600x2000", "5:4": "2000x1600", |
| 117 | "21:9": "2560x1088", |
| 118 | }, |
| 119 | "4K": { |
| 120 | "1:1": "2880x2880", "16:9": "3840x2160", "9:16": "2160x3840", |
| 121 | "3:2": "3520x2352", "2:3": "2352x3520", "4:3": "3264x2448", |
| 122 | "3:4": "2448x3264", "4:5": "2560x3200", "5:4": "3200x2560", |
| 123 | "21:9": "3840x1648", |
| 124 | }, |
| 125 | } |
| 126 | |
| 127 | DALL_E_2_SIZE_BY_IMAGE_SIZE = { |
| 128 | "512px": "512x512", |
| 129 | "1K": "1024x1024", |
| 130 | "2K": "1024x1024", |
| 131 | "4K": "1024x1024", |
| 132 | } |
| 133 | |
| 134 | VALID_ASPECT_RATIOS = list(LEGACY_COMPAT_ASPECT_RATIO_TO_SIZE.keys()) |
| 135 | |
| 136 | # image_size -> quality mapping |
| 137 | IMAGE_SIZE_TO_QUALITY = { |
| 138 | "512px": "low", |
| 139 | "1K": "medium", |
| 140 | "2K": "high", |
| 141 | "4K": "high", |
| 142 | } |
| 143 | |
| 144 | DEFAULT_MODEL = "gpt-image-2" |
| 145 | |
| 146 | GPT_IMAGE_2_MIN_PIXELS = 655_360 |
| 147 | GPT_IMAGE_2_MAX_PIXELS = 8_294_400 |
| 148 | GPT_IMAGE_2_MAX_EDGE = 3840 |
| 149 | GPT_IMAGE_2_MAX_RATIO = 3 |
| 150 | |
| 151 | GPT_IMAGE_OUTPUT_FORMATS = {"png", "jpeg", "webp"} |
| 152 | GPT_IMAGE_OUTPUT_EXTENSIONS = { |
| 153 | "png": ".png", |
| 154 | "jpeg": ".jpg", |
| 155 | "webp": ".webp", |
| 156 | } |
| 157 | OPENAI_SIZE_PRESETS = { |
| 158 | "auto", |
| 159 | "legacy", |
| 160 | "gpt-image", |
| 161 | "gpt-image-legacy", |
| 162 | "gpt-image-2", |
| 163 | "dall-e-2", |
| 164 | "dalle-2", |
| 165 | } |
| 166 | OPENAI_RESPONSE_FORMATS = {"auto", "b64_json", "url", "omit"} |
| 167 | OPENAI_QUALITY_VALUES = { |
| 168 | "auto", |
| 169 | "omit", |
| 170 | "low", |
| 171 | "medium", |
| 172 | "high", |
| 173 | "standard", |
| 174 | "hd", |
| 175 | } |
| 176 | GPT_IMAGE_BACKGROUNDS = {"auto", "opaque", "transparent"} |
| 177 | GPT_IMAGE_MODERATION_VALUES = {"auto", "low"} |
| 178 | DEFAULT_BASE_URL = "https://api.openai.com/v1" |
| 179 | |
| 180 | # Signals to image_gen.py that this backend can accept a reference_image |
| 181 | # (image-to-image edit via /v1/images/edits). Other backends omit this marker. |
| 182 | SUPPORTS_REFERENCE_IMAGE = True |
| 183 | |
| 184 | |
| 185 | def _field(value, name: str): |
| 186 | """Read a response field from either an SDK object or a dict.""" |
| 187 | if isinstance(value, Mapping): |
| 188 | return value.get(name) |
| 189 | return getattr(value, name, None) |
| 190 | |
| 191 | |
| 192 | def _normalized_model(model: str) -> str: |
| 193 | return (model or "").strip().lower() |
| 194 | |
| 195 | |
| 196 | def _is_gpt_image_model(model: str) -> bool: |
| 197 | return _normalized_model(model).startswith("gpt-image-") |
| 198 | |
| 199 | |
| 200 | def _is_gpt_image_2(model: str) -> bool: |
| 201 | return _normalized_model(model).startswith("gpt-image-2") |
| 202 | |
| 203 | |
| 204 | def _is_dall_e_2(model: str) -> bool: |
| 205 | return _normalized_model(model) == "dall-e-2" |
| 206 | |
| 207 | |
| 208 | def _parse_size(size: str) -> tuple[int, int]: |
| 209 | try: |
| 210 | width_s, height_s = size.lower().split("x", 1) |
| 211 | return int(width_s), int(height_s) |
| 212 | except Exception as exc: |
| 213 | raise ValueError(f"Invalid image size '{size}'. Expected WIDTHxHEIGHT.") from exc |
| 214 | |
| 215 | |
| 216 | def _validate_gpt_image_2_size(size: str) -> None: |
| 217 | width, height = _parse_size(size) |
| 218 | total_pixels = width * height |
| 219 | long_edge = max(width, height) |
| 220 | short_edge = min(width, height) |
| 221 | |
| 222 | errors = [] |
| 223 | if long_edge > GPT_IMAGE_2_MAX_EDGE: |
| 224 | errors.append(f"max edge {long_edge}px exceeds {GPT_IMAGE_2_MAX_EDGE}px") |
| 225 | if width % 16 != 0 or height % 16 != 0: |
| 226 | errors.append("both edges must be multiples of 16px") |
| 227 | if long_edge / short_edge > GPT_IMAGE_2_MAX_RATIO: |
| 228 | errors.append("long:short edge ratio must not exceed 3:1") |
| 229 | if not (GPT_IMAGE_2_MIN_PIXELS <= total_pixels <= GPT_IMAGE_2_MAX_PIXELS): |
| 230 | errors.append( |
| 231 | f"total pixels {total_pixels:,} must be between " |
| 232 | f"{GPT_IMAGE_2_MIN_PIXELS:,} and {GPT_IMAGE_2_MAX_PIXELS:,}" |
| 233 | ) |
| 234 | if errors: |
| 235 | raise ValueError(f"Invalid gpt-image-2 size '{size}': {', '.join(errors)}") |
| 236 | |
| 237 | |
| 238 | def _select_size( |
| 239 | model: str, |
| 240 | aspect_ratio: str, |
| 241 | image_size: str, |
| 242 | size_preset: str | None = None, |
| 243 | ) -> str: |
| 244 | """Select a model-compatible size while preserving legacy fallbacks.""" |
| 245 | preset = size_preset or "auto" |
| 246 | if preset in {"gpt-image-2"} or (preset == "auto" and _is_gpt_image_2(model)): |
| 247 | size = GPT_IMAGE_2_SIZES[image_size][aspect_ratio] |
| 248 | _validate_gpt_image_2_size(size) |
| 249 | return size |
| 250 | if preset in {"gpt-image", "gpt-image-legacy"} or ( |
| 251 | preset == "auto" and _is_gpt_image_model(model) |
| 252 | ): |
| 253 | return GPT_IMAGE_LEGACY_ASPECT_RATIO_TO_SIZE[aspect_ratio] |
| 254 | if preset in {"dall-e-2", "dalle-2"} or (preset == "auto" and _is_dall_e_2(model)): |
| 255 | return DALL_E_2_SIZE_BY_IMAGE_SIZE[image_size] |
| 256 | return LEGACY_COMPAT_ASPECT_RATIO_TO_SIZE[aspect_ratio] |
| 257 | |
| 258 | |
| 259 | def _supports_response_format(model: str) -> bool: |
| 260 | """GPT Image models always return base64; DALL-E/compatible models may need this.""" |
| 261 | return not _is_gpt_image_model(model) |
| 262 | |
| 263 | |
| 264 | def _read_env_choice(name: str, allowed: set[str]) -> str | None: |
| 265 | value = os.environ.get(name) |
| 266 | if value is None or not value.strip(): |
| 267 | return None |
| 268 | normalized = value.strip().lower() |
| 269 | if normalized not in allowed: |
| 270 | allowed_list = ", ".join(sorted(allowed)) |
| 271 | raise ValueError(f"Invalid {name}='{value}'. Supported: {allowed_list}") |
| 272 | return normalized |
| 273 | |
| 274 | |
| 275 | def _read_env_int(name: str, minimum: int, maximum: int) -> int | None: |
| 276 | value = os.environ.get(name) |
| 277 | if value is None or not value.strip(): |
| 278 | return None |
| 279 | try: |
| 280 | parsed = int(value) |
| 281 | except ValueError as exc: |
| 282 | raise ValueError(f"Invalid {name}='{value}'. Expected integer {minimum}-{maximum}.") from exc |
| 283 | if not (minimum <= parsed <= maximum): |
| 284 | raise ValueError(f"Invalid {name}={parsed}. Expected integer {minimum}-{maximum}.") |
| 285 | return parsed |
| 286 | |
| 287 | |
| 288 | def _gpt_image_options(model: str) -> tuple[dict, str]: |
| 289 | """Read optional GPT Image request parameters from environment.""" |
| 290 | output_format = _read_env_choice("OPENAI_OUTPUT_FORMAT", GPT_IMAGE_OUTPUT_FORMATS) |
| 291 | output_ext = GPT_IMAGE_OUTPUT_EXTENSIONS[output_format] if output_format else ".png" |
| 292 | options = {} |
| 293 | if output_format: |
| 294 | options["output_format"] = output_format |
| 295 | |
| 296 | output_compression = _read_env_int("OPENAI_OUTPUT_COMPRESSION", 0, 100) |
| 297 | if output_compression is not None: |
| 298 | if output_format not in {"jpeg", "webp"}: |
| 299 | raise ValueError( |
| 300 | "OPENAI_OUTPUT_COMPRESSION is only supported when " |
| 301 | "OPENAI_OUTPUT_FORMAT is jpeg or webp." |
| 302 | ) |
| 303 | options["output_compression"] = output_compression |
| 304 | |
| 305 | background = _read_env_choice("OPENAI_BACKGROUND", GPT_IMAGE_BACKGROUNDS) |
| 306 | if background: |
| 307 | if _is_gpt_image_2(model) and background == "transparent": |
| 308 | raise ValueError("gpt-image-2 does not support OPENAI_BACKGROUND=transparent.") |
| 309 | options["background"] = background |
| 310 | |
| 311 | moderation = _read_env_choice("OPENAI_MODERATION", GPT_IMAGE_MODERATION_VALUES) |
| 312 | if moderation: |
| 313 | options["moderation"] = moderation |
| 314 | |
| 315 | return options, output_ext |
| 316 | |
| 317 | |
| 318 | def _image_generations_url(base_url: str | None) -> str: |
| 319 | base = (base_url or DEFAULT_BASE_URL).rstrip("/") |
| 320 | if base.endswith("/images/generations"): |
| 321 | return base |
| 322 | return f"{base}/images/generations" |
| 323 | |
| 324 | |
| 325 | def _image_edits_url(base_url: str | None) -> str: |
| 326 | base = (base_url or DEFAULT_BASE_URL).rstrip("/") |
| 327 | if base.endswith("/images/edits"): |
| 328 | return base |
| 329 | if base.endswith("/images/generations"): |
| 330 | # Swap the sibling endpoint rather than appending to a full URL. |
| 331 | return base[: -len("/generations")] + "/edits" |
| 332 | return f"{base}/images/edits" |
| 333 | |
| 334 | |
| 335 | def _read_size_preset() -> str | None: |
| 336 | """Read the optional size mapping preset for OpenAI-compatible providers.""" |
| 337 | return _read_env_choice("OPENAI_SIZE_PRESET", OPENAI_SIZE_PRESETS) |
| 338 | |
| 339 | |
| 340 | def _read_response_format() -> str | None: |
| 341 | """Read the optional response_format override.""" |
| 342 | return _read_env_choice("OPENAI_RESPONSE_FORMAT", OPENAI_RESPONSE_FORMATS) |
| 343 | |
| 344 | |
| 345 | def _read_quality(image_size: str) -> str | None: |
| 346 | """Resolve the quality field for OpenAI-compatible requests.""" |
| 347 | quality = _read_env_choice("OPENAI_QUALITY", OPENAI_QUALITY_VALUES) |
| 348 | if quality == "omit": |
| 349 | return None |
| 350 | if quality and quality != "auto": |
| 351 | return quality |
| 352 | return IMAGE_SIZE_TO_QUALITY.get(image_size, "auto") |
| 353 | |
| 354 | |
| 355 | def _apply_response_format(request: dict, model: str) -> None: |
| 356 | """Apply response_format while preserving the existing default behavior.""" |
| 357 | response_format = _read_response_format() |
| 358 | if response_format == "omit": |
| 359 | return |
| 360 | if response_format in {"b64_json", "url"}: |
| 361 | request["response_format"] = response_format |
| 362 | return |
| 363 | if _supports_response_format(model): |
| 364 | request["response_format"] = "b64_json" |
| 365 | |
| 366 | |
| 367 | def _post_image_generation(api_key: str, base_url: str | None, request: dict) -> dict: |
| 368 | headers = { |
| 369 | "Authorization": f"Bearer {api_key}", |
| 370 | "Content-Type": "application/json", |
| 371 | } |
| 372 | response = requests.post( |
| 373 | _image_generations_url(base_url), |
| 374 | headers=headers, |
| 375 | json=request, |
| 376 | timeout=300, |
| 377 | ) |
| 378 | if not response.ok: |
| 379 | raise http_error(response, "OpenAI image generation") |
| 380 | try: |
| 381 | return response.json() |
| 382 | except ValueError as exc: |
| 383 | raise RuntimeError("OpenAI image generation returned invalid JSON.") from exc |
| 384 | |
| 385 | |
| 386 | def _post_image_edit(api_key: str, base_url: str | None, |
| 387 | data: dict, image_path: str) -> dict: |
| 388 | headers = {"Authorization": f"Bearer {api_key}"} |
| 389 | # GPT Image models take the image list field 'image[]'; dall-e-2 and other |
| 390 | # OpenAI-compatible edit models use the singular 'image'. |
| 391 | model = str(data.get("model", "")) |
| 392 | field = "image[]" if _is_gpt_image_model(model) else "image" |
| 393 | mime_type = mimetypes.guess_type(image_path)[0] or "application/octet-stream" |
| 394 | # Let requests build the multipart/form-data body (and its Content-Type |
| 395 | # boundary) from files=; do not set Content-Type by hand. |
| 396 | with open(image_path, "rb") as image_file: |
| 397 | response = requests.post( |
| 398 | _image_edits_url(base_url), |
| 399 | headers=headers, |
| 400 | data=data, |
| 401 | files=[(field, (Path(image_path).name, image_file, mime_type))], |
| 402 | timeout=300, |
| 403 | ) |
| 404 | if not response.ok: |
| 405 | raise http_error(response, "OpenAI image edit") |
| 406 | try: |
| 407 | return response.json() |
| 408 | except ValueError as exc: |
| 409 | raise RuntimeError("OpenAI image edit returned invalid JSON.") from exc |
| 410 | |
| 411 | |
| 412 | # ╔══════════════════════════════════════════════════════════════════╗ |
| 413 | # ║ Image Generation ║ |
| 414 | # ╚══════════════════════════════════════════════════════════════════╝ |
| 415 | |
| 416 | def _generate_image(api_key: str, prompt: str, |
| 417 | aspect_ratio: str = "1:1", image_size: str = "1K", |
| 418 | output_dir: str = None, filename: str = None, |
| 419 | model: str = DEFAULT_MODEL, base_url: str = None) -> str: |
| 420 | """ |
| 421 | Image generation via OpenAI-compatible API. |
| 422 | |
| 423 | Maps aspect_ratio to OpenAI's size parameter, and image_size to quality. |
| 424 | |
| 425 | Returns: |
| 426 | Path of the saved image file |
| 427 | |
| 428 | Raises: |
| 429 | RuntimeError: When generation fails |
| 430 | """ |
| 431 | # Map parameters |
| 432 | size_preset = _read_size_preset() |
| 433 | size = _select_size(model, aspect_ratio, image_size, size_preset) |
| 434 | quality = _read_quality(image_size) |
| 435 | output_ext = ".png" |
| 436 | request = { |
| 437 | "prompt": prompt, |
| 438 | "model": model, |
| 439 | "size": size, |
| 440 | "n": 1, |
| 441 | } |
| 442 | if quality is not None: |
| 443 | request["quality"] = quality |
| 444 | if _is_gpt_image_model(model): |
| 445 | gpt_options, output_ext = _gpt_image_options(model) |
| 446 | request.update(gpt_options) |
| 447 | _apply_response_format(request, model) |
| 448 | |
| 449 | mode_label = f"Proxy: {base_url}" if base_url else "OpenAI API" |
| 450 | print(f"[OpenAI - {mode_label}]") |
| 451 | print(f" Model: {model}") |
| 452 | print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}") |
| 453 | print(f" Size: {size} (from aspect_ratio={aspect_ratio})") |
| 454 | if size_preset and size_preset != "auto": |
| 455 | print(f" Size Preset: {size_preset}") |
| 456 | if quality is not None: |
| 457 | print(f" Quality: {quality} (from image_size={image_size})") |
| 458 | else: |
| 459 | print(" Quality: omitted") |
| 460 | if request.get("response_format"): |
| 461 | print(f" Response: {request['response_format']}") |
| 462 | elif _read_response_format() == "omit": |
| 463 | print(" Response: omitted") |
| 464 | if request.get("output_format"): |
| 465 | print(f" Format: {request['output_format']}") |
| 466 | if request.get("output_compression") is not None: |
| 467 | print(f" Compression: {request['output_compression']}") |
| 468 | if request.get("background"): |
| 469 | print(f" Background: {request['background']}") |
| 470 | if request.get("moderation"): |
| 471 | print(f" Moderation: {request['moderation']}") |
| 472 | print() |
| 473 | |
| 474 | start_time = time.time() |
| 475 | print(f" [..] Generating...", end="", flush=True) |
| 476 | |
| 477 | # Heartbeat thread |
| 478 | heartbeat_stop = threading.Event() |
| 479 | |
| 480 | def _heartbeat(): |
| 481 | while not heartbeat_stop.is_set(): |
| 482 | heartbeat_stop.wait(5) |
| 483 | if not heartbeat_stop.is_set(): |
| 484 | elapsed = time.time() - start_time |
| 485 | print(f" {elapsed:.0f}s...", end="", flush=True) |
| 486 | |
| 487 | hb_thread = threading.Thread(target=_heartbeat, daemon=True) |
| 488 | hb_thread.start() |
| 489 | |
| 490 | try: |
| 491 | resp = _post_image_generation(api_key, base_url, request) |
| 492 | finally: |
| 493 | heartbeat_stop.set() |
| 494 | hb_thread.join(timeout=1) |
| 495 | |
| 496 | elapsed = time.time() - start_time |
| 497 | print(f"\n [DONE] Image generated ({elapsed:.1f}s)") |
| 498 | |
| 499 | data = _field(resp, "data") if resp is not None else None |
| 500 | if data: |
| 501 | path = resolve_output_path(prompt, output_dir, filename, output_ext) |
| 502 | first_image = data[0] |
| 503 | b64_json = _field(first_image, "b64_json") |
| 504 | image_url = _field(first_image, "url") |
| 505 | if b64_json: |
| 506 | image_data = base64.b64decode(b64_json) |
| 507 | return save_image_bytes(image_data, path) |
| 508 | if image_url: |
| 509 | return download_image(image_url, path) |
| 510 | |
| 511 | raise RuntimeError("No image was generated. The server may have refused the request.") |
| 512 | |
| 513 | |
| 514 | def _edit_image(api_key: str, prompt: str, reference_image: str, |
| 515 | aspect_ratio: str = "1:1", image_size: str = "1K", |
| 516 | output_dir: str = None, filename: str = None, |
| 517 | model: str = DEFAULT_MODEL, base_url: str = None) -> str: |
| 518 | """ |
| 519 | Image-to-image edit via the OpenAI-compatible /v1/images/edits endpoint. |
| 520 | |
| 521 | Sends the reference image plus the prompt (used as the edit instruction). |
| 522 | Mirrors _generate_image's size/quality/format handling but posts |
| 523 | multipart/form-data instead of JSON. |
| 524 | |
| 525 | Returns: |
| 526 | Path of the saved image file |
| 527 | """ |
| 528 | size_preset = _read_size_preset() |
| 529 | size = _select_size(model, aspect_ratio, image_size, size_preset) |
| 530 | quality = None if _is_dall_e_2(model) else _read_quality(image_size) |
| 531 | output_ext = ".png" |
| 532 | request = { |
| 533 | "prompt": prompt, |
| 534 | "model": model, |
| 535 | "size": size, |
| 536 | "n": 1, |
| 537 | } |
| 538 | if quality is not None: |
| 539 | request["quality"] = quality |
| 540 | if _is_gpt_image_model(model): |
| 541 | gpt_options, output_ext = _gpt_image_options(model) |
| 542 | request.update(gpt_options) |
| 543 | _apply_response_format(request, model) |
| 544 | |
| 545 | mode_label = f"Proxy: {base_url}" if base_url else "OpenAI API" |
| 546 | print(f"[OpenAI - {mode_label}]") |
| 547 | print(f" Mode: edit (image-to-image)") |
| 548 | print(f" Model: {model}") |
| 549 | print(f" Reference: {reference_image}") |
| 550 | print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}") |
| 551 | print(f" Size: {size} (from aspect_ratio={aspect_ratio})") |
| 552 | if size_preset and size_preset != "auto": |
| 553 | print(f" Size Preset: {size_preset}") |
| 554 | if quality is not None: |
| 555 | print(f" Quality: {quality} (from image_size={image_size})") |
| 556 | else: |
| 557 | print(" Quality: omitted") |
| 558 | if request.get("response_format"): |
| 559 | print(f" Response: {request['response_format']}") |
| 560 | elif _read_response_format() == "omit": |
| 561 | print(" Response: omitted") |
| 562 | if request.get("output_format"): |
| 563 | print(f" Format: {request['output_format']}") |
| 564 | if request.get("output_compression") is not None: |
| 565 | print(f" Compression: {request['output_compression']}") |
| 566 | if request.get("background"): |
| 567 | print(f" Background: {request['background']}") |
| 568 | if request.get("moderation"): |
| 569 | print(f" Moderation: {request['moderation']}") |
| 570 | print() |
| 571 | |
| 572 | start_time = time.time() |
| 573 | print(f" [..] Editing...", end="", flush=True) |
| 574 | |
| 575 | heartbeat_stop = threading.Event() |
| 576 | |
| 577 | def _heartbeat(): |
| 578 | while not heartbeat_stop.is_set(): |
| 579 | heartbeat_stop.wait(5) |
| 580 | if not heartbeat_stop.is_set(): |
| 581 | elapsed = time.time() - start_time |
| 582 | print(f" {elapsed:.0f}s...", end="", flush=True) |
| 583 | |
| 584 | hb_thread = threading.Thread(target=_heartbeat, daemon=True) |
| 585 | hb_thread.start() |
| 586 | |
| 587 | try: |
| 588 | resp = _post_image_edit(api_key, base_url, request, reference_image) |
| 589 | finally: |
| 590 | heartbeat_stop.set() |
| 591 | hb_thread.join(timeout=1) |
| 592 | |
| 593 | elapsed = time.time() - start_time |
| 594 | print(f"\n [DONE] Image edited ({elapsed:.1f}s)") |
| 595 | |
| 596 | data = _field(resp, "data") if resp is not None else None |
| 597 | if data: |
| 598 | path = resolve_output_path(prompt, output_dir, filename, output_ext) |
| 599 | first_image = data[0] |
| 600 | b64_json = _field(first_image, "b64_json") |
| 601 | image_url = _field(first_image, "url") |
| 602 | if b64_json: |
| 603 | image_data = base64.b64decode(b64_json) |
| 604 | return save_image_bytes(image_data, path) |
| 605 | if image_url: |
| 606 | return download_image(image_url, path) |
| 607 | |
| 608 | raise RuntimeError("No image was returned. The server may have refused the edit request.") |
| 609 | |
| 610 | |
| 611 | # ╔══════════════════════════════════════════════════════════════════╗ |
| 612 | # ║ Public Entry Point ║ |
| 613 | # ╚══════════════════════════════════════════════════════════════════╝ |
| 614 | |
| 615 | def generate(prompt: str, |
| 616 | aspect_ratio: str = "1:1", image_size: str = "1K", |
| 617 | output_dir: str = None, filename: str = None, |
| 618 | model: str = None, max_retries: int = MAX_RETRIES, |
| 619 | reference_image: str = None) -> str: |
| 620 | """ |
| 621 | OpenAI-compatible image generation (or image-to-image edit) with retry. |
| 622 | |
| 623 | Reads credentials from the current process environment or a `.env` file: |
| 624 | OPENAI_API_KEY |
| 625 | OPENAI_BASE_URL |
| 626 | OPENAI_MODEL (optional override) |
| 627 | |
| 628 | Args: |
| 629 | prompt: Prompt text (edit instruction when reference_image is set) |
| 630 | aspect_ratio: Aspect ratio, mapped to OpenAI size |
| 631 | image_size: Image size, mapped to OpenAI quality |
| 632 | output_dir: Output directory |
| 633 | filename: Output filename (without extension) |
| 634 | model: Model name (default: gpt-image-2) |
| 635 | max_retries: Maximum number of retries |
| 636 | reference_image: Optional source image path. When set, the request |
| 637 | goes to /v1/images/edits (image-to-image) instead of |
| 638 | /v1/images/generations. |
| 639 | |
| 640 | Returns: |
| 641 | Path of the saved image file |
| 642 | """ |
| 643 | api_key = os.environ.get("OPENAI_API_KEY") |
| 644 | base_url = os.environ.get("OPENAI_BASE_URL") |
| 645 | |
| 646 | if not api_key: |
| 647 | raise ValueError( |
| 648 | "No API key found. Set OPENAI_API_KEY in the current environment or a .env file." |
| 649 | ) |
| 650 | |
| 651 | if model is None: |
| 652 | model = os.environ.get("OPENAI_MODEL") or DEFAULT_MODEL |
| 653 | |
| 654 | image_size = normalize_image_size(image_size) |
| 655 | |
| 656 | if aspect_ratio not in LEGACY_COMPAT_ASPECT_RATIO_TO_SIZE: |
| 657 | supported = list(LEGACY_COMPAT_ASPECT_RATIO_TO_SIZE.keys()) |
| 658 | raise ValueError( |
| 659 | f"Unsupported aspect ratio '{aspect_ratio}' for OpenAI backend. " |
| 660 | f"Supported: {supported}" |
| 661 | ) |
| 662 | |
| 663 | last_error = None |
| 664 | for attempt in range(max_retries + 1): |
| 665 | try: |
| 666 | if reference_image is not None: |
| 667 | return _edit_image(api_key, prompt, reference_image, |
| 668 | aspect_ratio, image_size, output_dir, |
| 669 | filename, model, base_url) |
| 670 | return _generate_image(api_key, prompt, |
| 671 | aspect_ratio, image_size, output_dir, |
| 672 | filename, model, base_url) |
| 673 | except Exception as e: |
| 674 | last_error = e |
| 675 | if attempt < max_retries and is_rate_limit_error(e): |
| 676 | delay = retry_delay(attempt, rate_limited=True) |
| 677 | print(f"\n [WARN] Rate limit hit (attempt {attempt + 1}/{max_retries + 1}). " |
| 678 | f"Waiting {delay}s before retry...") |
| 679 | time.sleep(delay) |
| 680 | elif attempt < max_retries: |
| 681 | delay = retry_delay(attempt, rate_limited=False) |
| 682 | print(f"\n [WARN] Error (attempt {attempt + 1}/{max_retries + 1}): {e}. " |
| 683 | f"Retrying in {delay}s...") |
| 684 | time.sleep(delay) |
| 685 | else: |
| 686 | break |
| 687 | |
| 688 | raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}") |
| 689 |