Coverage for widitapp/diffusion/diffusion_utils.py: 100.00%

1# Modified from OpenAI's diffusion repos 

2# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py 

3# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion 

4# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py 

5 

6import torch as th 

7import numpy as np 

8 

9 

10def normal_kl(mean1, logvar1, mean2, logvar2): 

11 """ 

12 Compute the KL divergence between two gaussians. 

13 Shapes are automatically broadcasted, so batches can be compared to 

14 scalars, among other use cases. 

15 """ 

16 tensor = None 

17 for obj in (mean1, logvar1, mean2, logvar2): 

18 if isinstance(obj, th.Tensor): 

19 tensor = obj 

20 break 

21 assert tensor is not None, "at least one argument must be a Tensor" 

22 

23 # Force variances to be Tensors. Broadcasting helps convert scalars to 

24 # Tensors, but it does not work for th.exp(). 

25 logvar1, logvar2 = [ 

26 x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) 

27 for x in (logvar1, logvar2) 

28 ] 

29 

30 return 0.5 * ( 

31 -1.0 

32 + logvar2 

33 - logvar1 

34 + th.exp(logvar1 - logvar2) 

35 + ((mean1 - mean2) ** 2) * th.exp(-logvar2) 

36 ) 

37 

38 

39def approx_standard_normal_cdf(x): 

40 """ 

41 A fast approximation of the cumulative distribution function of the 

42 standard normal. 

43 """ 

44 return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) 

45 

46 

47def continuous_gaussian_log_likelihood(x, *, means, log_scales): 

48 """ 

49 Compute the log-likelihood of a continuous Gaussian distribution. 

50 :param x: the targets 

51 :param means: the Gaussian mean Tensor. 

52 :param log_scales: the Gaussian log stddev Tensor. 

53 :return: a tensor like x of log probabilities (in nats). 

54 """ 

55 centered_x = x - means 

56 inv_stdv = th.exp(-log_scales) 

57 normalized_x = centered_x * inv_stdv 

58 log_probs = th.distributions.Normal(th.zeros_like(x), th.ones_like(x)).log_prob(normalized_x) 

59 return log_probs 

60 

61 

62def discretized_gaussian_log_likelihood(x, *, means, log_scales): 

63 """ 

64 Compute the log-likelihood of a Gaussian distribution discretizing to a 

65 given image. 

66 :param x: the target images. It is assumed that this was uint8 values, 

67 rescaled to the range [-1, 1]. 

68 :param means: the Gaussian mean Tensor. 

69 :param log_scales: the Gaussian log stddev Tensor. 

70 :return: a tensor like x of log probabilities (in nats). 

71 """ 

72 assert x.shape == means.shape == log_scales.shape 

73 centered_x = x - means 

74 inv_stdv = th.exp(-log_scales) 

75 plus_in = inv_stdv * (centered_x + 1.0 / 255.0) 

76 cdf_plus = approx_standard_normal_cdf(plus_in) 

77 min_in = inv_stdv * (centered_x - 1.0 / 255.0) 

78 cdf_min = approx_standard_normal_cdf(min_in) 

79 log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12)) 

80 log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12)) 

81 cdf_delta = cdf_plus - cdf_min 

82 log_probs = th.where( 

83 x < -0.999, 

84 log_cdf_plus, 

85 th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))), 

86 ) 

87 assert log_probs.shape == x.shape 

88 return log_probs