added more samplers

Author: divamgupta

backends/stable_diffusion/schedulers/k_euler.py

@@ -0,0 +1,53 @@
+import numpy as np
+
+
+def get_alphas_cumprod(beta_start=0.00085, beta_end=0.0120, n_training_steps=1000):
+ betas = np.linspace(beta_start ** 0.5, beta_end ** 0.5, n_training_steps, dtype=np.float32) ** 2
+ alphas = 1.0 - betas
+ alphas_cumprod = np.cumprod(alphas, axis=0)
+ return alphas_cumprod
+
+
+
+class KEulerSampler():
+ def __init__(self, n_inference_steps=50, n_training_steps=1000):
+ pass
+
+ def set_timesteps(self, n_inference_steps, n_training_steps=1000): #dg
+ timesteps = np.linspace(n_training_steps - 1, 0, n_inference_steps)
+
+ alphas_cumprod = get_alphas_cumprod(n_training_steps=n_training_steps)
+ sigmas = ((1 - alphas_cumprod) / alphas_cumprod) ** 0.5
+ log_sigmas = np.log(sigmas)
+ log_sigmas = np.interp(timesteps, range(n_training_steps), log_sigmas)
+ sigmas = np.exp(log_sigmas)
+ sigmas = np.append(sigmas, 0)
+ 
+ self.sigmas = sigmas
+ self.initial_scale = sigmas.max()
+ self.timesteps = timesteps
+ self.n_inference_steps = n_inference_steps
+ self.n_training_steps = n_training_steps
+ self.step_count = 0
+
+
+ def get_input_scale(self, step_count=None):
+ if step_count is None:
+ step_count = self.step_count
+ sigma = self.sigmas[step_count]
+ return 1 / (sigma ** 2 + 1) ** 0.5
+
+ def add_noise(self, latent, noise, idx ): #dg
+ for i in idx:
+ assert idx[0] == i
+ sc = self.sigmas[idx[0]]
+ return latent + noise*sc
+
+
+ def step(self, output , t , latents , seed=None ): #dg
+ 
+
+ sigma_from = self.sigmas[t]
+ sigma_to = self.sigmas[t + 1]
+ latents += output * (sigma_to - sigma_from)
+ return {"prev_sample": latents } #latents

backends/stable_diffusion/schedulers/k_euler_ancestral.py

@@ -0,0 +1,61 @@
+import numpy as np
+
+
+def get_alphas_cumprod(beta_start=0.00085, beta_end=0.0120, n_training_steps=1000):
+ betas = np.linspace(beta_start ** 0.5, beta_end ** 0.5, n_training_steps, dtype=np.float32) ** 2
+ alphas = 1.0 - betas
+ alphas_cumprod = np.cumprod(alphas, axis=0)
+ return alphas_cumprod
+
+
+class KEulerAncestralSampler():
+ def __init__(self):
+ pass
+
+
+ def set_timesteps(self, n_inference_steps, n_training_steps=1000): #dg
+ timesteps = np.linspace(n_training_steps - 1, 0, n_inference_steps)
+
+ alphas_cumprod = get_alphas_cumprod(n_training_steps=n_training_steps)
+ sigmas = ((1 - alphas_cumprod) / alphas_cumprod) ** 0.5
+ log_sigmas = np.log(sigmas)
+ log_sigmas = np.interp(timesteps, range(n_training_steps), log_sigmas)
+ sigmas = np.exp(log_sigmas)
+ sigmas = np.append(sigmas, 0)
+ 
+ self.sigmas = sigmas
+ self.initial_scale = sigmas.max()
+ self.timesteps = timesteps
+ self.n_inference_steps = n_inference_steps
+ self.n_training_steps = n_training_steps
+
+
+ def get_input_scale(self, step_count=None):
+ sigma = self.sigmas[step_count]
+ return 1 / (sigma ** 2 + 1) ** 0.5
+
+ # def set_strength(self, strength=1):
+ # start_step = self.n_inference_steps - int(self.n_inference_steps * strength)
+ # self.timesteps = np.linspace(self.n_training_steps - 1, 0, self.n_inference_steps)
+ # self.timesteps = self.timesteps[start_step:]
+ # self.initial_scale = self.sigmas[start_step]
+ # self.step_count = start_step
+
+
+ def add_noise(self, latent, noise, idx ): #dg
+ for i in idx:
+ assert idx[0] == i
+ sc = self.sigmas[idx[0]]
+ return latent + noise*sc
+
+ def step(self, output , t , latents , seed ): #dg
+ 
+
+ sigma_from = self.sigmas[t]
+ sigma_to = self.sigmas[t + 1]
+ sigma_up = sigma_to * (1 - (sigma_to ** 2 / sigma_from ** 2)) ** 0.5
+ sigma_down = sigma_to ** 2 / sigma_from
+ latents += output * (sigma_down - sigma_from)
+ noise = np.random.RandomState(seed).normal(size=latents.shape).astype('float32')
+ latents += noise * sigma_up
+ return {"prev_sample": latents } #latents

backends/stable_diffusion/schedulers/scheduling_ddim.py

@@ -130,6 +130,11 @@ def __init__(
  self.config = self
  self.set_format(tensor_format=tensor_format)
 
+ self.initial_scale = 1
+
+ def get_input_scale(self, _ ):
+ return 1
+
  def _get_variance(self, timestep, prev_timestep):
  alpha_prod_t = self.alphas_cumprod[timestep]
  alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
@@ -160,11 +165,15 @@ def step(
  model_output: Union[ np.ndarray],
  timestep: int,
  sample: Union[ np.ndarray],
+ seed=None,
  eta: float = 0.0,
  use_clipped_model_output: bool = False,
  generator=None,
  return_dict: bool = True,
  ) -> Union[ Tuple]:
+
+ timestep = self.timesteps[timestep]
+
  """
  Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
  process from the learned model outputs (most often the predicted noise).
@@ -249,6 +258,11 @@ def add_noise(
  noise: Union[ np.ndarray],
  timesteps: Union[np.ndarray],
  ) -> Union[ np.ndarray]:
+
+ 
+ timesteps = np.array([ self.timesteps[i] for i in timesteps ])
+
+
  sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5
  sqrt_alpha_prod = self.match_shape(sqrt_alpha_prod, original_samples)
  sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5

backends/stable_diffusion/schedulers/scheduling_lms_discrete.py

@@ -91,6 +91,12 @@ def __init__(
  self.timestep_values = timestep_values
  self.tensor_format = tensor_format
 
+ self.initial_scale = self.sigmas[0]
+
+ def get_input_scale(self, step_count=None):
+ sigma = self.sigmas[step_count]
+ return 1 / (sigma ** 2 + 1) ** 0.5
+
  def get_lms_coefficient(self, order, t, current_order):
  """
  Compute a linear multistep coefficient.
@@ -138,6 +144,7 @@ def step(
  model_output: Union[torch.FloatTensor, np.ndarray],
  timestep: int,
  sample: Union[torch.FloatTensor, np.ndarray],
+ seed=None,
  order: int = 4,
  return_dict: bool = True,
  ) -> Union[ Tuple]:

backends/stable_diffusion/schedulers/scheduling_pndm.py

@@ -132,6 +132,11 @@ def __init__(
  self.config = self
  self.set_format(tensor_format=tensor_format)
 
+ self.initial_scale = 1
+
+ def get_input_scale(self, _ ):
+ return 1
+
  def set_timesteps(self, num_inference_steps: int, offset: int = 0):
  """
  Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
@@ -175,8 +180,12 @@ def step(
  model_output: Union[ np.ndarray],
  timestep: int,
  sample: Union[ np.ndarray],
+ seed=None,
  return_dict: bool = True,
  ) -> Union[SchedulerOutput, Tuple]:
+
+ timestep = self.timesteps[timestep]
+
  """
  Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
  process from the learned model outputs (most often the predicted noise).
@@ -354,6 +363,9 @@ def add_noise(
  noise: Union[ np.ndarray],
  timesteps: Union[ np.ndarray],
  ) :
+
+ timesteps = np.array([ self.timesteps[i] for i in timesteps ])
+
  # mps requires indices to be in the same device, so we use cpu as is the default with cuda
  timesteps = timesteps.to(self.alphas_cumprod.device)
  sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5

backends/stable_diffusion/stable_diffusion.py

@@ -42,6 +42,8 @@
 from schedulers.scheduling_ddim import DDIMScheduler
 from schedulers.scheduling_lms_discrete import LMSDiscreteScheduler
 from schedulers.scheduling_pndm import PNDMScheduler
+from schedulers.k_euler_ancestral import KEulerAncestralSampler
+from schedulers.k_euler import KEulerSampler
 
 def process_inp_img(input_image):
  input_image = Image.open(input_image)
@@ -128,6 +130,12 @@ def get_scheduler(name):
  skip_prk_steps = True,
  tensor_format="np")
 
+ if name == "k_euler_ancestral":
+ return KEulerAncestralSampler()
+
+ if name == "k_euler":
+ return KEulerSampler()
+
 
 def dummy_callback(state="" , progress=-1):
  pass
@@ -283,18 +291,14 @@ def prepare_init_latent(self , sd_run):
  # latent_np = latent_np * np.float64(self.scheduler.init_noise_sigma)
  sd_run.latent = latent_np
 
- if isinstance(self.scheduler, LMSDiscreteScheduler):
- sd_run.latent = sd_run.latent * self.scheduler.sigmas[0]
+ sd_run.latent = sd_run.latent * self.scheduler.initial_scale
 
  else:
  latent = self.get_encoded_img(sd_run , sd_run.input_image_processed )
  sd_run.encoded_img_unmasked = np.copy(latent)
 
- if isinstance(self.scheduler, LMSDiscreteScheduler):
- start_timestep = np.array([self.t_to_i(sd_run.start_timestep)] * sd_run.batch_size, dtype=np.int64 )
- else:
- start_timestep = np.array([sd_run.start_timestep] * sd_run.batch_size, dtype=np.int64 )
-
+ start_timestep = np.array([self.t_to_i(sd_run.start_timestep)] * sd_run.batch_size, dtype=np.int64 )
+ 
  noise = self.get_noise(sd_run.seed , latent.shape )
  latent = self.scheduler.add_noise(latent, noise, start_timestep )
  sd_run.latent = latent
@@ -344,12 +348,8 @@ def get_unet_out(self, sd_run):
  np.repeat( (1 - sd_run.processed_mask_downscaled ), sd_run.batch_size , axis=0) , 
  sd_run.encoded_masked_img 
  ], axis=-1)
-
- if isinstance(self.scheduler, LMSDiscreteScheduler):
- sigma = self.scheduler.sigmas[self.t_to_i(t)]
- latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5)
-
- # latent_model_input = self.scheduler.scale_model_input(sd_run.latent, t)
+ 
+ latent_model_input = latent_model_input * self.scheduler.get_input_scale(self.t_to_i(t))
 
  if sd_run.combine_unet_run:
  latent_combined = np.concatenate([latent_model_input,latent_model_input])
@@ -374,22 +374,16 @@ def get_next_latent(self, sd_run ):
  eta = 0.0 # should be between 0 and 1, but 0 for now
  extra_step_kwargs["eta"] = eta
 
- if isinstance(self.scheduler, LMSDiscreteScheduler):
- latents = self.scheduler.step(noise_pred, self.t_to_i(t), sd_run.latent, **extra_step_kwargs)["prev_sample"]
- else:
- latents = self.scheduler.step(noise_pred, t, sd_run.latent, **extra_step_kwargs)["prev_sample"]
+ step_seed = sd_run.seed + 10000 + self.t_to_i(t)*4242
+ latents = self.scheduler.step(noise_pred, self.t_to_i(t), sd_run.latent , seed=step_seed , **extra_step_kwargs)["prev_sample"]
+ 
 
  if sd_run.do_masking:
 
  latent_proper = np.copy(sd_run.encoded_img_unmasked)
 
  noise = self.get_noise(sd_run.seed , latent_proper.shape )
-
- if isinstance(self.scheduler, LMSDiscreteScheduler):
- latent_proper = self.scheduler.add_noise(latent_proper, noise, np.array([self.t_to_i(sd_run.current_t)] * sd_run.batch_size, dtype=np.int64 ) )
- else:
- latent_proper = self.scheduler.add_noise(latent_proper, noise, np.array([sd_run.current_t] * sd_run.batch_size, dtype=np.int64 ) )
-
+ latent_proper = self.scheduler.add_noise(latent_proper, noise, np.array([self.t_to_i(sd_run.current_t)] * sd_run.batch_size, dtype=np.int64 ) )
 
  latents = (latent_proper * sd_run.processed_mask_downscaled) + (latents * (1 - sd_run.processed_mask_downscaled))
 
@@ -413,7 +407,7 @@ def generate(
  mask_image=None,
  negative_prompt="",
  input_image_strength=0.5,
- scheduler='pndm',
+ scheduler='k_euler',
  tdict_path=None, # if none then it will just use current one
  dtype='float16',
  mode="txt2img" # txt2img , img2img, inpaint_15

electron_app/src/bridge.js

@@ -92,11 +92,11 @@ function start_bridge() {
 
  if(RESTART_BACKEND_ON_CLOSE){
  // dialog.showMessageBox( win , { message: "Error in backend : " + last_few_err }); // this is non blocking 
- if(!(last_few_err.includes("leaked semaphore objects to clean up at shutdown"))){
- // this leaked semaphore issue just happens sometimes. so for now lets just silently restart 
- dialog.showMessageBox( { message: "Error in backend : " + last_few_err });
- }
- 
+ // if(!(last_few_err.includes("leaked semaphore objects to clean up at shutdown"))){
+ //  // this leaked semaphore issue just happens sometimes. so for now lets just silently restart 
+ //  dialog.showMessageBox( { message: "Error in backend : " + last_few_err });
+ // }
+ dialog.showMessageBox( { message: "Error in backend : " + last_few_err });
  return start_bridge()
  }
  else{