strips/stable-diffusion-rpc
2
0
Fork 0
mirror of https://github.com/GayPizzaSpecifications/stable-diffusion-rpc.git synced 2025-08-05 14:31:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add support for BPE tokenization.

Browse Source
This commit is contained in:
Alex Zenla
2023-04-23 18:49:52 -07:00
parent 0fe35cd976
commit f61fe6a18f
9 changed files with 562 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
Sources/StableDiffusionServer/TokenizerService.swift Normal file
View File

@ -0,0 +1,19 @@
import Foundation
import GRPC
import StableDiffusionCore
import StableDiffusionProtos
class TokenizerServiceProvider: SdTokenizerServiceAsyncProvider {
private let modelManager: ModelManager
init(modelManager: ModelManager) {
self.modelManager = modelManager
}
func tokenize(request: SdTokenizeRequest, context _: GRPCAsyncServerCallContext) async throws -> SdTokenizeResponse {
guard let state = await modelManager.getModelState(name: request.modelName) else {
throw SdCoreError.modelNotFound
}
return try await state.tokenize(request)
}
}

3
Sources/StableDiffusionServer/main.swift
View File

@ -34,7 +34,8 @@ struct ServerCommand: ParsableCommand {
_ = Server.insecure(group: group)
.withServiceProviders([
ModelServiceProvider(modelManager: modelManager),
ImageGenerationServiceProvider(modelManager: modelManager)
ImageGenerationServiceProvider(modelManager: modelManager),
TokenizerServiceProvider(modelManager: modelManager)
])
.bind(host: bindHost, port: bindPort)