Interrupt after current generation

7 files changed, 17 insertions, 8 deletions

diff --git a/modules/call_queue.py b/modules/call_queue.py
index ddf0d573..01c6d17f 100644
--- a/modules/call_queue.py
+++ b/modules/call_queue.py
@@ -78,6 +78,7 @@ def wrap_gradio_call(func, extra_outputs=None, add_stats=False):
 

         shared.state.skipped = False

         shared.state.interrupted = False

+        shared.state.interrupted_next = False

         shared.state.job_count = 0

 

         if not add_stats:

diff --git a/modules/img2img.py b/modules/img2img.py
index 52cb577a..31f8c2aa 100644
--- a/modules/img2img.py
+++ b/modules/img2img.py
@@ -49,7 +49,7 @@ def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args, to_scale=Fal
         if state.skipped:

             state.skipped = False

 

-        if state.interrupted:

+        if state.interrupted or state.interrupted_next:

             break

 

         try:

diff --git a/modules/processing.py b/modules/processing.py
index 40598f5c..e7eecd66 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -819,7 +819,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             if state.skipped:

                 state.skipped = False

 

-            if state.interrupted:

+            if state.interrupted or state.interrupted_next:

                 break

 

             sd_models.reload_model_weights()  # model can be changed for example by refiner

diff --git a/modules/shared_options.py b/modules/shared_options.py
index 32bf7353..4638ef06 100644
--- a/modules/shared_options.py
+++ b/modules/shared_options.py
@@ -113,6 +113,7 @@ options_templates.update(options_section(('system', "System"), {
     "disable_mmap_load_safetensors": OptionInfo(False, "Disable memmapping for loading .safetensors files.").info("fixes very slow loading speed in some cases"),

     "hide_ldm_prints": OptionInfo(True, "Prevent Stability-AI's ldm/sgm modules from printing noise to console."),

     "dump_stacks_on_signal": OptionInfo(False, "Print stack traces before exiting the program with ctrl+c."),

+    "interrupt_after_current": OptionInfo(False, "Interrupt generation after current image is finished on batch processing"),

 }))

 

 options_templates.update(options_section(('API', "API"), {

diff --git a/modules/shared_state.py b/modules/shared_state.py
index a68789cc..c72c3f63 100644
--- a/modules/shared_state.py
+++ b/modules/shared_state.py
@@ -12,6 +12,7 @@ log = logging.getLogger(__name__)
 class State:

     skipped = False

     interrupted = False

+    interrupted_next = False

     job = ""

     job_no = 0

     job_count = 0

@@ -76,8 +77,12 @@ class State:
         log.info("Received skip request")

 

     def interrupt(self):

-        self.interrupted = True

-        log.info("Received interrupt request")

+        if shared.opts.interrupt_after_current and self.job_count > 1:

+            self.interrupted_next = True

+            log.info("Received interrupt request, interrupt after current job")

+        else:

+            self.interrupted = True

+            log.info("Received interrupt request")

 

     def nextjob(self):

         if shared.opts.live_previews_enable and shared.opts.show_progress_every_n_steps == -1:

@@ -91,6 +96,7 @@ class State:
         obj = {

             "skipped": self.skipped,

             "interrupted": self.interrupted,

+            "interrupted_next": self.interrupted_next,

             "job": self.job,

             "job_count": self.job_count,

             "job_timestamp": self.job_timestamp,

@@ -114,6 +120,7 @@ class State:
         self.id_live_preview = 0

         self.skipped = False

         self.interrupted = False

+        self.interrupted_next = False

         self.textinfo = None

         self.job = job

         devices.torch_gc()

diff --git a/scripts/loopback.py b/scripts/loopback.py
index 2d5feaf9..ad921269 100644
--- a/scripts/loopback.py
+++ b/scripts/loopback.py
@@ -95,7 +95,7 @@ class Script(scripts.Script):
                 processed = processing.process_images(p)

 

                 # Generation cancelled.

-                if state.interrupted:

+                if state.interrupted or state.interrupted_next:

                     break

 

                 if initial_seed is None:

@@ -122,8 +122,8 @@ class Script(scripts.Script):
 

             p.inpainting_fill = original_inpainting_fill

 

-            if state.interrupted:

-                    break

+            if state.interrupted or state.interrupted_next:

+                break

 

         if len(history) > 1:

             grid = images.image_grid(history, rows=1)

diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py
index 0dc255bc..495008ad 100644
--- a/scripts/xyz_grid.py
+++ b/scripts/xyz_grid.py
@@ -688,7 +688,7 @@ class Script(scripts.Script):
         grid_infotext = [None] * (1 + len(zs))

 

         def cell(x, y, z, ix, iy, iz):

-            if shared.state.interrupted:

+            if shared.state.interrupted or state.interrupted_next:

                 return Processed(p, [], p.seed, "")

 

             pc = copy(p)