From af0d43c4536cfe63b1010eea2f2edc28dbdf8419 Mon Sep 17 00:00:00 2001
From: toor <jack0210_@hotmail.com>
Date: Tue, 18 Aug 2026 23:56:45 +0800
Subject: [PATCH] Support URL image responses

Download upstream image URLs in addition to Base64 responses.
---
 agent-image-t2i.py | 37 ++++++++++++++++++++++++++++++++-----
 1 file changed, 32 insertions(+), 5 deletions(-)

diff --git a/agent-image-t2i.py b/agent-image-t2i.py
index 53cbce6..0272fa7 100644
--- a/agent-image-t2i.py
+++ b/agent-image-t2i.py
@@ -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)