如何使用 Core ML 运行 Stable Diffusion

Core ML 是 Apple 框架支持的模型格式和机器学习库。如果您有兴趣在 macOS 或 iOS/iPadOS 应用中运行 Stable Diffusion 模型，本指南将向您展示如何将现有的 PyTorch 检查点转换为 Core ML 格式，并使用 Python 或 Swift 进行推理。

Core ML 模型可以利用 Apple 设备中所有可用的计算引擎：CPU、GPU 和 Apple 神经引擎（或 ANE，一种在 Apple 芯片 Mac 和现代 iPhone/iPad 中可用的张量优化加速器）。根据模型和运行它的设备，Core ML 也可以混合和匹配计算引擎，例如，模型的某些部分可能在 CPU 上运行，而其他部分在 GPU 上运行。

Stable Diffusion Core ML 检查点

Stable Diffusion 权重（或检查点）以 PyTorch 格式存储，因此您需要将它们转换为 Core ML 格式，然后才能在原生应用中使用它们。

幸运的是，Apple 工程师开发了一个基于 diffusers 的转换工具，用于将 PyTorch 检查点转换为 Core ML。

不过，在转换模型之前，请花一点时间浏览 Hugging Face Hub – 您感兴趣的模型很可能已经以 Core ML 格式提供

• Apple 组织包括 Stable Diffusion 1.4、1.5、2.0 base 和 2.1 base 版本
• coreml community 包括自定义微调模型
• 使用此过滤器返回所有可用的 Core ML 检查点

如果您找不到您感兴趣的模型，我们建议您按照 Apple 提供的 Converting Models to Core ML 说明进行操作。

选择要使用的 Core ML 变体

Stable Diffusion 模型可以转换为不同的 Core ML 变体，用于不同的目的

• 所使用的注意力模块的类型。注意力操作用于“关注”图像表示中不同区域之间的关系，并理解图像和文本表示是如何关联的。注意力操作是计算和内存密集型的，因此存在不同的实现，这些实现考虑了不同设备的硬件特性。对于 Core ML Stable Diffusion 模型，有两种注意力变体：

  • split_einsum（由 Apple 引入）针对 ANE 设备进行了优化，ANE 设备在现代 iPhone、iPad 和 M 系列计算机中可用。
  • “原始”注意力（diffusers 中使用的基本实现）仅与 CPU/GPU 兼容，与 ANE 不兼容。使用 original 注意力在 CPU + GPU 上运行模型可能比在 ANE 上更快。请参阅 此性能基准测试 以及社区提供的 其他衡量标准 以了解更多详细信息。

• 支持的推理框架。

  • packages 适用于 Python 推理。这可以用于在尝试将转换后的 Core ML 模型集成到原生应用之前对其进行测试，或者如果您想探索 Core ML 性能但不需要支持原生应用。例如，具有 Web UI 的应用程序可以完美地使用 Python Core ML 后端。
  • Swift 代码需要 compiled 模型。Hub 中的 compiled 模型将大型 UNet 模型权重拆分为多个文件，以与 iOS 和 iPadOS 设备兼容。这对应于 --chunk-unet 转换选项。如果您想支持原生应用，则需要选择 compiled 变体。

官方 Core ML Stable Diffusion 模型 包括这些变体，但社区模型可能有所不同

coreml-stable-diffusion-v1-4
├── README.md
├── original
│   ├── compiled
│   └── packages
└── split_einsum
    ├── compiled
    └── packages

您可以下载并使用您需要的变体，如下所示。

Python 中的 Core ML 推理

安装以下库以在 Python 中运行 Core ML 推理

pip install huggingface_hub
pip install git+https://github.com/apple/ml-stable-diffusion

下载模型检查点

要在 Python 中运行推理，请使用存储在 packages 文件夹中的版本之一，因为 compiled 版本仅与 Swift 兼容。您可以选择是否要使用 original 或 split_einsum 注意力。

以下是如何从 Hub 下载 original 注意力变体到名为 models 的目录的方法

from huggingface_hub import snapshot_download
from pathlib import Path

repo_id = "apple/coreml-stable-diffusion-v1-4"
variant = "original/packages"

model_path = Path("./models") / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
snapshot_download(repo_id, allow_patterns=f"{variant}/*", local_dir=model_path, local_dir_use_symlinks=False)
print(f"Model downloaded at {model_path}")

推理

下载模型快照后，您可以使用 Apple 的 Python 脚本对其进行测试。

python -m python_coreml_stable_diffusion.pipeline --prompt "a photo of an astronaut riding a horse on mars" -i ./models/coreml-stable-diffusion-v1-4_original_packages/original/packages -o </path/to/output/image> --compute-unit CPU_AND_GPU --seed 93

使用 -i 标志将下载的检查点的路径传递给脚本。--compute-unit 指示您想要允许用于推理的硬件。它必须是以下选项之一：ALL、CPU_AND_GPU、CPU_ONLY、CPU_AND_NE。您还可以提供可选的输出路径和用于可重现性的种子。

推理脚本假定您正在使用 Stable Diffusion 模型的原始版本 CompVis/stable-diffusion-v1-4。如果您使用另一个模型，则必须在推理命令行中指定其 Hub id，使用 --model-version 选项。这适用于已支持的模型以及您自己训练或微调的自定义模型。

例如，如果您想使用 stable-diffusion-v1-5/stable-diffusion-v1-5

python -m python_coreml_stable_diffusion.pipeline --prompt "a photo of an astronaut riding a horse on mars" --compute-unit ALL -o output --seed 93 -i models/coreml-stable-diffusion-v1-5_original_packages --model-version stable-diffusion-v1-5/stable-diffusion-v1-5

Swift 中的 Core ML 推理

在 Swift 中运行推理比在 Python 中稍快，因为模型已编译为 mlmodelc 格式。这在应用启动时加载模型时很明显，但如果您在之后运行多次生成，则不应明显。

下载

要在 Mac 上的 Swift 中运行推理，您需要 compiled 检查点版本之一。我们建议您使用类似于前面示例的 Python 代码在本地下载它们，但使用 compiled 变体之一

from huggingface_hub import snapshot_download
from pathlib import Path

repo_id = "apple/coreml-stable-diffusion-v1-4"
variant = "original/compiled"

model_path = Path("./models") / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
snapshot_download(repo_id, allow_patterns=f"{variant}/*", local_dir=model_path, local_dir_use_symlinks=False)
print(f"Model downloaded at {model_path}")

推理

要运行推理，请克隆 Apple 的仓库

git clone https://github.com/apple/ml-stable-diffusion
cd ml-stable-diffusion

然后使用 Apple 的命令行工具 Swift Package Manager

swift run StableDiffusionSample --resource-path models/coreml-stable-diffusion-v1-4_original_compiled --compute-units all "a photo of an astronaut riding a horse on mars"

您必须在 --resource-path 中指定上一步下载的检查点之一，因此请确保它包含扩展名为 .mlmodelc 的已编译 Core ML bundles。--compute-units 必须是以下值之一：all、cpuOnly、cpuAndGPU、cpuAndNeuralEngine。

有关更多详细信息，请参阅 Apple 仓库中的说明。

支持的 Diffusers 功能

Core ML 模型和推理代码不支持 🧨 Diffusers 的许多功能、选项和灵活性。以下是一些需要记住的限制：

• Core ML 模型仅适用于推理。它们不能用于训练或微调。
• 只有两个调度器已移植到 Swift，即 Stable Diffusion 使用的默认调度器和 DPMSolverMultistepScheduler，我们已将其从 diffusers 实现移植到 Swift。我们建议您使用 DPMSolverMultistepScheduler，因为它在大约一半的步骤中产生相同的质量。
• 推理代码中提供了负面提示、无分类器引导缩放和图像到图像任务。深度引导、ControlNet 和潜在升频器等高级功能尚不可用。

Apple 的 转换和推理仓库 以及我们自己的 swift-coreml-diffusers 仓库旨在作为技术演示器，以使其他开发人员能够在其基础上进行构建。

如果您对任何缺失的功能有强烈的感觉，请随时打开功能请求，或者，更好的是，提交贡献 PR 🙂。

原生 Diffusers Swift 应用

在您自己的 Apple 硬件上运行 Stable Diffusion 的一个简单方法是使用我们的开源 Swift 代码库，它基于 diffusers 和 Apple 的转换及推理代码库。你可以学习代码，使用 Xcode 编译它，并根据自己的需求进行调整。为了方便起见，App Store 中还有一个独立的 Mac 应用程序，这样你就可以直接使用它，而无需处理代码或 IDE。如果您是一名开发者，并且已经确定 Core ML 是构建您的 Stable Diffusion 应用程序的最佳解决方案，那么您可以使用本指南的其余部分来开始您的项目。我们迫不及待地想看看您将构建出什么 🙂。