876 lines
41 KiB
Python
876 lines
41 KiB
Python
"""套图生成服务：图像 provider（豆包 Seedream / 通义万相 / RightAPI 中转）+ 任务执行器。
|
||
|
||
平移自 image-suite-studio/services/generator.py，适配点：
|
||
- 存储走本仓 services/storage 抽象（本地 data/media 兜底 / 七牛），save_bytes 直接返回可访问 URL；
|
||
- 参考图解析复用 storage.local_path 与 storage.download_bytes（防盗链 Referer 由 api/proxy.guess_referer 提供）；
|
||
- run_suite 新增可选 on_image_ok 回调：每张成功即通知调用方落库（product_assets/generated），
|
||
本模块不接触数据库，保持「套图引擎不依赖商品模型」的独立产品化裁剪能力。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import base64
|
||
import io
|
||
import logging
|
||
import mimetypes
|
||
import re
|
||
from dataclasses import dataclass
|
||
from typing import Awaitable, Callable
|
||
|
||
import httpx
|
||
from PIL import Image, ImageOps
|
||
|
||
from config import get_settings
|
||
from services.prompts import build_context, build_prompt, type_name
|
||
from services.tasks import (
|
||
IMG_FAILED,
|
||
IMG_OK,
|
||
TASK_DONE,
|
||
TASK_FAILED,
|
||
TASK_PARTIAL,
|
||
TASK_RUNNING,
|
||
Task,
|
||
TaskImage,
|
||
)
|
||
from services.watermark import apply_watermark
|
||
from services.storage import download_bytes, get_storage, local_path, read_local_bytes
|
||
|
||
log = logging.getLogger("suite.generator")
|
||
|
||
# 单张成功回调：(image) -> None；由 API 层注入用于回写商品素材
|
||
OnImageOk = Callable[[TaskImage], Awaitable[None]]
|
||
|
||
|
||
class ApiError(RuntimeError):
|
||
"""带 HTTP 状态码的 API 错误（用于区分可重试的网关/限流错误）。"""
|
||
|
||
def __init__(self, message: str, status: int = 0):
|
||
super().__init__(message)
|
||
self.status = status
|
||
|
||
|
||
_HTML_TITLE_RE = re.compile(r"<title[^>]*>(.*?)</title>", re.IGNORECASE | re.DOTALL)
|
||
|
||
|
||
def _raise_api_error(resp, provider: str):
|
||
"""HTTP 错误时抛出带 API 错误码/信息的异常（响应体里有真正的失败原因）。"""
|
||
if resp.is_success:
|
||
return
|
||
text = resp.text or ""
|
||
if "<html" in text[:300].lower() or text.lstrip()[:15].lower().startswith("<!doctype"):
|
||
# Cloudflare/网关错误页：取 <title> 作摘要，避免整段 HTML 进错误信息
|
||
m = _HTML_TITLE_RE.search(text)
|
||
detail = (re.sub(r"\s+", " ", m.group(1)).strip() if m else "") or "网关返回 HTML 错误页（上游/CDN 故障）"
|
||
raise ApiError(f"{provider} API HTTP {resp.status_code} — {detail}", resp.status_code)
|
||
try:
|
||
body = resp.json()
|
||
except Exception: # noqa: BLE001
|
||
raise ApiError(f"{provider} API HTTP {resp.status_code} — {text[:200]}", resp.status_code)
|
||
# 兼容两种错误形状：顶层 {code, message}（豆包/通义）与嵌套 {error: {code, message}}（new_api 系中转）
|
||
err = body.get("error") if isinstance(body.get("error"), dict) else body
|
||
detail = f"{err.get('code', '')}: {err.get('message', '')}".strip(': ')
|
||
raise ApiError(f"{provider} API HTTP {resp.status_code} — {detail or '无错误详情'}", resp.status_code)
|
||
|
||
# 参考图选择：material 用第 2 张（背面/细节），其余用第 1 张（正面）
|
||
TYPE_REF_INDEX = {
|
||
"material": 1,
|
||
}
|
||
DEFAULT_REF_COUNT = 2 # 每次生图最多带的参考图数（正面 1 张 + 背面/细节 1 张）
|
||
|
||
|
||
def _image_size(provider: str, ratio: str, is_wan: bool = True, model: str = "") -> str:
|
||
"""平台比例 → provider 尺寸参数。3:4 竖版（Ozon），1:1 方图（国内）。"""
|
||
if provider == "doubao":
|
||
return {"3:4": "1536x2048", "4:5": "1536x1920"}.get(ratio, "2048x2048")
|
||
if provider == "rightapi":
|
||
# gpt-image 自定义尺寸约束：16 的倍数、长短边比 ≤ 3:1（1536x2048 合法）
|
||
return {"3:4": "1536x2048", "4:5": "1536x1920"}.get(ratio, "2048x2048")
|
||
if provider == "pcore":
|
||
# 竖图必须与 Ozon 的 3:4 规格一致。1024x1536 实际是 2:3，会使成图明显偏瘦高。
|
||
return {"3:4": "1152x1536", "4:5": "1280x1600"}.get(ratio, "1024x1024")
|
||
# tongyi：万象与千问的 size 语法相同（* 分隔），档位不同
|
||
# wan2.6 系列总像素限制在 [1280², 1440²]，wan2.7 的 1536*2048/2048*2048 会超限
|
||
if model.startswith("wan2.6"):
|
||
return {"3:4": "1152*1536", "4:5": "1280*1600"}.get(ratio, "1440*1440")
|
||
if ratio == "3:4":
|
||
return "1536*2048" if is_wan else "768*1024"
|
||
if ratio == "4:5":
|
||
return "1536*1920" if is_wan else "960*1200"
|
||
return "2048*2048" if is_wan else "1024*1024"
|
||
|
||
_DOUBAO_ANTI_AI = (
|
||
"authentic real-world photography, natural imperfections, genuine texture, "
|
||
"no synthetic look, no CGI quality, no heavy post-processing"
|
||
)
|
||
|
||
DEFAULT_NEGATIVE_PROMPT = (
|
||
"AI-generated look, artificial, CGI quality, 3D render, synthetic texture, "
|
||
"plastic skin, mannequin-like, too perfect, oversaturated, HDR, heavy vignette, "
|
||
"low resolution, blurry, deformed, bad anatomy, overexposed, underexposed, grainy, "
|
||
"watermark, text distortion, bad typography, overlapping text, cheap look, cartoon"
|
||
)
|
||
|
||
|
||
# ── 参考图解析 ────────────────────────────────────────────────────────────
|
||
|
||
def _bytes_to_data_uri(data: bytes, mime: str) -> str:
|
||
return f"data:{mime};base64,{base64.b64encode(data).decode()}"
|
||
|
||
|
||
async def _resolve_ref_bytes(url: str) -> tuple[bytes, str]:
|
||
"""参考图 URL → (bytes, mime)。本地 media 文件直读磁盘；远程 URL 带 Referer 下载。
|
||
|
||
生图 API 的服务器无法访问 127.0.0.1，代理 URL 也不能直接透传，
|
||
所以统一在本地解析成原始字节再进请求体（data URI 或 multipart）。
|
||
"""
|
||
if url.startswith("data:"):
|
||
head, _, b64 = url.partition(",")
|
||
mime = head[5:].split(";", 1)[0] or "image/jpeg"
|
||
return base64.b64decode(b64), mime
|
||
path = local_path(url)
|
||
if path is not None:
|
||
mime = mimetypes.guess_type(path.name)[0] or "image/jpeg"
|
||
return read_local_bytes(path), mime
|
||
if url.startswith(("http://", "https://")):
|
||
from api.proxy import guess_referer
|
||
data, ctype = await download_bytes(url, referer=guess_referer(url))
|
||
if not ctype.startswith("image/"):
|
||
ctype = "image/jpeg"
|
||
return data, ctype
|
||
raise FileNotFoundError(f"无法解析参考图: {url}")
|
||
|
||
|
||
async def _resolve_ref(url: str) -> str:
|
||
data, mime = await _resolve_ref_bytes(url)
|
||
return _bytes_to_data_uri(data, mime)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class OutpaintCanvas:
|
||
"""用于受限扩图的预处理结果。
|
||
|
||
``image_data_url`` 与 ``mask_data_url`` 尺寸严格一致：蒙版不透明的区域是原图，
|
||
透明区域才允许供应商生成。最终仍会把 ``original`` 覆盖回结果，防止供应商在
|
||
蒙版区域有轻微重绘。
|
||
"""
|
||
|
||
image_data_url: str
|
||
mask_data_url: str
|
||
original: Image.Image
|
||
box: tuple[int, int, int, int]
|
||
size: tuple[int, int]
|
||
|
||
def restore_original(self, generated: bytes) -> bytes:
|
||
"""恢复原图区域，确保输出中原图像素不受模型影响。"""
|
||
with Image.open(io.BytesIO(generated)) as result:
|
||
result = ImageOps.exif_transpose(result).convert("RGBA")
|
||
if result.size != self.size:
|
||
# 供应商偶发不遵从 size 时仍输出固定 3:4；随后会精确覆盖原图区域。
|
||
result = ImageOps.fit(result, self.size, method=Image.Resampling.LANCZOS)
|
||
result.alpha_composite(self.original, self.box[:2])
|
||
out = io.BytesIO()
|
||
result.save(out, format="PNG")
|
||
return out.getvalue()
|
||
|
||
|
||
def _png_data_url(image: Image.Image) -> str:
|
||
out = io.BytesIO()
|
||
image.save(out, format="PNG")
|
||
return _bytes_to_data_uri(out.getvalue(), "image/png")
|
||
|
||
|
||
async def prepare_outpaint_canvas(url: str, size: str) -> OutpaintCanvas:
|
||
"""把单张参考图置于目标画布，并生成仅开放新增边缘的 PNG alpha 蒙版。"""
|
||
try:
|
||
target_width, target_height = (int(value) for value in size.lower().split("x", 1))
|
||
except (TypeError, ValueError) as exc:
|
||
raise ValueError(f"仅扩图需要合法的宽高尺寸，收到: {size}") from exc
|
||
if target_width <= 0 or target_height <= 0:
|
||
raise ValueError(f"仅扩图需要正数宽高尺寸，收到: {size}")
|
||
|
||
data, _ = await _resolve_ref_bytes(url)
|
||
with Image.open(io.BytesIO(data)) as opened:
|
||
source = ImageOps.exif_transpose(opened).convert("RGBA")
|
||
if source.width <= 0 or source.height <= 0:
|
||
raise ValueError("参考底图尺寸无效")
|
||
scale = min(target_width / source.width, target_height / source.height)
|
||
source_width = max(1, round(source.width * scale))
|
||
source_height = max(1, round(source.height * scale))
|
||
original = source.resize((source_width, source_height), Image.Resampling.LANCZOS)
|
||
|
||
x = (target_width - original.width) // 2
|
||
y = (target_height - original.height) // 2
|
||
canvas = Image.new("RGBA", (target_width, target_height), (0, 0, 0, 0))
|
||
canvas.alpha_composite(original, (x, y))
|
||
# OpenAI/GPT Image 兼容蒙版约定：透明像素是可编辑区域；原图区域不透明。
|
||
mask = Image.new("RGBA", (target_width, target_height), (0, 0, 0, 0))
|
||
mask.paste((255, 255, 255, 255), (x, y, x + original.width, y + original.height))
|
||
return OutpaintCanvas(
|
||
image_data_url=_png_data_url(canvas),
|
||
mask_data_url=_png_data_url(mask),
|
||
original=original,
|
||
box=(x, y, x + original.width, y + original.height),
|
||
size=(target_width, target_height),
|
||
)
|
||
|
||
|
||
# ── Provider：豆包 Seedream（火山方舟）────────────────────────────────────
|
||
|
||
async def generate_doubao(prompt: str, ref_images: list[str], size: str = "2048x2048", model: str | None = None) -> bytes:
|
||
s = get_settings()
|
||
if not s.ark_api_key:
|
||
raise RuntimeError("未配置 ARK_API_KEY（.env）")
|
||
body = {
|
||
"model": model or s.ark_image_model,
|
||
"prompt": prompt.rstrip(". ") + ". " + _DOUBAO_ANTI_AI,
|
||
"size": size,
|
||
"response_format": "url",
|
||
"watermark": False,
|
||
"n": 1,
|
||
}
|
||
if ref_images:
|
||
body["image"] = [await _resolve_ref(u) for u in ref_images]
|
||
async with httpx.AsyncClient(timeout=s.request_timeout, verify=False) as client:
|
||
resp = await client.post(
|
||
s.ark_base_url,
|
||
headers={"Authorization": f"Bearer {s.ark_api_key}", "Content-Type": "application/json"},
|
||
json=body,
|
||
)
|
||
_raise_api_error(resp, "豆包")
|
||
img_url = resp.json()["data"][0]["url"]
|
||
dl = await client.get(img_url, timeout=s.request_timeout)
|
||
dl.raise_for_status()
|
||
return dl.content
|
||
|
||
|
||
# ── Provider：通义万相 / 千问（DashScope）────────────────────────────────
|
||
|
||
def _is_wan_model(model: str) -> bool:
|
||
return model.lower().startswith("wan")
|
||
|
||
|
||
def _is_t2i_model(model: str) -> bool:
|
||
"""纯文生图模型（如 wan2.6-t2i）：不接受参考图，商品一致性只能靠文案描述。"""
|
||
return "t2i" in model.lower()
|
||
|
||
|
||
async def _tongyi_poll_task(client: httpx.AsyncClient, key: str, task_id: str, max_wait: int) -> str:
|
||
poll_url = "https://dashscope.aliyuncs.com/api/v1/tasks/" + task_id
|
||
elapsed, interval = 0, 3
|
||
while elapsed < max_wait:
|
||
resp = await client.get(poll_url, headers={"Authorization": f"Bearer {key}"}, timeout=30)
|
||
resp.raise_for_status()
|
||
result = resp.json()
|
||
status = result.get("output", {}).get("task_status", "")
|
||
if status == "SUCCEEDED":
|
||
choices = result["output"].get("choices", [])
|
||
if choices:
|
||
content = choices[0].get("message", {}).get("content", [])
|
||
if content:
|
||
return content[0].get("image", "")
|
||
results = result["output"].get("results", [])
|
||
if results:
|
||
return results[0].get("url") or results[0].get("b64_image", "")
|
||
raise RuntimeError(f"通义任务成功但无结果: {result}")
|
||
if status in ("FAILED", "UNKNOWN"):
|
||
raise RuntimeError(f"通义任务失败: {result}")
|
||
await asyncio.sleep(interval)
|
||
elapsed += interval
|
||
interval = min(interval + 2, 10)
|
||
raise TimeoutError(f"通义异步任务超时 ({max_wait}s): task_id={task_id}")
|
||
|
||
|
||
async def generate_tongyi(prompt: str, ref_images: list[str], size: str = "2048*2048", model: str | None = None) -> bytes:
|
||
s = get_settings()
|
||
if not s.dashscope_api_key:
|
||
raise RuntimeError("未配置 DASHSCOPE_API_KEY（.env）")
|
||
model = model or s.dashscope_model
|
||
is_wan = _is_wan_model(model)
|
||
url = s.dashscope_base_url or (
|
||
"https://dashscope.aliyuncs.com/api/v1/services/aigc/image-generation/generation"
|
||
if is_wan
|
||
else "https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"
|
||
)
|
||
|
||
# t2i 模型不接受参考图：content 只有文本，商品一致性依赖 prompt 里的标题/卖点描述
|
||
content: list[dict] = []
|
||
if not _is_t2i_model(model):
|
||
content = [{"image": await _resolve_ref(u)} for u in ref_images]
|
||
content.append({"text": prompt})
|
||
|
||
params = {"size": size, "n": 1, "watermark": False}
|
||
if not is_wan:
|
||
params["prompt_extend"] = False
|
||
params["negative_prompt"] = DEFAULT_NEGATIVE_PROMPT[:500]
|
||
|
||
headers = {"Authorization": f"Bearer {s.dashscope_api_key}", "Content-Type": "application/json"}
|
||
if is_wan:
|
||
headers["X-DashScope-Async"] = "enable"
|
||
|
||
body = {"model": model, "input": {"messages": [{"role": "user", "content": content}]}, "parameters": params}
|
||
|
||
async with httpx.AsyncClient(timeout=s.request_timeout, verify=False) as client:
|
||
resp = await client.post(url, headers=headers, json=body)
|
||
_raise_api_error(resp, "通义")
|
||
data = resp.json()
|
||
if is_wan:
|
||
task_id = data.get("output", {}).get("task_id", "")
|
||
if not task_id:
|
||
raise RuntimeError(f"通义万象未返回 task_id: {data}")
|
||
img_url = await _tongyi_poll_task(client, s.dashscope_api_key, task_id, s.poll_max_wait)
|
||
if img_url.startswith("data:") or len(img_url) > 500:
|
||
return base64.b64decode(img_url.split(",", 1)[-1] if "," in img_url else img_url)
|
||
dl = await client.get(img_url, timeout=s.request_timeout)
|
||
dl.raise_for_status()
|
||
return dl.content
|
||
img_url = data["output"]["choices"][0]["message"]["content"][0]["image"]
|
||
dl = await client.get(img_url, timeout=s.request_timeout)
|
||
dl.raise_for_status()
|
||
return dl.content
|
||
|
||
|
||
# ── Provider：RightAPI（gpt-image / nano-banana，OpenAI 兼容中转）──────────
|
||
|
||
# 可重试的状态码：中转限流/网关抖动（该中转限流时返回 Cloudflare 502 而非 429）
|
||
RETRYABLE_STATUS = {429, 500, 502, 503, 504}
|
||
|
||
|
||
async def _rightapi_poll_task(client: httpx.AsyncClient, headers: dict, origin: str,
|
||
task_id: str, max_wait: int) -> dict:
|
||
"""轮询站点级任务查询接口 GET /v1/tasks/{task_id}（不带 /draw 前缀）。
|
||
|
||
实测要点：
|
||
- 完成响应**没有** status:"completed" 字段，完成判定 = 响应里出现 data；
|
||
- progress 基本不动（0~2），不能当进度条依据；
|
||
- 失败态 = status 为 failed / error / cancelled。
|
||
"""
|
||
poll_url = f"{origin}/v1/tasks/{task_id}"
|
||
elapsed, interval = 0, 3
|
||
while elapsed < max_wait:
|
||
resp = await client.get(poll_url, headers=headers, timeout=30)
|
||
_raise_api_error(resp, "RightAPI")
|
||
result = resp.json()
|
||
status = result.get("status", "")
|
||
if status in ("failed", "error", "cancelled"):
|
||
err = result.get("error") or {}
|
||
raise RuntimeError(f"RightAPI 任务失败: {err.get('message') or result}")
|
||
if "data" in result:
|
||
return result
|
||
await asyncio.sleep(interval)
|
||
elapsed += interval
|
||
interval = min(interval + 2, 10)
|
||
raise TimeoutError(f"RightAPI 异步任务超时 ({max_wait}s): task_id={task_id}")
|
||
|
||
|
||
async def _rightapi_poll_task_once(client: httpx.AsyncClient, headers: dict, origin: str, task_id: str) -> dict | None:
|
||
"""持久化 Worker 的单次查询版本；未完成返回 None，不占住 Worker 等待。"""
|
||
response = await client.get(f"{origin}/v1/tasks/{task_id}", headers=headers, timeout=30)
|
||
_raise_api_error(response, "RightAPI")
|
||
result = response.json()
|
||
status = result.get("status", "")
|
||
if status in ("failed", "error", "cancelled"):
|
||
error = result.get("error") or {}
|
||
raise RuntimeError(f"RightAPI 任务失败: {error.get('message') or result}")
|
||
return result if "data" in result else None
|
||
|
||
|
||
def _rightapi_extract_image(result: dict) -> tuple[str | None, str | None]:
|
||
"""从轮询完成结果里取 (kind, payload)：kind ∈ url | b64，未取到返回 (None, None)。
|
||
|
||
完成形状为 Images 协议：{"created":..., "data":[{"url": "..."}]}（实测只见 url）。
|
||
"""
|
||
data = result.get("data") or []
|
||
if data:
|
||
item = data[0] or {}
|
||
url = item.get("url") or ""
|
||
if url:
|
||
return ("url", url)
|
||
b64 = item.get("b64_json") or ""
|
||
if b64:
|
||
return ("b64", b64)
|
||
return (None, None)
|
||
|
||
|
||
async def _rightapi_request(s, prompt: str, ref_images: list[str], size: str, model: str) -> bytes:
|
||
"""RightAPI 各模型：统一走 /v1/images/generations（异步）。
|
||
|
||
官方协议（2026-07 起统一异步）：
|
||
- POST /draw/v1/images/generations，请求体固定带 async:true，参考图放 image 数组（data-URI）；
|
||
- 返回 task_id 后轮询 GET /v1/tasks/{task_id}（站点级，不带 /draw）；
|
||
- 参数只有 model/prompt/n/size/imageSize/image/async；不传 quality/output_format/input_fidelity。
|
||
参考图沿用现选图逻辑（≤2 张，image 数组）。单张 1-5 分钟，轮询上限 poll_max_wait 兜底。
|
||
"""
|
||
base = s.rightapi_base_url.rstrip("/")
|
||
# 任务查询是站点级接口，不带 /draw：从 base 里拆出 origin（https://rightapi.ai/draw → https://rightapi.ai）
|
||
origin = base.split("/draw", 1)[0].rstrip("/") or base
|
||
headers = {"Authorization": f"Bearer {s.rightapi_api_key}"}
|
||
|
||
body = {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"n": 1,
|
||
"size": size,
|
||
"async": True,
|
||
}
|
||
if ref_images:
|
||
body["image"] = [await _resolve_ref(u) for u in ref_images]
|
||
|
||
async with httpx.AsyncClient(timeout=max(s.request_timeout, 600), verify=False) as client:
|
||
resp = await client.post(
|
||
f"{base}/v1/images/generations",
|
||
headers={**headers, "Content-Type": "application/json"},
|
||
json=body,
|
||
)
|
||
_raise_api_error(resp, "RightAPI")
|
||
submitted = resp.json()
|
||
task_id = submitted.get("task_id") or ""
|
||
if task_id:
|
||
result = await _rightapi_poll_task(client, headers, origin, task_id, s.poll_max_wait)
|
||
else:
|
||
# 极端兜底：个别中转可能同步返回 data（文档不保证，但防御处理）
|
||
result = submitted
|
||
|
||
kind, payload = _rightapi_extract_image(result)
|
||
if kind == "b64" and payload:
|
||
return base64.b64decode(payload.split(",", 1)[-1] if "," in payload else payload)
|
||
if kind == "url" and payload:
|
||
dl = await client.get(payload, timeout=s.request_timeout)
|
||
dl.raise_for_status()
|
||
return dl.content
|
||
raise RuntimeError(f"RightAPI 任务完成但没有图片数据: {result}")
|
||
|
||
|
||
async def generate_rightapi(prompt: str, ref_images: list[str], size: str = "2048x2048", model: str | None = None) -> bytes:
|
||
"""带重试的 RightAPI 入口：429/5xx/超时按递增间隔重试。
|
||
|
||
实测该中转对同 key 连续请求有分钟级冷却（成功一张后紧接着的请求会被网关秒拒 502），
|
||
60s → 120s → 240s 的退避基本能等到窗口放开。
|
||
"""
|
||
s = get_settings()
|
||
if not s.rightapi_api_key:
|
||
raise RuntimeError("未配置 RIGHTAPI_API_KEY（.env）")
|
||
model = model or s.rightapi_image_model
|
||
attempts = max(1, s.rightapi_max_retries)
|
||
|
||
last_exc: Exception | None = None
|
||
for i in range(attempts):
|
||
try:
|
||
return await _rightapi_request(s, prompt, ref_images, size, model)
|
||
except ApiError as exc:
|
||
last_exc = exc
|
||
if exc.status not in RETRYABLE_STATUS:
|
||
raise # 参数错误等不可重试，立即失败
|
||
except (httpx.TimeoutException, httpx.TransportError) as exc:
|
||
last_exc = exc # 网络抖动/超时可重试
|
||
if i == attempts - 1:
|
||
break
|
||
wait = s.rightapi_retry_wait * (2 ** i)
|
||
log.warning("RightAPI 第 %d/%d 次请求失败（%s），%ds 后重试", i + 1, attempts, last_exc, wait)
|
||
await asyncio.sleep(wait)
|
||
raise last_exc # type: ignore[misc]
|
||
|
||
|
||
# ── Provider：Pcore（pcore.ai，OpenAI 兼容异步图片协议）──────────────────
|
||
|
||
async def _pcore_poll_task(client: httpx.AsyncClient, headers: dict, base: str,
|
||
kind: str, task_id: str, max_wait: int) -> dict:
|
||
"""轮询 Pcore 任务：GET {base}/v1/images/{kind}/{task_id}，kind ∈ generations | edits。
|
||
|
||
状态机：queued / in_progress → 继续；completed / succeeded → 返回；
|
||
failed / error / cancelled → 抛错。
|
||
"""
|
||
poll_url = f"{base}/v1/images/{kind}/{task_id}"
|
||
elapsed, interval = 0, 3
|
||
while elapsed < max_wait:
|
||
resp = await client.get(poll_url, headers=headers, timeout=30)
|
||
_raise_api_error(resp, "Pcore")
|
||
result = resp.json()
|
||
status = result.get("status", "")
|
||
if status in ("failed", "error", "cancelled"):
|
||
err = result.get("error") or {}
|
||
raise RuntimeError(f"Pcore 任务失败: {err.get('message') or result}")
|
||
if status in ("completed", "succeeded"):
|
||
return result
|
||
await asyncio.sleep(interval)
|
||
elapsed += interval
|
||
interval = min(interval + 2, 10)
|
||
raise TimeoutError(f"Pcore 异步任务超时 ({max_wait}s): task_id={task_id}")
|
||
|
||
|
||
async def _pcore_poll_task_once(client: httpx.AsyncClient, headers: dict, base: str,
|
||
kind: str, task_id: str) -> dict | None:
|
||
"""持久化 Worker 的单次查询版本；未完成返回 None。"""
|
||
response = await client.get(f"{base}/v1/images/{kind}/{task_id}", headers=headers, timeout=30)
|
||
_raise_api_error(response, "Pcore")
|
||
result = response.json()
|
||
status = result.get("status", "")
|
||
if status in ("failed", "error", "cancelled"):
|
||
error = result.get("error") or {}
|
||
raise RuntimeError(f"Pcore 任务失败: {error.get('message') or result}")
|
||
return result if status in ("completed", "succeeded") else None
|
||
|
||
|
||
async def _pcore_request(s, prompt: str, ref_images: list[str], size: str, model: str) -> bytes:
|
||
"""Pcore：无参考图走 /v1/images/generations，带参考图走 /v1/images/edits（JSON，data-URI）。
|
||
|
||
两者均为异步：提交返回任务 id → 轮询对应资源路径 → 完成后 data[0].url 取图。
|
||
完成形状与 RightAPI 同为 Images 协议，复用 _rightapi_extract_image。
|
||
"""
|
||
base = s.pcore_base_url.rstrip("/")
|
||
headers = {"Authorization": f"Bearer {s.pcore_api_key}"}
|
||
|
||
kind = "edits" if ref_images else "generations"
|
||
body: dict = {"model": model, "prompt": prompt, "n": 1, "size": size, "async": True}
|
||
if ref_images:
|
||
# Pcore 的 gpt-image-2 通道实际要求 OpenAI 兼容对象数组。
|
||
# 虽然通用文档描述字符串数组也可用，该通道会在异步任务中校验 images[].image_url。
|
||
body["images"] = [{"image_url": await _resolve_ref(u)} for u in ref_images]
|
||
|
||
async with httpx.AsyncClient(timeout=max(s.request_timeout, 600), verify=False) as client:
|
||
resp = await client.post(
|
||
f"{base}/v1/images/{kind}",
|
||
headers={**headers, "Content-Type": "application/json"},
|
||
json=body,
|
||
)
|
||
_raise_api_error(resp, "Pcore")
|
||
submitted = resp.json()
|
||
# new_api 系中转常见「HTTP 200 + 错误体」，需主动检查 error 字段
|
||
if isinstance(submitted, dict) and submitted.get("error"):
|
||
err = submitted["error"] or {}
|
||
raise RuntimeError(f"Pcore 提交失败: {err.get('message') or submitted}")
|
||
task_id = submitted.get("task_id") or submitted.get("id") or ""
|
||
if not task_id:
|
||
raise RuntimeError(f"Pcore 提交未返回任务 id: {submitted}")
|
||
result = await _pcore_poll_task(client, headers, base, kind, task_id, s.poll_max_wait)
|
||
|
||
img_kind, payload = _rightapi_extract_image(result)
|
||
if img_kind == "b64" and payload:
|
||
return base64.b64decode(payload.split(",", 1)[-1] if "," in payload else payload)
|
||
if img_kind == "url" and payload:
|
||
dl = await client.get(payload, timeout=s.request_timeout)
|
||
dl.raise_for_status()
|
||
return dl.content
|
||
raise RuntimeError(f"Pcore 任务完成但没有图片数据: {result}")
|
||
|
||
|
||
async def generate_pcore(prompt: str, ref_images: list[str], size: str = "1536x2048", model: str | None = None) -> bytes:
|
||
"""带重试的 Pcore 入口：429/5xx/超时按递增间隔重试（与 RightAPI 同策略）。"""
|
||
s = get_settings()
|
||
if not s.pcore_api_key:
|
||
raise RuntimeError("未配置 PCORE_API_KEY（.env）")
|
||
model = model or "gpt-image-2"
|
||
attempts = max(1, s.pcore_max_retries)
|
||
|
||
last_exc: Exception | None = None
|
||
for i in range(attempts):
|
||
try:
|
||
return await _pcore_request(s, prompt, ref_images, size, model)
|
||
except ApiError as exc:
|
||
last_exc = exc
|
||
if exc.status not in RETRYABLE_STATUS:
|
||
raise
|
||
# 通道不存在（模型未开通）重试无意义，立即失败
|
||
if "no available channel" in str(exc).lower() or "model_not_found" in str(exc).lower():
|
||
raise
|
||
except (httpx.TimeoutException, httpx.TransportError) as exc:
|
||
last_exc = exc # 网络抖动/超时可重试
|
||
if i == attempts - 1:
|
||
break
|
||
wait = s.pcore_retry_wait * (2 ** i)
|
||
log.warning("Pcore 第 %d/%d 次请求失败（%s），%ds 后重试", i + 1, attempts, last_exc, wait)
|
||
await asyncio.sleep(wait)
|
||
raise last_exc # type: ignore[misc]
|
||
|
||
|
||
# ── 可恢复的逐步 provider 调用（持久化 Worker 使用）──────────────────────
|
||
|
||
@dataclass(frozen=True)
|
||
class GenerationAdvance:
|
||
"""一次 Worker 推进的结果：data 有值时已完成，否则持久化外部 task id 后稍后再查。"""
|
||
|
||
data: bytes | None = None
|
||
provider_task_id: str | None = None
|
||
provider_task_kind: str | None = None
|
||
|
||
|
||
async def _download_provider_image(client: httpx.AsyncClient, result: dict, provider: str) -> bytes:
|
||
kind, payload = _rightapi_extract_image(result)
|
||
if kind == "b64" and payload:
|
||
return base64.b64decode(payload.split(",", 1)[-1] if "," in payload else payload)
|
||
if kind == "url" and payload:
|
||
response = await client.get(payload, timeout=get_settings().request_timeout)
|
||
_raise_api_error(response, provider)
|
||
return response.content
|
||
raise RuntimeError(f"{provider} 任务完成但没有图片数据: {result}")
|
||
|
||
|
||
async def _poll_tongyi_once(client: httpx.AsyncClient, key: str, task_id: str) -> tuple[str, str | None]:
|
||
response = await client.get(
|
||
"https://dashscope.aliyuncs.com/api/v1/tasks/" + task_id,
|
||
headers={"Authorization": f"Bearer {key}"}, timeout=30,
|
||
)
|
||
_raise_api_error(response, "通义")
|
||
result = response.json()
|
||
status = result.get("output", {}).get("task_status", "")
|
||
if status in ("FAILED", "UNKNOWN"):
|
||
raise RuntimeError(f"通义任务失败: {result}")
|
||
if status != "SUCCEEDED":
|
||
return "waiting", None
|
||
choices = result.get("output", {}).get("choices", [])
|
||
if choices:
|
||
content = choices[0].get("message", {}).get("content", [])
|
||
if content:
|
||
return "done", content[0].get("image", "")
|
||
results = result.get("output", {}).get("results", [])
|
||
if results:
|
||
return "done", results[0].get("url") or results[0].get("b64_image", "")
|
||
raise RuntimeError(f"通义任务成功但无结果: {result}")
|
||
|
||
|
||
async def advance_generation(
|
||
provider: str,
|
||
prompt: str,
|
||
ref_images: list[str],
|
||
*,
|
||
size: str,
|
||
model: str,
|
||
provider_task_id: str | None = None,
|
||
provider_task_kind: str | None = None,
|
||
edit_mask: str | None = None,
|
||
) -> GenerationAdvance:
|
||
"""推进一次图片生成，不在本函数内无限轮询。
|
||
|
||
对异步供应商：首次仅提交并返回 task id；后续调用只轮询该 id。这样 Worker 重启后
|
||
能继续外部任务，且在等待期间可以领取其他任务。豆包和同步千问仍在一次调用内完成。
|
||
"""
|
||
settings = get_settings()
|
||
if edit_mask and provider != "pcore":
|
||
raise RuntimeError("带蒙版的仅扩图当前仅支持 Pcore gpt-image-2")
|
||
if provider == "doubao":
|
||
return GenerationAdvance(data=await generate_doubao(prompt, ref_images, size=size, model=model))
|
||
|
||
if provider == "tongyi":
|
||
if not settings.dashscope_api_key:
|
||
raise RuntimeError("未配置 DASHSCOPE_API_KEY（.env）")
|
||
is_wan = _is_wan_model(model)
|
||
if not is_wan:
|
||
return GenerationAdvance(data=await generate_tongyi(prompt, ref_images, size=size, model=model))
|
||
headers = {"Authorization": f"Bearer {settings.dashscope_api_key}", "Content-Type": "application/json"}
|
||
if provider_task_id:
|
||
async with httpx.AsyncClient(timeout=60, verify=False) as client:
|
||
state, payload = await _poll_tongyi_once(client, settings.dashscope_api_key, provider_task_id)
|
||
if state == "waiting":
|
||
return GenerationAdvance(provider_task_id=provider_task_id, provider_task_kind="tongyi_wan")
|
||
if payload and (payload.startswith("data:") or len(payload) > 500):
|
||
return GenerationAdvance(data=base64.b64decode(payload.split(",", 1)[-1] if "," in payload else payload))
|
||
if not payload:
|
||
raise RuntimeError("通义任务成功但图片地址为空")
|
||
response = await client.get(payload, timeout=settings.request_timeout)
|
||
_raise_api_error(response, "通义")
|
||
return GenerationAdvance(data=response.content)
|
||
url = settings.dashscope_base_url or "https://dashscope.aliyuncs.com/api/v1/services/aigc/image-generation/generation"
|
||
content: list[dict] = []
|
||
if not _is_t2i_model(model):
|
||
content = [{"image": await _resolve_ref(item)} for item in ref_images]
|
||
content.append({"text": prompt})
|
||
body = {
|
||
"model": model,
|
||
"input": {"messages": [{"role": "user", "content": content}]},
|
||
"parameters": {"size": size, "n": 1, "watermark": False, "negative_prompt": DEFAULT_NEGATIVE_PROMPT[:500]},
|
||
}
|
||
async with httpx.AsyncClient(timeout=settings.request_timeout, verify=False) as client:
|
||
response = await client.post(url, headers={**headers, "X-DashScope-Async": "enable"}, json=body)
|
||
_raise_api_error(response, "通义")
|
||
task_id = response.json().get("output", {}).get("task_id", "")
|
||
if not task_id:
|
||
raise RuntimeError(f"通义万象未返回 task_id: {response.json()}")
|
||
return GenerationAdvance(provider_task_id=task_id, provider_task_kind="tongyi_wan")
|
||
|
||
if provider == "rightapi":
|
||
if not settings.rightapi_api_key:
|
||
raise RuntimeError("未配置 RIGHTAPI_API_KEY（.env）")
|
||
base = settings.rightapi_base_url.rstrip("/")
|
||
origin = base.split("/draw", 1)[0].rstrip("/") or base
|
||
headers = {"Authorization": f"Bearer {settings.rightapi_api_key}"}
|
||
async with httpx.AsyncClient(timeout=max(settings.request_timeout, 600), verify=False) as client:
|
||
if provider_task_id:
|
||
result = await _rightapi_poll_task_once(client, headers, origin, provider_task_id)
|
||
if result is None:
|
||
return GenerationAdvance(provider_task_id=provider_task_id, provider_task_kind="rightapi")
|
||
return GenerationAdvance(data=await _download_provider_image(client, result, "RightAPI"))
|
||
body: dict = {"model": model, "prompt": prompt, "n": 1, "size": size, "async": True}
|
||
if ref_images:
|
||
body["image"] = [await _resolve_ref(item) for item in ref_images]
|
||
response = await client.post(
|
||
f"{base}/v1/images/generations", headers={**headers, "Content-Type": "application/json"}, json=body,
|
||
)
|
||
_raise_api_error(response, "RightAPI")
|
||
submitted = response.json()
|
||
task_id = submitted.get("task_id") or ""
|
||
if task_id:
|
||
return GenerationAdvance(provider_task_id=task_id, provider_task_kind="rightapi")
|
||
return GenerationAdvance(data=await _download_provider_image(client, submitted, "RightAPI"))
|
||
|
||
if provider == "pcore":
|
||
if not settings.pcore_api_key:
|
||
raise RuntimeError("未配置 PCORE_API_KEY（.env）")
|
||
base = settings.pcore_base_url.rstrip("/")
|
||
headers = {"Authorization": f"Bearer {settings.pcore_api_key}"}
|
||
async with httpx.AsyncClient(timeout=max(settings.request_timeout, 600), verify=False) as client:
|
||
if provider_task_id:
|
||
kind = provider_task_kind or "edits"
|
||
result = await _pcore_poll_task_once(client, headers, base, kind, provider_task_id)
|
||
if result is None:
|
||
return GenerationAdvance(provider_task_id=provider_task_id, provider_task_kind=kind)
|
||
return GenerationAdvance(data=await _download_provider_image(client, result, "Pcore"))
|
||
kind = "edits" if ref_images else "generations"
|
||
body: dict = {"model": model, "prompt": prompt, "n": 1, "size": size, "async": True}
|
||
if ref_images:
|
||
# 与 _pcore_request 保持同一经 gpt-image-2 实际验证的形状，避免 Worker 路径漂移。
|
||
body["images"] = [{"image_url": await _resolve_ref(item)} for item in ref_images]
|
||
if edit_mask:
|
||
# Pcore 的图片参数使用 JSON {image_url}；蒙版沿用相同的 OpenAI 兼容对象形式。
|
||
# 透明区域允许编辑，不透明区域保留。若供应商轻微偏移，Worker 会二次覆盖原图。
|
||
body["mask"] = {"image_url": edit_mask}
|
||
response = await client.post(
|
||
f"{base}/v1/images/{kind}", headers={**headers, "Content-Type": "application/json"}, json=body,
|
||
)
|
||
_raise_api_error(response, "Pcore")
|
||
submitted = response.json()
|
||
if isinstance(submitted, dict) and submitted.get("error"):
|
||
error = submitted["error"] or {}
|
||
raise RuntimeError(f"Pcore 提交失败: {error.get('message') or submitted}")
|
||
task_id = submitted.get("task_id") or submitted.get("id") or ""
|
||
if not task_id:
|
||
raise RuntimeError(f"Pcore 提交未返回任务 id: {submitted}")
|
||
return GenerationAdvance(provider_task_id=task_id, provider_task_kind=kind)
|
||
|
||
raise RuntimeError(f"未知 provider: {provider}")
|
||
|
||
|
||
GENERATORS = {
|
||
"doubao": generate_doubao,
|
||
"tongyi": generate_tongyi,
|
||
"rightapi": generate_rightapi,
|
||
"pcore": generate_pcore,
|
||
}
|
||
|
||
|
||
# ── 任务执行器 ────────────────────────────────────────────────────────────
|
||
|
||
def _order_refs(refs: list[str], type_id: str) -> list[str]:
|
||
"""参考图槽位选择 + 截断：material 偏好第 2 张，其余用第 1 张。"""
|
||
preferred = TYPE_REF_INDEX.get(type_id)
|
||
if preferred is not None and len(refs) > preferred:
|
||
refs = [refs[preferred]] + [r for i, r in enumerate(refs) if i != preferred]
|
||
return refs[:DEFAULT_REF_COUNT]
|
||
|
||
|
||
def _refs_for_job(images: list[dict], job: dict) -> list[str]:
|
||
"""无状态路径：按方案项选参考图。
|
||
|
||
优先 variant_name 精确匹配（「主图·粉色」用粉色那张 SKU 图）；
|
||
匹配不到则回退 main 组第一张（再退到任意第一张）。
|
||
"""
|
||
variant = job.get("variant_name")
|
||
if variant:
|
||
matched = [i["url"] for i in images if i.get("variant_name") == variant]
|
||
if matched:
|
||
return matched[:DEFAULT_REF_COUNT]
|
||
mains = [i["url"] for i in images if i.get("group_key") == "main"]
|
||
others = [i["url"] for i in images if i.get("group_key") != "main"]
|
||
pool = mains or others or [i["url"] for i in images]
|
||
if not pool:
|
||
raise RuntimeError("任务没有参考图")
|
||
return _order_refs(pool, job.get("kind", ""))
|
||
|
||
|
||
# 串行生成队列：所有用户共享同一批 API key，并发生成会触发中转限流
|
||
# （rightapi 同 key 分钟级冷却）；同一时间只跑一个任务，其余保持 pending 排队。
|
||
_GEN_LOCK = asyncio.Lock()
|
||
|
||
|
||
async def _store_image_bytes(task: Task, index: int, data: bytes, version: str = "result") -> str:
|
||
"""生成图字节 → 存储落盘，返回可访问 URL。按 PNG 魔数定扩展名（部分中转不遵守 output_format）。"""
|
||
is_png = data[:8] == b"\x89PNG\r\n\x1a\n"
|
||
ext = ".png" if is_png else ".jpg"
|
||
ctype = "image/png" if is_png else "image/jpeg"
|
||
return await get_storage().save_bytes(data, f"suites/{task.id}/{version}/{index}{ext}", ctype)
|
||
|
||
|
||
async def run_suite(task: Task, on_image_ok: OnImageOk | None = None) -> None:
|
||
"""后台执行套图任务：排队 → 逐张生成 → 落盘 → 更新内存状态；单张失败不中断。
|
||
|
||
on_image_ok：每张成功落盘后的回调（API 层用于实时回写 product_assets/generated），
|
||
回调异常只记日志，不影响任务本身。
|
||
"""
|
||
settings = get_settings()
|
||
provider_name = task.provider or settings.image_provider
|
||
generator = GENERATORS.get(provider_name)
|
||
if generator is None:
|
||
task.status = TASK_FAILED
|
||
task.error = f"未知 provider: {provider_name}"
|
||
return
|
||
|
||
ctx = build_context(task.context or {}, fallback_name="")
|
||
model = task.model or {
|
||
"tongyi": settings.dashscope_model,
|
||
"rightapi": settings.rightapi_image_model,
|
||
}.get(provider_name, settings.ark_image_model)
|
||
is_wan = provider_name == "tongyi" and _is_wan_model(model)
|
||
size = _image_size(provider_name, task.ratio, is_wan=is_wan, model=model)
|
||
jobs = [dict(j) for j in task.plan]
|
||
|
||
async with _GEN_LOCK:
|
||
task.status = TASK_RUNNING
|
||
ok, failed = 0, 0
|
||
failures: list[str] = []
|
||
for job in jobs:
|
||
type_id = job["kind"]
|
||
image = TaskImage(type_id=type_id, name=job.get("title") or type_name(type_id))
|
||
task.images.append(image)
|
||
try:
|
||
# 提示词按模型家族分发：国产主体参考 / gpt edits 保真 / google 主体保持
|
||
# 单张 AI 编辑可传入 prompt_extra：任务展示名保留「AI 生图」，
|
||
# 但不把中文展示名拼进模型 prompt，避免生成图内出现中文。
|
||
prompt_extra = job.get("prompt_extra") or job
|
||
prompt = build_prompt(
|
||
provider_name, model, type_id, ctx, task.style_set, task.lang,
|
||
extra=prompt_extra, style_prompt=task.style_prompt, requirements=task.requirements,
|
||
)
|
||
refs = _refs_for_job(list(task.ref_images or []), job)
|
||
data = await generator(prompt, refs, size=size, model=model)
|
||
# 水印：AI 出图返回后、落盘前的后处理（失败不阻断，内部返回原图）
|
||
wm = task.watermark or {}
|
||
if wm.get("enabled"):
|
||
image.source_url = await _store_image_bytes(task, len(task.images), data, version="original")
|
||
data = apply_watermark(data, wm)
|
||
image.url = await _store_image_bytes(task, len(task.images), data, version="watermarked")
|
||
image.watermark_config = wm
|
||
else:
|
||
image.url = await _store_image_bytes(task, len(task.images), data)
|
||
image.status = IMG_OK
|
||
ok += 1
|
||
if on_image_ok is not None and image.url:
|
||
try:
|
||
await on_image_ok(image)
|
||
except Exception: # noqa: BLE001
|
||
log.exception("套图 %s 第 %d 张回写商品素材失败", task.id, len(task.images))
|
||
except Exception as exc: # noqa: BLE001
|
||
log.exception("套图 %s 类型 %s 生成失败", task.id, type_id)
|
||
image.status = IMG_FAILED # 默认 pending，失败显式置 failed
|
||
image.error = str(exc)[:500]
|
||
failures.append(f"{job.get('title') or type_name(type_id)}：{str(exc)[:200]}")
|
||
failed += 1
|
||
|
||
task.status = TASK_DONE if failed == 0 else (TASK_PARTIAL if ok > 0 else TASK_FAILED)
|
||
if failed:
|
||
uniq = list(dict.fromkeys(failures)) # 去重保序
|
||
detail = "；".join(uniq[:6])
|
||
if len(uniq) > 6:
|
||
detail += f"；…等共 {failed} 张失败"
|
||
if ok == 0:
|
||
task.error = f"全部生成失败。{detail}"
|
||
else:
|
||
task.error = f"部分生成失败（{failed} 张）。{detail}"
|