From 46a9c008262f605f5544525c84f6d4a9419dcb3b Mon Sep 17 00:00:00 2001
From: Hiraku <dev@hiraku.tw>
Date: Fri, 16 Dec 2022 17:34:34 +0800
Subject: [PATCH] add --increment-seed argument

---
 swift/StableDiffusionCLI/main.swift | 49 +++++++++++++++++++----------
 1 file changed, 33 insertions(+), 16 deletions(-)

diff --git a/swift/StableDiffusionCLI/main.swift b/swift/StableDiffusionCLI/main.swift
index 7343cf2b..3aa10ea2 100644
--- a/swift/StableDiffusionCLI/main.swift
+++ b/swift/StableDiffusionCLI/main.swift
@@ -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)")
@@ -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(