added preview option

4 files changed, 67 insertions, 10 deletions

diff --git a/modules/processing.py b/modules/processing.py
index e8923a7a..e615ffdc 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -176,6 +176,11 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
                 shared.state.job = f"Batch {n+1} out of {p.n_iter}"

 

             samples_ddim = p.sample(x=x, conditioning=c, unconditional_conditioning=uc)

+            if state.interrupted:

+

+                # if we are interruped, sample returns just noise

+                # use the image collected previously in sampler loop

+                samples_ddim = shared.state.current_latent

 

             x_samples_ddim = p.sd_model.decode_first_stage(samples_ddim)

             x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)

diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index 896e8b3f..ff7e686e 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -42,6 +42,8 @@ def p_sample_ddim_hook(sampler_wrapper, x_dec, cond, ts, *args, **kwargs):
         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

+

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

 

 

@@ -141,6 +143,9 @@ class KDiffusionSampler:
         self.func = getattr(k_diffusion.sampling, self.funcname)

         self.model_wrap_cfg = CFGDenoiser(self.model_wrap)

 

+    def callback_state(self, d):

+        state.current_latent = d["denoised"]

+

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

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

         sigmas = self.model_wrap.get_sigmas(p.steps)

@@ -157,7 +162,7 @@ class KDiffusionSampler:
         if hasattr(k_diffusion.sampling, 'trange'):

             k_diffusion.sampling.trange = lambda *args, **kwargs: extended_trange(*args, **kwargs)

 

-        return self.func(self.model_wrap_cfg, xi, sigma_sched, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False)

+        return self.func(self.model_wrap_cfg, xi, sigma_sched, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False, callback=self.callback_state)

 

     def sample(self, p, x, conditioning, unconditional_conditioning):

         sigmas = self.model_wrap.get_sigmas(p.steps)

@@ -166,6 +171,6 @@ class KDiffusionSampler:
         if hasattr(k_diffusion.sampling, 'trange'):

             k_diffusion.sampling.trange = lambda *args, **kwargs: extended_trange(*args, **kwargs)

 

-        samples_ddim = self.func(self.model_wrap_cfg, x, sigmas, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False)

+        samples_ddim = self.func(self.model_wrap_cfg, x, sigmas, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False, callback=self.callback_state)

         return samples_ddim

 

diff --git a/modules/shared.py b/modules/shared.py
index d57aba37..e9c88e31 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -39,6 +39,7 @@ 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)

 

+

 class State:

     interrupted = False

     job = ""

@@ -46,6 +47,8 @@ class State:
     job_count = 0

     sampling_step = 0

     sampling_steps = 0

+    current_latent = None

+    current_image = None

 

     def interrupt(self):

         self.interrupted = True

@@ -99,6 +102,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 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 1df74070..8e7a3ee4 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -9,6 +9,8 @@ import sys
 import time

 import traceback

 

+import numpy as np

+import torch

 from PIL import Image

 

 import gradio as gr

@@ -119,6 +121,9 @@ def wrap_gradio_call(func):
             print("Arguments:", args, kwargs, file=sys.stderr)

             print(traceback.format_exc(), file=sys.stderr)

 

+            shared.state.job = ""

+            shared.state.job_count = 0

+

             res = [None, '', f"<div class='error'>{plaintext_to_html(type(e).__name__+': '+str(e))}</div>"]

 

         elapsed = time.perf_counter() - t

@@ -134,11 +139,9 @@ def wrap_gradio_call(func):
 

 

 def check_progress_call():

-    if not opts.show_progressbar:

-        return ""

 

     if shared.state.job_count == 0:

-        return ""

+        return "", gr_show(False), gr_show(False)

 

     progress = 0

 

@@ -149,9 +152,29 @@ def check_progress_call():
 

     progress = min(progress, 1)

 

-    progressbar = f"""<div class='progressDiv'><div class='progress' style="width:{progress * 100}%">{str(int(progress*100))+"%" if progress > 0.01 else ""}</div></div>"""

+    progressbar = ""

+    if opts.show_progressbar:

+        progressbar = f"""<div class='progressDiv'><div class='progress' style="width:{progress * 100}%">{str(int(progress*100))+"%" if progress > 0.01 else ""}</div></div>"""

+

+    image = gr_show(False)

+    preview_visibility = gr_show(False)

+

+    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:

+            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)

 

-    return f"<span style='display: none'>{time.time()}</span><p>{progressbar}</p>"

+        image = shared.state.current_image

+

+        if image is None or progress >= 1:

+            image = gr.update(value=None)

+        else:

+            preview_visibility = gr_show(True)

+

+    return f"<span style='display: none'>{time.time()}</span><p>{progressbar}</p>", preview_visibility, image

 

 

 def roll_artist(prompt):

@@ -204,6 +227,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
 

             with gr.Column(variant='panel'):

                 with gr.Group():

+                    txt2img_preview = gr.Image(elem_id='txt2img_preview', visible=False)

                     txt2img_gallery = gr.Gallery(label='Output', elem_id='txt2img_gallery')

 

 

@@ -251,8 +275,9 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
 

             check_progress.click(

                 fn=check_progress_call,

+                show_progress=False,

                 inputs=[],

-                outputs=[progressbar],

+                outputs=[progressbar, txt2img_preview, txt2img_preview],

             )

 

 

@@ -337,13 +362,16 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
 

             with gr.Column(variant='panel'):

                 with gr.Group():

+                    img2img_preview = gr.Image(elem_id='img2img_preview', visible=False)

                     img2img_gallery = gr.Gallery(label='Output', elem_id='img2img_gallery')

 

                 with gr.Group():

                     with gr.Row():

-                        interrupt = gr.Button('Interrupt')

                         save = gr.Button('Save')

+                        img2img_send_to_img2img = gr.Button('Send to img2img')

+                        img2img_send_to_inpaint = gr.Button('Send to inpaint')

                         img2img_send_to_extras = gr.Button('Send to extras')

+                        interrupt = gr.Button('Interrupt')

 

                 progressbar = gr.HTML(elem_id="progressbar")

 

@@ -426,8 +454,9 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
 

             check_progress.click(

                 fn=check_progress_call,

+                show_progress=False,

                 inputs=[],

-                outputs=[progressbar],

+                outputs=[progressbar, img2img_preview, img2img_preview],

             )

 

             interrupt.click(

@@ -463,6 +492,20 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                 outputs=[init_img_with_mask],

             )

 

+            img2img_send_to_img2img.click(

+                fn=lambda x: image_from_url_text(x),

+                _js="extract_image_from_gallery",

+                inputs=[img2img_gallery],

+                outputs=[init_img],

+            )

+

+            img2img_send_to_inpaint.click(

+                fn=lambda x: image_from_url_text(x),

+                _js="extract_image_from_gallery",

+                inputs=[img2img_gallery],

+                outputs=[init_img_with_mask],

+            )

+

     with gr.Blocks(analytics_enabled=False) as extras_interface:

         with gr.Row().style(equal_height=False):

             with gr.Column(variant='panel'):