Files
image-suite-studio/server/services/generator.py
T
2026-08-19 17:05:23 +08:00

488 lines
22 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""套图生成服务:图像 provider(豆包 Seedream / 通义万相)+ 任务执行器。
Provider 调用方式移植自 ecommerce-image-suite/scripts/generate.py
- doubao:火山方舟 images/generations,同步返回 URL;参考图走 image 字段(data URI
- tongyi:wan* 万象模型走异步任务轮询;qwen* 走同步 multimodal-generation
"""
from __future__ import annotations
import asyncio
import base64
import logging
import mimetypes
import re
from uuid import UUID
import httpx
from sqlalchemy import select
from config import get_settings
from db import get_session_factory
from models import Product, ProductAsset, Suite, SuiteImage, SUITE_RUNNING, SUITE_DONE, SUITE_PARTIAL, SUITE_FAILED, STATUS_OK, STATUS_FAILED
from services import storage
from services.prompt import build_prompt, build_context, type_name, wrap_prompt_for_gpt_edits
log = logging.getLogger("suite.generator")
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()
detail = f"{body.get('code', '')}: {body.get('message', '')}".strip(': ')
except Exception: # noqa: BLE001
detail = text[:200]
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/WB),1:1 方图(国内)。"""
if provider == "doubao":
return "1536x2048" if ratio == "3:4" else "2048x2048"
if provider == "rightapi":
# gpt-image 自定义尺寸约束:16 的倍数、长短边比 ≤ 3:1(1536x2048 合法)
return "1536x2048" if ratio == "3:4" else "2048x2048"
# tongyi:万象与千问的 size 语法相同(* 分隔),档位不同
# wan2.6 系列总像素限制在 [1280², 1440²]wan2.7 的 1536*2048/2048*2048 会超限
if model.startswith("wan2.6"):
return "1152*1536" if ratio == "3:4" else "1440*1440"
if ratio == "3:4":
return "1536*2048" if is_wan else "768*1024"
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 = storage.local_path(url)
if path is not None:
mime = mimetypes.guess_type(path.name)[0] or "image/jpeg"
return path.read_bytes(), mime
if url.startswith(("http://", "https://")):
from api.proxy import guess_referer
data, ctype = await storage.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)
# ── 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
# ── ProviderRightAPIgpt-imageOpenAI 兼容中转)──────────────────────
# 可重试的状态码:中转限流/网关抖动(该中转限流时返回 Cloudflare 502 而非 429
RETRYABLE_STATUS = {429, 500, 502, 503, 504}
# 中转对 input_fidelity 参数的支持探测:None=未探测,True=支持,False=不支持(已降级)
_rightapi_fidelity_supported: bool | None = None
async def _rightapi_request(s, prompt: str, ref_images: list[str], size: str, model: str) -> bytes:
"""gpt-image 系列:有参考图走 /v1/images/editsmultipart),无参考图走 /v1/images/generations。
OpenAI 兼容协议:响应固定 b64_json(不支持 response_format 参数,传了报 400);
同步调用无任务轮询,高质量档单张 1-5 分钟,超时按文档建议兜底 600s。
input_fidelity=high 强制高保真保留输入图细节(商品一致性关键参数,仅 edits 端点);
中转若不认该参数(400),自动去掉重试并记住,后续请求不再带。
"""
global _rightapi_fidelity_supported
base = s.rightapi_base_url.rstrip("/")
headers = {"Authorization": f"Bearer {s.rightapi_api_key}"}
use_fidelity = bool(ref_images) and s.rightapi_input_fidelity and _rightapi_fidelity_supported is not False
async with httpx.AsyncClient(timeout=max(s.request_timeout, 600), verify=False) as client:
common = {
"model": model,
"prompt": prompt,
"size": size,
"quality": s.rightapi_image_quality,
"output_format": "jpeg", # 与落盘 .jpg 后缀一致
"n": 1,
}
if use_fidelity:
common["input_fidelity"] = s.rightapi_input_fidelity
if ref_images:
files = []
for i, u in enumerate(ref_images):
data, mime = await _resolve_ref_bytes(u)
files.append(("image[]", (f"ref-{i + 1}.{mime.split('/')[-1]}", data, mime)))
resp = await client.post(f"{base}/v1/images/edits", headers=headers, files=files, data=common)
# 中转不认 input_fidelity:去掉参数重试一次(仅一次探测)
if resp.status_code == 400 and use_fidelity and "input_fidelity" in resp.text:
_rightapi_fidelity_supported = False
log.warning("RightAPI 不支持 input_fidelity 参数,已自动去掉并降级(后续请求不再带)")
common.pop("input_fidelity", None)
resp = await client.post(f"{base}/v1/images/edits", headers=headers, files=files, data=common)
elif resp.is_success and use_fidelity:
_rightapi_fidelity_supported = True
else:
resp = await client.post(
f"{base}/v1/images/generations",
headers={**headers, "Content-Type": "application/json"},
json=common,
)
_raise_api_error(resp, "RightAPI")
item = resp.json()["data"][0]
b64 = item.get("b64_json") or ""
if b64:
return base64.b64decode(b64.split(",", 1)[-1] if "," in b64 else b64)
img_url = item.get("url") or ""
if not img_url:
raise RuntimeError(f"RightAPI 响应里没有图片数据: {item}")
dl = await client.get(img_url, timeout=s.request_timeout)
dl.raise_for_status()
return dl.content
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]
GENERATORS = {"doubao": generate_doubao, "tongyi": generate_tongyi, "rightapi": generate_rightapi}
# ── 任务执行器 ────────────────────────────────────────────────────────────
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", ""))
async def _select_ref_images(db, product_id: UUID, type_id: str) -> list[str]:
"""商品路径:主图组前几张。转存完成的用本地文件,未完成的直接用源站 URL。"""
assets = (await db.scalars(
select(ProductAsset).where(
ProductAsset.product_id == product_id,
ProductAsset.group_key == "main",
ProductAsset.type == "img",
).order_by(ProductAsset.sort_order)
)).all()
refs = [a.stored_url or a.source_url for a in assets if (a.stored_url or a.source_url)]
if not refs:
raise RuntimeError("商品没有可用参考图(未采集主图)")
return _order_refs(refs, type_id)
async def run_suite(suite_id: str) -> None:
"""后台执行套图任务:逐张生成 → 落盘 → 记录;单张失败不中断。
两条路径:
- 无状态(product_id 为空):上下文与参考图来自请求自带的 context / ref_images
- 商品路径(兼容旧流程):从 product + product_assets 取
"""
settings = get_settings()
async with get_session_factory()() as db:
suite = await db.get(Suite, UUID(suite_id))
if suite is None:
return
product = None
if suite.product_id:
product = await db.get(Product, suite.product_id)
if product is None:
suite.status = SUITE_FAILED
suite.error = "商品不存在"
await db.commit()
return
suite.status = SUITE_RUNNING
await db.commit()
provider_name = suite.provider or settings.image_provider
generator = GENERATORS.get(provider_name)
if generator is None:
suite.status = SUITE_FAILED
suite.error = f"未知 provider: {provider_name}"
await db.commit()
return
raw = suite.context if not product else (product.raw or {})
ctx = build_context(raw or {}, fallback_name=product.name if product else "")
model = suite.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, suite.ratio, is_wan=is_wan, model=model)
# 任务列表:方案(逐张)优先,旧路径按 types
if suite.plan:
jobs = [dict(j) for j in suite.plan]
else:
jobs = [
{"kind": t, "title": type_name(t), "detail": "", "prompt_hint": "", "variant_name": None}
for t in (suite.types or [])
]
ok, failed = 0, 0
failures: list[str] = []
for job in jobs:
type_id = job["kind"]
image_row = SuiteImage(
suite_id=suite.id,
type_id=type_id,
name=job.get("title") or type_name(type_id),
status=STATUS_FAILED,
)
db.add(image_row)
await db.flush()
try:
prompt = build_prompt(
type_id, ctx, suite.style_set, suite.lang,
extra=job, style_prompt=suite.style_prompt, requirements=suite.requirements,
)
# gpt-image edits 语义:商品冻结契约前置,防止风格词改商品
if provider_name == "rightapi":
prompt = wrap_prompt_for_gpt_edits(prompt)
if product:
refs = await _select_ref_images(db, product.id, type_id)
else:
refs = _refs_for_job(list(suite.ref_images or []), job)
data = await generator(prompt, refs, size=size, model=model)
# 部分中转不遵守 output_format(要 jpeg 回 PNG),按魔数定扩展名
ext = ".png" if data[:8] == b"\x89PNG\r\n\x1a\n" else ".jpg"
key = storage.write_bytes(data, key_prefix=f"suites/{suite.id}", ext=ext)
image_row.stored_url = storage.public_url(key)
image_row.status = STATUS_OK
ok += 1
except Exception as exc: # noqa: BLE001
log.exception("套图 %s 类型 %s 生成失败", suite_id, type_id)
err = str(exc)[:500]
image_row.error = err
failures.append(f"{job.get('title') or type_name(type_id)}{err[:200]}")
failed += 1
await db.commit()
suite.status = SUITE_DONE if failed == 0 else (SUITE_PARTIAL if ok > 0 else SUITE_FAILED)
if failed:
uniq = list(dict.fromkeys(failures)) # 去重保序
detail = "".join(uniq[:6])
if len(uniq) > 6:
detail += f";…等共 {failed} 张失败"
if ok == 0:
suite.error = f"全部生成失败。{detail}"
else:
suite.error = f"部分生成失败({failed} 张)。{detail}"
from datetime import datetime, timezone
suite.finished_at = datetime.now(timezone.utc)
if product:
product.stage = "generated" # 商品路径才有的阶段升级
await db.commit()