From f211c498b9d496222aa36360193e0f690053dfa4 Mon Sep 17 00:00:00 2001 From: JohannesGaessler Date: Thu, 8 Sep 2022 15:37:13 +0200 Subject: [PATCH] More informative progress printing --- modules/img2img.py | 2 ++ modules/scripts.py | 3 +++ modules/sd_samplers.py | 6 ++++-- modules/shared.py | 31 +++++++++++++++++++++++++++++++ modules/txt2img.py | 3 +++ 5 files changed, 43 insertions(+), 2 deletions(-) diff --git a/modules/img2img.py b/modules/img2img.py index c23923050..b5f3def3b 100644 --- a/modules/img2img.py +++ b/modules/img2img.py @@ -52,6 +52,7 @@ def img2img(prompt: str, init_img, init_img_with_mask, steps: int, sampler_index inpainting_mask_invert=inpainting_mask_invert, extra_generation_params={"Denoising Strength": denoising_strength} ) + print(f"\nimg2img: {prompt}", file=shared.progress_print_out) if is_loopback: output_images, info = None, None @@ -168,5 +169,6 @@ def img2img(prompt: str, init_img, init_img_with_mask, steps: int, sampler_index if processed is None: processed = process_images(p) + shared.total_tqdm.clear() return processed.images, processed.js(), plaintext_to_html(processed.info) diff --git a/modules/scripts.py b/modules/scripts.py index 89a0618de..74591baba 100644 --- a/modules/scripts.py +++ b/modules/scripts.py @@ -6,6 +6,7 @@ import modules.ui as ui import gradio as gr from modules.processing import StableDiffusionProcessing +from modules import shared class Script: filename = None @@ -137,6 +138,8 @@ class ScriptRunner: script_args = args[script.args_from:script.args_to] processed = script.run(p, *script_args) + shared.total_tqdm.clear() + return processed diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index 140b5dead..230222068 100644 --- a/modules/sd_samplers.py +++ b/modules/sd_samplers.py @@ -70,13 +70,14 @@ def extended_tdqm(sequence, *args, desc=None, **kwargs): state.sampling_steps = len(sequence) state.sampling_step = 0 - for x in tqdm.tqdm(sequence, *args, desc=state.job, **kwargs): + for x in tqdm.tqdm(sequence, *args, desc=state.job, file=shared.progress_print_out, **kwargs): if state.interrupted: break yield x state.sampling_step += 1 + shared.total_tqdm.update() ldm.models.diffusion.ddim.tqdm = lambda *args, desc=None, **kwargs: extended_tdqm(*args, desc=desc, **kwargs) @@ -146,13 +147,14 @@ def extended_trange(count, *args, **kwargs): state.sampling_steps = count state.sampling_step = 0 - for x in tqdm.trange(count, *args, desc=state.job, **kwargs): + for x in tqdm.trange(count, *args, desc=state.job, file=shared.progress_print_out, **kwargs): if state.interrupted: break yield x state.sampling_step += 1 + shared.total_tqdm.update() class KDiffusionSampler: diff --git a/modules/shared.py b/modules/shared.py index 85318d7ec..07b288c20 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -1,9 +1,11 @@ +import sys import argparse import json import os import gradio as gr import torch +import tqdm import modules.artists from modules.paths import script_path, sd_path @@ -124,6 +126,7 @@ class Options: "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_progress_every_n_steps": OptionInfo(0, "Show show image creation progress every N sampling steps. Set 0 to disable.", gr.Slider, {"minimum": 0, "maximum": 32, "step": 1}), + "multiple_tqdm": OptionInfo(True, "Add a second progress bar to the console that shows progress for an entire job. Broken in PyCharm console."), "face_restoration_model": OptionInfo(None, "Face restoration model", gr.Radio, lambda: {"choices": [x.name() for x in face_restorers]}), "code_former_weight": OptionInfo(0.5, "CodeFormer weight parameter; 0 = maximum effect; 1 = minimum effect", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}), } @@ -165,4 +168,32 @@ sd_upscalers = [] sd_model = None +progress_print_out = sys.stdout + +class TotalTQDM: + def __init__(self): + self._tqdm = None + + def reset(self): + self._tqdm = tqdm.tqdm( + desc="Total progress", + total=state.job_count * state.sampling_steps, + position=1, + file=progress_print_out + ) + + def update(self): + if not opts.multiple_tqdm: + return + if self._tqdm is None: + self.reset() + self._tqdm.update() + + def clear(self): + if self._tqdm is not None: + self._tqdm.close() + self._tqdm = None + + +total_tqdm = TotalTQDM() diff --git a/modules/txt2img.py b/modules/txt2img.py index fd81ff0f9..410a7a7be 100644 --- a/modules/txt2img.py +++ b/modules/txt2img.py @@ -25,6 +25,7 @@ def txt2img(prompt: str, negative_prompt: str, steps: int, sampler_index: int, r tiling=tiling, ) + print(f"\ntxt2img: {prompt}", file=shared.progress_print_out) processed = modules.scripts.scripts_txt2img.run(p, *args) if processed is not None: @@ -32,5 +33,7 @@ def txt2img(prompt: str, negative_prompt: str, steps: int, sampler_index: int, r else: processed = process_images(p) + shared.total_tqdm.clear() + return processed.images, processed.js(), plaintext_to_html(processed.info)