|
|
|
|
@ -4,6 +4,8 @@ |
|
|
|
|
import base64 |
|
|
|
|
import os |
|
|
|
|
import time |
|
|
|
|
import urllib.error |
|
|
|
|
import urllib.request |
|
|
|
|
from pathlib import Path |
|
|
|
|
from uuid import uuid4 |
|
|
|
|
|
|
|
|
|
@ -56,6 +58,7 @@ class ImageGenerationResponse(BaseModel): |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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) |
|
|
|
|
@ -69,6 +72,28 @@ def generate_with_retry(client: OpenAI, *, model: str, prompt: str, size: str, n |
|
|
|
|
time.sleep(2 ** (attempt - 1)) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def download_image(url: str) -> bytes: |
|
|
|
|
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)) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@app.get("/health") |
|
|
|
|
def health_check(): |
|
|
|
|
return {"status": "ok"} |
|
|
|
|
@ -93,13 +118,15 @@ def generate_images(request: ImageGenerationRequest): |
|
|
|
|
|
|
|
|
|
images = [] |
|
|
|
|
for item in response.data: |
|
|
|
|
if not item.b64_json: |
|
|
|
|
raise HTTPException(status_code=502, detail="上游图片服务未返回 Base64 图片数据") |
|
|
|
|
|
|
|
|
|
try: |
|
|
|
|
raw = base64.b64decode(item.b64_json) |
|
|
|
|
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="上游图片数据格式无效") from e |
|
|
|
|
raise HTTPException(status_code=502, detail="上游 Base64 图片数据格式无效") from e |
|
|
|
|
|
|
|
|
|
filename = f"{model}-{uuid4().hex}.png" |
|
|
|
|
(OUTPUT_DIR / filename).write_bytes(raw) |
|
|
|
|
|