rework #3722 to not introduce duplicate code

1 files changed, 3 insertions, 16 deletions

diff --git a/webui.py b/webui.py
index ade7334b..29530872 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