Native Stable Diffusion inference on iOS / macOS using MPSGraph
MIT License
Native Diffusion runs Stable Diffusion models locally on macOS / iOS devices, in Swift, using the MPSGraph framework (not Python).
This is the Swift Package Manager wrapper of Maple Diffusion. It adds image-to-image, Swift Package Manager package, and convenient ways to use the code, like Combine publishers and async/await versions. It also supports downloading weights from any local or remote URL, including the app bundle itself.
Would not be possible without
Get started in 10 minutes
Make it do what you want
Built to be fun to code with
Built for end-user speed and great user experience
Commercial use allowed
In its simplest form it's as simple as one line:
let image = try? await Diffusion.generate(localOrRemote: modelUrl, prompt: "cat astronaut")
You can give it a local or remote URL or both. If remote, the downloaded weights are saved for later.
The single line version is currently limited in terms of parameters.
See examples/SingleLineDiffusion
for a working example.
Let's add some UI. Here's an entire working image generator app in a single SwiftUI view:
struct ContentView: View {
// 1
@StateObject var sd = Diffusion()
@State var prompt = ""
@State var image : CGImage?
@State var imagePublisher = Diffusion.placeholderPublisher
@State var progress : Double = 0
var anyProgress : Double { sd.loadingProgress < 1 ? sd.loadingProgress : progress }
var body: some View {
VStack {
DiffusionImage(image: $image, progress: $progress)
Spacer()
TextField("Prompt", text: $prompt)
// 3
.onSubmit { self.imagePublisher = sd.generate(prompt: prompt) }
.disabled(!sd.isModelReady)
ProgressView(value: anyProgress)
.opacity(anyProgress == 1 || anyProgress == 0 ? 0 : 1)
}
.task {
// 2
let path = URL(string: "http://localhost:8080/Diffusion.zip")!
try! await sd.prepModels(remoteURL: path)
}
// 4
.onReceive(imagePublisher) { r in
self.image = r.image
self.progress = r.progress
}
.frame(minWidth: 200, minHeight: 200)
}
}
Here's what it does
Diffusion
objectSee examples/SimpleDiffusion
for a working example.
DiffusionImage
An optional SwiftUI view that is specialized for diffusion:
Add https://github.com/mortenjust/native-diffusion
in the "Swift Package Manager" tab in Xcode
Native Diffusion splits the weights into a binary format that is different from the typical CKPT format. It uses many small files which it then (optionally) swaps in and out of memory, enabling it to run on both macOS and iOS. You can use the converter script in the package to convert your own CKPT file.
By downloading this zip file, you accept the creative license from StabilityAI.
Download ZIP. Please don't use this URL in your software.
We'll get back to what to do with it in a second.
ckpt
fileDownload a Stable Diffusion model checkpoint to a folder, e.g. ~/Downloads/sd
(sd-v1-5.ckpt
, or some derivation thereof)
Setup & install Python with PyTorch, if you haven't already.
# Grab the converter script
cd ~/Downloads/sd
curl https://raw.githubusercontent.com/mortenjust/maple-diffusion/main/Converter%20Script/maple-convert.py > maple-convert.py
# may need to install conda first https://github.com/conda-forge/miniforge#homebrew
conda deactivate
conda remove -n native-diffusion --all
conda create -n native-diffusion python=3.10
conda activate native-diffusion
pip install torch typing_extensions numpy Pillow requests pytorch_lightning
./native-convert.py ~/Downloads/sd-v1-4.ckpt
The script will create a new folder called bins
. We'll get back to what to do with it in a second.
Yes. Just copy the alpha*
files from the standard conversion. This repo will include these files in the future. See this issue.
Yes. Simply pass in an initImage
to your SampleInput
when generating.
You may need to regenerate the model files with the python script in the repo. This happens if you converted your ckpt model file before we added image2image.
Yes! A rough roadmap:
See Issues for smaller contributions.
If you're making changes to the MPSGraph part of the codebase, consider making your contributions to the single-file repo and then integrate the changes in the wrapped file in this repo.
On my MacBook Pro M1 Max, I get ~0.3s/step, which is significantly faster than any Python/PyTorch/Tensorflow installation I've tried.
On an iPhone it should take a minute or two.
To attain usable performance without tripping over iOS's 4GB memory limit, Native Diffusion relies internally on FP16 (NHWC) tensors, operator fusion from MPSGraph, and a truly pitiable degree of swapping models to device storage.
Not yet. Would love some help on this. See above.
Feel free to post an issue!