Add Euler Ancestral Discrete Scheduler

swift/StableDiffusion/pipeline/EulerAncestralDiscreteScheduler.swift

@@ -0,0 +1,120 @@
+import Accelerate
+import CoreML
+
+@available(iOS 16.2, macOS 13.1, *)
+public final class EulerAncestralDiscreteScheduler: Scheduler {
+ public let trainStepCount: Int
+ public let inferenceStepCount: Int
+ public let betas: [Float]
+ public let timeSteps: [Int]
+ public let alphas: [Float]
+ public let alphasCumProd: [Float]
+ public let sigmas: [Float]
+ public let initNoiseSigma: Float
+ private(set) var randomSource: RandomSource
+ public private(set) var modelOutputs: [MLShapedArray<Float32>] = []
+
+ public init(
+ randomSource: RandomSource,
+ stepCount: Int = 50,
+ trainStepCount: Int = 1000,
+ betaStart: Float = 0.0001,
+ betaEnd: Float = 0.02,
+ betaSchedule: BetaSchedule = .scaledLinear,
+ timestepSpacing: TimestepSpacing = .leading,
+ stepsOffset: Int = 1
+ ) {
+ self.randomSource = randomSource
+ self.trainStepCount = trainStepCount
+ inferenceStepCount = stepCount
+
+ switch betaSchedule {
+ case .linear:
+ betas = linspace(betaStart, betaEnd, trainStepCount)
+ case .scaledLinear:
+ betas = linspace(sqrt(betaStart), sqrt(betaEnd), trainStepCount).map { $0 * $0 }
+ }
+
+ alphas = betas.map { 1.0 - $0 }
+ var alphasCumProd = alphas
+ for i in 1 ..< alphasCumProd.count {
+ alphasCumProd[i] *= alphasCumProd[i - 1]
+ }
+ self.alphasCumProd = alphasCumProd
+
+ var sigmas = vForce.sqrt(
+ vDSP.divide(
+ vDSP.subtract(
+ [Float](repeating: 1, count: alphasCumProd.count),
+ alphasCumProd
+ ),
+ alphasCumProd
+ )
+ )
+
+ var timeSteps = [Float](repeating: 0.0, count: stepCount)
+ switch timestepSpacing {
+ case .linspace:
+ timeSteps = linspace(0, Float(trainStepCount - 1), stepCount)
+ case .leading:
+ let stepRatio = trainStepCount / stepCount
+ timeSteps = (0 ..< stepCount).map { Float($0 * stepRatio + stepsOffset) }
+ case .trailing:
+ let stepRatio = trainStepCount / stepCount
+ timeSteps = (1 ... stepCount).map { Float($0 * stepRatio - 1) }
+ }
+ timeSteps.reverse()
+ var sigmasInt = [Float](repeating: 0.0, count: timeSteps.count)
+ vDSP_vlint(&sigmas, &timeSteps, vDSP_Stride(1), &sigmasInt, vDSP_Stride(1), vDSP_Length(timeSteps.count), vDSP_Length(sigmas.count))
+ sigmasInt.append(0.0)
+ initNoiseSigma = sigmasInt.max()!
+ self.timeSteps = timeSteps.map { Int($0) }
+ self.sigmas = sigmasInt
+ }
+
+ public func step(
+ output: MLShapedArray<Float32>,
+ timeStep: Int,
+ sample: MLShapedArray<Float32>
+ ) -> MLShapedArray<Float32> {
+ let stepIndex = timeSteps.firstIndex(of: timeStep) ?? timeSteps.count - 1
+ let sigma = sigmas[stepIndex]
+ let predOriginalSample = weightedSum([1.0, Double(-1.0 * sigma)], [sample, output])
+
+ let sigmaFrom = sigmas[stepIndex]
+ let sigmaTo = sigmas[stepIndex + 1]
+ let sigmaUp = sqrt(pow(sigmaTo, 2.0) * (pow(sigmaFrom, 2.0) - pow(sigmaTo, 2.0)) / pow(sigmaFrom, 2.0))
+ let sigmaDown = sqrt(pow(sigmaTo, 2.0) - pow(sigmaUp, 2.0))
+
+ // Convert to an ODE derivative:
+ let derivative = weightedSum([Double(1.0 / sigma), Double(-1.0 / sigma)], [sample, predOriginalSample])
+ let dt = sigmaDown - sigma
+ let prevSample = weightedSum([1.0, Double(dt)], [sample, derivative])
+ let noise = MLShapedArray<Float32>(
+ converting: randomSource.normalShapedArray(
+ output.shape,
+ mean: 0.0,
+ stdev: 1.0
+ )
+ )
+ return weightedSum([1.0, Double(sigmaUp)], [prevSample, noise])
+ }
+
+ public func scaleModelInput(
+ sample: MLShapedArray<Float32>,
+ timeStep: Int
+ ) -> MLShapedArray<Float32> {
+ let stepIndex = timeSteps.firstIndex(of: timeStep) ?? timeSteps.count - 1
+ let sigma = sigmas[stepIndex]
+ let scale = sqrt(pow(sigma, 2.0) + 1.0)
+ let scalarCount = sample.scalarCount
+ return MLShapedArray(unsafeUninitializedShape: sample.shape) { scalars, _ in
+ assert(scalars.count == scalarCount)
+ sample.withUnsafeShapedBufferPointer { sample, _, _ in
+ for i in 0 ..< scalarCount {
+ scalars.initializeElement(at: i, to: sample[i] / scale)
+ }
+ }
+ }
+ }
+}

swift/StableDiffusion/pipeline/Scheduler.swift

@@ -47,6 +47,11 @@ public protocol Scheduler {
  timeStep t: Int,
  sample s: MLShapedArray<Float32>
  ) -> MLShapedArray<Float32>
+
+ func scaleModelInput(
+ sample: MLShapedArray<Float32>,
+ timeStep: Int
+ ) -> MLShapedArray<Float32>
 }
 
 @available(iOS 16.2, macOS 13.1, *)
@@ -100,6 +105,13 @@ public extension Scheduler {
 
  return noisySamples
  }
+
+ func scaleModelInput(
+ sample: MLShapedArray<Float32>,
+ timeStep: Int
+ ) -> MLShapedArray<Float32> {
+ return sample
+ }
 }
 
 // MARK: - Timesteps
@@ -125,6 +137,13 @@ public enum BetaSchedule {
  case scaledLinear
 }
 
+@available(iOS 16.2, macOS 13.1, *)
+public enum TimestepSpacing {
+ case linspace
+ case leading
+ case trailing
+}
+
 // MARK: - PNDMScheduler
 
 /// A scheduler used to compute a de-noised image

swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift

@@ -13,6 +13,8 @@ public enum StableDiffusionScheduler {
  case pndmScheduler
  /// Scheduler that uses a second order DPM-Solver++ algorithm
  case dpmSolverMultistepScheduler
+
+ case eulerAncestralDiscreteScheduler
 }
 
 /// RNG compatible with StableDiffusionPipeline
@@ -229,6 +231,9 @@ public struct StableDiffusionPipeline: StableDiffusionPipelineProtocol {
  switch config.schedulerType {
  case .pndmScheduler: return PNDMScheduler(stepCount: config.stepCount)
  case .dpmSolverMultistepScheduler: return DPMSolverMultistepScheduler(stepCount: config.stepCount, timeStepSpacing: config.schedulerTimestepSpacing)
+ case .eulerAncestralDiscreteScheduler: return EulerAncestralDiscreteScheduler(
+ randomSource: randomSource(from: config.rngType, seed: config.seed),
+ stepCount: config.stepCount)
  }
  }
 
@@ -258,10 +263,16 @@ public struct StableDiffusionPipeline: StableDiffusionPipelineProtocol {
 
  // Expand the latents for classifier-free guidance
  // and input to the Unet noise prediction model
- let latentUnetInput = latents.map {
+ var latentUnetInput = latents.map {
  MLShapedArray<Float32>(concatenating: [$0, $0], alongAxis: 0)
  }
 
+ for i in 0..<config.imageCount {
+ latentUnetInput[i] = scheduler[i].scaleModelInput(
+ sample: latentUnetInput[i],
+ timeStep: t)
+ }
+
  // Before Unet, execute controlNet and add the output into Unet inputs
  let additionalResiduals = try controlNet?.execute(
  latents: latentUnetInput,

swift/StableDiffusion/pipeline/StableDiffusionXLPipeline.swift

@@ -176,6 +176,9 @@ public struct StableDiffusionXLPipeline: StableDiffusionPipelineProtocol {
  switch config.schedulerType {
  case .pndmScheduler: return PNDMScheduler(stepCount: config.stepCount)
  case .dpmSolverMultistepScheduler: return DPMSolverMultistepScheduler(stepCount: config.stepCount, timeStepSpacing: config.schedulerTimestepSpacing)
+ case .eulerAncestralDiscreteScheduler: return EulerAncestralDiscreteScheduler(
+ randomSource: randomSource(from: config.rngType, seed: config.seed),
+ stepCount: config.stepCount)
  }
  }
 
@@ -208,10 +211,17 @@ public struct StableDiffusionXLPipeline: StableDiffusionPipelineProtocol {
  for (step,t) in timeSteps.enumerated() {
  // Expand the latents for classifier-free guidance
  // and input to the Unet noise prediction model
- let latentUnetInput = latents.map {
+ var latentUnetInput = latents.map {
  MLShapedArray<Float32>(concatenating: [$0, $0], alongAxis: 0)
  }
 
+ for i in 0..<config.imageCount {
+ latentUnetInput[i] = scheduler[i].scaleModelInput(
+ sample: latentUnetInput[i],
+ timeStep: t
+ )
+ }
+
  // Switch to refiner if specified
  if let refiner = unetRefiner, step == refinerStartStep {
  unet.unloadResources()

swift/StableDiffusionTests/EulerAncestralDiscreteSchedulerTests.swift

@@ -0,0 +1,36 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import XCTest
+import CoreML
+@testable import StableDiffusion
+
+@available(iOS 16.2, macOS 14.0, *)
+final class EulerAncestralDiscreteSchedulerTests: XCTestCase {
+ func testScal() throws {
+ let scheduler = EulerAncestralDiscreteScheduler(randomSource: TorchRandomSource(seed: 0), stepCount: 25)
+ let scaledSample = scheduler.scaleModelInput(
+ sample: MLShapedArray(
+ scalars: [0.1, 0.2, 0.3, 0.4, 0.5],
+ shape: [5]
+ ),
+ timeStep: 960)
+ XCTAssertTrue(true)
+ }
+
+ func testLinspaceStep() throws {
+ let scheduler = EulerAncestralDiscreteScheduler(randomSource: TorchRandomSource(seed: 0), stepCount: 25)
+ let preSample = scheduler.step(
+ output: MLShapedArray<Float32>(repeating: 0.5, shape: [5]),
+ timeStep: 960,
+ sample: MLShapedArray<Float32>(repeating: 0.5, shape: [5])
+ )
+ XCTAssertEqual(
+ preSample,
+ MLShapedArray(
+ scalars: [18.519714, -15.129302, -49.71263, 0.6798735, -29.640398],
+ shape: [5]
+ )
+ )
+ }
+}