feat: refactor

2 files changed, 13 insertions, 7 deletions

diff --git a/modules/shared_state.py b/modules/shared_state.py
index c72c3f63..532fdcd8 100644
--- a/modules/shared_state.py
+++ b/modules/shared_state.py
@@ -77,12 +77,12 @@ class State:
         log.info("Received skip request")

 

     def interrupt(self):

-        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")

+        self.interrupted = True

+        log.info("Received interrupt request")

+

+    def interrupt_next(self):

+        self.interrupted_next = True

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

 

     def nextjob(self):

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

diff --git a/modules/ui.py b/modules/ui.py
index bcf39199..c30093d7 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -216,8 +216,14 @@ class Toprow:
                         outputs=[],

                     )

 

+                    def interrupt_fn():

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

+                            shared.state.interrupt_next()

+                        else:

+                            shared.state.interrupt()

+

                     self.interrupt.click(

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

+                        fn=interrupt_fn,

                         inputs=[],

                         outputs=[],

                     )