v prediction for sd2

Author: divamgupta

backends/stable_diffusion/schedulers/scheduling_ddim.py

@@ -88,6 +88,7 @@ def __init__(
  clip_sample: bool = True,
  set_alpha_to_one: bool = True,
  tensor_format: str = "pt",
+ prediction_type: str = "epsilon"
  ):
  if trained_betas is not None:
  self.betas = np.asarray(trained_betas)
@@ -115,6 +116,7 @@ def __init__(
  self.clip_sample = clip_sample
  self.set_alpha_to_one = set_alpha_to_one
  self.tensor_format = tensor_format
+ self.prediction_type = prediction_type
 
  # At every step in ddim, we are looking into the previous alphas_cumprod
  # For the final step, there is no previous alphas_cumprod because we are already at 0
@@ -217,8 +219,14 @@ def step(
 
  # 3. compute predicted original sample from predicted noise also called
  # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
- pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
-
+ if self.config.prediction_type == "epsilon":
+ pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
+ pred_epsilon = model_output
+ elif self.config.prediction_type == "v_prediction":
+ pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
+ pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
+ else:
+ raise ValueError("Unknown prediction_type")
  # 4. Clip "predicted x_0"
  if self.config.clip_sample:
  pred_original_sample = self.clip(pred_original_sample, -1, 1)
@@ -230,10 +238,10 @@ def step(
 
  if use_clipped_model_output:
  # the model_output is always re-derived from the clipped x_0 in Glide
- model_output = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)
+ pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)
 
  # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
- pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * model_output
+ pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * pred_epsilon
 
  # 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
  prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction

backends/stable_diffusion/stable_diffusion.py

@@ -111,7 +111,21 @@ def get_scheduler(name):
  set_alpha_to_one=False,
  # steps_offset= 1,
  trained_betas= None,
- tensor_format="np"
+ tensor_format="np",
+ )
+
+ if name == "ddim_v":
+ return DDIMScheduler(
+ beta_start=0.00085,
+ beta_end=0.012,
+ beta_schedule="scaled_linear",
+ clip_sample= False,
+ num_train_timesteps= 1000,
+ set_alpha_to_one=False,
+ # steps_offset= 1,
+ trained_betas= None,
+ tensor_format="np",
+ prediction_type="v_prediction"
  )
 
  if name == "lmsd":

backends/stable_diffusion/tests.py

@@ -287,15 +287,14 @@ def test_sd2_2():
 def test_sd2_4():
 
  img = sd.generate(
- prompt="A Tree" , 
+ prompt="a tree" , 
  img_height=512, 
  img_width=512, 
- seed=1, 
- num_steps=10,
+ seed=13, 
+ num_steps=30,
  tdict_path="/Volumes/ext_drive_1/sd_data_models/v2-1_768-nonema-pruned.tdict",
- batch_size=1,
  dtype="float32",
- scheduler='pndm',
+ scheduler='ddim_v',
  mode="txt2img" )
 
  # sd2_a_cat_111_test_sd2_3.png