feat(project): init
This commit is contained in:
992
scripts/generate_image.py
Normal file
992
scripts/generate_image.py
Normal file
@@ -0,0 +1,992 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import binascii
|
||||
from dataclasses import dataclass
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
import re
|
||||
import socket
|
||||
import sys
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from runtime_config import (
|
||||
ConfigError,
|
||||
DEFAULT_BASE_URL,
|
||||
DEFAULT_MODEL,
|
||||
ROOT_DIR,
|
||||
build_api_url,
|
||||
infer_request_style,
|
||||
load_runtime_config,
|
||||
normalize_base_url,
|
||||
normalize_token,
|
||||
parse_env_file,
|
||||
runtime_config_summary,
|
||||
save_base_url as save_runtime_base_url,
|
||||
save_token as save_runtime_token,
|
||||
)
|
||||
|
||||
|
||||
SUPPORTED_ASPECT_RATIOS = (
|
||||
"1:1",
|
||||
"16:9",
|
||||
"9:16",
|
||||
"4:3",
|
||||
"3:4",
|
||||
"3:2",
|
||||
"2:3",
|
||||
"5:4",
|
||||
"4:5",
|
||||
"2:1",
|
||||
"1:2",
|
||||
"21:9",
|
||||
"9:21",
|
||||
)
|
||||
PIXEL_SIZE_PATTERN = re.compile(r"^(?P<width>\d{2,5})x(?P<height>\d{2,5})$")
|
||||
OPENAI_RATIO_SIZE_MAP = {
|
||||
"1:1": "1024x1024",
|
||||
"16:9": "1536x864",
|
||||
"9:16": "864x1536",
|
||||
"4:3": "1536x1152",
|
||||
"3:4": "1152x1536",
|
||||
"3:2": "1536x1024",
|
||||
"2:3": "1024x1536",
|
||||
"5:4": "1280x1024",
|
||||
"4:5": "1024x1280",
|
||||
"2:1": "1536x768",
|
||||
"1:2": "768x1536",
|
||||
"21:9": "1680x720",
|
||||
"9:21": "720x1680",
|
||||
}
|
||||
SIZE_ENV_VARS = (
|
||||
"GPT_IMAGE_2_SIZE",
|
||||
"IMAGE_SIZE",
|
||||
"OPENAI_IMAGE_SIZE",
|
||||
)
|
||||
EDIT_SIZE_ENV_VARS = (
|
||||
"GPT_IMAGE_2_EDIT_SIZE",
|
||||
"IMAGE_EDIT_SIZE",
|
||||
"OPENAI_IMAGE_EDIT_SIZE",
|
||||
)
|
||||
EDIT_QUALITY_ENV_VARS = (
|
||||
"GPT_IMAGE_2_EDIT_QUALITY",
|
||||
"IMAGE_EDIT_QUALITY",
|
||||
"OPENAI_IMAGE_EDIT_QUALITY",
|
||||
)
|
||||
EDIT_BACKGROUND_ENV_VARS = (
|
||||
"GPT_IMAGE_2_EDIT_BACKGROUND",
|
||||
"IMAGE_EDIT_BACKGROUND",
|
||||
"OPENAI_IMAGE_EDIT_BACKGROUND",
|
||||
)
|
||||
EDIT_OUTPUT_FORMAT_ENV_VARS = (
|
||||
"GPT_IMAGE_2_EDIT_OUTPUT_FORMAT",
|
||||
"IMAGE_EDIT_OUTPUT_FORMAT",
|
||||
"OPENAI_IMAGE_EDIT_OUTPUT_FORMAT",
|
||||
)
|
||||
EDIT_INPUT_FIDELITY_ENV_VARS = (
|
||||
"GPT_IMAGE_2_EDIT_INPUT_FIDELITY",
|
||||
"IMAGE_EDIT_INPUT_FIDELITY",
|
||||
"OPENAI_IMAGE_EDIT_INPUT_FIDELITY",
|
||||
)
|
||||
REQUEST_STYLE_CHOICES = ("auto", "apicodex", "openai")
|
||||
REQUEST_RETRY_ATTEMPTS = max(1, int(os.environ.get("IMAGE_RETRY_ATTEMPTS", "2")))
|
||||
REQUEST_RETRY_DELAY_SECONDS = max(0.0, float(os.environ.get("IMAGE_RETRY_DELAY_SECONDS", "2")))
|
||||
|
||||
|
||||
class ImageGenerationError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class ImageRequestError(ImageGenerationError):
|
||||
def __init__(self, message: str, *, status: int | None = None, timeout: bool = False):
|
||||
super().__init__(message)
|
||||
self.status = status
|
||||
self.timeout = timeout
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ImageInputFile:
|
||||
filename: str
|
||||
mime_type: str
|
||||
data: bytes
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RequestTarget:
|
||||
name: str
|
||||
base_url: str
|
||||
token: str
|
||||
request_style: str
|
||||
source: str
|
||||
|
||||
|
||||
def normalize_size_value(raw_size: str) -> str:
|
||||
value = raw_size.strip().lower().replace(" ", "")
|
||||
match = PIXEL_SIZE_PATTERN.fullmatch(value)
|
||||
if match:
|
||||
width = int(match.group("width"))
|
||||
height = int(match.group("height"))
|
||||
if width < 64 or height < 64:
|
||||
raise ImageGenerationError(f"Unsupported size: {raw_size}")
|
||||
return f"{width}x{height}"
|
||||
return value
|
||||
|
||||
|
||||
def is_supported_size(size: str) -> bool:
|
||||
return size in SUPPORTED_ASPECT_RATIOS or PIXEL_SIZE_PATTERN.fullmatch(size) is not None
|
||||
|
||||
|
||||
def resolve_api_size(requested_size: str, request_style: str) -> str:
|
||||
if request_style == "openai" and requested_size in OPENAI_RATIO_SIZE_MAP:
|
||||
return OPENAI_RATIO_SIZE_MAP[requested_size]
|
||||
return requested_size
|
||||
|
||||
|
||||
def lookup_env_setting(env_map: dict[str, str], env_vars: tuple[str, ...]) -> str | None:
|
||||
for env_name in env_vars:
|
||||
value = os.environ.get(env_name, "").strip()
|
||||
if value:
|
||||
return value
|
||||
for env_name in env_vars:
|
||||
value = env_map.get(env_name, "").strip()
|
||||
if value:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def resolve_requested_size(
|
||||
size: str | None,
|
||||
request_style: str,
|
||||
has_refs: bool,
|
||||
env_map: dict[str, str],
|
||||
) -> tuple[str, str]:
|
||||
if size and size.strip():
|
||||
requested_size = normalize_size_value(size)
|
||||
else:
|
||||
env_vars = EDIT_SIZE_ENV_VARS if has_refs else SIZE_ENV_VARS
|
||||
env_size = lookup_env_setting(env_map, env_vars)
|
||||
requested_size = normalize_size_value(env_size) if env_size else "1:1"
|
||||
|
||||
if not is_supported_size(requested_size):
|
||||
raise ImageGenerationError(f"Unsupported size: {requested_size}")
|
||||
return requested_size, resolve_api_size(requested_size, request_style)
|
||||
|
||||
|
||||
def resolve_optional_edit_field(
|
||||
explicit_value: str | None,
|
||||
env_map: dict[str, str],
|
||||
env_vars: tuple[str, ...],
|
||||
) -> str | None:
|
||||
if explicit_value is not None and explicit_value.strip():
|
||||
return explicit_value.strip()
|
||||
value = lookup_env_setting(env_map, env_vars)
|
||||
return value.strip() if value else None
|
||||
|
||||
|
||||
def build_env_template() -> str:
|
||||
return "\n".join(
|
||||
[
|
||||
"# Required: apicodex-compatible image API token",
|
||||
"IMAGE_API_KEY=your_apicodex_token_here",
|
||||
"",
|
||||
"# Optional: host root or full /v1 path are both accepted",
|
||||
f"IMAGE_BASE_URL={DEFAULT_BASE_URL}/v1",
|
||||
"",
|
||||
"# Optional: defaults shown below",
|
||||
f"IMAGE_MODEL={DEFAULT_MODEL}",
|
||||
"IMAGE_REQUEST_STYLE=auto",
|
||||
"IMAGE_OUTPUT_DIR=./output",
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def write_env_template(
|
||||
*,
|
||||
env_file: str | Path | None = None,
|
||||
overwrite: bool = False,
|
||||
) -> Path:
|
||||
execution_root = Path.cwd()
|
||||
if env_file:
|
||||
target_path = Path(env_file).expanduser()
|
||||
if not target_path.is_absolute():
|
||||
target_path = (execution_root / target_path).resolve()
|
||||
else:
|
||||
target_path = execution_root / ".env.example"
|
||||
|
||||
if target_path.exists() and not overwrite:
|
||||
raise ConfigError(f"Config template already exists: {target_path}. Use --force to overwrite.")
|
||||
|
||||
target_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
target_path.write_text(build_env_template(), encoding="utf-8")
|
||||
return target_path
|
||||
|
||||
def read_json_response(request: urllib.request.Request, timeout: int) -> dict[str, Any]:
|
||||
try:
|
||||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||||
body = response.read().decode("utf-8")
|
||||
except urllib.error.HTTPError as error:
|
||||
body = error.read().decode("utf-8", errors="replace")
|
||||
raise ImageRequestError(f"HTTP {error.code}: {body}", status=error.code) from error
|
||||
except urllib.error.URLError as error:
|
||||
if is_timeout_error(error):
|
||||
raise ImageRequestError(f"Timeout: {error}", timeout=True) from error
|
||||
raise ImageRequestError(f"Network error: {error}") from error
|
||||
except (TimeoutError, socket.timeout) as error:
|
||||
raise ImageRequestError(f"Timeout: {error}", timeout=True) from error
|
||||
|
||||
if not body:
|
||||
return {}
|
||||
try:
|
||||
return json.loads(body)
|
||||
except json.JSONDecodeError as error:
|
||||
raise ImageGenerationError(f"API returned non-JSON response: {body[:500]}") from error
|
||||
|
||||
|
||||
def request_json(method: str, url: str, token: str, payload: dict[str, Any] | None = None, timeout: int = 300) -> dict[str, Any]:
|
||||
data = None
|
||||
headers = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Content-Type": "application/json",
|
||||
"User-Agent": "gpt-image-2-generator/1.0",
|
||||
}
|
||||
if payload is not None:
|
||||
data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
||||
|
||||
request = urllib.request.Request(url, data=data, headers=headers, method=method)
|
||||
return read_json_response(request, timeout)
|
||||
|
||||
|
||||
def request_multipart_json(
|
||||
url: str,
|
||||
token: str,
|
||||
fields: dict[str, Any],
|
||||
files: list[ImageInputFile],
|
||||
image_field_name: str = "image[]",
|
||||
timeout: int = 300,
|
||||
) -> dict[str, Any]:
|
||||
boundary = f"----gpt-image-2-{uuid.uuid4().hex}"
|
||||
body = bytearray()
|
||||
|
||||
def add(value: str | bytes) -> None:
|
||||
if isinstance(value, str):
|
||||
body.extend(value.encode("utf-8"))
|
||||
else:
|
||||
body.extend(value)
|
||||
|
||||
for name, value in fields.items():
|
||||
add(f"--{boundary}\r\n")
|
||||
add(f'Content-Disposition: form-data; name="{name}"\r\n\r\n')
|
||||
add(str(value))
|
||||
add("\r\n")
|
||||
|
||||
for image_file in files:
|
||||
add(f"--{boundary}\r\n")
|
||||
add(
|
||||
f'Content-Disposition: form-data; name="{image_field_name}"; '
|
||||
f'filename="{image_file.filename}"\r\n'
|
||||
)
|
||||
add(f"Content-Type: {image_file.mime_type}\r\n\r\n")
|
||||
add(image_file.data)
|
||||
add("\r\n")
|
||||
|
||||
add(f"--{boundary}--\r\n")
|
||||
headers = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Content-Type": f"multipart/form-data; boundary={boundary}",
|
||||
"Content-Length": str(len(body)),
|
||||
"User-Agent": "gpt-image-2-generator/1.0",
|
||||
}
|
||||
request = urllib.request.Request(url, data=bytes(body), headers=headers, method="POST")
|
||||
return read_json_response(request, timeout)
|
||||
|
||||
|
||||
def is_timeout_error(error: urllib.error.URLError) -> bool:
|
||||
reason = getattr(error, "reason", None)
|
||||
if isinstance(reason, (TimeoutError, socket.timeout)):
|
||||
return True
|
||||
return "timed out" in str(reason).lower() or "timeout" in str(error).lower()
|
||||
|
||||
|
||||
def is_failover_eligible_error(error: BaseException) -> bool:
|
||||
if isinstance(error, ImageRequestError):
|
||||
if error.status in {502, 503, 504, 524}:
|
||||
return True
|
||||
if error.timeout:
|
||||
return True
|
||||
message = str(error).lower()
|
||||
return "timed out" in message or "timeout" in message
|
||||
|
||||
|
||||
def read_first_config_line(path: Path) -> str:
|
||||
if not path.exists():
|
||||
return ""
|
||||
for raw_line in path.read_text(encoding="utf-8-sig").splitlines():
|
||||
line = raw_line.strip()
|
||||
if line and not line.startswith("#"):
|
||||
return line
|
||||
return ""
|
||||
|
||||
|
||||
def load_skill_fallback_target(runtime_config: Any) -> RequestTarget | None:
|
||||
skill_config_dir = ROOT_DIR / "config"
|
||||
base_url = normalize_base_url(read_first_config_line(skill_config_dir / "base_url.txt"))
|
||||
token = normalize_token(read_first_config_line(skill_config_dir / "token.txt"))
|
||||
request_style_raw = read_first_config_line(skill_config_dir / "request_style.txt").strip().lower()
|
||||
|
||||
if not base_url or not token:
|
||||
return None
|
||||
|
||||
request_style = request_style_raw if request_style_raw in {"apicodex", "openai"} else infer_request_style(base_url)
|
||||
primary_token = normalize_token(runtime_config.token or "")
|
||||
if runtime_config.base_url == base_url and primary_token == token:
|
||||
return None
|
||||
|
||||
return RequestTarget(
|
||||
name="skill-fallback",
|
||||
base_url=base_url,
|
||||
token=token,
|
||||
request_style=request_style,
|
||||
source=str(skill_config_dir),
|
||||
)
|
||||
|
||||
|
||||
def build_request_targets(runtime_config: Any, request_style: str) -> list[RequestTarget]:
|
||||
targets = [
|
||||
RequestTarget(
|
||||
name="primary",
|
||||
base_url=runtime_config.base_url,
|
||||
token=runtime_config.token or "",
|
||||
request_style=request_style,
|
||||
source=runtime_config.base_url_source,
|
||||
)
|
||||
]
|
||||
fallback_target = load_skill_fallback_target(runtime_config)
|
||||
if fallback_target is not None:
|
||||
targets.append(fallback_target)
|
||||
return targets
|
||||
|
||||
|
||||
def format_request_error(error: BaseException) -> str:
|
||||
if isinstance(error, ImageRequestError):
|
||||
if error.status is not None:
|
||||
return f"HTTP {error.status}"
|
||||
if error.timeout:
|
||||
return "timeout"
|
||||
return error.__class__.__name__
|
||||
|
||||
|
||||
def should_retry_request(error: BaseException) -> bool:
|
||||
if isinstance(error, ImageRequestError):
|
||||
if error.status in {429, 500, 502, 503, 504, 524}:
|
||||
return True
|
||||
if error.timeout:
|
||||
return True
|
||||
message = str(error).lower()
|
||||
return "timed out" in message or "timeout" in message
|
||||
|
||||
|
||||
def submit_request_for_target(
|
||||
target: RequestTarget,
|
||||
*,
|
||||
runtime_config: Any,
|
||||
env_map: dict[str, str],
|
||||
prompt: str,
|
||||
size: str | None,
|
||||
n: int,
|
||||
has_refs: bool,
|
||||
image_files: list[ImageInputFile],
|
||||
timeout: int,
|
||||
quality: str | None,
|
||||
background: str | None,
|
||||
output_format: str | None,
|
||||
input_fidelity: str | None,
|
||||
) -> tuple[dict[str, Any], str, str, str]:
|
||||
generations_api_url = build_api_url(target.base_url, "images/generations")
|
||||
edits_api_url = build_api_url(target.base_url, "images/edits")
|
||||
payload: dict[str, Any] = {
|
||||
"model": runtime_config.model,
|
||||
"prompt": prompt,
|
||||
}
|
||||
requested_size, api_size = resolve_requested_size(size, target.request_style, has_refs, env_map)
|
||||
payload["size"] = api_size
|
||||
if target.request_style == "apicodex" or n != 1:
|
||||
payload["n"] = n
|
||||
|
||||
if has_refs:
|
||||
edit_quality = resolve_optional_edit_field(quality, env_map, EDIT_QUALITY_ENV_VARS)
|
||||
edit_background = resolve_optional_edit_field(background, env_map, EDIT_BACKGROUND_ENV_VARS)
|
||||
edit_output_format = resolve_optional_edit_field(output_format, env_map, EDIT_OUTPUT_FORMAT_ENV_VARS)
|
||||
edit_input_fidelity = resolve_optional_edit_field(input_fidelity, env_map, EDIT_INPUT_FIDELITY_ENV_VARS)
|
||||
if edit_quality:
|
||||
payload["quality"] = edit_quality
|
||||
if edit_background:
|
||||
payload["background"] = edit_background
|
||||
if edit_output_format:
|
||||
payload["output_format"] = edit_output_format
|
||||
if edit_input_fidelity:
|
||||
payload["input_fidelity"] = edit_input_fidelity
|
||||
image_field_name = "image[]" if target.request_style == "apicodex" or len(image_files) != 1 else "image"
|
||||
response = request_multipart_json(
|
||||
edits_api_url,
|
||||
target.token,
|
||||
payload,
|
||||
image_files,
|
||||
image_field_name=image_field_name,
|
||||
timeout=timeout,
|
||||
)
|
||||
else:
|
||||
response = request_json("POST", generations_api_url, target.token, payload, timeout=timeout)
|
||||
|
||||
return response, target.request_style, requested_size, api_size
|
||||
|
||||
|
||||
def safe_filename(filename: str, fallback: str, mime_type: str) -> str:
|
||||
candidate = re.sub(r"[^A-Za-z0-9._-]+", "_", Path(filename).name).strip("._")
|
||||
if not candidate:
|
||||
candidate = fallback
|
||||
if not Path(candidate).suffix:
|
||||
candidate += mimetypes.guess_extension(mime_type) or ".png"
|
||||
return candidate
|
||||
|
||||
|
||||
def read_data_uri_image(value: str, index: int) -> ImageInputFile:
|
||||
match = re.match(r"^data:(?P<mime>image/[^;,]+)(?P<params>(?:;[^,]+)*),(?P<data>.+)$", value, flags=re.DOTALL)
|
||||
if not match or ";base64" not in match.group("params").lower():
|
||||
raise ImageGenerationError("Reference image data URI must be a base64 image")
|
||||
|
||||
mime_type = match.group("mime")
|
||||
try:
|
||||
data = base64.b64decode(match.group("data"), validate=True)
|
||||
except (ValueError, binascii.Error) as error:
|
||||
raise ImageGenerationError("Reference image data URI contains invalid base64") from error
|
||||
return ImageInputFile(
|
||||
filename=safe_filename("", f"reference-{index}", mime_type),
|
||||
mime_type=mime_type,
|
||||
data=data,
|
||||
)
|
||||
|
||||
|
||||
def download_reference_image(url: str, index: int, timeout: int) -> ImageInputFile:
|
||||
request = urllib.request.Request(url, headers={"User-Agent": "gpt-image-2-generator/1.0"})
|
||||
try:
|
||||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||||
data = response.read()
|
||||
content_type = response.headers.get("Content-Type")
|
||||
except urllib.error.HTTPError as error:
|
||||
body = error.read().decode("utf-8", errors="replace")
|
||||
raise ImageGenerationError(f"Reference image URL returned HTTP {error.code}: {body[:500]}") from error
|
||||
except urllib.error.URLError as error:
|
||||
raise ImageGenerationError(f"Reference image URL could not be downloaded: {error}") from error
|
||||
|
||||
mime_type = (content_type or "").split(";")[0].strip().lower()
|
||||
if not mime_type.startswith("image/"):
|
||||
mime_type = mimetypes.guess_type(urllib.parse.urlparse(url).path)[0] or "image/png"
|
||||
filename = safe_filename(Path(urllib.parse.urlparse(url).path).name, f"reference-{index}", mime_type)
|
||||
return ImageInputFile(filename=filename, mime_type=mime_type, data=data)
|
||||
|
||||
|
||||
def read_local_reference_image(image_ref: str, index: int) -> ImageInputFile:
|
||||
value = image_ref.strip()
|
||||
path = Path(value).expanduser()
|
||||
if not path.is_absolute():
|
||||
cwd_path = Path.cwd() / path
|
||||
root_path = ROOT_DIR / path
|
||||
path = cwd_path if cwd_path.exists() else root_path
|
||||
|
||||
if not path.exists():
|
||||
raise ImageGenerationError(f"Reference image not found: {image_ref}")
|
||||
|
||||
mime_type = mimetypes.guess_type(str(path))[0] or "image/png"
|
||||
return ImageInputFile(
|
||||
filename=safe_filename(path.name, f"reference-{index}", mime_type),
|
||||
mime_type=mime_type,
|
||||
data=path.read_bytes(),
|
||||
)
|
||||
|
||||
|
||||
def image_ref_to_input_file(image_ref: str, index: int, timeout: int) -> ImageInputFile:
|
||||
value = image_ref.strip()
|
||||
lowered = value.lower()
|
||||
if lowered.startswith("data:"):
|
||||
return read_data_uri_image(value, index)
|
||||
if lowered.startswith(("http://", "https://")):
|
||||
return download_reference_image(value, index, timeout)
|
||||
return read_local_reference_image(value, index)
|
||||
|
||||
|
||||
def save_json(path: Path, data: dict[str, Any]) -> None:
|
||||
path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
|
||||
|
||||
def unique_path(path: Path) -> Path:
|
||||
if not path.exists():
|
||||
return path
|
||||
for index in range(2, 1000):
|
||||
candidate = path.with_name(f"{path.stem}-{index}{path.suffix}")
|
||||
if not candidate.exists():
|
||||
return candidate
|
||||
raise ImageGenerationError(f"Cannot create unique output path for {path}")
|
||||
|
||||
|
||||
def save_b64_image(image_b64: str, output_dir: Path, stem: str, index: int, mime: str | None = None) -> Path:
|
||||
extension = mimetypes.guess_extension(mime or "image/png") or ".png"
|
||||
output_path = unique_path(output_dir / f"{stem}-{index}{extension}")
|
||||
output_path.write_bytes(base64.b64decode(image_b64))
|
||||
return output_path
|
||||
|
||||
|
||||
def save_data_uri(data_uri: str, output_dir: Path, stem: str, index: int) -> Path:
|
||||
match = re.match(r"^data:(?P<mime>[^;]+);base64,(?P<data>.+)$", data_uri, flags=re.DOTALL)
|
||||
if not match:
|
||||
raise ImageGenerationError("Unsupported data URI image format")
|
||||
return save_b64_image(match.group("data"), output_dir, stem, index, match.group("mime"))
|
||||
|
||||
|
||||
def extract_b64_items(data: Any) -> list[dict[str, Any]]:
|
||||
items: list[dict[str, Any]] = []
|
||||
|
||||
def visit(value: Any) -> None:
|
||||
if isinstance(value, dict):
|
||||
if isinstance(value.get("b64_json"), str):
|
||||
items.append(value)
|
||||
return
|
||||
for child in value.values():
|
||||
visit(child)
|
||||
elif isinstance(value, list):
|
||||
for child in value:
|
||||
visit(child)
|
||||
|
||||
visit(data)
|
||||
return items
|
||||
|
||||
|
||||
def extract_image_values(data: Any) -> list[str]:
|
||||
values: list[str] = []
|
||||
|
||||
def add(value: Any) -> None:
|
||||
if isinstance(value, str) and value.startswith(("http://", "https://", "data:image/")):
|
||||
values.append(value)
|
||||
elif isinstance(value, list):
|
||||
for item in value:
|
||||
add(item)
|
||||
|
||||
candidates = []
|
||||
if isinstance(data, dict):
|
||||
candidates.append(data.get("url"))
|
||||
nested_data = data.get("data")
|
||||
if isinstance(nested_data, list):
|
||||
candidates.extend(item.get("url") for item in nested_data if isinstance(item, dict))
|
||||
if isinstance(nested_data, dict):
|
||||
result = nested_data.get("result")
|
||||
if isinstance(result, dict):
|
||||
images = result.get("images")
|
||||
if isinstance(images, list):
|
||||
candidates.extend(item.get("url") for item in images if isinstance(item, dict))
|
||||
candidates.append(result.get("url"))
|
||||
result = data.get("result")
|
||||
if isinstance(result, dict):
|
||||
images = result.get("images")
|
||||
if isinstance(images, list):
|
||||
candidates.extend(item.get("url") for item in images if isinstance(item, dict))
|
||||
|
||||
for candidate in candidates:
|
||||
add(candidate)
|
||||
return values
|
||||
|
||||
|
||||
def extract_task_id(data: dict[str, Any]) -> str | None:
|
||||
for key in ("task_id", "id"):
|
||||
value = data.get(key)
|
||||
if isinstance(value, str) and value:
|
||||
return value
|
||||
|
||||
nested_data = data.get("data")
|
||||
if isinstance(nested_data, dict):
|
||||
for key in ("task_id", "id"):
|
||||
value = nested_data.get(key)
|
||||
if isinstance(value, str) and value:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def get_task_status(task: dict[str, Any]) -> str | None:
|
||||
nested_data = task.get("data")
|
||||
if isinstance(nested_data, dict) and isinstance(nested_data.get("status"), str):
|
||||
return nested_data["status"]
|
||||
status = task.get("status")
|
||||
return status if isinstance(status, str) else None
|
||||
|
||||
|
||||
def get_task_error(task: dict[str, Any]) -> str:
|
||||
nested_data = task.get("data")
|
||||
error = nested_data.get("error") if isinstance(nested_data, dict) else task.get("error")
|
||||
if isinstance(error, dict):
|
||||
return str(error.get("message") or error)
|
||||
if error:
|
||||
return str(error)
|
||||
return "task failed"
|
||||
|
||||
|
||||
def guess_download_extension(url: str, content_type: str | None) -> str:
|
||||
suffix = Path(urllib.parse.urlparse(url).path).suffix.lower()
|
||||
if suffix in {".png", ".jpg", ".jpeg", ".webp", ".gif"}:
|
||||
return suffix
|
||||
if content_type:
|
||||
return mimetypes.guess_extension(content_type.split(";")[0].strip()) or ".png"
|
||||
return ".png"
|
||||
|
||||
|
||||
def download_image(url: str, output_dir: Path, stem: str, index: int, timeout: int) -> Path:
|
||||
request = urllib.request.Request(url, headers={"User-Agent": "gpt-image-2-generator/1.0"})
|
||||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||||
content_type = response.headers.get("Content-Type")
|
||||
extension = guess_download_extension(url, content_type)
|
||||
output_path = unique_path(output_dir / f"{stem}-{index}{extension}")
|
||||
output_path.write_bytes(response.read())
|
||||
return output_path
|
||||
|
||||
|
||||
def save_image_values(values: list[str], output_dir: Path, stem: str, timeout: int) -> list[Path]:
|
||||
paths: list[Path] = []
|
||||
for index, value in enumerate(values, start=1):
|
||||
if value.startswith("data:image/"):
|
||||
paths.append(save_data_uri(value, output_dir, stem, index))
|
||||
else:
|
||||
paths.append(download_image(value, output_dir, stem, index, timeout))
|
||||
return paths
|
||||
|
||||
|
||||
def generate_image(
|
||||
prompt: str,
|
||||
size: str | None = None,
|
||||
image_refs: list[str] | None = None,
|
||||
n: int = 1,
|
||||
output_dir: str | Path | None = None,
|
||||
timeout: int = 300,
|
||||
initial_poll_delay: float = 10,
|
||||
poll_interval: float = 5,
|
||||
max_wait: float = 300,
|
||||
config_dir: str | Path | None = None,
|
||||
env_file: str | Path | None = None,
|
||||
base_url: str | None = None,
|
||||
model: str | None = None,
|
||||
request_style: str | None = None,
|
||||
quality: str | None = None,
|
||||
background: str | None = None,
|
||||
output_format: str | None = None,
|
||||
input_fidelity: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
prompt = prompt.strip()
|
||||
if not prompt:
|
||||
raise ImageGenerationError("Prompt is required")
|
||||
if n < 1:
|
||||
raise ImageGenerationError("n must be >= 1")
|
||||
|
||||
runtime_config = load_runtime_config(
|
||||
execution_root=Path.cwd(),
|
||||
config_dir=config_dir,
|
||||
env_file=env_file,
|
||||
base_url_override=base_url,
|
||||
model_override=model,
|
||||
output_dir_override=output_dir,
|
||||
require_token=True,
|
||||
)
|
||||
env_map = parse_env_file(runtime_config.env_file)
|
||||
resolved_request_style = request_style.strip().lower() if request_style else runtime_config.request_style
|
||||
if resolved_request_style == "auto":
|
||||
resolved_request_style = runtime_config.request_style
|
||||
if resolved_request_style not in {"apicodex", "openai"}:
|
||||
raise ImageGenerationError(
|
||||
f"Unsupported request style: {resolved_request_style}. Use one of: {', '.join(REQUEST_STYLE_CHOICES)}"
|
||||
)
|
||||
|
||||
refs = [ref.strip() for ref in (image_refs or []) if ref.strip()]
|
||||
output_path = runtime_config.output_dir
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||||
stem = f"gpt-image-2-{timestamp}"
|
||||
request_type = "edit" if refs else "generation"
|
||||
image_files = [image_ref_to_input_file(ref, index, timeout) for index, ref in enumerate(refs, start=1)] if refs else []
|
||||
targets = build_request_targets(runtime_config, resolved_request_style)
|
||||
response: dict[str, Any] | None = None
|
||||
selected_target: RequestTarget | None = None
|
||||
requested_size = ""
|
||||
api_size = ""
|
||||
last_error: BaseException | None = None
|
||||
for target_index, target in enumerate(targets):
|
||||
for attempt in range(1, REQUEST_RETRY_ATTEMPTS + 1):
|
||||
try:
|
||||
response, resolved_request_style, requested_size, api_size = submit_request_for_target(
|
||||
target,
|
||||
runtime_config=runtime_config,
|
||||
env_map=env_map,
|
||||
prompt=prompt,
|
||||
size=size,
|
||||
n=n,
|
||||
has_refs=bool(refs),
|
||||
image_files=image_files,
|
||||
timeout=timeout,
|
||||
quality=quality,
|
||||
background=background,
|
||||
output_format=output_format,
|
||||
input_fidelity=input_fidelity,
|
||||
)
|
||||
selected_target = target
|
||||
break
|
||||
except BaseException as error:
|
||||
last_error = error
|
||||
if attempt < REQUEST_RETRY_ATTEMPTS and should_retry_request(error):
|
||||
print(
|
||||
f"request attempt {attempt}/{REQUEST_RETRY_ATTEMPTS} to {target.base_url} failed with {format_request_error(error)}; retrying in {REQUEST_RETRY_DELAY_SECONDS:g}s",
|
||||
file=sys.stderr,
|
||||
)
|
||||
time.sleep(REQUEST_RETRY_DELAY_SECONDS)
|
||||
continue
|
||||
next_target = targets[target_index + 1] if target_index + 1 < len(targets) else None
|
||||
if next_target is not None and is_failover_eligible_error(error):
|
||||
print(
|
||||
f"request to {target.base_url} failed with {format_request_error(error)}; switching to fallback {next_target.base_url}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
break
|
||||
raise
|
||||
if response is not None and selected_target is not None:
|
||||
break
|
||||
if response is None or selected_target is None:
|
||||
raise ImageGenerationError(f"Image request failed: {last_error}")
|
||||
|
||||
task_api_base = build_api_url(selected_target.base_url, "tasks/")
|
||||
response_path = output_path / f"{stem}-response.json"
|
||||
save_json(response_path, response)
|
||||
|
||||
image_paths: list[Path] = []
|
||||
revised_prompts: list[str] = []
|
||||
|
||||
b64_items = extract_b64_items(response)
|
||||
if b64_items:
|
||||
for index, item in enumerate(b64_items, start=1):
|
||||
image_paths.append(save_b64_image(item["b64_json"], output_path, stem, index))
|
||||
if isinstance(item.get("revised_prompt"), str):
|
||||
revised_prompts.append(item["revised_prompt"])
|
||||
return {
|
||||
"images": [str(path) for path in image_paths],
|
||||
"response_path": str(response_path),
|
||||
"revised_prompts": revised_prompts,
|
||||
"request_type": request_type,
|
||||
"mode": "sync-b64",
|
||||
"base_url": selected_target.base_url,
|
||||
"model": runtime_config.model,
|
||||
"request_style": resolved_request_style,
|
||||
"requested_size": requested_size,
|
||||
"api_size": api_size,
|
||||
"output_dir": str(output_path),
|
||||
}
|
||||
|
||||
image_values = extract_image_values(response)
|
||||
if image_values:
|
||||
image_paths = save_image_values(image_values, output_path, stem, timeout)
|
||||
return {
|
||||
"images": [str(path) for path in image_paths],
|
||||
"response_path": str(response_path),
|
||||
"revised_prompts": revised_prompts,
|
||||
"request_type": request_type,
|
||||
"mode": "sync-url",
|
||||
"base_url": selected_target.base_url,
|
||||
"model": runtime_config.model,
|
||||
"request_style": resolved_request_style,
|
||||
"requested_size": requested_size,
|
||||
"api_size": api_size,
|
||||
"output_dir": str(output_path),
|
||||
}
|
||||
|
||||
task_id = extract_task_id(response)
|
||||
if not task_id:
|
||||
raise ImageGenerationError(f"Unexpected response. Saved raw response to {response_path}")
|
||||
|
||||
time.sleep(initial_poll_delay)
|
||||
deadline = time.time() + max_wait
|
||||
last_task: dict[str, Any] | None = None
|
||||
while time.time() <= deadline:
|
||||
last_task = request_json(
|
||||
"GET",
|
||||
urllib.parse.urljoin(task_api_base, urllib.parse.quote(task_id)),
|
||||
runtime_config.token or "",
|
||||
timeout=timeout,
|
||||
)
|
||||
status = get_task_status(last_task)
|
||||
if status == "completed":
|
||||
task_path = output_path / f"{stem}-task.json"
|
||||
save_json(task_path, last_task)
|
||||
|
||||
b64_items = extract_b64_items(last_task)
|
||||
if b64_items:
|
||||
for index, item in enumerate(b64_items, start=1):
|
||||
image_paths.append(save_b64_image(item["b64_json"], output_path, stem, index))
|
||||
if isinstance(item.get("revised_prompt"), str):
|
||||
revised_prompts.append(item["revised_prompt"])
|
||||
else:
|
||||
image_paths = save_image_values(extract_image_values(last_task), output_path, stem, timeout)
|
||||
|
||||
if not image_paths:
|
||||
raise ImageGenerationError(f"Task completed but no image was found. Saved task response to {task_path}")
|
||||
return {
|
||||
"images": [str(path) for path in image_paths],
|
||||
"response_path": str(response_path),
|
||||
"task_response_path": str(task_path),
|
||||
"revised_prompts": revised_prompts,
|
||||
"request_type": request_type,
|
||||
"mode": "async-task",
|
||||
"base_url": selected_target.base_url,
|
||||
"model": runtime_config.model,
|
||||
"request_style": resolved_request_style,
|
||||
"requested_size": requested_size,
|
||||
"api_size": api_size,
|
||||
"output_dir": str(output_path),
|
||||
}
|
||||
if status == "failed":
|
||||
raise ImageGenerationError(get_task_error(last_task))
|
||||
time.sleep(poll_interval)
|
||||
|
||||
if last_task is not None:
|
||||
task_path = output_path / f"{stem}-task-timeout.json"
|
||||
save_json(task_path, last_task)
|
||||
raise ImageGenerationError(f"Task polling timed out. Saved last task response to {task_path}")
|
||||
raise ImageGenerationError("Task polling timed out before the first task response")
|
||||
|
||||
|
||||
def save_token(token: str, config_dir: str | Path | None = None) -> Path:
|
||||
return save_runtime_token(token, execution_root=Path.cwd(), config_dir=config_dir)
|
||||
|
||||
|
||||
def save_base_url(base_url: str, config_dir: str | Path | None = None) -> Path:
|
||||
return save_runtime_base_url(base_url, execution_root=Path.cwd(), config_dir=config_dir)
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description="Generate or edit images with GPT Image 2 via apicodex-style or OpenAI-compatible APIs")
|
||||
parser.add_argument("--save-token", metavar="TOKEN", help="Save an API token to <cwd>/config/token.txt and exit")
|
||||
parser.add_argument("--save-base-url", metavar="URL", help="Save an API base URL to <cwd>/config/base_url.txt and exit")
|
||||
parser.add_argument("--print-env-template", action="store_true", help="Print a recommended .env template and exit")
|
||||
parser.add_argument("--init-config", action="store_true", help="Write a recommended env template file and exit")
|
||||
parser.add_argument("--prompt", required=False, default=None, help="Text prompt")
|
||||
parser.add_argument("--size", default=None, help="Size such as 2:3 or 1024x1536. If omitted, tries env defaults before falling back to 1:1.")
|
||||
parser.add_argument("--n", default=1, type=int, help="Number of images")
|
||||
parser.add_argument("--image-ref", action="append", default=[], help="Reference image URL, data URI, or local path")
|
||||
parser.add_argument("--output-dir", default=None, help="Output directory. Defaults to <cwd>/outputs.")
|
||||
parser.add_argument("--timeout", default=300, type=int, help="HTTP timeout in seconds")
|
||||
parser.add_argument("--initial-poll-delay", default=10, type=float, help="Seconds to wait before polling async tasks")
|
||||
parser.add_argument("--poll-interval", default=5, type=float, help="Async task poll interval in seconds")
|
||||
parser.add_argument("--max-wait", default=300, type=float, help="Maximum async task wait time in seconds")
|
||||
parser.add_argument("--config-dir", help="Config directory. Defaults to <cwd>/config.")
|
||||
parser.add_argument("--env-file", help="Env file. Defaults to <cwd>/.env.")
|
||||
parser.add_argument("--base-url", help="Override the API base URL for this request.")
|
||||
parser.add_argument("--model", help="Override the model for this request.")
|
||||
parser.add_argument("--request-style", choices=REQUEST_STYLE_CHOICES, default=None, help="Force request style. Defaults to auto/inferred from base URL.")
|
||||
parser.add_argument("--quality", default=None, help="Optional image edit quality value.")
|
||||
parser.add_argument("--background", default=None, help="Optional image edit background value.")
|
||||
parser.add_argument("--output-format", default=None, help="Optional image edit output format value.")
|
||||
parser.add_argument("--input-fidelity", default=None, help="Optional image edit input fidelity value.")
|
||||
parser.add_argument("--show-config", action="store_true", help="Print the resolved runtime configuration and exit")
|
||||
parser.add_argument("--force", action="store_true", help="Allow overwriting files when used with --init-config")
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = build_parser().parse_args()
|
||||
|
||||
if args.save_token:
|
||||
try:
|
||||
path = save_token(args.save_token, config_dir=args.config_dir)
|
||||
print(f"Token saved to {path}")
|
||||
return 0
|
||||
except (ConfigError, Exception) as error:
|
||||
print(f"ERROR: {error}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
if args.save_base_url:
|
||||
try:
|
||||
path = save_base_url(args.save_base_url, config_dir=args.config_dir)
|
||||
print(f"Base URL saved to {path}")
|
||||
return 0
|
||||
except (ConfigError, Exception) as error:
|
||||
print(f"ERROR: {error}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
if args.print_env_template:
|
||||
print(build_env_template(), end="")
|
||||
return 0
|
||||
|
||||
if args.init_config:
|
||||
try:
|
||||
path = write_env_template(env_file=args.env_file or ".env.example", overwrite=args.force)
|
||||
print(f"Config template written to {path}")
|
||||
return 0
|
||||
except (ConfigError, Exception) as error:
|
||||
print(f"ERROR: {error}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
if args.show_config:
|
||||
try:
|
||||
config = load_runtime_config(
|
||||
execution_root=Path.cwd(),
|
||||
config_dir=args.config_dir,
|
||||
env_file=args.env_file,
|
||||
base_url_override=args.base_url,
|
||||
model_override=args.model,
|
||||
output_dir_override=args.output_dir,
|
||||
require_token=False,
|
||||
)
|
||||
print(json.dumps(runtime_config_summary(config), ensure_ascii=False, indent=2))
|
||||
return 0
|
||||
except (ConfigError, Exception) as error:
|
||||
print(f"ERROR: {error}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
if not args.prompt:
|
||||
print(
|
||||
"ERROR: --prompt is required (unless using --save-token / --save-base-url / --print-env-template / --init-config / --show-config)",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 1
|
||||
|
||||
try:
|
||||
result = generate_image(
|
||||
prompt=args.prompt,
|
||||
size=args.size,
|
||||
image_refs=args.image_ref,
|
||||
n=args.n,
|
||||
output_dir=args.output_dir,
|
||||
timeout=args.timeout,
|
||||
initial_poll_delay=args.initial_poll_delay,
|
||||
poll_interval=args.poll_interval,
|
||||
max_wait=args.max_wait,
|
||||
config_dir=args.config_dir,
|
||||
env_file=args.env_file,
|
||||
base_url=args.base_url,
|
||||
model=args.model,
|
||||
request_style=args.request_style,
|
||||
quality=args.quality,
|
||||
background=args.background,
|
||||
output_format=args.output_format,
|
||||
input_fidelity=args.input_fidelity,
|
||||
)
|
||||
except (ConfigError, Exception) as error:
|
||||
print(f"ERROR: {error}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user