Add NvRandomSource (#257)

Draw Things 1.20230913.1 released with NVIDIA GPU Compatible mode. See
example: https://twitter.com/drawthingsapp/status/1702387783166476689

This is the implementation that should be compatible with the
ml-stable-diffusion project.

Author: liuliu

swift/StableDiffusion/pipeline/NvRandomSource.swift

@@ -0,0 +1,91 @@
+import Foundation
+import CoreML
+
+/// A random source consistent with NVIDIA curandom
+///
+///  This implementation references to:
+///  https://github.com/dsnz/random/blob/master/philox.py for Philox_M4_32 configuration.
+///
+@available(iOS 16.2, macOS 13.1, *)
+struct NvRandomSource: RandomSource {
+  public let seed: UInt64
+  private var offset: UInt32
+
+  /// Initialize with a random seed
+  ///
+  /// - Parameters
+  ///     - seed: Seed for underlying Philox M4 32 generator
+  /// - Returns random source
+  public init(seed: UInt32) {
+    self.seed = UInt64(seed)
+    offset = 0
+  }
+
+  static private let PHILOX_M4_32: (UInt32, UInt32) = (0xD251_1F53, 0xCD9E_8D57)
+  static private let PHILOX_W_32: (UInt32, UInt32) = (0x9E37_79B9, 0xBB67_AE85)
+
+  static private func philox4Round(counter: inout [[UInt32]], key: [[UInt32]]) {
+    for i in 0..<counter[0].count {
+      let v1: UInt64 = UInt64(counter[0][i]) * UInt64(PHILOX_M4_32.0)
+      let v2: UInt64 = UInt64(counter[2][i]) * UInt64(PHILOX_M4_32.1)
+      counter[0][i] = UInt32(v2 >> 32) ^ counter[1][i] ^ key[0][i]
+      counter[1][i] = UInt32(v2 & 0xffff_ffff)
+      counter[2][i] = UInt32(v1 >> 32) ^ counter[3][i] ^ key[1][i]
+      counter[3][i] = UInt32(v1 & 0xffff_ffff)
+    }
+  }
+
+  static private func philox4Bumpkey(key: inout [[UInt32]]) {
+    for (i, element) in key[0].enumerated() {
+      key[0][i] = element &+ PHILOX_W_32.0
+    }
+    for (i, element) in key[1].enumerated() {
+      key[1][i] = element &+ PHILOX_W_32.1
+    }
+  }
+
+  static private func philox4_32(counter: inout [[UInt32]], key: inout [[UInt32]], rounds: Int = 10) {
+    for _ in 0..<(rounds - 1) {
+      philox4Round(counter: &counter, key: key)
+      philox4Bumpkey(key: &key)
+    }
+    philox4Round(counter: &counter, key: key)
+  }
+
+  private func boxMuller(_ counter1: [UInt32], _ counter2: [UInt32], mean: Double, stdev: Double) -> [Double] {
+    // Box-Muller transform
+    return zip(counter1, counter2).map {
+      let u: Double = Double($0) / 4294967296.0 + (1.0 / 8589934592.0)
+      let v: Double = Double($1) * (.pi / 2147483648.0) + (.pi / 4294967296.0)
+      let radius = stdev * sqrt(-2.0 * log(u))
+      return radius * sin(v) + mean
+    }
+  }
+
+  private mutating func normalArray(count: Int, mean: Double, stdev: Double) -> [Double] {
+    var counter: [[UInt32]] = [
+      Array(repeating: offset, count: count),
+      Array(repeating: 0, count: count),
+      Array(0..<UInt32(count)),
+      Array(repeating: 0, count: count),
+    ]
+    offset += 1
+    var key: [[UInt32]] = [
+      Array(repeating: UInt32(seed & 0xffff_ffff), count: count),
+      Array(repeating: UInt32(seed >> 32), count: count),
+    ]
+    Self.philox4_32(counter: &counter, key: &key)
+    return boxMuller(counter[0], counter[1], mean: mean, stdev: stdev)
+  }
+
+  /// Generates a random value from a normal distribution with given mean and standard deviation.
+  mutating func nextNormal(mean: Double = 0.0, stdev: Double = 1.0) -> Double {
+    return normalArray(count: 1, mean: mean, stdev: stdev)[0]
+  }
+
+  /// Generate a shaped array with scalars from a normal distribution with given mean and standard deviation.
+  mutating func normalShapedArray(_ shape: [Int], mean: Double = 0.0, stdev: Double = 1.0) -> MLShapedArray<Double> {
+    let count = shape.reduce(1, *)
+    return .init(scalars: normalArray(count: count, mean: mean, stdev: stdev), shape: shape)
+  }
+}

swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift

@@ -21,6 +21,8 @@ public enum StableDiffusionRNG {
     case numpyRNG
     /// RNG that matches PyTorch CPU implementation.
     case torchRNG
+    /// RNG that matches PyTorch CUDA implementation.
+    case nvidiaRNG
 }
 
 public enum PipelineError: String, Swift.Error {
@@ -398,6 +400,8 @@ extension StableDiffusionPipelineProtocol {
             return NumPyRandomSource(seed: seed)
         case .torchRNG:
             return TorchRandomSource(seed: seed)
+        case .nvidiaRNG:
+            return NvRandomSource(seed: seed)
         }
     }
 

swift/StableDiffusionCLI/main.swift

@@ -74,7 +74,7 @@ struct StableDiffusionSample: ParsableCommand {
     @Option(help: "Scheduler to use, one of {pndm, dpmpp}")
     var scheduler: SchedulerOption = .pndm
 
-    @Option(help: "Random number generator to use, one of {numpy, torch}")
+    @Option(help: "Random number generator to use, one of {numpy, torch, nvidia}")
     var rng: RNGOption = .numpy
 
     @Option(
@@ -336,11 +336,12 @@ enum SchedulerOption: String, ExpressibleByArgument {
 
 @available(iOS 16.2, macOS 13.1, *)
 enum RNGOption: String, ExpressibleByArgument {
-    case numpy, torch
+    case numpy, torch, nvidia
     var stableDiffusionRNG: StableDiffusionRNG {
         switch self {
         case .numpy: return .numpyRNG
         case .torch: return .torchRNG
+        case .nvidia: return .nvidiaRNG
         }
     }
 }