181 lines
6.0 KiB
Python
181 lines
6.0 KiB
Python
"""阶段 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])
|