chore(phase1): pin baseline toolchain and add bench harness scaffold
Phase 1 of the modernization plan: freeze the currently-working environment so later refactors have a measured reference point. - Pin python-coreml-stable-diffusion to commit e5d960c4 (the one already installed in the maintainer's apple_env), plus torch==2.0.1, coremltools==8.2 and numpy<1.25 to match the only env that loads ComfyUI successfully (Comfy's checkpoint-safe-loading branch in utils.py is gated on torch>=2.4, so newer torch + numpy 1.23 breaks at import). - Mirror the same pins in requirements.txt and commit uv.lock for reproducible installs. - Add requires-comfyui pinning ComfyUI to ab541335 (the validated commit). - Fix tests/unit/test_chunks.py fixture: get_model_config() now takes a ModelVersion argument; pass ModelVersion.SD15 (the previously-broken test was the only Phase 1 production-code change required). - Add the Phase 1 baseline harness: bench/run.py (direct Core ML UNet latency, deterministic), bench/scripts/convert_sd15.py (one-command conversion bypassing the node graph), bench/scripts/smoke_image.py (POSTs the existing e2e workflow to a local ComfyUI server and saves the Core ML image), bench/env/capture.sh (env snapshot), bench/prompts.json (fixed prompt set). - Ignore apple_env/, comfy_env/, and bench/scripts/*.log. Tests: 20/20 unit pass (test_chunks + test_controlnet).
This commit is contained in:
aszc-dev
@@ -0,0 +1,87 @@
|
||||
#!/usr/bin/env python3
|
||||
"""bench/scripts/convert_sd15.py — Phase 1 baseline UNet conversion.
|
||||
|
||||
Converts a SD1.5 checkpoint to a Core ML UNet (.mlmodelc) for the bench harness.
|
||||
Bypasses the ComfyUI node graph and calls coreml_suite.converter directly so
|
||||
the conversion is reproducible from a single command.
|
||||
|
||||
Run from the repo root with the project venv:
|
||||
.venv/bin/python bench/scripts/convert_sd15.py
|
||||
|
||||
Environment overrides:
|
||||
COMFY_DIR=... (default: ../..)
|
||||
CKPT_NAME=v1-5-pruned-emaonly.safetensors
|
||||
ATTN=SPLIT_EINSUM (SPLIT_EINSUM | SPLIT_EINSUM_V2 | ORIGINAL)
|
||||
HEIGHT=512 WIDTH=512 BATCH_SIZE=1
|
||||
CONTROLNET=0
|
||||
"""
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
log = logging.getLogger("convert_sd15")
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
COMFY_DIR = Path(os.environ.get("COMFY_DIR", REPO_ROOT.parents[1])).resolve()
|
||||
|
||||
if str(COMFY_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(COMFY_DIR))
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
import folder_paths # noqa: E402 (ComfyUI module — needs sys.path above)
|
||||
|
||||
# Ensure ComfyUI's folder_paths is pointed at the real install for checkpoints.
|
||||
folder_paths.base_path = str(COMFY_DIR)
|
||||
folder_paths.add_model_folder_path("checkpoints", str(COMFY_DIR / "models" / "checkpoints"))
|
||||
folder_paths.add_model_folder_path("unet", str(COMFY_DIR / "models" / "unet"))
|
||||
|
||||
from coreml_suite import converter # noqa: E402
|
||||
from coreml_suite.config import ModelVersion # noqa: E402
|
||||
|
||||
|
||||
def main() -> int:
|
||||
ckpt_name = os.environ.get("CKPT_NAME", "v1-5-pruned-emaonly.safetensors")
|
||||
attn = os.environ.get("ATTN", "SPLIT_EINSUM")
|
||||
height = int(os.environ.get("HEIGHT", "512"))
|
||||
width = int(os.environ.get("WIDTH", "512"))
|
||||
batch_size = int(os.environ.get("BATCH_SIZE", "1"))
|
||||
controlnet = os.environ.get("CONTROLNET", "0") not in ("0", "false", "False", "")
|
||||
|
||||
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
||||
if not ckpt_path:
|
||||
log.error("checkpoint not found: %s under %s", ckpt_name, COMFY_DIR / "models" / "checkpoints")
|
||||
return 2
|
||||
|
||||
attn_suffix = {"SPLIT_EINSUM": "se", "SPLIT_EINSUM_V2": "se2", "ORIGINAL": "orig"}[attn]
|
||||
cn_suffix = "_cn" if controlnet else ""
|
||||
stem = ckpt_name.split(".")[0]
|
||||
out_name = f"{stem}_{batch_size}x{width}x{height}{cn_suffix}_{attn_suffix}"
|
||||
unet_out_path = converter.get_out_path("unet", out_name)
|
||||
|
||||
log.info("repo_root=%s comfy_dir=%s", REPO_ROOT, COMFY_DIR)
|
||||
log.info("ckpt=%s out_name=%s", ckpt_path, out_name)
|
||||
log.info("attn=%s size=%dx%d batch=%d controlnet=%s", attn, width, height, batch_size, controlnet)
|
||||
|
||||
converter.convert(
|
||||
ckpt_path=ckpt_path,
|
||||
model_version=ModelVersion.SD15,
|
||||
unet_out_path=unet_out_path,
|
||||
batch_size=batch_size,
|
||||
sample_size=(height // 8, width // 8),
|
||||
controlnet_support=controlnet,
|
||||
lora_weights=[],
|
||||
attn_impl=attn,
|
||||
config_path=None,
|
||||
)
|
||||
|
||||
target_path = converter.compile_model(out_path=unet_out_path, out_name=out_name, submodule_name="unet")
|
||||
log.info("compiled: %s", target_path)
|
||||
print(target_path)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,163 @@
|
||||
#!/usr/bin/env python3
|
||||
"""bench/scripts/smoke_image.py — Phase 1 known-good-workflow image smoke test.
|
||||
|
||||
Posts a Stable Diffusion 1.5 workflow (CLIP → Core ML UNet → VAE → PNG) to a
|
||||
locally running ComfyUI server, waits for the queue to drain, then copies the
|
||||
generated CoreML + MPS reference images into bench/results/smoke/<gitsha>/
|
||||
and writes a small report.json with the PSNR between them. This is the Phase 1
|
||||
"known-good workflow still produces an image [M2-ANE]" artifact.
|
||||
|
||||
Prereqs:
|
||||
- ComfyUI server running locally on $COMFY_HOST:$COMFY_PORT (default
|
||||
http://localhost:8188), started against the project .venv so it picks up
|
||||
the pinned ml-stable-diffusion + Core ML Suite nodes.
|
||||
- The Core ML UNet has already been converted (see convert_sd15.py); the
|
||||
workflow's Core ML Converter node will skip if the .mlmodelc exists.
|
||||
|
||||
Run from the repo root:
|
||||
.venv/bin/python bench/scripts/smoke_image.py
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
COMFY_DIR = Path(os.environ.get("COMFY_DIR", REPO_ROOT.parents[1])).resolve()
|
||||
COMFY_HOST = os.environ.get("COMFY_HOST", "localhost")
|
||||
COMFY_PORT = int(os.environ.get("COMFY_PORT", "8188"))
|
||||
COMFY_URL = f"http://{COMFY_HOST}:{COMFY_PORT}"
|
||||
|
||||
CKPT_NAME = os.environ.get("CKPT_NAME", "v1-5-pruned-emaonly.safetensors")
|
||||
WORKFLOW_PATH = Path(os.environ.get(
|
||||
"WORKFLOW_PATH",
|
||||
REPO_ROOT / "tests" / "integration" / "workflows" / "e2e-1.5-basic-conversion.json",
|
||||
))
|
||||
TIMEOUT_S = int(os.environ.get("TIMEOUT_S", "1800")) # 30min cap
|
||||
|
||||
|
||||
def git_sha() -> str:
|
||||
try:
|
||||
return subprocess.check_output(
|
||||
["git", "-C", str(REPO_ROOT), "rev-parse", "--short", "HEAD"],
|
||||
stderr=subprocess.DEVNULL,
|
||||
).decode().strip()
|
||||
except Exception:
|
||||
return "nogit"
|
||||
|
||||
|
||||
def http_get_json(path: str) -> dict:
|
||||
with urllib.request.urlopen(f"{COMFY_URL}{path}", timeout=60) as r:
|
||||
return json.loads(r.read().decode())
|
||||
|
||||
|
||||
def http_post_json(path: str, payload: dict) -> dict:
|
||||
data = json.dumps(payload).encode("utf-8")
|
||||
req = urllib.request.Request(
|
||||
f"{COMFY_URL}{path}", data=data,
|
||||
headers={"Content-Type": "application/json"}, method="POST",
|
||||
)
|
||||
# Long timeout: first POST blocks while the server warms model loaders +
|
||||
# the Core ML compile-on-load pass (~60-90s on a cold cache).
|
||||
with urllib.request.urlopen(req, timeout=300) as r:
|
||||
return json.loads(r.read().decode())
|
||||
|
||||
|
||||
def psnr(a: np.ndarray, b: np.ndarray) -> float:
|
||||
mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2))
|
||||
if mse == 0:
|
||||
return 100.0
|
||||
return 20.0 * float(np.log10(255.0 / np.sqrt(mse)))
|
||||
|
||||
|
||||
def find_image(out_dir: Path, prefix: str) -> Path | None:
|
||||
matches = sorted(out_dir.glob(f"{prefix}_*.png"), reverse=True)
|
||||
return matches[0] if matches else None
|
||||
|
||||
|
||||
def main() -> int:
|
||||
sha = git_sha()
|
||||
out_root = REPO_ROOT / "bench" / "results" / "smoke" / sha
|
||||
out_root.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
workflow = json.loads(WORKFLOW_PATH.read_text())
|
||||
# Point both checkpoint nodes at the maintainer's available SD1.5 ckpt.
|
||||
for nid in ("4", "10"):
|
||||
if nid in workflow:
|
||||
workflow[nid]["inputs"]["ckpt_name"] = CKPT_NAME
|
||||
# Fixed seed for reproducibility (Phase 1 baseline).
|
||||
seed = int(os.environ.get("SEED", "42"))
|
||||
for nid in ("3", "11"):
|
||||
if nid in workflow and "seed" in workflow[nid].get("inputs", {}):
|
||||
workflow[nid]["inputs"]["seed"] = seed
|
||||
# Phase 1 only needs a Core ML image. The reference MPS branch (nodes
|
||||
# 3/8/9) trips a known MPS f16/f32 mismatch on torch 2.0.1 + macOS 26+,
|
||||
# which would crash the whole server. Strip those nodes so the queue
|
||||
# only runs the Core ML path; SKIP_MPS=0 to opt back in.
|
||||
if os.environ.get("SKIP_MPS", "1") not in ("0", "false", "False"):
|
||||
for nid in ("3", "8", "9"):
|
||||
workflow.pop(nid, None)
|
||||
|
||||
print(f"[smoke] git={sha} server={COMFY_URL} ckpt={CKPT_NAME} seed={seed}", file=sys.stderr)
|
||||
|
||||
resp = http_post_json("/prompt", {"prompt": workflow})
|
||||
print(f"[smoke] queued: {resp}", file=sys.stderr)
|
||||
|
||||
t0 = time.time()
|
||||
while True:
|
||||
if time.time() - t0 > TIMEOUT_S:
|
||||
print(f"[smoke] TIMEOUT after {TIMEOUT_S}s waiting for queue drain", file=sys.stderr)
|
||||
return 2
|
||||
try:
|
||||
q = http_get_json("/prompt")
|
||||
remaining = q.get("exec_info", {}).get("queue_remaining", -1)
|
||||
if remaining == 0:
|
||||
break
|
||||
except (urllib.error.URLError, json.JSONDecodeError) as exc:
|
||||
print(f"[smoke] poll error: {exc}", file=sys.stderr)
|
||||
time.sleep(2)
|
||||
|
||||
comfy_out = COMFY_DIR / "output"
|
||||
coreml_png = find_image(comfy_out, "E2E-1.5-CoreML")
|
||||
if coreml_png is None:
|
||||
print(f"[smoke] missing Core ML image under {comfy_out}", file=sys.stderr)
|
||||
return 3
|
||||
coreml_dst = out_root / coreml_png.name
|
||||
shutil.copy2(coreml_png, coreml_dst)
|
||||
|
||||
mps_png = find_image(comfy_out, "E2E-1.5-MPS")
|
||||
mps_dst = None
|
||||
psnr_db = None
|
||||
if mps_png is not None:
|
||||
mps_dst = out_root / mps_png.name
|
||||
shutil.copy2(mps_png, mps_dst)
|
||||
a = np.array(Image.open(coreml_dst).convert("RGB"))
|
||||
b = np.array(Image.open(mps_dst).convert("RGB"))
|
||||
if a.shape == b.shape:
|
||||
psnr_db = round(psnr(a, b), 2)
|
||||
|
||||
report = {
|
||||
"git_sha": sha,
|
||||
"seed": seed,
|
||||
"ckpt_name": CKPT_NAME,
|
||||
"workflow": str(WORKFLOW_PATH.relative_to(REPO_ROOT)),
|
||||
"coreml_image": str(coreml_dst.relative_to(REPO_ROOT)),
|
||||
"mps_image": str(mps_dst.relative_to(REPO_ROOT)) if mps_dst else None,
|
||||
"psnr_db": psnr_db,
|
||||
"wall_seconds": round(time.time() - t0, 1),
|
||||
}
|
||||
(out_root / "report.json").write_text(json.dumps(report, indent=2))
|
||||
print(json.dumps(report, indent=2))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user