You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
98 lines
4.1 KiB
98 lines
4.1 KiB
"""GPT Image 和 Grok Image 文生图的核心逻辑。"""
|
|
|
|
import base64
|
|
import os
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
from fastapi import HTTPException
|
|
from openai import APIConnectionError, APIStatusError, APITimeoutError, OpenAI
|
|
|
|
API_BASE_URL = os.getenv("IMAGE_API_BASE_URL", "https://api.slomerex.xyz/v1")
|
|
OUTPUT_DIR = Path("output")
|
|
MAX_RETRIES = 3
|
|
RETRYABLE_STATUS_CODES = {408, 429, 500, 502, 503, 504}
|
|
|
|
PROVIDERS = {
|
|
"grok": {
|
|
"api_key_env": "GROK_API_KEY",
|
|
"model": "grok-imagine-image-lite",
|
|
},
|
|
"gpt": {
|
|
"api_key_env": "GPT_API_KEY",
|
|
"model": "gpt-image-2-1K",
|
|
},
|
|
}
|
|
|
|
|
|
def _generate_with_retry(client: OpenAI, *, model: str, prompt: str, size: str, n: int):
|
|
"""调用图片接口，并对可恢复的上游错误进行指数退避重试。"""
|
|
for attempt in range(1, MAX_RETRIES + 1):
|
|
try:
|
|
return client.images.generate(model=model, prompt=prompt, size=size, n=n)
|
|
except APIStatusError as e:
|
|
if e.status_code not in RETRYABLE_STATUS_CODES or attempt == MAX_RETRIES:
|
|
raise HTTPException(status_code=502, detail=f"上游图片服务错误: HTTP {e.status_code}") from e
|
|
time.sleep(2 ** (attempt - 1))
|
|
except (APIConnectionError, APITimeoutError) as e:
|
|
if attempt == MAX_RETRIES:
|
|
raise HTTPException(status_code=504, detail="上游图片服务连接或请求超时") from e
|
|
time.sleep(2 ** (attempt - 1))
|
|
|
|
|
|
def _download_image(url: str) -> bytes:
|
|
"""从 URL 或 data URI 下载图片。"""
|
|
if url.startswith("data:image/"):
|
|
try:
|
|
_, encoded = url.split(",", 1)
|
|
return base64.b64decode(encoded, validate=True)
|
|
except (ValueError, TypeError) as e:
|
|
raise HTTPException(status_code=502, detail="图片 data URL 格式无效") from e
|
|
|
|
request = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
|
|
for attempt in range(1, MAX_RETRIES + 1):
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=120) as response:
|
|
return response.read()
|
|
except urllib.error.HTTPError as e:
|
|
if e.code not in RETRYABLE_STATUS_CODES or attempt == MAX_RETRIES:
|
|
raise HTTPException(status_code=502, detail=f"图片 URL 下载失败: HTTP {e.code}") from e
|
|
except (urllib.error.URLError, TimeoutError) as e:
|
|
if attempt == MAX_RETRIES:
|
|
raise HTTPException(status_code=504, detail="图片 URL 下载超时或网络连接失败") from e
|
|
time.sleep(2 ** (attempt - 1))
|
|
|
|
|
|
def generate_images(*, provider: str, model: str | None, prompt: str, size: str, n: int, profile_name: str):
|
|
"""生成图片并保存到本地，返回结果字典。"""
|
|
config = PROVIDERS[provider]
|
|
api_key = os.getenv(config["api_key_env"])
|
|
if not api_key:
|
|
raise HTTPException(status_code=500, detail=f"未配置环境变量 {config['api_key_env']}")
|
|
|
|
resolved_model = model or config["model"]
|
|
profile_dir = OUTPUT_DIR / profile_name
|
|
profile_dir.mkdir(parents=True, exist_ok=True)
|
|
client = OpenAI(base_url=API_BASE_URL, api_key=api_key, timeout=120.0, max_retries=0)
|
|
response = _generate_with_retry(client, model=resolved_model, prompt=prompt, size=size, n=n)
|
|
|
|
images = []
|
|
for item in response.data:
|
|
try:
|
|
if item.b64_json:
|
|
raw = base64.b64decode(item.b64_json, validate=True)
|
|
elif item.url:
|
|
raw = _download_image(item.url)
|
|
else:
|
|
raise HTTPException(status_code=502, detail="上游图片服务未返回 b64_json 或 url")
|
|
except ValueError as e:
|
|
raise HTTPException(status_code=502, detail="上游 Base64 图片数据格式无效") from e
|
|
|
|
filename = f"{resolved_model}-{uuid4().hex}.png"
|
|
(profile_dir / filename).write_bytes(raw)
|
|
images.append({"filename": filename, "url": f"/output/{profile_name}/{filename}", "bytes": len(raw)})
|
|
|
|
return {"provider": provider, "model": resolved_model, "images": images}
|
|
|