SD XL support (#67)

* Support SDXL base model.

* Add support for 4.5 bit SDXL model.

* Update package version

Author: pcuenca

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

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

Diffusion/Common/ModelInfo.swift

@@ -43,14 +43,19 @@ struct ModelInfo {
     /// Are weights quantized? This is only used to decide whether to use `reduceMemory`
     let quantized: Bool
 
+    /// Whether this is a Stable Diffusion XL model
+    // TODO: retrieve from remote config
+    let isXL: Bool
+    
     //TODO: refactor all these properties
     init(modelId: String, modelVersion: String,
          originalAttentionSuffix: String = "original_compiled",
          splitAttentionSuffix: String = "split_einsum_compiled",
          splitAttentionV2Suffix: String = "split_einsum_v2_compiled",
          supportsEncoder: Bool = false,
          supportsAttentionV2: Bool = false,
-         quantized: Bool = false) {
+         quantized: Bool = false,
+         isXL: Bool = false) {
         self.modelId = modelId
         self.modelVersion = modelVersion
         self.originalAttentionSuffix = originalAttentionSuffix
@@ -59,6 +64,7 @@ struct ModelInfo {
         self.supportsEncoder = supportsEncoder
         self.supportsAttentionV2 = supportsAttentionV2
         self.quantized = quantized
+        self.isXL = isXL
     }
 }
 
@@ -165,6 +171,21 @@ extension ModelInfo {
         modelVersion: "OFA-Sys/small-stable-diffusion-v0"
     )
 
+    static let xl = ModelInfo(
+        modelId: "apple/coreml-stable-diffusion-xl-base",
+        modelVersion: "Stable Diffusion XL base",
+        supportsEncoder: true,
+        isXL: true
+    )
+    
+    static let xlmbp = ModelInfo(
+        modelId: "apple/coreml-stable-diffusion-mixed-bit-palettization",
+        modelVersion: "Stable Diffusion XL base [4.5 bit]",
+        supportsEncoder: true,
+        quantized: true,
+        isXL: true
+    )
+    
     static let MODELS: [ModelInfo] = {
         if deviceSupportsQuantization {
             return [
@@ -176,15 +197,15 @@ extension ModelInfo {
                 ModelInfo.v2Palettized,
                 ModelInfo.v21Base,
                 ModelInfo.v21Palettized,
-                ModelInfo.ofaSmall
+                ModelInfo.xl,
+                ModelInfo.xlmbp
             ]
         } else {
             return [
                 ModelInfo.v14Base,
                 ModelInfo.v15Base,
                 ModelInfo.v2Base,
                 ModelInfo.v21Base,
-                ModelInfo.ofaSmall
             ]
         }
     }()

Diffusion/Common/Pipeline/Pipeline.swift

@@ -40,7 +40,7 @@ struct GenerationResult {
 }
 
 class Pipeline {
-    let pipeline: StableDiffusionPipeline
+    let pipeline: StableDiffusionPipelineProtocol
     let maxSeed: UInt32
 
     var progress: StableDiffusionProgress? = nil {
@@ -52,7 +52,7 @@ class Pipeline {
 
     private var canceled = false
 
-    init(_ pipeline: StableDiffusionPipeline, maxSeed: UInt32 = UInt32.max) {
+    init(_ pipeline: StableDiffusionPipelineProtocol, maxSeed: UInt32 = UInt32.max) {
         self.pipeline = pipeline
         self.maxSeed = maxSeed
     }

Diffusion/Common/Pipeline/PipelineLoader.swift

@@ -161,15 +161,26 @@ extension PipelineLoader {
         state = .readyOnDisk
     }
 
-    func load(url: URL) async throws -> StableDiffusionPipeline {
+    func load(url: URL) async throws -> StableDiffusionPipelineProtocol {
         let beginDate = Date()
         let configuration = MLModelConfiguration()
         configuration.computeUnits = computeUnits
-        let pipeline = try StableDiffusionPipeline(resourcesAt: url,
-                                                   controlNet: [],
-                                                   configuration: configuration,
-                                                   disableSafety: false,
-                                                   reduceMemory: model.reduceMemory)
+        let pipeline: StableDiffusionPipelineProtocol
+        if model.isXL {
+            if #available(macOS 14.0, iOS 17.0, *) {
+                pipeline = try StableDiffusionXLPipeline(resourcesAt: url,
+                                                       configuration: configuration,
+                                                       reduceMemory: model.reduceMemory)
+            } else {
+                throw "Stable Diffusion XL requires macOS 14"
+            }
+        } else {
+            pipeline = try StableDiffusionPipeline(resourcesAt: url,
+                                                       controlNet: [],
+                                                       configuration: configuration,
+                                                       disableSafety: false,
+                                                       reduceMemory: model.reduceMemory)
+        }
         try pipeline.loadResources()
         print("Pipeline loaded in \(Date().timeIntervalSince(beginDate))")
         state = .loaded