优化为4b 模型
This commit is contained in:
@@ -0,0 +1,157 @@
|
||||
#!/usr/bin/env python3
|
||||
"""9B vs 4B 模型对比：3 张图 × 5 种发型 = 15 张对比图。
|
||||
|
||||
第一次运行：PLAN_TAG=A （9B 模型，需先切工作流到 .bak）
|
||||
第二次运行：PLAN_TAG=Bplus（4B 模型，需切回 4B 工作流）
|
||||
"""
|
||||
import json, os, time, base64, subprocess, urllib.request
|
||||
from datetime import datetime
|
||||
|
||||
API_BASE = "http://127.0.0.1:8187"
|
||||
TOKEN = "dev-shared-secret-2026"
|
||||
IMAGES = [
|
||||
"/home/ubuntu/hair/image/girl_img/girl1.jpg",
|
||||
"/home/ubuntu/hair/image/girl_img/girl7.jpg",
|
||||
"/home/ubuntu/hair/image/girl_img/girl13.jpg",
|
||||
]
|
||||
PLAN_TAG = os.getenv("PLAN_TAG", "A")
|
||||
OUT_DIR = f"/home/ubuntu/hair/benchmark_out/9b_vs_4b_{PLAN_TAG}"
|
||||
RESULT_JSON = os.path.join(OUT_DIR, "results.json")
|
||||
REPORT_HTML = os.path.join(OUT_DIR, "report.html")
|
||||
HAIR_STYLES = [(1, "椭圆"), (2, "花瓣"), (3, "心形"), (4, "直线"), (5, "波浪")]
|
||||
WARMUP = os.getenv("WARMUP", "1") == "1" # 第一次调用做预热
|
||||
|
||||
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 = "----CmpTest" + 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(image_path, hair_style):
|
||||
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 extract_image(data, img_idx, hs):
|
||||
"""从响应提取 grown_image_base64 并保存为 jpg。"""
|
||||
if not data or not data.get("data"):
|
||||
return None
|
||||
items = data["data"].get("results") or []
|
||||
if not items:
|
||||
return None
|
||||
b64 = items[0].get("grown_image_base64") or ""
|
||||
if b64.startswith("data:"):
|
||||
b64 = b64.split(",", 1)[1]
|
||||
if not b64:
|
||||
return None
|
||||
fname = f"img{img_idx}_style{hs}.jpg"
|
||||
path = os.path.join(OUT_DIR, fname)
|
||||
with open(path, "wb") as f:
|
||||
f.write(base64.b64decode(b64))
|
||||
return path
|
||||
|
||||
|
||||
def run_test():
|
||||
results = []
|
||||
print(f"9B vs 4B 对比测试 — Plan={PLAN_TAG}")
|
||||
print(f"图片: {[os.path.basename(p) for p in IMAGES]}")
|
||||
print(f"发型: {[(s,n) for s,n in HAIR_STYLES]}")
|
||||
print(f"总调用: {len(IMAGES)*len(HAIR_STYLES)}\n")
|
||||
|
||||
# 预热请求（避免第一次冷启动计入统计）
|
||||
if WARMUP:
|
||||
print("[warmup] 预热请求 (girl13, style=1) ...")
|
||||
t0 = time.time()
|
||||
_, warmup_time, _ = call_iface2(IMAGES[2], 1)
|
||||
print(f" warmup: {warmup_time:.1f}s\n")
|
||||
|
||||
for img_idx, img_path in enumerate(IMAGES, start=1):
|
||||
img_name = os.path.basename(img_path)
|
||||
for hs, hs_name in HAIR_STYLES:
|
||||
print(f"[img{img_idx}/{len(IMAGES)}] {img_name} style={hs}({hs_name}) ...")
|
||||
|
||||
vram_before, _ = get_vram()
|
||||
data, elapsed, error = call_iface2(img_path, hs)
|
||||
vram_after, _ = get_vram()
|
||||
|
||||
code = data.get("code", -1) if data else -1
|
||||
ok = (code == 0)
|
||||
img_saved = extract_image(data, img_idx, hs) if ok else None
|
||||
|
||||
record = {
|
||||
"plan": PLAN_TAG,
|
||||
"image_idx": img_idx,
|
||||
"image_name": img_name,
|
||||
"image_path": img_path,
|
||||
"hair_style": hs,
|
||||
"hair_style_name": hs_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,
|
||||
"saved_image_path": img_saved,
|
||||
}
|
||||
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_saved else 'none'}")
|
||||
|
||||
with open(RESULT_JSON, "w") as f:
|
||||
json.dump(results, f, indent=2, ensure_ascii=False)
|
||||
|
||||
print(f"\n✅ 完成！{RESULT_JSON}")
|
||||
print(f"📊 成功: {sum(1 for r in results if r['success'])}/{len(results)}")
|
||||
if any(r['success'] for r in results):
|
||||
avg = sum(r['elapsed_s'] for r in results if r['success']) / sum(1 for r in results if r['success'])
|
||||
print(f"⏱ 平均: {avg:.2f}s")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_test()
|
||||
Reference in New Issue
Block a user