From 149784202cca8612b43629c601ee27cfda64e623 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sun, 30 Oct 2022 09:10:22 +0300
Subject: [PATCH] rework #3722 to not introduce duplicate code

---
 modules/api/api.py | 43 +++++++++++++------------------------------
 modules/shared.py  | 22 +++++++++++++++++++---
 webui.py           | 17 ++---------------
 3 files changed, 34 insertions(+), 48 deletions(-)

diff --git a/modules/api/api.py b/modules/api/api.py
index 5c5b210f5..6c06d4499 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -9,31 +9,6 @@ from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusion
 from modules.sd_samplers import all_samplers
 from modules.extras import run_extras, run_pnginfo
 
-# copy from wrap_gradio_gpu_call of webui.py
-# because queue lock will be acquired in api handlers
-# and time start needs to be set
-# the function has been modified into two parts
-
-def before_gpu_call():
-    devices.torch_gc()
-
-    shared.state.sampling_step = 0
-    shared.state.job_count = -1
-    shared.state.job_no = 0
-    shared.state.job_timestamp = shared.state.get_job_timestamp()
-    shared.state.current_latent = None
-    shared.state.current_image = None
-    shared.state.current_image_sampling_step = 0
-    shared.state.skipped = False
-    shared.state.interrupted = False
-    shared.state.textinfo = None
-    shared.state.time_start = time.time()
-
-def after_gpu_call():
-    shared.state.job = ""
-    shared.state.job_count = 0
-
-    devices.torch_gc()
 
 def upscaler_to_index(name: str):
     try:
@@ -41,8 +16,10 @@ def upscaler_to_index(name: str):
     except:
         raise HTTPException(status_code=400, detail=f"Invalid upscaler, needs to be on of these: {' , '.join([x.name for x in sd_upscalers])}")
 
+
 sampler_to_index = lambda name: next(filter(lambda row: name.lower() == row[1].name.lower(), enumerate(all_samplers)), None)
 
+
 def setUpscalers(req: dict):
     reqDict = vars(req)
     reqDict['extras_upscaler_1'] = upscaler_to_index(req.upscaler_1)
@@ -51,6 +28,7 @@ def setUpscalers(req: dict):
     reqDict.pop('upscaler_2')
     return reqDict
 
+
 class Api:
     def __init__(self, app, queue_lock):
         self.router = APIRouter()
@@ -78,10 +56,13 @@ class Api:
         )
         p = StableDiffusionProcessingTxt2Img(**vars(populate))
         # Override object param
-        before_gpu_call()
+
+        shared.state.begin()
+
         with self.queue_lock:
             processed = process_images(p)
-        after_gpu_call()
+
+        shared.state.end()
 
         b64images = list(map(encode_pil_to_base64, processed.images))
 
@@ -119,11 +100,13 @@ class Api:
             imgs = [img] * p.batch_size
 
         p.init_images = imgs
-        # Override object param
-        before_gpu_call()
+
+        shared.state.begin()
+
         with self.queue_lock:
             processed = process_images(p)
-        after_gpu_call()
+
+        shared.state.end()
 
         b64images = list(map(encode_pil_to_base64, processed.images))
 
diff --git a/modules/shared.py b/modules/shared.py
index f7b0990c5..e4f163c11 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -144,9 +144,6 @@ class State:
         self.sampling_step = 0
         self.current_image_sampling_step = 0
 
-    def get_job_timestamp(self):
-        return datetime.datetime.now().strftime("%Y%m%d%H%M%S")  # shouldn't this return job_timestamp?
-
     def dict(self):
         obj = {
             "skipped": self.skipped,
@@ -160,6 +157,25 @@ class State:
 
         return obj
 
+    def begin(self):
+        self.sampling_step = 0
+        self.job_count = -1
+        self.job_no = 0
+        self.job_timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+        self.current_latent = None
+        self.current_image = None
+        self.current_image_sampling_step = 0
+        self.skipped = False
+        self.interrupted = False
+        self.textinfo = None
+
+        devices.torch_gc()
+
+    def end(self):
+        self.job = ""
+        self.job_count = 0
+
+        devices.torch_gc()
 
 state = State()
 
diff --git a/webui.py b/webui.py
index ade7334bf..29530872c 100644
--- a/webui.py
+++ b/webui.py
@@ -46,26 +46,13 @@ def wrap_queued_call(func):
 
 def wrap_gradio_gpu_call(func, extra_outputs=None):
     def f(*args, **kwargs):
-        devices.torch_gc()
 
-        shared.state.sampling_step = 0
-        shared.state.job_count = -1
-        shared.state.job_no = 0
-        shared.state.job_timestamp = shared.state.get_job_timestamp()
-        shared.state.current_latent = None
-        shared.state.current_image = None
-        shared.state.current_image_sampling_step = 0
-        shared.state.skipped = False
-        shared.state.interrupted = False
-        shared.state.textinfo = None
+        shared.state.begin()
 
         with queue_lock:
             res = func(*args, **kwargs)
 
-        shared.state.job = ""
-        shared.state.job_count = 0
-
-        devices.torch_gc()
+        shared.state.end()
 
         return res