Multi-Model Selection

* Added selection of available models via the UI
* Overhauled the README to reflect all the new changes

Author: FahimF

Diffusion/Support/AppState.swift

@@ -24,6 +24,14 @@ class AppState: ObservableObject {
 		}
 	}
 
+	var currentModel: String = "" {
+		didSet {
+			Task {
+				await load(model: currentModel)
+			}
+		}
+	}
+	
 	private init() {
 		NSLog("*** AppState initialized")
 		// Does the model path exist?
@@ -53,19 +61,10 @@ class AppState: ObservableObject {
 			state = .error("No models found under model directory: \(dir.path)")
 			return
 		}
-		Task {
-			do {
-				try await load(model: model)
-			} catch {
-				NSLog("Error loading model: \(error)")
-				DispatchQueue.main.async {
-					self.state = .error(error.localizedDescription)
-				}
-			}
-		}
+		currentModel = model
 	}
 
-	func load(model: String) async throws {
+	func load(model: String) async {
 		NSLog("*** Loading model: \(model)")
 		let dir = modelDir.appending(component: model, directoryHint: .isDirectory)
 		let fm = FileManager.default
@@ -80,11 +79,18 @@ class AppState: ObservableObject {
 		// .all works for v1.4, but not for v1.5
 		configuration.computeUnits = .cpuAndGPU
 		// TODO: measure performance on different devices
-		let pipeline = try StableDiffusionPipeline(resourcesAt: dir, configuration: configuration, disableSafety: true)
-		NSLog("Pipeline loaded in \(Date().timeIntervalSince(beginDate))")
-		DispatchQueue.main.async {
-			self.pipeline = Pipeline(pipeline)
-			self.state = .ready("Ready")
+		do {
+			let pipeline = try StableDiffusionPipeline(resourcesAt: dir, configuration: configuration, disableSafety: true)
+			NSLog("Pipeline loaded in \(Date().timeIntervalSince(beginDate))")
+			DispatchQueue.main.async {
+				self.pipeline = Pipeline(pipeline)
+				self.state = .ready("Ready")
+			}
+		} catch {
+			NSLog("Error loading model: \(error)")
+			DispatchQueue.main.async {
+				self.state = .error(error.localizedDescription)
+			}
 		}
 	}
 }

Diffusion/Views/MainAppView.swift

@@ -71,39 +71,59 @@ struct MainAppView: View {
                 .buttonStyle(.borderedProminent)
 				.disabled(isBusy)
             }
+			
 			Spacer().frame(height: 16)
+			
 			HStack(alignment: .top) {
 				VStack(alignment: .leading) {
 					Group {
+						Picker("Model", selection: $context.currentModel) {
+							ForEach(context.models, id: \.self) { s in
+								Text(s).tag(s)
+							}
+						}
+						
+						Spacer().frame(height: 16)
+						
 						Picker("Scheduler", selection: $scheduler) {
 							ForEach(StableDiffusionScheduler.allCases, id: \.self) { s in
 								Text(s.rawValue).tag(s)
 							}
 						}
-						Text("")
+						
+						Spacer().frame(height: 16)
+						
 						Text("Guidance Scale: \(String(format: "%.1f", guidance))")
 						Slider(value: $guidance, in: 0...15, step: 0.1, label: {},
 							minimumValueLabel: {Text("0")},
 							maximumValueLabel: {Text("15")})
-						Text("")
+						
+						Spacer().frame(height: 16)
 					}
 					Group {
 						Text("Number of Inference Steps: \(String(format: "%.0f", steps))")
 						Slider(value: $steps, in: 1...300, step: 1, label: {},
 							minimumValueLabel: {Text("1")},
 							maximumValueLabel: {Text("300")})
-						Text("")
+						
+						Spacer().frame(height: 16)
+						
 						Text("Number of Images: \(String(format: "%.0f", numImages))")
 						Slider(value: $numImages, in: 1...8, step: 1, label: {},
 							minimumValueLabel: {Text("1")},
 							maximumValueLabel: {Text("8")})
-						Text("")
+						
+						Spacer().frame(height: 16)
+					}
+					Group {
+						Text("Safety Check On?")
+						Toggle("", isOn: $safetyOn)
+						
+						Spacer().frame(height: 16)
+						
+						Text("Seed")
+						TextField("", value: $seed, format: .number)
 					}
-					Text("Safety Check On?")
-					Toggle("", isOn: $safetyOn)
-					Text("")
-					Text("Seed")
-					TextField("", value: $seed, format: .number)
 //					Group {
 //						Text("Image Width")
 //						Slider(value: $width, in: 64...2048, step: 8, label: {},

README.md

@@ -1,14 +1,10 @@
 #  Diffusion
 
-This is a simple app that shows how to integrate Apple's [Core ML Stable Diffusion implementation](https://github.com/apple/ml-stable-diffusion) in a native Swift UI application. It can be used for faster iteration, or as sample code for other use cases.
+This app integrates Apple's [Core ML Stable Diffusion implementation](https://github.com/apple/ml-stable-diffusion) in a native Swift UI application. It can be used for simple image generation, or as a test platform for trying out new features and functionality. You'd have to compile the code yourself though since there is no binary release — but compiling the code should be fairly straightforward since everything is self-contained 🙂
 
 This is what it looks like:
 ![App Screenshot](assets/screenshot.jpg)
 
-On first launch, the application downloads a zipped archive with a Core ML version of Runway's Stable Diffusion v1.5, from [this location in the Hugging Face Hub](https://huggingface.co/pcuenq/coreml-stable-diffusion/tree/main). This process takes a while, as several GB of data have to be downloaded and unarchived.
-
-For faster inference, we use a very fast scheduler: [DPM-Solver++](https://github.com/LuChengTHU/dpm-solver) that we ported to Swift. Since this scheduler is still not available in Apple's GitHub repository, the application depends on the following fork instead: https://github.com/pcuenca/ml-stable-diffusion. Our Swift port is based on [Diffusers' DPMSolverMultistepScheduler](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py), with a number of simplifications.
-
 ## Installing
 
 ### The Project
@@ -38,14 +34,21 @@ The plan is to add the ability to switch between multiple modesl by selecting th
 
 ## Limitations
 
-- The UI does not expose a way to configure the scheduler, number of inference steps, or generation seed. These are all available in the underlying code.
-- A single model (Stable Diffusion v1.5) is considered. The Core ML compute units have been hardcoded to CPU and GPU, since that's what gives best results on my Mac (M1 Max MacBook Pro).
+- ~~The UI does not expose a way to configure the scheduler, number of inference steps, or generation seed. These are all available in the underlying code.~~
+- ~~A single model (Stable Diffusion v1.5) is considered.~~ 
+- The Core ML compute units have been hardcoded to CPU and GPU, since that's what gives best results on my Mac (M1 Max MacBook Pro).
 - Sometimes generation returns a `nil` image. This needs to be investigated.
 
 ## Next Steps
 
 - ~~Improve UI. Allow the user to select generation parameters.~~
 - ~~Implement negative prompts.~~
 - Implement other interesting schedulers.
-- Allow other models to run. Provide a recommended "compute units" configuration based on model and platform.
-- Explore other features (image to image, for example).
+- ~~Allow other models to run.~~ 
+- Provide a recommended "compute units" configuration based on model and platform.
+- Explore other features (image to image, for example)
+
+## Credit
+
+For faster inference, the project uses a very fast scheduler: [DPM-Solver++](https://github.com/LuChengTHU/dpm-solver) that was ported to Swift by the Hugging Face team/[Pedro Cuenca](https://github.com/pcuenca/). Since this scheduler is still not available in Apple's GitHub repository, the application depends on the following fork instead: https://github.com/pcuenca/ml-stable-diffusion. The Swift port is based on [Diffusers' DPMSolverMultistepScheduler](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py), with a number of simplifications.
+