Files
apple--ml-stable-diffusion/README.md
T
2026-07-13 10:42:24 +00:00

55 KiB
Raw Blame History

Note

本文档由 WeHub 基于上游 README 翻译整理,属于社区翻译,非官方中文文档。
English · 原始项目 · 上游 README
原作者、版权与许可证归属以原始项目及本仓库 LICENSE 文件为准。

Core ML Stable Diffusion

在 Apple Silicon 上使用 Core ML 运行 Stable Diffusion

[博客文章] [BibTeX]

本仓库包含:

  • python_coreml_stable_diffusion,一个 Python 包,用于将 PyTorch 模型转换为 Core ML 格式,并使用 Hugging Face diffusers 在 Python 中进行图像生成
  • StableDiffusion,一个 Swift 包,开发者可将其作为依赖项添加到 Xcode 项目中,以在应用中部署图像生成功能。该 Swift 包依赖由 python_coreml_stable_diffusion 生成的 Core ML 模型文件

如果你在安装或运行过程中遇到问题,请参阅 FAQ 部分。开始之前,请先参阅 系统要求 部分。

系统要求

详情(点击展开)

模型转换:

macOS Python coremltools
13.1 3.8 7.0

项目构建:

macOS Xcode Swift
13.1 14.3 5.8

目标设备运行时:

macOS iPadOS, iOS
13.1 16.2

目标设备运行时(含内存改进):

macOS iPadOS, iOS
14.0 17.0

目标设备硬件世代:

Mac iPad iPhone
M1 M1 A14

性能基准测试

详情(点击展开)

stabilityai/stable-diffusion-2-1-base (512x512)

设备 --compute-unit --attention-implementation 端到端延迟 (s) 扩散速度 (iter/s)
iPhone 12 Mini CPU_AND_NE SPLIT_EINSUM_V2 18.5* 1.44
iPhone 12 Pro Max CPU_AND_NE SPLIT_EINSUM_V2 15.4 1.45
iPhone 13 CPU_AND_NE SPLIT_EINSUM_V2 10.8* 2.53
iPhone 13 Pro Max CPU_AND_NE SPLIT_EINSUM_V2 10.4 2.55
iPhone 14 CPU_AND_NE SPLIT_EINSUM_V2 8.6 2.57
iPhone 14 Pro Max CPU_AND_NE SPLIT_EINSUM_V2 7.9 2.69
iPad Pro (M1) CPU_AND_NE SPLIT_EINSUM_V2 11.2 2.19
iPad Pro (M2) CPU_AND_NE SPLIT_EINSUM_V2 7.0 3.07
详情(点击展开)
  • 该基准测试由 Apple 与 Hugging Face 于 2023 年 8 月使用 iOS 17.0、iPadOS 17.0 和 macOS 14.0 Seed 8 的公开测试版完成。
  • 性能数据通过 Diffusers appbenchmark 分支采集
  • Swift 代码尚未完全优化,会引入最高约 10% 的额外开销,与 Core ML 模型执行无关。
  • 报告的是连续 5 次端到端执行的延迟中位数
  • 图像生成流程遵循标准配置:20 次推理步数、512x512 输出图像分辨率、77 文本 token 序列长度、无分类器引导(unet 的 batch size 为 2)。
  • 实际提示词长度不会影响性能,因为 Core ML 模型在转换时采用静态形状,无论输入文本的实际长度如何,都会对文本 token 序列中的全部 77 个元素(tokenizer.model_max_length)计算前向传播。
  • 权重压缩为 6 bit 精度。详情请参阅本节
  • GPU 与 Neural Engine 的激活值均为 float16 精度。
  • * 表示已启用 reduceMemory 选项,该选项会即时加载和卸载模型以避免内存不足。这会使端到端延迟最多增加 2 秒。
  • 在基准测试表中,我们报告每台设备上表现最佳的 --compute-unit--attention-implementation 值。前者不修改 Core ML 模型,可在运行时应用;后者会修改 Core ML 模型。请注意,表现最佳的计算单元因模型版本和硬件而异。
  • 请注意,本仓库中的性能优化(例如 --attention-implementation)通常适用于 Transformers,而非针对 Stable Diffusion 定制。通过自定义内核调优可能获得更好性能。因此,这些数据并不代表峰值硬件能力。
  • 由于模型架构变化,不同版本的 Stable Diffusion 性能可能有所不同。每项报告数据均针对该上下文中提及的特定模型版本。
  • 性能可能因其他应用增加系统负载或设备散热状态不佳等因素而波动。

stabilityai/stable-diffusion-xl-base-1.0-ios (768x768)

设备 --compute-unit --attention-implementation 端到端延迟 (s) 扩散速度 (iter/s)
iPhone 12 Pro CPU_AND_NE SPLIT_EINSUM 116* 0.50
iPhone 13 Pro Max CPU_AND_NE SPLIT_EINSUM 86* 0.68
iPhone 14 Pro Max CPU_AND_NE SPLIT_EINSUM 77* 0.83
iPhone 15 Pro Max CPU_AND_NE SPLIT_EINSUM 31 0.85
iPad Pro (M1) CPU_AND_NE SPLIT_EINSUM 36 0.69
iPad Pro (M2) CPU_AND_NE SPLIT_EINSUM 27 0.98
详情(点击展开)
  • 该基准测试由 Apple 与 Hugging Face 于 2023 年 9 月使用 iOS 17.0.2 和 iPadOS 17.0.2 完成。
  • 性能数据通过 Diffusers appbenchmark 分支采集
  • 报告的是连续 5 次端到端执行的延迟中位数
  • 图像生成流程遵循以下配置:20 次推理步数、768x768 输出图像分辨率、77 文本 token 序列长度、无分类器引导(unet 的 batch size 为 2)。
  • Unet.mlmodelc 按照此处 发布的 Mixed-Bit Palettization 算法方案压缩至 4.04 bit 精度
  • Unet.mlmodelc 外,所有模型均压缩为 16 bit 精度
  • madebyollin/sdxl-vae-fp16-fix@madebyollin 用作 VAEDecoder.mlmodelc 的源 PyTorch 模型,以便为 VAE 模型启用 float16 权重与激活值量化。
  • 由于 SPLIT_EINSUM_V2 的编译时间过长,选用 --attention-implementation SPLIT_EINSUM 替代
  • * 表示已启用 reduceMemory 选项,该选项会即时加载和卸载模型以避免内存不足。这会给端到端延迟带来显著额外开销。请注意,尽管扩散速度相同,iPad Pro (M1)iPhone 13 Pro Max 的端到端延迟仍存在差异。
  • 实际提示词长度不会影响性能,因为 Core ML 模型在转换时采用静态形状,无论输入文本的实际长度如何,都会对文本 token 序列中的全部 77 个元素(tokenizer.model_max_length)计算前向传播。
  • 在基准测试表中,我们报告每台设备上表现最佳的 --compute-unit--attention-implementation 值。前者不修改 Core ML 模型,可在运行时应用;后者会修改 Core ML 模型。请注意,表现最佳的计算单元因模型版本和硬件而异。
  • 请注意,本仓库中的性能优化(例如 --attention-implementation)通常适用于 Transformers,而非针对 Stable Diffusion 定制。通过自定义内核调优可能获得更好性能。因此,这些数据并不代表峰值硬件能力。
  • 由于模型架构变化,不同版本的 Stable Diffusion 性能可能有所不同。每项报告数据均针对该上下文中提及的特定模型版本。
  • 性能可能因其他应用增加系统负载或设备散热状态不佳等因素而波动。

stabilityai/stable-diffusion-xl-base-1.0 (1024x1024)

Device --compute-unit --attention-implementation End-to-End Latency (s) Diffusion Speed (iter/s)
MacBook Pro (M1 Max) CPU_AND_GPU ORIGINAL 46 0.46
MacBook Pro (M2 Max) CPU_AND_GPU ORIGINAL 37 0.57
Mac Studio (M1 Ultra) CPU_AND_GPU ORIGINAL 25 0.89
Mac Studio (M2 Ultra) CPU_AND_GPU ORIGINAL 20 1.11
详情(点击展开)
  • 该基准测试由 Apple 与 Hugging Face 于 2023 年 7 月使用 iOS 17.0、iPadOS 17.0 和 macOS 14.0 的公开测试版完成。
  • 性能数据通过运行 StableDiffusion Swift 流水线收集。
  • 报告的是连续 3 次端到端执行的延迟中位数。
  • 图像生成流程遵循标准配置:20 次推理步数、1024x1024 输出图像分辨率、无分类器引导(classifier-free guidanceunet 的 batch size 为 2)。
  • 权重与激活值采用 float16 精度。
  • 由于模型架构本身的变化,不同版本的 Stable Diffusion 性能可能有所差异。每个报告数值均针对该上下文中提及的特定模型版本。
  • 其他应用带来的系统负载增加或设备散热状态不佳等因素也可能影响性能。鉴于这些因素,我们不报告亚秒级的延迟波动。

权重压缩(6 位及以上)

详情(点击展开)

coremltools-7.0 支持用于 pruning, palettizationlinear 8-bit quantization. 的高级权重压缩技术。对于这些技术,coremltools.optimize.torch.* 包含需要在更高压缩率下通过微调以保持精度的 API;而 coremltools.optimize.coreml.* 包含在训练后应用且无需数据的 API。

我们演示了 coremltools.optimize.coreml.palettize_weights 中实现的无数据 post-training palettization 如何显著提升移动设备上 Stable Diffusion 的性能。该 API 实现了 Fast Exact k-Means 算法以实现最优权重聚类,从而生成更精确的调色板(palette)。在 conversion 期间使用 --quantize-nbits {2,4,6,8} 将对 unet 和 text_encoder 模型应用此压缩。

为获得最佳效果,若微调模型可行,我们建议使用 training-time palettization: coremltools.optimize.torch.palettization.DKMPalettizer。该 API 实现了 Differentiable k-Means (DKM) 可学习调色板算法。在本示例中,为简便起见并便于复现,我们仍采用训练后调色板化(post-training palettization)。

神经网络引擎(Neural Engine)可加速采用低位调色板化的模型:1、2、4、6 或 8 位。在 iOS 17 和 macOS 14 上,Core ML 模型的压缩权重可在运行时即时解压(just-in-time decompression,而非加载时预先解压 ahead-of-time decompression),以匹配激活张量的精度。这带来显著的内存节省,并使模型能在内存更小的设备上运行(例如 iPhone 12 Mini)。此外,压缩权重从内存中获取更快,可降低受内存带宽限制的层的延迟。即时解压行为取决于计算单元、层类型和硬件代次。

Weight Precision --compute-unit stabilityai/stable-diffusion-2-1-base 生成 "a high quality photo of a surfing dog"
6-bit cpuAndNeuralEngine
16-bit cpuAndNeuralEngine
16-bit cpuAndGPU

请注意,16 位(float16)与 6 位结果之间存在细微差异。这些差异与 float16 与 float32 之间的差异,或如上所示不同计算单元之间的差异相当。我们建议对 Stable Diffusion 进行调色板化时至少使用 6 位。更少的位数(1、2 和 4)将需要微调或 MBP 等高级调色板化技术。

资源:

高级权重压缩(低于 6 位)

详情(点击展开)

本节介绍一种名为 Mixed-Bit Palettization (MBP) 的高级压缩算法,它构建于 Post-Training Weight Palettization tools 之上,并使用来自 coremltools. 的 Weights Metadata API

MBP 通过为每一层选择神经网络引擎支持的 1、2、4、6 和 8 位位宽中的合适位数,构建逐层“调色板化配方”(palettization recipe),以在保持所需信号强度的同时实现最低平均位宽。信号强度通过将压缩模型输出与原始 float16 模型输出进行比较来衡量。在相同随机种子和文本提示下,计算去噪潜变量之间的 PSNR。压缩率取决于模型版本以及对信号损失(PSNR 下降)的容忍度,因为该算法具有自适应性。

3.41-bit 4.50-bit 6.55-bit 16-bit (original)

例如,原始 float16 stabilityai/stable-diffusion-xl-base-1.0 模型的信号强度约为 82 dB。对 Unet 模型简单应用 linear 8-bit quantization 会将信号降至约 65 dB。相比之下,应用 MBP 可实现平均 2.81 位量化,同时保持约 67 dB 的信号强度。与在模型转换期间使用 --quantize-nbits 相比,该技术通常能获得更好的结果,但需要进行一次“预分析”运行,在单块 GPU(mpscuda)上最多耗时数小时。

下图展示了 stabilityai/stable-diffusion-xl-base-1.0 的信号强度(PSNR,单位 dB)与模型体积缩减(占 float16 体积的百分比)之间的关系。{1,2,4,6,8}-bit 曲线通过使用固定比特数的调色板(palettization)逐步对更多层进行调色板量化而生成。各层按其对端到端信号强度的独立影响升序排列,从而尽可能推迟累积压缩带来的影响。混合比特曲线会在某层对端到端信号完整性的独立影响低于阈值时,立即回退到更高的比特数。请注意,在相同模型体积下,除 1 比特外,所有基于调色板量化的曲线均优于线性 8 比特量化。

以下是在其他模型版本上应用该技术的步骤:

Step 1 运行预分析脚本,生成具有不同信号强度的 "recipes"

python -m python_coreml_stable_diffusion.mixed_bit_compression_pre_analysis --model-version <model-version> -o <output-dir>

对于常用的基础模型,你可以在此处. 找到预计算的预分析结果。微调模型很可能会沿用其对应基础模型的 recipe,但这一点尚未经过验证。

Step 2 Step 1 生成的 JSON 文件会列出 "baselines",例如:

{
  "model_version": "stabilityai/stable-diffusion-xl-base-1.0",
  "baselines": {
    "original": 82.2,
    "linear_8bit": 66.025,
    "recipe_6.55_bit_mixedpalette": 79.9,
    "recipe_5.52_bit_mixedpalette": 78.2,
    "recipe_4.89_bit_mixedpalette": 76.8,
    "recipe_4.41_bit_mixedpalette": 75.5,
    "recipe_4.04_bit_mixedpalette": 73.2,
    "recipe_3.67_bit_mixedpalette": 72.2,
    "recipe_3.32_bit_mixedpalette": 71.4,
    "recipe_3.19_bit_mixedpalette": 70.4,
    "recipe_3.08_bit_mixedpalette": 69.6,
    "recipe_2.98_bit_mixedpalette": 68.6,
    "recipe_2.90_bit_mixedpalette": 67.8,
    "recipe_2.83_bit_mixedpalette": 67.0,
    "recipe_2.71_bit_mixedpalette": 66.3
  },
}

在这些 baseline 中,根据你期望的信号强度选择一种 recipe。我们建议根据使用场景将调色板量化至约 4 比特,即便较低比特值的信号完整性高于线性 8 比特量化 baseline。

最后,按如下方式将所选 recipe 应用于 float16 Core ML 模型:

python -m python_coreml_stable_diffusion.mixed_bit_compression_apply --mlpackage-path <path-to-float16-unet-mlpackage> -o <output-dir> --pre-analysis-json-path <path-to--pre-analysis-json> --selected-recipe <selected-recipe-string-key>

示例:<selected-recipe-string-key> 可以是 "recipe_4.50_bit_mixedpalette",它实现了平均 4.50 比特压缩(SDXL 从约 5.2GB 压缩到约 1.46GB)。请注意,信号强度与图文对齐(image-text alignment)并非直接对应。请务必验证你的 MBP 压缩模型变体能否针对测试提示词准确生成图像。

Activation Quantization

详情(点击展开)

在搭载 A17 Pro 或 M4 芯片的新硬件上(例如 iPhone 15 Pro),将激活和权重同时量化为 int8 可以利用 Neural Engine 上的优化计算,从而提升计算密集型模型的运行时延迟表现。

本节演示如何在 Stable Diffusion UNet 模型上,使用校准数据应用训练后激活量化(Post Training Activation Quantization, on Stable Diffusion UNet model.

与上文所述的混合比特调色板量化(Mixed-Bit PalettizationMBP类似,首先运行逐层分析,以确定哪些中间激活对 8 比特压缩更敏感。 敏感度较低的层进行权重与激活量化(W8A8),而敏感度较高的层仅进行权重量化(W8A16)。

以下是在其他模型版本上应用该技术的步骤:

Step 1 生成校准数据

python -m python_coreml_stable_diffusion.activation_quantization --model-version <model-version> --generate-calibration-data  -o <output-dir>

一组校准文本提示词会经由 StableDiffusionPipeline 运行,UNet 模型输入会被记录并以 pickle 文件形式保存在指定输出目录内的 calibration_data_<model-version> 文件夹中。

Step 2 运行逐层敏感度分析

python -m python_coreml_stable_diffusion.activation_quantization --model-version <model-version> --layerwise-sensitivity --calibration-nsamples <num-samples> -o <output-dir>

这会对模型中所有卷积(Convolutional)和注意力(AttentionEinsum)模块运行分析。 对每个模块,仅量化该层的权重与激活以生成压缩版本。 然后使用相同的随机种子和文本提示词,计算压缩模型与原始模型输出之间的 PSNR。

该分析在单块 GPU(cuda)上最多可能需要数小时。可减少用于量化模型的校准样本数量以加快流程。

生成的 JSON 文件如下所示:

{
  "conv": {
    "conv_in": 30.74,
    "down_blocks.0.attentions.0.proj_in": 38.93,
    "down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q": 48.15,
    "down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_k": 50.13,
    "down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_v": 45.70,
    "down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_out.0": 39.56,
    ...
  },
  "einsum": {
    "down_blocks.0.attentions.0.transformer_blocks.0.attn1.einsum": 25.34,
    "down_blocks.0.attentions.0.transformer_blocks.0.attn2.einsum": 31.76,
    "down_blocks.0.attentions.1.transformer_blocks.0.attn1.einsum": 23.40,
    "down_blocks.0.attentions.1.transformer_blocks.0.attn2.einsum": 31.56,
    ...
  },
  "model_version": "stabilityai/stable-diffusion-2-1-base"
}

Step 3 生成量化模型

使用校准数据和逐层敏感度分析结果,可按如下方式生成量化 CoreML 模型:

python -m python_coreml_stable_diffusion.activation_quantization --model-version <model-version> --quantize-pytorch --conv-psnr 38 --attn-psnr 26 -o <output-dir>

PSNR 阈值决定哪些层将进行激活量化。可调整该数值以在输出质量与推理延迟之间权衡。

Using Stable Diffusion 3

详情(点击展开)

Model Conversion

Stable Diffusion 3 运行时同时使用一些新模型与旧模型。对于文本编码器,可使用与之前类似的命令进行转换,并添加 --sd3-version 标志。

python -m python_coreml_stable_diffusion.torch2coreml --model-version stabilityai/stable-diffusion-3-medium --bundle-resources-for-swift-cli --convert-text-encoder --sd3-version -o <output-dir>

对于新模型(MMDiT、16 通道的新 VAE 以及 T5 文本编码器),有一系列新的 CLI 标志,它们利用了 DiffusionKit 仓库:

  • --sd3-version:指示转换器将其视为 Stable Diffusion 3 模型
  • --convert-mmdit:转换 MMDiT 模型
  • --convert-vae-decoder:转换新的 VAE 模型(若设置了 --sd3-version,将使用 16 通道版本)
  • --include-t5:下载并在转换中包含预转换的 T5 文本编码器

e.g.

python -m python_coreml_stable_diffusion.torch2coreml --model-version stabilityai/stable-diffusion-3-medium --bundle-resources-for-swift-cli --convert-vae-decoder --convert-mmdit  --include-t5 --sd3-version -o <output-dir>

要以 1024x1024 分辨率转换完整流水线,可使用以下命令:

python -m python_coreml_stable_diffusion.torch2coreml --model-version stabilityai/stable-diffusion-3-medium --bundle-resources-for-swift-cli --convert-text-encoder --convert-vae-decoder --convert-mmdit --include-t5 --sd3-version --latent-h 128 --latent-w 128 -o <output-dir>

请注意,MMDiT 模型体量较大,随着潜空间(latent)分辨率升高,转换所需的内存和时间会显著增加。

另外请注意,目前 MMDiT 模型需要 fp32,因此仅支持 CPU_AND_GPU 计算单元以及 ORIGINAL 注意力实现(本流水线的默认设置)。

Swift 推理

Stable Diffusion 3 的 Swift 推理与之前版本类似。唯一区别是应使用 --sd3 标志,以表明该模型为 Stable Diffusion 3 模型。

swift run StableDiffusionSample <prompt> --resource-path <output-mlpackages-directory/Resources> --output-path <output-dir> --compute-units cpuAndGPU --sd3

使用 Stable Diffusion XL

详情(点击展开)

模型转换

例如:

python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --convert-vae-decoder --convert-text-encoder --xl-version --model-version stabilityai/stable-diffusion-xl-base-1.0 --refiner-version stabilityai/stable-diffusion-xl-refiner-1.0 --bundle-resources-for-swift-cli --attention-implementation {ORIGINAL,SPLIT_EINSUM} -o <output-dir>
  • --xl-version:在指定 XL 模型时传给转换脚本的额外参数
  • --refiner-version:在指定 XL refiner 模型时传给转换脚本的额外参数,"Ensemble of Expert Denoisers" 推理时需要
  • --attention-implementation:在 Mac 上部署时,建议对 cpuAndGPU 使用 ORIGINAL
  • --attention-implementation:在 iPhone 与 iPad 上部署时,建议对 cpuAndNeuralEngine 使用 SPLIT_EINSUM
  • --attention-implementation:由于编译时间过长,不建议对 Stable Diffusion XL 使用 SPLIT_EINSUM_V2
  • 提示: 建议为 iOS 与 iPadOS 部署添加 --latent-h 96 --latent-w 96,这样会以 768x768 生成,而非默认的 1024x1024。
  • 提示: 由于原始 Stable Diffusion XL VAE 存在已知的 float16 溢出问题,模型转换脚本会强制使用 float32 精度. 通过 --custom-vae-version madebyollin/sdxl-vae-fp16-fix 使用自定义 VAE 版本(例如 @madebyollinmadebyollin/sdxl-vae-fp16-fix)可恢复 VAE 的默认 float16 精度。

Swift 推理

swift run StableDiffusionSample <prompt> --resource-path <output-mlpackages-directory/Resources> --output-path <output-dir> --compute-units {cpuAndGPU,cpuAndNeuralEngine} --xl
  • 仅需要 base 模型;refiner 模型为可选,若资源目录中提供则默认使用
  • XL 的 ControlNet 尚不支持

Python 推理

python -m python_coreml_stable_diffusion.pipeline --prompt <prompt> --compute-unit {CPU_AND_GPU,CPU_AND_NE} -o <output-dir> -i <output-mlpackages-directory/Resources> --model-version stabilityai/stable-diffusion-xl-base-1.0
  • refiner 模型尚不支持
  • XL 的 ControlNet 尚不支持

使用 ControlNet

详情(点击展开)

以下示例结果使用提示词 "a high quality photo of a surfing dog",并基于涂鸦(scribble,最左侧)进行条件生成:

ControlNet 允许用户将 Stable Diffusion 的图像生成条件化到边缘图、深度图、分割图、涂鸦和姿态等信号上。得益于 @ryu38 的贡献,Python CLI 与 Swift 包均支持 ControlNet 模型。关于如何将 Stable Diffusion 与 ControlNet 一起配置,请参阅本节

请注意,ControlNet 目前尚不支持 Stable Diffusion XL。

使用系统多语言文本编码器

详情(点击展开)

在 iOS 17 与 macOS 14 中,NaturalLanguage 框架引入了 NLContextualEmbedding,可为拉丁文(20 种语言)、西里尔文(4 种语言)以及 CJK(3 种语言)脚本提供基于 Transformer 的文本嵌入。WWDC23 专题演讲 Explore Natural Language multilingual models 演示了开发者如何利用这一强大的新模型训练下游任务,例如基于 Stable Diffusion 的多语言图像生成。

本仓库提供了复现该演示工作流的代码。该工作流有多种实现方式。示例如下:

步骤 1 整理包含目标语言的图文数据集。

步骤 2 预先计算 NLContextualEmbedding 值,并用这些嵌入向量替换数据集中的文本字符串。

步骤 3 从 Hugging Face Hub 选取与 StableDiffusionPipeline 兼容的基座模型,使用你的新数据集进行微调,并用预先计算的 NLContextualEmbedding 值替换默认的 text_encoder。

步骤 4 若要在不训练新层的情况下替换基座模型的 text_encoder,基座模型的 text_encoder.hidden_size 必须与 NLContextualEmbedding 一致。若不一致,则需要训练一个线性投影层以在两种维度之间映射。微调后,该线性层应按如下方式转换为 CoreML:

python -m python_coreml_stable_diffusion.multilingual_projection --input-path <path-to-projection-torchscript> --output-dir <output-dir>

上述命令会在 --output-dir 下生成 MultilingualTextEncoderProjection.mlmodelc 文件,该文件应与通过 --bundle-resources-for-swift-cli 生成的其余 Core ML 模型资源放在同一位置。

步骤 5 现在可通过在初始化流水线时将 useMultilingualTextEncoder 设为 true,或在 CLI 中设置 --use-multilingual-text-encoder,来调用多语言系统文本编码器。请注意,模型资源通过空中分发(over-the-air),首次调用会触发资源下载,体积小于 100MB。

资源:

使用 Hugging Face Hub 上现成的 Core ML 模型

点击展开

🤗 Hugging Face 已对以下模型执行了转换流程,并在 Hub 上公开提供了 Core ML 权重。若你希望转换 Hub 上尚不可用的 Stable Diffusion 版本,请参阅将模型转换为 Core ML

如果你想使用这些模型中的任意一个,可以下载权重,然后继续 使用 Python 生成图像使用 Swift 生成图像

每个模型仓库中都有多个变体。你可以使用 gitgit lfs 克隆整个仓库以下载所有变体,或选择性下载你需要的变体。

要使用 git 克隆仓库,请按以下流程操作:

步骤 1 为你的系统安装 git lfs 扩展。

git lfs 将大文件存储在主 git 仓库之外,在你克隆或检出后从相应服务器下载这些文件。它在大多数包管理器中均可用,详情请查看安装页面

步骤 2 运行以下命令一次以启用 git lfs

git lfs install

步骤 3 使用 git clone 下载包含所有模型变体的仓库副本。对于 Stable Diffusion 1.4 版本,你可以在终端中执行以下命令:

git clone https://huggingface.co/apple/coreml-stable-diffusion-v1-4

如果你希望下载特定变体而不是克隆整个仓库,可以使用 huggingface_hub Python 库。例如,要在 Python 中使用 ORIGINAL attention 实现进行生成(详情请阅读此节),可以使用以下辅助代码:

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}")

model_path 将是检查点保存在本地文件系统中的路径。更多详情,请参阅此文章

将模型转换为 Core ML

点击展开

步骤 1 创建 Python 环境并安装依赖:

conda create -n coreml_stable_diffusion python=3.8 -y
conda activate coreml_stable_diffusion
cd /path/to/cloned/ml-stable-diffusion/repository
pip install -e .

步骤 2 登录或注册你的 Hugging Face 账户,,生成 用户访问令牌(User Access Token,并在终端窗口中运行 huggingface-cli login 以使用该令牌配置 Hugging Face API 访问。

步骤 3Hugging Face Hub 上导航至你想使用的 Stable Diffusion 版本并接受其使用条款。默认模型版本为 CompVis/stable-diffusion-v1-4.。用户可按下一步所述更改模型版本。

步骤 4 在终端中执行以下命令以生成 Core ML 模型文件(.mlpackage

python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --convert-text-encoder --convert-vae-decoder --convert-safety-checker --model-version <model-version-string-from-hub> -o <output-mlpackages-directory>

警告: 此命令将从 Hugging Face 下载数 GB 的 PyTorch 检查点。请确保你使用的是 Wi-Fi 且拥有足够的磁盘空间。

在 M1 MacBook Pro 上通常需要 15-20 分钟。成功执行后,构成 Stable Diffusion 的 4 个神经网络模型将从 PyTorch 转换为 Core ML.mlpackage),并保存到指定的 <output-mlpackages-directory>。其他一些值得注意的参数:

  • --model-version:在 Hugging Face Hub 上发布的模型版本名称

  • --refiner-version:在 Hugging Face Hub. 上发布的 refiner 版本名称。此参数为可选;若指定,将转换 refiner unet 并与模型 unet 一并打包。

  • --bundle-resources-for-swift-cli:编译全部 4 个模型,并将文本分词所需资源一并打包到 <output-mlpackages-directory>/Resources 中,该文件应作为 Swift 包的输入。对于基于 diffusers 的 Python 流水线,此标志并非必需。不过,在 Python 中使用这些已编译模型将显著加快推理速度.

  • --quantize-nbits:使用全局最优 k-means 聚类算法将 unet 和 text_encoder 模型的权重量化至 2、4、6 或 8 比特。默认情况下,即使未指定此参数,所有模型也会被权重量化至 16 比特。详情请参阅此节及关于权重压缩的更多说明。

  • --chunk-unet:将 Unet 模型拆分为两个大致相等的块(每块权重均小于 1GB),以便在移动端部署。若权重未量化至 6 比特或更低(--quantize-nbits {2,4,6}),在 iOS 和 iPadOS 上部署到 Neural Engine 时必须使用此选项。macOS 上不需要。Swift CLI 可使用分块版和普通版 Unet 模型,但会优先使用前者。请注意,分块 unet 与 Python 流水线不兼容,因为 Python 流水线仅面向 macOS。

  • --attention-implementation:默认为 SPLIT_EINSUM,即 Deploying Transformers on the Apple Neural Engine. 中描述的实现。--attention-implementation SPLIT_EINSUM_V2 可为移动设备带来 10-30% 的性能提升,仍面向 Neural Engine。--attention-implementation ORIGINAL 将切换到另一种实现,适用于部分 Mac 设备上的 CPU 或 GPU 部署。更多说明请参阅 性能基准测试 一节。

  • --check-output-correctness:将原始 PyTorch 模型的输出与最终 Core ML 模型的输出进行对比。此标志会显著增加 RAM 占用,因此建议仅用于调试。

  • --convert-controlnet:转换此选项后指定的 ControlNet 模型。若像 --convert-controlnet lllyasviel/sd-controlnet-mlsd lllyasviel/sd-controlnet-depth 这样指定,也可转换多个模型。

  • --unet-support-controlnet:使转换后的 UNet 模型能够接收来自 ControlNet 的额外输入。使用 ControlNet 生成图像时需要此选项,并会以不同名称 *_control-unet.mlpackage 保存,与普通 UNet 不同。另一方面,该 UNet 模型在没有 ControlNet 的情况下无法工作。若仅需 txt2img,请使用普通 UNet。

  • --unet-batch-one:为 unet 使用批次大小(batch size)为 1;若你不想进行无分类器引导(classifier free guidance),即使用小于 1 的 guidance-scale,则需要此设置。

  • --convert-vae-encoder:文生图(text-to-image)应用不需要。图生图(image-to-image)应用需要,以便将输入图像映射到潜空间(latent space)。

使用 Python 进行图像生成

点击展开

使用基于 diffusers: 的示例 Python 流水线运行文生图生成。

python -m python_coreml_stable_diffusion.pipeline --prompt "a photo of an astronaut riding a horse on mars" -i <core-ml-model-directory> -o </path/to/output/image> --compute-unit ALL --seed 93

请参阅帮助菜单了解所有可用参数:python -m python_coreml_stable_diffusion.pipeline -h。一些值得注意的参数:

  • -i:应指向上文 将模型转换为 Core ML 章节第 4 步中的 -o 目录。若在转换期间指定了 --bundle-resources-for-swift-cli,则使用生成的 Resources 文件夹(其中包含已编译的 .mlmodelc 文件)。编译后的模型在首次使用后加载会快得多.
  • --model-version:若在将模型转换为 Core ML 时覆盖了默认模型版本,则需要在此处指定相同的模型版本。
  • --compute-unit:请注意,对于此特定实现,性能最佳的计算单元可能因硬件而异。CPU_AND_GPUCPU_AND_NE 可能比 ALL 更快。请参阅性能基准测试章节以获取更多指导。
  • --scheduler:若想尝试不同的调度器(scheduler),可在此处指定。可用选项请参阅帮助菜单。你还可以通过 --num-inference-steps 指定自定义推理步数,默认值为 50。
  • --controlnet:通过此选项指定的 ControlNet 模型用于图像生成。请按 --controlnet lllyasviel/sd-controlnet-mlsd lllyasviel/sd-controlnet-depth 格式使用此选项,并确保与 --controlnet-inputs 配合使用。
  • --controlnet-inputs:与各 ControlNet 模型对应的图像输入。请按 --controlnet 中模型的相同顺序提供图像路径,例如:--controlnet-inputs image_mlsd image_depth
  • --unet-batch-one:不要对提示词和负向提示词的 unet 预测进行批处理。这要求 unet 在转换时批次大小为 1,请参阅转换脚本中的 --unet-batch-one 选项。

使用 Swift 进行图像生成

点击展开

示例 CLI 用法

swift run StableDiffusionSample "a photo of an astronaut riding a horse on mars" --resource-path <output-mlpackages-directory>/Resources/ --seed 93 --output-path </path/to/output/image>

输出文件名将基于提示词和随机种子命名: 例如 </path/to/output/image>/a_photo_of_an_astronaut_riding_a_horse_on_mars.93.final.png

请使用 --help 标志了解批量生成等更多信息。

示例库用法

import StableDiffusion
...
let pipeline = try StableDiffusionPipeline(resourcesAt: resourceURL)
pipeline.loadResources()
let image = try pipeline.generateImages(prompt: prompt, seed: seed).first

在 iOS 上,构造 StableDiffusionPipeline 时,应将 reduceMemory 选项设置为 true

Swift 软件包详情

此 Swift 软件包包含两个产品:

  • StableDiffusion
  • StableDiffusionSample 命令行工具

这两个产品都需要提供 Core ML 模型和分词(tokenization)资源。通过目录路径指定资源时,该目录必须包含以下内容:

  • TextEncoder.mlmodelc 或 `TextEncoder2.mlmodelc(文本嵌入模型)
  • Unet.mlmodelcUnetChunk1.mlmodelcUnetChunk2.mlmodelc(去噪自编码器模型)
  • VAEDecoder.mlmodelc(图像解码器模型)
  • vocab.json(分词器词汇表文件)
  • merges.text(字节对编码(byte pair encoding)合并文件)

可选,用于图生图、修复(in-painting)或类似场景:

  • VAEEncoder.mlmodelc(图像编码器模型)

可选,还可包含某些版本 Stable Diffusion 所包含的安全检查器(safety checker)模型:

  • SafetyChecker.mlmodelc

可选,用于 SDXL refiner

  • UnetRefiner.mlmodelcrefiner unet 模型)

可选,用于 ControlNet

  • ControlledUNet.mlmodelcControlledUnetChunk1.mlmodelcControlledUnetChunk2.mlmodelc(已启用以接收 ControlNet 值)
  • controlnet/(包含 ControlNet 模型的目录)
    • LllyasvielSdControlnetMlsd.mlmodelc(例如,来自 lllyasviel/sd-controlnet-mlsd
    • LllyasvielSdControlnetDepth.mlmodelc(例如,来自 lllyasviel/sd-controlnet-depth
    • 你转换的其他模型

请注意,会首先检查 Unet 的分块(chunked)版本。仅当其不存在时才会加载完整的 Unet.mlmodelc。iOS 和 iPadOS 需要分块,macOS 则不需要。

示例 Swift 应用

点击展开

🤗 Hugging Face 在此库之上创建了一个开源演示应用。它使用原生 Swift 和 SwiftUI 编写,可在 macOS、iOS 和 iPadOS 上运行。你可以将代码作为应用的起点,或了解如何在自己的项目中集成此库。

Hugging Face 已将此应用上架 Mac App Store.

常见问题(FAQ

点击展开
Q1: ERROR: Failed building wheel for tokenizers or error: can't find Rust compiler

A1: 请查看此潜在解决方案.

Q2: RuntimeError: {NSLocalizedDescription = "Error computing NN outputs."

A2: 此错误可能有许多原因。在此场景下,当你的系统因其他应用而面临更大的内存压力时,很可能遇到该错误。降低其他应用的内存占用可能有助于缓解此问题。

Q3: 我的 Mac 有 8GB 内存,我正在使用示例命令将模型转换为 Core ML。进程因内存问题被终止。如何解决这个问题?

A3: 为尽量降低模型转换过程的内存占用,请改为执行以下命令:

python -m python_coreml_stable_diffusion.torch2coreml --convert-vae-encoder --model-version <model-version-string-from-hub> -o <output-mlpackages-directory> && \
python -m python_coreml_stable_diffusion.torch2coreml --convert-vae-decoder --model-version <model-version-string-from-hub> -o <output-mlpackages-directory> && \
python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --model-version <model-version-string-from-hub> -o <output-mlpackages-directory> && \
python -m python_coreml_stable_diffusion.torch2coreml --convert-text-encoder --model-version <model-version-string-from-hub> -o <output-mlpackages-directory> && \
python -m python_coreml_stable_diffusion.torch2coreml --convert-safety-checker --model-version <model-version-string-from-hub> -o <output-mlpackages-directory> &&

若你需要 --chunk-unet,可以在另一个独立命令中完成,该命令会复用先前导出的 Unet 模型并就地对其进行分块:

python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --chunk-unet -o <output-mlpackages-directory>
Q4: 我的 Mac 只有 8GB 内存,图像生成能在我的机器上运行吗?

A4: 可以!尤其是 --compute-unit CPU_AND_NE 选项,在其他应用程序负载合理的情况下应该可以正常工作。请注意,示例结果 中的部分内容是在配备 8GB 内存的 M2 MacBook Air 上生成的。

Q5: 每次使用 Python 流水线生成图像时,加载所有 Core ML 模型都需要 2-3 分钟。这是正常现象吗?

A5: 在指定特定计算单元(compute unit)后,.mlpackage.mlmodelc 模型在首次加载时都会进行编译(在 Core ML 术语中也称为“模型准备”,model preparation)。.mlpackage 不会缓存此编译产物,因此每次加载模型都会重新触发编译,可能耗时数分钟。另一方面,.mlmodelc 文件会缓存此编译产物,非首次加载时间可缩短至仅数秒。

若要受益于编译缓存,你可以在 Swift(默认)和 Python(得益于 @lopez-hector's contribution) 图像生成流水线)中,使用 .mlmodelc 资源替代 .mlpackage 资源。

Q6: 我想在移动应用中部署 StableDiffusion 这一 Swift 软件包。需要注意什么?

A6: 使用 Swift 进行图像生成 一节介绍了本软件包所需的最低 SDK 和操作系统版本,以及支持的设备型号。我们建议在部署目标中内存最少的设备上仔细测试该软件包。

StableDiffusion 中的图像生成过程在运行时,根据所选计算单元,峰值内存可能超过 2 GB。在 iPadOS 上,我们建议在配置中使用 .cpuAndNeuralEngine,并在构建 StableDiffusionPipeline 时选用 reduceMemory 选项,以尽量降低内存压力。

若应用在图像生成期间崩溃,可考虑添加 Increased Memory Limit 能力,以告知系统:在支持的设备上,应用的部分核心功能在超出默认内存限制时可能运行得更好。

在 iOS 上,根据 iPhone 型号、Stable Diffusion 模型版本、所选计算单元、系统负载以及应用设计,这可能仍不足以将应用的峰值内存控制在限制以内。请记住,由于设备在应用与 iOS 进程之间共享内存,单个应用占用过多内存可能影响整台设备的使用体验。

我们强烈建议按照高级权重压缩(低于 6 位) 中的方案压缩模型,以便在 iOS 上部署。这可将峰值 RAM 占用最多降低 75%(从 16 位降至 4 位),同时保持模型输出质量。

Q7: 如何使用同一套 Core ML 模型生成不同分辨率的图像?

A7: 当前版本的 python_coreml_stable_diffusion 并不开箱即用地支持单模型多分辨率。不过,开发者可以 fork 本项目,并利用 coremltools 的灵活形状 支持,通过 coremltools.EnumeratedShapes 扩展 torch2coreml 脚本。请注意,虽然 text_encoder 与图像分辨率无关,但 vae_decoderunet 模型的输入与输出取决于目标图像分辨率。

Q8: Core ML 与 PyTorch 生成的图像会完全相同吗?

A8: 如有需要,PyTorch 与 Core ML 生成的图像可以做到近似一致。但默认情况下并不保证如此。以下若干因素可能导致 PyTorch 与 Core ML 生成不同图像:

1. 随机数生成器行为

PyTorch 与 Core ML 结果可能不同的主要来源是随机数生成器(RNG) 行为。PyTorch 与 Numpy 的随机性来源不同。python_coreml_stable_diffusion 通常依赖 Numpy 进行 RNG(例如潜变量初始化),而 StableDiffusion Swift 库默认会复现该 RNG 行为。不过,基于 PyTorch 的流水线(如 Hugging Face diffusers)依赖 PyTorch 的 RNG 行为。得益于 @liuliu 的贡献,,可通过指定 --rng torch/cuda 来选择 torchRNG/cudaRNG 模式,从而在 Swift 中匹配 PyTorchCPU/GPU)的 RNG 行为。

2. PyTorch

“无法保证在 PyTorch 不同版本、各个提交或不同平台之间得到完全可复现的结果。此外,即使使用相同种子,CPU 与 GPU 执行的结果也可能不可复现。”来源).

3. 转换过程中的模型函数漂移

对应 PyTorch 与 Core ML 模型输出差异是潜在原因之一。转换过程中会测试信号完整性(通过向 python_coreml_stable_diffusion.torch2coreml 传入 --check-output-correctness 参数启用),并在随机输入上验证其高于最低 PSNR 值。请注意,这仅是健全性检查,并不能保证所有可能输入都达到该最低 PSNR。此外,在不同计算单元上执行同一 Core ML 模型时,结果也不保证完全相同。如本节 中的示例视觉结果所示,这通常不应成为主要差异来源。

4. 权重与激活的数据类型

将模型从 float32 量化到 float16 等较低精度数据类型时,即使使用同一 PyTorch 模型,生成图像在语义上也已知会略有差异。coremltools 生成的 Core ML 模型默认使用 float16 权重与激活除非显式覆盖.。这通常不应成为主要差异来源。

Q9: 模型文件非常大,如何避免应用二进制体积过大?

A9: 推荐做法是在应用首次启动时提示用户下载这些资源。这样应用二进制大小可与所部署的 Core ML 模型解耦。向用户披露下载大小极为重要,因为可能产生流量费用或占用存储空间,用户未必能够接受。

Q10: `Could not initialize NNPACK! Reason: Unsupported hardware`

A10: 在本仓库的上下文中,可以安全忽略此警告。

Q11: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect

A11: 在本仓库的上下文中,可以安全忽略此警告。

Q12: UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown

A12: 如果此警告紧跟在 zsh: killed python -m python_coreml_stable_diffusion.torch2coreml ... 之后出现,那么你的 Mac 极有可能在将模型转换为 Core ML 时耗尽了内存。请参阅上文中的 Q3 获取解决方案。

BibTeX 引用

@misc{stable-diffusion-coreml-apple-silicon,
title = {Stable Diffusion with Core ML on Apple Silicon},
author = {Atila Orhon and Michael Siracusa and Aseem Wadhwa},
year = {2022},
URL = {null}
}