Merge pull request #5976 from AbstractQbit/fast_preview

Add an option for faster low quality previews

2 files changed, 19 insertions, 9 deletions

diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index 8efe74df..27ef4ff8 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -106,20 +106,29 @@ def setup_img2img_steps(p, steps=None):
     return steps, t_enc

 

 

-def single_sample_to_image(sample):

-    x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0]

+def single_sample_to_image(sample, approximation=False):

+    if approximation:

+        # https://discuss.huggingface.co/t/decoding-latents-to-rgb-without-upscaling/23204/2

+        coefs = torch.tensor(

+            [[ 0.298,  0.207,  0.208],

+             [ 0.187,  0.286,  0.173],

+             [-0.158,  0.189,  0.264],

+             [-0.184, -0.271, -0.473]]).to(sample.device)

+        x_sample = torch.einsum("lxy,lr -> rxy", sample, coefs)

+    else:

+        x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[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 sample_to_image(samples, index=0):

-    return single_sample_to_image(samples[index])

+def sample_to_image(samples, index=0, approximation=False):

+    return single_sample_to_image(samples[index], approximation)

 

 

-def samples_to_image_grid(samples):

-    return images.image_grid([single_sample_to_image(sample) for sample in samples])

+def samples_to_image_grid(samples, approximation=False):

+    return images.image_grid([single_sample_to_image(sample, approximation) for sample in samples])

 

 

 def store_latent(decoded):

@@ -127,7 +136,7 @@ def store_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)

+            shared.state.current_image = sample_to_image(decoded, approximation=opts.show_progress_approximate)

 

 

 class InterruptedException(BaseException):

diff --git a/modules/shared.py b/modules/shared.py
index a75de535..eb3e5aec 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -212,9 +212,9 @@ class State:
 

         import modules.sd_samplers

         if opts.show_progress_grid:

-            self.current_image = modules.sd_samplers.samples_to_image_grid(self.current_latent)

+            self.current_image = modules.sd_samplers.samples_to_image_grid(self.current_latent, approximation=opts.show_progress_approximate)

         else:

-            self.current_image = modules.sd_samplers.sample_to_image(self.current_latent)

+            self.current_image = modules.sd_samplers.sample_to_image(self.current_latent, approximation=opts.show_progress_approximate)

 

         self.current_image_sampling_step = self.sampling_step

 

@@ -392,6 +392,7 @@ options_templates.update(options_section(('interrogate', "Interrogate Options"),
 options_templates.update(options_section(('ui', "User interface"), {

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

     "show_progress_every_n_steps": OptionInfo(0, "Show image creation progress every N sampling steps. Set to 0 to disable. Set to -1 to show after completion of batch.", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}),

+    "show_progress_approximate": OptionInfo(False, "Calculate small previews using fast linear approximation instead of VAE"),

     "show_progress_grid": OptionInfo(True, "Show previews of all images generated in a batch as a grid"),

     "return_grid": OptionInfo(True, "Show grid in results for web"),

     "do_not_show_images": OptionInfo(False, "Do not show any images in results for web"),