Files
stable-diffusion-rpc/Sources/StableDiffusionControl/main.swift

40 lines
1.2 KiB
Swift
Raw Normal View History

2023-04-22 14:52:27 -07:00
import Foundation
import GRPC
2023-04-22 14:52:27 -07:00
import NIO
import StableDiffusionProtos
2023-04-22 14:52:27 -07:00
import System
2023-04-22 16:32:54 -07:00
let client = try StableDiffusionClient(connectionTarget: .host("127.0.0.1", port: 4546), transportSecurity: .plaintext)
2023-04-22 14:52:27 -07:00
Task { @MainActor in
do {
2023-04-22 16:32:54 -07:00
let modelListResponse = try await client.modelService.listModels(.init())
print("Loading random model...")
let modelInfo = modelListResponse.availableModels.randomElement()!
2023-04-22 16:32:54 -07:00
_ = try await client.modelService.loadModel(.with { request in
2023-04-22 14:52:27 -07:00
request.modelName = modelInfo.name
})
print("Loaded random model.")
2023-04-22 14:52:27 -07:00
print("Generating image...")
let request = SdGenerateImagesRequest.with {
$0.modelName = modelInfo.name
$0.outputImageFormat = .png
2023-04-22 14:52:27 -07:00
$0.prompt = "cat"
$0.batchCount = 1
$0.batchSize = 1
2023-04-22 14:52:27 -07:00
}
let response = try await client.imageGenerationService.generateImages(request)
2023-04-22 16:32:54 -07:00
let image = response.images.first!
try image.data.write(to: URL(filePath: "output.png"))
2023-04-22 16:32:54 -07:00
print("Generated image to output.png")
exit(0)
2023-04-22 14:52:27 -07:00
} catch {
print(error)
exit(1)
}
}
2023-04-22 14:52:27 -07:00
dispatchMain()