add --increment-seed argument

apple · Dec 16, 2022 · 5e9fd04 · 5e9fd04
1 parent 66dde8d
commit 5e9fd04
Showing 1 changed file with 31 additions and 14 deletions.
diff --git a/swift/StableDiffusionCLI/main.swift b/swift/StableDiffusionCLI/main.swift
@@ -62,6 +62,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)")
@@ -83,22 +86,36 @@ struct StableDiffusionSample: ParsableCommand {
         let sampleTimer = SampleTimer()
         sampleTimer.start()
 
-        let images = try pipeline.generateImages(
-            prompt: prompt,
-            imageCount: imageCount,
-            stepCount: stepCount,
-            seed: seed,
-            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,
+                imageCount: imageCountPerBatch,
+                stepCount: stepCount,
+                seed: seed,
+                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(