619 lines
22 KiB
Python
619 lines
22 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import re
|
||
import time
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import requests
|
||
|
||
from gpt_image2_generate import (
|
||
DEFAULT_CONNECT_TIMEOUT,
|
||
DEFAULT_INITIAL_POLL_DELAY,
|
||
DEFAULT_MAX_POLL_WAIT,
|
||
DEFAULT_MODEL,
|
||
DEFAULT_PROMPT,
|
||
DEFAULT_PIXEL_SIZE,
|
||
DEFAULT_QUALITY,
|
||
DEFAULT_READ_TIMEOUT,
|
||
DEFAULT_POLL_CONNECT_TIMEOUT,
|
||
DEFAULT_POLL_INTERVAL,
|
||
DEFAULT_POLL_READ_TIMEOUT,
|
||
EmptyStreamResponseError,
|
||
ImageInput,
|
||
build_non_stream_payload,
|
||
build_request_session,
|
||
classify_stream_event,
|
||
dump_json,
|
||
env_int,
|
||
env_optional_int,
|
||
extract_api_error_message,
|
||
extract_stream_image_item,
|
||
extract_image_bytes_from_item,
|
||
is_event_stream_response,
|
||
is_gpt_image_model,
|
||
image_source_to_input,
|
||
iter_stream_event_records,
|
||
is_pixel_size,
|
||
load_env_file,
|
||
make_run_dir,
|
||
normalize_base_url,
|
||
normalize_size,
|
||
payload_to_form_fields,
|
||
resolve_optional_enum,
|
||
resolve_api_size,
|
||
resolve_min_output_edge,
|
||
resolve_request_style,
|
||
resolve_stream_options,
|
||
resolve_output_extension,
|
||
save_error_artifact,
|
||
save_http_response,
|
||
save_image_bytes,
|
||
save_stream_response_metadata,
|
||
SUPPORTED_MODERATION_LEVELS,
|
||
SUPPORTED_OUTPUT_FORMATS,
|
||
STREAM_COMPLETED_EVENT_TYPES,
|
||
STREAM_PARTIAL_EVENT_TYPES,
|
||
)
|
||
|
||
|
||
DEFAULT_OUTPUT_STEM = "gpt-image-2-multi"
|
||
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",
|
||
}
|
||
|
||
|
||
def _json_dumps(data: Any) -> str:
|
||
return json.dumps(data, ensure_ascii=False, indent=2)
|
||
|
||
|
||
def split_reference_values(raw_value: str) -> list[str]:
|
||
value = raw_value.strip()
|
||
if not value:
|
||
return []
|
||
|
||
if value.startswith("["):
|
||
try:
|
||
parsed = json.loads(value)
|
||
except json.JSONDecodeError:
|
||
inner = value[1:-1].strip()
|
||
if not inner:
|
||
return []
|
||
parts = [part.strip().strip('"').strip("'") for part in inner.split(",")]
|
||
return [part for part in parts if part]
|
||
|
||
if not isinstance(parsed, list):
|
||
raise RuntimeError("REFERENCE_IMAGES 必须是 JSON 数组")
|
||
return [str(item).strip() for item in parsed if str(item).strip()]
|
||
|
||
return [part.strip() for part in value.split("|") if part.strip()]
|
||
|
||
|
||
def collect_reference_sources() -> list[str]:
|
||
raw_list = os.getenv("REFERENCE_IMAGES", "").strip()
|
||
if raw_list:
|
||
refs = split_reference_values(raw_list)
|
||
else:
|
||
refs = []
|
||
for index in range(1, 17):
|
||
value = os.getenv(f"REFERENCE_IMAGE_{index}", "").strip()
|
||
if value:
|
||
refs.append(value)
|
||
if not refs:
|
||
single_candidates = [
|
||
os.getenv("REFERENCE_IMAGE", "").strip(),
|
||
os.getenv("REFERENCE_IMAGE_FILE", "").strip(),
|
||
os.getenv("REFERENCE_IMAGE_URL", "").strip(),
|
||
os.getenv("REFERENCE_IMAGE_DATA_URL", "").strip(),
|
||
]
|
||
refs.extend([value for value in single_candidates if value])
|
||
|
||
deduped: list[str] = []
|
||
seen: set[str] = set()
|
||
for ref in refs:
|
||
if ref not in seen:
|
||
deduped.append(ref)
|
||
seen.add(ref)
|
||
return deduped
|
||
|
||
|
||
def collect_mask_source() -> str | None:
|
||
candidates = {
|
||
"MASK_IMAGE": os.getenv("MASK_IMAGE", "").strip(),
|
||
"MASK_IMAGE_FILE": os.getenv("MASK_IMAGE_FILE", "").strip(),
|
||
"MASK_IMAGE_URL": os.getenv("MASK_IMAGE_URL", "").strip(),
|
||
"MASK_IMAGE_DATA_URL": os.getenv("MASK_IMAGE_DATA_URL", "").strip(),
|
||
}
|
||
provided = [value for value in candidates.values() if value]
|
||
if len(provided) > 1:
|
||
raise RuntimeError("MASK_IMAGE 相关变量只能设置一个")
|
||
return provided[0] if provided else None
|
||
|
||
|
||
def resolve_output_stem() -> str:
|
||
raw_value = os.getenv("IMAGE_OUTPUT_STEM", DEFAULT_OUTPUT_STEM).strip()
|
||
candidate = re.sub(r"[^A-Za-z0-9._-]+", "_", Path(raw_value).name).strip("._")
|
||
return candidate or DEFAULT_OUTPUT_STEM
|
||
|
||
|
||
def validate_requested_size(requested_size: str) -> None:
|
||
if not (is_pixel_size(requested_size) or requested_size in SUPPORTED_ASPECT_RATIOS):
|
||
raise RuntimeError(f"不支持的 IMAGE_SIZE: {requested_size}")
|
||
|
||
|
||
def build_payload(base_url: str, reference_sources: list[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))
|
||
validate_requested_size(requested_size)
|
||
request_style = resolve_request_style(base_url)
|
||
api_size = resolve_api_size(requested_size, request_style)
|
||
quality = os.getenv("IMAGE_QUALITY", DEFAULT_QUALITY).strip()
|
||
resolution = os.getenv("IMAGE_RESOLUTION", "").strip()
|
||
background = os.getenv("IMAGE_BACKGROUND", "").strip() or None
|
||
input_fidelity = os.getenv("IMAGE_INPUT_FIDELITY", "").strip() or None
|
||
output_format = resolve_optional_enum("IMAGE_OUTPUT_FORMAT", SUPPORTED_OUTPUT_FORMATS)
|
||
moderation = resolve_optional_enum("IMAGE_MODERATION", SUPPORTED_MODERATION_LEVELS)
|
||
output_compression = env_optional_int("IMAGE_OUTPUT_COMPRESSION")
|
||
stream_requested, partial_images = resolve_stream_options(model)
|
||
|
||
n = env_int("IMAGE_N", max(1, len(reference_sources)))
|
||
if n < 1:
|
||
raise RuntimeError("IMAGE_N 必须大于等于 1")
|
||
if moderation and not is_gpt_image_model(model):
|
||
raise RuntimeError(f"当前 model 不支持 moderation 参数: {model}")
|
||
if output_compression is not None:
|
||
if output_compression < 0 or output_compression > 100:
|
||
raise RuntimeError("IMAGE_OUTPUT_COMPRESSION 仅支持 0-100")
|
||
if output_format not in {"jpeg", "webp"}:
|
||
raise RuntimeError("IMAGE_OUTPUT_COMPRESSION 仅适用于 IMAGE_OUTPUT_FORMAT=jpeg/webp")
|
||
|
||
payload: dict[str, Any] = {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"n": n,
|
||
"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"
|
||
if background:
|
||
payload["background"] = background
|
||
if input_fidelity:
|
||
payload["input_fidelity"] = input_fidelity
|
||
if output_format:
|
||
payload["output_format"] = output_format
|
||
if output_compression is not None:
|
||
payload["output_compression"] = output_compression
|
||
if moderation:
|
||
payload["moderation"] = moderation
|
||
if stream_requested:
|
||
payload["stream"] = True
|
||
if partial_images is not None:
|
||
payload["partial_images"] = partial_images
|
||
|
||
mask_source = collect_mask_source()
|
||
if mask_source and not reference_sources:
|
||
raise RuntimeError("MASK_IMAGE 只能和参考图一起使用")
|
||
|
||
request_mode = "edit" if reference_sources else "generate"
|
||
return payload, model, requested_size, request_mode, mask_source, background, input_fidelity
|
||
|
||
|
||
def build_edit_files(
|
||
request_style: str,
|
||
reference_inputs: list[ImageInput],
|
||
mask_input: ImageInput | None,
|
||
) -> list[tuple[str, tuple[str, bytes, str]]]:
|
||
field_name = "image" if request_style == "openai" and len(reference_inputs) == 1 else "image[]"
|
||
files: list[tuple[str, tuple[str, bytes, str]]] = [
|
||
(field_name, (ref.filename, ref.data, ref.mime_type))
|
||
for ref in reference_inputs
|
||
]
|
||
if mask_input is not None:
|
||
files.append(("mask", (mask_input.filename, mask_input.data, mask_input.mime_type)))
|
||
return files
|
||
|
||
|
||
def extract_task_id(data: Any) -> str | None:
|
||
if not isinstance(data, dict):
|
||
return None
|
||
for key in ("task_id", "id"):
|
||
value = data.get(key)
|
||
if isinstance(value, str) and value:
|
||
return value
|
||
nested = data.get("data")
|
||
if isinstance(nested, dict):
|
||
for key in ("task_id", "id"):
|
||
value = nested.get(key)
|
||
if isinstance(value, str) and value:
|
||
return value
|
||
return None
|
||
|
||
|
||
def get_task_error(task: dict[str, Any]) -> str:
|
||
nested = task.get("data")
|
||
error = nested.get("error") if isinstance(nested, 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 extract_image_items(data: Any) -> list[dict[str, Any]]:
|
||
if isinstance(data, list):
|
||
return [item for item in data if isinstance(item, dict)]
|
||
|
||
if isinstance(data, dict):
|
||
for key in ("images", "data"):
|
||
nested = data.get(key)
|
||
if isinstance(nested, list):
|
||
return [item for item in nested if isinstance(item, dict)]
|
||
|
||
result = data.get("result")
|
||
if isinstance(result, dict):
|
||
images = result.get("images")
|
||
if isinstance(images, list):
|
||
return [item for item in images if isinstance(item, dict)]
|
||
|
||
return []
|
||
|
||
|
||
def poll_task_images(
|
||
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,
|
||
) -> list[dict[str, Any]]:
|
||
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 [item for item in images if isinstance(item, dict)]
|
||
|
||
if status == "failed":
|
||
raise RuntimeError(get_task_error(task))
|
||
|
||
time.sleep(poll_interval)
|
||
|
||
|
||
def save_output_images(
|
||
items: list[dict[str, Any]],
|
||
session: requests.Session,
|
||
run_dir: Path,
|
||
stem: str,
|
||
min_output_edge: int | None,
|
||
) -> list[str]:
|
||
paths: list[str] = []
|
||
for index, item in enumerate(items, start=1):
|
||
image_bytes = extract_image_bytes_from_item(item, session)
|
||
output_path = run_dir / f"{stem}-{index}.png"
|
||
save_image_bytes(image_bytes, output_path, min_output_edge=min_output_edge)
|
||
paths.append(str(output_path))
|
||
return paths
|
||
|
||
|
||
def handle_streaming_multi_response(
|
||
response: requests.Response,
|
||
session: requests.Session,
|
||
run_dir: Path,
|
||
stem: str,
|
||
min_output_edge: int | None,
|
||
) -> tuple[list[str], list[str]]:
|
||
partial_dir = run_dir / "partials"
|
||
partial_dir.mkdir(parents=True, exist_ok=True)
|
||
partial_paths: list[str] = []
|
||
image_paths: list[str] = []
|
||
record_count = 0
|
||
|
||
save_stream_response_metadata(run_dir / "post-response.json", response)
|
||
response.raise_for_status()
|
||
for record in iter_stream_event_records(run_dir / "stream-events.jsonl", response):
|
||
record_count += 1
|
||
payload = record.get("data")
|
||
if not isinstance(payload, dict):
|
||
continue
|
||
|
||
event_kind = classify_stream_event(record, payload)
|
||
image_item = extract_stream_image_item(payload)
|
||
if event_kind is None or image_item is None:
|
||
continue
|
||
|
||
image_bytes = extract_image_bytes_from_item(image_item, session)
|
||
output_format = payload.get("output_format") if isinstance(payload.get("output_format"), str) else None
|
||
|
||
if event_kind == "partial":
|
||
partial_index = payload.get("partial_image_index")
|
||
suffix = resolve_output_extension(Path(stem), output_format)
|
||
if isinstance(partial_index, int):
|
||
filename = f"{stem}-partial-{partial_index:02d}-event-{record['index']:03d}{suffix}"
|
||
else:
|
||
filename = f"{stem}-partial-{len(partial_paths):02d}-event-{record['index']:03d}{suffix}"
|
||
partial_path = partial_dir / filename
|
||
save_image_bytes(image_bytes, partial_path)
|
||
partial_paths.append(str(partial_path))
|
||
continue
|
||
|
||
if event_kind == "completed":
|
||
suffix = resolve_output_extension(Path(stem), output_format)
|
||
output_path = run_dir / f"{stem}-{len(image_paths) + 1}{suffix}"
|
||
save_image_bytes(image_bytes, output_path, min_output_edge=min_output_edge)
|
||
image_paths.append(str(output_path))
|
||
|
||
if record_count == 0:
|
||
raise EmptyStreamResponseError("流式响应为空，未收到任何 SSE 事件")
|
||
if not image_paths:
|
||
raise RuntimeError("流式响应结束，但没有收到 completed 图片事件")
|
||
|
||
return image_paths, partial_paths
|
||
|
||
|
||
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")
|
||
|
||
reference_sources = collect_reference_sources()
|
||
payload, model, requested_size, request_mode, mask_source, background, input_fidelity = build_payload(
|
||
base_url,
|
||
reference_sources,
|
||
)
|
||
run_dir = make_run_dir()
|
||
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)
|
||
stream_requested = bool(payload.get("stream"))
|
||
endpoint_path = "/v1/images/edits" if request_mode == "edit" else "/v1/images/generations"
|
||
min_output_edge = resolve_min_output_edge(requested_size, payload.get("resolution"), request_style)
|
||
prompt = payload["prompt"]
|
||
stem = f"{resolve_output_stem()}-{datetime.now().strftime('%Y%m%d-%H%M%S-%f')}"
|
||
|
||
print(
|
||
"开始生成图片:",
|
||
f"mode={request_mode}",
|
||
f"request_style={request_style}",
|
||
f"model={model}",
|
||
f"refs={len(reference_sources)}",
|
||
f"n={payload['n']}",
|
||
f"size={requested_size}",
|
||
f"resolution={payload.get('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"stream={stream_requested}",
|
||
f"partial_images={payload.get('partial_images') if stream_requested else 'none'}",
|
||
f"moderation={payload.get('moderation') or 'default'}",
|
||
f"endpoint={base_url}{endpoint_path}",
|
||
f"responses={run_dir}",
|
||
)
|
||
|
||
dump_json(
|
||
run_dir / "request.json",
|
||
{
|
||
"base_url": base_url,
|
||
"endpoint_path": endpoint_path,
|
||
"request_mode": request_mode,
|
||
"request_style": request_style,
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"size": requested_size,
|
||
"resolution": payload.get("resolution"),
|
||
"n": payload["n"],
|
||
"reference_sources": reference_sources,
|
||
"mask_source": mask_source,
|
||
"background": background,
|
||
"input_fidelity": input_fidelity,
|
||
"min_output_edge": min_output_edge,
|
||
"output_dir": str(run_dir),
|
||
"request_timeout_seconds": request_timeout,
|
||
"max_poll_wait_seconds": max_poll_wait,
|
||
"initial_poll_delay_seconds": initial_poll_delay,
|
||
"poll_interval_seconds": poll_interval,
|
||
"payload": payload,
|
||
},
|
||
)
|
||
|
||
try:
|
||
if request_mode == "edit":
|
||
reference_inputs = [
|
||
image_source_to_input(source, f"reference-{index}", session)
|
||
for index, source in enumerate(reference_sources, start=1)
|
||
]
|
||
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_inputs, mask_input),
|
||
headers={"Authorization": f"Bearer {api_key}"},
|
||
stream=stream_requested,
|
||
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",
|
||
},
|
||
stream=stream_requested,
|
||
timeout=(DEFAULT_CONNECT_TIMEOUT, request_timeout),
|
||
)
|
||
|
||
if stream_requested and is_event_stream_response(resp):
|
||
try:
|
||
image_paths, partial_paths = handle_streaming_multi_response(
|
||
resp,
|
||
session,
|
||
run_dir,
|
||
stem,
|
||
min_output_edge=min_output_edge,
|
||
)
|
||
dump_json(
|
||
run_dir / "result.json",
|
||
{
|
||
"images": image_paths,
|
||
"partial_images": partial_paths,
|
||
"response_path": str(run_dir / "post-response.json"),
|
||
"event_log_path": str(run_dir / "stream-events.jsonl"),
|
||
"request_mode": request_mode,
|
||
"request_style": request_style,
|
||
"model": model,
|
||
"size": requested_size,
|
||
"n": payload["n"],
|
||
"reference_count": len(reference_sources),
|
||
"stream": True,
|
||
},
|
||
)
|
||
print(_json_dumps({"images": image_paths, "partial_images": partial_paths, "output_dir": str(run_dir)}))
|
||
return
|
||
except EmptyStreamResponseError:
|
||
fallback_payload = build_non_stream_payload(payload)
|
||
print("上游返回空 EventStream，自动回退为非流式请求重试一次")
|
||
if request_mode == "edit":
|
||
resp = session.post(
|
||
f"{base_url}{endpoint_path}",
|
||
data=payload_to_form_fields(fallback_payload),
|
||
files=build_edit_files(request_style, reference_inputs, mask_input),
|
||
headers={"Authorization": f"Bearer {api_key}"},
|
||
timeout=(DEFAULT_CONNECT_TIMEOUT, request_timeout),
|
||
)
|
||
else:
|
||
resp = session.post(
|
||
f"{base_url}{endpoint_path}",
|
||
json=fallback_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")
|
||
items = extract_image_items(response_data)
|
||
if items:
|
||
image_paths = save_output_images(items, session, run_dir, stem, min_output_edge)
|
||
dump_json(
|
||
run_dir / "result.json",
|
||
{
|
||
"images": image_paths,
|
||
"response_path": str(run_dir / "post-response.json"),
|
||
"request_mode": request_mode,
|
||
"request_style": request_style,
|
||
"model": model,
|
||
"size": requested_size,
|
||
"n": payload["n"],
|
||
"reference_count": len(reference_sources),
|
||
},
|
||
)
|
||
print(_json_dumps({"images": image_paths, "output_dir": str(run_dir)}))
|
||
return
|
||
|
||
task_id = extract_task_id(response_data)
|
||
if task_id:
|
||
items = poll_task_images(
|
||
task_id,
|
||
base_url,
|
||
{"Authorization": f"Bearer {api_key}"},
|
||
session,
|
||
run_dir,
|
||
max_poll_wait,
|
||
initial_poll_delay,
|
||
poll_interval,
|
||
)
|
||
image_paths = save_output_images(items, session, run_dir, stem, min_output_edge)
|
||
dump_json(
|
||
run_dir / "result.json",
|
||
{
|
||
"images": image_paths,
|
||
"response_path": str(run_dir / "post-response.json"),
|
||
"request_mode": request_mode,
|
||
"request_style": request_style,
|
||
"model": model,
|
||
"size": requested_size,
|
||
"n": payload["n"],
|
||
"reference_count": len(reference_sources),
|
||
},
|
||
)
|
||
print(_json_dumps({"images": image_paths, "output_dir": str(run_dir)}))
|
||
return
|
||
|
||
error_message = extract_api_error_message(data)
|
||
if error_message:
|
||
raise RuntimeError(error_message)
|
||
|
||
raise RuntimeError(f"unexpected response: {data}")
|
||
except BaseException as exc:
|
||
save_error_artifact(run_dir / "error.json", exc, {"output_dir": str(run_dir)})
|
||
raise
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|