#!/usr/bin/env python3 """接口2女性5种发型对比测试：Plan B vs Plan B+。 测试 hair_style=1..5（椭圆/花瓣/心形/直线/波浪），每种发型一次，记录耗时+显存+生成图。 """ import json, os, time, base64, subprocess, urllib.request, urllib.error from datetime import datetime API_BASE = "http://127.0.0.1:8187" TOKEN = "dev-shared-secret-2026" GIRL_IMG = "/home/ubuntu/hair/image/girl_img/girl13.jpg" PLAN_TAG = os.getenv("PLAN_TAG", "Bplus") # Bplus / B OUT_DIR = f"/home/ubuntu/hair/benchmark_out/iface2_female_{PLAN_TAG}" RESULT_JSON = os.path.join(OUT_DIR, "results.json") REPORT_HTML = os.path.join(OUT_DIR, "report.html") HAIR_STYLE_NAMES = { "1": "椭圆", "2": "花瓣", "3": "心形", "4": "直线", "5": "波浪", } os.makedirs(OUT_DIR, exist_ok=True) def get_vram(): try: out = subprocess.check_output( ["nvidia-smi", "--query-gpu=memory.used,memory.free", "--format=csv,noheader,nounits"], text=True, timeout=5).strip() used, free = out.split(",") return int(used.strip()), int(free.strip()) except Exception: return -1, -1 def _multipart(fields, files=None): boundary = "----If2Test" + str(int(time.time() * 1000)) parts = [] for k, v in fields.items(): parts.append(f"--{boundary}\r\n".encode()) parts.append(f'Content-Disposition: form-data; name="{k}"\r\n\r\n'.encode()) parts.append(str(v).encode()) parts.append(b"\r\n") if files: for field_name, (filename, data, mime) in files.items(): parts.append(f"--{boundary}\r\n".encode()) parts.append(f'Content-Disposition: form-data; name="{field_name}"; filename="{filename}"\r\n'.encode()) parts.append(f"Content-Type: {mime}\r\n\r\n".encode()) parts.append(data) parts.append(b"\r\n") parts.append(f"--{boundary}--\r\n".encode()) return b"".join(parts), boundary def call_iface2_female(image_path, hair_style): """接口2 女性 + 指定发型（走 swapHair 路径）""" with open(image_path, "rb") as f: img_data = f.read() fields = {"gender": "female", "hair_style": str(hair_style), "prompt": "填充遮罩区域的头发，皮肤加一点磨皮"} body, boundary = _multipart( fields, {"image_file": (os.path.basename(image_path), img_data, "image/jpeg")}) req = urllib.request.Request(f"{API_BASE}/api/v1/hair/grow", data=body, method="POST") req.add_header("Content-Type", f"multipart/form-data; boundary={boundary}") req.add_header("X-Internal-Token", TOKEN) t0 = time.time() try: with urllib.request.urlopen(req, timeout=600) as resp: data = json.loads(resp.read()) return data, time.time() - t0, None except Exception as e: return {}, time.time() - t0, str(e) def run_test(): results = [] print(f"接口2女性5种发型测试 — Plan={PLAN_TAG}") print(f"图片：{GIRL_IMG}\n") for hs in range(1, 6): name = HAIR_STYLE_NAMES[str(hs)] print(f"\n[{hs}/5] hair_style={hs} ({name}) ...") vram_before, vram_free_before = get_vram() data, elapsed, error = call_iface2_female(GIRL_IMG, hs) vram_after, vram_free_after = get_vram() code = data.get("code", -1) if data else -1 ok = (code == 0) # 保存生成图（如果有） img_path = None if ok and data.get("data"): try: items = data["data"].get("results") or [] if items and isinstance(items, list): first = items[0] b64 = first.get("grown_image_base64") or "" if b64.startswith("data:"): b64 = b64.split(",", 1)[1] if b64: img_path = os.path.join(OUT_DIR, f"hair_style_{hs}_{name}.jpg") with open(img_path, "wb") as f: f.write(base64.b64decode(b64)) except Exception as e: print(f" 保存图片失败: {e}") record = { "plan": PLAN_TAG, "hair_style": hs, "hair_style_name": name, "timestamp": datetime.now().strftime("%H:%M:%S"), "elapsed_s": round(elapsed, 2), "success": ok, "code": code, "error": error, "vram_before_mb": vram_before, "vram_after_mb": vram_after, "vram_delta_mb": vram_after - vram_before, "image_path": img_path, } results.append(record) status = "✓" if ok else "✗" print(f" → {status} code={code} time={elapsed:.1f}s " f"vram={vram_before}→{vram_after}MB (Δ{vram_after-vram_before:+d}) " f"img={'saved' if img_path else 'none'}") with open(RESULT_JSON, "w") as f: json.dump(results, f, indent=2, ensure_ascii=False) generate_html(results) generate_html(results) print(f"\n✅ 完成！{RESULT_JSON}") print(f"📄 报告：{REPORT_HTML}") def generate_html(results): total = len(results) success = sum(1 for r in results if r["success"]) times = [r["elapsed_s"] for r in results if r["success"]] avg_time = sum(times) / len(times) if times else 0 html = f""" 接口2女5种发型 - Plan {PLAN_TAG}

📊 接口2女性5种发型测试报告

Plan {PLAN_TAG} ｜ {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ｜ {total} 次调用
{success}
成功
{total-success}
失败
{avg_time:.1f}s
平均耗时
详细结果
""" for r in results: status = '✓' if r["success"] else f'✗ {str(r.get("error",""))[:30]}' html += f'' \ f'' \ f'' \ f'' html += '
发型名称耗时显存前显存后Δ状态
style {r["hair_style"]}{r["hair_style_name"]}
{r["elapsed_s"]:.1f}s		{r["vram_before_mb"]}MB{r["vram_after_mb"]}MB{r["vram_delta_mb"]:+d}{status}
' # 图片画廊 saved = [r for r in results if r.get("image_path") and os.path.isfile(r["image_path"])] if saved: html += '
生成图片
' for r in saved: rel = os.path.relpath(r["image_path"], OUT_DIR) html += f'
{r["hair_style_name"]} ({r["elapsed_s"]:.1f}s)
' html += '
' html += '' with open(REPORT_HTML, "w", encoding="utf-8") as f: f.write(html) if __name__ == "__main__": run_test()