DiffusionKit enables on-device image generation for Apple Silicon, targeting developers and researchers who want to leverage diffusion models locally. It provides tools to convert PyTorch models to Core ML format and perform inference using MLX, offering a streamlined path for integrating advanced image generation capabilities into macOS and iOS applications.

How It Works

DiffusionKit facilitates the conversion of PyTorch diffusion models (like Stable Diffusion 3 and FLUX) into Apple's Core ML format, optimizing them for efficient on-device execution. It then utilizes the MLX framework, designed for machine learning on Apple hardware, to perform the actual image generation inference. This approach leverages Apple's specialized hardware (Neural Engine) for accelerated performance and privacy-preserving local processing.

Quick Start & Requirements

• Installation: pip install -e . within a Python 3.11 conda environment.
• Prerequisites: Hugging Face Hub credentials (requires accepting terms for gated models like SD3 and FLUX.1-dev).
• Model Conversion: Python scripts are provided to convert PyTorch checkpoints to .mlpackage format.
• CLI Usage: diffusionkit-cli --prompt "a photo of a cat" --output-path <output_path>
• Python Usage: Import DiffusionPipeline or FluxPipeline from diffusionkit.mlx.
• Swift Usage: Leverages ml-stable-diffusion for Core ML backend, with ongoing development for a holistic Swift package.
• Documentation: https://github.com/argmaxinc/DiffusionKit

Highlighted Details

• Supports on-device inference for Stable Diffusion 3 and FLUX models.
• Includes tools for converting PyTorch models to Core ML (.mlpackage).
• Offers both a Python package and a Swift package for inference.
• CLI and Python API for easy integration and experimentation.

Maintenance & Community

• Developed by Argmax, Inc.
• Citation information provided for academic use.

Licensing & Compatibility

• Released under the MIT License, permitting commercial use and integration into closed-source projects.

Limitations & Caveats

• Stable Diffusion 3 and FLUX.1-dev require accepting terms and granting specific Hugging Face token permissions for gated model access.
• The Swift package for MLX inference is marked as under development (🚧).