From 618493714f75570ac4c4c298abe94f9b7305079d Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Tue, 28 Jun 2022 14:34:39 -0700 Subject: [PATCH] clamp the sigma coming out of the log normal distribution --- denoising_diffusion_pytorch/elucidated_diffusion.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/denoising_diffusion_pytorch/elucidated_diffusion.py b/denoising_diffusion_pytorch/elucidated_diffusion.py index c918950..114dc3e 100644 --- a/denoising_diffusion_pytorch/elucidated_diffusion.py +++ b/denoising_diffusion_pytorch/elucidated_diffusion.py @@ -99,7 +99,8 @@ class ElucidatedDiffusion(nn.Module): # noise distribution def noise_distribution(self, batch_size): - return (self.P_mean + self.P_std * torch.randn((batch_size,), device = self.device)).exp() + sigmas = (self.P_mean + self.P_std * torch.randn((batch_size,), device = self.device)).exp() + return sigmas.clamp(min = self.sigma_min, max =self.sigma_max) def loss_weight(self, sigma): return (sigma ** 2 + self.sigma_data ** 2) * (sigma * self.sigma_data) ** -2