From 8ddb7dd7b7fcf2b09bcfd66c2427a459a3011f3d Mon Sep 17 00:00:00 2001
From: Bingsu <ks2515@naver.com>
Date: Fri, 26 May 2023 11:51:46 +0900
Subject: [PATCH] feat: use shared.state

---
 scripts/!adetailer.py |  6 +++++-
 sd_webui/shared.py    | 26 +++++++++++++++++++++++++-
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/scripts/!adetailer.py b/scripts/!adetailer.py
index 3d31171..7cc7216 100644
--- a/scripts/!adetailer.py
+++ b/scripts/!adetailer.py
@@ -39,7 +39,7 @@ from sd_webui.processing import (
     create_infotext,
     process_images,
 )
-from sd_webui.shared import cmd_opts, opts
+from sd_webui.shared import cmd_opts, opts, state
 
 try:
     from rich import print
@@ -451,6 +451,9 @@ class AfterDetailerScript(scripts.Script):
 
             `True` if image was processed, `False` otherwise.
         """
+        if state.interrupted:
+            return False
+
         i = p._idx
 
         i2i = self.get_i2i_p(p, args, pp.image)
@@ -488,6 +491,7 @@ class AfterDetailerScript(scripts.Script):
 
         steps = len(masks)
         processed = None
+        state.job_count += steps
 
         if is_mediapipe:
             print(f"mediapipe: {steps} detected.")
diff --git a/sd_webui/shared.py b/sd_webui/shared.py
index e71260a..18b0cd0 100644
--- a/sd_webui/shared.py
+++ b/sd_webui/shared.py
@@ -7,6 +7,29 @@ if TYPE_CHECKING:
     from dataclasses import dataclass
     from typing import Any, Callable
 
+    import torch
+    from PIL import Image
+
+    @dataclass
+    class State:
+        skipped: bool = False
+        interrupted: bool = False
+        job: str = ""
+        job_no: int = 0
+        job_count: int = 0
+        processing_has_refined_job_count: bool = False
+        job_timestamp: str = "0"
+        sampling_step: int = 0
+        sampling_steps: int = 0
+        current_latent: torch.Tensor | None = None
+        current_image: Image.Image | None = None
+        current_image_sampling_step: int = 0
+        id_live_preview: int = 0
+        textinfo: str | None = None
+        time_start: float | None = None
+        need_restart: bool = False
+        server_start: float | None = None
+
     @dataclass
     class OptionInfo:
         default: Any = None
@@ -37,6 +60,7 @@ if TYPE_CHECKING:
 
     opts = Option()
     cmd_opts = argparse.Namespace()
+    state = State()
 
 else:
-    from modules.shared import OptionInfo, cmd_opts, opts
+    from modules.shared import OptionInfo, cmd_opts, opts, state