strips/stable-diffusion-rpc
2
0
Fork 0
mirror of https://github.com/GayPizzaSpecifications/stable-diffusion-rpc.git synced 2025-10-11 10:29:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

Job management and preparation for multi-hosting.

Browse Source
This commit is contained in:
Alex Zenla
2023-05-08 16:06:07 -07:00
parent a2d9e14f3a
commit ace2c07aa1
30 changed files with 3879 additions and 2307 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
Sources/StableDiffusionServer/ImageGenerationService.swift
View File

@ -4,9 +4,11 @@ import StableDiffusionCore
import StableDiffusionProtos
class ImageGenerationServiceProvider: SdImageGenerationServiceAsyncProvider {
private let jobManager: JobManager
private let modelManager: ModelManager
init(modelManager: ModelManager) {
init(jobManager: JobManager, modelManager: ModelManager) {
self.jobManager = jobManager
self.modelManager = modelManager
}
@ -14,13 +16,25 @@ class ImageGenerationServiceProvider: SdImageGenerationServiceAsyncProvider {
guard let state = await modelManager.getModelState(name: request.modelName) else {
throw SdCoreError.modelNotFound
}
return try await state.generate(request)
let job = await jobManager.create()
DispatchQueue.main.async {
Task {
await self.jobManager.updateJobQueued(job)
}
}
return try await state.generate(request, job: job)
}
func generateImagesStreaming(request: SdGenerateImagesRequest, responseStream: GRPCAsyncResponseStreamWriter<SdGenerateImagesStreamUpdate>, context _: GRPCAsyncServerCallContext) async throws {
guard let state = await modelManager.getModelState(name: request.modelName) else {
throw SdCoreError.modelNotFound
}
try await state.generateStreaming(request, stream: responseStream)
let job = await jobManager.create()
DispatchQueue.main.async {
Task {
await self.jobManager.updateJobQueued(job)
}
}
_ = try await state.generate(request, job: job, stream: responseStream)
}
}

37
Sources/StableDiffusionServer/JobService.swift Normal file
View File

@ -0,0 +1,37 @@
import Foundation
import GRPC
import StableDiffusionCore
import StableDiffusionProtos
class JobServiceProvider: SdJobServiceAsyncProvider {
private let jobManager: JobManager
init(jobManager: JobManager) {
self.jobManager = jobManager
}
func getJob(request: SdGetJobRequest, context _: GRPCAsyncServerCallContext) async throws -> SdGetJobResponse {
var response = SdGetJobResponse()
guard let job = await jobManager.job(id: request.id) else {
throw SdCoreError.jobNotFound
}
response.job = job
return response
}
func cancelJob(request _: SdCancelJobRequest, context _: GRPCAsyncServerCallContext) async throws -> SdCancelJobResponse {
throw SdCoreError.notImplemented
}
func streamJobUpdates(request: SdStreamJobUpdatesRequest, responseStream: GRPCAsyncResponseStreamWriter<SdJobUpdate>, context _: GRPCAsyncServerCallContext) async throws {
let isFilteredById = request.id != 0
for await job in await jobManager.jobUpdatePublisher {
if isFilteredById, job.id != request.id {
continue
}
var update = SdJobUpdate()
update.job = job
try await responseStream.send(update)
}
}
}

8
Sources/StableDiffusionServer/main.swift
View File

@ -16,8 +16,9 @@ struct ServerCommand: ParsableCommand {
var bindPort: Int = 4546
mutating func run() throws {
let jobManager = JobManager()
let modelsDirectoryURL = URL(filePath: modelsDirectoryPath)
let modelManager = ModelManager(modelBaseURL: modelsDirectoryURL)
let modelManager = ModelManager(modelBaseURL: modelsDirectoryURL, jobManager: jobManager)
let semaphore = DispatchSemaphore(value: 0)
Task {
@ -34,8 +35,9 @@ struct ServerCommand: ParsableCommand {
_ = Server.insecure(group: group)
.withServiceProviders([
ModelServiceProvider(modelManager: modelManager),
ImageGenerationServiceProvider(modelManager: modelManager),
TokenizerServiceProvider(modelManager: modelManager)
ImageGenerationServiceProvider(jobManager: jobManager, modelManager: modelManager),
TokenizerServiceProvider(modelManager: modelManager),
JobServiceProvider(jobManager: jobManager)
])
.bind(host: bindHost, port: bindPort)