save code

2026-03-12 09:48:49 +00:00
parent 49c0ac8d2a
commit 42f83d3619
4 changed files with 990 additions and 0 deletions
@@ -0,0 +1,203 @@
+#!/usr/bin/env python3
+"""
+ComfyUI 换装服务
+- 接受3个base64图片（模特、上衣、裤子）
+- 上传到 ComfyUI，运行工作流
+- 返回生成结果 base64 图片
+"""
+
+import asyncio
+import base64
+import io
+import json
+import time
+import uuid
+from pathlib import Path
+
+import httpx
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import FileResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
+
+COMFYUI_URL = "http://127.0.0.1:8188"
+# COMFYUI_URL = "http://117.50.80.187:47697"
+WORKFLOW_PATH = Path(__file__).parent / "change2_2_0308.json"
+
+# 节点 ID 映射
+NODE_MODEL = "11"   # 模特
+NODE_SHIRT = "10"   # 上衣
+NODE_PANTS = "4"    # 裤子
+NODE_OUTPUT = "33"  # 输出
+
+app = FastAPI(title="ComfyUI 换装服务")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+class TryOnRequest(BaseModel):
+    model_image: str   # base64 模特图片
+    shirt_image: str   # base64 上衣图片
+    pants_image: str   # base64 裤子图片
+
+
+class TryOnResponse(BaseModel):
+    result_image: str  # base64 结果图片
+
+
+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 upload_image(client: httpx.AsyncClient, image_bytes: bytes, filename: str) -> str:
+    """上传图片到 ComfyUI，返回服务器文件名"""
+    files = {
+        "image": (filename, io.BytesIO(image_bytes), "image/jpeg"),
+    }
+    data = {"overwrite": "true"}
+    resp = await client.post(f"{COMFYUI_URL}/upload/image", files=files, data=data)
+    resp.raise_for_status()
+    result = resp.json()
+    return result["name"]
+
+
+async def queue_prompt(client: httpx.AsyncClient, workflow: dict) -> str:
+    """提交工作流到队列，返回 prompt_id"""
+    payload = {"prompt": workflow, "client_id": str(uuid.uuid4())}
+    resp = await client.post(f"{COMFYUI_URL}/prompt", json=payload)
+    resp.raise_for_status()
+    return resp.json()["prompt_id"]
+
+
+async def wait_for_result(client: httpx.AsyncClient, prompt_id: str, timeout: int = 300) -> dict:
+    """轮询历史记录，等待任务完成，返回输出节点数据"""
+    deadline = time.time() + timeout
+    while time.time() < deadline:
+        resp = await client.get(f"{COMFYUI_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_as_base64(client: httpx.AsyncClient, filename: str, subfolder: str = "", type_: str = "output") -> str:
+    """从 ComfyUI 下载图片并编码为 base64"""
+    params = {"filename": filename, "subfolder": subfolder, "type": type_}
+    resp = await client.get(f"{COMFYUI_URL}/view", params=params)
+    resp.raise_for_status()
+    content_type = resp.headers.get("content-type", "image/png")
+    b64 = base64.b64encode(resp.content).decode()
+    return f"data:{content_type};base64,{b64}"
+
+
+@app.post("/try-on", response_model=TryOnResponse)
+async def try_on(req: TryOnRequest):
+    """
+    换装接口
+    - model_image: 模特图片 base64
+    - shirt_image: 上衣图片 base64
+    - pants_image: 裤子图片 base64
+    返回 result_image: 换装结果 base64
+    """
+    # 加载工作流模板
+    workflow = json.loads(WORKFLOW_PATH.read_text(encoding="utf-8"))
+
+    async with httpx.AsyncClient(timeout=60.0) as client:
+        # 解码三张图片
+        try:
+            model_bytes = decode_base64_image(req.model_image)
+            shirt_bytes = decode_base64_image(req.shirt_image)
+            pants_bytes = decode_base64_image(req.pants_image)
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=f"图片解码失败: {e}")
+
+        # 生成唯一文件名，避免缓存冲突
+        uid = uuid.uuid4().hex[:8]
+        model_fname = f"model_{uid}.jpg"
+        shirt_fname = f"shirt_{uid}.jpg"
+        pants_fname = f"pants_{uid}.jpg"
+
+        # 并发上传三张图片
+        try:
+            model_name, shirt_name, pants_name = await asyncio.gather(
+                upload_image(client, model_bytes, model_fname),
+                upload_image(client, shirt_bytes, shirt_fname),
+                upload_image(client, pants_bytes, pants_fname),
+            )
+        except httpx.HTTPError as e:
+            raise HTTPException(status_code=502, detail=f"上传图片到 ComfyUI 失败: {e}")
+
+        # 修改工作流节点的图片输入
+        workflow[NODE_MODEL]["inputs"]["image"] = model_name
+        workflow[NODE_SHIRT]["inputs"]["image"] = shirt_name
+        workflow[NODE_PANTS]["inputs"]["image"] = pants_name
+
+        # 提交工作流
+        try:
+            prompt_id = await queue_prompt(client, workflow)
+        except httpx.HTTPError as e:
+            raise HTTPException(status_code=502, detail=f"提交工作流失败: {e}")
+
+        # 等待结果（最多5分钟）
+        async with httpx.AsyncClient(timeout=30.0) as poll_client:
+            outputs = await wait_for_result(poll_client, prompt_id, timeout=300)
+
+        # 从节点33的输出获取图片
+        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) as dl_client:
+            result_b64 = await fetch_image_as_base64(dl_client, filename, subfolder, type_)
+
+    return TryOnResponse(result_image=result_b64)
+
+
+@app.get("/health")
+async def health():
+    """健康检查"""
+    return {"status": "ok", "comfyui": COMFYUI_URL}
+
+
+# 挂载静态文件（前端页面）
+static_dir = Path(__file__).parent / "static"
+static_dir.mkdir(exist_ok=True)
+print(f"静态文件目录: {static_dir}")
+app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")
+
+
+@app.get("/")
+async def index():
+    return FileResponse(str(static_dir / "index.html"))
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=47698, log_level="info")