Add button to skip the current iteration

4 files changed, 21 insertions, 0 deletions

diff --git a/modules/img2img.py b/modules/img2img.py
index da212d72..e60b7e0f 100644
--- a/modules/img2img.py
+++ b/modules/img2img.py
@@ -32,6 +32,10 @@ def process_batch(p, input_dir, output_dir, args):
 

     for i, image in enumerate(images):

         state.job = f"{i+1} out of {len(images)}"

+        if state.skipped:

+            state.skipped = False

+            state.interrupted = False

+            continue

 

         if state.interrupted:

             break

diff --git a/modules/processing.py b/modules/processing.py
index d814d5ac..6805039c 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -355,6 +355,10 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
             state.job_count = p.n_iter

 

         for n in range(p.n_iter):

+            if state.skipped:

+                state.skipped = False

+                state.interrupted = False

+            

             if state.interrupted:

                 break

 

diff --git a/modules/shared.py b/modules/shared.py
index 864e772c..7f802bd9 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -84,6 +84,7 @@ def selected_hypernetwork():
 

 

 class State:

+    skipped = False

     interrupted = False

     job = ""

     job_no = 0

@@ -96,6 +97,10 @@ class State:
     current_image_sampling_step = 0

     textinfo = None

 

+    def skip(self):

+        self.skipped = True

+        self.interrupted = True

+

     def interrupt(self):

         self.interrupted = True

 

diff --git a/modules/ui.py b/modules/ui.py
index 4f18126f..e3e62fdd 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -191,6 +191,7 @@ def wrap_gradio_call(func, extra_outputs=None):
         # last item is always HTML

         res[-1] += f"<div class='performance'><p class='time'>Time taken: <wbr>{elapsed_text}</p>{vram_html}</div>"

 

+        shared.state.skipped = False

         shared.state.interrupted = False

         shared.state.job_count = 0

 

@@ -411,9 +412,16 @@ def create_toprow(is_img2img):
 

         with gr.Column(scale=1):

             with gr.Row():

+                skip = gr.Button('Skip', elem_id=f"{id_part}_skip")

                 interrupt = gr.Button('Interrupt', elem_id=f"{id_part}_interrupt")

                 submit = gr.Button('Generate', elem_id=f"{id_part}_generate", variant='primary')

 

+                skip.click(

+                    fn=lambda: shared.state.skip(),

+                    inputs=[],

+                    outputs=[],

+                )

+

                 interrupt.click(

                     fn=lambda: shared.state.interrupt(),

                     inputs=[],