From ab9a2e717d54531b93fa0b48cd2811bc6fe58c0f Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Fri, 13 Sep 2024 19:18:54 -0500
Subject: [PATCH 01/20] comment on redundancy

---
 scripts/spartan/worker.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/scripts/spartan/worker.py b/scripts/spartan/worker.py
index f4eed0a..f9fe383 100644
--- a/scripts/spartan/worker.py
+++ b/scripts/spartan/worker.py
@@ -361,6 +361,7 @@ class Worker:
                     mode = 'img2img'  # for use in checking script compat
                     images = []
                     for image in init_images:
+                        # looks redundant when encode_pil...() could be used, but it does not support all file formats. E.g. AVIF
                         buffer = io.BytesIO()
                         image.save(buffer, format="PNG")
                         image = 'data:image/png;base64,' + str(base64.b64encode(buffer.getvalue()), 'utf-8')

From b2118c8cae580a1790a66903a17d5a7d026014fc Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Mon, 16 Sep 2024 22:55:42 -0500
Subject: [PATCH 02/20] remove refresh button, updates will be automatic

---
 javascript/distributed.js | 15 ++++++++++++++-
 scripts/spartan/ui.py     |  4 ++--
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/javascript/distributed.js b/javascript/distributed.js
index 5827386..e72f81f 100644
--- a/javascript/distributed.js
+++ b/javascript/distributed.js
@@ -1,4 +1,17 @@
 
 function confirm_restart_workers(_) {
     return confirm('Restart remote workers?')
-}
\ No newline at end of file
+}
+
+// live updates
+function update() {
+    try {
+        let refresh_button = document.getElementById('distributed-refresh-status')
+        refresh_button.click()
+    } catch (e) {
+        if (!(e instanceof TypeError)) {
+            throw e
+        }
+    }
+}
+setInterval(update, 1500)
\ No newline at end of file
diff --git a/scripts/spartan/ui.py b/scripts/spartan/ui.py
index 5aac350..88b03e4 100644
--- a/scripts/spartan/ui.py
+++ b/scripts/spartan/ui.py
@@ -236,8 +236,8 @@ class UI:
                         info='top-most message is newest'
                     )
 
-                    refresh_status_btn = gradio.Button(value='Refresh 🔄', size='sm')
-                    refresh_status_btn.click(self.status_btn, inputs=[], outputs=[jobs, status, logs])
+                    refresh_status_btn = gradio.Button(value='Refresh 🔄', size='sm', elem_id='distributed-refresh-status', visible=False)
+                    refresh_status_btn.click(self.status_btn, inputs=[], outputs=[jobs, status, logs], show_progress=False)
 
                     status_tab.select(fn=self.status_btn, inputs=[], outputs=[jobs, status, logs])
                     components += [status, jobs, logs, refresh_status_btn]

From 3e1a895c80cfe16b6dc21a442fb8bca8320b7a00 Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Tue, 17 Sep 2024 02:23:47 -0500
Subject: [PATCH 03/20] fix state inconsistency at startup, connection error
 handling

---
 scripts/spartan/worker.py | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/scripts/spartan/worker.py b/scripts/spartan/worker.py
index f9fe383..b47cf4b 100644
--- a/scripts/spartan/worker.py
+++ b/scripts/spartan/worker.py
@@ -641,6 +641,7 @@ class Worker:
             return []
 
     def load_options(self, model, vae=None):
+        failure_msg = f"failed to load options for worker '{self.label}'"
         if self.master:
             return
 
@@ -654,17 +655,25 @@ class Worker:
         if vae is not None:
             payload['sd_vae'] = vae
 
-        self.set_state(State.WORKING)
+        state_cache = self.state
+        self.set_state(State.WORKING, expect_cycle=True) # may already be WORKING if called by worker.request()
         start = time.time()
-        response = self.session.post(
-            self.full_url("options"),
-            json=payload
-        )
+        try:
+            response = self.session.post(
+                self.full_url("options"),
+                json=payload
+            )
+        except requests.exceptions.RequestException:
+            self.set_state(State.UNAVAILABLE)
+            logger.debug(f"{failure_msg} (connection error... OOM?)")
+            return
+
         elapsed = time.time() - start
-        self.set_state(State.IDLE)
+        if state_cache != State.WORKING: # see above comment, this lets caller determine when worker is IDLE
+            self.set_state(State.IDLE)
 
         if response.status_code != 200:
-            logger.debug(f"failed to load options for worker '{self.label}'")
+            logger.debug(failure_msg)
         else:
             logger.debug(f"worker '{self.label}' loaded weights in {elapsed:.2f}s")
             self.loaded_model = model_name

From 8ac2b3c190fec1a323f6c742437d049bf334ca1a Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Tue, 17 Sep 2024 02:29:26 -0500
Subject: [PATCH 04/20] increase visibility

---
 scripts/spartan/worker.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scripts/spartan/worker.py b/scripts/spartan/worker.py
index b47cf4b..c82d80a 100644
--- a/scripts/spartan/worker.py
+++ b/scripts/spartan/worker.py
@@ -665,7 +665,7 @@ class Worker:
             )
         except requests.exceptions.RequestException:
             self.set_state(State.UNAVAILABLE)
-            logger.debug(f"{failure_msg} (connection error... OOM?)")
+            logger.error(f"{failure_msg} (connection error... OOM?)")
             return
 
         elapsed = time.time() - start

From f5331b0bdb0ab1f44ebb847e82bbe7fadd79805a Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Thu, 19 Sep 2024 01:16:36 -0500
Subject: [PATCH 05/20] update for both tabs only when visible

---
 javascript/distributed.js | 10 ++++++++--
 scripts/spartan/world.py  |  3 ---
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/javascript/distributed.js b/javascript/distributed.js
index e72f81f..ba54cae 100644
--- a/javascript/distributed.js
+++ b/javascript/distributed.js
@@ -6,8 +6,14 @@ function confirm_restart_workers(_) {
 // live updates
 function update() {
     try {
-        let refresh_button = document.getElementById('distributed-refresh-status')
-        refresh_button.click()
+        let currentTab = get_uiCurrentTabContent()
+        let buttons = document.querySelectorAll('#distributed-refresh-status')
+        for(let i = 0; i < buttons.length; i++) {
+            if(currentTab.contains(buttons[i])) {
+                buttons[i].click()
+                break
+            }
+        }
     } catch (e) {
         if (!(e instanceof TypeError)) {
             throw e
diff --git a/scripts/spartan/world.py b/scripts/spartan/world.py
index f08faf3..bae3605 100644
--- a/scripts/spartan/world.py
+++ b/scripts/spartan/world.py
@@ -103,7 +103,6 @@ class World:
     def __repr__(self):
         return f"{len(self._workers)} workers"
 
-
     def default_batch_size(self) -> int:
         """the amount of images/total images requested that a worker would compute if conditions were perfect and
         each worker generated at the same speed. assumes one batch only"""
@@ -393,7 +392,6 @@ class World:
             self.initialized = True
             logger.debug("world initialized!")
 
-
     def get_workers(self):
         filtered: List[Worker] = []
         for worker in self._workers:
@@ -753,7 +751,6 @@ class World:
         for worker in self._workers:
             worker.restart()
 
-
     def inject_model_dropdown_handler(self):
         if self.config().get('enabled', False): # TODO avoid access from config()
             return

From ce2d3254d47e7b13a0eef5ff3d842ff78b67c299 Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Sun, 22 Sep 2024 11:59:00 -0500
Subject: [PATCH 06/20] use InputAccordion for main Accordion

---
 scripts/spartan/ui.py | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/scripts/spartan/ui.py b/scripts/spartan/ui.py
index 88b03e4..2535107 100644
--- a/scripts/spartan/ui.py
+++ b/scripts/spartan/ui.py
@@ -9,6 +9,7 @@ from .shared import logger, LOG_LEVEL, gui_handler
 from .worker import State
 from modules.call_queue import queue_lock
 from modules import progress
+from modules.ui_components import InputAccordion
 
 worker_select_dropdown = None
 
@@ -210,15 +211,9 @@ class UI:
         components = []
 
         with gradio.Blocks(variant='compact'):  # Group() and Box() remove spacing
-            with gradio.Accordion(label='Distributed', open=False):
-                # https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/6109#issuecomment-1403315784
-                main_toggle = gradio.Checkbox(  # main on/off ext. toggle
-                    elem_id='enable',
-                    label='Enable',
-                    value=self.world.enabled if self.world.enabled is not None else True,
-                    interactive=True
-                )
+            with InputAccordion(label='Distributed', open=False, value=self.world.config().get('enabled', False), elem_id='enable') as main_toggle:
                 main_toggle.input(self.main_toggle_btn)
+                setattr(main_toggle.accordion, 'do_not_save_to_config', True) # InputAccordion is really a CheckBox
                 components.append(main_toggle)
 
                 with gradio.Tab('Status') as status_tab:

From 429adb26598b0d5fbcf656d16659d02debf3c4c2 Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Sun, 22 Sep 2024 15:22:32 -0500
Subject: [PATCH 07/20] allow t2i and i2i to be enabled/disabled separately
 from one another

---
 scripts/distributed.py     | 15 +++++++++++----
 scripts/spartan/pmodels.py |  1 +
 scripts/spartan/ui.py      | 16 ++++++++++++----
 scripts/spartan/world.py   |  3 +++
 4 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/scripts/distributed.py b/scripts/distributed.py
index 492d4d4..524e73b 100644
--- a/scripts/distributed.py
+++ b/scripts/distributed.py
@@ -61,7 +61,7 @@ class DistributedScript(scripts.Script):
         return scripts.AlwaysVisible
 
     def ui(self, is_img2img):
-        extension_ui = UI(world=self.world)
+        extension_ui = UI(world=self.world, is_img2img=is_img2img)
         # root, api_exposed = extension_ui.create_ui()
         components = extension_ui.create_ui()
 
@@ -218,8 +218,12 @@ class DistributedScript(scripts.Script):
     # p's type is
     # "modules.processing.StableDiffusionProcessing*"
     def before_process(self, p, *args):
-        if not self.world.enabled:
-            logger.debug("extension is disabled")
+        is_img2img = getattr(p, 'init_images', False)
+        if is_img2img and self.world.enabled_i2i is False:
+            logger.debug("extension is disabled for i2i")
+            return
+        elif not is_img2img and self.world.enabled is False:
+            logger.debug("extension is disabled for t2i")
             return
         self.world.update(p)
 
@@ -352,7 +356,10 @@ class DistributedScript(scripts.Script):
         return
 
     def postprocess(self, p, processed, *args):
-        if not self.world.enabled:
+        is_img2img = getattr(p, 'init_images', False)
+        if is_img2img and self.world.enabled_i2i is False:
+            return
+        elif not is_img2img and self.world.enabled is False:
             return
 
         if self.master_start is not None:
diff --git a/scripts/spartan/pmodels.py b/scripts/spartan/pmodels.py
index e349dc8..10a5e0a 100644
--- a/scripts/spartan/pmodels.py
+++ b/scripts/spartan/pmodels.py
@@ -41,5 +41,6 @@ class ConfigModel(BaseModel):
     )
     job_timeout: Optional[int] = Field(default=3)
     enabled: Optional[bool] = Field(description="Whether the extension as a whole should be active or disabled", default=True)
+    enabled_i2i: Optional[bool] = Field(description="Same as above but for image to image", default=True)
     complement_production: Optional[bool] = Field(description="Whether to generate complementary images to prevent under-utilizing hardware", default=True)
     step_scaling: Optional[bool] = Field(description="Whether to downscale requested steps in order to meet time constraints", default=False)
diff --git a/scripts/spartan/ui.py b/scripts/spartan/ui.py
index 2535107..82ff504 100644
--- a/scripts/spartan/ui.py
+++ b/scripts/spartan/ui.py
@@ -17,9 +17,10 @@ worker_select_dropdown = None
 class UI:
     """extension user interface related things"""
 
-    def __init__(self, world):
+    def __init__(self, world, is_img2img):
         self.world = world
         self.original_model_dropdown_handler = opts.data_labels.get('sd_model_checkpoint').onchange
+        self.is_img2img = is_img2img
 
     # handlers
     @staticmethod
@@ -186,11 +187,15 @@ class UI:
         self.world.save_config()
 
     def main_toggle_btn(self):
-        self.world.enabled = not self.world.enabled
+        if self.is_img2img:
+            self.world.enabled_i2i = not self.world.enabled_i2i
+        else:
+            self.world.enabled = not self.world.enabled
+
         self.world.save_config()
 
         # restore vanilla sdwui handler for model dropdown if extension is disabled or inject if otherwise
-        if not self.world.enabled:
+        if not self.world.enabled and not self.world.enabled_i2i:
             model_dropdown = opts.data_labels.get('sd_model_checkpoint')
             if self.original_model_dropdown_handler is not None:
                 model_dropdown.onchange = self.original_model_dropdown_handler
@@ -209,9 +214,12 @@ class UI:
     def create_ui(self):
         """creates the extension UI and returns relevant components"""
         components = []
+        elem_id = 'enabled'
+        if self.is_img2img:
+            elem_id += '_i2i'
 
         with gradio.Blocks(variant='compact'):  # Group() and Box() remove spacing
-            with InputAccordion(label='Distributed', open=False, value=self.world.config().get('enabled', False), elem_id='enable') as main_toggle:
+            with InputAccordion(label='Distributed', open=False, value=self.world.config().get(elem_id), elem_id=elem_id) as main_toggle:
                 main_toggle.input(self.main_toggle_btn)
                 setattr(main_toggle.accordion, 'do_not_save_to_config', True) # InputAccordion is really a CheckBox
                 components.append(main_toggle)
diff --git a/scripts/spartan/world.py b/scripts/spartan/world.py
index bae3605..73cc742 100644
--- a/scripts/spartan/world.py
+++ b/scripts/spartan/world.py
@@ -91,6 +91,7 @@ class World:
         self.verify_remotes = verify_remotes
         self.thin_client_mode = False
         self.enabled = True
+        self.enabled_i2i = True
         self.is_dropdown_handler_injected = False
         self.complement_production = True
         self.step_scaling = False
@@ -671,6 +672,7 @@ class World:
         sh.benchmark_payload = Benchmark_Payload(**config.benchmark_payload)
         self.job_timeout = config.job_timeout
         self.enabled = config.enabled
+        self.enabled_i2i = config.enabled_i2i
         self.complement_production = config.complement_production
         self.step_scaling = config.step_scaling
 
@@ -686,6 +688,7 @@ class World:
             benchmark_payload=sh.benchmark_payload,
             job_timeout=self.job_timeout,
             enabled=self.enabled,
+            enabled_i2i=self.enabled_i2i,
             complement_production=self.complement_production,
             step_scaling=self.step_scaling
         )

From 897259231971fe24c011c730effef109f6e7ec04 Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Sun, 22 Sep 2024 15:50:24 -0500
Subject: [PATCH 08/20] fix distributed inpainting for certain file formats

---
 scripts/spartan/worker.py | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/scripts/spartan/worker.py b/scripts/spartan/worker.py
index c82d80a..28616b9 100644
--- a/scripts/spartan/worker.py
+++ b/scripts/spartan/worker.py
@@ -15,6 +15,7 @@ from modules.shared import cmd_opts
 from modules.shared import state as master_state
 from . import shared as sh
 from .shared import logger, warmup_samples, LOG_LEVEL
+from PIL import Image
 
 try:
     from webui import server_name
@@ -40,6 +41,13 @@ class State(Enum):
     DISABLED = 5
 
 
+# looks redundant when encode_pil...() could be used, but it does not support all file formats. E.g. AVIF
+def pil_to_64(image: Image) -> str:
+     buffer = io.BytesIO()
+     image.save(buffer, format="PNG")
+     return 'data:image/png;base64,' + str(base64.b64encode(buffer.getvalue()), 'utf-8')
+
+
 class Worker:
     """
         This class represents a worker node in a distributed computing setup.
@@ -361,11 +369,7 @@ class Worker:
                     mode = 'img2img'  # for use in checking script compat
                     images = []
                     for image in init_images:
-                        # looks redundant when encode_pil...() could be used, but it does not support all file formats. E.g. AVIF
-                        buffer = io.BytesIO()
-                        image.save(buffer, format="PNG")
-                        image = 'data:image/png;base64,' + str(base64.b64encode(buffer.getvalue()), 'utf-8')
-                        images.append(image)
+                        images.append(pil_to_64(image))
                     payload['init_images'] = images
 
                 alwayson_scripts = payload.get('alwayson_scripts', None)  # key may not always exist, benchmarking being one example
@@ -402,9 +406,7 @@ class Worker:
                 # if an image mask is present
                 image_mask = payload.get('image_mask', None)
                 if image_mask is not None:
-                    image_b64 = encode_pil_to_base64(image_mask)
-                    image_b64 = str(image_b64, 'utf-8')
-                    payload['mask'] = image_b64
+                    payload['mask'] = pil_to_64(image_mask)
                     del payload['image_mask']
 
                 # see if there is anything else wrong with serializing to payload

From 4023845658e2d1a7b93953e77b875adc3a61c70a Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Sun, 22 Sep 2024 19:09:20 -0500
Subject: [PATCH 09/20] adetailer support via injecting results earlier which
 should also help other post-processing based extensions. However, more
 accurate infotext and autosave temporarily broken from this

---
 scripts/distributed.py | 90 +++++++++++++++++++++++-------------------
 1 file changed, 49 insertions(+), 41 deletions(-)

diff --git a/scripts/distributed.py b/scripts/distributed.py
index 524e73b..839d93b 100644
--- a/scripts/distributed.py
+++ b/scripts/distributed.py
@@ -13,6 +13,7 @@ import time
 from threading import Thread
 from typing import List
 import gradio
+from torchvision.transforms import ToTensor
 import urllib3
 from PIL import Image
 from modules import processing
@@ -71,7 +72,7 @@ class DistributedScript(scripts.Script):
         # return some components that should be exposed to the api
         return components
 
-    def add_to_gallery(self, processed, p):
+    def add_to_gallery(self, pp, p):
         """adds generated images to the image gallery after waiting for all workers to finish"""
 
         def processed_inject_image(image, info_index, save_path_override=None, grid=False, response=None):
@@ -101,17 +102,20 @@ class DistributedScript(scripts.Script):
                 processed_inject_image(image=image, info_index=0, response=response)
                 return
 
-            processed.all_seeds.append(seed)
-            processed.all_subseeds.append(subseed)
-            processed.all_negative_prompts.append(negative_prompt)
-            processed.all_prompts.append(pos_prompt)
-            processed.images.append(image)  # actual received image
+            p.seeds.append(seed)
+            p.subseeds.append(subseed)
+            p.negative_prompts.append(negative_prompt)
+            p.prompts.append(pos_prompt)
+
+            transform = ToTensor()
+            pp.images.append(transform(image))  # actual received image
+
 
             # generate info-text string
             # modules.ui_common -> update_generation_info renders to html below gallery
             images_per_batch = p.n_iter * p.batch_size
             # zero-indexed position of image in total batch (so including master results)
-            true_image_pos = len(processed.images) - 1
+            true_image_pos = len(pp.images) - 1
             num_remote_images = images_per_batch * p.batch_size
             if p.n_iter > 1:  # if splitting by batch count
                 num_remote_images *= p.n_iter - 1
@@ -120,28 +124,28 @@ class DistributedScript(scripts.Script):
                          f"info-index: {info_index}")
 
             if self.world.thin_client_mode:
-                p.all_negative_prompts = processed.all_negative_prompts
+                p.all_negative_prompts = pp.all_negative_prompts
 
-            try:
-                info_text = image_info_post['infotexts'][i]
-            except IndexError:
-                if not grid:
-                    logger.warning(f"image {true_image_pos + 1} was missing info-text")
-                info_text = processed.infotexts[0]
-            info_text += f", Worker Label: {job.worker.label}"
-            processed.infotexts.append(info_text)
+            # try:
+            #     info_text = image_info_post['infotexts'][i]
+            # except IndexError:
+            #     if not grid:
+            #         logger.warning(f"image {true_image_pos + 1} was missing info-text")
+            #     info_text = processed.infotexts[0]
+            # info_text += f", Worker Label: {job.worker.label}"
+            # processed.infotexts.append(info_text)
 
             # automatically save received image to local disk if desired
-            if cmd_opts.distributed_remotes_autosave:
-                save_image(
-                    image=image,
-                    path=p.outpath_samples if save_path_override is None else save_path_override,
-                    basename="",
-                    seed=seed,
-                    prompt=pos_prompt,
-                    info=info_text,
-                    extension=opts.samples_format
-                )
+            #if cmd_opts.distributed_remotes_autosave:
+            #    save_image(
+            #        image=image,
+            #        path=p.outpath_samples if save_path_override is None else save_path_override,
+            #        basename="",
+            #        seed=seed,
+            #        prompt=pos_prompt,
+            #        info=info_text,
+            #        extension=opts.samples_format
+            #    )
 
         # get master ipm by estimating based on worker speed
         master_elapsed = time.time() - self.master_start
@@ -197,22 +201,11 @@ class DistributedScript(scripts.Script):
             logger.critical("couldn't collect any responses, the extension will have no effect")
             return
 
-        # generate and inject grid
-        if opts.return_grid and len(processed.images) > 1:
-            grid = image_grid(processed.images, len(processed.images))
-            processed_inject_image(
-                image=grid,
-                info_index=0,
-                save_path_override=p.outpath_grids,
-                grid=True,
-                response=donor_worker.response
-            )
-
         # cleanup after we're doing using all the responses
         for worker in self.world.get_workers():
             worker.response = None
 
-        p.batch_size = len(processed.images)
+        p.batch_size = len(pp.images)
         return
 
     # p's type is
@@ -351,11 +344,26 @@ class DistributedScript(scripts.Script):
         self.master_start = time.time()
 
         # generate images assigned to local machine
-        p.do_not_save_grid = True  # don't generate grid from master as we are doing this later.
+        # p.do_not_save_grid = True  # don't generate grid from master as we are doing this later.
         self.runs_since_init += 1
         return
 
-    def postprocess(self, p, processed, *args):
+    # def postprocess(self, p, processed, *args):
+    #     is_img2img = getattr(p, 'init_images', False)
+    #     if is_img2img and self.world.enabled_i2i is False:
+    #         return
+    #     elif not is_img2img and self.world.enabled is False:
+    #         return
+
+    #     if self.master_start is not None:
+    #         self.add_to_gallery(p=p, processed=processed)
+
+    #     # restore process_images_inner if it was monkey-patched
+    #     processing.process_images_inner = self.original_process_images_inner
+    #     # save any dangling state to prevent load_config in next iteration overwriting it
+    #     self.world.save_config()
+
+    def postprocess_batch_list(self, p, pp, *args, **kwargs):
         is_img2img = getattr(p, 'init_images', False)
         if is_img2img and self.world.enabled_i2i is False:
             return
@@ -363,7 +371,7 @@ class DistributedScript(scripts.Script):
             return
 
         if self.master_start is not None:
-            self.add_to_gallery(p=p, processed=processed)
+            self.add_to_gallery(p=p, pp=pp)
 
         # restore process_images_inner if it was monkey-patched
         processing.process_images_inner = self.original_process_images_inner

From 9de268702df0dd4a6eb4708d1d5724bf2967c211 Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Mon, 23 Sep 2024 11:28:37 -0500
Subject: [PATCH 10/20] return control to webui when adetailer is set to skip
 img2img pass

---
 scripts/distributed.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/scripts/distributed.py b/scripts/distributed.py
index 839d93b..c1d68ae 100644
--- a/scripts/distributed.py
+++ b/scripts/distributed.py
@@ -206,6 +206,7 @@ class DistributedScript(scripts.Script):
             worker.response = None
 
         p.batch_size = len(pp.images)
+        webui_state.textinfo = ""
         return
 
     # p's type is
@@ -231,6 +232,14 @@ class DistributedScript(scripts.Script):
                 continue
             title = script.title()
 
+            if title == "ADetailer":
+                adetailer_args = p.script_args[script.args_from:script.args_to]
+
+                # InputAccordion main toggle, skip img2img toggle
+                if adetailer_args[0] and adetailer_args[1]:
+                    logger.debug(f"adetailer is skipping img2img, returning control to wui")
+                    return
+
             # check for supported scripts
             if title == "ControlNet":
                 # grab all controlnet units

From 32bc086ec655e9b27d4716034191b210fe3b1290 Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Tue, 24 Sep 2024 05:21:16 -0500
Subject: [PATCH 11/20] prevent ext. toggle inversion if ui desyncs

---
 scripts/spartan/ui.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/scripts/spartan/ui.py b/scripts/spartan/ui.py
index 82ff504..b87c724 100644
--- a/scripts/spartan/ui.py
+++ b/scripts/spartan/ui.py
@@ -13,7 +13,6 @@ from modules.ui_components import InputAccordion
 
 worker_select_dropdown = None
 
-
 class UI:
     """extension user interface related things"""
 
@@ -186,11 +185,15 @@ class UI:
             worker.session.auth = (user, password)
         self.world.save_config()
 
-    def main_toggle_btn(self):
+    def main_toggle_btn(self, state):
         if self.is_img2img:
-            self.world.enabled_i2i = not self.world.enabled_i2i
+            if self.world.enabled_i2i == state: # just prevents a redundant config save if ui desyncs
+                return
+            self.world.enabled_i2i = state
         else:
-            self.world.enabled = not self.world.enabled
+            if self.world.enabled == state:
+                return
+            self.world.enabled = state
 
         self.world.save_config()
 
@@ -220,7 +223,7 @@ class UI:
 
         with gradio.Blocks(variant='compact'):  # Group() and Box() remove spacing
             with InputAccordion(label='Distributed', open=False, value=self.world.config().get(elem_id), elem_id=elem_id) as main_toggle:
-                main_toggle.input(self.main_toggle_btn)
+                main_toggle.input(self.main_toggle_btn, inputs=[main_toggle])
                 setattr(main_toggle.accordion, 'do_not_save_to_config', True) # InputAccordion is really a CheckBox
                 components.append(main_toggle)
 

From 8a810b7714f1c4ba43f64a50b7e363367462dbed Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Fri, 27 Sep 2024 20:21:11 -0500
Subject: [PATCH 12/20] once again inject infotext from remote results into
 gallery

---
 scripts/distributed.py   | 70 +++++++++++++---------------------------
 scripts/spartan/world.py |  1 +
 2 files changed, 24 insertions(+), 47 deletions(-)

diff --git a/scripts/distributed.py b/scripts/distributed.py
index c1d68ae..059d9e2 100644
--- a/scripts/distributed.py
+++ b/scripts/distributed.py
@@ -75,9 +75,9 @@ class DistributedScript(scripts.Script):
     def add_to_gallery(self, pp, p):
         """adds generated images to the image gallery after waiting for all workers to finish"""
 
-        def processed_inject_image(image, info_index, save_path_override=None, grid=False, response=None):
-            image_params: json = response['parameters']
-            image_info_post: json = json.loads(response["info"])  # image info known after processing
+        def processed_inject_image(image, info_index, save_path_override=None, grid=False, job=None):
+            image_params: json = job.worker.response['parameters']
+            image_info_post: json = json.loads(job.worker.response["info"])  # image info known after processing
             num_response_images = image_params["batch_size"] * image_params["n_iter"]
 
             seed = None
@@ -99,7 +99,7 @@ class DistributedScript(scripts.Script):
             except IndexError:
                 # like with controlnet masks, there isn't always full post-gen info, so we use the first images'
                 logger.debug(f"Image at index {i} for '{job.worker.label}' was missing some post-generation data")
-                processed_inject_image(image=image, info_index=0, response=response)
+                processed_inject_image(image=image, info_index=0, job=job)
                 return
 
             p.seeds.append(seed)
@@ -110,12 +110,10 @@ class DistributedScript(scripts.Script):
             transform = ToTensor()
             pp.images.append(transform(image))  # actual received image
 
-
-            # generate info-text string
             # modules.ui_common -> update_generation_info renders to html below gallery
             images_per_batch = p.n_iter * p.batch_size
-            # zero-indexed position of image in total batch (so including master results)
-            true_image_pos = len(pp.images) - 1
+            # zero-indexed position of image in gallery (so including master/local results)
+            true_image_pos = (len(pp.images) - 1) + (not p.do_not_save_grid)
             num_remote_images = images_per_batch * p.batch_size
             if p.n_iter > 1:  # if splitting by batch count
                 num_remote_images *= p.n_iter - 1
@@ -126,26 +124,8 @@ class DistributedScript(scripts.Script):
             if self.world.thin_client_mode:
                 p.all_negative_prompts = pp.all_negative_prompts
 
-            # try:
-            #     info_text = image_info_post['infotexts'][i]
-            # except IndexError:
-            #     if not grid:
-            #         logger.warning(f"image {true_image_pos + 1} was missing info-text")
-            #     info_text = processed.infotexts[0]
-            # info_text += f", Worker Label: {job.worker.label}"
-            # processed.infotexts.append(info_text)
-
-            # automatically save received image to local disk if desired
-            #if cmd_opts.distributed_remotes_autosave:
-            #    save_image(
-            #        image=image,
-            #        path=p.outpath_samples if save_path_override is None else save_path_override,
-            #        basename="",
-            #        seed=seed,
-            #        prompt=pos_prompt,
-            #        info=info_text,
-            #        extension=opts.samples_format
-            #    )
+            # saves final position of image in gallery so that we can later modify the correct infotext
+            job.gallery_map.append(true_image_pos)
 
         # get master ipm by estimating based on worker speed
         master_elapsed = time.time() - self.master_start
@@ -195,15 +175,12 @@ class DistributedScript(scripts.Script):
                 image = Image.open(io.BytesIO(image_bytes))
 
                 # inject image
-                processed_inject_image(image=image, info_index=i, response=job.worker.response)
+                processed_inject_image(image=image, info_index=i, job=job)
 
         if donor_worker is None:
             logger.critical("couldn't collect any responses, the extension will have no effect")
             return
 
-        # cleanup after we're doing using all the responses
-        for worker in self.world.get_workers():
-            worker.response = None
 
         p.batch_size = len(pp.images)
         webui_state.textinfo = ""
@@ -357,21 +334,6 @@ class DistributedScript(scripts.Script):
         self.runs_since_init += 1
         return
 
-    # def postprocess(self, p, processed, *args):
-    #     is_img2img = getattr(p, 'init_images', False)
-    #     if is_img2img and self.world.enabled_i2i is False:
-    #         return
-    #     elif not is_img2img and self.world.enabled is False:
-    #         return
-
-    #     if self.master_start is not None:
-    #         self.add_to_gallery(p=p, processed=processed)
-
-    #     # restore process_images_inner if it was monkey-patched
-    #     processing.process_images_inner = self.original_process_images_inner
-    #     # save any dangling state to prevent load_config in next iteration overwriting it
-    #     self.world.save_config()
-
     def postprocess_batch_list(self, p, pp, *args, **kwargs):
         is_img2img = getattr(p, 'init_images', False)
         if is_img2img and self.world.enabled_i2i is False:
@@ -382,6 +344,20 @@ class DistributedScript(scripts.Script):
         if self.master_start is not None:
             self.add_to_gallery(p=p, pp=pp)
 
+
+    def postprocess(self, p, processed, *args):
+        for job in self.world.jobs:
+            if job.worker.response is not None:
+                for i, v in enumerate(job.gallery_map):
+                    infotext = json.loads(job.worker.response['info'])['infotexts'][i]
+                    logger.debug(f"replacing image {v}'s infotext with\n"
+                                    f"> '{infotext}'")
+                    infotext += f", Worker Label: {job.worker.label}"
+                    processed.infotexts[v] = infotext
+
+        # cleanup
+        for worker in self.world.get_workers():
+            worker.response = None
         # restore process_images_inner if it was monkey-patched
         processing.process_images_inner = self.original_process_images_inner
         # save any dangling state to prevent load_config in next iteration overwriting it
diff --git a/scripts/spartan/world.py b/scripts/spartan/world.py
index 73cc742..fc6e4cd 100644
--- a/scripts/spartan/world.py
+++ b/scripts/spartan/world.py
@@ -46,6 +46,7 @@ class Job:
         self.complementary: bool = False
         self.step_override = None
         self.thread = None
+        self.gallery_map: List[int] = []
 
     def __str__(self):
         prefix = ''

From 7349ffd506677e11e692154280593d5a846b366f Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Sun, 29 Sep 2024 23:08:52 -0500
Subject: [PATCH 13/20] beginning of refactoring processed_inject_image

---
 scripts/distributed.py | 90 ++++++++++++++++++++++++------------------
 1 file changed, 51 insertions(+), 39 deletions(-)

diff --git a/scripts/distributed.py b/scripts/distributed.py
index 059d9e2..07b9b8f 100644
--- a/scripts/distributed.py
+++ b/scripts/distributed.py
@@ -72,54 +72,62 @@ class DistributedScript(scripts.Script):
         # return some components that should be exposed to the api
         return components
 
-    def add_to_gallery(self, pp, p):
-        """adds generated images to the image gallery after waiting for all workers to finish"""
+    def api_to_internal(self, job):
+        # takes worker response received from api and returns parsed objects in internal sdwui format. E.g. all_seeds
 
-        def processed_inject_image(image, info_index, save_path_override=None, grid=False, job=None):
-            image_params: json = job.worker.response['parameters']
-            image_info_post: json = json.loads(job.worker.response["info"])  # image info known after processing
-            num_response_images = image_params["batch_size"] * image_params["n_iter"]
-
-            seed = None
-            subseed = None
-            negative_prompt = None
-            pos_prompt = None
+        image_params: json = job.worker.response['parameters']
+        image_info_post: json = json.loads(job.worker.response["info"])  # image info known after processing
+        all_seeds, all_subseeds, all_negative_prompts, all_prompts, images = [], [], [], [], []
 
+        for i in range(len(job.worker.response["images"])):
             try:
-                if num_response_images > 1:
-                    seed = image_info_post['all_seeds'][info_index]
-                    subseed = image_info_post['all_subseeds'][info_index]
-                    negative_prompt = image_info_post['all_negative_prompts'][info_index]
-                    pos_prompt = image_info_post['all_prompts'][info_index]
-                else:
-                    seed = image_info_post['seed']
-                    subseed = image_info_post['subseed']
-                    negative_prompt = image_info_post['negative_prompt']
-                    pos_prompt = image_info_post['prompt']
+                if image_params["batch_size"] * image_params["n_iter"] > 1:
+                    all_seeds.append(image_info_post['all_seeds'][i])
+                    all_subseeds.append(image_info_post['all_subseeds'][i])
+                    all_negative_prompts.append(image_info_post['all_negative_prompts'][i])
+                    all_prompts.append(image_info_post['all_prompts'][i])
+                else: # only a single image received
+                    all_seeds.append(image_info_post['seed'])
+                    all_subseeds.append(image_info_post['subseed'])
+                    all_negative_prompts.append(image_info_post['negative_prompt'])
+                    all_prompts.append(image_info_post['prompt'])
             except IndexError:
-                # like with controlnet masks, there isn't always full post-gen info, so we use the first images'
-                logger.debug(f"Image at index {i} for '{job.worker.label}' was missing some post-generation data")
-                processed_inject_image(image=image, info_index=0, job=job)
-                return
-
-            p.seeds.append(seed)
-            p.subseeds.append(subseed)
-            p.negative_prompts.append(negative_prompt)
-            p.prompts.append(pos_prompt)
+                # # like with controlnet masks, there isn't always full post-gen info, so we use the first images'
+                # logger.debug(f"Image at index {info_index} for '{job.worker.label}' was missing some post-generation data")
+                # self.processed_inject_image(image=image, info_index=0, job=job, p=p)
+                # return
+                logger.critical(f"Image at index {i} for '{job.worker.label}' was missing some post-generation data")
+                continue
 
+            # parse image
+            image_bytes = base64.b64decode(job.worker.response["images"][i])
+            image = Image.open(io.BytesIO(image_bytes))
             transform = ToTensor()
-            pp.images.append(transform(image))  # actual received image
+            images.append(transform(image))
+
+        return all_seeds, all_subseeds, all_negative_prompts, all_prompts, images
+
+    def processed_inject_image(self, job, p, pp):
+        all_seeds, all_subseeds, all_negative_prompts, all_prompts, images = self.api_to_internal(job)
+
+        p.seeds.extend(all_seeds)
+        p.subseeds.extend(all_subseeds)
+        p.negative_prompts.extend(all_negative_prompts)
+        p.prompts.extend(all_prompts)
+
+        for i, image in enumerate(images):
+            pp.images.append(image) # add one image to the gallery
 
             # modules.ui_common -> update_generation_info renders to html below gallery
             images_per_batch = p.n_iter * p.batch_size
             # zero-indexed position of image in gallery (so including master/local results)
             true_image_pos = (len(pp.images) - 1) + (not p.do_not_save_grid)
-            num_remote_images = images_per_batch * p.batch_size
+            num_remote_images = images_per_batch * p.batch_size # the **expected** amount of remote images
             if p.n_iter > 1:  # if splitting by batch count
                 num_remote_images *= p.n_iter - 1
 
-            logger.debug(f"image {true_image_pos + 1}/{self.world.p.batch_size * p.n_iter}, "
-                         f"info-index: {info_index}")
+            logger.debug(f"image {true_image_pos + 1}/{(self.world.p.batch_size * p.n_iter) + (not p.do_not_save_grid) + 1}, "
+                         f"info-index: fix me")
 
             if self.world.thin_client_mode:
                 p.all_negative_prompts = pp.all_negative_prompts
@@ -127,6 +135,9 @@ class DistributedScript(scripts.Script):
             # saves final position of image in gallery so that we can later modify the correct infotext
             job.gallery_map.append(true_image_pos)
 
+    def add_to_gallery(self, pp, p):
+        """adds generated images to the image gallery after waiting for all workers to finish"""
+
         # get master ipm by estimating based on worker speed
         master_elapsed = time.time() - self.master_start
         logger.debug(f"Took master {master_elapsed:.2f}s")
@@ -170,12 +181,13 @@ class DistributedScript(scripts.Script):
                 continue
 
             # visibly add work from workers to the image gallery
-            for i in range(0, len(images)):
-                image_bytes = base64.b64decode(images[i])
-                image = Image.open(io.BytesIO(image_bytes))
+            # for i in range(0, len(images)):
+            #     image_bytes = base64.b64decode(images[i])
+            #     image = Image.open(io.BytesIO(image_bytes))
 
-                # inject image
-                processed_inject_image(image=image, info_index=i, job=job)
+            #     # inject image
+            #     self.processed_inject_image(image=image, info_index=i, job=job, p=p, pp=pp)
+            self.processed_inject_image(job, p, pp)
 
         if donor_worker is None:
             logger.critical("couldn't collect any responses, the extension will have no effect")

From d2d772b810d45ff613e5527130d07171995d08b8 Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Sun, 29 Sep 2024 23:42:32 -0500
Subject: [PATCH 14/20] minor commenting and changes to naming

---
 scripts/distributed.py | 36 +++++++++++++++---------------------
 1 file changed, 15 insertions(+), 21 deletions(-)

diff --git a/scripts/distributed.py b/scripts/distributed.py
index 07b9b8f..d361507 100644
--- a/scripts/distributed.py
+++ b/scripts/distributed.py
@@ -25,7 +25,7 @@ from modules.shared import state as webui_state
 from scripts.spartan.control_net import pack_control_net
 from scripts.spartan.shared import logger
 from scripts.spartan.ui import UI
-from scripts.spartan.world import World, State
+from scripts.spartan.world import World, State, Job
 
 old_sigint_handler = signal.getsignal(signal.SIGINT)
 old_sigterm_handler = signal.getsignal(signal.SIGTERM)
@@ -72,7 +72,7 @@ class DistributedScript(scripts.Script):
         # return some components that should be exposed to the api
         return components
 
-    def api_to_internal(self, job):
+    def api_to_internal(self, job) -> ([], [], [], [], []):
         # takes worker response received from api and returns parsed objects in internal sdwui format. E.g. all_seeds
 
         image_params: json = job.worker.response['parameters']
@@ -107,7 +107,8 @@ class DistributedScript(scripts.Script):
 
         return all_seeds, all_subseeds, all_negative_prompts, all_prompts, images
 
-    def processed_inject_image(self, job, p, pp):
+    def inject_job(self, job: Job, p, pp):
+        """Adds the work completed by one Job via its worker response to the processing and postprocessing objects"""
         all_seeds, all_subseeds, all_negative_prompts, all_prompts, images = self.api_to_internal(job)
 
         p.seeds.extend(all_seeds)
@@ -115,7 +116,7 @@ class DistributedScript(scripts.Script):
         p.negative_prompts.extend(all_negative_prompts)
         p.prompts.extend(all_prompts)
 
-        for i, image in enumerate(images):
+        for image in images:
             pp.images.append(image) # add one image to the gallery
 
             # modules.ui_common -> update_generation_info renders to html below gallery
@@ -126,6 +127,7 @@ class DistributedScript(scripts.Script):
             if p.n_iter > 1:  # if splitting by batch count
                 num_remote_images *= p.n_iter - 1
 
+            # TODO slightly off from changes
             logger.debug(f"image {true_image_pos + 1}/{(self.world.p.batch_size * p.n_iter) + (not p.do_not_save_grid) + 1}, "
                          f"info-index: fix me")
 
@@ -135,8 +137,8 @@ class DistributedScript(scripts.Script):
             # saves final position of image in gallery so that we can later modify the correct infotext
             job.gallery_map.append(true_image_pos)
 
-    def add_to_gallery(self, pp, p):
-        """adds generated images to the image gallery after waiting for all workers to finish"""
+    def update_gallery(self, pp, p):
+        """adds all remotely generated images to the image gallery after waiting for all workers to finish"""
 
         # get master ipm by estimating based on worker speed
         master_elapsed = time.time() - self.master_start
@@ -153,8 +155,7 @@ class DistributedScript(scripts.Script):
         logger.debug("all worker request threads returned")
         webui_state.textinfo = "Distributed - injecting images"
 
-        # some worker which we know has a good response that we can use for generating the grid
-        donor_worker = None
+        received_images = False
         for job in self.world.jobs:
             if job.worker.response is None or job.batch_size < 1 or job.worker.master:
                 continue
@@ -165,8 +166,7 @@ class DistributedScript(scripts.Script):
                 if (job.batch_size * p.n_iter) < len(images):
                     logger.debug(f"requested {job.batch_size} image(s) from '{job.worker.label}', got {len(images)}")
 
-                if donor_worker is None:
-                    donor_worker = job.worker
+                received_images = True
             except KeyError:
                 if job.batch_size > 0:
                     logger.warning(f"Worker '{job.worker.label}' had no images")
@@ -180,20 +180,14 @@ class DistributedScript(scripts.Script):
                     logger.exception(e)
                 continue
 
-            # visibly add work from workers to the image gallery
-            # for i in range(0, len(images)):
-            #     image_bytes = base64.b64decode(images[i])
-            #     image = Image.open(io.BytesIO(image_bytes))
+            # adding the images in
+            self.inject_job(job, p, pp)
 
-            #     # inject image
-            #     self.processed_inject_image(image=image, info_index=i, job=job, p=p, pp=pp)
-            self.processed_inject_image(job, p, pp)
-
-        if donor_worker is None:
+        # TODO fix controlnet masks returned via api having no generation info
+        if received_images is False:
             logger.critical("couldn't collect any responses, the extension will have no effect")
             return
 
-
         p.batch_size = len(pp.images)
         webui_state.textinfo = ""
         return
@@ -354,7 +348,7 @@ class DistributedScript(scripts.Script):
             return
 
         if self.master_start is not None:
-            self.add_to_gallery(p=p, pp=pp)
+            self.update_gallery(p=p, pp=pp)
 
 
     def postprocess(self, p, processed, *args):

From 5eb1af9f73bc9affcba61368b81ceacc0497a8a6 Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Tue, 1 Oct 2024 03:08:10 -0500
Subject: [PATCH 15/20] refactoring, fix for n_iter > 1

---
 scripts/distributed.py         | 28 +++++++++++-----------------
 scripts/spartan/control_net.py |  3 +++
 scripts/spartan/world.py       | 26 +++++++++++++++-----------
 3 files changed, 29 insertions(+), 28 deletions(-)

diff --git a/scripts/distributed.py b/scripts/distributed.py
index d361507..9992729 100644
--- a/scripts/distributed.py
+++ b/scripts/distributed.py
@@ -116,26 +116,19 @@ class DistributedScript(scripts.Script):
         p.negative_prompts.extend(all_negative_prompts)
         p.prompts.extend(all_prompts)
 
-        for image in images:
-            pp.images.append(image) # add one image to the gallery
-
+        num_local = self.world.p.n_iter * self.world.p.batch_size + opts.return_grid
+        num_injected = len(pp.images) - self.world.p.batch_size
+        for i, image in enumerate(images):
             # modules.ui_common -> update_generation_info renders to html below gallery
-            images_per_batch = p.n_iter * p.batch_size
-            # zero-indexed position of image in gallery (so including master/local results)
-            true_image_pos = (len(pp.images) - 1) + (not p.do_not_save_grid)
-            num_remote_images = images_per_batch * p.batch_size # the **expected** amount of remote images
-            if p.n_iter > 1:  # if splitting by batch count
-                num_remote_images *= p.n_iter - 1
-
-            # TODO slightly off from changes
-            logger.debug(f"image {true_image_pos + 1}/{(self.world.p.batch_size * p.n_iter) + (not p.do_not_save_grid) + 1}, "
-                         f"info-index: fix me")
 
+            # TODO probably shouldn't be here
             if self.world.thin_client_mode:
                 p.all_negative_prompts = pp.all_negative_prompts
 
-            # saves final position of image in gallery so that we can later modify the correct infotext
-            job.gallery_map.append(true_image_pos)
+            gallery_index = num_local + num_injected + i # zero-indexed point of image in total gallery
+            job.gallery_map.append(gallery_index) # so we know where to edit infotext
+            pp.images.append(image)
+            logger.debug(f"image {gallery_index + 1}/{self.world.num_total()}")
 
     def update_gallery(self, pp, p):
         """adds all remotely generated images to the image gallery after waiting for all workers to finish"""
@@ -341,6 +334,9 @@ class DistributedScript(scripts.Script):
         return
 
     def postprocess_batch_list(self, p, pp, *args, **kwargs):
+        if p.n_iter != kwargs['batch_number'] + 1: # skip if not the final batch
+            return
+
         is_img2img = getattr(p, 'init_images', False)
         if is_img2img and self.world.enabled_i2i is False:
             return
@@ -356,8 +352,6 @@ class DistributedScript(scripts.Script):
             if job.worker.response is not None:
                 for i, v in enumerate(job.gallery_map):
                     infotext = json.loads(job.worker.response['info'])['infotexts'][i]
-                    logger.debug(f"replacing image {v}'s infotext with\n"
-                                    f"> '{infotext}'")
                     infotext += f", Worker Label: {job.worker.label}"
                     processed.infotexts[v] = infotext
 
diff --git a/scripts/spartan/control_net.py b/scripts/spartan/control_net.py
index dc55cd9..0a58d42 100644
--- a/scripts/spartan/control_net.py
+++ b/scripts/spartan/control_net.py
@@ -58,6 +58,9 @@ def pack_control_net(cn_units) -> dict:
                 unit['mask'] = mask_b64  # mikubill
                 unit['mask_image'] = mask_b64  # forge
 
+        # TODO serialize ResizeMode and other unserializable attributes
+        # https://github.com/Mikubill/sd-webui-controlnet/wiki/API#examples-1
+
         # avoid returning duplicate detection maps since master should return the same one
         unit['save_detected_map'] = False
         # remove anything unserializable
diff --git a/scripts/spartan/world.py b/scripts/spartan/world.py
index fc6e4cd..a1e5ee1 100644
--- a/scripts/spartan/world.py
+++ b/scripts/spartan/world.py
@@ -114,7 +114,7 @@ class World:
     def size(self) -> int:
         """
         Returns:
-            int: The number of nodes currently registered in the world.
+            int: The number of nodes currently registered in the world and in a valid state
         """
         return len(self.get_workers())
 
@@ -274,7 +274,7 @@ class World:
                 gradio.Info("Distributed: benchmarking complete!")
                 self.save_config()
 
-    def get_current_output_size(self) -> int:
+    def num_requested(self) -> int:
         """
         returns how many images would be returned from all jobs
         """
@@ -286,6 +286,11 @@ class World:
 
         return num_images
 
+    def num_total(self) -> int:
+        """how many images should be returned in total, so including local generations and its grid"""
+
+        return self.num_requested() * self.p.n_iter + shared.opts.return_grid
+
     def speed_summary(self) -> str:
         """
         Returns string listing workers by their ipm in descending order.
@@ -472,7 +477,7 @@ class World:
         #######################
 
         # when total number of requested images was not cleanly divisible by world size then we tack the remainder on
-        remainder_images = self.p.batch_size - self.get_current_output_size()
+        remainder_images = self.p.batch_size - self.num_requested()
         if remainder_images >= 1:
             logger.debug(f"The requested number of images({self.p.batch_size}) was not cleanly divisible by the number of realtime nodes({len(self.realtime_jobs())}) resulting in {remainder_images} that will be redistributed")
 
@@ -551,7 +556,7 @@ class World:
         else:
             logger.debug("complementary image production is disabled")
 
-        logger.info(self.distro_summary(payload))
+        logger.info(self.distro_summary())
 
         if self.thin_client_mode is True or self.master_job().batch_size == 0:
             # save original process_images_inner for later so we can restore once we're done
@@ -576,18 +581,17 @@ class World:
                 del self.jobs[last]
             last -= 1
 
-    def distro_summary(self, payload):
-        # iterations = dict(payload)['n_iter']
-        iterations = self.p.n_iter
-        num_returning = self.get_current_output_size()
+    def distro_summary(self):
+        num_returning = self.num_requested()
         num_complementary = num_returning - self.p.batch_size
+
         distro_summary = "Job distribution:\n"
-        distro_summary += f"{self.p.batch_size} * {iterations} iteration(s)"
+        distro_summary += f"{self.p.batch_size} * {self.p.n_iter} iteration(s)"
         if num_complementary > 0:
             distro_summary += f" + {num_complementary} complementary"
-        distro_summary += f": {num_returning} images total\n"
+        distro_summary += f": {num_returning * self.p.n_iter} images total\n"
         for job in self.jobs:
-            distro_summary += f"'{job.worker.label}' - {job.batch_size * iterations} image(s) @ {job.worker.avg_ipm:.2f} ipm\n"
+            distro_summary += f"'{job.worker.label}' - {job.batch_size * self.p.n_iter} image(s) @ {job.worker.avg_ipm:.2f} ipm\n"
         return distro_summary
 
     def config(self) -> dict:

From 979d559803224392c502d565e74489ce4c5b2eef Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Sat, 5 Oct 2024 03:28:22 -0500
Subject: [PATCH 16/20] fix for controlnet

---
 scripts/spartan/control_net.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/scripts/spartan/control_net.py b/scripts/spartan/control_net.py
index 0a58d42..a029aea 100644
--- a/scripts/spartan/control_net.py
+++ b/scripts/spartan/control_net.py
@@ -1,9 +1,12 @@
+# https://github.com/Mikubill/sd-webui-controlnet/wiki/API#examples-1
+
 import copy
 from PIL import Image
 from modules.api.api import encode_pil_to_base64
 from scripts.spartan.shared import logger
 import numpy as np
 import json
+import enum
 
 
 def np_to_b64(image: np.ndarray):
@@ -58,13 +61,14 @@ def pack_control_net(cn_units) -> dict:
                 unit['mask'] = mask_b64  # mikubill
                 unit['mask_image'] = mask_b64  # forge
 
-        # TODO serialize ResizeMode and other unserializable attributes
-        # https://github.com/Mikubill/sd-webui-controlnet/wiki/API#examples-1
+
+        # serialize all enums
+        for k in unit.keys():
+            if isinstance(unit[k], enum.Enum):
+                unit[k] = unit[k].value
 
         # avoid returning duplicate detection maps since master should return the same one
         unit['save_detected_map'] = False
-        # remove anything unserializable
-        del unit['input_mode']
 
     try:
         json.dumps(controlnet)

From 23ec02c2294b01a03b1d36a377eb51040b83235d Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Sat, 5 Oct 2024 03:32:42 -0500
Subject: [PATCH 17/20] fix for controlnet

---
 scripts/distributed.py   | 2 +-
 scripts/spartan/world.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/scripts/distributed.py b/scripts/distributed.py
index 9992729..b3c10f6 100644
--- a/scripts/distributed.py
+++ b/scripts/distributed.py
@@ -128,7 +128,7 @@ class DistributedScript(scripts.Script):
             gallery_index = num_local + num_injected + i # zero-indexed point of image in total gallery
             job.gallery_map.append(gallery_index) # so we know where to edit infotext
             pp.images.append(image)
-            logger.debug(f"image {gallery_index + 1}/{self.world.num_total()}")
+            logger.debug(f"image {gallery_index + 1}/{self.world.num_gallery()}")
 
     def update_gallery(self, pp, p):
         """adds all remotely generated images to the image gallery after waiting for all workers to finish"""
diff --git a/scripts/spartan/world.py b/scripts/spartan/world.py
index a1e5ee1..9ff7ceb 100644
--- a/scripts/spartan/world.py
+++ b/scripts/spartan/world.py
@@ -286,8 +286,8 @@ class World:
 
         return num_images
 
-    def num_total(self) -> int:
-        """how many images should be returned in total, so including local generations and its grid"""
+    def num_gallery(self) -> int:
+        """How many images should appear in the gallery. This includes local generations and a grid(if enabled)"""
 
         return self.num_requested() * self.p.n_iter + shared.opts.return_grid
 

From 622827aab5e6de28591e86c50ca02d15f7875e19 Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Fri, 11 Oct 2024 06:11:10 -0500
Subject: [PATCH 18/20] refactoring

---
 scripts/spartan/pmodels.py | 6 +++---
 scripts/spartan/shared.py  | 2 +-
 scripts/spartan/world.py   | 4 ++--
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/scripts/spartan/pmodels.py b/scripts/spartan/pmodels.py
index 10a5e0a..48fe7a6 100644
--- a/scripts/spartan/pmodels.py
+++ b/scripts/spartan/pmodels.py
@@ -29,13 +29,13 @@ class Worker_Model(BaseModel):
         default=False
     )
     state: Optional[Any] = Field(default=1, description="The last known state of this worker")
-    user: Optional[str] = Field(description="The username to be used when authenticating with this worker")
-    password: Optional[str] = Field(description="The password to be used when authenticating with this worker")
+    user: Optional[str] = Field(description="The username to be used when authenticating with this worker", default=None)
+    password: Optional[str] = Field(description="The password to be used when authenticating with this worker", default=None)
     pixel_cap: Optional[int] = Field(default=-1, description="Max amount of pixels to allow one worker to handle at the same time. -1 means there is no limit")
 
 class ConfigModel(BaseModel):
     workers: List[Dict[str, Worker_Model]]
-    benchmark_payload: Dict = Field(
+    benchmark_payload: Benchmark_Payload = Field(
         default=Benchmark_Payload,
         description='the payload used when benchmarking a node'
     )
diff --git a/scripts/spartan/shared.py b/scripts/spartan/shared.py
index 5400464..ead9f65 100644
--- a/scripts/spartan/shared.py
+++ b/scripts/spartan/shared.py
@@ -65,7 +65,7 @@ samples = 3  # number of times to benchmark worker after warmup benchmarks are c
 
 
 class BenchmarkPayload(BaseModel):
-    validate_assignment = True
+    # validate_assignment = True
     prompt: str = Field(default="A herd of cows grazing at the bottom of a sunny valley")
     negative_prompt: str = Field(default="")
     steps: int = Field(default=20)
diff --git a/scripts/spartan/world.py b/scripts/spartan/world.py
index 9ff7ceb..3b432b9 100644
--- a/scripts/spartan/world.py
+++ b/scripts/spartan/world.py
@@ -674,7 +674,7 @@ class World:
 
             self.add_worker(**fields)
 
-        sh.benchmark_payload = Benchmark_Payload(**config.benchmark_payload)
+        sh.benchmark_payload = Benchmark_Payload(**config.benchmark_payload.dict())
         self.job_timeout = config.job_timeout
         self.enabled = config.enabled
         self.enabled_i2i = config.enabled_i2i
@@ -699,7 +699,7 @@ class World:
         )
 
         with open(self.config_path, 'w+') as config_file:
-            config_file.write(config.json(indent=3))
+            config_file.write(config.model_dump_json(indent=3))
             logger.debug(f"config saved")
 
     def ping_remotes(self, indiscriminate: bool = False):

From 64bc137d65bbb2f9bdfc381c5065be861260ef2a Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Fri, 25 Oct 2024 21:55:49 -0500
Subject: [PATCH 19/20] update thin-client mode

---
 scripts/distributed.py    | 13 +++----------
 scripts/spartan/worker.py |  1 +
 scripts/spartan/world.py  | 36 +++++++++++++++++++++++++++++-------
 3 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/scripts/distributed.py b/scripts/distributed.py
index b3c10f6..9ff6f7b 100644
--- a/scripts/distributed.py
+++ b/scripts/distributed.py
@@ -116,19 +116,14 @@ class DistributedScript(scripts.Script):
         p.negative_prompts.extend(all_negative_prompts)
         p.prompts.extend(all_prompts)
 
-        num_local = self.world.p.n_iter * self.world.p.batch_size + opts.return_grid
+        num_local = self.world.p.n_iter * self.world.p.batch_size + (opts.return_grid - self.world.thin_client_mode)
         num_injected = len(pp.images) - self.world.p.batch_size
         for i, image in enumerate(images):
             # modules.ui_common -> update_generation_info renders to html below gallery
-
-            # TODO probably shouldn't be here
-            if self.world.thin_client_mode:
-                p.all_negative_prompts = pp.all_negative_prompts
-
             gallery_index = num_local + num_injected + i # zero-indexed point of image in total gallery
             job.gallery_map.append(gallery_index) # so we know where to edit infotext
             pp.images.append(image)
-            logger.debug(f"image {gallery_index + 1}/{self.world.num_gallery()}")
+            logger.debug(f"image {gallery_index + 1 + self.world.thin_client_mode}/{self.world.num_gallery()}")
 
     def update_gallery(self, pp, p):
         """adds all remotely generated images to the image gallery after waiting for all workers to finish"""
@@ -328,13 +323,11 @@ class DistributedScript(scripts.Script):
             p.batch_size = self.world.master_job().batch_size
         self.master_start = time.time()
 
-        # generate images assigned to local machine
-        # p.do_not_save_grid = True  # don't generate grid from master as we are doing this later.
         self.runs_since_init += 1
         return
 
     def postprocess_batch_list(self, p, pp, *args, **kwargs):
-        if p.n_iter != kwargs['batch_number'] + 1: # skip if not the final batch
+        if not self.world.thin_client_mode and p.n_iter != kwargs['batch_number'] + 1: # skip if not the final batch
             return
 
         is_img2img = getattr(p, 'init_images', False)
diff --git a/scripts/spartan/worker.py b/scripts/spartan/worker.py
index 28616b9..7bff6f3 100644
--- a/scripts/spartan/worker.py
+++ b/scripts/spartan/worker.py
@@ -609,6 +609,7 @@ class Worker:
                 self.full_url("memory"),
                 timeout=3
             )
+            self.response = response
             return response.status_code == 200
 
         except requests.exceptions.ConnectionError as e:
diff --git a/scripts/spartan/world.py b/scripts/spartan/world.py
index 3b432b9..92f153b 100644
--- a/scripts/spartan/world.py
+++ b/scripts/spartan/world.py
@@ -21,6 +21,9 @@ from .worker import Worker, State
 from modules.call_queue import wrap_queued_call, queue_lock
 from modules import processing
 from modules import progress
+from modules.scripts import PostprocessBatchListArgs
+from torchvision.transforms import ToPILImage
+from modules.images import image_grid
 
 
 class NotBenchmarked(Exception):
@@ -562,15 +565,31 @@ class World:
             # save original process_images_inner for later so we can restore once we're done
             logger.debug(f"bypassing local generation completely")
             def process_images_inner_bypass(p) -> processing.Processed:
+                p.seeds, p.subseeds, p.negative_prompts, p.prompts = [], [], [], []
+                pp = PostprocessBatchListArgs(images=[])
+                self.p.scripts.postprocess_batch_list(p, pp)
+
                 processed = processing.Processed(p, [], p.seed, info="")
-                processed.all_prompts = []
-                processed.all_seeds = []
-                processed.all_subseeds = []
-                processed.all_negative_prompts = []
-                processed.infotexts = []
-                processed.prompt = None
+                processed.all_prompts = p.prompts
+                processed.all_seeds = p.seeds
+                processed.all_subseeds = p.subseeds
+                processed.all_negative_prompts = p.negative_prompts
+                processed.images = pp.images
+                processed.infotexts = [''] * self.num_requested()
+
+                transform = ToPILImage()
+                for i, image in enumerate(processed.images):
+                    processed.images[i] = transform(image)
+
 
                 self.p.scripts.postprocess(p, processed)
+
+                # generate grid if enabled
+                if shared.opts.return_grid and len(processed.images) > 1:
+                    grid = image_grid(processed.images, len(processed.images))
+                    processed.images.insert(0, grid)
+                    processed.infotexts.insert(0, processed.infotexts[0])
+
                 return processed
             processing.process_images_inner = process_images_inner_bypass
 
@@ -699,7 +718,7 @@ class World:
         )
 
         with open(self.config_path, 'w+') as config_file:
-            config_file.write(config.model_dump_json(indent=3))
+            config_file.write(config.json(indent=3))
             logger.debug(f"config saved")
 
     def ping_remotes(self, indiscriminate: bool = False):
@@ -749,6 +768,9 @@ class World:
                     worker.set_state(State.IDLE, expect_cycle=True)
                 else:
                     msg = f"worker '{worker.label}' is unreachable"
+                    if worker.response.status_code is not None:
+                        msg += f" <{worker.response.status_code}>"
+
                     logger.info(msg)
                     gradio.Warning("Distributed: "+msg)
                     worker.set_state(State.UNAVAILABLE)

From 6179dee7ed741ccaff2a878cb77217127957f39f Mon Sep 17 00:00:00 2001
From: papuSpartan <mcgpapu@gmail.com>
Date: Sat, 26 Oct 2024 21:45:00 -0500
Subject: [PATCH 20/20] update changelog

---
 CHANGELOG.md | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 83f3d94..82459e8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,10 +1,23 @@
 # Change Log
 Formatting: [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), [Semantic Versioning](https://semver.org/spec/v2.0.0.html)
 
+## [2.3.0] - 2024-10-26
+
+## Added
+- Compatibility for some extensions which mostly only do postprocessing (e.g. Adetailer)
+- Separate toggle state for img2img tab so txt2img can be enabled and t2i disabled or vice versa
+
+## Changed
+- Status tab will now automatically refresh
+- Main toggle is now in the form of an InputAccordion
+
+## Fixed
+- An issue affecting controlnet and inpainting
+- Toggle state sometimes desyncing when the page was refreshed
+
 ## [2.2.2] - 2024-8-30
 
 ### Fixed
-
 - Unavailable state sometimes being ignored
 
 ## [2.2.1] - 2024-5-16
@@ -82,4 +95,4 @@ Formatting: [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), [Semantic
 - Worker randomly disconnecting when under high load due to handling a previous request
 
 ### Removed
-- Certain superfluous warnings in logs related to third party extensions
\ No newline at end of file
+- Certain superfluous warnings in logs related to third party extensions