608 lines
21 KiB
Python
608 lines
21 KiB
Python
from __future__ import annotations
|
||
|
||
import base64
|
||
import json
|
||
import mimetypes
|
||
import os
|
||
import traceback
|
||
import time
|
||
from dataclasses import dataclass
|
||
from datetime import datetime
|
||
from io import BytesIO
|
||
from pathlib import Path
|
||
from typing import Any
|
||
from urllib.parse import urlparse
|
||
|
||
import requests
|
||
from PIL import Image, ImageFile
|
||
|
||
|
||
DEFAULT_PROMPT = "高端电子产品海报，干净背景，商业摄影，细节清晰"
|
||
DEFAULT_MODEL = "gpt-image-2"
|
||
DEFAULT_PIXEL_SIZE = "4096x4096"
|
||
DEFAULT_QUALITY = "high"
|
||
DEFAULT_OUTPUT_FILE = "output.png"
|
||
DEFAULT_RESPONSE_DIR = "outputs/image-demo-runs"
|
||
DEFAULT_CONNECT_TIMEOUT = 10
|
||
DEFAULT_READ_TIMEOUT = 900
|
||
DEFAULT_IMAGE_FETCH_TIMEOUT = 120
|
||
DEFAULT_POLL_CONNECT_TIMEOUT = 10
|
||
DEFAULT_POLL_READ_TIMEOUT = 60
|
||
DEFAULT_POLL_INTERVAL = 5
|
||
DEFAULT_INITIAL_POLL_DELAY = 10
|
||
DEFAULT_MAX_POLL_WAIT = 1800
|
||
DIRECT_HOSTS = {"vip.auto-code.net","api.toskaxy.xyz","api.aoixx.com"}
|
||
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",
|
||
}
|
||
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",
|
||
}
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class ImageInput:
|
||
filename: str
|
||
mime_type: str
|
||
data: bytes
|
||
|
||
|
||
def load_env_file(path: str = ".env") -> None:
|
||
env_path = Path(path)
|
||
if not env_path.exists():
|
||
return
|
||
|
||
for raw_line in env_path.read_text(encoding="utf-8").splitlines():
|
||
line = raw_line.strip()
|
||
if not line or line.startswith("#") or "=" not in line:
|
||
continue
|
||
|
||
key, value = line.split("=", 1)
|
||
key = key.strip()
|
||
value = value.strip().strip('"').strip("'")
|
||
if key and key not in os.environ:
|
||
os.environ[key] = value
|
||
|
||
|
||
def is_pixel_size(size: str) -> bool:
|
||
parts = size.lower().split("x")
|
||
return len(parts) == 2 and all(part.isdigit() for part in parts)
|
||
|
||
|
||
def normalize_size(size: str) -> str:
|
||
value = size.strip().lower().replace(" ", "")
|
||
if is_pixel_size(value):
|
||
width, height = value.split("x", 1)
|
||
return f"{int(width)}x{int(height)}"
|
||
return value
|
||
|
||
|
||
def is_solid_black_image(img: Image.Image) -> bool:
|
||
rgb = img.convert("RGB")
|
||
return all(lo == 0 and hi == 0 for lo, hi in rgb.getextrema())
|
||
|
||
|
||
def env_int(name: str, default: int) -> int:
|
||
raw_value = os.getenv(name, "").strip()
|
||
if not raw_value:
|
||
return default
|
||
|
||
try:
|
||
return int(raw_value)
|
||
except ValueError as exc:
|
||
raise RuntimeError(f"{name} 不是有效整数: {raw_value}") from exc
|
||
|
||
|
||
def make_run_dir() -> Path:
|
||
response_root = Path(os.getenv("IMAGE_RESPONSE_DIR", DEFAULT_RESPONSE_DIR)).expanduser()
|
||
run_id = datetime.now().strftime("%Y%m%d-%H%M%S-%f")
|
||
run_dir = response_root / run_id
|
||
run_dir.mkdir(parents=True, exist_ok=False)
|
||
return run_dir
|
||
|
||
|
||
def dump_json(path: Path, data: Any) -> None:
|
||
path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
|
||
|
||
def save_http_response(path: Path, response: requests.Response) -> dict[str, Any]:
|
||
record: dict[str, Any] = {
|
||
"url": response.url,
|
||
"status_code": response.status_code,
|
||
"headers": dict(response.headers),
|
||
}
|
||
|
||
try:
|
||
record["json"] = response.json()
|
||
except ValueError:
|
||
record["body_text"] = response.text
|
||
|
||
dump_json(path, record)
|
||
return record
|
||
|
||
|
||
def save_error_artifact(path: Path, error: BaseException, extra: dict[str, Any] | None = None) -> None:
|
||
payload: dict[str, Any] = {
|
||
"error_type": type(error).__name__,
|
||
"error_message": str(error),
|
||
"traceback": traceback.format_exc(),
|
||
}
|
||
if extra:
|
||
payload.update(extra)
|
||
dump_json(path, payload)
|
||
|
||
|
||
def resolve_output_path(run_dir: Path, output_file: str) -> Path:
|
||
output_path = Path(output_file).expanduser()
|
||
if output_path.is_absolute():
|
||
return output_path
|
||
return run_dir / output_path
|
||
|
||
|
||
def normalize_base_url(base_url: str) -> str:
|
||
normalized = base_url.rstrip("/")
|
||
if normalized.endswith("/v1"):
|
||
return normalized[:-3]
|
||
return normalized
|
||
|
||
|
||
def build_request_session(base_url: str) -> requests.Session:
|
||
session = requests.Session()
|
||
hostname = urlparse(base_url).hostname or ""
|
||
if hostname in DIRECT_HOSTS:
|
||
session.trust_env = False
|
||
print(f"直连 {hostname}，忽略本地代理环境变量")
|
||
return session
|
||
|
||
|
||
def resolve_request_style(base_url: str) -> str:
|
||
request_style = os.getenv("IMAGE_REQUEST_STYLE", "auto").strip().lower()
|
||
if request_style in {"openai", "apicodex"}:
|
||
return request_style
|
||
|
||
hostname = (urlparse(base_url).hostname or "").lower()
|
||
if hostname.endswith("openai.com"):
|
||
return "openai"
|
||
return "apicodex"
|
||
|
||
|
||
def resolve_api_size(size: str, request_style: str) -> str:
|
||
if request_style == "openai" and size in OPENAI_RATIO_SIZE_MAP:
|
||
return OPENAI_RATIO_SIZE_MAP[size]
|
||
return size
|
||
|
||
|
||
def resolve_image_source(prefix: str) -> str | None:
|
||
candidates = {
|
||
prefix: os.getenv(prefix, "").strip(),
|
||
f"{prefix}_FILE": os.getenv(f"{prefix}_FILE", "").strip(),
|
||
f"{prefix}_URL": os.getenv(f"{prefix}_URL", "").strip(),
|
||
f"{prefix}_DATA_URL": os.getenv(f"{prefix}_DATA_URL", "").strip(),
|
||
}
|
||
provided = [(name, value) for name, value in candidates.items() if value]
|
||
if len(provided) > 1:
|
||
raise RuntimeError(
|
||
f"{prefix} 相关变量只能设置一个，当前同时设置了: "
|
||
+ ", ".join(name for name, _ in provided)
|
||
)
|
||
if not provided:
|
||
return None
|
||
return provided[0][1]
|
||
|
||
|
||
def describe_image_source(source: str | None, fallback_name: str) -> dict[str, str] | None:
|
||
if source is None:
|
||
return None
|
||
|
||
value = source.strip()
|
||
if not value:
|
||
return None
|
||
|
||
lowered = value.lower()
|
||
if lowered.startswith("data:image/"):
|
||
source_kind = "data_url"
|
||
filename = fallback_name
|
||
elif lowered.startswith(("http://", "https://")):
|
||
source_kind = "url"
|
||
filename = Path(urlparse(value).path).name or fallback_name
|
||
else:
|
||
source_kind = "file"
|
||
filename = Path(value).expanduser().name or fallback_name
|
||
|
||
return {
|
||
"source": value,
|
||
"kind": source_kind,
|
||
"filename": filename,
|
||
}
|
||
|
||
|
||
def safe_filename(filename: str, fallback: str, mime_type: str) -> str:
|
||
candidate = 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 data_url_to_image_input(data_url: str, fallback: str) -> ImageInput:
|
||
if not data_url.startswith("data:image/") or "," not in data_url:
|
||
raise RuntimeError("data URL 参考图格式无效")
|
||
|
||
header, data_b64 = data_url.split(",", 1)
|
||
if ";base64" not in header.lower():
|
||
raise RuntimeError("data URL 参考图必须使用 base64 编码")
|
||
|
||
mime_type = header[5:].split(";", 1)[0] or "image/png"
|
||
try:
|
||
data = base64.b64decode(data_b64, validate=True)
|
||
except ValueError as exc:
|
||
raise RuntimeError("data URL 参考图 base64 无法解码") from exc
|
||
|
||
return ImageInput(
|
||
filename=safe_filename("", fallback, mime_type),
|
||
mime_type=mime_type,
|
||
data=data,
|
||
)
|
||
|
||
|
||
def local_file_to_image_input(file_path: str, fallback: str) -> ImageInput:
|
||
path = Path(file_path).expanduser()
|
||
if not path.is_absolute():
|
||
path = Path.cwd() / path
|
||
if not path.exists():
|
||
raise FileNotFoundError(f"找不到参考图文件: {file_path}")
|
||
|
||
mime_type = mimetypes.guess_type(path.name)[0] or "image/png"
|
||
return ImageInput(
|
||
filename=safe_filename(path.name, fallback, mime_type),
|
||
mime_type=mime_type,
|
||
data=path.read_bytes(),
|
||
)
|
||
|
||
|
||
def fetch_url_to_image_input(image_url: str, fallback: str, session: requests.Session) -> ImageInput:
|
||
resp = session.get(image_url, timeout=(DEFAULT_CONNECT_TIMEOUT, DEFAULT_IMAGE_FETCH_TIMEOUT))
|
||
resp.raise_for_status()
|
||
|
||
mime_type = (resp.headers.get("Content-Type") or "").split(";", 1)[0].strip()
|
||
if not mime_type.startswith("image/"):
|
||
mime_type = mimetypes.guess_type(urlparse(image_url).path)[0] or "image/png"
|
||
|
||
filename = safe_filename(Path(urlparse(image_url).path).name, fallback, mime_type)
|
||
return ImageInput(filename=filename, mime_type=mime_type, data=resp.content)
|
||
|
||
|
||
def image_source_to_input(source: str, fallback: str, session: requests.Session) -> ImageInput:
|
||
value = source.strip()
|
||
lowered = value.lower()
|
||
if lowered.startswith("data:image/"):
|
||
return data_url_to_image_input(value, fallback)
|
||
if lowered.startswith(("http://", "https://")):
|
||
return fetch_url_to_image_input(value, fallback, session)
|
||
return local_file_to_image_input(value, fallback)
|
||
|
||
|
||
def build_payload(base_url: str) -> tuple[dict[str, Any], str, str, str, str, str | None, str | None]:
|
||
model = os.getenv("IMAGE_MODEL", DEFAULT_MODEL).strip()
|
||
prompt = os.getenv("IMAGE_PROMPT", DEFAULT_PROMPT).strip() or DEFAULT_PROMPT
|
||
requested_size = normalize_size(os.getenv("IMAGE_SIZE", DEFAULT_PIXEL_SIZE))
|
||
resolution = os.getenv("IMAGE_RESOLUTION", "").strip()
|
||
quality = os.getenv("IMAGE_QUALITY", DEFAULT_QUALITY).strip()
|
||
output_file = os.getenv("OUTPUT_FILE", DEFAULT_OUTPUT_FILE).strip()
|
||
request_style = resolve_request_style(base_url)
|
||
if not (is_pixel_size(requested_size) or requested_size in SUPPORTED_ASPECT_RATIOS):
|
||
raise RuntimeError(f"不支持的 IMAGE_SIZE: {requested_size}")
|
||
api_size = resolve_api_size(requested_size, request_style)
|
||
|
||
payload: dict[str, Any] = {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"n": int(os.getenv("IMAGE_N", "1")),
|
||
"size": api_size,
|
||
"quality": quality,
|
||
}
|
||
if not (request_style == "openai" and model.startswith("gpt-image")):
|
||
payload["response_format"] = "b64_json"
|
||
|
||
if not is_pixel_size(requested_size) and request_style != "openai":
|
||
payload["resolution"] = resolution or "4K"
|
||
|
||
reference_source = resolve_image_source("REFERENCE_IMAGE")
|
||
mask_source = resolve_image_source("MASK_IMAGE")
|
||
if mask_source and not reference_source:
|
||
raise RuntimeError("MASK_IMAGE 只能和 REFERENCE_IMAGE 一起使用")
|
||
|
||
request_mode = "edit" if reference_source else "generate"
|
||
return payload, model, requested_size, output_file, request_mode, reference_source, mask_source
|
||
|
||
|
||
def payload_to_form_fields(payload: dict[str, Any]) -> dict[str, str]:
|
||
fields: dict[str, str] = {}
|
||
for key, value in payload.items():
|
||
if value is None:
|
||
continue
|
||
fields[key] = str(value)
|
||
return fields
|
||
|
||
|
||
def build_edit_files(
|
||
request_style: str,
|
||
reference_input: ImageInput,
|
||
mask_input: ImageInput | None,
|
||
) -> list[tuple[str, tuple[str, bytes, str]]]:
|
||
image_field_name = "image" if request_style == "openai" else "image[]"
|
||
files: list[tuple[str, tuple[str, bytes, str]]] = [
|
||
(
|
||
image_field_name,
|
||
(reference_input.filename, reference_input.data, reference_input.mime_type),
|
||
)
|
||
]
|
||
if mask_input is not None:
|
||
files.append(("mask", (mask_input.filename, mask_input.data, mask_input.mime_type)))
|
||
return files
|
||
|
||
|
||
def resolve_min_output_edge(requested_size: str, resolution: str | None, request_style: str) -> int | None:
|
||
if is_pixel_size(requested_size):
|
||
width, height = requested_size.split("x", 1)
|
||
min_edge = max(int(width), int(height))
|
||
return min_edge if min_edge >= 3840 else None
|
||
if resolution and resolution.strip().upper() == "4K" and request_style != "openai":
|
||
return 3840
|
||
return None
|
||
|
||
|
||
def save_image_bytes(image_bytes: bytes, output_path: Path, min_output_edge: int | None = None) -> tuple[int, int]:
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
previous_truncated_setting = ImageFile.LOAD_TRUNCATED_IMAGES
|
||
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
||
try:
|
||
with Image.open(BytesIO(image_bytes)) as img:
|
||
img.load()
|
||
width, height = img.size
|
||
if is_solid_black_image(img):
|
||
raise RuntimeError(
|
||
"上游返回了疑似纯黑占位图 "
|
||
f"({width}x{height})，通常表示请求没有走到正常生成链路。"
|
||
"请检查 OPENAI_BASE_URL、IMAGE_REQUEST_STYLE、model 和 size，"
|
||
"必要时切换到 vip.auto-code.net。"
|
||
)
|
||
|
||
if min_output_edge is not None and max(width, height) < min_output_edge:
|
||
raise RuntimeError(f"生成结果未达到预期尺寸下限 {min_output_edge}px：{width}x{height}")
|
||
finally:
|
||
ImageFile.LOAD_TRUNCATED_IMAGES = previous_truncated_setting
|
||
|
||
output_path.write_bytes(image_bytes)
|
||
print(f"已保存 {output_path} ({width}x{height})")
|
||
|
||
return width, height
|
||
|
||
|
||
def extract_image_bytes_from_item(item: dict[str, Any], session: requests.Session) -> bytes:
|
||
if item.get("b64_json"):
|
||
return base64.b64decode(item["b64_json"])
|
||
|
||
url_value = item.get("url")
|
||
if isinstance(url_value, list) and url_value:
|
||
image_url = url_value[0]
|
||
elif isinstance(url_value, str):
|
||
image_url = url_value
|
||
else:
|
||
image_url = None
|
||
|
||
if image_url:
|
||
resp = session.get(image_url, timeout=(DEFAULT_CONNECT_TIMEOUT, DEFAULT_IMAGE_FETCH_TIMEOUT))
|
||
resp.raise_for_status()
|
||
return resp.content
|
||
|
||
raise RuntimeError(f"响应里没有可用的图片数据: {item}")
|
||
|
||
|
||
def poll_task(
|
||
task_id: str,
|
||
base_url: str,
|
||
headers: dict[str, str],
|
||
session: requests.Session,
|
||
run_dir: Path,
|
||
max_wait_seconds: int,
|
||
initial_poll_delay: int,
|
||
poll_interval: int,
|
||
) -> bytes:
|
||
task_url = f"{base_url}/v1/tasks/{task_id}"
|
||
deadline = time.monotonic() + max_wait_seconds
|
||
poll_index = 0
|
||
|
||
time.sleep(initial_poll_delay)
|
||
while True:
|
||
if time.monotonic() >= deadline:
|
||
raise TimeoutError(f"任务轮询超时: task_id={task_id}")
|
||
|
||
poll_index += 1
|
||
task_resp = session.get(
|
||
task_url,
|
||
headers=headers,
|
||
timeout=(DEFAULT_POLL_CONNECT_TIMEOUT, DEFAULT_POLL_READ_TIMEOUT),
|
||
)
|
||
task_record = save_http_response(run_dir / f"poll-{poll_index:03d}.json", task_resp)
|
||
task_resp.raise_for_status()
|
||
task = task_record.get("json")
|
||
if not isinstance(task, dict):
|
||
raise RuntimeError(f"任务响应不是 JSON 对象: {task_record}")
|
||
|
||
data = task.get("data")
|
||
if not isinstance(data, dict):
|
||
raise RuntimeError(f"任务响应格式异常: {task}")
|
||
|
||
status = data.get("status")
|
||
print(f"任务轮询 #{poll_index}: status={status}")
|
||
if status == "completed":
|
||
result = data.get("result", {})
|
||
images = result.get("images", [])
|
||
if not images:
|
||
raise RuntimeError(f"任务已完成，但没有返回图片: {task}")
|
||
return extract_image_bytes_from_item(images[0], session)
|
||
|
||
if status == "failed":
|
||
error = data.get("error", {})
|
||
if isinstance(error, dict):
|
||
message = error.get("message") or str(error)
|
||
else:
|
||
message = str(error)
|
||
raise RuntimeError(message or "任务失败")
|
||
|
||
time.sleep(poll_interval)
|
||
|
||
|
||
def main() -> None:
|
||
load_env_file()
|
||
|
||
api_key = os.getenv("OPENAI_API_KEY") or os.getenv("API_KEY")
|
||
base_url = normalize_base_url(
|
||
os.getenv("OPENAI_BASE_URL") or os.getenv("BASE_URL") or "https://api.toskaxy.xyz"
|
||
)
|
||
if not api_key:
|
||
raise RuntimeError("未找到 API Key，请先在 .env 中设置 OPENAI_API_KEY")
|
||
|
||
payload, model, size, output_file, request_mode, reference_source, mask_source = build_payload(base_url)
|
||
resolution = payload.get("resolution")
|
||
run_dir = make_run_dir()
|
||
output_path = resolve_output_path(run_dir, output_file)
|
||
session = build_request_session(base_url)
|
||
request_style = resolve_request_style(base_url)
|
||
request_timeout = env_int("IMAGE_REQUEST_TIMEOUT", DEFAULT_READ_TIMEOUT)
|
||
max_poll_wait = env_int("IMAGE_MAX_POLL_WAIT", DEFAULT_MAX_POLL_WAIT)
|
||
initial_poll_delay = env_int("IMAGE_INITIAL_POLL_DELAY", DEFAULT_INITIAL_POLL_DELAY)
|
||
poll_interval = env_int("IMAGE_POLL_INTERVAL", DEFAULT_POLL_INTERVAL)
|
||
endpoint_path = "/v1/images/edits" if request_mode == "edit" else "/v1/images/generations"
|
||
min_output_edge = resolve_min_output_edge(size, resolution, request_style)
|
||
reference_image_info = describe_image_source(reference_source, "reference-image")
|
||
mask_image_info = describe_image_source(mask_source, "mask-image")
|
||
|
||
print(
|
||
"开始生成图片:",
|
||
f"mode={request_mode}",
|
||
f"request_style={request_style}",
|
||
f"model={model}",
|
||
f"size={size}",
|
||
f"resolution={resolution or 'none'}",
|
||
f"request_timeout={request_timeout}s",
|
||
f"max_poll_wait={max_poll_wait}s",
|
||
f"initial_poll_delay={initial_poll_delay}s",
|
||
f"poll_interval={poll_interval}s",
|
||
f"endpoint={base_url}{endpoint_path}",
|
||
f"output={output_path}",
|
||
f"responses={run_dir}",
|
||
)
|
||
|
||
dump_json(run_dir / "request.json", {
|
||
"base_url": base_url,
|
||
"endpoint_path": endpoint_path,
|
||
"request_mode": request_mode,
|
||
"model": model,
|
||
"size": size,
|
||
"resolution": resolution,
|
||
"output_file": str(output_path),
|
||
"request_timeout_seconds": request_timeout,
|
||
"max_poll_wait_seconds": max_poll_wait,
|
||
"initial_poll_delay_seconds": initial_poll_delay,
|
||
"poll_interval_seconds": poll_interval,
|
||
"request_style": request_style,
|
||
"min_output_edge": min_output_edge,
|
||
"reference_image_present": bool(reference_source),
|
||
"reference_image_source": reference_image_info["source"] if reference_image_info else None,
|
||
"reference_image_source_kind": reference_image_info["kind"] if reference_image_info else None,
|
||
"reference_image_filename": reference_image_info["filename"] if reference_image_info else None,
|
||
"mask_image_present": bool(mask_source),
|
||
"mask_image_source": mask_image_info["source"] if mask_image_info else None,
|
||
"mask_image_source_kind": mask_image_info["kind"] if mask_image_info else None,
|
||
"mask_image_filename": mask_image_info["filename"] if mask_image_info else None,
|
||
"payload": payload,
|
||
})
|
||
|
||
try:
|
||
if request_mode == "edit":
|
||
reference_input = image_source_to_input(reference_source or "", "reference-image", session)
|
||
mask_input = (
|
||
image_source_to_input(mask_source, "mask-image", session)
|
||
if mask_source
|
||
else None
|
||
)
|
||
resp = session.post(
|
||
f"{base_url}{endpoint_path}",
|
||
data=payload_to_form_fields(payload),
|
||
files=build_edit_files(request_style, reference_input, mask_input),
|
||
headers={"Authorization": f"Bearer {api_key}"},
|
||
timeout=(DEFAULT_CONNECT_TIMEOUT, request_timeout),
|
||
)
|
||
else:
|
||
resp = session.post(
|
||
f"{base_url}{endpoint_path}",
|
||
json=payload,
|
||
headers={
|
||
"Authorization": f"Bearer {api_key}",
|
||
"Content-Type": "application/json",
|
||
},
|
||
timeout=(DEFAULT_CONNECT_TIMEOUT, request_timeout),
|
||
)
|
||
post_record = save_http_response(run_dir / "post-response.json", resp)
|
||
resp.raise_for_status()
|
||
data = post_record.get("json")
|
||
if not isinstance(data, dict):
|
||
raise RuntimeError(f"post response is not a JSON object: {post_record}")
|
||
|
||
response_data = data.get("data")
|
||
if isinstance(response_data, list) and response_data:
|
||
image_bytes = extract_image_bytes_from_item(response_data[0], session)
|
||
save_image_bytes(image_bytes, output_path, min_output_edge=min_output_edge)
|
||
return
|
||
|
||
if isinstance(response_data, dict) and response_data.get("id"):
|
||
image_bytes = poll_task(
|
||
str(response_data["id"]),
|
||
base_url,
|
||
{"Authorization": f"Bearer {api_key}"},
|
||
session,
|
||
run_dir,
|
||
max_poll_wait,
|
||
initial_poll_delay,
|
||
poll_interval,
|
||
)
|
||
save_image_bytes(image_bytes, output_path, min_output_edge=min_output_edge)
|
||
return
|
||
|
||
raise RuntimeError(f"unexpected response: {data}")
|
||
except BaseException as exc:
|
||
save_error_artifact(run_dir / "error.json", exc, {"output_file": str(output_path)})
|
||
raise
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|