save code
This commit is contained in:
+54
-23
@@ -28,6 +28,7 @@ import oss2
|
||||
|
||||
# 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_1_ex.json"
|
||||
|
||||
# WORKFLOW_PATH = Path(__file__).parent / "change_fast2.json"
|
||||
@@ -97,31 +98,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:
|
||||
@@ -136,10 +160,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
|
||||
|
||||
@@ -154,7 +178,9 @@ async def try_on(req: TryOnRequest):
|
||||
"""
|
||||
|
||||
print(f"接收到换装请求,正在处理...")
|
||||
# 加载工作流模板
|
||||
|
||||
comfyui_url = await pick_comfyui_url()
|
||||
|
||||
workflow = json.loads(WORKFLOW_PATH.read_text(encoding="utf-8"))
|
||||
|
||||
desc = (req.change_desc or "").strip()
|
||||
@@ -162,42 +188,35 @@ 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)
|
||||
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"
|
||||
|
||||
# 并发上传两张图片
|
||||
try:
|
||||
model_name, shirt_name = await asyncio.gather(
|
||||
upload_image(client, model_bytes, model_fname),
|
||||
upload_image(client, shirt_bytes, shirt_fname),
|
||||
upload_image(client, comfyui_url, model_bytes, model_fname),
|
||||
upload_image(client, comfyui_url, shirt_bytes, shirt_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
|
||||
|
||||
# 提交工作流
|
||||
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} 的输出")
|
||||
@@ -211,9 +230,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))
|
||||
@@ -245,7 +263,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},
|
||||
}
|
||||
|
||||
|
||||
# 挂载静态文件(前端页面)
|
||||
|
||||
Reference in New Issue
Block a user