Split out worker related things to a separate service definition.

This commit is contained in:
2023-05-08 22:12:24 -07:00
parent ace2c07aa1
commit 2e5a37ea4b
28 changed files with 1271 additions and 359 deletions

View File

@ -0,0 +1,40 @@
import Foundation
import GRPC
import StableDiffusionCore
import StableDiffusionProtos
class ImageGenerationServiceProvider: SdImageGenerationServiceAsyncProvider {
private let jobManager: JobManager
private let modelManager: ModelManager
init(jobManager: JobManager, modelManager: ModelManager) {
self.jobManager = jobManager
self.modelManager = modelManager
}
func generateImages(request: SdGenerateImagesRequest, context _: GRPCAsyncServerCallContext) async throws -> SdGenerateImagesResponse {
guard let state = await modelManager.getModelState(name: request.modelName) else {
throw SdCoreError.modelNotFound
}
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
}
let job = await jobManager.create()
DispatchQueue.main.async {
Task {
await self.jobManager.updateJobQueued(job)
}
}
_ = try await state.generate(request, job: job, stream: responseStream)
}
}