Merge pull request #158 from JohannesGaessler/progress-printing

More informative progress printing

5 files changed, 43 insertions, 2 deletions

diff --git a/modules/img2img.py b/modules/img2img.py
index a4acdb22..0c91ef3f 100644
--- a/modules/img2img.py
+++ b/modules/img2img.py
@@ -55,6 +55,7 @@ def img2img(prompt: str, init_img, init_img_with_mask, steps: int, sampler_index
             "Denoising strength change factor": denoising_strength_change_factor

         }

     )

+    print(f"\nimg2img: {prompt}", file=shared.progress_print_out)

 

     if is_loopback:

         output_images, info = None, None

@@ -171,5 +172,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 89a0618d..74591bab 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 140b5dea..23022206 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 85318d7e..07b288c2 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 fd81ff0f..410a7a7b 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)