引入3d模型

art-agent/backend/app/services/style_restorer.py

@@ -0,0 +1,180 @@
+"""阶段 3：对 Trellis 抽帧结果做 IP-Adapter + ControlNet 重绘，还原原画风。
+
+输入：
+- original_image：用户上传的原图（提供风格 / IP-Adapter 参考）
+- structure_image：Stage 2 抽出的 normal 帧（提供几何 / ControlNet 结构约束）
+
+输出：
+- 单张保留原画风且视角正确的 PNG，落盘到 GENERATED_DIR/ 下
+
+模型：chigozienri/ip_adapter-sdxl-controlnet-depth（配置见 config.IMAGE_MODELS）
+- $0.07/张，~72s，硬件 L40S
+- `image` 字段 = IP-Adapter 参考图（风格源）
+- `controlnet_input` 字段 = 结构参考图（当前一期传 normal 帧作为 depth 近似）
+
+调参建议（可在 IMAGE_MODELS 默认值基础上覆盖）：
+- scale (IP-Adapter 权重): 0.7-0.8，越高越贴近原图风格
+- controlnet_conditioning_scale (结构权重): 0.7-0.9，越高越保几何
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import uuid
+from pathlib import Path
+from typing import Any
+
+import httpx
+
+from app.config import get_image_model_config, get_style_restore_model_id
+from app.services.image_gen import GENERATED_DIR, _make_replicate_client, to_data_uri
+
+logger = logging.getLogger(__name__)
+
+_replicate_client = _make_replicate_client()
+
+
+async def _download_image(url: str, prefix: str) -> Path:
+    """把 Replicate 返回的图片下载到 generated/。"""
+    filename = f"{prefix}_{uuid.uuid4().hex[:12]}.png"
+    filepath = GENERATED_DIR / filename
+    proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("HTTP_PROXY")
+    async with httpx.AsyncClient(proxy=proxy, timeout=httpx.Timeout(180.0)) as client:
+        resp = await client.get(url, follow_redirects=True)
+        resp.raise_for_status()
+        filepath.write_bytes(resp.content)
+    return filepath
+
+
+def _first_url(output: Any) -> str | None:
+    """Replicate 输出有时是 list[FileOutput]、有时是单个；统一取第一张。"""
+    if output is None:
+        return None
+    if isinstance(output, list):
+        if not output:
+            return None
+        return str(output[0])
+    return str(output)
+
+
+async def restore(
+    original_image: str,
+    structure_image: str,
+    azimuth: int = 0,
+    elevation: int = 0,
+    params: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """对单张 normal 帧做 IP-Adapter + ControlNet 重绘。
+
+    Args:
+        original_image: 用户原图，本地路径 / data URI / URL 均可
+        structure_image: Stage 2 抽出的 normal 帧，同上
+        azimuth / elevation: 仅用于日志与返回元数据，不参与模型调用
+        params: 覆盖默认参数的字典（scale / controlnet_conditioning_scale / prompt 等）
+
+    Returns:
+        {
+          "success": bool,
+          "url": "/generated/restyled_xxx.png" | None,
+          "azimuth": int,
+          "elevation": int,
+          "error": str | None,
+        }
+    """
+    config = get_image_model_config(get_style_restore_model_id())
+    replicate_model_id = config["model_id"]
+    merged_params = dict(config.get("default_params", {}))
+    if params:
+        merged_params.update(params)
+
+    GENERATED_DIR.mkdir(parents=True, exist_ok=True)
+
+    try:
+        image_uri = to_data_uri(original_image)
+        structure_uri = to_data_uri(structure_image)
+
+        input_params: dict[str, Any] = {
+            "image": image_uri,                      # IP-Adapter 参考图 = 用户原图
+            "controlnet_input": structure_uri,       # 结构参考 = normal 帧
+            **merged_params,
+        }
+
+        logger.info(
+            "Stage 3 重绘 az=%d el=%d: scale=%s, cn_scale=%s",
+            azimuth, elevation,
+            merged_params.get("scale"),
+            merged_params.get("controlnet_conditioning_scale"),
+        )
+
+        output = await _replicate_client.async_run(
+            replicate_model_id, input=input_params, wait=False
+        )
+
+        url = _first_url(output)
+        if not url or not (url.startswith("http://") or url.startswith("https://")):
+            return {
+                "success": False,
+                "url": None,
+                "azimuth": azimuth,
+                "elevation": elevation,
+                "error": f"模型未返回有效图片 URL: {str(output)[:200]}",
+            }
+
+        local = await _download_image(url, prefix=f"restyled_az{azimuth}_el{elevation}")
+        return {
+            "success": True,
+            "url": f"/generated/{local.name}",
+            "azimuth": azimuth,
+            "elevation": elevation,
+            "error": None,
+        }
+
+    except Exception as e:
+        detail = str(e) or f"{type(e).__name__}: {repr(e)}"
+        logger.error(
+            "Stage 3 重绘失败 az=%d el=%d: %s", azimuth, elevation, detail, exc_info=True
+        )
+        return {
+            "success": False,
+            "url": None,
+            "azimuth": azimuth,
+            "elevation": elevation,
+            "error": detail,
+        }
+
+
+async def restore_batch(
+    original_image: str,
+    frames: list[dict[str, Any]],
+    params: dict[str, Any] | None = None,
+    concurrency: int = 3,
+) -> list[dict[str, Any]]:
+    """对一批 frames 做 Stage 3 重绘，带并发控制。
+
+    Args:
+        original_image: 用户原图（所有帧共用同一张风格参考）
+        frames: [{"az", "el", "normal_png", ...}, ...]，normal_png 作为结构参考
+        params: 透传给 restore() 的覆盖参数
+        concurrency: 最大并发数；Replicate 并发打太猛容易撞 rate limit，3 比较稳
+
+    Returns:
+        列表，顺序与 frames 对齐，元素与 restore() 的返回一致。
+    """
+    if not frames:
+        return []
+
+    sem = asyncio.Semaphore(max(1, concurrency))
+
+    async def _one(f: dict[str, Any]) -> dict[str, Any]:
+        async with sem:
+            return await restore(
+                original_image=original_image,
+                structure_image=f["normal_png"],
+                azimuth=int(f.get("az", 0)),
+                elevation=int(f.get("el", 0)),
+                params=params,
+            )
+
+    return await asyncio.gather(*[_one(f) for f in frames])