diff options
| -rw-r--r-- | modules/scripts.py | 35 | ||||
| -rw-r--r-- | modules/sd_samplers_cfg_denoiser.py | 6 | ||||
| -rw-r--r-- | modules/sd_samplers_timesteps.py | 6 | ||||
| -rw-r--r-- | modules/sd_samplers_timesteps_impl.py | 18 |
4 files changed, 43 insertions, 22 deletions
diff --git a/modules/scripts.py b/modules/scripts.py index d4a9da94..cbdac2b5 100644 --- a/modules/scripts.py +++ b/modules/scripts.py @@ -239,6 +239,8 @@ class Script: """
Calls callback before a component is created. The callback function is called with a single argument of type OnComponent.
+ May be called in show() or ui() - but it may be too late in latter as some components may already be created.
+
This function is an alternative to before_component in that it also cllows to run before a component is created, but
it doesn't require to be called for every created component - just for the one you need.
"""
@@ -445,6 +447,28 @@ class ScriptRunner: self.scripts.append(script)
self.selectable_scripts.append(script)
+ self.apply_on_before_component_callbacks()
+
+ def apply_on_before_component_callbacks(self):
+ for script in self.scripts:
+ on_before = script.on_before_component_elem_id or []
+ on_after = script.on_after_component_elem_id or []
+
+ for elem_id, callback in on_before:
+ if elem_id not in self.on_before_component_elem_id:
+ self.on_before_component_elem_id[elem_id] = []
+
+ self.on_before_component_elem_id[elem_id].append((callback, script))
+
+ for elem_id, callback in on_after:
+ if elem_id not in self.on_after_component_elem_id:
+ self.on_after_component_elem_id[elem_id] = []
+
+ self.on_after_component_elem_id[elem_id].append((callback, script))
+
+ on_before.clear()
+ on_after.clear()
+
def create_script_ui(self, script):
import modules.api.models as api_models
@@ -555,16 +579,7 @@ class ScriptRunner: self.infotext_fields.append((dropdown, lambda x: gr.update(value=x.get('Script', 'None'))))
self.infotext_fields.extend([(script.group, onload_script_visibility) for script in self.selectable_scripts])
- for script in self.scripts:
- for elem_id, callback in script.on_before_component_elem_id or []:
- items = self.on_before_component_elem_id.get(elem_id, [])
- items.append((callback, script))
- self.on_before_component_elem_id[elem_id] = items
-
- for elem_id, callback in script.on_after_component_elem_id or []:
- items = self.on_after_component_elem_id.get(elem_id, [])
- items.append((callback, script))
- self.on_after_component_elem_id[elem_id] = items
+ self.apply_on_before_component_callbacks()
return self.inputs
diff --git a/modules/sd_samplers_cfg_denoiser.py b/modules/sd_samplers_cfg_denoiser.py index 113425b2..bc9b97e4 100644 --- a/modules/sd_samplers_cfg_denoiser.py +++ b/modules/sd_samplers_cfg_denoiser.py @@ -56,6 +56,7 @@ class CFGDenoiser(torch.nn.Module): self.sampler = sampler
self.model_wrap = None
self.p = None
+ self.mask_before_denoising = False
@property
def inner_model(self):
@@ -104,7 +105,7 @@ class CFGDenoiser(torch.nn.Module): assert not is_edit_model or all(len(conds) == 1 for conds in conds_list), "AND is not supported for InstructPix2Pix checkpoint (unless using Image CFG scale = 1.0)"
- if self.mask is not None:
+ if self.mask_before_denoising and self.mask is not None:
x = self.init_latent * self.mask + self.nmask * x
batch_size = len(conds_list)
@@ -206,6 +207,9 @@ class CFGDenoiser(torch.nn.Module): else:
denoised = self.combine_denoised(x_out, conds_list, uncond, cond_scale)
+ if not self.mask_before_denoising and self.mask is not None:
+ denoised = self.init_latent * self.mask + self.nmask * denoised
+
self.sampler.last_latent = self.get_pred_x0(torch.cat([x_in[i:i + 1] for i in denoised_image_indexes]), torch.cat([x_out[i:i + 1] for i in denoised_image_indexes]), sigma)
if opts.live_preview_content == "Prompt":
diff --git a/modules/sd_samplers_timesteps.py b/modules/sd_samplers_timesteps.py index 16572c7e..c1f534ed 100644 --- a/modules/sd_samplers_timesteps.py +++ b/modules/sd_samplers_timesteps.py @@ -49,12 +49,12 @@ class CFGDenoiserTimesteps(CFGDenoiser): super().__init__(sampler)
self.alphas = shared.sd_model.alphas_cumprod
+ self.mask_before_denoising = True
def get_pred_x0(self, x_in, x_out, sigma):
- ts = int(sigma.item())
+ ts = sigma.to(dtype=int)
- s_in = x_in.new_ones([x_in.shape[0]])
- a_t = self.alphas[ts].item() * s_in
+ a_t = self.alphas[ts][:, None, None, None]
sqrt_one_minus_at = (1 - a_t).sqrt()
pred_x0 = (x_in - sqrt_one_minus_at * x_out) / a_t.sqrt()
diff --git a/modules/sd_samplers_timesteps_impl.py b/modules/sd_samplers_timesteps_impl.py index d32e3521..a72daafd 100644 --- a/modules/sd_samplers_timesteps_impl.py +++ b/modules/sd_samplers_timesteps_impl.py @@ -16,16 +16,17 @@ def ddim(model, x, timesteps, extra_args=None, callback=None, disable=None, eta= sigmas = eta * np.sqrt((1 - alphas_prev.cpu().numpy()) / (1 - alphas.cpu()) * (1 - alphas.cpu() / alphas_prev.cpu().numpy()))
extra_args = {} if extra_args is None else extra_args
- s_in = x.new_ones([x.shape[0]])
+ s_in = x.new_ones((x.shape[0]))
+ s_x = x.new_ones((x.shape[0], 1, 1, 1))
for i in tqdm.trange(len(timesteps) - 1, disable=disable):
index = len(timesteps) - 1 - i
e_t = model(x, timesteps[index].item() * s_in, **extra_args)
- a_t = alphas[index].item() * s_in
- a_prev = alphas_prev[index].item() * s_in
- sigma_t = sigmas[index].item() * s_in
- sqrt_one_minus_at = sqrt_one_minus_alphas[index].item() * s_in
+ a_t = alphas[index].item() * s_x
+ a_prev = alphas_prev[index].item() * s_x
+ sigma_t = sigmas[index].item() * s_x
+ sqrt_one_minus_at = sqrt_one_minus_alphas[index].item() * s_x
pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
dir_xt = (1. - a_prev - sigma_t ** 2).sqrt() * e_t
@@ -47,13 +48,14 @@ def plms(model, x, timesteps, extra_args=None, callback=None, disable=None): extra_args = {} if extra_args is None else extra_args
s_in = x.new_ones([x.shape[0]])
+ s_x = x.new_ones((x.shape[0], 1, 1, 1))
old_eps = []
def get_x_prev_and_pred_x0(e_t, index):
# select parameters corresponding to the currently considered timestep
- a_t = alphas[index].item() * s_in
- a_prev = alphas_prev[index].item() * s_in
- sqrt_one_minus_at = sqrt_one_minus_alphas[index].item() * s_in
+ a_t = alphas[index].item() * s_x
+ a_prev = alphas_prev[index].item() * s_x
+ sqrt_one_minus_at = sqrt_one_minus_alphas[index].item() * s_x
# current prediction for x_0
pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
|