Skip to content

Commit

Permalink
add --increment-seed argument
Browse files Browse the repository at this point in the history
  • Loading branch information
hirakujira committed Dec 29, 2022
1 parent c1fd82a commit 46a9c00
Showing 1 changed file with 33 additions and 16 deletions.
49 changes: 33 additions & 16 deletions swift/StableDiffusionCLI/main.swift
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,9 @@ struct StableDiffusionSample: ParsableCommand {
@Flag(help: "Reduce memory usage")
var reduceMemory: Bool = false

@Flag(help: "Increse random seed by 1 for each image")
var incrementSeed: Bool = false

mutating func run() throws {
guard FileManager.default.fileExists(atPath: resourcePath) else {
throw RunError.resources("Resource path does not exist \(resourcePath)")
Expand All @@ -89,24 +92,38 @@ struct StableDiffusionSample: ParsableCommand {
let sampleTimer = SampleTimer()
sampleTimer.start()

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()
let loops = incrementSeed ? imageCount : 1
let imageCountPerBatch = incrementSeed ? 1 : imageCount

for i in 0 ..< loops {
if (incrementSeed) {
log("Generating image \(i+1) of \(imageCount) with seed \(seed)\n")
log("\n")
}
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(
Expand Down

0 comments on commit 46a9c00

Please sign in to comment.