fix(image): support stream and moderation options

gpt_image2_generate.py

@@ -32,6 +32,10 @@ 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_MODERATION_LEVELS = {"auto", "low"}
+SUPPORTED_OUTPUT_FORMATS = {"png", "jpeg", "webp"}
+STREAM_COMPLETED_EVENT_TYPES = {"image_generation.completed", "image_edit.completed"}
+STREAM_PARTIAL_EVENT_TYPES = {"image_generation.partial_image", "image_edit.partial_image"}
 SUPPORTED_ASPECT_RATIOS = {
     "1:1",
     "16:9",
@@ -71,6 +75,10 @@ class ImageInput:
     data: bytes
 
 
+class EmptyStreamResponseError(RuntimeError):
+    pass
+
+
 def load_env_file(path: str = ".env") -> None:
     env_path = Path(path)
     if not env_path.exists():
@@ -117,6 +125,28 @@ def env_int(name: str, default: int) -> int:
         raise RuntimeError(f"{name} 不是有效整数: {raw_value}") from exc
 
 
+def env_optional_int(name: str) -> int | None:
+    raw_value = os.getenv(name, "").strip()
+    if not raw_value:
+        return None
+
+    try:
+        return int(raw_value)
+    except ValueError as exc:
+        raise RuntimeError(f"{name} 不是有效整数: {raw_value}") from exc
+
+
+def env_bool(name: str, default: bool = False) -> bool:
+    raw_value = os.getenv(name, "").strip().lower()
+    if not raw_value:
+        return default
+    if raw_value in {"1", "true", "yes", "y", "on"}:
+        return True
+    if raw_value in {"0", "false", "no", "n", "off"}:
+        return False
+    raise RuntimeError(f"{name} 不是有效布尔值: {raw_value}")
+
+
 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")
@@ -145,6 +175,16 @@ def save_http_response(path: Path, response: requests.Response) -> dict[str, Any
     return record
 
 
+def save_stream_response_metadata(path: Path, response: requests.Response) -> dict[str, Any]:
+    record: dict[str, Any] = {
+        "url": response.url,
+        "status_code": response.status_code,
+        "headers": dict(response.headers),
+    }
+    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__,
@@ -196,6 +236,11 @@ def resolve_api_size(size: str, request_style: str) -> str:
     return size
 
 
+def is_gpt_image_model(model: str) -> bool:
+    normalized = model.strip().lower()
+    return normalized.startswith("gpt-image") or normalized == "chatgpt-image-latest"
+
+
 def resolve_image_source(prefix: str) -> str | None:
     candidates = {
         prefix: os.getenv(prefix, "").strip(),
@@ -249,6 +294,16 @@ def safe_filename(filename: str, fallback: str, mime_type: str) -> str:
     return candidate
 
 
+def resolve_output_extension(path: Path, output_format: str | None = None) -> str:
+    if path.suffix:
+        return path.suffix
+    if output_format == "jpeg":
+        return ".jpg"
+    if output_format in SUPPORTED_OUTPUT_FORMATS:
+        return f".{output_format}"
+    return ".png"
+
+
 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 参考图格式无效")
@@ -307,17 +362,56 @@ def image_source_to_input(source: str, fallback: str, session: requests.Session)
     return local_file_to_image_input(value, fallback)
 
 
+def resolve_optional_enum(name: str, allowed: set[str]) -> str | None:
+    raw_value = os.getenv(name, "").strip().lower()
+    if not raw_value:
+        return None
+    if raw_value not in allowed:
+        allowed_text = ", ".join(sorted(allowed))
+        raise RuntimeError(f"{name} 仅支持: {allowed_text}；当前值: {raw_value}")
+    return raw_value
+
+
+def resolve_stream_options(model: str) -> tuple[bool, int | None]:
+    stream_requested = env_bool("IMAGE_STREAM", False)
+    partial_images = env_optional_int("IMAGE_PARTIAL_IMAGES")
+
+    if partial_images is not None:
+        if partial_images < 0 or partial_images > 3:
+            raise RuntimeError("IMAGE_PARTIAL_IMAGES 仅支持 0-3")
+        stream_requested = True
+
+    if stream_requested and not is_gpt_image_model(model):
+        raise RuntimeError(f"当前 model 不支持流式图片输出: {model}")
+
+    return stream_requested, partial_images
+
+
 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()
+    background = os.getenv("IMAGE_BACKGROUND", "").strip().lower() or None
+    input_fidelity = os.getenv("IMAGE_INPUT_FIDELITY", "").strip().lower() 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")
     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)
+    stream_requested, partial_images = resolve_stream_options(model)
+
+    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,
@@ -331,6 +425,20 @@ def build_payload(base_url: str) -> tuple[dict[str, Any], str, str, str, str, st
 
     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
 
     reference_source = resolve_image_source("REFERENCE_IMAGE")
     mask_source = resolve_image_source("MASK_IMAGE")
@@ -346,7 +454,10 @@ def payload_to_form_fields(payload: dict[str, Any]) -> dict[str, str]:
     for key, value in payload.items():
         if value is None:
             continue
-        fields[key] = str(value)
+        if isinstance(value, bool):
+            fields[key] = str(value).lower()
+        else:
+            fields[key] = str(value)
     return fields
 
 
@@ -425,6 +536,148 @@ def extract_image_bytes_from_item(item: dict[str, Any], session: requests.Sessio
     raise RuntimeError(f"响应里没有可用的图片数据: {item}")
 
 
+def classify_stream_event(record: dict[str, Any], payload: dict[str, Any]) -> str | None:
+    candidates = [
+        str(payload.get("type") or ""),
+        str(record.get("event") or ""),
+        str(payload.get("event") or ""),
+        str(payload.get("status") or ""),
+    ]
+    normalized = [candidate.strip().lower() for candidate in candidates if candidate]
+
+    for candidate in normalized:
+        if candidate in STREAM_PARTIAL_EVENT_TYPES or candidate.endswith("partial_image"):
+            return "partial"
+        if candidate in STREAM_COMPLETED_EVENT_TYPES or candidate.endswith("completed"):
+            return "completed"
+    return None
+
+
+def extract_stream_image_item(payload: dict[str, Any]) -> dict[str, Any] | None:
+    if payload.get("b64_json") or payload.get("url"):
+        return payload
+
+    data = payload.get("data")
+    if isinstance(data, dict):
+        if data.get("b64_json") or data.get("url"):
+            return data
+
+    images = payload.get("images")
+    if isinstance(images, list):
+        for item in images:
+            if isinstance(item, dict) and (item.get("b64_json") or item.get("url")):
+                return item
+
+    result = payload.get("result")
+    if isinstance(result, dict):
+        result_images = result.get("images")
+        if isinstance(result_images, list):
+            for item in result_images:
+                if isinstance(item, dict) and (item.get("b64_json") or item.get("url")):
+                    return item
+
+    if isinstance(data, dict):
+        nested_images = data.get("images")
+        if isinstance(nested_images, list):
+            for item in nested_images:
+                if isinstance(item, dict) and (item.get("b64_json") or item.get("url")):
+                    return item
+
+    return None
+
+
+def extract_api_error_message(data: Any) -> str | None:
+    if not isinstance(data, dict):
+        return None
+
+    error = data.get("error")
+    if isinstance(error, dict):
+        message = error.get("message")
+        if isinstance(message, str) and message.strip():
+            return message.strip()
+        return str(error)
+    if error:
+        return str(error)
+
+    nested = data.get("data")
+    if isinstance(nested, dict):
+        nested_error = nested.get("error")
+        if isinstance(nested_error, dict):
+            message = nested_error.get("message")
+            if isinstance(message, str) and message.strip():
+                return message.strip()
+            return str(nested_error)
+        if nested_error:
+            return str(nested_error)
+
+    return None
+
+
+def is_event_stream_response(response: requests.Response) -> bool:
+    content_type = response.headers.get("Content-Type", "")
+    return "text/event-stream" in content_type.lower()
+
+
+def iter_sse_events(response: requests.Response):
+    event_name = ""
+    data_lines: list[str] = []
+
+    def flush() -> tuple[str, str] | None:
+        nonlocal event_name, data_lines
+        if event_name or data_lines:
+            event = (event_name or "message", "\n".join(data_lines))
+        else:
+            event = None
+        event_name = ""
+        data_lines = []
+        return event
+
+    for raw_line in response.iter_lines(decode_unicode=True):
+        if raw_line is None:
+            continue
+
+        line = raw_line if isinstance(raw_line, str) else raw_line.decode("utf-8", errors="replace")
+        if line == "":
+            event = flush()
+            if event is not None:
+                yield event
+            continue
+        if line.startswith(":"):
+            continue
+
+        field, _, value = line.partition(":")
+        if value.startswith(" "):
+            value = value[1:]
+
+        if field == "event":
+            event_name = value
+        elif field == "data":
+            data_lines.append(value)
+
+    event = flush()
+    if event is not None:
+        yield event
+
+
+def iter_stream_event_records(path: Path, response: requests.Response):
+    with path.open("w", encoding="utf-8") as handle:
+        for index, (event_name, data_text) in enumerate(iter_sse_events(response), start=1):
+            payload: Any
+            try:
+                payload = json.loads(data_text)
+            except json.JSONDecodeError:
+                payload = data_text
+
+            record = {
+                "index": index,
+                "event": event_name,
+                "data": payload,
+            }
+            handle.write(json.dumps(record, ensure_ascii=False) + "\n")
+            handle.flush()
+            yield record
+
+
 def poll_task(
     task_id: str,
     base_url: str,
@@ -480,6 +733,75 @@ def poll_task(
         time.sleep(poll_interval)
 
 
+def handle_streaming_single_response(
+    response: requests.Response,
+    session: requests.Session,
+    run_dir: Path,
+    output_path: Path,
+    min_output_edge: int | None,
+) -> None:
+    save_stream_response_metadata(run_dir / "post-response.json", response)
+    response.raise_for_status()
+    partial_dir = run_dir / "partials"
+    partial_dir.mkdir(parents=True, exist_ok=True)
+    partial_paths: list[str] = []
+    final_event: dict[str, Any] | None = None
+    record_count = 0
+
+    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(output_path, output_format)
+            if isinstance(partial_index, int):
+                filename = f"{output_path.stem}-partial-{partial_index:02d}-event-{record['index']:03d}{suffix}"
+            else:
+                filename = f"{output_path.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":
+            save_image_bytes(image_bytes, output_path, min_output_edge=min_output_edge)
+            final_event = payload
+
+    if record_count == 0:
+        raise EmptyStreamResponseError("流式响应为空，未收到任何 SSE 事件")
+    if final_event is None:
+        raise RuntimeError("流式响应结束，但没有收到最终 completed 图片事件")
+
+    dump_json(
+        run_dir / "result.json",
+        {
+            "image": str(output_path),
+            "partial_images": partial_paths,
+            "response_path": str(run_dir / "post-response.json"),
+            "event_log_path": str(run_dir / "stream-events.jsonl"),
+            "stream": True,
+        },
+    )
+
+
+def build_non_stream_payload(payload: dict[str, Any]) -> dict[str, Any]:
+    fallback_payload = dict(payload)
+    fallback_payload.pop("stream", None)
+    fallback_payload.pop("partial_images", None)
+    return fallback_payload
+
+
 def main() -> None:
     load_env_file()
 
@@ -500,6 +822,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(size, resolution, request_style)
     reference_image_info = describe_image_source(reference_source, "reference-image")
@@ -516,6 +839,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"output={output_path}",
         f"responses={run_dir}",
@@ -559,6 +885,7 @@ def main() -> None:
                 data=payload_to_form_fields(payload),
                 files=build_edit_files(request_style, reference_input, mask_input),
                 headers={"Authorization": f"Bearer {api_key}"},
+                stream=stream_requested,
                 timeout=(DEFAULT_CONNECT_TIMEOUT, request_timeout),
             )
         else:
@@ -569,8 +896,41 @@ 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:
+                handle_streaming_single_response(
+                    resp,
+                    session,
+                    run_dir,
+                    output_path,
+                    min_output_edge=min_output_edge,
+                )
+                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_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=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")
@@ -597,6 +957,10 @@ def main() -> None:
             save_image_bytes(image_bytes, output_path, min_output_edge=min_output_edge)
             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_file": str(output_path)})