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