DiffusionKit

在 Apple Silicon 上使用 Core ML 和 MLX 运行扩散模型

本仓库包括：

diffusionkit，一个用于将 PyTorch 模型转换为 Core ML 格式并使用 Python 中的 MLX 进行图像生成的 Python 包
DiffusionKit，一个用于使用 Core ML 和 MLX 在设备上进行扩散模型推理的 Swift 包

安装

以下安装步骤适用于：

MLX 推理
PyTorch 到 Core ML 模型转换

Python 环境设置

conda create -n diffusionkit python=3.11 -y
conda activate diffusionkit
cd /path/to/diffusionkit/repo
pip install -e .

Hugging Face Hub 凭证

点击展开

Stable Diffusion 3 要求用户在下载检查点之前接受条款。一旦您接受了条款，请使用您的 Hugging Face hub READ 令牌登录，如下所示：

[!重要] 如果使用细粒度令牌，还需要编辑权限以允许"读取您可以访问的所有公共受限仓库的内容"

huggingface-cli login --token YOUR_HF_HUB_TOKEN

将模型从 PyTorch 转换为 Core ML

点击展开

步骤 1: 按照上一节的安装步骤操作

步骤 2: 确认您已接受 StabilityAI 许可条款并在 HuggingFace 令牌上允许受限访问

步骤 3: 准备去噪模型 (MMDiT) Core ML 模型文件 (.mlpackage)

python -m python.src.diffusionkit.tests.torch2coreml.test_mmdit --sd3-ckpt-path stabilityai/stable-diffusion-3-medium --model-version 2b -o <output-mlpackages-directory> --latent-size {64, 128}

步骤 4: 准备 VAE 解码器 Core ML 模型文件 (.mlpackage)

python -m python.src.diffusionkit.tests.torch2coreml.test_vae --sd3-ckpt-path stabilityai/stable-diffusion-3-medium -o <output-mlpackages-directory> --latent-size {64, 128}

注意：

--sd3-ckpt-path 可以是任何 HuggingFace 仓库的路径（例如 stabilityai/stable-diffusion-3-medium）或本地 sd3_medium.safetensors 文件的路径

使用 Python MLX 生成图像

点击展开

命令行界面

最简单的方式：

diffusionkit-cli --prompt "a photo of a cat" --output-path </path/to/output/image.png>

一些值得注意的可选参数：

用于结果的可重复性，使用 --seed
用于图像到图像，使用 --image-path（输入图像的路径）和 --denoise（0 到 1 之间的值）
在 SD3 中启用 T5 编码器，使用 --t5（FLUX 必须使用 T5，无论此参数如何）
不同的分辨率，使用 --height 和 --width
使用本地检查点，使用 --local-ckpt </path/to/ckpt.safetensors>（例如 ~/models/stable-diffusion-3-medium/sd3_medium.safetensors）

请参考帮助菜单了解所有可用参数：diffusionkit-cli -h。

代码

对于 Stable Diffusion 3：

from diffusionkit.mlx import DiffusionPipeline
pipeline = DiffusionPipeline(
  model="argmaxinc/stable-diffusion",
  shift=3.0,
  use_t5=False,
  model_version="stable-diffusion-3-medium",
  low_memory_mode=True,
  a16=True,
  w16=True,
)

对于 FLUX：

from diffusionkit.mlx import FluxPipeline
pipeline = FluxPipeline(
  model="argmaxinc/stable-diffusion",
  shift=1.0,
  model_version="FLUX.1-schnell",
  low_memory_mode=True,
  a16=True,
  w16=True,
)

最后，要生成图像，使用 generate_image() 函数：

HEIGHT = 512
WIDTH = 512
NUM_STEPS = 4  #  FLUX.1-schnell 为 4，SD3 为 50
CFG_WEIGHT = 0. # FLUX.1-schnell 为 0.，SD3 为 5.

image, _ = pipeline.generate_image(
  "a photo of a cat",
  cfg_weight=CFG_WEIGHT,
  num_steps=NUM_STEPS,
  latent_size=(HEIGHT // 8, WIDTH // 8),
)

一些值得注意的可选参数：

对于图像到图像，使用 image_path（输入图像的路径）和 denoise（0 到 1 之间的值）输入变量。
对于种子，使用 seed 输入变量。
对于负面提示，使用 negative_text 输入变量。

生成的 image 可以通过以下方式保存：

image.save("path/to/save.png")

使用 Swift 生成图像

点击展开

Core ML Swift

Apple Core ML Stable Diffusion 是 DiffusionKit 的初始 Core ML 后端。Stable Diffusion 3 的支持已上传到该仓库，同时我们正在构建全面的 Swift 推理包。

MLX Swift

🚧

许可证

DiffusionKit 根据 MIT 许可证发布。有关更多详细信息，请参阅 LICENSE。

引用

如果您使用 DiffusionKit 做了一些很酷的事情或只是觉得它很有用，请给我们发送一个消息至 info@takeargmax.com！

如果您将 DiffusionKit 用于学术工作，以下是 BibTeX：

@misc{diffusionkit-argmax,
   title = {DiffusionKit},
   author = {Argmax, Inc.},
   year = {2024},
   URL = {https://github.com/argmaxinc/DiffusionKit}
}