From 3d15e58b0a30f2ef1e731f9e429f4d3cf1c259c5 Mon Sep 17 00:00:00 2001
From: Anthony Fu <anthonyfu117@hotmail.com>
Date: Mon, 16 Oct 2023 15:00:17 +0800
Subject: [PATCH] feat: refactor

---
 modules/shared_state.py | 12 ++++++------
 modules/ui.py           |  8 +++++++-
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/modules/shared_state.py b/modules/shared_state.py
index c72c3f63c..532fdcd8d 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 bcf391997..c30093d75 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=[],
                     )