DDIM masked sampling to support webui "inpainting"

ldm/models/diffusion/ddim.py

@@ -221,7 +221,7 @@ def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):
 
  @torch.no_grad()
  def decode(self, x_latent, cond, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None,
- use_original_steps=False):
+ use_original_steps=False, z_mask = None, x0=None):
 
  timesteps = np.arange(self.ddpm_num_timesteps) if use_original_steps else self.ddim_timesteps
  timesteps = timesteps[:t_start]
@@ -235,6 +235,13 @@ def decode(self, x_latent, cond, t_start, unconditional_guidance_scale=1.0, unco
  for i, step in enumerate(iterator):
  index = total_steps - i - 1
  ts = torch.full((x_latent.shape[0],), step, device=x_latent.device, dtype=torch.long)
+
+ if z_mask is not None and i < total_steps - 2:
+ assert x0 is not None
+ img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass?
+ mask_inv = 1. - z_mask
+ x_dec = (img_orig * mask_inv) + (z_mask * x_dec)
+
  x_dec, _ = self.p_sample_ddim(x_dec, cond, ts, index=index, use_original_steps=use_original_steps,
  unconditional_guidance_scale=unconditional_guidance_scale,
  unconditional_conditioning=unconditional_conditioning)