358 lines
13 KiB
Python
358 lines
13 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
虚拟换装网关 —— 统一 FastAPI 服务（HTTPS / 443）
|
||
|
||
路由规则（与原 change_app.py 保持一致）：
|
||
kuzi_img 存在 → service3 逻辑（模特 + 上衣 + 裤子）
|
||
cloth_img 存在（无裤子）→ service2 逻辑（模特 + 上衣）
|
||
仅 human_img → service1 逻辑（单图编辑）
|
||
"""
|
||
|
||
import asyncio
|
||
import base64
|
||
import io
|
||
import json
|
||
import os
|
||
import tempfile
|
||
import time
|
||
import uuid
|
||
from pathlib import Path
|
||
|
||
import httpx
|
||
import oss2
|
||
import uvicorn
|
||
from fastapi import FastAPI, HTTPException, Request
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from fastapi.responses import FileResponse, JSONResponse
|
||
from fastapi.staticfiles import StaticFiles
|
||
from pydantic import BaseModel
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# ComfyUI 配置
|
||
# ---------------------------------------------------------------------------
|
||
COMFYUI_URL = "http://112.126.94.241:28188"
|
||
COMFYUI_URL_BACKUP = "http://112.126.94.241:38188"
|
||
|
||
_root = Path(__file__).parent
|
||
COMFYUI_USER = (_root / "user.txt").read_text(encoding="utf-8").strip()
|
||
COMFYUI_PASS = (_root / "password.txt").read_text(encoding="utf-8").strip()
|
||
COMFYUI_AUTH = (COMFYUI_USER, COMFYUI_PASS)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 工作流路径
|
||
# ---------------------------------------------------------------------------
|
||
WORKFLOW1 = _root / "change1" / "change1.json" # 单图编辑
|
||
WORKFLOW2 = _root / "change2" / "change2_1_ex.json" # 模特 + 上衣
|
||
WORKFLOW3 = _root / "change3" / "0604-change_fast3.json" # 模特 + 上衣 + 裤子
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 工作流节点 ID
|
||
# ---------------------------------------------------------------------------
|
||
NODE_MODEL = "11" # 模特图（LoadImage）
|
||
NODE_SHIRT = "10" # 上衣图（LoadImage）
|
||
NODE_PANTS = "4" # 裤子图（LoadImage）
|
||
NODE_OUTPUT = "33" # 输出（SaveImage）
|
||
NODE_PROMPT = "31" # 文本提示（PrimitiveStringMultiline）
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# OSS 配置
|
||
# ---------------------------------------------------------------------------
|
||
OSS_KEY_ID = "LTAI5tGp1sLzedqxihcNC1eb"
|
||
OSS_KEY_SECRET = "IFZE1b8YYreCP6zfA6GaZ9uBT678qO"
|
||
OSS_ENDPOINT = "oss-cn-beijing.aliyuncs.com"
|
||
OSS_BUCKET = "xiangsilian"
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# HTTPS 证书路径
|
||
# ---------------------------------------------------------------------------
|
||
SSL_CERT = "/etc/letsencrypt/live/xiangsilian.com/fullchain.pem"
|
||
SSL_KEY = "/etc/letsencrypt/live/xiangsilian.com/privkey.pem"
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# FastAPI 应用
|
||
# ---------------------------------------------------------------------------
|
||
app = FastAPI(title="虚拟换装服务")
|
||
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
|
||
@app.exception_handler(HTTPException)
|
||
async def http_exception_handler(request: Request, exc: HTTPException):
|
||
"""将 HTTPException 统一转换为业务错误格式"""
|
||
return JSONResponse(
|
||
status_code=exc.status_code,
|
||
content={"ret": -1, "state": -1, "msg": exc.detail},
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# OSS 工具
|
||
# ---------------------------------------------------------------------------
|
||
def upload_to_oss(image_path: str, object_name: str) -> str | None:
|
||
auth = oss2.Auth(OSS_KEY_ID, OSS_KEY_SECRET)
|
||
bucket = oss2.Bucket(auth, OSS_ENDPOINT, OSS_BUCKET)
|
||
try:
|
||
bucket.put_object_from_file(object_name, image_path)
|
||
url = bucket.sign_url("GET", object_name, 3600 * 24 * 365 * 10)
|
||
print(f"OSS 上传成功: {url}")
|
||
return url
|
||
except Exception as e:
|
||
print(f"OSS 上传失败: {e}")
|
||
return None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# ComfyUI 工具函数
|
||
# ---------------------------------------------------------------------------
|
||
def decode_base64_image(b64_str: str) -> bytes:
|
||
"""解码 base64 图片，自动去除 data:image/... 前缀"""
|
||
if "," in b64_str:
|
||
b64_str = b64_str.split(",", 1)[1]
|
||
return base64.b64decode(b64_str)
|
||
|
||
|
||
async def check_comfyui_alive(base_url: str, timeout: float = 3.0) -> bool:
|
||
try:
|
||
async with httpx.AsyncClient(timeout=timeout, auth=COMFYUI_AUTH) as client:
|
||
resp = await client.get(f"{base_url}/system_stats")
|
||
return resp.status_code < 500
|
||
except Exception as e:
|
||
print(f"ComfyUI 健康检查失败 {base_url}: {e}")
|
||
return False
|
||
|
||
|
||
async def pick_comfyui_url() -> str:
|
||
if await check_comfyui_alive(COMFYUI_URL):
|
||
print(f"使用正式服务器: {COMFYUI_URL}")
|
||
return COMFYUI_URL
|
||
print(f"正式服务器不可用，切换备份: {COMFYUI_URL_BACKUP}")
|
||
if await check_comfyui_alive(COMFYUI_URL_BACKUP):
|
||
print(f"使用备份服务器: {COMFYUI_URL_BACKUP}")
|
||
return COMFYUI_URL_BACKUP
|
||
raise HTTPException(status_code=503, detail="正式与备份 ComfyUI 服务器均不可用")
|
||
|
||
|
||
async def upload_image(client: httpx.AsyncClient, base_url: str, image_bytes: bytes, filename: str) -> str:
|
||
files = {"image": (filename, io.BytesIO(image_bytes), "image/jpeg")}
|
||
resp = await client.post(f"{base_url}/upload/image", files=files, data={"overwrite": "true"})
|
||
resp.raise_for_status()
|
||
return resp.json()["name"]
|
||
|
||
|
||
async def queue_prompt(client: httpx.AsyncClient, base_url: str, workflow: dict) -> str:
|
||
payload = {"prompt": workflow, "client_id": str(uuid.uuid4())}
|
||
resp = await client.post(f"{base_url}/prompt", json=payload)
|
||
resp.raise_for_status()
|
||
return resp.json()["prompt_id"]
|
||
|
||
|
||
async def wait_for_result(
|
||
client: httpx.AsyncClient, base_url: str, prompt_id: str, timeout: int = 300
|
||
) -> dict:
|
||
deadline = time.time() + timeout
|
||
while time.time() < deadline:
|
||
resp = await client.get(f"{base_url}/history/{prompt_id}")
|
||
resp.raise_for_status()
|
||
history = resp.json()
|
||
if prompt_id in history:
|
||
entry = history[prompt_id]
|
||
status = entry.get("status", {})
|
||
if status.get("completed"):
|
||
return entry.get("outputs", {})
|
||
if status.get("status_str") == "error":
|
||
messages = status.get("messages", [])
|
||
raise HTTPException(status_code=500, detail=f"ComfyUI 工作流执行出错: {messages}")
|
||
await asyncio.sleep(2)
|
||
raise HTTPException(status_code=504, detail="等待 ComfyUI 超时（300秒）")
|
||
|
||
|
||
async def fetch_image_bytes(
|
||
client: httpx.AsyncClient, base_url: str, filename: str,
|
||
subfolder: str = "", type_: str = "output"
|
||
) -> bytes:
|
||
params = {"filename": filename, "subfolder": subfolder, "type": type_}
|
||
resp = await client.get(f"{base_url}/view", params=params)
|
||
resp.raise_for_status()
|
||
return resp.content
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 通用工作流执行
|
||
# ---------------------------------------------------------------------------
|
||
async def run_workflow(
|
||
workflow_path: Path,
|
||
node_images: dict[str, bytes],
|
||
change_desc: str = "",
|
||
) -> str:
|
||
"""
|
||
执行 ComfyUI 工作流并返回 OSS URL。
|
||
|
||
node_images: {节点ID: 图片字节}，按顺序并发上传并注入工作流。
|
||
"""
|
||
comfyui_url = await pick_comfyui_url()
|
||
workflow = json.loads(workflow_path.read_text(encoding="utf-8"))
|
||
|
||
desc = (change_desc or "").strip()
|
||
if desc:
|
||
workflow[NODE_PROMPT]["inputs"] = {"value": desc}
|
||
|
||
uid = uuid.uuid4().hex[:8]
|
||
|
||
async with httpx.AsyncClient(timeout=60.0, auth=COMFYUI_AUTH) as client:
|
||
node_ids = list(node_images.keys())
|
||
upload_tasks = [
|
||
upload_image(client, comfyui_url, node_images[nid], f"{nid}_{uid}.jpg")
|
||
for nid in node_ids
|
||
]
|
||
try:
|
||
names = await asyncio.gather(*upload_tasks)
|
||
except httpx.HTTPError as e:
|
||
raise HTTPException(status_code=502, detail=f"上传图片到 ComfyUI 失败: {e}")
|
||
|
||
for node_id, name in zip(node_ids, names):
|
||
workflow[node_id]["inputs"]["image"] = name
|
||
|
||
try:
|
||
prompt_id = await queue_prompt(client, comfyui_url, workflow)
|
||
except httpx.HTTPError as e:
|
||
raise HTTPException(status_code=502, detail=f"提交工作流失败: {e}")
|
||
|
||
async with httpx.AsyncClient(timeout=30.0, auth=COMFYUI_AUTH) as poll_client:
|
||
outputs = await wait_for_result(poll_client, comfyui_url, prompt_id, timeout=300)
|
||
|
||
node_output = outputs.get(NODE_OUTPUT)
|
||
if not node_output:
|
||
raise HTTPException(status_code=500, detail=f"工作流未返回节点 {NODE_OUTPUT} 的输出")
|
||
|
||
images = node_output.get("images", [])
|
||
if not images:
|
||
raise HTTPException(status_code=500, detail="输出节点没有图片")
|
||
|
||
img_info = images[0]
|
||
filename = img_info["filename"]
|
||
subfolder = img_info.get("subfolder", "")
|
||
type_ = img_info.get("type", "output")
|
||
|
||
async with httpx.AsyncClient(timeout=30.0, auth=COMFYUI_AUTH) as dl_client:
|
||
result_bytes = await fetch_image_bytes(dl_client, comfyui_url, filename, subfolder, type_)
|
||
|
||
suffix = Path(filename).suffix or ".png"
|
||
object_name = f"tryon/{uuid.uuid4().hex}{suffix}"
|
||
with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
|
||
tmp.write(result_bytes)
|
||
tmp_path = tmp.name
|
||
|
||
try:
|
||
result_url = upload_to_oss(tmp_path, object_name)
|
||
finally:
|
||
os.unlink(tmp_path)
|
||
|
||
if not result_url:
|
||
raise HTTPException(status_code=500, detail="上传结果图片到 OSS 失败")
|
||
|
||
return result_url
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# API 模型
|
||
# ---------------------------------------------------------------------------
|
||
class ChangeClothRequest(BaseModel):
|
||
human_img: str
|
||
cloth_img: str = ""
|
||
kuzi_img: str = ""
|
||
change_desc: str = ""
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 接口
|
||
# ---------------------------------------------------------------------------
|
||
@app.post("/change_cloth_base64")
|
||
async def change_cloth_base64(req: ChangeClothRequest):
|
||
print(f"收到请求，字段: human_img={'有' if req.human_img else '无'}, "
|
||
f"cloth_img={'有' if req.cloth_img else '无'}, "
|
||
f"kuzi_img={'有' if req.kuzi_img else '无'}")
|
||
|
||
if not req.human_img:
|
||
raise HTTPException(status_code=400, detail="human_img is required")
|
||
|
||
try:
|
||
model_bytes = decode_base64_image(req.human_img)
|
||
except Exception as e:
|
||
raise HTTPException(status_code=400, detail=f"human_img 解码失败: {e}")
|
||
|
||
if req.kuzi_img:
|
||
try:
|
||
shirt_bytes = decode_base64_image(req.cloth_img)
|
||
pants_bytes = decode_base64_image(req.kuzi_img)
|
||
except Exception as e:
|
||
raise HTTPException(status_code=400, detail=f"图片解码失败: {e}")
|
||
result_url = await run_workflow(
|
||
WORKFLOW3,
|
||
{NODE_MODEL: model_bytes, NODE_SHIRT: shirt_bytes, NODE_PANTS: pants_bytes},
|
||
req.change_desc,
|
||
)
|
||
elif req.cloth_img:
|
||
try:
|
||
shirt_bytes = decode_base64_image(req.cloth_img)
|
||
except Exception as e:
|
||
raise HTTPException(status_code=400, detail=f"cloth_img 解码失败: {e}")
|
||
result_url = await run_workflow(
|
||
WORKFLOW2,
|
||
{NODE_MODEL: model_bytes, NODE_SHIRT: shirt_bytes},
|
||
req.change_desc,
|
||
)
|
||
else:
|
||
result_url = await run_workflow(
|
||
WORKFLOW1,
|
||
{NODE_MODEL: model_bytes},
|
||
req.change_desc,
|
||
)
|
||
|
||
return {"ret": 0, "state": 0, "msg": "success", "data": result_url}
|
||
|
||
|
||
@app.get("/")
|
||
async def root():
|
||
return FileResponse(str(_root / "static" / "index.html"))
|
||
|
||
|
||
@app.get("/health")
|
||
async def health():
|
||
primary_ok = await check_comfyui_alive(COMFYUI_URL)
|
||
backup_ok = await check_comfyui_alive(COMFYUI_URL_BACKUP)
|
||
active = COMFYUI_URL if primary_ok else (COMFYUI_URL_BACKUP if backup_ok else None)
|
||
return {
|
||
"status": "ok" if active else "degraded",
|
||
"comfyui_active": active,
|
||
"comfyui_primary": {"url": COMFYUI_URL, "alive": primary_ok},
|
||
"comfyui_backup": {"url": COMFYUI_URL_BACKUP, "alive": backup_ok},
|
||
}
|
||
|
||
|
||
app.mount("/static", StaticFiles(directory=str(_root / "static")), name="static")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 启动入口
|
||
# ---------------------------------------------------------------------------
|
||
if __name__ == "__main__":
|
||
use_https = os.path.exists(SSL_CERT) and os.path.exists(SSL_KEY)
|
||
|
||
if use_https:
|
||
print(f"启动 HTTPS 服务，端口 443，证书: {SSL_CERT}")
|
||
uvicorn.run(
|
||
app,
|
||
host="0.0.0.0",
|
||
port=443,
|
||
ssl_certfile=SSL_CERT,
|
||
ssl_keyfile=SSL_KEY,
|
||
log_level="info",
|
||
)
|
||
else:
|
||
print("⚠ 证书文件不存在，降级为 HTTP，端口 28888")
|
||
uvicorn.run(app, host="0.0.0.0", port=28888, log_level="info")
|