fffiloni commited on
Commit
08926b4
1 Parent(s): 1469496

Create util.py

Browse files
Files changed (1) hide show
  1. utils/util.py +283 -0
utils/util.py ADDED
@@ -0,0 +1,283 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # adopted from
2
+ # https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
3
+ # and
4
+ # https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
5
+ # and
6
+ # https://github.com/openai/guided-diffusion/blob/0ba878e517b276c45d1195eb29f6f5f72659a05b/guided_diffusion/nn.py
7
+ #
8
+ # thanks!
9
+ import importlib
10
+
11
+ import os
12
+ import math
13
+ import torch
14
+ import torch.nn as nn
15
+ import numpy as np
16
+ from einops import repeat
17
+
18
+
19
+ def instantiate_from_config(config):
20
+ if not "target" in config:
21
+ if config == '__is_first_stage__':
22
+ return None
23
+ elif config == "__is_unconditional__":
24
+ return None
25
+ raise KeyError("Expected key `target` to instantiate.")
26
+ return get_obj_from_str(config["target"])(**config.get("params", dict()))
27
+
28
+
29
+ def get_obj_from_str(string, reload=False):
30
+ module, cls = string.rsplit(".", 1)
31
+ if reload:
32
+ module_imp = importlib.import_module(module)
33
+ importlib.reload(module_imp)
34
+ return getattr(importlib.import_module(module, package=None), cls)
35
+
36
+
37
+ def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
38
+ if schedule == "linear":
39
+ betas = (
40
+ torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2
41
+ )
42
+
43
+ elif schedule == "cosine":
44
+ timesteps = (
45
+ torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s
46
+ )
47
+ alphas = timesteps / (1 + cosine_s) * np.pi / 2
48
+ alphas = torch.cos(alphas).pow(2)
49
+ alphas = alphas / alphas[0]
50
+ betas = 1 - alphas[1:] / alphas[:-1]
51
+ betas = np.clip(betas, a_min=0, a_max=0.999)
52
+
53
+ elif schedule == "sqrt_linear":
54
+ betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64)
55
+ elif schedule == "sqrt":
56
+ betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5
57
+ else:
58
+ raise ValueError(f"schedule '{schedule}' unknown.")
59
+ return betas.numpy()
60
+
61
+
62
+ def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True):
63
+ if ddim_discr_method == 'uniform':
64
+ c = num_ddpm_timesteps // num_ddim_timesteps
65
+ ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))
66
+ elif ddim_discr_method == 'quad':
67
+ ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int)
68
+ else:
69
+ raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"')
70
+
71
+ # assert ddim_timesteps.shape[0] == num_ddim_timesteps
72
+ # add one to get the final alpha values right (the ones from first scale to data during sampling)
73
+ steps_out = ddim_timesteps + 1
74
+ if verbose:
75
+ print(f'Selected timesteps for ddim sampler: {steps_out}')
76
+ return steps_out
77
+
78
+
79
+ def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=True):
80
+ # select alphas for computing the variance schedule
81
+ alphas = alphacums[ddim_timesteps]
82
+ alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
83
+
84
+ # according the the formula provided in https://arxiv.org/abs/2010.02502
85
+ sigmas = eta * np.sqrt((1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev))
86
+ if verbose:
87
+ print(f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}')
88
+ print(f'For the chosen value of eta, which is {eta}, '
89
+ f'this results in the following sigma_t schedule for ddim sampler {sigmas}')
90
+ return sigmas, alphas, alphas_prev
91
+
92
+
93
+ def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
94
+ """
95
+ Create a beta schedule that discretizes the given alpha_t_bar function,
96
+ which defines the cumulative product of (1-beta) over time from t = [0,1].
97
+ :param num_diffusion_timesteps: the number of betas to produce.
98
+ :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
99
+ produces the cumulative product of (1-beta) up to that
100
+ part of the diffusion process.
101
+ :param max_beta: the maximum beta to use; use values lower than 1 to
102
+ prevent singularities.
103
+ """
104
+ betas = []
105
+ for i in range(num_diffusion_timesteps):
106
+ t1 = i / num_diffusion_timesteps
107
+ t2 = (i + 1) / num_diffusion_timesteps
108
+ betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
109
+ return np.array(betas)
110
+
111
+
112
+ def extract_into_tensor(a, t, x_shape):
113
+ b, *_ = t.shape
114
+ out = a.gather(-1, t)
115
+ return out.reshape(b, *((1,) * (len(x_shape) - 1)))
116
+
117
+
118
+ def checkpoint(func, inputs, params, flag):
119
+ """
120
+ Evaluate a function without caching intermediate activations, allowing for
121
+ reduced memory at the expense of extra compute in the backward pass.
122
+ :param func: the function to evaluate.
123
+ :param inputs: the argument sequence to pass to `func`.
124
+ :param params: a sequence of parameters `func` depends on but does not
125
+ explicitly take as arguments.
126
+ :param flag: if False, disable gradient checkpointing.
127
+ """
128
+ if flag:
129
+ args = tuple(inputs) + tuple(params)
130
+ return CheckpointFunction.apply(func, len(inputs), *args)
131
+ else:
132
+ return func(*inputs)
133
+
134
+
135
+ class CheckpointFunction(torch.autograd.Function):
136
+ @staticmethod
137
+ def forward(ctx, run_function, length, *args):
138
+ ctx.run_function = run_function
139
+ ctx.input_tensors = list(args[:length])
140
+ ctx.input_params = list(args[length:])
141
+
142
+ with torch.no_grad():
143
+ output_tensors = ctx.run_function(*ctx.input_tensors)
144
+ return output_tensors
145
+
146
+ @staticmethod
147
+ def backward(ctx, *output_grads):
148
+ ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
149
+ with torch.enable_grad():
150
+ # Fixes a bug where the first op in run_function modifies the
151
+ # Tensor storage in place, which is not allowed for detach()'d
152
+ # Tensors.
153
+ shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
154
+ output_tensors = ctx.run_function(*shallow_copies)
155
+ input_grads = torch.autograd.grad(
156
+ output_tensors,
157
+ ctx.input_tensors + ctx.input_params,
158
+ output_grads,
159
+ allow_unused=True,
160
+ )
161
+ del ctx.input_tensors
162
+ del ctx.input_params
163
+ del output_tensors
164
+ return (None, None) + input_grads
165
+
166
+
167
+ def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
168
+ """
169
+ Create sinusoidal timestep embeddings.
170
+ :param timesteps: a 1-D Tensor of N indices, one per batch element.
171
+ These may be fractional.
172
+ :param dim: the dimension of the output.
173
+ :param max_period: controls the minimum frequency of the embeddings.
174
+ :return: an [N x dim] Tensor of positional embeddings.
175
+ """
176
+ if not repeat_only:
177
+ half = dim // 2
178
+ freqs = torch.exp(
179
+ -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
180
+ ).to(device=timesteps.device)
181
+ args = timesteps[:, None].float() * freqs[None]
182
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
183
+ if dim % 2:
184
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
185
+ else:
186
+ embedding = repeat(timesteps, 'b -> b d', d=dim)
187
+ return embedding
188
+
189
+
190
+ def zero_module(module):
191
+ """
192
+ Zero out the parameters of a module and return it.
193
+ """
194
+ for p in module.parameters():
195
+ p.detach().zero_()
196
+ return module
197
+
198
+
199
+ def scale_module(module, scale):
200
+ """
201
+ Scale the parameters of a module and return it.
202
+ """
203
+ for p in module.parameters():
204
+ p.detach().mul_(scale)
205
+ return module
206
+
207
+
208
+ def mean_flat(tensor):
209
+ """
210
+ Take the mean over all non-batch dimensions.
211
+ """
212
+ return tensor.mean(dim=list(range(1, len(tensor.shape))))
213
+
214
+
215
+ def normalization(channels):
216
+ """
217
+ Make a standard normalization layer.
218
+ :param channels: number of input channels.
219
+ :return: an nn.Module for normalization.
220
+ """
221
+ return GroupNorm32(32, channels)
222
+
223
+
224
+ # PyTorch 1.7 has SiLU, but we support PyTorch 1.5.
225
+ class SiLU(nn.Module):
226
+ def forward(self, x):
227
+ return x * torch.sigmoid(x)
228
+
229
+
230
+ class GroupNorm32(nn.GroupNorm):
231
+ def forward(self, x):
232
+ return super().forward(x.float()).type(x.dtype)
233
+
234
+ def conv_nd(dims, *args, **kwargs):
235
+ """
236
+ Create a 1D, 2D, or 3D convolution module.
237
+ """
238
+ if dims == 1:
239
+ return nn.Conv1d(*args, **kwargs)
240
+ elif dims == 2:
241
+ return nn.Conv2d(*args, **kwargs)
242
+ elif dims == 3:
243
+ return nn.Conv3d(*args, **kwargs)
244
+ raise ValueError(f"unsupported dimensions: {dims}")
245
+
246
+
247
+ def linear(*args, **kwargs):
248
+ """
249
+ Create a linear module.
250
+ """
251
+ return nn.Linear(*args, **kwargs)
252
+
253
+
254
+ def avg_pool_nd(dims, *args, **kwargs):
255
+ """
256
+ Create a 1D, 2D, or 3D average pooling module.
257
+ """
258
+ if dims == 1:
259
+ return nn.AvgPool1d(*args, **kwargs)
260
+ elif dims == 2:
261
+ return nn.AvgPool2d(*args, **kwargs)
262
+ elif dims == 3:
263
+ return nn.AvgPool3d(*args, **kwargs)
264
+ raise ValueError(f"unsupported dimensions: {dims}")
265
+
266
+
267
+ class HybridConditioner(nn.Module):
268
+
269
+ def __init__(self, c_concat_config, c_crossattn_config):
270
+ super().__init__()
271
+ self.concat_conditioner = instantiate_from_config(c_concat_config)
272
+ self.crossattn_conditioner = instantiate_from_config(c_crossattn_config)
273
+
274
+ def forward(self, c_concat, c_crossattn):
275
+ c_concat = self.concat_conditioner(c_concat)
276
+ c_crossattn = self.crossattn_conditioner(c_crossattn)
277
+ return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]}
278
+
279
+
280
+ def noise_like(shape, device, repeat=False):
281
+ repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
282
+ noise = lambda: torch.randn(shape, device=device)
283
+ return repeat_noise() if repeat else noise()