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()