Formatting, linting, and hopefully a CI build.

.github/workflows/macos.yml

@@ -0,0 +1,31 @@
+name: macOS
+on: [push]
+jobs:
+  build:
+    runs-on: macos-12
+    steps:
+      - name: Checkout Repository
+        uses: actions/checkout@v2
+      - name: Build Executable
+        run: swift build -c release --arch arm64 --arch x86_64
+      - name: Copy Executable
+        run: cp .build/apple/Products/Release/StableDiffusionServer StableDiffusionServer
+      - name: Archive Executable
+        uses: actions/upload-artifact@v2
+        with:
+          name: StableDiffusionServer
+          path: StableDiffusionServer
+  format:
+    runs-on: macos-12
+    steps:
+      - name: Checkout Repository
+        uses: actions/checkout@v2
+      - name: Swift Format
+        run: swiftformat --lint Package.swift Sources
+  lint:
+    runs-on: macos-12
+    steps:
+      - name: Checkout Repository
+        uses: actions/checkout@v2
+      - name: Swift Lint
+        run: swiftlint Package.swift Sources

.swift-version

@@ -0,0 +1 @@
+5.6

.swiftformat

@@ -0,0 +1,4 @@
+--indent 4
+--disable trailingCommas
+--exclude "Sources/StableDiffusionProtos/*.pb.swift"
+--exclude "Sources/StableDiffusionProtos/*.grpc.swift"

.swiftlint.yml

@@ -0,0 +1,4 @@
+line_length: 180
+excluded:
+- Sources/StableDiffusionProtos/*.pb.swift
+- Sources/StableDiffusionProtos/*.grpc.swift

Package.swift

@@ -16,17 +16,22 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-argument-parser", from: "1.2.0")
     ],
     targets: [
+        .target(name: "StableDiffusionProtos", dependencies: [
+            .product(name: "SwiftProtobuf", package: "swift-protobuf"),
+            .product(name: "GRPC", package: "grpc-swift")
+        ]),
+        .target(name: "StableDiffusionCore", dependencies: [
+            .product(name: "StableDiffusion", package: "ml-stable-diffusion"),
+            .target(name: "StableDiffusionProtos")
+        ]),
         .executableTarget(name: "StableDiffusionServer", dependencies: [
             .product(name: "StableDiffusion", package: "ml-stable-diffusion"),
             .product(name: "SwiftProtobuf", package: "swift-protobuf"),
             .product(name: "GRPC", package: "grpc-swift"),
             .target(name: "StableDiffusionProtos"),
+            .target(name: "StableDiffusionCore"),
             .product(name: "ArgumentParser", package: "swift-argument-parser")
         ]),
-        .target(name: "StableDiffusionProtos", dependencies: [
-            .product(name: "SwiftProtobuf", package: "swift-protobuf"),
-            .product(name: "GRPC", package: "grpc-swift")
-        ]),
         .executableTarget(name: "TestStableDiffusionClient", dependencies: [
             .target(name: "StableDiffusionProtos"),
             .product(name: "GRPC", package: "grpc-swift")

Sources/StableDiffusionCore/Errors.swift

@@ -1,6 +1,6 @@
 import Foundation
 
-enum SdServerError: Error {
+public enum SdCoreError: Error {
     case modelNotLoaded
     case imageEncode
     case modelNotFound

Sources/StableDiffusionCore/ImageExtensions.swift

@@ -1,22 +1,22 @@
-import Foundation
 import CoreImage
+import Foundation
 import UniformTypeIdentifiers
 
 extension CGImage {
     func toPngData() throws -> Data {
         guard let data = CFDataCreateMutable(nil, 0) else {
-            throw SdServerError.imageEncode
+            throw SdCoreError.imageEncode
         }
 
         guard let destination = CGImageDestinationCreateWithData(data, "public.png" as CFString, 1, nil) else {
-            throw SdServerError.imageEncode
+            throw SdCoreError.imageEncode
         }
 
         CGImageDestinationAddImage(destination, self, nil)
         if CGImageDestinationFinalize(destination) {
             return data as Data
         } else {
-            throw SdServerError.imageEncode
+            throw SdCoreError.imageEncode
         }
     }
 }

Sources/StableDiffusionCore/ModelManager.swift

@@ -2,10 +2,10 @@ import Foundation
 import StableDiffusion
 import StableDiffusionProtos
 
-actor ModelManager {
-    private var modelInfos: [String : SdModelInfo] = [:]
-    private var modelUrls: [String : URL] = [:]
-    private var modelStates: [String : ModelState] = [:]
+public actor ModelManager {
+    private var modelInfos: [String: SdModelInfo] = [:]
+    private var modelUrls: [String: URL] = [:]
+    private var modelStates: [String: ModelState] = [:]
 
     private let modelBaseURL: URL
 
@@ -13,7 +13,7 @@ actor ModelManager {
         self.modelBaseURL = modelBaseURL
     }
 
-    func reloadModels() throws {
+    public func reloadModels() throws {
         modelInfos.removeAll()
         modelStates.removeAll()
         let contents = try FileManager.default.contentsOfDirectory(at: modelBaseURL.resolvingSymlinksInPath(), includingPropertiesForKeys: [.isDirectoryKey])
@@ -25,12 +25,12 @@ actor ModelManager {
         }
     }
 
-    func listModels() -> [SdModelInfo] {
-        return Array(modelInfos.values)
+    public func listModels() -> [SdModelInfo] {
+        Array(modelInfos.values)
     }
 
-    func getModelState(name: String) -> ModelState? {
-        return modelStates[name]
+    public func getModelState(name: String) -> ModelState? {
+        modelStates[name]
     }
 
     private func addModel(url: URL) throws {

Sources/StableDiffusionCore/ModelState.swift

@@ -1,18 +1,18 @@
-import Foundation
-import StableDiffusionProtos
-import StableDiffusion
 import CoreML
+import Foundation
+import StableDiffusion
+import StableDiffusionProtos
 
-actor ModelState {
+public actor ModelState {
     private let url: URL
-    private var pipeline: StableDiffusionPipeline? = nil
-    private var tokenizer: BPETokenizer? = nil
+    private var pipeline: StableDiffusionPipeline?
+    private var tokenizer: BPETokenizer?
 
-    init(url: URL) throws {
+    public init(url: URL) throws {
         self.url = url
     }
 
-    func load() throws {
+    public func load() throws {
         let config = MLModelConfiguration()
         config.computeUnits = .all
         pipeline = try StableDiffusionPipeline(
@@ -27,9 +27,9 @@ actor ModelState {
         tokenizer = try BPETokenizer(mergesAt: mergesUrl, vocabularyAt: vocabUrl)
     }
 
-    func generate(_ request: SdGenerateImagesRequest) throws -> SdGenerateImagesResponse {
+    public func generate(_ request: SdGenerateImagesRequest) throws -> SdGenerateImagesResponse {
         guard let pipeline else {
-            throw SdServerError.modelNotLoaded
+            throw SdCoreError.modelNotLoaded
         }
 
         var pipelineConfig = StableDiffusionPipeline.Configuration(prompt: request.prompt)

Sources/StableDiffusionServer/ImageGenerationService.swift

@@ -1,5 +1,6 @@
 import Foundation
 import GRPC
+import StableDiffusionCore
 import StableDiffusionProtos
 
 class ImageGenerationServiceProvider: SdImageGenerationServiceAsyncProvider {
@@ -9,9 +10,9 @@ class ImageGenerationServiceProvider: SdImageGenerationServiceAsyncProvider {
         self.modelManager = modelManager
     }
 
-    func generateImage(request: SdGenerateImagesRequest, context: GRPCAsyncServerCallContext) async throws -> SdGenerateImagesResponse {
+    func generateImage(request: SdGenerateImagesRequest, context _: GRPCAsyncServerCallContext) async throws -> SdGenerateImagesResponse {
         guard let state = await modelManager.getModelState(name: request.modelName) else {
-            throw SdServerError.modelNotFound
+            throw SdCoreError.modelNotFound
         }
         return try await state.generate(request)
     }

Sources/StableDiffusionServer/ModelService.swift

@@ -1,5 +1,6 @@
 import Foundation
 import GRPC
+import StableDiffusionCore
 import StableDiffusionProtos
 
 class ModelServiceProvider: SdModelServiceAsyncProvider {
@@ -9,21 +10,21 @@ class ModelServiceProvider: SdModelServiceAsyncProvider {
         self.modelManager = modelManager
     }
 
-    func listModels(request: SdListModelsRequest, context: GRPCAsyncServerCallContext) async throws -> SdListModelsResponse {
+    func listModels(request _: SdListModelsRequest, context _: GRPCAsyncServerCallContext) async throws -> SdListModelsResponse {
         let models = await modelManager.listModels()
         var response = SdListModelsResponse()
         response.models.append(contentsOf: models)
         return response
     }
 
-    func reloadModels(request: SdReloadModelsRequest, context: GRPCAsyncServerCallContext) async throws -> SdReloadModelsResponse {
+    func reloadModels(request _: SdReloadModelsRequest, context _: GRPCAsyncServerCallContext) async throws -> SdReloadModelsResponse {
         try await modelManager.reloadModels()
         return SdReloadModelsResponse()
     }
 
-    func loadModel(request: SdLoadModelRequest, context: GRPCAsyncServerCallContext) async throws -> SdLoadModelResponse {
+    func loadModel(request: SdLoadModelRequest, context _: GRPCAsyncServerCallContext) async throws -> SdLoadModelResponse {
         guard let state = await modelManager.getModelState(name: request.modelName) else {
-            throw SdServerError.modelNotFound
+            throw SdCoreError.modelNotFound
         }
         try await state.load()
         return SdLoadModelResponse()

Sources/StableDiffusionServer/main.swift

@@ -1,7 +1,8 @@
-import Foundation
 import ArgumentParser
+import Foundation
 import GRPC
 import NIO
+import StableDiffusionCore
 import System
 
 struct ServerCommand: ParsableCommand {
@@ -35,6 +36,6 @@ struct ServerCommand: ParsableCommand {
 
         dispatchMain()
     }
-
 }
+
 ServerCommand.main()

Sources/TestStableDiffusionClient/main.swift

@@ -1,8 +1,8 @@
 import Foundation
-import StableDiffusionProtos
-import NIO
-import System
 import GRPC
+import NIO
+import StableDiffusionProtos
+import System
 
 let group = PlatformSupport.makeEventLoopGroup(loopCount: 1)
 defer {
@@ -43,4 +43,5 @@ Task { @MainActor in
         exit(1)
     }
 }
+
 dispatchMain()