Merge remote-tracking branch 'origin/master'

3 files changed, 25 insertions, 16 deletions

diff --git a/modules/processing.py b/modules/processing.py
index 8223423a..0a4b6198 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -79,7 +79,7 @@ class StableDiffusionProcessing:
         self.paste_to = None

         self.color_corrections = None

         self.denoising_strength: float = 0

-

+        self.sampler_noise_scheduler_override = None

         self.ddim_discretize = opts.ddim_discretize

         self.s_churn = opts.s_churn

         self.s_tmin = opts.s_tmin

@@ -130,7 +130,7 @@ class Processed:
         self.s_tmin = p.s_tmin

         self.s_tmax = p.s_tmax

         self.s_noise = p.s_noise

-        

+        self.sampler_noise_scheduler_override = p.sampler_noise_scheduler_override

         self.prompt = self.prompt if type(self.prompt) != list else self.prompt[0]

         self.negative_prompt = self.negative_prompt if type(self.negative_prompt) != list else self.negative_prompt[0]

         self.seed = int(self.seed if type(self.seed) != list else self.seed[0])

diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index dff89c09..92522214 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -290,7 +290,10 @@ class KDiffusionSampler:
     def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning, steps=None):

         steps, t_enc = setup_img2img_steps(p, steps)

 

-        sigmas = self.model_wrap.get_sigmas(steps)

+        if p.sampler_noise_scheduler_override:

+          sigmas = p.sampler_noise_scheduler_override(steps)

+        else:

+          sigmas = self.model_wrap.get_sigmas(steps)

 

         noise = noise * sigmas[steps - t_enc - 1]

         xi = x + noise

@@ -306,7 +309,10 @@ class KDiffusionSampler:
     def sample(self, p, x, conditioning, unconditional_conditioning, steps=None):

         steps = steps or p.steps

 

-        sigmas = self.model_wrap.get_sigmas(steps)

+        if p.sampler_noise_scheduler_override:

+          sigmas = p.sampler_noise_scheduler_override(steps)

+        else:

+          sigmas = self.model_wrap.get_sigmas(steps)

         x = x * sigmas[0]

 

         extra_params_kwargs = self.initialize(p)

diff --git a/modules/swinir_model.py b/modules/swinir_model.py
index 41fda5a7..9bd454c6 100644
--- a/modules/swinir_model.py
+++ b/modules/swinir_model.py
@@ -5,6 +5,7 @@ import numpy as np
 import torch
 from PIL import Image
 from basicsr.utils.download_util import load_file_from_url
+from tqdm import tqdm
 
 from modules import modelloader
 from modules.paths import models_path
@@ -122,18 +123,20 @@ def inference(img, model, tile, tile_overlap, window_size, scale):
     E = torch.zeros(b, c, h * sf, w * sf, dtype=torch.half, device=device).type_as(img)
     W = torch.zeros_like(E, dtype=torch.half, device=device)
 
-    for h_idx in h_idx_list:
-        for w_idx in w_idx_list:
-            in_patch = img[..., h_idx: h_idx + tile, w_idx: w_idx + tile]
-            out_patch = model(in_patch)
-            out_patch_mask = torch.ones_like(out_patch)
-
-            E[
-            ..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
-            ].add_(out_patch)
-            W[
-            ..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
-            ].add_(out_patch_mask)
+    with tqdm(total=len(h_idx_list) * len(w_idx_list), desc="SwinIR tiles") as pbar:
+        for h_idx in h_idx_list:
+            for w_idx in w_idx_list:
+                in_patch = img[..., h_idx: h_idx + tile, w_idx: w_idx + tile]
+                out_patch = model(in_patch)
+                out_patch_mask = torch.ones_like(out_patch)
+
+                E[
+                ..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
+                ].add_(out_patch)
+                W[
+                ..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
+                ].add_(out_patch_mask)
+                pbar.update(1)
     output = E.div_(W)
 
     return output