972 lines
34 KiB
Python
972 lines
34 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_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",
|
||
"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
|
||
|
||
|
||
class EmptyStreamResponseError(RuntimeError):
|
||
pass
|
||
|
||
|
||
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 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")
|
||
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_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__,
|
||
"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 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(),
|
||
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 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 参考图格式无效")
|
||
|
||
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 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,
|
||
"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"
|
||
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")
|
||
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
|
||
if isinstance(value, bool):
|
||
fields[key] = str(value).lower()
|
||
else:
|
||
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 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,
|
||
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 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()
|
||
|
||
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)
|
||
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")
|
||
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"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}",
|
||
)
|
||
|
||
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}"},
|
||
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:
|
||
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")
|
||
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
|
||
|
||
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)})
|
||
raise
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|