diff --git a/.github/workflows/tier0.yml b/.github/workflows/tier0.yml new file mode 100644 index 0000000..75dead2 --- /dev/null +++ b/.github/workflows/tier0.yml @@ -0,0 +1,33 @@ +name: Tier 0 — Unit (Linux) + +on: + push: + branches: [main] + pull_request: + +# Minimal-deps run: Tier 0 must work without ComfyUI, coremltools, or +# python_coreml_stable_diffusion (Linux CI image won't have them). The +# in-tree purity gate (tests/unit/test_tier0_purity.py) double-checks +# that the suite hasn't started leaking framework imports. +jobs: + unit: + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install Tier 0 deps + run: | + python -m pip install --upgrade pip + # Pins mirror pyproject (Phase 1 baseline). Tier 0 only needs + # torch + numpy + pytest; everything else is Mac-only. + python -m pip install \ + "torch==2.0.1" "numpy<1.25" \ + "pytest>=8" "pytest-xdist" + + - name: Run Tier 0 + run: pytest -m unit tests/ -v diff --git a/.github/workflows/tier1.yml b/.github/workflows/tier1.yml new file mode 100644 index 0000000..bb063ea --- /dev/null +++ b/.github/workflows/tier1.yml @@ -0,0 +1,31 @@ +name: Tier 1 — Smoke (macOS-ARM) + +on: + push: + branches: [main] + pull_request: + # Gate behind the run-tier1 label too, so external PRs that touch + # only docs don't burn a minute of macOS-ARM time. Maintainers can + # always re-run via the run-tier1 label. + types: [opened, synchronize, reopened, labeled] + +jobs: + smoke: + if: | + github.event_name == 'push' || + github.event.action != 'labeled' || + contains(github.event.pull_request.labels.*.name, 'run-tier1') + runs-on: macos-14 # M1, Apple Silicon hosted runner + timeout-minutes: 20 + steps: + - uses: actions/checkout@v4 + + - uses: astral-sh/setup-uv@v3 + with: + enable-cache: true + + - name: uv sync + run: uv sync --no-install-project + + - name: Run Tier 1 (synthetic micro-UNet smoke) + run: uv run pytest -m smoke tests/ -v diff --git a/.github/workflows/tier2.yml b/.github/workflows/tier2.yml new file mode 100644 index 0000000..a7ec67e --- /dev/null +++ b/.github/workflows/tier2.yml @@ -0,0 +1,63 @@ +name: Tier 2 — M2 / ANE (self-hosted) + +on: + pull_request: + types: [labeled] + schedule: + # Nightly at 04:00 UTC (~05/06 in PL). Keeps the M2 path honest + # without burning the runner on every PR. + - cron: "0 4 * * *" + workflow_dispatch: + +jobs: + m2: + if: | + github.event_name == 'schedule' || + github.event_name == 'workflow_dispatch' || + (github.event_name == 'pull_request' && + contains(github.event.pull_request.labels.*.name, 'run-m2')) + # Self-hosted Mac registered by the maintainer. See docs/ci-m2.md + # for runner setup and required model paths. + runs-on: [self-hosted, macOS, ARM64, coreml] + timeout-minutes: 60 + steps: + - uses: actions/checkout@v4 + + - name: uv sync + run: uv sync + + - name: Start ComfyUI server (background) + env: + COMFY_DIR: ${{ env.COMFY_DIR }} + run: | + cd "$COMFY_DIR" + nohup "$GITHUB_WORKSPACE/.venv/bin/python" main.py --port 8188 --cpu-vae > /tmp/comfyui-ci.log 2>&1 & + # Wait for readiness, fail fast if it never comes up. + for _ in $(seq 1 60); do + if grep -q "To see the GUI" /tmp/comfyui-ci.log 2>/dev/null; then + echo "comfy ready"; exit 0 + fi + sleep 2 + done + echo "comfy failed to start"; tail -100 /tmp/comfyui-ci.log; exit 1 + + - name: Run Tier 2 (m2 marker) + run: uv run pytest -m m2 tests/ -v + + - name: Run bench harness + run: | + uv run python bench/run.py \ + --model "$COMFY_DIR/models/unet/v1-5-pruned-emaonly_1x512x512_se_unet.mlmodelc" \ + --compute-units CPU_AND_NE CPU_AND_GPU \ + --repeats 30 --assumed-steps 20 + + - name: Upload bench results + uses: actions/upload-artifact@v4 + if: always() + with: + name: bench-results + path: bench/results/*.json + + - name: Stop ComfyUI server + if: always() + run: pkill -f "main.py.*8188" || true diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..8a2c480 --- /dev/null +++ b/Makefile @@ -0,0 +1,76 @@ +# ComfyUI-CoreMLSuite — tiered test/bench dispatcher (Phase 4). +# +# Tiers (see MODERNIZATION_SPEC.md): +# Tier 0 (unit): framework-free pure-logic tests, run anywhere in seconds. +# Tier 1 (smoke): macOS-ARM, no ANE, no full model — converts a synthetic +# micro-UNet to catch coremltools / ml-stable-diffusion +# API breakage in minutes. +# Tier 2 (m2): real ANE on Apple Silicon; integration + bench. +# +# COMFY_DIR defaults to the canonical custom-node layout (two dirs up from here). +# PY defaults to the project's uv-managed venv interpreter. + +COMFY_DIR ?= $(realpath $(CURDIR)/../..) +PY ?= $(CURDIR)/.venv/bin/python +PYTEST ?= $(PY) -m pytest + +UNAME_S := $(shell uname -s) +UNAME_M := $(shell uname -m) +IS_MACOS_ARM := $(filter Darwin,$(UNAME_S))$(filter arm64,$(UNAME_M)) + +.PHONY: help test-unit test-smoke test-m2 bench bench-rerun ci-tier0 ci-tier1 clean check-macos-arm + +help: + @echo "ComfyUI-CoreMLSuite — make targets" + @echo "" + @echo " test-unit Tier 0: pure-logic pytest, runs anywhere, seconds" + @echo " test-smoke Tier 1: synthetic micro-UNet ct.convert + load (macOS-ARM, minutes)" + @echo " test-m2 Tier 2: pytest -m m2 against a real Core ML UNet (Apple Silicon + ANE)" + @echo " bench Run bench/run.py against a converted .mlmodelc" + @echo "" + @echo "Vars: COMFY_DIR (default: $(COMFY_DIR)), PY (default: $(PY))" + +check-macos-arm: + @if [ -z "$(IS_MACOS_ARM)" ]; then \ + echo "this target requires macOS on Apple Silicon (got $(UNAME_S)/$(UNAME_M))"; \ + exit 2; \ + fi + +# Tier 0 — Linux-safe pure logic. Should not import comfy/coremltools. +test-unit: + $(PYTEST) -m unit tests/ + +# Tier 1 — macOS-ARM smoke. Converts a synthetic UNet through coremltools to +# catch API breakage without needing a real SD checkpoint or the ANE. +test-smoke: check-macos-arm + $(PYTEST) -m smoke tests/ + +# Tier 2 — full Apple Silicon path: integration + m2 golden + bench. +# Requires a converted .mlmodelc (see bench/scripts/convert_sd15.py). +test-m2: check-macos-arm + $(PYTEST) -m m2 tests/ + +# Bench harness. Override MODEL=/path/to/.mlmodelc for an explicit model. +MODEL ?= $(COMFY_DIR)/models/unet/v1-5-pruned-emaonly_1x512x512_se_unet.mlmodelc +COMPUTE_UNITS ?= CPU_AND_NE CPU_AND_GPU +REPEATS ?= 30 +ASSUMED_STEPS ?= 20 +bench: check-macos-arm + $(PY) bench/run.py \ + --model "$(MODEL)" \ + --compute-units $(COMPUTE_UNITS) \ + --repeats $(REPEATS) \ + --assumed-steps $(ASSUMED_STEPS) + +# What CI actually invokes — same as test-unit but echoes the env capture +# alongside so failed runs land with diagnostics. +ci-tier0: + @echo "## env (Tier 0)" && $(PY) --version && uv pip freeze --python "$(PY)" 2>/dev/null | head -50 || true + $(MAKE) test-unit + +ci-tier1: check-macos-arm + @echo "## env (Tier 1)" && $(PY) --version && uv pip freeze --python "$(PY)" 2>/dev/null | head -50 || true + $(MAKE) test-smoke + +clean: + rm -rf .pytest_cache tests/m2/_latest_generated.png pytestdebug.log diff --git a/docs/ci-m2.md b/docs/ci-m2.md new file mode 100644 index 0000000..4cf7912 --- /dev/null +++ b/docs/ci-m2.md @@ -0,0 +1,90 @@ +# Self-hosted M2 runner — setup + +Tier 2 (ANE + integration + bench) runs on a self-hosted GitHub Actions +runner registered against the maintainer's M-series Mac. Hosted macOS +runners on GitHub do not expose the Apple Neural Engine, so the ANE +half of the matrix has to live on real hardware. + +## One-time runner setup + +1. **Install dependencies on the Mac.** Python 3.11.x (matching the + `requires-python` pin), `uv`, `git`, plus the ComfyUI checkout at the + path the workflow expects (default: `$HOME/dev/ComfyUI`). The + workflow reads `COMFY_DIR` from the runner's env. + + ```bash + brew install python@3.11 uv git + ``` + +2. **Register the runner.** From the repo Settings → Actions → Runners + → New self-hosted runner, follow the macOS-ARM instructions. Add the + labels exactly: `self-hosted`, `macOS`, `ARM64`, `coreml` (the + workflow `runs-on` clause requires all four). + + ```bash + mkdir ~/actions-runner && cd ~/actions-runner + curl -O -L https://github.com/actions/runner/releases/download/v2.317.0/actions-runner-osx-arm64-2.317.0.tar.gz + tar xzf actions-runner-osx-arm64-2.317.0.tar.gz + ./config.sh --url https://github.com// \ + --token \ + --labels self-hosted,macOS,ARM64,coreml \ + --name "$(hostname)-m2" + ./svc.sh install && ./svc.sh start # run as a launchd service + ``` + +3. **Persist `COMFY_DIR` for the runner.** The workflow needs to know + where the ComfyUI checkout lives. Add it to the runner's `.env`: + + ```bash + echo 'COMFY_DIR=/Users//dev/ComfyUI' >> ~/actions-runner/.env + ``` + +4. **Pre-convert the baseline SD1.5 model.** The bench step expects + `$COMFY_DIR/models/unet/v1-5-pruned-emaonly_1x512x512_se_unet.mlmodelc`. + Run the conversion once manually: + + ```bash + cd $GITHUB_WORKSPACE + uv run python bench/scripts/convert_sd15.py + ``` + + Re-runs of the same combination are a no-op; the converter skips when + the .mlmodelc already exists. + +## Triggers + +The Tier 2 workflow (`.github/workflows/tier2.yml`) runs: + +- **On PR label `run-m2`** — maintainers add the label to opt a PR + into the ANE lane (the runner is not free; default off). +- **Nightly at 04:00 UTC** via `schedule:`. +- **Manually** via the workflow_dispatch button. + +## Artifacts + +- Bench JSON/MD are uploaded as `bench-results`. +- M2 golden image regressions surface as test failures in + `tests/m2/test_golden_image.py`; diff PNG is written next to the + golden under `tests/m2/_latest_generated.png` (gitignored). + +## When the runner is down + +If the maintainer's Mac is offline, the workflow queues until the +runner comes back. Cancel a stuck run from the Actions UI; the gate is +not blocking by default (Tier 0 + Tier 1 carry PR status). Tier 2 is +"good to merge once it goes green," not "blocked until then." + +## Replacing the integration e2e + +The legacy `tests/integration/test_basic_conversion_1_5.py` checked +CoreML output against an MPS reference image at PSNR > 25 dB. That +reference path is broken on macOS 26.x with torch 2.0.1 (see Phase 1 +Gate report). Phase 4 moves the same coverage to +`tests/m2/test_golden_image.py`, which: + +- runs the Core ML pipeline only (no MPS reference), +- asserts SHA256 against `tests/m2/goldens/sd15_seed42.sha256`, +- falls back to PSNR ≥ 40 dB if the hash drifts. + +This removes the human-eyeball dependency: a regression is now a +numerical fail, not a "looks different to me." diff --git a/tests/integration/test_basic_conversion_1_5.py b/tests/integration/test_basic_conversion_1_5.py deleted file mode 100644 index 1107885..0000000 --- a/tests/integration/test_basic_conversion_1_5.py +++ /dev/null @@ -1,85 +0,0 @@ -import json -import os - -import pytest -import requests -import torch -from PIL import Image -import numpy as np - -from folder_paths import get_save_image_path, get_output_directory - -IMAGE_PREFIX = "E2E-1.5" -IMAGE_PREFIX_CML = f"{IMAGE_PREFIX}-CoreML" -IMAGE_PREFIX_MPS = f"{IMAGE_PREFIX}-MPS" - - -class OutputImageRepository: - def __init__(self, name_prefix): - self.name_prefix = name_prefix - - def list_images(self): - full_output_folder, _, _, _, _ = get_save_image_path( - self.name_prefix, get_output_directory(), 512, 512 - ) - return full_output_folder, os.listdir(full_output_folder) - - def delete_images(self): - full_output_folder, images = self.list_images() - for image in images: - os.remove(os.path.join(full_output_folder, image)) - - def get_latest_image(self, prefix): - full_output_folder, images = self.list_images() - for image in sorted(images, reverse=True): - if image.startswith(prefix): - return os.path.join(full_output_folder, image) - return None - - -@pytest.fixture(scope="function") -def output_image_repository(): - repo = OutputImageRepository(IMAGE_PREFIX) - yield repo - repo.delete_images() - - -def test_basic_conversion_1_5(output_image_repository): - with open("tests/integration/workflows/e2e-1.5-basic-conversion.json") as f: - prompt = json.load(f) - prompt = randomize_seed_in_prompt(prompt) - queue_prompt(prompt) - - coreml_img_path = output_image_repository.get_latest_image(IMAGE_PREFIX_CML) - mps_img_path = output_image_repository.get_latest_image(IMAGE_PREFIX_MPS) - - coreml_image = Image.open(coreml_img_path) - mps_image = Image.open(mps_img_path) - - assert psnr(np.array(coreml_image), np.array(mps_image)) > 25 - - -def psnr(img1, img2): - mse = np.mean((img1 - img2) ** 2) - if mse == 0: - return 100 - PIXEL_MAX = 255.0 - return 20 * np.log10(PIXEL_MAX / np.sqrt(mse)) - - -def queue_prompt(prompt: dict): - p = {"prompt": prompt} - data = json.dumps(p).encode("utf-8") - req = requests.post("http://localhost:8188/prompt", data=data) - assert req.status_code == 200 - while True: - req = requests.get("http://localhost:8188/prompt") - if req.json()["exec_info"]["queue_remaining"] == 0: - break - - -def randomize_seed_in_prompt(prompt): - seed = torch.random.seed() - prompt["3"]["inputs"]["seed"] = seed - prompt["11"]["inputs"]["seed"] = seed - return prompt diff --git a/tests/smoke/__init__.py b/tests/smoke/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/smoke/test_synthetic_unet.py b/tests/smoke/test_synthetic_unet.py new file mode 100644 index 0000000..b5f074c --- /dev/null +++ b/tests/smoke/test_synthetic_unet.py @@ -0,0 +1,122 @@ +"""Tier 1 smoke: convert a synthetic micro-UNet through coremltools and load +it back with python_coreml_stable_diffusion's CoreMLModel. + +Purpose: catch API breakage in coremltools / ml-stable-diffusion *without* +needing a real SD checkpoint, the ANE, or a converted .mlmodelc on disk. +Runs in minutes on a hosted macOS-ARM runner (no Apple internal stuff). + +What it asserts: + - coremltools.convert still accepts the call shape we use today + - the resulting .mlpackage round-trips through CoreMLModel + - expected_inputs exposes the input names/shapes we declared + - calling the model returns the named output (`noise_pred`) + +Auto-skips on non-Apple-Silicon hosts so Tier 0 CI on Linux ignores it. +""" +import platform +import shutil + +import numpy as np +import pytest +import torch +import torch.nn as nn + + +pytestmark = pytest.mark.skipif( + platform.system() != "Darwin" or platform.machine() != "arm64", + reason="Tier 1 requires macOS on Apple Silicon", +) + + +# Tiny shapes — large enough to exercise conv2d + linear + addition kernels in +# coremltools, small enough that conversion finishes in seconds on CPU. +SAMPLE_SHAPE = (1, 4, 8, 8) +TIMESTEP_SHAPE = (1,) +ENCODER_SHAPE = (1, 64, 1, 4) # matches SD's transposed encoder_hidden_states layout +OUT_NAME = "noise_pred" + + +class TinyUNet(nn.Module): + """Minimal UNet-shaped graph: conv -> add(time+context) -> conv. + + Not a real diffusion model. Just enough op variety to exercise the + PyTorch -> MIL frontend in coremltools and confirm we can still wire + the inputs/outputs the way ml-stable-diffusion expects. + """ + + def __init__(self): + super().__init__() + self.conv_in = nn.Conv2d(4, 8, kernel_size=3, padding=1) + self.conv_out = nn.Conv2d(8, 4, kernel_size=3, padding=1) + self.time_proj = nn.Linear(1, 8) + self.text_proj = nn.Linear(64, 8) + + def forward(self, sample, timestep, encoder_hidden_states): + h = self.conv_in(sample) + t_emb = self.time_proj(timestep.unsqueeze(-1)).view(1, 8, 1, 1) + c_emb = self.text_proj(encoder_hidden_states.squeeze(2).mean(-1)).view(1, 8, 1, 1) + h = h + t_emb + c_emb + return self.conv_out(h) + + +@pytest.fixture(scope="module") +def tiny_mlpackage(tmp_path_factory): + """Convert TinyUNet once per test session and reuse the .mlpackage.""" + import coremltools as ct + + torch.manual_seed(0) + model = TinyUNet().eval() + example = ( + torch.randn(*SAMPLE_SHAPE), + torch.randn(*TIMESTEP_SHAPE), + torch.randn(*ENCODER_SHAPE), + ) + traced = torch.jit.trace(model, example) + + mlmodel = ct.convert( + traced, + inputs=[ + ct.TensorType(name="sample", shape=SAMPLE_SHAPE, dtype=np.float16), + ct.TensorType(name="timestep", shape=TIMESTEP_SHAPE, dtype=np.float16), + ct.TensorType(name="encoder_hidden_states", shape=ENCODER_SHAPE, dtype=np.float16), + ], + outputs=[ct.TensorType(name=OUT_NAME, dtype=np.float16)], + compute_units=ct.ComputeUnit.CPU_ONLY, + compute_precision=ct.precision.FLOAT16, + convert_to="mlprogram", + minimum_deployment_target=ct.target.macOS13, + ) + + out_dir = tmp_path_factory.mktemp("tiny_unet") + pkg_path = out_dir / "tiny.mlpackage" + mlmodel.save(str(pkg_path)) + yield pkg_path + shutil.rmtree(out_dir, ignore_errors=True) + + +def test_coremltools_convert_round_trips_via_coreml_model(tiny_mlpackage): + from python_coreml_stable_diffusion.coreml_model import CoreMLModel + + model = CoreMLModel(str(tiny_mlpackage), "CPU_ONLY", "packages") + + # expected_inputs is the contract our wrappers depend on. Lock the shape + # of the dict + a sample entry. + expected = dict(model.expected_inputs) + assert set(expected.keys()) == {"sample", "timestep", "encoder_hidden_states"} + assert tuple(expected["sample"]["shape"]) == SAMPLE_SHAPE + assert tuple(expected["timestep"]["shape"]) == TIMESTEP_SHAPE + assert tuple(expected["encoder_hidden_states"]["shape"]) == ENCODER_SHAPE + + # Forward pass: drive the model the way CoreMLModelWrapper does. + rng = np.random.default_rng(0) + inputs = { + "sample": rng.standard_normal(SAMPLE_SHAPE).astype(np.float16), + "timestep": rng.standard_normal(TIMESTEP_SHAPE).astype(np.float16), + "encoder_hidden_states": rng.standard_normal(ENCODER_SHAPE).astype(np.float16), + } + out = model(**inputs) + assert isinstance(out, dict), f"unexpected output type: {type(out)}" + assert OUT_NAME in out, f"missing output {OUT_NAME!r}; got {sorted(out)}" + assert out[OUT_NAME].shape == SAMPLE_SHAPE, ( + f"output shape drift: got {out[OUT_NAME].shape}, expected {SAMPLE_SHAPE}" + )