save code

This commit is contained in:
Your Name
2026-05-05 23:02:59 +08:00
parent 79d8c6d17f
commit 0b9b5d21ff
2 changed files with 109 additions and 47 deletions
+55 -24
View File
@@ -27,6 +27,7 @@ from pydantic import BaseModel
# COMFYUI_URL = "http://127.0.0.1:8188"
COMFYUI_URL = "http://112.126.94.241:28188"
COMFYUI_URL_BACKUP = "http://112.126.94.241:38188"
WORKFLOW_PATH = Path(__file__).parent / "change2_2_0308.json"
# WORKFLOW_PATH = Path(__file__).parent / "change_fast3.json"
@@ -88,31 +89,54 @@ def decode_base64_image(b64_str: str) -> bytes:
return base64.b64decode(b64_str)
async def upload_image(client: httpx.AsyncClient, image_bytes: bytes, filename: str) -> str:
async def check_comfyui_alive(base_url: str, timeout: float = 3.0) -> bool:
"""检查指定的 ComfyUI 服务器是否可用"""
try:
async with httpx.AsyncClient(timeout=timeout) as client:
resp = await client.get(f"{base_url}/system_stats")
return resp.status_code == 200
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:
"""上传图片到 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 = await client.post(f"{base_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:
async def queue_prompt(client: httpx.AsyncClient, base_url: str, workflow: dict) -> str:
"""提交工作流到队列,返回 prompt_id"""
payload = {"prompt": workflow, "client_id": str(uuid.uuid4())}
resp = await client.post(f"{COMFYUI_URL}/prompt", json=payload)
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, prompt_id: str, timeout: int = 300) -> dict:
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"{COMFYUI_URL}/history/{prompt_id}")
resp = await client.get(f"{base_url}/history/{prompt_id}")
resp.raise_for_status()
history = resp.json()
if prompt_id in history:
@@ -127,10 +151,10 @@ async def wait_for_result(client: httpx.AsyncClient, prompt_id: str, timeout: in
raise HTTPException(status_code=504, detail="等待 ComfyUI 超时(300秒)")
async def fetch_image_bytes(client: httpx.AsyncClient, filename: str, subfolder: str = "", type_: str = "output") -> bytes:
async def fetch_image_bytes(client: httpx.AsyncClient, base_url: str, filename: str, subfolder: str = "", type_: str = "output") -> bytes:
"""从 ComfyUI 下载图片,返回原始字节"""
params = {"filename": filename, "subfolder": subfolder, "type": type_}
resp = await client.get(f"{COMFYUI_URL}/view", params=params)
resp = await client.get(f"{base_url}/view", params=params)
resp.raise_for_status()
return resp.content
@@ -145,7 +169,9 @@ async def try_on(req: TryOnRequest):
返回 result_image: 换装结果 base64
"""
print(f"接收到换装请求,正在处理...")
# 加载工作流模板
comfyui_url = await pick_comfyui_url()
workflow = json.loads(WORKFLOW_PATH.read_text(encoding="utf-8"))
desc = (req.change_desc or "").strip()
@@ -153,7 +179,6 @@ async def try_on(req: TryOnRequest):
workflow[NODE_PROMPT]["inputs"] = {"value": desc}
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)
@@ -161,38 +186,32 @@ async def try_on(req: TryOnRequest):
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),
upload_image(client, comfyui_url, model_bytes, model_fname),
upload_image(client, comfyui_url, shirt_bytes, shirt_fname),
upload_image(client, comfyui_url, 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)
prompt_id = await queue_prompt(client, comfyui_url, 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)
outputs = await wait_for_result(poll_client, comfyui_url, prompt_id, timeout=300)
# 从节点33的输出获取图片
node_output = outputs.get(NODE_OUTPUT)
if not node_output:
raise HTTPException(status_code=500, detail=f"工作流未返回节点 {NODE_OUTPUT} 的输出")
@@ -206,9 +225,8 @@ async def try_on(req: TryOnRequest):
subfolder = img_info.get("subfolder", "")
type_ = img_info.get("type", "output")
# 下载结果图片
async with httpx.AsyncClient(timeout=30.0) as dl_client:
result_bytes = await fetch_image_bytes(dl_client, filename, subfolder, type_)
result_bytes = await fetch_image_bytes(dl_client, comfyui_url, filename, subfolder, type_)
# 缩放到 768x1024
# img = Image.open(io.BytesIO(result_bytes))
@@ -240,7 +258,20 @@ async def try_on(req: TryOnRequest):
@app.get("/health")
async def health():
"""健康检查"""
return {"status": "ok", "comfyui": COMFYUI_URL}
primary_ok = await check_comfyui_alive(COMFYUI_URL)
backup_ok = await check_comfyui_alive(COMFYUI_URL_BACKUP)
if primary_ok:
active = COMFYUI_URL
elif backup_ok:
active = COMFYUI_URL_BACKUP
else:
active = 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},
}
# 挂载静态文件(前端页面)