DiffusionKit: 在苹果芯片上实现扩散模型的本地推理

DiffusionKit:在苹果芯片上释放扩散模型的力量

在人工智能和机器学习领域,扩散模型已经成为生成高质量图像的主流技术。然而,这些模型通常需要强大的GPU资源才能高效运行。随着苹果芯片的出现,在本地设备上运行复杂的AI模型成为可能。DiffusionKit应运而生,它是一个专为苹果芯片优化的扩散模型推理工具包,旨在让开发者和研究人员能够在Mac电脑上轻松运行和部署扩散模型。

DiffusionKit的核心功能

DiffusionKit主要由两个核心组件构成:

diffusionkit: 这是一个Python包,其主要功能包括:
- 将PyTorch模型转换为Core ML格式
- 使用MLX(苹果的机器学习加速框架)在Python环境中进行图像生成
DiffusionKit: 这是一个Swift包,专门用于:
- 使用Core ML和MLX在设备上进行扩散模型的推理

这种双重结构使得DiffusionKit能够满足不同场景的需求,无论是模型转换、实验还是在iOS应用中集成。

安装与环境配置

要开始使用DiffusionKit,首先需要设置适当的Python环境:

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

对于需要使用Stable Diffusion 3等模型的用户,还需要配置Hugging Face Hub的凭证:

huggingface-cli login --token YOUR_HF_HUB_TOKEN

注意:使用细粒度token时,需要额外授权"Read access to contents of all public gated repos you can access"。

从PyTorch到Core ML的模型转换

DiffusionKit提供了一套完整的工具,用于将PyTorch模型转换为Core ML格式。这个过程主要包括以下步骤:

准备去噪模型(MMDiT)的Core ML文件:

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}

准备VAE解码器的Core ML文件:

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}

这些命令会生成.mlpackage文件,这些文件可以直接在Core ML环境中使用。

使用Python MLX进行图像生成

DiffusionKit提供了简单易用的CLI和Python API,方便用户进行图像生成。

CLI使用示例

最基本的使用方式如下:

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

此外,CLI还支持多种高级选项,如设置随机种子、进行图像到图像的转换、调整图像分辨率等。

Python代码示例

对于需要更多控制的场景,可以直接使用Python API:

from diffusionkit.mlx import DiffusionPipeline

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

HEIGHT = 512
WIDTH = 512
NUM_STEPS = 50
CFG_WEIGHT = 5.0

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

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