forked from extern/easydiffusion
163 lines
6.2 KiB
Diff
163 lines
6.2 KiB
Diff
diff --git a/optimizedSD/ddpm.py b/optimizedSD/ddpm.py
|
|
index 79058bc..a473411 100644
|
|
--- a/optimizedSD/ddpm.py
|
|
+++ b/optimizedSD/ddpm.py
|
|
@@ -564,12 +564,12 @@ class UNet(DDPM):
|
|
unconditional_guidance_scale=unconditional_guidance_scale,
|
|
callback=callback, img_callback=img_callback)
|
|
|
|
+ yield from samples
|
|
+
|
|
if(self.turbo):
|
|
self.model1.to("cpu")
|
|
self.model2.to("cpu")
|
|
|
|
- return samples
|
|
-
|
|
@torch.no_grad()
|
|
def plms_sampling(self, cond,b, img,
|
|
ddim_use_original_steps=False,
|
|
@@ -608,10 +608,10 @@ class UNet(DDPM):
|
|
old_eps.append(e_t)
|
|
if len(old_eps) >= 4:
|
|
old_eps.pop(0)
|
|
- if callback: callback(i)
|
|
- if img_callback: img_callback(pred_x0, i)
|
|
+ if callback: yield from callback(i)
|
|
+ if img_callback: yield from img_callback(pred_x0, i)
|
|
|
|
- return img
|
|
+ yield from img_callback(img, len(iterator)-1)
|
|
|
|
@torch.no_grad()
|
|
def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
|
|
@@ -740,13 +740,13 @@ class UNet(DDPM):
|
|
unconditional_guidance_scale=unconditional_guidance_scale,
|
|
unconditional_conditioning=unconditional_conditioning)
|
|
|
|
- if callback: callback(i)
|
|
- if img_callback: img_callback(x_dec, i)
|
|
+ if callback: yield from callback(i)
|
|
+ if img_callback: yield from img_callback(x_dec, i)
|
|
|
|
if mask is not None:
|
|
- return x0 * mask + (1. - mask) * x_dec
|
|
+ x_dec = x0 * mask + (1. - mask) * x_dec
|
|
|
|
- return x_dec
|
|
+ yield from img_callback(x_dec, len(iterator)-1)
|
|
|
|
|
|
@torch.no_grad()
|
|
@@ -820,12 +820,12 @@ class UNet(DDPM):
|
|
|
|
|
|
d = to_d(x, sigma_hat, denoised)
|
|
- if callback: callback(i)
|
|
- if img_callback: img_callback(x, i)
|
|
+ if callback: yield from callback(i)
|
|
+ if img_callback: yield from img_callback(x, i)
|
|
dt = sigmas[i + 1] - sigma_hat
|
|
# Euler method
|
|
x = x + d * dt
|
|
- return x
|
|
+ yield from img_callback(x, len(sigmas)-1)
|
|
|
|
@torch.no_grad()
|
|
def euler_ancestral_sampling(self,ac,x, S, cond, unconditional_conditioning = None, unconditional_guidance_scale = 1,extra_args=None, callback=None, disable=None, img_callback=None):
|
|
@@ -852,14 +852,14 @@ class UNet(DDPM):
|
|
denoised = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
|
|
|
|
sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1])
|
|
- if callback: callback(i)
|
|
- if img_callback: img_callback(x, i)
|
|
+ if callback: yield from callback(i)
|
|
+ if img_callback: yield from img_callback(x, i)
|
|
d = to_d(x, sigmas[i], denoised)
|
|
# Euler method
|
|
dt = sigma_down - sigmas[i]
|
|
x = x + d * dt
|
|
x = x + torch.randn_like(x) * sigma_up
|
|
- return x
|
|
+ yield from img_callback(x, len(sigmas)-1)
|
|
|
|
|
|
|
|
@@ -892,8 +892,8 @@ class UNet(DDPM):
|
|
denoised = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
|
|
|
|
d = to_d(x, sigma_hat, denoised)
|
|
- if callback: callback(i)
|
|
- if img_callback: img_callback(x, i)
|
|
+ if callback: yield from callback(i)
|
|
+ if img_callback: yield from img_callback(x, i)
|
|
dt = sigmas[i + 1] - sigma_hat
|
|
if sigmas[i + 1] == 0:
|
|
# Euler method
|
|
@@ -913,7 +913,7 @@ class UNet(DDPM):
|
|
d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
|
|
d_prime = (d + d_2) / 2
|
|
x = x + d_prime * dt
|
|
- return x
|
|
+ yield from img_callback(x, len(sigmas)-1)
|
|
|
|
|
|
@torch.no_grad()
|
|
@@ -944,8 +944,8 @@ class UNet(DDPM):
|
|
e_t_uncond, e_t = (x_in + eps * c_out).chunk(2)
|
|
denoised = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
|
|
|
|
- if callback: callback(i)
|
|
- if img_callback: img_callback(x, i)
|
|
+ if callback: yield from callback(i)
|
|
+ if img_callback: yield from img_callback(x, i)
|
|
|
|
d = to_d(x, sigma_hat, denoised)
|
|
# Midpoint method, where the midpoint is chosen according to a rho=3 Karras schedule
|
|
@@ -966,7 +966,7 @@ class UNet(DDPM):
|
|
|
|
d_2 = to_d(x_2, sigma_mid, denoised_2)
|
|
x = x + d_2 * dt_2
|
|
- return x
|
|
+ yield from img_callback(x, len(sigmas)-1)
|
|
|
|
|
|
@torch.no_grad()
|
|
@@ -994,8 +994,8 @@ class UNet(DDPM):
|
|
|
|
|
|
sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1])
|
|
- if callback: callback(i)
|
|
- if img_callback: img_callback(x, i)
|
|
+ if callback: yield from callback(i)
|
|
+ if img_callback: yield from img_callback(x, i)
|
|
d = to_d(x, sigmas[i], denoised)
|
|
# Midpoint method, where the midpoint is chosen according to a rho=3 Karras schedule
|
|
sigma_mid = ((sigmas[i] ** (1 / 3) + sigma_down ** (1 / 3)) / 2) ** 3
|
|
@@ -1016,7 +1016,7 @@ class UNet(DDPM):
|
|
d_2 = to_d(x_2, sigma_mid, denoised_2)
|
|
x = x + d_2 * dt_2
|
|
x = x + torch.randn_like(x) * sigma_up
|
|
- return x
|
|
+ yield from img_callback(x, len(sigmas)-1)
|
|
|
|
|
|
@torch.no_grad()
|
|
@@ -1042,8 +1042,8 @@ class UNet(DDPM):
|
|
e_t_uncond, e_t = (x_in + eps * c_out).chunk(2)
|
|
denoised = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
|
|
|
|
- if callback: callback(i)
|
|
- if img_callback: img_callback(x, i)
|
|
+ if callback: yield from callback(i)
|
|
+ if img_callback: yield from img_callback(x, i)
|
|
|
|
d = to_d(x, sigmas[i], denoised)
|
|
ds.append(d)
|
|
@@ -1054,4 +1054,4 @@ class UNet(DDPM):
|
|
cur_order = min(i + 1, order)
|
|
coeffs = [linear_multistep_coeff(cur_order, sigmas.cpu(), i, j) for j in range(cur_order)]
|
|
x = x + sum(coeff * d for coeff, d in zip(coeffs, reversed(ds)))
|
|
- return x
|
|
+ yield from img_callback(x, len(sigmas)-1)
|