From 34fe2c5078a2bee6033a202a9a550f60831bb99d Mon Sep 17 00:00:00 2001
From: unknown <mcgpapu@gmail.com>
Date: Wed, 14 Jun 2023 11:01:10 -0500
Subject: [PATCH] support inpaint

---
 scripts/extension.py      | 13 ++++++++++++-
 scripts/spartan/Worker.py |  9 +++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/scripts/extension.py b/scripts/extension.py
index 34deb30..6eb0796 100644
--- a/scripts/extension.py
+++ b/scripts/extension.py
@@ -161,7 +161,14 @@ class Script(scripts.Script):
                     donor_worker = job.worker
             except KeyError:
                 if job.batch_size > 0:
-                    logger.warning(f"Worker '{job.worker.uuid}' had nothing")
+                    logger.warning(f"Worker '{job.worker.uuid}' had no images")
+                continue
+            except TypeError as e:
+                if job.worker.response is None:
+                    logger.error(f"worker '{job.worker.uuid}' had no response")
+                else:
+                    logger.exception(e)
+                continue
 
             injected_to_iteration = 0
             images_per_iteration = Script.world.get_current_output_size()
@@ -179,6 +186,10 @@ class Script(scripts.Script):
                 else:
                     injected_to_iteration += 1
 
+        if donor_worker is None:
+            logger.critical("couldn't collect any responses, distributed will do nothing")
+            return
+
         # generate and inject grid
         if opts.return_grid:
             grid = processing.images.image_grid(processed.images, len(processed.images))
diff --git a/scripts/spartan/Worker.py b/scripts/spartan/Worker.py
index 2648657..83d8475 100644
--- a/scripts/spartan/Worker.py
+++ b/scripts/spartan/Worker.py
@@ -16,6 +16,7 @@ import json
 import base64
 import queue
 from modules.shared import state as master_state
+from modules.api.api import encode_pil_to_base64
 
 
 class InvalidWorkerResponse(Exception):
@@ -322,6 +323,14 @@ class Worker:
                         images.append(image)
                     payload['init_images'] = images
 
+                # 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
+                    del payload['image_mask']
+
                 # see if there is anything else wrong with serializing to payload
                 try:
                     json.dumps(payload)