aszc-dev
ef2a18cff3
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).
88 lines
3.3 KiB
Python
88 lines
3.3 KiB
Python
#!/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())
|