Support URL image responses

Download upstream image URLs in addition to Base64 responses.
main
toor 1 month ago committed by Jack
parent 9bf2a97d00
commit af0d43c453
  1. 37
      agent-image-t2i.py

@ -4,6 +4,8 @@
import base64 import base64
import os import os
import time import time
import urllib.error
import urllib.request
from pathlib import Path from pathlib import Path
from uuid import uuid4 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): def generate_with_retry(client: OpenAI, *, model: str, prompt: str, size: str, n: int):
"""调用图片接口,并对可恢复的上游错误进行指数退避重试。"""
for attempt in range(1, MAX_RETRIES + 1): for attempt in range(1, MAX_RETRIES + 1):
try: try:
return client.images.generate(model=model, prompt=prompt, size=size, n=n) 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)) 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") @app.get("/health")
def health_check(): def health_check():
return {"status": "ok"} return {"status": "ok"}
@ -93,13 +118,15 @@ def generate_images(request: ImageGenerationRequest):
images = [] images = []
for item in response.data: for item in response.data:
if not item.b64_json:
raise HTTPException(status_code=502, detail="上游图片服务未返回 Base64 图片数据")
try: 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: 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" filename = f"{model}-{uuid4().hex}.png"
(OUTPUT_DIR / filename).write_bytes(raw) (OUTPUT_DIR / filename).write_bytes(raw)

Loading…
Cancel
Save