Support URL image responses
Download upstream image URLs in addition to Base64 responses.
This commit is contained in:
+32
-5
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user