add --increment-seed argument

pull/67/head
Hiraku 2 years ago
parent c1fd82a96d
commit 46a9c00826

@ -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(

Loading…
Cancel
Save