fix for live progress breaking lowvram and medvram optimizations

5 files changed, 33 insertions, 16 deletions

diff --git a/README.md b/README.md
index e18c7c4d..99c19d3a 100644
--- a/README.md
+++ b/README.md
@@ -33,6 +33,9 @@ A browser interface based on Gradio library for Stable Diffusion.
 - Running custom code from UI

 - Mouseover hints fo most UI elements

 - Possible to change defaults/mix/max/step values for UI elements via text config

+- Random artist button

+- Tiling support: UI checkbox to create images that can be tiled like textures

+- Progress bar and live image generation preview

 

 ## Installing and running

 

diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index ff7e686e..e8bc5be2 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -1,8 +1,8 @@
 from collections import namedtuple

-

-import ldm.models.diffusion.ddim

+import numpy as np

 import torch

 import tqdm

+from PIL import Image

 

 import k_diffusion.sampling

 import ldm.models.diffusion.ddim

@@ -37,12 +37,28 @@ samplers = [
 samplers_for_img2img = [x for x in samplers if x.name != 'PLMS']

 

 

+def sample_to_image(samples):

+    x_sample = shared.sd_model.decode_first_stage(samples[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 = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)

+    x_sample = x_sample.astype(np.uint8)

+    return Image.fromarray(x_sample)

+

+

+def store_latent(decoded):

+    state.current_latent = decoded

+

+    if opts.show_progress_every_n_steps > 0 and shared.state.sampling_step % opts.show_progress_every_n_steps == 0:

+        if not shared.parallel_processing_allowed:

+            shared.state.current_image = sample_to_image(decoded)

+

+

 def p_sample_ddim_hook(sampler_wrapper, x_dec, cond, ts, *args, **kwargs):

     if sampler_wrapper.mask is not None:

         img_orig = sampler_wrapper.sampler.model.q_sample(sampler_wrapper.init_latent, ts)

         x_dec = img_orig * sampler_wrapper.mask + sampler_wrapper.nmask * x_dec

 

-    state.current_latent = x_dec

+    store_latent(x_dec)

 

     return sampler_wrapper.orig_p_sample_ddim(x_dec, cond, ts, *args, **kwargs)

 

@@ -144,7 +160,7 @@ class KDiffusionSampler:
         self.model_wrap_cfg = CFGDenoiser(self.model_wrap)

 

     def callback_state(self, d):

-        state.current_latent = d["denoised"]

+        store_latent(d["denoised"])

 

     def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning):

         t_enc = int(min(p.denoising_strength, 0.999) * p.steps)

diff --git a/modules/shared.py b/modules/shared.py
index 5eac3317..49f17145 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -38,7 +38,7 @@ cpu = torch.device("cpu")
 gpu = torch.device("cuda")

 device = gpu if torch.cuda.is_available() else cpu

 batch_cond_uncond = cmd_opts.always_batch_cond_uncond or not (cmd_opts.lowvram or cmd_opts.medvram)

-

+parallel_processing_allowed = not cmd_opts.lowvram and not cmd_opts.medvram

 

 class State:

     interrupted = False

@@ -49,7 +49,8 @@ class State:
     sampling_steps = 0

     current_latent = None

     current_image = None

-    current_progress_index = 0

+    current_image_sampling_step = 0

+

 

     def interrupt(self):

         self.interrupted = True

@@ -57,6 +58,7 @@ class State:
     def nextjob(self):

         self.job_no += 1

         self.sampling_step = 0

+        self.current_image_sampling_step = 0

 

 

 state = State()

@@ -103,7 +105,7 @@ class Options:
         "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}),

         "show_progressbar": OptionInfo(True, "Show progressbar"),

-        "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}),

+        "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}),

     }

 

     def __init__(self):

diff --git a/modules/ui.py b/modules/ui.py
index fb3c4d33..92d8bcdd 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -160,13 +160,11 @@ def check_progress_call():
     preview_visibility = gr_show(False)

 

     if opts.show_progress_every_n_steps > 0:

-        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 = 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 = x_sample.astype(np.uint8)

-            shared.state.current_image = Image.fromarray(x_sample)

+        if shared.parallel_processing_allowed:

 

+            if shared.state.sampling_step - shared.state.current_image_sampling_step >= opts.show_progress_every_n_steps and shared.state.current_latent is not None:

+                shared.state.current_image = modules.sd_samplers.sample_to_image(shared.state.current_latent)

+                shared.state.current_image_sampling_step = shared.state.sampling_step

 

         image = shared.state.current_image

 

@@ -175,8 +173,6 @@ def check_progress_call():
         else:

             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

 

 

diff --git a/webui.py b/webui.py
index b6037365..d20ff38f 100644
--- a/webui.py
+++ b/webui.py
@@ -127,7 +127,7 @@ def wrap_gradio_gpu_call(func):
         shared.state.job_no = 0

         shared.state.current_latent = None

         shared.state.current_image = None

-        shared.state.current_progress_index = 0

+        shared.state.current_image_sampling_step = 0

 

         with queue_lock:

             res = func(*args, **kwargs)