aljaz/DALLE2-pytorch
1
1
Fork 0
mirror of https://github.com/lucidrains/DALLE2-pytorch.git synced 2025-12-19 09:44:19 +01:00
Code Issues Packages Projects Releases Wiki Activity

in ddim, noise should be predicted after x0 is maybe clipped, thanks to @lukovnikov for pointing this out in another repository

Browse Source
This commit is contained in:
Phil Wang
2022-09-01 09:40:47 -07:00
parent 8bbc956ff1
commit 0d82dff9c5
2 changed files with 24 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
dalle2_pytorch/dalle2_pytorch.py
View File

@@ -1281,12 +1281,14 @@ class DiffusionPrior(nn.Module):
pred = self.net.forward_with_cond_scale(image_embed, time_cond, self_cond = self_cond, cond_scale = cond_scale, **text_cond) pred = self.net.forward_with_cond_scale(image_embed, time_cond, self_cond = self_cond, cond_scale = cond_scale, **text_cond)
# derive x0
if self.predict_x_start: if self.predict_x_start:
x_start = pred x_start = pred
pred_noise = self.noise_scheduler.predict_noise_from_start(image_embed, t = time_cond, x0 = pred)
else: else:
x_start = self.noise_scheduler.predict_start_from_noise(image_embed, t = time_cond, noise = pred) x_start = self.noise_scheduler.predict_start_from_noise(image_embed, t = time_cond, noise = pred_noise)
pred_noise = pred
# clip x0 before maybe predicting noise
if not self.predict_x_start: if not self.predict_x_start:
x_start.clamp_(-1., 1.) x_start.clamp_(-1., 1.)
@@ -1294,6 +1296,13 @@ class DiffusionPrior(nn.Module):
if self.predict_x_start and self.sampling_clamp_l2norm: if self.predict_x_start and self.sampling_clamp_l2norm:
x_start = self.l2norm_clamp_embed(x_start) x_start = self.l2norm_clamp_embed(x_start)
# predict noise
if self.predict_x_start:
pred_noise = self.noise_scheduler.predict_noise_from_start(image_embed, t = time_cond, x0 = x_start)
else:
pred_noise = pred
if time_next < 0: if time_next < 0:
image_embed = x_start image_embed = x_start
continue continue
@@ -2897,16 +2906,25 @@ class Decoder(nn.Module):
pred, _ = self.parse_unet_output(learned_variance, unet_output) pred, _ = self.parse_unet_output(learned_variance, unet_output)
# predict x0
if predict_x_start: if predict_x_start:
x_start = pred x_start = pred
pred_noise = noise_scheduler.predict_noise_from_start(img, t = time_cond, x0 = pred)
else: else:
x_start = noise_scheduler.predict_start_from_noise(img, t = time_cond, noise = pred) x_start = noise_scheduler.predict_start_from_noise(img, t = time_cond, noise = pred)
pred_noise = pred
# maybe clip x0
if clip_denoised: if clip_denoised:
x_start = self.dynamic_threshold(x_start) x_start = self.dynamic_threshold(x_start)
# predict noise
if predict_x_start:
pred_noise = noise_scheduler.predict_noise_from_start(img, t = time_cond, x0 = pred)
else:
pred_noise = pred
c1 = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt() c1 = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
c2 = ((1 - alpha_next) - torch.square(c1)).sqrt() c2 = ((1 - alpha_next) - torch.square(c1)).sqrt()
noise = torch.randn_like(img) if not is_last_timestep else 0. noise = torch.randn_like(img) if not is_last_timestep else 0.

2
dalle2_pytorch/version.py
View File

@@ -1 +1 @@
__version__ = '1.10.4' __version__ = '1.10.5'