"""Python 3.8+: pip install requests pillow. API Key and model come from env.""" import argparse import base64 import io import os import re from contextlib import ExitStack from pathlib import Path import requests from PIL import Image def validate_options(model, size, quality, count, timeout): if isinstance(count, bool) or not isinstance(count, int) or not 1 <= count <= 10: raise ValueError("n must be an integer between 1 and 10") if isinstance(timeout, bool) or not isinstance(timeout, int) or timeout <= 0: raise ValueError("--timeout must be a positive number of seconds") qualities = ["auto", "low", "medium", "high"] if model.startswith("gpt-image-2.5-"): qualities += ["xhigh", "max"] if quality not in qualities: raise ValueError("Unsupported quality for the selected model") if size == "auto": return None if not isinstance(size, str) or not re.fullmatch(r"[1-9][0-9]*x[1-9][0-9]*", size): raise ValueError("size must be auto or positive WIDTHxHEIGHT using lowercase x") dimensions = tuple(int(value) for value in size.split("x")) if model == "gpt-image-2": width, height = dimensions if width % 16 or height % 16 or max(dimensions) > 3840 or max(dimensions) > 3 * min(dimensions) or not 655360 <= width * height <= 8294400: raise ValueError("size violates the selected model's pixel, alignment or aspect-ratio constraints") return dimensions def main(): parser = argparse.ArgumentParser() parser.add_argument("mode", choices=["generate", "edit"]) parser.add_argument("--prompt", required=True) parser.add_argument("--image", action="append", default=[]) parser.add_argument("--mask") parser.add_argument("--size", default="1024x1024") parser.add_argument("--quality", default="medium") parser.add_argument("--format", choices=["png", "jpeg"], default="png") parser.add_argument("--n", type=int, default=1) parser.add_argument("--out", default="image-results") parser.add_argument("--timeout", type=int, default=300, help="Read timeout in seconds; a timeout does not mean generation was cancelled") args = parser.parse_args() key, model = os.environ["YUNQIAI_API_KEY"], os.environ["YUNQIAI_MODEL"] if model not in ["gpt-image-2", "gpt-image-2.5-flare", "gpt-image-2.5-sunburst"]: raise ValueError("Choose an exact Images API model ID") expected_size = validate_options(model, args.size, args.quality, args.n, args.timeout) body = {"model": model, "prompt": args.prompt, "size": args.size, "quality": args.quality, "output_format": args.format, "n": args.n} route = "/v1/images/generations" with ExitStack() as stack: request = {"json": body} if args.mode == "edit": if not args.image: raise ValueError("At least one --image is required") route = "/v1/images/edits" if model.startswith("gpt-image-2.5-"): route += "?api-version=2025-04-01-preview" files = [] for name in args.image: path = Path(name) with Image.open(path) as image: image.verify() mime = {"PNG": "image/png", "JPEG": "image/jpeg", "WEBP": "image/webp"}.get(image.format) if not mime: raise ValueError("Reference must be PNG, JPEG or WebP") files.append(("image[]" if len(args.image) > 1 else "image", (path.name, stack.enter_context(path.open("rb")), mime))) if args.mask: mask_path = Path(args.mask) with Image.open(mask_path) as mask, Image.open(args.image[0]) as first: if mask.format != "PNG" or "A" not in mask.getbands() or mask.size != first.size: raise ValueError("Mask must be PNG with alpha and match the first image dimensions") files.append(("mask", (mask_path.name, stack.enter_context(mask_path.open("rb")), "image/png"))) request = {"data": {k: str(v) for k, v in body.items()}, "files": files} elif args.image or args.mask: raise ValueError("Use edit mode for reference images or masks") try: response = requests.post("https://www.yunqiai.chat" + route, headers={"Authorization": "Bearer " + key}, allow_redirects=False, timeout=(15, args.timeout), **request) except requests.exceptions.RequestException as error: raise RuntimeError("Generation state is unknown after a network failure. Do not automatically repeat the POST; check platform request records first") from error if not 200 <= response.status_code < 300: raise RuntimeError("HTTP " + str(response.status_code) + ": " + response.text[:4000]) data = response.json() if data.get("error") or not data.get("data"): raise RuntimeError("No image data returned") output = Path(args.out) output.mkdir(parents=True, exist_ok=True) failures = [] for index, item in enumerate(data["data"]): try: if item.get("b64_json"): raw = base64.b64decode(item["b64_json"], validate=True) elif item.get("url"): download = requests.get(item["url"], timeout=(15, 300)) download.raise_for_status() raw = download.content else: raise ValueError("No b64_json or URL") with Image.open(io.BytesIO(raw)) as image: image.load() extension = {"JPEG": "jpg", "PNG": "png", "WEBP": "webp"}[image.format] path = output / ("image-" + str(index) + "." + extension) path.write_bytes(raw) print(str(path), image.format, image.size) if expected_size and image.size != expected_size: failures.append({"index": index, "error": "Returned dimensions differ from size", "requested": expected_size, "actual": image.size}) if image.format != {"png": "PNG", "jpeg": "JPEG"}[args.format]: failures.append({"index": index, "error": "Returned image format differs from output_format", "requested": args.format, "actual": image.format}) except Exception as error: failures.append({"index": index, "error": str(error)}) if failures: raise RuntimeError("Image decode/download failures: " + str(failures)) if len(data["data"]) != args.n: raise RuntimeError("Returned image count differs from requested n") if __name__ == "__main__": main()