SDXL for iOS, and refiner for Mac (huggingface#83)

* Bring model definitions over from benchmark branch

* PrompTextField uses location from PipelineLoader

This could be merged in main.

* Fix potential model mismatch.

* Disable previews only on iOS

* Fix macOS build

* Update package

* Enable SDXL on big iPhones

* Update model description

* Refiner model

* Use Karras by default for XL models

* Update URL

Author: pcuenca

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.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" : "94814cfa41935efd8151a43758360a54e4e3c5d5"
       }
     },
     {

Diffusion/Common/ModelInfo.swift

@@ -107,6 +107,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,33 +174,56 @@ 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 xlWithRefiner = ModelInfo(
+        modelId: "apple/coreml-stable-diffusion-xl-base-with-refiner",
+        modelVersion: "SDXL with refiner (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: "apple/coreml-stable-diffusion-xl-base-ios",
+        modelVersion: "SDXL base (768, iOS) [4 bit]",
+        supportsEncoder: false,
+        quantized: true,
+        isXL: true
+    )
+
     static let MODELS: [ModelInfo] = {
         if deviceSupportsQuantization {
-            return [
+            var models = [
                 ModelInfo.v14Base,
                 ModelInfo.v14Palettized,
                 ModelInfo.v15Base,
                 ModelInfo.v15Palettized,
                 ModelInfo.v2Base,
                 ModelInfo.v2Palettized,
                 ModelInfo.v21Base,
-                ModelInfo.v21Palettized,
-                ModelInfo.xl,
-                ModelInfo.xlmbp
+                ModelInfo.v21Palettized
             ]
+            if runningOnMac {
+                models.append(contentsOf: [
+                    ModelInfo.xl,
+                    ModelInfo.xlWithRefiner,
+                    ModelInfo.xlmbp
+                ])
+            } else {
+                models.append(ModelInfo.xlmbpChunked)
+            }
+            return models
         } else {
             return [
                 ModelInfo.v14Base,

Diffusion/Common/Pipeline/Pipeline.swift

@@ -91,6 +91,7 @@ class Pipeline {
         if isXL {
             config.encoderScaleFactor = 0.13025
             config.decoderScaleFactor = 0.13025
+            config.schedulerTimestepSpacing = .karras
         }
 
         // Evenly distribute previews based on inference steps

Diffusion/Common/State.swift

@@ -66,7 +66,7 @@ class GenerationContext: ObservableObject {
     @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

@@ -18,7 +18,8 @@ struct DiffusionApp: App {
 }
 
 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,9 +9,14 @@
 import SwiftUI
 import Combine
 
-let model = deviceSupportsQuantization ? ModelInfo.v21Palettized : ModelInfo.v21Base
+func iosModel() -> ModelInfo {
+    guard deviceSupportsQuantization else { return ModelInfo.v21Base }
+    if deviceHas6GBOrMore { return ModelInfo.xlmbpChunked }
+    return ModelInfo.v21Palettized
+}
 
 struct LoadingView: View {
+
     @StateObject var generation = GenerationContext()
 
     @State private var preparationPhase = "Downloading…"
@@ -40,7 +45,7 @@ struct LoadingView: View {
         .environmentObject(generation)
         .onAppear {
             Task.init {
-                let loader = PipelineLoader(model: model)
+                let loader = PipelineLoader(model: iosModel())
                 stateSubscriber = loader.statePublisher.sink { state in
                     DispatchQueue.main.async {
                         switch state {

Diffusion/Views/TextToImage.swift

@@ -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: iosModel().modelVersion)
                 Button("Generate") {
                     submit()
                 }