fix for progress bar

This commit is contained in:
AUTOMATIC 2022-09-06 20:26:09 +03:00
parent fd66199769
commit 7e5b8becf9
3 changed files with 7 additions and 2 deletions

View File

@ -49,6 +49,7 @@ class State:
sampling_steps = 0 sampling_steps = 0
current_latent = None current_latent = None
current_image = None current_image = None
current_progress_index = 0
def interrupt(self): def interrupt(self):
self.interrupted = True self.interrupted = True
@ -102,7 +103,7 @@ class Options:
"random_artist_categories": OptionInfo([], "Allowed categories for random artists selection when using the Roll button", gr.CheckboxGroup, {"choices": artist_db.categories()}), "random_artist_categories": OptionInfo([], "Allowed categories for random artists selection when using the Roll button", gr.CheckboxGroup, {"choices": artist_db.categories()}),
"upscale_at_full_resolution_padding": OptionInfo(16, "Inpainting at full resolution: padding, in pixels, for the masked region.", gr.Slider, {"minimum": 0, "maximum": 128, "step": 4}), "upscale_at_full_resolution_padding": OptionInfo(16, "Inpainting at full resolution: padding, in pixels, for the masked region.", gr.Slider, {"minimum": 0, "maximum": 128, "step": 4}),
"show_progressbar": OptionInfo(True, "Show progressbar"), "show_progressbar": OptionInfo(True, "Show progressbar"),
"show_progress_every_n_steps": OptionInfo(0, "Show show image creation progress every N steps. Set 0 to disable.", gr.Slider, {"minimum": 0, "maximum": 32, "step": 1}), "show_progress_every_n_steps": OptionInfo(0, "Show show image creation progress every N progress pudates. Set 0 to disable.", gr.Slider, {"minimum": 0, "maximum": 32, "step": 1}),
} }
def __init__(self): def __init__(self):

View File

@ -160,13 +160,14 @@ def check_progress_call():
preview_visibility = gr_show(False) preview_visibility = gr_show(False)
if opts.show_progress_every_n_steps > 0: if opts.show_progress_every_n_steps > 0:
if (shared.state.sampling_step-1) % opts.show_progress_every_n_steps == 0 and shared.state.current_latent is not None: if shared.state.current_progress_index % opts.show_progress_every_n_steps == 0 and shared.state.current_latent is not None:
x_sample = shared.sd_model.decode_first_stage(shared.state.current_latent[0:1].type(shared.sd_model.dtype))[0] x_sample = shared.sd_model.decode_first_stage(shared.state.current_latent[0:1].type(shared.sd_model.dtype))[0]
x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0) x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)
x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2) x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
x_sample = x_sample.astype(np.uint8) x_sample = x_sample.astype(np.uint8)
shared.state.current_image = Image.fromarray(x_sample) shared.state.current_image = Image.fromarray(x_sample)
image = shared.state.current_image image = shared.state.current_image
if image is None or progress >= 1: if image is None or progress >= 1:
@ -174,6 +175,8 @@ def check_progress_call():
else: else:
preview_visibility = gr_show(True) preview_visibility = gr_show(True)
shared.state.current_progress_index += 1
return f"<span style='display: none'>{time.time()}</span><p>{progressbar}</p>", preview_visibility, image return f"<span style='display: none'>{time.time()}</span><p>{progressbar}</p>", preview_visibility, image

View File

@ -127,6 +127,7 @@ def wrap_gradio_gpu_call(func):
shared.state.job_no = 0 shared.state.job_no = 0
shared.state.current_latent = None shared.state.current_latent = None
shared.state.current_image = None shared.state.current_image = None
shared.state.current_progress_index = 0
with queue_lock: with queue_lock:
res = func(*args, **kwargs) res = func(*args, **kwargs)