coreml-diffusion 0.1.0__tar.gz

coreml_diffusion-0.1.0/.github/workflows/publish-pypi.yml

@@ -0,0 +1,29 @@
+name: Publish to PyPI
+
+# Trusted Publishing (OIDC) — no API token stored. Configure a PyPI publisher
+# for this repo/workflow first (see https://docs.pypi.org/trusted-publishers/).
+# Fires when a GitHub Release is published; the release tag must match the
+# version in pyproject.toml (e.g. tag v0.1.0 -> version 0.1.0).
+on:
+  release:
+    types: [published]
+  workflow_dispatch:
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write # required for Trusted Publishing
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+
+      - name: Build
+        run: uv build
+
+      - name: Publish
+        run: uv publish

coreml_diffusion-0.1.0/.github/workflows/tier0.yml

@@ -0,0 +1,31 @@
+name: Tier 0 — Unit (Linux)
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+# Framework-free unit lane. The in-tree purity gate
+# (tests/unit/test_tier0_purity.py) enforces that the unit tests do not import
+# coremltools / diffusers / any ComfyUI module.
+jobs:
+  unit:
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+
+      - name: uv sync
+        run: uv sync
+
+      - name: Lint (ruff)
+        run: |
+          uv run ruff check .
+          uv run ruff format --check .
+
+      - name: Run Tier 0
+        run: uv run pytest -m unit tests/ -v

coreml_diffusion-0.1.0/.github/workflows/tier1.yml

@@ -0,0 +1,25 @@
+name: Tier 1 — Smoke (macOS)
+
+# Converts a synthetic micro-UNet through coremltools on a hosted Apple Silicon
+# runner — catches coremltools API breakage without a real checkpoint or the ANE.
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  smoke:
+    runs-on: macos-14
+    timeout-minutes: 20
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+
+      - name: uv sync
+        run: uv sync
+
+      - name: Run Tier 1
+        run: uv run pytest -m smoke tests/ -v

coreml_diffusion-0.1.0/.gitignore

@@ -0,0 +1,8 @@
+.venv/
+__pycache__/
+*.pyc
+*.mlpackage/
+*.mlmodelc/
+dist/
+build/
+*.egg-info/

coreml_diffusion-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023-2026 Adrian Szczepański
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

coreml_diffusion-0.1.0/PKG-INFO

@@ -0,0 +1,98 @@
+Metadata-Version: 2.4
+Name: coreml-diffusion
+Version: 0.1.0
+Summary: Convert diffusion-model checkpoints (SD1.5/SDXL) to Core ML for Apple Neural Engine — framework-free, ComfyUI-independent.
+Project-URL: Homepage, https://github.com/aszc-dev/coreml-diffusion
+Project-URL: Repository, https://github.com/aszc-dev/coreml-diffusion
+Project-URL: Issues, https://github.com/aszc-dev/coreml-diffusion/issues
+Author-email: Adrian Szczepański <hi@aszc.dev>
+License-Expression: MIT
+License-File: LICENSE
+Keywords: ane,apple-neural-engine,comfyui,core-ml,coreml,diffusers,diffusion,sdxl,stable-diffusion
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: MacOS
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Multimedia :: Graphics :: Graphics Conversion
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Typing :: Typed
+Requires-Python: <3.13,>=3.12
+Requires-Dist: coremltools<10,>=9
+Requires-Dist: diffusers>=0.30
+Requires-Dist: numpy<3,>=2
+Requires-Dist: omegaconf>=2.3
+Requires-Dist: peft>=0.13
+Requires-Dist: torch<2.8,>=2.7
+Requires-Dist: transformers>=4.44
+Description-Content-Type: text/markdown
+
+# coreml-diffusion
+
+Convert diffusion-model checkpoints into Core ML `.mlpackage` artifacts for the
+Apple Neural Engine (ANE) — framework-free and independent of ComfyUI.
+
+`coreml-diffusion` takes a single-file Stable Diffusion checkpoint and produces a
+Core ML UNet you can run on-device (macOS/iOS) or load back into ComfyUI via
+[ComfyUI-CoreMLSuite](https://github.com/aszc-dev/ComfyUI-CoreMLSuite), which
+depends on this package for its conversion path.
+
+## Positioning
+
+The niche is **diffusion models on the Apple Neural Engine via Core ML** — inside
+ComfyUI and on-device. ANE is the differentiator: low-power, GPU-free, embeddable
+in a Swift/iOS app. This is about feasibility and power efficiency for SD1.5/SDXL
+on ANE, not a raw-throughput claim against desktop GPUs.
+
+Supported today: SD1.5 and SDXL (verified). SDXL refiner and LCM convert but are
+not yet golden-verified (experimental). The scope is diffusion architectures
+generally, not Stable Diffusion specifically.
+
+## Install
+
+```sh
+uv pip install coreml-diffusion          # from PyPI (planned)
+uv pip install -e .                       # from a checkout
+```
+
+Requires Python 3.12 and (for conversion) `coremltools` 9 — conversion runs on
+macOS; the package imports and its CLI parse on any platform.
+
+## CLI
+
+```sh
+coreml-diffusion convert \
+    --ckpt path/to/model.safetensors \
+    --model-version SD15 \
+    --out unet.mlpackage \
+    --height 512 --width 512 \
+    --attn-impl SPLIT_EINSUM \
+    --quantize none
+```
+
+Options: `--batch-size`, `--controlnet`, `--lora PATH[:STRENGTH]` (repeatable),
+`--config` (original-config YAML). `--quantize {none,8,6,4}` applies k-means
+weight palettization. Run `coreml-diffusion convert --help` for the full list.
+
+The output `.mlpackage` is the deliverable: load it natively in Swift/Core ML, or
+through ComfyUI-CoreMLSuite.
+
+## Library
+
+```python
+import coreml_diffusion
+from coreml_diffusion import ModelVersion
+
+coreml_diffusion.convert(
+    "model.safetensors", ModelVersion.SD15, "unet.mlpackage",
+    height=512, width=512, attn_impl="SPLIT_EINSUM",
+)
+```
+
+Discovery API (`list_model_versions`, `list_attention_impls`, `list_quant_modes`,
+`CONTRACT_VERSION`) reports what this build can convert; the identifiers are an
+additive-only contract (removing/renaming one is a major bump).
+
+## License
+
+MIT

coreml_diffusion-0.1.0/README.md

@@ -0,0 +1,69 @@
+# coreml-diffusion
+
+Convert diffusion-model checkpoints into Core ML `.mlpackage` artifacts for the
+Apple Neural Engine (ANE) — framework-free and independent of ComfyUI.
+
+`coreml-diffusion` takes a single-file Stable Diffusion checkpoint and produces a
+Core ML UNet you can run on-device (macOS/iOS) or load back into ComfyUI via
+[ComfyUI-CoreMLSuite](https://github.com/aszc-dev/ComfyUI-CoreMLSuite), which
+depends on this package for its conversion path.
+
+## Positioning
+
+The niche is **diffusion models on the Apple Neural Engine via Core ML** — inside
+ComfyUI and on-device. ANE is the differentiator: low-power, GPU-free, embeddable
+in a Swift/iOS app. This is about feasibility and power efficiency for SD1.5/SDXL
+on ANE, not a raw-throughput claim against desktop GPUs.
+
+Supported today: SD1.5 and SDXL (verified). SDXL refiner and LCM convert but are
+not yet golden-verified (experimental). The scope is diffusion architectures
+generally, not Stable Diffusion specifically.
+
+## Install
+
+```sh
+uv pip install coreml-diffusion          # from PyPI (planned)
+uv pip install -e .                       # from a checkout
+```
+
+Requires Python 3.12 and (for conversion) `coremltools` 9 — conversion runs on
+macOS; the package imports and its CLI parse on any platform.
+
+## CLI
+
+```sh
+coreml-diffusion convert \
+    --ckpt path/to/model.safetensors \
+    --model-version SD15 \
+    --out unet.mlpackage \
+    --height 512 --width 512 \
+    --attn-impl SPLIT_EINSUM \
+    --quantize none
+```
+
+Options: `--batch-size`, `--controlnet`, `--lora PATH[:STRENGTH]` (repeatable),
+`--config` (original-config YAML). `--quantize {none,8,6,4}` applies k-means
+weight palettization. Run `coreml-diffusion convert --help` for the full list.
+
+The output `.mlpackage` is the deliverable: load it natively in Swift/Core ML, or
+through ComfyUI-CoreMLSuite.
+
+## Library
+
+```python
+import coreml_diffusion
+from coreml_diffusion import ModelVersion
+
+coreml_diffusion.convert(
+    "model.safetensors", ModelVersion.SD15, "unet.mlpackage",
+    height=512, width=512, attn_impl="SPLIT_EINSUM",
+)
+```
+
+Discovery API (`list_model_versions`, `list_attention_impls`, `list_quant_modes`,
+`CONTRACT_VERSION`) reports what this build can convert; the identifiers are an
+additive-only contract (removing/renaming one is a major bump).
+
+## License
+
+MIT

coreml_diffusion-0.1.0/coreml_diffusion/__init__.py

@@ -0,0 +1,108 @@
+"""coreml_diffusion — framework-free Core ML diffusion conversion.
+
+Converts diffusion-model checkpoints (SD1.5/SDXL today) into Core ML
+``.mlpackage`` artifacts for Apple Neural Engine, with no ComfyUI dependency.
+Usable as a library, via the ``coreml-diffusion`` CLI, or embedded in on-device
+(iOS) tooling. The public surface is the discovery API below plus ``convert``,
+``compose_out_name`` and ``ModelVersion``.
+
+This package MUST stay free of ``comfy`` / ``folder_paths`` / ``comfy_extras``;
+``import coreml_diffusion`` works in a comfy-free environment.
+
+Discovery contract (consumed by ComfyUI-CoreMLSuite): the node populates its
+dropdowns by calling ``list_*`` here, so installing a newer ``coreml_diffusion``
+surfaces new conversion types in the old node with no Suite change and no Suite
+version bump. The identifiers returned here are an ADDITIVE-ONLY contract:
+
+- adding an identifier, or promoting EXPERIMENTAL -> VERIFIED  => minor bump
+- removing/renaming an identifier, or demoting VERIFIED        => MAJOR bump + note
+
+because a saved workflow JSON references these strings verbatim.
+"""
+
+from enum import Enum
+
+from coreml_diffusion.attention import ATTENTION_IMPLEMENTATIONS
+from coreml_diffusion.model_version import ModelVersion
+from coreml_diffusion.naming import (
+    QUANT_NBITS_VALUES,
+    compose_out_name,
+    lora_names_from_params,
+)
+
+__all__ = [
+    "ModelVersion",
+    "Status",
+    "list_model_versions",
+    "list_attention_impls",
+    "list_quant_modes",
+    "CONTRACT_VERSION",
+    "compose_out_name",
+    "lora_names_from_params",
+    "convert",
+]
+
+
+class Status(Enum):
+    VERIFIED = "verified"  # has a golden anchor + passing [M2-ANE] check
+    EXPERIMENTAL = "experimental"  # convertible, not yet anchored/verified
+
+
+# Single source of truth for which conversions the Suite may surface. The Suite
+# gates on this status, NOT on a hardcoded node list: promoting a model to
+# VERIFIED expands the node's dropdown with no Suite change.
+#
+# Keyed by the ModelVersion MEMBER so ``list_model_versions`` can emit ``.name``
+# ("SD15", "SDXL"). The node reverses the dropdown string via ``ModelVersion[...]``
+# (name lookup, nodes.py), so emitting ``.value`` ("sd15") would raise KeyError on
+# every saved workflow. See seam.md §5.
+_MODEL_STATUS = {
+    ModelVersion.SD15: Status.VERIFIED,
+    ModelVersion.SDXL: Status.VERIFIED,
+    ModelVersion.SDXL_REFINER: Status.EXPERIMENTAL,  # -> VERIFIED after a refiner golden anchor
+    ModelVersion.LCM: Status.EXPERIMENTAL,  # -> VERIFIED after E-LCM golden anchor
+}
+
+
+def list_model_versions(include_experimental: bool = False) -> list[str]:
+    """Model versions by ``.name`` (e.g. ``["SD15", "SDXL"]``).
+
+    Returns VERIFIED versions only by default — the converter node calls this
+    plainly. A power-user/CLI path may pass ``include_experimental=True`` to also
+    list convertible-but-unanchored versions.
+    """
+    return [
+        version.name
+        for version, status in _MODEL_STATUS.items()
+        if status is Status.VERIFIED
+        or (include_experimental and status is Status.EXPERIMENTAL)
+    ]
+
+
+def list_attention_impls() -> list[str]:
+    """Supported attention implementations, e.g. ``["SPLIT_EINSUM", ...]``."""
+    return list(ATTENTION_IMPLEMENTATIONS)
+
+
+def list_quant_modes() -> list[str]:
+    """Palettization modes, e.g. ``["none", "8", "6", "4"]`` ("none" = unquantized)."""
+    return list(QUANT_NBITS_VALUES)
+
+
+# Discovery-contract version. Bump per the additive-only rules in this module's
+# docstring and CONVERTER_EXTRACTION_SPEC.md "Interface contract".
+CONTRACT_VERSION = "1.0"
+
+
+def __getattr__(name):
+    """Lazily expose the heavy conversion entrypoint.
+
+    ``convert`` pulls coremltools/diffusers, so importing it eagerly would drag
+    the Mac/heavy stack into every ``import coreml_diffusion`` and break the
+    Tier-0 (Linux, framework-free) lane. Resolve it only on first access.
+    """
+    if name == "convert":
+        from coreml_diffusion.convert import convert as _convert
+
+        return _convert
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

coreml_diffusion-0.1.0/coreml_diffusion/attention.py

@@ -0,0 +1,5 @@
+ATTENTION_IMPLEMENTATIONS = (
+    "SPLIT_EINSUM",
+    "SPLIT_EINSUM_V2",
+    "ORIGINAL",
+)

coreml_diffusion-0.1.0/coreml_diffusion/cli.py

@@ -0,0 +1,114 @@
+"""Command-line entry point for coreml_diffusion.
+
+Mirrors ``coreml_diffusion.convert`` so the package can produce a Core ML
+``.mlpackage`` with no ComfyUI involved — for headless and on-device (iOS)
+conversion workflows. The heavy import (coremltools/diffusers, pulled by
+``convert``) is deferred into the handler, so ``--help`` and argument parsing
+stay light and the arg→call mapping is testable on plain Linux.
+
+Example:
+    coreml-diffusion convert --ckpt model.safetensors --model-version SD15 \\
+        --out unet.mlpackage --height 512 --width 512 --attn-impl SPLIT_EINSUM
+"""
+
+import argparse
+
+import coreml_diffusion
+
+
+def _parse_lora(spec):
+    """Parse a ``PATH:STRENGTH`` lora spec into ``(path, float_strength)``.
+
+    Strength defaults to 1.0 when omitted. ``rsplit`` on the last ':' so Windows
+    drive letters / colons in the path survive.
+    """
+    path, sep, strength = spec.rpartition(":")
+    if not sep:
+        return spec, 1.0
+    return path, float(strength)
+
+
+def _convert_cmd(args):
+    sample_size = (args.height // 8, args.width // 8)
+    lora_weights = [_parse_lora(spec) for spec in (args.lora or [])]
+    coreml_diffusion.convert(
+        args.ckpt,
+        coreml_diffusion.ModelVersion[args.model_version],
+        args.out,
+        batch_size=args.batch_size,
+        sample_size=sample_size,
+        controlnet_support=args.controlnet,
+        lora_weights=lora_weights or None,
+        attn_impl=args.attn_impl,
+        config_path=args.config,
+        quantize_nbits=args.quantize,
+    )
+
+
+def build_parser():
+    parser = argparse.ArgumentParser(
+        prog="coreml-diffusion",
+        description="Convert diffusion checkpoints to Core ML for Apple Neural Engine.",
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    conv = sub.add_parser("convert", help="Convert a checkpoint's UNet to a .mlpackage")
+    conv.add_argument(
+        "--ckpt", required=True, help="Path to the source .safetensors checkpoint"
+    )
+    conv.add_argument(
+        "--model-version",
+        required=True,
+        # include experimental: the CLI is the power-user path. Experimental
+        # versions (LCM, SDXL_REFINER) convert but are not golden-verified.
+        choices=coreml_diffusion.list_model_versions(include_experimental=True),
+        help="Model architecture (verified: SD15, SDXL; experimental otherwise)",
+    )
+    conv.add_argument("--out", required=True, help="Output .mlpackage path to write")
+    conv.add_argument(
+        "--height", type=int, default=512, help="Target image height (default 512)"
+    )
+    conv.add_argument(
+        "--width", type=int, default=512, help="Target image width (default 512)"
+    )
+    conv.add_argument(
+        "--batch-size", type=int, default=1, help="Batch size (default 1)"
+    )
+    conv.add_argument(
+        "--attn-impl",
+        choices=coreml_diffusion.list_attention_impls(),
+        default=coreml_diffusion.list_attention_impls()[0],
+        help="Attention implementation (default SPLIT_EINSUM)",
+    )
+    conv.add_argument(
+        "--controlnet",
+        action="store_true",
+        help="Add ControlNet residual inputs to the converted UNet",
+    )
+    conv.add_argument(
+        "--lora",
+        action="append",
+        metavar="PATH[:STRENGTH]",
+        help="LoRA to fuse before conversion; repeatable. STRENGTH defaults to 1.0",
+    )
+    conv.add_argument(
+        "--config", default=None, help="Optional original-config YAML path"
+    )
+    conv.add_argument(
+        "--quantize",
+        choices=coreml_diffusion.list_quant_modes(),
+        default="none",
+        help="K-means weight palettization bits (default none = unquantized)",
+    )
+    conv.set_defaults(func=_convert_cmd)
+    return parser
+
+
+def main(argv=None):
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    return args.func(args)
+
+
+if __name__ == "__main__":
+    main()

coreml_diffusion-0.1.0/coreml_diffusion/conversion/__init__.py

@@ -0,0 +1,9 @@
+"""Core ML conversion helpers.
+
+The conversion approach originates from Apple's ml-stable-diffusion
+(https://github.com/apple/ml-stable-diffusion). This implementation has since
+diverged: it runs natively on diffusers' UNet2DConditionModel with its own
+SPLIT_EINSUM / SPLIT_EINSUM_V2 attention processors and no longer depends on
+that package. The intent is to keep iterating on these methods independently
+while tracking current tooling.
+"""