190 lines
6.5 KiB
Python
190 lines
6.5 KiB
Python
"""3D Mesh 生成服务 — 阶段 1：单图 → textured .glb + 360° 预渲染视频。
|
||
|
||
当前支持：
|
||
- Trellis (firtoz/trellis)：~30s，$0.041/次，自带 color_video + normal_video + .glb
|
||
- Hunyuan3D-2 (tencent/hunyuan3d-2)：~127s，$0.12/次，仅 .glb 无预渲染视频（二期启用）
|
||
"""
|
||
|
||
import logging
|
||
import os
|
||
import uuid
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import httpx
|
||
|
||
from app.config import get_view_transform_model_config
|
||
from app.services.image_gen import _make_replicate_client, GENERATED_DIR, to_data_uri
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
_replicate_client = _make_replicate_client()
|
||
|
||
SUPPORTED_EXTS = {
|
||
"model_file": ".glb",
|
||
"color_video": ".mp4",
|
||
"normal_video": ".mp4",
|
||
"combined_video": ".mp4",
|
||
"gaussian_ply": ".ply",
|
||
"mesh": ".glb",
|
||
}
|
||
|
||
|
||
async def _download_asset(url: str, ext: str, prefix: str) -> Path:
|
||
"""下载远程资源到 generated/ 目录，返回本地路径。"""
|
||
filename = f"{prefix}_{uuid.uuid4().hex[:12]}{ext}"
|
||
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(300.0)) as client:
|
||
resp = await client.get(url, follow_redirects=True)
|
||
resp.raise_for_status()
|
||
filepath.write_bytes(resp.content)
|
||
return filepath
|
||
|
||
|
||
def _to_url(value: Any) -> str | None:
|
||
"""把 Replicate FileOutput / str / None 统一转为 URL 字符串。"""
|
||
if value is None:
|
||
return None
|
||
s = str(value)
|
||
if s.startswith("http://") or s.startswith("https://"):
|
||
return s
|
||
return None
|
||
|
||
|
||
async def generate_with_trellis(
|
||
image_path: str,
|
||
params: dict[str, Any] | None = None,
|
||
) -> dict[str, Any]:
|
||
"""调用 Trellis 做单图 3D 重建，下载所有输出资产到本地。
|
||
|
||
Args:
|
||
image_path: 输入图（本地路径 / data URI / URL）
|
||
params: 覆盖 default_params 的可选参数字典
|
||
|
||
Returns:
|
||
{
|
||
"success": bool,
|
||
"glb_path": "/generated/trellis_xxx.glb" | None,
|
||
"color_video_path": "/generated/trellis_xxx_color.mp4" | None,
|
||
"normal_video_path": "/generated/trellis_xxx_normal.mp4" | None,
|
||
"combined_video_path": ...,
|
||
"raw_output_keys": [...], # 调试用
|
||
"error": str | None,
|
||
}
|
||
"""
|
||
config = get_view_transform_model_config("trellis")
|
||
replicate_model_id = config["model_id"]
|
||
default_params = dict(config.get("default_params", {}))
|
||
if params:
|
||
default_params.update(params)
|
||
|
||
GENERATED_DIR.mkdir(parents=True, exist_ok=True)
|
||
|
||
try:
|
||
image_uri = to_data_uri(image_path)
|
||
# firtoz/trellis:e8f6c452... 这个部署版的 schema：
|
||
# - 输入必须用 `images` 数组（传 `image` 单数会 422）
|
||
# - 输出字段 `color_video` 和 `normal_video` 实测恒为 null
|
||
# - 实际可用的视频只有 `combined_video`（color+normal 左右并排）
|
||
# 因此 Phase 2 抽帧时从 combined_video 切左半得 color、切右半得 normal。
|
||
input_params: dict[str, Any] = {
|
||
"images": [image_uri],
|
||
"randomize_seed": False,
|
||
"seed": 0,
|
||
**default_params,
|
||
}
|
||
|
||
logger.info(
|
||
"调用 Trellis 进行 3D 重建：model=%s, params=%s",
|
||
replicate_model_id,
|
||
{k: v for k, v in input_params.items() if k != "images"},
|
||
)
|
||
output = await _replicate_client.async_run(
|
||
replicate_model_id, input=input_params, wait=False
|
||
)
|
||
|
||
if not isinstance(output, dict):
|
||
return {
|
||
"success": False,
|
||
"error": f"Trellis 返回非字典输出：{type(output).__name__} / {str(output)[:200]}",
|
||
"raw_output_keys": [],
|
||
}
|
||
|
||
raw_keys = list(output.keys())
|
||
logger.info("Trellis 返回字段：%s", raw_keys)
|
||
|
||
# 按已知字段挨个下载
|
||
result: dict[str, Any] = {
|
||
"success": True,
|
||
"raw_output_keys": raw_keys,
|
||
"error": None,
|
||
}
|
||
download_map = [
|
||
("model_file", "glb_path", ".glb", "trellis"),
|
||
("color_video", "color_video_path", ".mp4", "trellis_color"),
|
||
("normal_video", "normal_video_path", ".mp4", "trellis_normal"),
|
||
("combined_video", "combined_video_path", ".mp4", "trellis_combined"),
|
||
]
|
||
|
||
for src_key, dst_key, ext, prefix in download_map:
|
||
if src_key not in output:
|
||
result[dst_key] = None
|
||
continue
|
||
url = _to_url(output[src_key])
|
||
if not url:
|
||
logger.warning("Trellis 字段 %s 不是有效 URL：%s", src_key, output[src_key])
|
||
result[dst_key] = None
|
||
continue
|
||
try:
|
||
local_path = await _download_asset(url, ext, prefix)
|
||
result[dst_key] = f"/generated/{local_path.name}"
|
||
logger.info("已下载 %s → %s", src_key, local_path.name)
|
||
except Exception as e:
|
||
logger.error("下载 %s 失败：%s", src_key, e)
|
||
result[dst_key] = None
|
||
|
||
# 阶段 2 抽帧的优先级：color_video > combined_video（切左半）
|
||
# 只要其中一个可用即算成功
|
||
has_color_source = bool(result.get("color_video_path")) or bool(
|
||
result.get("combined_video_path")
|
||
)
|
||
if not has_color_source:
|
||
result["success"] = False
|
||
result["error"] = (
|
||
"Trellis 未返回 color_video 或 combined_video，无法进入阶段 2 抽帧；"
|
||
f"实际返回字段：{raw_keys}"
|
||
)
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
detail = str(e) or f"{type(e).__name__}: {repr(e)}"
|
||
logger.error("Trellis 调用失败：%s", detail, exc_info=True)
|
||
return {
|
||
"success": False,
|
||
"error": detail,
|
||
"raw_output_keys": [],
|
||
}
|
||
|
||
|
||
async def generate_mesh(
|
||
image_path: str,
|
||
model_id: str = "trellis",
|
||
params: dict[str, Any] | None = None,
|
||
) -> dict[str, Any]:
|
||
"""统一入口：按 model_id 分发到具体实现。"""
|
||
if model_id == "trellis":
|
||
return await generate_with_trellis(image_path, params)
|
||
if model_id == "hunyuan3d":
|
||
return {
|
||
"success": False,
|
||
"error": "Hunyuan3D-2 (一期未启用，待二期接入自定义渲染)",
|
||
"raw_output_keys": [],
|
||
}
|
||
return {
|
||
"success": False,
|
||
"error": f"未知 mesh 模型：{model_id}",
|
||
"raw_output_keys": [],
|
||
}
|