3218485270
f493bde8a9
- 缓存热点批量流程（有采集缓存不触发 Google） - 简报不足直接从采集缓存生成（轻量补齐） - 三图合成（模特/印花/底图）+ 底图压缩 <2MB - 热点去重→风格去重自动切换 + 不适合类目 review 兜底 - 透明背景（background=transparent）+ 提示词清洗（敏感词/背景描述） - 任务前 basemap 校验 + 模板国家校验 + 模特任务级分配
243 lines
12 KiB
Python
243 lines
12 KiB
Python
"""OpenAI 图像后端（生成设计稿 + 三图合成）。
|
||
|
||
支持 OpenAI images/generations（文生图）与 images/edits（img2img 多参考图），
|
||
兼容两类网关返回：
|
||
1) 同步：{"data": [{"b64_json" | "url"}]}
|
||
2) 异步任务（ai-media.vip 等重负载自动转异步）：
|
||
提交返回 202 + {"object":"image.task","task_id","poll_url","poll_after_ms"}，
|
||
需轮询 GET {base}/images/tasks/{task_id} 直到 succeeded，再从成功响应取图。
|
||
|
||
所有请求跳过环境代理（NO_PROXY），适配用户挂 VPN 时直连国内/自建网关。
|
||
"""
|
||
import base64
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
|
||
import requests
|
||
import os
|
||
|
||
# 模型调用一律直连：用户常开 VPN（系统代理），网关多为国内/自建，走代理会被拦截或变慢。
|
||
# 环境变量级 NO_PROXY 双保险（requests/urllib3 均读取），Google 采集（pytrends）不受影响。
|
||
os.environ.setdefault('NO_PROXY', '*')
|
||
os.environ.setdefault('no_proxy', '*')
|
||
|
||
from .base import ImageBackend
|
||
|
||
# 直连策略：忽略环境代理（用户挂 VPN 时代理会拦截国内/自建网关的请求）
|
||
NO_PROXY = {"http": None, "https": None}
|
||
|
||
_FINAL_STATUS = ("succeeded", "completed", "done")
|
||
_FAIL_STATUS = ("failed", "error")
|
||
|
||
|
||
def _shrink_blob_to_2mb(img_path: str, blob: bytes, max_bytes: int = 2 * 1024 * 1024) -> bytes:
|
||
"""把大图压缩到 <max_bytes（默认 2MB）再上传：
|
||
- 尺寸过大先缩放（合成输入 1600x2200 内足够）；
|
||
- 有透明通道 → PNG 优化；否则 → JPEG 循环降质。
|
||
失败时返回原 blob（不阻塞流程）。"""
|
||
try:
|
||
import io
|
||
from PIL import Image
|
||
img = Image.open(io.BytesIO(blob))
|
||
max_w, max_h = 1600, 2200
|
||
if img.width > max_w or img.height > max_h:
|
||
img.thumbnail((max_w, max_h))
|
||
has_alpha = img.mode in ("RGBA", "LA")
|
||
if not has_alpha:
|
||
img = img.convert("RGB")
|
||
for q in (88, 70, 50, 35):
|
||
buf = io.BytesIO()
|
||
if has_alpha:
|
||
img.save(buf, format="PNG", optimize=True)
|
||
else:
|
||
img.save(buf, format="JPEG", quality=q)
|
||
if buf.tell() <= max_bytes:
|
||
return buf.getvalue()
|
||
# 保底：最小质量 PNG/JPEG
|
||
buf = io.BytesIO()
|
||
if has_alpha:
|
||
img.save(buf, format="PNG", optimize=True)
|
||
else:
|
||
img.save(buf, format="JPEG", quality=30)
|
||
return buf.getvalue()
|
||
except Exception as e: # noqa: BLE001
|
||
print(f"[img] 图片压缩失败（用原图）: {e}")
|
||
return blob
|
||
|
||
|
||
def _save_from_response(j: dict, out_path: str) -> str:
|
||
"""从同步/异步最终响应提取图片（data[].b64_json 或 url）并保存。"""
|
||
data_item = (j.get("data") or [{}])[0]
|
||
b64 = data_item.get("b64_json")
|
||
if b64:
|
||
Path(out_path).parent.mkdir(parents=True, exist_ok=True)
|
||
Path(out_path).write_bytes(base64.b64decode(b64))
|
||
return out_path
|
||
url = data_item.get("url")
|
||
if not url:
|
||
raise RuntimeError(f"图像 API 返回无 b64_json/url: {str(j)[:200]}")
|
||
img_resp = requests.get(url, proxies=NO_PROXY, timeout=120)
|
||
img_resp.raise_for_status()
|
||
Path(out_path).parent.mkdir(parents=True, exist_ok=True)
|
||
Path(out_path).write_bytes(img_resp.content)
|
||
return out_path
|
||
|
||
|
||
def _wait_task(base_url: str, headers: dict, task_id: str, poll_after_ms: int,
|
||
timeout: int = 60) -> dict:
|
||
"""轮询异步图像任务直到完成；uncertain（结果暂时不确定）继续等，不重复提交。
|
||
timeout 默认 60s：异步路径不稳定（ai-media 网关常 uncertain/task not found），
|
||
超时即抛异常由调用方重试同步提交。"""
|
||
deadline = time.time() + timeout
|
||
last = ""
|
||
while time.time() < deadline:
|
||
time.sleep(max(int(poll_after_ms or 2000), 2000) / 1000.0)
|
||
try:
|
||
resp = requests.get(f"{base_url}/images/tasks/{task_id}", headers=headers,
|
||
timeout=60, proxies=NO_PROXY)
|
||
if resp.status_code >= 400:
|
||
continue
|
||
j = resp.json()
|
||
except Exception as e: # noqa: BLE001
|
||
print(f"[img] 任务轮询异常（重试）: {e}")
|
||
continue
|
||
st = j.get("status")
|
||
if st != last:
|
||
print(f"[img] 异步任务 {task_id[:20]}… status={st}")
|
||
last = st
|
||
if st in _FINAL_STATUS:
|
||
return j
|
||
if st in _FAIL_STATUS:
|
||
raise RuntimeError(f"图像异步任务失败: {j.get('error')} {j.get('error_code')}")
|
||
# queued / running / uncertain → 继续等
|
||
raise RuntimeError(f"图像异步任务超时（{timeout}s）")
|
||
|
||
|
||
def _resolve_task_or_sync(j: dict, base_url: str, headers: dict, out_path: str) -> str:
|
||
"""提交后的统一处理：异步任务则轮询，然后提取图片保存。"""
|
||
if j.get("object") == "image.task" or j.get("task_id") or j.get("id", "").startswith("imgtask"):
|
||
tid = j.get("task_id") or j.get("id") or ""
|
||
if not tid:
|
||
raise RuntimeError(f"异步任务无 task_id: {str(j)[:200]}")
|
||
j = _wait_task(base_url, headers, tid, int(j.get("poll_after_ms") or 2000))
|
||
return _save_from_response(j, out_path)
|
||
|
||
|
||
class OpenAIImageBackend(ImageBackend):
|
||
name = "openai"
|
||
|
||
def __init__(self):
|
||
self._cfg: dict = {}
|
||
|
||
def bind_config(self, cfg: dict):
|
||
self._cfg = cfg or {}
|
||
|
||
def _base_url(self) -> str:
|
||
return str(self._cfg.get("base_url") or "https://api.openai.com/v1").rstrip("/")
|
||
|
||
def print(self, prompt: str, base_image: str, out_path: str, negative: str = "",
|
||
extra_images=None, size: str = "") -> str:
|
||
"""img2img 编辑：参考图 base_image（+ 可选 extra_images 多参考图）按 prompt 生成新图。
|
||
|
||
- 平铺服装图：base_image=平铺衣服底图(图3)，extra_images=[印花设计稿(图2)]
|
||
- 三图模特合成：base_image=模特图(图1)，extra_images=[印花设计稿(图2), 平铺底图(图3)]
|
||
提交顺序即图1→图2→图3，与提示词中的图片角色一一对应。
|
||
size: 显式尺寸覆盖（如 "1504x2000"）；留空用配置 size（默认 1024x1024）。
|
||
"""
|
||
cfg = self._cfg
|
||
api_key = cfg.get("api_key", "")
|
||
model = cfg.get("model", "gpt-image-1")
|
||
if not api_key:
|
||
raise RuntimeError("compose.api_key 未配置，无法调用 OpenAI 图像后端。")
|
||
|
||
full_prompt = prompt + (f"\nNegative: {negative}" if negative else "")
|
||
base_url = self._base_url()
|
||
headers = {"Authorization": f"Bearer {api_key}"}
|
||
# 必须把文件读成 bytes 再提交（句柄随 with 关闭会导致 "read of closed file"）；
|
||
# 输入图（模特/设计/底图）先压缩校验：>2MB 压缩到 <2MB 再上传，避免大图拖垮网关导致超时
|
||
file_payloads = []
|
||
for img_path in [base_image] + list(extra_images or []):
|
||
blob = Path(img_path).read_bytes()
|
||
if len(blob) > 2 * 1024 * 1024:
|
||
shrunk = _shrink_blob_to_2mb(img_path, blob)
|
||
print(f"[img] 输入图压缩: {Path(img_path).name} {len(blob)//1024}KB → {len(shrunk)//1024}KB（<2MB）")
|
||
blob = shrunk
|
||
file_payloads.append((Path(img_path).name, blob))
|
||
files = [("image", (name, blob, "image/png")) for name, blob in file_payloads]
|
||
data = {
|
||
"prompt": full_prompt,
|
||
"n": 1,
|
||
"size": size or cfg.get("size", "1024x1024"),
|
||
"model": model,
|
||
"execution_mode": "sync", # 强制同步（ai-media 等网关默认重负载转异步任务，不稳定）
|
||
}
|
||
# 提交重试：异步路径不稳定 → 失败重试同步提交（最多 3 次）；
|
||
# 内容政策拦截（content_policy_violation）多为网关误判 → 等待后重试
|
||
last_err: Optional[str] = None
|
||
for attempt in range(3):
|
||
resp = requests.post(f"{base_url}/images/edits", headers=headers, files=files,
|
||
data=data, timeout=300, proxies=NO_PROXY)
|
||
if resp.status_code >= 400:
|
||
body = resp.text or ""
|
||
if "content_policy" in body and attempt < 2:
|
||
print(f"[img] 内容政策拦截（可能误判），等待后重试 {attempt + 2}/3")
|
||
time.sleep(2)
|
||
continue
|
||
if attempt == 0:
|
||
data.pop("execution_mode", None) # 网关不认识该参数（官方 OpenAI）→ 去掉重试
|
||
continue
|
||
raise RuntimeError(f"图像 API {resp.status_code}: {body[:300]}")
|
||
try:
|
||
return _resolve_task_or_sync(resp.json(), base_url, headers, out_path)
|
||
except Exception as e: # noqa: BLE001
|
||
last_err = str(e)
|
||
print(f"[img] 第 {attempt + 1} 次提交异步失败，重试同步提交: {e}")
|
||
raise RuntimeError(f"图像合成多次提交均失败: {last_err}")
|
||
|
||
def generate(self, prompt: str, out_path: str, negative: str = "", size: str = "") -> str:
|
||
"""纯文生图：生成白底纯印花设计稿（standalone pure print design）。
|
||
size: 显式尺寸覆盖（印花设计统一 1024x1024）；留空用配置 size。
|
||
background: 配置 compose.background="transparent" 时传 background 参数 → 透明背景 PNG
|
||
（gpt-image-1/2 等模型支持；网关不支持该参数时会被忽略或由网关兜底）。"""
|
||
cfg = self._cfg
|
||
api_key = cfg.get("api_key", "")
|
||
model = cfg.get("model", "gpt-image-1")
|
||
if not api_key:
|
||
raise RuntimeError("compose.api_key 未配置，无法调用 OpenAI 图像后端。")
|
||
|
||
full_prompt = prompt + (f"\nNegative: {negative}" if negative else "")
|
||
base_url = self._base_url()
|
||
headers = {"Authorization": f"Bearer {api_key}"}
|
||
data = {
|
||
"prompt": full_prompt,
|
||
"n": 1,
|
||
"size": size or cfg.get("size", "1024x1024"),
|
||
"model": model,
|
||
"response_format": "b64_json",
|
||
"execution_mode": "sync", # 强制同步（ai-media 等网关默认重负载转异步任务，不稳定）
|
||
}
|
||
bg = str(cfg.get("background") or "").strip()
|
||
if bg:
|
||
data["background"] = bg # 如 "transparent"（透明背景 PNG）
|
||
last_err: Optional[str] = None
|
||
for attempt in range(3):
|
||
resp = requests.post(f"{base_url}/images/generations", headers=headers, json=data,
|
||
timeout=300, proxies=NO_PROXY)
|
||
if resp.status_code >= 400:
|
||
body = resp.text or ""
|
||
if "content_policy" in body and attempt < 2:
|
||
print(f"[img] 内容政策拦截（可能误判），等待后重试 {attempt + 2}/3")
|
||
time.sleep(2)
|
||
continue
|
||
if attempt == 0:
|
||
data.pop("execution_mode", None) # 网关不认识该参数（官方 OpenAI）→ 去掉重试
|
||
continue
|
||
raise RuntimeError(f"图像 API {resp.status_code}: {body[:300]}")
|
||
try:
|
||
return _resolve_task_or_sync(resp.json(), base_url, headers, out_path)
|
||
except Exception as e: # noqa: BLE001
|
||
last_err = str(e)
|
||
print(f"[img] 第 {attempt + 1} 次生成异步失败，重试同步提交: {e}")
|
||
raise RuntimeError(f"图像生成多次提交均失败: {last_err}")
|