ozon-seller-kit/server/services/image_edit.py

"""多模型图生图 / 图像编辑服务。

按模型路由到不同的 DashScope 调用形态：
- wanx2.1-imageedit                     → ImageSynthesis（function=description_edit）
- wan2.6-image / qwen-image-edit(-plus) → MultiModalConversation（多模态生成）

职责单一：把请求翻译成 DashScope 调用并返回结果。当前不落盘、不存储图片；
后续接入七牛云时，可在返回结果 URL 之后增加「下载并转存」步骤，而不改 api 层与前端契约。
"""
from __future__ import annotations

import asyncio
import base64
from http import HTTPStatus

import httpx
from fastapi import HTTPException

from config import get_settings
from schemas.image_edit import ImageEditRequest, ImageEditResponse, ImageEditResult

# 服务端代理下载结果图时的体积上限，超出则跳过 base64（前端回退用 url）
_MAX_PROXY_BYTES = 20 * 1024 * 1024

# 走 ImageSynthesis 的模型（其余统一走 MultiModalConversation）
_WANX_IMAGEEDIT = "wanx2.1-imageedit"


def _download_to_data_url(url: str) -> str:
    """把生成结果图下载并转成 data URL。

    目的：阿里云结果 URL 未开放 CORS，前端直接绘制到 canvas 会导致画布被污染、
    无法 toDataURL 导出；由服务端代下载可彻底规避。失败时返回空串，前端回退用 url。
    """
    try:
        with httpx.Client(timeout=60.0, follow_redirects=True) as client:
            resp = client.get(url)
            resp.raise_for_status()
        content_type = (resp.headers.get("content-type") or "image/png").split(";")[0].strip()
        if not content_type.startswith("image/"):
            content_type = "image/png"
        payload = resp.content
        if len(payload) > _MAX_PROXY_BYTES:
            return ""
        b64 = base64.b64encode(payload).decode("ascii")
        return f"data:{content_type};base64,{b64}"
    except Exception:  # noqa: BLE001 - 下载失败不影响主流程
        return ""


def _apply_base_url(settings) -> None:
    """仅华北2（北京）业务空间需要配置 base_http_api_url。"""
    import dashscope

    url = (settings.dashscope_base_http_api_url or "").strip()
    if url:
        dashscope.base_http_api_url = url


def _fail(rsp, label: str) -> None:
    code = getattr(rsp, "code", None)
    message = getattr(rsp, "message", "")
    status = getattr(rsp, "status_code", "?")
    raise HTTPException(status_code=502, detail=f"{label}失败（HTTP {status}，code={code}）：{message}")


def _to_results(urls: list[str]) -> list[ImageEditResult]:
    return [ImageEditResult(url=u, image_base64=_download_to_data_url(u)) for u in urls]


def _call_wanx(req: ImageEditRequest, api_key: str) -> ImageEditResponse:
    """wanx2.1-imageedit：ImageSynthesis 同步调用。"""
    from dashscope import ImageSynthesis

    kwargs: dict = {
        "api_key": api_key,
        "model": "wanx2.1-imageedit",
        "function": req.function,
        "prompt": req.prompt,
        "base_image_url": req.base_image,
        "n": req.n,
    }
    if req.mask_image:
        kwargs["mask_image_url"] = req.mask_image
    if req.size:
        kwargs["size"] = req.size
    if req.seed is not None:
        kwargs["seed"] = req.seed
    if req.style:
        kwargs["style"] = req.style
    if req.prompt_extend:
        kwargs["prompt_extend"] = True
    if req.strength is not None:
        # strength 走 SDK 的 **kwargs，最终落入请求体的 parameters
        kwargs["strength"] = req.strength

    try:
        rsp = ImageSynthesis.call(**kwargs)
    except Exception as exc:  # noqa: BLE001 - SDK 抛错类型不统一，统一转为 502
        raise HTTPException(status_code=502, detail=f"DashScope 调用异常：{exc}") from exc

    if rsp.status_code != HTTPStatus.OK:
        _fail(rsp, "图生图")

    urls: list[str] = []
    for item in getattr(rsp.output, "results", None) or []:
        url = item.get("url") if isinstance(item, dict) else getattr(item, "url", None)
        if url:
            urls.append(url)

    if not urls:
        raise HTTPException(status_code=502, detail="模型未返回结果图片 URL")

    usage = getattr(rsp, "usage", None)
    image_count = int(getattr(usage, "image_count", 0) or 0) if usage else len(urls)
    return ImageEditResponse(
        task_id=getattr(rsp.output, "task_id", ""),
        results=_to_results(urls),
        image_count=image_count or len(urls),
        request_id=getattr(rsp, "request_id", ""),
    )


def _call_multimodal(req: ImageEditRequest, api_key: str) -> ImageEditResponse:
    """wan2.6-image / qwen-image-edit(-plus[-快照])：MultiModalConversation。"""
    from dashscope import MultiModalConversation

    messages = [{"role": "user", "content": [{"image": req.base_image}, {"text": req.prompt}]}]
    kwargs: dict = {
        "api_key": api_key,
        "model": req.model,
        "messages": messages,
        "prompt_extend": req.prompt_extend,
    }
    if req.n:
        kwargs["n"] = req.n
    if req.size:
        kwargs["size"] = req.size

    try:
        rsp = MultiModalConversation.call(**kwargs)
    except Exception as exc:  # noqa: BLE001
        raise HTTPException(status_code=502, detail=f"DashScope 调用异常：{exc}") from exc

    if rsp.status_code != HTTPStatus.OK:
        _fail(rsp, "图生图")

    urls: list[str] = []
    try:
        content = rsp.output.choices[0].message.content
    except (AttributeError, IndexError, KeyError, TypeError):
        content = []
    if isinstance(content, list):
        for item in content:
            img = item.get("image") if isinstance(item, dict) else None
            if img:
                urls.append(img)

    if not urls:
        raise HTTPException(status_code=502, detail="模型未返回结果图片")

    return ImageEditResponse(
        task_id="",
        results=_to_results(urls),
        image_count=len(urls),
        request_id=getattr(rsp, "request_id", ""),
    )


def _call_dashscope(req: ImageEditRequest, api_key: str) -> ImageEditResponse:
    """同步调用 DashScope（在线程池中执行），按模型路由。"""
    if req.model == _WANX_IMAGEEDIT:
        return _call_wanx(req, api_key)
    return _call_multimodal(req, api_key)


async def edit_image(req: ImageEditRequest) -> ImageEditResponse:
    """入口：校验密钥后在独立线程中执行同步 SDK 调用。"""
    settings = get_settings()
    api_key = (settings.dashscope_api_key or "").strip()
    if not api_key:
        raise HTTPException(status_code=500, detail="未配置 DASHSCOPE_API_KEY，请写入 .env")

    _apply_base_url(settings)
    return await asyncio.to_thread(_call_dashscope, req, api_key)