From 97708c3cbc1c557714f4227a1452c68a22181a7c Mon Sep 17 00:00:00 2001
From: Siavosh Kasravi <siavosh.kasravi@gmail.com>
Date: Fri, 28 Jun 2024 14:48:55 +0330
Subject: [PATCH] Enhancement - ability to change image count in the queue

---
 ui/index.html               |  2 +-
 ui/media/css/main.css       |  4 ++++
 ui/media/js/main.js         | 14 +++++++++++++-
 ui/media/js/task-manager.js |  2 ++
 4 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/ui/index.html b/ui/index.html
index a8c6685f..5553c507 100644
--- a/ui/index.html
+++ b/ui/index.html
@@ -156,7 +156,7 @@
                     <div><table>
                     <tr><b class="settings-subheader">Image Settings</b></tr>
                     <tr class="pl-5"><td><label for="seed">Seed:</label></td><td><input id="seed" name="seed" size="10" value="0" onkeypress="preventNonNumericalInput(event)" inputmode="numeric"> <input id="random_seed" name="random_seed" type="checkbox" checked><label for="random_seed">Random</label></td></tr>
-                    <tr class="pl-5"><td><label for="num_outputs_total">Number of Images:</label></td>
+                    <tr class="pl-5"><td><label for="num_outputs_total">Count of Images:</label></td>
                         <td><input id="num_outputs_total" name="num_outputs_total" value="1" type="number" value="1" min="1" step="1" onkeypres"="preventNonNumericalInput(event)" inputmode="numeric">
                         <label><small>(total)</small></label>
                         <input id="num_outputs_parallel" name="num_outputs_parallel" value="1" type="number" value="1" min="1" step="1" onkeypress="preventNonNumericalInput(event)" inputmode="numeric">
diff --git a/ui/media/css/main.css b/ui/media/css/main.css
index 009c13d5..5ab2506a 100644
--- a/ui/media/css/main.css
+++ b/ui/media/css/main.css
@@ -2040,3 +2040,7 @@ div#enlarge-buttons {
     padding-top: 6pt;
     color: var(--small-label-color);
 }
+
+.imgCount {
+    width: 42pt;
+}
\ No newline at end of file
diff --git a/ui/media/js/main.js b/ui/media/js/main.js
index bff87b3b..bfe73faf 100644
--- a/ui/media/js/main.js
+++ b/ui/media/js/main.js
@@ -1163,7 +1163,7 @@ function createTask(task) {
     }
 
     taskConfig += `<div class="taskConfigData">${createTaskConfig(task)}</span></div></div>`
-
+    let imgCount = task.numOutputsTotal
     let taskEntry = document.createElement("div")
     taskEntry.id = `imageTaskContainer-${Date.now()}`
     taskEntry.className = "imageTaskContainer"
@@ -1172,6 +1172,8 @@ function createTask(task) {
                                 <div class="taskStatusLabel">Enqueued</div>
                                 <button class="secondaryButton stopTask"><i class="fa-solid fa-xmark"></i> Cancel</button>
                                 <button class="tertiaryButton useSettings"><i class="fa-solid fa-redo"></i> Use these settings</button>
+                                <label for="count">Count:</label>
+                                <input class="imgCount" type="number" id="imgCount" name="count" min="1" inputmode="numeric" value=${imgCount}>
                                 <div class="preview-prompt"></div>
                                 <div class="taskConfig">${taskConfig}</div>
                                 <div class="outputMsg"></div>
@@ -1234,6 +1236,16 @@ function createTask(task) {
     task["previewPrompt"] = taskEntry.querySelector(".preview-prompt")
     task["progressBar"] = taskEntry.querySelector(".progress-bar")
     task["stopTask"] = taskEntry.querySelector(".stopTask")
+    task["imgCount"] = taskEntry.querySelector(".imgCount")
+
+    task["imgCount"].addEventListener("change", (e) => {
+        e.stopPropagation()
+        task.batchCount = e.target.value / (task.numOutputsTotal / task.batchCount)
+        task.numOutputsTotal = e.target.value
+    })
+    task["imgCount"].addEventListener("click", (e) => {
+        e.stopPropagation(); // Prevent propagation of the click event to the parent
+    });
 
     task["stopTask"].addEventListener("click", (e) => {
         e.stopPropagation()
diff --git a/ui/media/js/task-manager.js b/ui/media/js/task-manager.js
index 71fd7803..85047616 100644
--- a/ui/media/js/task-manager.js
+++ b/ui/media/js/task-manager.js
@@ -165,12 +165,14 @@ function getTaskUpdater(task, reqBody, outputContainer) {
                     task["taskStatusLabel"].innerText = "Waiting"
                     task["taskStatusLabel"].classList.add("waitingTaskLabel")
                     task["taskStatusLabel"].classList.remove("activeTaskLabel")
+                    task["imgCount"].setAttribute("disabled", true)
                     break
                 case SD.TaskStatus.processing:
                 case SD.TaskStatus.completed:
                     task["taskStatusLabel"].innerText = "Processing"
                     task["taskStatusLabel"].classList.add("activeTaskLabel")
                     task["taskStatusLabel"].classList.remove("waitingTaskLabel")
+                    task["imgCount"].setAttribute("disabled", true)
                     break
                 case SD.TaskStatus.stopped:
                     break