Fix coeff in ddim_sample

mathpluscode · mathpluscode · commit 5a54cffcaab5 · 2023-03-26T17:11:19.000+01:00
diff --git a/imgx/diffusion/gaussian_diffusion.py b/imgx/diffusion/gaussian_diffusion.py
@@ -539,7 +539,7 @@ def ddim_sample(
         alphas_cumprod_prev = extract_and_expand(
             self.alphas_cumprod_prev, t, x_t.ndim
         )
-        coeff_start = alphas_cumprod_prev
+        coeff_start = jnp.sqrt(alphas_cumprod_prev)
         log_variance = (
             extract_and_expand(self.posterior_log_variance_clipped, t, x_t.ndim)
             * eta

Original file line number	Diff line number	Diff line change
`@@ -539,7 +539,7 @@ def ddim_sample(`
`539`	`539`	`alphas_cumprod_prev = extract_and_expand(`
`540`	`540`	`self.alphas_cumprod_prev, t, x_t.ndim`
`541`	`541`	`)`
`542`		`- coeff_start = alphas_cumprod_prev`
	`542`	`+ coeff_start = jnp.sqrt(alphas_cumprod_prev)`
`543`	`543`	`log_variance = (`
`544`	`544`	`extract_and_expand(self.posterior_log_variance_clipped, t, x_t.ndim)`
`545`	`545`	`* eta`