feat: 初始化项目,并且接近完成 ozon 部分
This commit is contained in:
@@ -0,0 +1,316 @@
|
||||
"""套图生成服务:图像 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
|
||||
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
|
||||
|
||||
log = logging.getLogger("suite.generator")
|
||||
|
||||
# 参考图选择: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) -> str:
|
||||
"""平台比例 → provider 尺寸参数。3:4 竖版(Ozon/WB),1:1 方图(国内)。"""
|
||||
if provider == "doubao":
|
||||
return "1536x2048" if ratio == "3:4" else "2048x2048"
|
||||
# tongyi:万象与千问的 size 语法相同(* 分隔),档位不同
|
||||
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(url: str) -> str:
|
||||
"""参考图 URL → data URI。本地 media 文件直读磁盘;远程 URL 带 Referer 下载。
|
||||
|
||||
生图 API 的服务器无法访问 127.0.0.1,代理 URL 也不能直接透传,
|
||||
所以统一在本地解析成 base64 data URI 再进请求体。
|
||||
"""
|
||||
if url.startswith("data:"):
|
||||
return url
|
||||
path = storage.local_path(url)
|
||||
if path is not None:
|
||||
mime = mimetypes.guess_type(path.name)[0] or "image/jpeg"
|
||||
return _bytes_to_data_uri(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 _bytes_to_data_uri(data, ctype)
|
||||
raise FileNotFoundError(f"无法解析参考图: {url}")
|
||||
|
||||
|
||||
# ── Provider:豆包 Seedream(火山方舟)────────────────────────────────────
|
||||
|
||||
async def generate_doubao(prompt: str, ref_images: list[str], size: str = "2048x2048") -> bytes:
|
||||
s = get_settings()
|
||||
if not s.ark_api_key:
|
||||
raise RuntimeError("未配置 ARK_API_KEY(.env)")
|
||||
body = {
|
||||
"model": 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,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
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")
|
||||
|
||||
|
||||
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") -> bytes:
|
||||
s = get_settings()
|
||||
if not s.dashscope_api_key:
|
||||
raise RuntimeError("未配置 DASHSCOPE_API_KEY(.env)")
|
||||
is_wan = _is_wan_model(s.dashscope_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"
|
||||
)
|
||||
|
||||
content: list[dict] = [{"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": s.dashscope_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)
|
||||
resp.raise_for_status()
|
||||
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
|
||||
|
||||
|
||||
GENERATORS = {"doubao": generate_doubao, "tongyi": generate_tongyi}
|
||||
|
||||
|
||||
# ── 任务执行器 ────────────────────────────────────────────────────────────
|
||||
|
||||
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 "")
|
||||
size = _image_size(provider_name, suite.ratio, is_wan=_is_wan_model(settings.dashscope_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
|
||||
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)
|
||||
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)
|
||||
key = storage.write_bytes(data, key_prefix=f"suites/{suite.id}", ext=".jpg")
|
||||
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)
|
||||
image_row.error = str(exc)[:500]
|
||||
failed += 1
|
||||
await db.commit()
|
||||
|
||||
suite.status = SUITE_DONE if failed == 0 else (SUITE_PARTIAL if ok > 0 else SUITE_FAILED)
|
||||
if failed and not ok:
|
||||
suite.error = "全部生成失败,请检查 API Key / 参考图"
|
||||
from datetime import datetime, timezone
|
||||
suite.finished_at = datetime.now(timezone.utc)
|
||||
if product:
|
||||
product.stage = "generated" # 商品路径才有的阶段升级
|
||||
await db.commit()
|
||||
Reference in New Issue
Block a user