Argmax Model Optimization Toolkit for Diffusion Models.
Project description
DiffusionKit
Run Diffusion Models on Apple Silicon with Core ML and MLX
This repository comprises
diffusionkit
, a Python package for converting PyTorch models to Core ML format and performing image generation with MLX in PythonDiffusionKit
, a Swift package for on-device inference of diffusion models using Core ML and MLX
Installation
The following installation steps are required for:
- MLX inference
- PyTorch to Core ML model conversion
Python Environment Setup
conda create -n diffusionkit python=3.11 -y
conda activate diffusionkit
cd /path/to/diffusionkit/repo
pip install -e .
Hugging Face Hub Credentials
Stable Diffusion 3 requires users to accept the terms before downloading the checkpoint. Once you accept the terms, sign in with your Hugging Face hub READ token as below:
huggingface-cli login --token YOUR_HF_HUB_TOKEN
Converting Models from PyTorch to Core ML
Click to expand
Step 1: Follow the installation steps from the previous section
Step 2: Prepare the denoise model (MMDiT) Core ML model files (.mlpackage
)
python -m tests.torch2coreml.test_mmdit --sd3-ckpt-path <path-to-sd3-mmdit.safetensors> --model-version {2b} -o <output-mlpackages-directory> --latent-size {64, 128}
Step 3: Prepare the VAE Decoder Core ML model files (.mlpackage
)
python -m tests.torch2coreml.test_vae --sd3-ckpt-path <path-to-sd3-mmdit.safetensors> -o <output-mlpackages-directory> --latent-size {64, 128}
Image Generation with Python MLX
Click to expand
For simple text-to-image in float16 precision:
diffusionkit-cli --prompt "a photo of a cat" --output-path </path/to/output/image.png> --seed 0 --w16 --a16
Some notable optional arguments:
- For image-to-image, use
--image-path
(path to input image) and--denoise
(value between 0. and 1.) - T5 text embeddings, use
--t5
- For different resolutions, use
--height
and--width
Please refer to the help menu for all available arguments: diffusionkit-cli -h
.
Image Generation with Swift
Click to expand
Core ML Swift
Apple Core ML Stable Diffusion is the initial Core ML backend for DiffusionKit. Stable Diffusion 3 support is upstreamed to that repository while we build the holistic Swift inference package.
MLX Swift
🚧
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for diffusionkit-0.1.0-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 921af144a9a34e963ce2b0c80cb4f7205e67f87691b844154cbeb640e02056f7 |
|
MD5 | 093e4a1c6a2e4dde6dc1ac8749a2dd28 |
|
BLAKE2b-256 | 761fae8b144191e7dbe48f2a7667b0bb8c15d07c5d31e8ae5b0e133c3700390c |