Benchmark replication

Diffusion-macOS/Capabilities.swift

@@ -10,6 +10,8 @@ import Foundation
 
 let runningOnMac = true
 let deviceHas6GBOrMore = true
+let deviceHas8GBOrMore = true
+let BENCHMARK = false
 
 let deviceSupportsQuantization = {
     if #available(macOS 14, *) {

Diffusion-macOS/ContentView.swift

@@ -66,7 +66,7 @@ struct ContentView: View {
     @StateObject var generation = GenerationContext()
 
     func toolbar() -> any View {
-        if case .complete(let prompt, let cgImage, _, _) = generation.state, let cgImage = cgImage {
+        if case .complete(let prompt, let cgImage, _, _, _) = generation.state, let cgImage = cgImage {
             // TODO: share seed too
             return ShareButtons(image: cgImage, name: prompt)
         } else {

Diffusion-macOS/GeneratedImageView.swift

@@ -42,7 +42,7 @@ struct GeneratedImageView: View {
                     .buttonStyle(.plain)
                 }
             })
-        case .complete(_, let image, _, _):
+        case .complete(_, let image, _, _, _):
             guard let theImage = image else {
                 return AnyView(Image(systemName: "exclamationmark.triangle").resizable())
             }

Diffusion-macOS/StatusView.swift

@@ -23,7 +23,7 @@ struct StatusView: View {
                 if result.userCanceled {
                     generation.state = .userCanceled
                 } else {
-                    generation.state = .complete(generation.positivePrompt, result.image, result.lastSeed, result.interval)
+                    generation.state = .complete(generation.positivePrompt, result.image, result.lastSeed, result.interval, result.itsPerSecond)
                 }
             } catch {
                 generation.state = .failed(error)
@@ -75,7 +75,7 @@ struct StatusView: View {
                 Text("Generating \(Int(round(100*fraction)))%")
                 Spacer()
             }
-        case .complete(_, let image, let lastSeed, let interval):
+        case .complete(_, let image, let lastSeed, let interval, let _):
             guard let _ = image else {
                 return HStack {
                     Text("Safety checker triggered, please try a different prompt or seed.")

Diffusion.xcodeproj/project.pbxproj

@@ -16,9 +16,16 @@
 		8CEEB7D92A54C88C00C23829 /* DiffusionImage.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8CEEB7D82A54C88C00C23829 /* DiffusionImage.swift */; };
 		8CEEB7DA2A54C88C00C23829 /* DiffusionImage.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8CEEB7D82A54C88C00C23829 /* DiffusionImage.swift */; };
 		EB067F872992E561004D1AD9 /* HelpContent.swift in Sources */ = {isa = PBXBuildFile; fileRef = EB067F862992E561004D1AD9 /* HelpContent.swift */; };
-		EB25B3D62A3A2DC4000E25A1 /* StableDiffusion in Frameworks */ = {isa = PBXBuildFile; productRef = EB25B3D52A3A2DC4000E25A1 /* StableDiffusion */; };
-		EB25B3D82A3A2DD5000E25A1 /* StableDiffusion in Frameworks */ = {isa = PBXBuildFile; productRef = EB25B3D72A3A2DD5000E25A1 /* StableDiffusion */; };
 		EB560F0429A3C20800C0F8B8 /* Capabilities.swift in Sources */ = {isa = PBXBuildFile; fileRef = EB560F0329A3C20800C0F8B8 /* Capabilities.swift */; };
+		EB9BE0152AC425D300468918 /* StableDiffusion in Frameworks */ = {isa = PBXBuildFile; productRef = EB9BE0142AC425D300468918 /* StableDiffusion */; };
+		EB9BE0172AC425DA00468918 /* StableDiffusion in Frameworks */ = {isa = PBXBuildFile; productRef = EB9BE0162AC425DA00468918 /* StableDiffusion */; };
+		EBB56B382AB87932005361A6 /* merges.txt in Resources */ = {isa = PBXBuildFile; fileRef = EBB56B342AB87559005361A6 /* merges.txt */; };
+		EBB56B392AB87932005361A6 /* UnetChunk2.mlmodelc in Resources */ = {isa = PBXBuildFile; fileRef = EBB56B302AB87559005361A6 /* UnetChunk2.mlmodelc */; };
+		EBB56B3A2AB87932005361A6 /* UnetChunk1.mlmodelc in Resources */ = {isa = PBXBuildFile; fileRef = EBB56B362AB87559005361A6 /* UnetChunk1.mlmodelc */; };
+		EBB56B3B2AB87932005361A6 /* VAEDecoder.mlmodelc in Resources */ = {isa = PBXBuildFile; fileRef = EBB56B332AB87559005361A6 /* VAEDecoder.mlmodelc */; };
+		EBB56B3C2AB87932005361A6 /* vocab.json in Resources */ = {isa = PBXBuildFile; fileRef = EBB56B352AB87559005361A6 /* vocab.json */; };
+		EBB56B3D2AB87932005361A6 /* TextEncoder.mlmodelc in Resources */ = {isa = PBXBuildFile; fileRef = EBB56B312AB87559005361A6 /* TextEncoder.mlmodelc */; };
+		EBB56B3F2AB87932005361A6 /* TextEncoder2.mlmodelc in Resources */ = {isa = PBXBuildFile; fileRef = EBB56B322AB87559005361A6 /* TextEncoder2.mlmodelc */; };
 		EBB5BA5329425BEE003A2A5B /* PipelineLoader.swift in Sources */ = {isa = PBXBuildFile; fileRef = EBB5BA5229425BEE003A2A5B /* PipelineLoader.swift */; };
 		EBB5BA5A29426E06003A2A5B /* Downloader.swift in Sources */ = {isa = PBXBuildFile; fileRef = EBB5BA5929426E06003A2A5B /* Downloader.swift */; };
 		EBB5BA5D294504DE003A2A5B /* ZIPFoundation in Frameworks */ = {isa = PBXBuildFile; productRef = EBB5BA5C294504DE003A2A5B /* ZIPFoundation */; };
@@ -79,6 +86,13 @@
 		EB33A51E2954E1BC00B16357 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist; path = Info.plist; sourceTree = "<group>"; };
 		EB560F0329A3C20800C0F8B8 /* Capabilities.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Capabilities.swift; sourceTree = "<group>"; };
 		EB560F0529A4090D00C0F8B8 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist; path = Info.plist; sourceTree = "<group>"; };
+		EBB56B302AB87559005361A6 /* UnetChunk2.mlmodelc */ = {isa = PBXFileReference; lastKnownFileType = wrapper; path = UnetChunk2.mlmodelc; sourceTree = "<group>"; };
+		EBB56B312AB87559005361A6 /* TextEncoder.mlmodelc */ = {isa = PBXFileReference; lastKnownFileType = wrapper; path = TextEncoder.mlmodelc; sourceTree = "<group>"; };
+		EBB56B322AB87559005361A6 /* TextEncoder2.mlmodelc */ = {isa = PBXFileReference; lastKnownFileType = wrapper; path = TextEncoder2.mlmodelc; sourceTree = "<group>"; };
+		EBB56B332AB87559005361A6 /* VAEDecoder.mlmodelc */ = {isa = PBXFileReference; lastKnownFileType = wrapper; path = VAEDecoder.mlmodelc; sourceTree = "<group>"; };
+		EBB56B342AB87559005361A6 /* merges.txt */ = {isa = PBXFileReference; lastKnownFileType = text; path = merges.txt; sourceTree = "<group>"; };
+		EBB56B352AB87559005361A6 /* vocab.json */ = {isa = PBXFileReference; lastKnownFileType = text.json; path = vocab.json; sourceTree = "<group>"; };
+		EBB56B362AB87559005361A6 /* UnetChunk1.mlmodelc */ = {isa = PBXFileReference; lastKnownFileType = wrapper; path = UnetChunk1.mlmodelc; sourceTree = "<group>"; };
 		EBB5BA5229425BEE003A2A5B /* PipelineLoader.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = PipelineLoader.swift; sourceTree = "<group>"; };
 		EBB5BA5929426E06003A2A5B /* Downloader.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Downloader.swift; sourceTree = "<group>"; };
 		EBDD7DB22973200200C1C4B2 /* Utils.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Utils.swift; sourceTree = "<group>"; };
@@ -116,8 +130,8 @@
 			isa = PBXFrameworksBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
-				EB25B3D62A3A2DC4000E25A1 /* StableDiffusion in Frameworks */,
 				EBB5BA5D294504DE003A2A5B /* ZIPFoundation in Frameworks */,
+				EB9BE0152AC425D300468918 /* StableDiffusion in Frameworks */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
@@ -140,7 +154,7 @@
 			buildActionMask = 2147483647;
 			files = (
 				F155203C297118E700DC009B /* CompactSlider in Frameworks */,
-				EB25B3D82A3A2DD5000E25A1 /* StableDiffusion in Frameworks */,
+				EB9BE0172AC425DA00468918 /* StableDiffusion in Frameworks */,
 				EBDD7DAF29731FB300C1C4B2 /* ZIPFoundation in Frameworks */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
@@ -171,6 +185,20 @@
 			path = Diffusion/Common;
 			sourceTree = "<group>";
 		};
+		EBB56B212AB874C0005361A6 /* Model */ = {
+			isa = PBXGroup;
+			children = (
+				EBB56B342AB87559005361A6 /* merges.txt */,
+				EBB56B312AB87559005361A6 /* TextEncoder.mlmodelc */,
+				EBB56B322AB87559005361A6 /* TextEncoder2.mlmodelc */,
+				EBB56B362AB87559005361A6 /* UnetChunk1.mlmodelc */,
+				EBB56B302AB87559005361A6 /* UnetChunk2.mlmodelc */,
+				EBB56B332AB87559005361A6 /* VAEDecoder.mlmodelc */,
+				EBB56B352AB87559005361A6 /* vocab.json */,
+			);
+			path = Model;
+			sourceTree = "<group>";
+		};
 		EBB5BA5129425B07003A2A5B /* Pipeline */ = {
 			isa = PBXGroup;
 			children = (
@@ -187,6 +215,7 @@
 				EBDD7DBA2976F03600C1C4B2 /* debug.xcconfig */,
 				EBE4438729488DCA00CDA605 /* README.md */,
 				EBE443892948953600CDA605 /* LICENSE */,
+				EBB56B212AB874C0005361A6 /* Model */,
 				EBE755FF293E910800806B32 /* Packages */,
 				8CF53E022A44AE0400E6358B /* Common */,
 				EBE755C7293E37DD00806B32 /* Diffusion */,
@@ -318,7 +347,7 @@
 			name = Diffusion;
 			packageProductDependencies = (
 				EBB5BA5C294504DE003A2A5B /* ZIPFoundation */,
-				EB25B3D52A3A2DC4000E25A1 /* StableDiffusion */,
+				EB9BE0142AC425D300468918 /* StableDiffusion */,
 			);
 			productName = Diffusion;
 			productReference = EBE755C5293E37DD00806B32 /* Diffusion.app */;
@@ -378,7 +407,7 @@
 			packageProductDependencies = (
 				F155203B297118E700DC009B /* CompactSlider */,
 				EBDD7DAE29731FB300C1C4B2 /* ZIPFoundation */,
-				EB25B3D72A3A2DD5000E25A1 /* StableDiffusion */,
+				EB9BE0162AC425DA00468918 /* StableDiffusion */,
 			);
 			productName = "Diffusion-macOS";
 			productReference = F15520212971093300DC009B /* Diffusers.app */;
@@ -422,7 +451,7 @@
 			packageReferences = (
 				EBB5BA5B294504DE003A2A5B /* XCRemoteSwiftPackageReference "ZIPFoundation" */,
 				F155203A297118E600DC009B /* XCRemoteSwiftPackageReference "CompactSlider" */,
-				EB25B3D42A3A2DC4000E25A1 /* XCRemoteSwiftPackageReference "ml-stable-diffusion" */,
+				EB9BE0132AC4256F00468918 /* XCRemoteSwiftPackageReference "ml-stable-diffusion" */,
 			);
 			productRefGroup = EBE755C6293E37DD00806B32 /* Products */;
 			projectDirPath = "";
@@ -441,7 +470,14 @@
 			isa = PBXResourcesBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				EBB56B3B2AB87932005361A6 /* VAEDecoder.mlmodelc in Resources */,
+				EBB56B382AB87932005361A6 /* merges.txt in Resources */,
+				EBB56B3F2AB87932005361A6 /* TextEncoder2.mlmodelc in Resources */,
+				EBB56B392AB87932005361A6 /* UnetChunk2.mlmodelc in Resources */,
+				EBB56B3A2AB87932005361A6 /* UnetChunk1.mlmodelc in Resources */,
+				EBB56B3D2AB87932005361A6 /* TextEncoder.mlmodelc in Resources */,
 				EBE755D1293E37DD00806B32 /* Preview Assets.xcassets in Resources */,
+				EBB56B3C2AB87932005361A6 /* vocab.json in Resources */,
 				EBE755CD293E37DD00806B32 /* Assets.xcassets in Resources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
@@ -963,7 +999,7 @@
 /* End XCConfigurationList section */
 
 /* Begin XCRemoteSwiftPackageReference section */
-		EB25B3D42A3A2DC4000E25A1 /* XCRemoteSwiftPackageReference "ml-stable-diffusion" */ = {
+		EB9BE0132AC4256F00468918 /* XCRemoteSwiftPackageReference "ml-stable-diffusion" */ = {
 			isa = XCRemoteSwiftPackageReference;
 			repositoryURL = "https://github.com/apple/ml-stable-diffusion";
 			requirement = {
@@ -994,14 +1030,14 @@
 			isa = XCSwiftPackageProductDependency;
 			productName = StableDiffusion;
 		};
-		EB25B3D52A3A2DC4000E25A1 /* StableDiffusion */ = {
+		EB9BE0142AC425D300468918 /* StableDiffusion */ = {
 			isa = XCSwiftPackageProductDependency;
-			package = EB25B3D42A3A2DC4000E25A1 /* XCRemoteSwiftPackageReference "ml-stable-diffusion" */;
+			package = EB9BE0132AC4256F00468918 /* XCRemoteSwiftPackageReference "ml-stable-diffusion" */;
 			productName = StableDiffusion;
 		};
-		EB25B3D72A3A2DD5000E25A1 /* StableDiffusion */ = {
+		EB9BE0162AC425DA00468918 /* StableDiffusion */ = {
 			isa = XCSwiftPackageProductDependency;
-			package = EB25B3D42A3A2DC4000E25A1 /* XCRemoteSwiftPackageReference "ml-stable-diffusion" */;
+			package = EB9BE0132AC4256F00468918 /* XCRemoteSwiftPackageReference "ml-stable-diffusion" */;
 			productName = StableDiffusion;
 		};
 		EBB5BA5C294504DE003A2A5B /* ZIPFoundation */ = {

Diffusion.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved

@@ -15,7 +15,7 @@
       "location" : "https://github.com/apple/ml-stable-diffusion",
       "state" : {
         "branch" : "main",
-        "revision" : "ce8ee78e28613d8a2e4c8b56932b236cb57e7e20"
+        "revision" : "d0d05ce44d50d6eae7378d2945ffae90d3ab43ed"
       }
     },
     {

Diffusion/Common/ModelInfo.swift

@@ -83,6 +83,7 @@ extension ModelInfo {
     static var defaultComputeUnits: MLComputeUnits { defaultAttention.defaultComputeUnits }
 
     var bestAttention: AttentionVariant {
+        if BENCHMARK { return .splitEinsum }
         if !runningOnMac && supportsAttentionV2 { return .splitEinsumV2 }
         return ModelInfo.defaultAttention
     }
@@ -107,6 +108,7 @@ extension ModelInfo {
     var reduceMemory: Bool {
         // Enable on iOS devices, except when using quantization
         if runningOnMac { return false }
+        if isXL { return !deviceHas8GBOrMore }
         return !(quantized && deviceHas6GBOrMore)
     }
 }
@@ -173,14 +175,22 @@ extension ModelInfo {
 
     static let xl = ModelInfo(
         modelId: "apple/coreml-stable-diffusion-xl-base",
-        modelVersion: "Stable Diffusion XL base",
+        modelVersion: "SDXL base (1024, macOS)",
         supportsEncoder: true,
         isXL: true
     )
 
     static let xlmbp = ModelInfo(
         modelId: "apple/coreml-stable-diffusion-mixed-bit-palettization",
-        modelVersion: "Stable Diffusion XL base [4.5 bit]",
+        modelVersion: "SDXL base (1024, macOS) [4.5 bit]",
+        supportsEncoder: true,
+        quantized: true,
+        isXL: true
+    )
+
+    static let xlmbpChunked = ModelInfo(
+        modelId: "coremlfiles/coreml-stable-diffusion-mixed-bit-palettization",
+        modelVersion: "SDXL base (768, iOS) [4 bit]",
         supportsEncoder: true,
         quantized: true,
         isXL: true
@@ -198,7 +208,8 @@ extension ModelInfo {
                 ModelInfo.v21Base,
                 ModelInfo.v21Palettized,
                 ModelInfo.xl,
-                ModelInfo.xlmbp
+                ModelInfo.xlmbp,
+                ModelInfo.xlmbpChunked,
             ]
         } else {
             return [

Diffusion/Common/Pipeline/PipelineLoader.swift

@@ -93,7 +93,12 @@ extension PipelineLoader {
 
     var packagesFilename: String { (filename as NSString).deletingPathExtension }
 
-    var compiledURL: URL { downloadedURL.deletingLastPathComponent().appendingPathComponent(packagesFilename)  }
+    var compiledURL: URL {
+        guard BENCHMARK else { return downloadedURL.deletingLastPathComponent().appendingPathComponent(packagesFilename) }
+
+        // Model files must be part of the bundle when benchmarking
+        return Bundle.main.resourceURL!
+    }
 
     var downloaded: Bool {
         return FileManager.default.fileExists(atPath: downloadedURL.path)

Diffusion/Common/State.swift

@@ -17,7 +17,7 @@ let DEFAULT_PROMPT = "Labrador in the style of Vermeer"
 enum GenerationState {
     case startup
     case running(StableDiffusionProgress?)
-    case complete(String, CGImage?, UInt32, TimeInterval?)
+    case complete(String, CGImage?, UInt32, TimeInterval?, Double?)
     case userCanceled
     case failed(Error)
 }
@@ -62,11 +62,11 @@ class GenerationContext: ObservableObject {
     @Published var negativePrompt = ""
 
     // FIXME: Double to support the slider component
-    @Published var steps = 25.0
+    @Published var steps = 20.0
     @Published var numImages = 1.0
     @Published var seed: UInt32 = 0
     @Published var guidanceScale = 7.5
-    @Published var previews = 5.0
+    @Published var previews = runningOnMac ? 5.0 : 0.0
     @Published var disableSafety = false
     @Published var previewImage: CGImage? = nil
 

Diffusion/Common/Views/PromptTextField.swift

@@ -28,27 +28,13 @@ struct PromptTextField: View {
         ModelInfo.from(modelVersion: $model.wrappedValue)
     }
 
-    private var filename: String? {
-        let variant = modelInfo?.bestAttention ?? .original
-        return modelInfo?.modelURL(for: variant).lastPathComponent
+    private var pipelineLoader: PipelineLoader? {
+        guard let modelInfo = modelInfo else { return nil }
+        return PipelineLoader(model: modelInfo)
     }
-
-    private var downloadedURL: URL? {
-        if let filename = filename {
-            return PipelineLoader.models.appendingPathComponent(filename)
-        }
-        return nil
-    }
-
-    private var packagesFilename: String? {
-        (filename as NSString?)?.deletingPathExtension
-    }
-
+
     private var compiledURL: URL? {
-        if let packagesFilename = packagesFilename {
-            return downloadedURL?.deletingLastPathComponent().appendingPathComponent(packagesFilename)
-        }
-        return nil
+        return pipelineLoader?.compiledURL
     }
 
     private var textColor: Color {

Diffusion/DiffusionApp.swift

@@ -17,8 +17,10 @@ struct DiffusionApp: App {
     }
 }
 
+let BENCHMARK = true
 let runningOnMac = ProcessInfo.processInfo.isMacCatalystApp
-let deviceHas6GBOrMore = ProcessInfo.processInfo.physicalMemory > 5924000000   // Different devices report different amounts, so approximate
+let deviceHas6GBOrMore = ProcessInfo.processInfo.physicalMemory > 5910000000   // Reported by iOS 17 beta (21A5319a) on iPhone 13 Pro: 5917753344
+let deviceHas8GBOrMore = ProcessInfo.processInfo.physicalMemory > 7900000000   // Reported by iOS 17.0.2 on iPhone 15 Pro Max: 8021032960
 
 let deviceSupportsQuantization = {
     if #available(iOS 17, *) {

Diffusion/Views/Loading.swift

@@ -9,7 +9,7 @@
 import SwiftUI
 import Combine
 
-let model = deviceSupportsQuantization ? ModelInfo.v21Palettized : ModelInfo.v21Base
+let model = BENCHMARK ? ModelInfo.xlmbpChunked : deviceSupportsQuantization ? ModelInfo.v21Palettized : ModelInfo.v21Base
 
 struct LoadingView: View {
     @StateObject var generation = GenerationContext()
@@ -60,6 +60,7 @@ struct LoadingView: View {
                 }
                 do {
                     generation.pipeline = try await loader.prepare()
+                    print("Did load model \(loader.model)")
                     self.currentView = .textToImage
                 } catch {
                     self.currentView = .error("Could not load model, error: \(error)")

Diffusion/Views/TextToImage.swift

@@ -77,7 +77,7 @@ struct ImageWithPlaceholder: View {
                 }
                 ProgressView(label, value: fraction, total: 1).padding()
             })
-        case .complete(let lastPrompt, let image, _, let interval):
+        case .complete(let lastPrompt, let image, _, let interval, let itsPerSecond):
             guard let theImage = image else {
                 return AnyView(Image(systemName: "exclamationmark.triangle").resizable())
             }
@@ -87,7 +87,7 @@ struct ImageWithPlaceholder: View {
                 VStack {
                     imageView.resizable().clipShape(RoundedRectangle(cornerRadius: 20))
                     HStack {
-                        let intervalString = String(format: "Time: %.1fs", interval ?? 0)
+                        let intervalString = String(format: "%.1fs, %.2f it/s", interval ?? 0, itsPerSecond ?? 0)
                         Rectangle().fill(.clear).overlay(Text(intervalString).frame(maxWidth: .infinity, alignment: .leading).padding(.leading))
                         Rectangle().fill(.clear).overlay(
                             HStack {
@@ -114,7 +114,7 @@ struct TextToImage: View {
             generation.state = .running(nil)
             do {
                 let result = try await generation.generate()
-                generation.state = .complete(generation.positivePrompt, result.image, result.lastSeed, result.interval)
+                generation.state = .complete(generation.positivePrompt, result.image, result.lastSeed, result.interval, result.itsPerSecond)
             } catch {
                 generation.state = .failed(error)
             }
@@ -124,7 +124,7 @@ struct TextToImage: View {
     var body: some View {
         VStack {
             HStack {
-                PromptTextField(text: $generation.positivePrompt, isPositivePrompt: true, model: deviceSupportsQuantization ? ModelInfo.v21Palettized.modelVersion : ModelInfo.v21Base.modelVersion)
+                PromptTextField(text: $generation.positivePrompt, isPositivePrompt: true, model: model.modelVersion)
                 Button("Generate") {
                     submit()
                 }