|
|
@ -68,6 +68,9 @@ struct StableDiffusionSample: ParsableCommand {
|
|
|
|
@Flag(help: "Reduce memory usage")
|
|
|
|
@Flag(help: "Reduce memory usage")
|
|
|
|
var reduceMemory: Bool = false
|
|
|
|
var reduceMemory: Bool = false
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@Flag(help: "Increse random seed by 1 for each image")
|
|
|
|
|
|
|
|
var incrementSeed: Bool = false
|
|
|
|
|
|
|
|
|
|
|
|
mutating func run() throws {
|
|
|
|
mutating func run() throws {
|
|
|
|
guard FileManager.default.fileExists(atPath: resourcePath) else {
|
|
|
|
guard FileManager.default.fileExists(atPath: resourcePath) else {
|
|
|
|
throw RunError.resources("Resource path does not exist \(resourcePath)")
|
|
|
|
throw RunError.resources("Resource path does not exist \(resourcePath)")
|
|
|
@ -89,24 +92,38 @@ struct StableDiffusionSample: ParsableCommand {
|
|
|
|
let sampleTimer = SampleTimer()
|
|
|
|
let sampleTimer = SampleTimer()
|
|
|
|
sampleTimer.start()
|
|
|
|
sampleTimer.start()
|
|
|
|
|
|
|
|
|
|
|
|
let images = try pipeline.generateImages(
|
|
|
|
let loops = incrementSeed ? imageCount : 1
|
|
|
|
prompt: prompt,
|
|
|
|
let imageCountPerBatch = incrementSeed ? 1 : imageCount
|
|
|
|
negativePrompt: negativePrompt,
|
|
|
|
|
|
|
|
imageCount: imageCount,
|
|
|
|
for i in 0 ..< loops {
|
|
|
|
stepCount: stepCount,
|
|
|
|
if (incrementSeed) {
|
|
|
|
seed: seed,
|
|
|
|
log("Generating image \(i+1) of \(imageCount) with seed \(seed)\n")
|
|
|
|
guidanceScale: guidanceScale,
|
|
|
|
log("\n")
|
|
|
|
scheduler: scheduler.stableDiffusionScheduler
|
|
|
|
|
|
|
|
) { progress in
|
|
|
|
|
|
|
|
sampleTimer.stop()
|
|
|
|
|
|
|
|
handleProgress(progress,sampleTimer)
|
|
|
|
|
|
|
|
if progress.stepCount != progress.step {
|
|
|
|
|
|
|
|
sampleTimer.start()
|
|
|
|
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return true
|
|
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
_ = try saveImages(images, logNames: true)
|
|
|
|
let images = try pipeline.generateImages(
|
|
|
|
|
|
|
|
prompt: prompt,
|
|
|
|
|
|
|
|
negativePrompt: negativePrompt,
|
|
|
|
|
|
|
|
imageCount: imageCount,
|
|
|
|
|
|
|
|
stepCount: stepCount,
|
|
|
|
|
|
|
|
seed: seed,
|
|
|
|
|
|
|
|
guidanceScale: guidanceScale,
|
|
|
|
|
|
|
|
scheduler: scheduler.stableDiffusionScheduler
|
|
|
|
|
|
|
|
) { progress in
|
|
|
|
|
|
|
|
sampleTimer.stop()
|
|
|
|
|
|
|
|
handleProgress(progress,sampleTimer)
|
|
|
|
|
|
|
|
if progress.stepCount != progress.step {
|
|
|
|
|
|
|
|
sampleTimer.start()
|
|
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
return true
|
|
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
_ = try saveImages(images, logNames: true)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if (incrementSeed) {
|
|
|
|
|
|
|
|
seed += 1
|
|
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func handleProgress(
|
|
|
|
func handleProgress(
|
|
|
|