mirror of
https://github.com/GayPizzaSpecifications/stable-diffusion-rpc.git
synced 2025-08-05 22:41:30 +00:00
Split out worker related things to a separate service definition.
This commit is contained in:
40
Sources/StableDiffusionNode/ImageGenerationService.swift
Normal file
40
Sources/StableDiffusionNode/ImageGenerationService.swift
Normal 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)
|
||||
}
|
||||
}
|
Reference in New Issue
Block a user