bot/image-generate
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(image): support stream and moderation options

Browse Source
This commit is contained in:
shay7sev
2026-05-25 18:42:47 +08:00
parent ca2fa1b03d
commit 7b9b96e332
9 changed files with 609 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

151
gpt_image2_generate_multi.py
View File

@@ -22,24 +22,41 @@ from gpt_image2_generate import (
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,
)
@@ -151,11 +168,21 @@ def build_payload(base_url: str, reference_sources: list[str]) -> tuple[dict[str
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 = os.getenv("IMAGE_OUTPUT_FORMAT", "").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,
@@ -174,6 +201,14 @@ def build_payload(base_url: str, reference_sources: list[str]) -> tuple[dict[str
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:
@@ -309,6 +344,61 @@ def save_output_images(
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()
@@ -331,6 +421,7 @@ def main() -> None:
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"]
@@ -349,6 +440,9 @@ def main() -> None:
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}",
)
@@ -391,6 +485,7 @@ def main() -> None:
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:
@@ -401,9 +496,59 @@ def main() -> None:
"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")
@@ -459,6 +604,10 @@ def main() -> None:
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)})