图片生成与编辑

import base64
import binascii
import os
import sys
from pathlib import Path
from openai import OpenAI, APIConnectionError, APIStatusError

mode = sys.argv[1] if len(sys.argv) == 2 else ""
if mode not in ("generate", "edit"):
    raise SystemExit("Usage: python images.py generate|edit")

with OpenAI(base_url=os.environ["API_BASE_URL"].rstrip("/"),
            api_key=os.environ["API_KEY"], timeout=180.0, max_retries=0) as client:
    try:
        if mode == "generate":
            result = client.images.generate(
                model=os.environ["API_MODEL"], n=1,
                prompt="A small blue ceramic cup on a plain white background.",
            )
        else:
            with Path(os.environ["API_IMAGE_PATH"]).open("rb") as source:
                result = client.images.edit(
                    model=os.environ["API_MODEL"], image=source, n=1,
                    prompt="Change the cup to green, preserving the composition.",
                )
        if not result.data:
            raise ValueError("No image payload.")
        for index, item in enumerate(result.data, start=1):
            if item.b64_json:
                data = base64.b64decode(item.b64_json, validate=True)
                if data.startswith(b"\x89PNG\r\n\x1a\n"):
                    extension = "png"
                elif data.startswith(b"\xff\xd8\xff"):
                    extension = "jpg"
                elif data.startswith(b"RIFF") and data[8:12] == b"WEBP":
                    extension = "webp"
                else:
                    raise ValueError("Unsupported image bytes.")
                output = Path(f"image-{index}.{extension}")
                with output.open("xb") as target:
                    target.write(data)
                print(f"Saved {output}")
            elif item.url:
                # Do not send the gateway API key to this URL.
                print(item.url)
            else:
                raise ValueError("Image entry has no payload.")
    except APIStatusError as error:
        raise SystemExit(f"Request failed (HTTP {error.status_code}). Check the request logs.")
    except APIConnectionError:
        raise SystemExit("Connection failed; do not automatically resubmit generation.")
    except (ValueError, binascii.Error, OSError):
        raise SystemExit("Cannot save image. Check the response format and output files.")