diff --git a/.github/workflows/macos.yml b/.github/workflows/macos.yml new file mode 100644 index 0000000..fb04adf --- /dev/null +++ b/.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 diff --git a/.swift-version b/.swift-version new file mode 100644 index 0000000..2df33d7 --- /dev/null +++ b/.swift-version @@ -0,0 +1 @@ +5.6 diff --git a/.swiftformat b/.swiftformat new file mode 100644 index 0000000..2af4fa1 --- /dev/null +++ b/.swiftformat @@ -0,0 +1,4 @@ +--indent 4 +--disable trailingCommas +--exclude "Sources/StableDiffusionProtos/*.pb.swift" +--exclude "Sources/StableDiffusionProtos/*.grpc.swift" diff --git a/.swiftlint.yml b/.swiftlint.yml new file mode 100644 index 0000000..90ebdc1 --- /dev/null +++ b/.swiftlint.yml @@ -0,0 +1,4 @@ +line_length: 180 +excluded: +- Sources/StableDiffusionProtos/*.pb.swift +- Sources/StableDiffusionProtos/*.grpc.swift diff --git a/Package.swift b/Package.swift index d5b8c44..cac9c4f 100644 --- a/Package.swift +++ b/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") diff --git a/Sources/StableDiffusionServer/Errors.swift b/Sources/StableDiffusionCore/Errors.swift similarity index 72% rename from Sources/StableDiffusionServer/Errors.swift rename to Sources/StableDiffusionCore/Errors.swift index 7fd124e..5868f76 100644 --- a/Sources/StableDiffusionServer/Errors.swift +++ b/Sources/StableDiffusionCore/Errors.swift @@ -1,6 +1,6 @@ import Foundation -enum SdServerError: Error { +public enum SdCoreError: Error { case modelNotLoaded case imageEncode case modelNotFound diff --git a/Sources/StableDiffusionServer/ImageExtensions.swift b/Sources/StableDiffusionCore/ImageExtensions.swift similarity index 76% rename from Sources/StableDiffusionServer/ImageExtensions.swift rename to Sources/StableDiffusionCore/ImageExtensions.swift index fc9ea97..d242107 100644 --- a/Sources/StableDiffusionServer/ImageExtensions.swift +++ b/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 } } } diff --git a/Sources/StableDiffusionServer/ModelManager.swift b/Sources/StableDiffusionCore/ModelManager.swift similarity index 80% rename from Sources/StableDiffusionServer/ModelManager.swift rename to Sources/StableDiffusionCore/ModelManager.swift index c6d1b37..67fe366 100644 --- a/Sources/StableDiffusionServer/ModelManager.swift +++ b/Sources/StableDiffusionCore/ModelManager.swift @@ -2,18 +2,18 @@ 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 - + public init(modelBaseURL: URL) { 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]) @@ -24,15 +24,15 @@ 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 { var info = SdModelInfo() info.name = url.lastPathComponent @@ -42,10 +42,10 @@ actor ModelManager { modelUrls[info.name] = url modelStates[info.name] = try ModelState(url: url) } - + private func getModelAttention(_ url: URL) -> String? { let unetMetadataURL = url.appending(components: "Unet.mlmodelc", "metadata.json") - + struct ModelMetadata: Decodable { let mlProgramOperationTypeHistogram: [String: Int] } diff --git a/Sources/StableDiffusionServer/ModelState.swift b/Sources/StableDiffusionCore/ModelState.swift similarity index 79% rename from Sources/StableDiffusionServer/ModelState.swift rename to Sources/StableDiffusionCore/ModelState.swift index 6d02421..f0ef37c 100644 --- a/Sources/StableDiffusionServer/ModelState.swift +++ b/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( @@ -26,20 +26,20 @@ actor ModelState { let vocabUrl = url.appending(component: "vocab.json") 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) pipelineConfig.negativePrompt = request.negativePrompt pipelineConfig.seed = UInt32.random(in: 0 ..< UInt32.max) - + var response = SdGenerateImagesResponse() for _ in 0 ..< request.imageCount { let images = try pipeline.generateImages(configuration: pipelineConfig) - + for cgImage in images { guard let cgImage else { continue } var image = SdImage() diff --git a/Sources/StableDiffusionServer/ImageGenerationService.swift b/Sources/StableDiffusionServer/ImageGenerationService.swift index 6a02b5c..0b33c2c 100644 --- a/Sources/StableDiffusionServer/ImageGenerationService.swift +++ b/Sources/StableDiffusionServer/ImageGenerationService.swift @@ -1,17 +1,18 @@ import Foundation import GRPC +import StableDiffusionCore import StableDiffusionProtos class ImageGenerationServiceProvider: SdImageGenerationServiceAsyncProvider { private let modelManager: ModelManager - + init(modelManager: ModelManager) { 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) } diff --git a/Sources/StableDiffusionServer/ModelService.swift b/Sources/StableDiffusionServer/ModelService.swift index 140ca44..96af18a 100644 --- a/Sources/StableDiffusionServer/ModelService.swift +++ b/Sources/StableDiffusionServer/ModelService.swift @@ -1,29 +1,30 @@ import Foundation import GRPC +import StableDiffusionCore import StableDiffusionProtos class ModelServiceProvider: SdModelServiceAsyncProvider { private let modelManager: ModelManager - + init(modelManager: ModelManager) { 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() diff --git a/Sources/StableDiffusionServer/main.swift b/Sources/StableDiffusionServer/main.swift index 9c7c219..560a791 100644 --- a/Sources/StableDiffusionServer/main.swift +++ b/Sources/StableDiffusionServer/main.swift @@ -1,17 +1,18 @@ -import Foundation import ArgumentParser +import Foundation import GRPC import NIO +import StableDiffusionCore import System struct ServerCommand: ParsableCommand { @Option(name: .shortAndLong, help: "Path to models directory") var modelsDirectoryPath: String = "models" - + mutating func run() throws { let modelsDirectoryURL = URL(filePath: modelsDirectoryPath) let modelManager = ModelManager(modelBaseURL: modelsDirectoryURL) - + let semaphore = DispatchSemaphore(value: 0) Task { print("Loading initial models...") @@ -32,9 +33,9 @@ struct ServerCommand: ParsableCommand { ImageGenerationServiceProvider(modelManager: modelManager) ]) .bind(host: "0.0.0.0", port: 4546) - + dispatchMain() } - } + ServerCommand.main() diff --git a/Sources/TestStableDiffusionClient/main.swift b/Sources/TestStableDiffusionClient/main.swift index 8093c0b..58c4ba6 100644 --- a/Sources/TestStableDiffusionClient/main.swift +++ b/Sources/TestStableDiffusionClient/main.swift @@ -1,18 +1,18 @@ import Foundation -import StableDiffusionProtos -import NIO -import System import GRPC +import NIO +import StableDiffusionProtos +import System let group = PlatformSupport.makeEventLoopGroup(loopCount: 1) defer { - try? group.syncShutdownGracefully() + try? group.syncShutdownGracefully() } let channel = try GRPCChannelPool.with( - target: .host("localhost", port: 4546), - transportSecurity: .plaintext, - eventLoopGroup: group + target: .host("localhost", port: 4546), + transportSecurity: .plaintext, + eventLoopGroup: group ) let modelService = SdModelServiceAsyncClient(channel: channel) @@ -27,14 +27,14 @@ Task { @MainActor in request.modelName = modelInfo.name }) print("Loaded model.") - + print("Generating image...") let request = SdGenerateImagesRequest.with { $0.modelName = modelInfo.name $0.prompt = "cat" $0.imageCount = 1 } - + let response = try await imageGeneratorService.generateImage(request) print("Generated image.") print(response) @@ -43,4 +43,5 @@ Task { @MainActor in exit(1) } } + dispatchMain()