DiffusionKit：在 Mac 上跑 Diffusion 扩散模型的新选择

一款让 AI 画图"跑在本地"的工具

想象一下这样的场景：你有一台 M2 MacBook Air，想给文章配一张图，但不想把数据传到云端——隐私顾虑、网络延迟、API 费用，每一条都让人头疼。传统方案要么靠云端 API，要么在服务器上部署庞大的 PyTorch 环境，两条路都不轻松。

DiffusionKit 正是来解决这个问题的：它让你直接在 Apple Silicon Mac 上跑 Stable Diffusion 3、FLUX.1 等主流扩散模型，用的是苹果自研的 MLX 框架（对标 CUDA 的 Apple GPU 加速库）和 Core ML 推理引擎。702 颗 GitHub Stars，已经有 Piper 这样的开发者开始在生产环境用它做本地 AI 图像生成。

背景：为什么 Apple Silicon 需要专门的 AI 推理方案

Apple Silicon（M 系列芯片）在 CPU 和统一内存架构上有显著优势，但 AI 推理一直是短板——NVIDIA 有 CUDA，AMD 有 ROCm，而 Apple 需要自己的生态。2023 年苹果开源了 MLX（类似 NumPy 的数组框架，专为 Apple 芯片优化），DiffusionKit 就是建立在 MLX 之上的扩散模型推理引擎。

项目由 Argmax, Inc. 开发维护，这是一家专注于模型优化的公司，曾为多家科技公司提供模型压缩和部署服务。DiffusionKit 的核心价值是把 PyTorch 训练好的模型"翻译"成 Core ML 格式，然后在 Mac 上高效推理。整个工具链从模型转换到图像生成，形成了完整的闭环。

技术架构：Python + Swift 双生态

DiffusionKit 实际上包含两个子项目：

Python 层：diffusionkit 包

python/src/diffusionkit/
├── mlx/          # MLX 推理后端（主力）
│   ├── diffusionkit.mlx.clip.py      # CLIP 文本编码器
│   ├── diffusionkit.mlx.t5.py        # T5 编码器（SD3 专用）
│   ├── diffusionkit.mlx.mmdit.py     # MMDiT 主模型（SD3 核心）
│   ├── diffusionkit.mlx.vae.py       # VAE 解码器
│   ├── diffusionkit.mlx.tokenizer.py # 文本分词
│   ├── diffusionkit.mlx.sampler.py   # 采样器（调度）
│   └── diffusionkit.mlx.scripts/generate_images.py  # CLI 入口
├── torch/        # PyTorch 模型结构定义（用于加载权重）
└── tests/        # 转换和推理测试

推理流程（以 SD3 为例）：

1. 文本编码：用户输入 → CLIP（快速）或 T5（高质量）→ 向量嵌入
2. MMDiT 主推理：在潜空间（latent space）执行多轮去噪迭代，默认 50 步
3. VAE 解码：潜空间向量 → 还原为 512×512 RGB 图像

Swift 层：DiffusionKit Swift Package

Swift 包处于早期阶段（README 标注 🚧），主要用于纯 iOS/macOS 原生应用场景，无需 Python 环境即可集成扩散模型推理。

核心功能

CLI 使用示例

# 安装
pip install -e .
huggingface-cli login --token YOUR_HF_TOKEN

# 最简文生图
diffusionkit-cli --prompt "a photo of a cat" --output-path cat.png

# 图生图（保持构图，略改风格）
diffusionkit-cli --image-path input.png --denoise 0.7 --prompt "a photo of a cat"

# 高分辨率生成
diffusionkit-cli --height 1024 --width 1024 --prompt "landscape"

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,    # 使用 16 位激活
    w16=True,    # 使用 16 位权重
)

image, _ = pipeline.generate_image(
    "a photo of a cat",
    cfg_weight=5.0,
    num_steps=50,
    latent_size=(64, 64),  # 64*8 = 512
)
image.save("cat.png")

部署体验：macOS 用户友好，其他平台受限

优点：

• Apple Silicon 用户体验顺畅，MLX 自动调用 GPU 加速
• conda + pip 安装流程标准，文档清晰
• 提供预转换的 MLX 模型（argmaxinc/mlx-stable-diffusion-3-medium），跳过 Core ML 转换步骤
• 支持 MacBook 本地运行，隐私性好

缺点：

• Apple Silicon 专属：不支持 NVIDIA GPU、Linux 或 Windows，这是最大限制
• 无 Docker 支持：无法容器化部署，不适合服务器批量推理
• 无 Web UI：需要命令行或 Python 编程，对非技术用户不友好
• 首次部署耗时：需 HuggingFace 认证、协议授权、模型下载，首次约需 30-60 分钟

局限与争议

1. Apple Silicon 生态锁定：对于已有 NVIDIA 算力基础设施的团队，引入 DiffusionKit 意味着引入新的硬件依赖
2. Swift 包尚在开发：README 标注 🚧，iOS/macOS 原生集成还不成熟
3. 性能对比不明：官方未提供与 PyTorch 在 NVIDIA GPU 上的性能基准测试，实际生成速度取决于具体模型和分辨率
4. 授权协议门槛：SD3 和 FLUX.1 都需要用户在 HuggingFace 上单独同意授权协议，不可跳过

行业意义与趋势

DiffusionKit 代表了一个重要趋势：端侧 AI 推理的崛起。随着 M3/M4 芯片内存容量提升（最高 128GB 统一内存），本地运行 7B 参数模型已成为现实。DiffusionKit 通过 MLX 把苹果芯片的 AI 算力直接暴露给开发者，是 Apple 芯片 AI 开源生态的重要拼图。

对于内容创作者、独立开发者和隐私敏感型应用，DiffusionKit 提供了一个无需云端、随用随跑的 AI 画图方案。随着苹果持续投资 MLX 框架，这类工具的能力边界还会继续扩展。