Coverage for widitapp/diffusion/respace.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 numpy as np 

7import torch as th 

8 

9from .gaussian_diffusion import GaussianDiffusion 

10 

11 

12def space_timesteps(num_timesteps, section_counts): 

13 """ 

14 Create a list of timesteps to use from an original diffusion process, 

15 given the number of timesteps we want to take from equally-sized portions 

16 of the original process. 

17 For example, if there's 300 timesteps and the section counts are [10,15,20] 

18 then the first 100 timesteps are strided to be 10 timesteps, the second 100 

19 are strided to be 15 timesteps, and the final 100 are strided to be 20. 

20 If the stride is a string starting with "ddim", then the fixed striding 

21 from the DDIM paper is used, and only one section is allowed. 

22 :param num_timesteps: the number of diffusion steps in the original 

23 process to divide up. 

24 :param section_counts: either a list of numbers, or a string containing 

25 comma-separated numbers, indicating the step count 

26 per section. As a special case, use "ddimN" where N 

27 is a number of steps to use the striding from the 

28 DDIM paper. 

29 :return: a set of diffusion steps from the original process to use. 

30 """ 

31 if isinstance(section_counts, str): 

32 if section_counts.startswith("ddim"): 

33 desired_count = int(section_counts[len("ddim") :]) 

34 for i in range(1, num_timesteps): 

35 if len(range(0, num_timesteps, i)) == desired_count: 

36 return set(range(0, num_timesteps, i)) 

37 raise ValueError( 

38 f"cannot create exactly {num_timesteps} steps with an integer stride" 

39 ) 

40 section_counts = [int(x) for x in section_counts.split(",")] 

41 size_per = num_timesteps // len(section_counts) 

42 extra = num_timesteps % len(section_counts) 

43 start_idx = 0 

44 all_steps = [] 

45 for i, section_count in enumerate(section_counts): 

46 size = size_per + (1 if i < extra else 0) 

47 if size < section_count: 

48 raise ValueError( 

49 f"cannot divide section of {size} steps into {section_count}" 

50 ) 

51 if section_count <= 1: 

52 frac_stride = 1 

53 else: 

54 frac_stride = (size - 1) / (section_count - 1) 

55 cur_idx = 0.0 

56 taken_steps = [] 

57 for _ in range(section_count): 

58 taken_steps.append(start_idx + round(cur_idx)) 

59 cur_idx += frac_stride 

60 all_steps += taken_steps 

61 start_idx += size 

62 return set(all_steps) 

63 

64 

65class SpacedDiffusion(GaussianDiffusion): 

66 """ 

67 A diffusion process which can skip steps in a base diffusion process. 

68 :param use_timesteps: a collection (sequence or set) of timesteps from the 

69 original diffusion process to retain. 

70 :param kwargs: the kwargs to create the base diffusion process. 

71 """ 

72 

73 def __init__(self, use_timesteps, **kwargs): 

74 self.use_timesteps = set(use_timesteps) 

75 self.timestep_map = [] 

76 self.original_num_steps = len(kwargs["betas"]) 

77 

78 base_diffusion = GaussianDiffusion(**kwargs) # pylint: disable=missing-kwoa 

79 last_alpha_cumprod = 1.0 

80 new_betas = [] 

81 for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod): 

82 if i in self.use_timesteps: 

83 new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) 

84 last_alpha_cumprod = alpha_cumprod 

85 self.timestep_map.append(i) 

86 kwargs["betas"] = np.array(new_betas) 

87 super().__init__(**kwargs) 

88 

89 def p_mean_variance( 

90 self, model, *args, **kwargs 

91 ): # pylint: disable=signature-differs 

92 return super().p_mean_variance(self._wrap_model(model), *args, **kwargs) 

93 

94 def training_losses( 

95 self, model, *args, **kwargs 

96 ): # pylint: disable=signature-differs 

97 return super().training_losses(self._wrap_model(model), *args, **kwargs) 

98 

99 def condition_mean(self, cond_fn, *args, **kwargs): 

100 return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs) 

101 

102 def condition_score(self, cond_fn, *args, **kwargs): 

103 return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs) 

104 

105 def _wrap_model(self, model): 

106 if isinstance(model, _WrappedModel): 

107 return model 

108 return _WrappedModel( 

109 model, self.timestep_map, self.original_num_steps 

110 ) 

111 

112 def _scale_timesteps(self, t): 

113 # Scaling is done by the wrapped model. 

114 return t 

115 

116 

117class _WrappedModel: 

118 def __init__(self, model, timestep_map, original_num_steps): 

119 self.model = model 

120 self.timestep_map = timestep_map 

121 # self.rescale_timesteps = rescale_timesteps 

122 self.original_num_steps = original_num_steps 

123 

124 def __call__(self, x, ts, **kwargs): 

125 map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype) 

126 new_ts = map_tensor[ts] 

127 # if self.rescale_timesteps: 

128 # new_ts = new_ts.float() * (1000.0 / self.original_num_steps) 

129 return self.model(x, new_ts, **kwargs)