anonymous commited on
Commit
4fcfd85
1 Parent(s): fd216ce

add controlnet

Browse files
ControlNet/ldm/data/__init__.py ADDED
File without changes
ControlNet/ldm/data/util.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+ from ldm.modules.midas.api import load_midas_transform
4
+
5
+
6
+ class AddMiDaS(object):
7
+ def __init__(self, model_type):
8
+ super().__init__()
9
+ self.transform = load_midas_transform(model_type)
10
+
11
+ def pt2np(self, x):
12
+ x = ((x + 1.0) * .5).detach().cpu().numpy()
13
+ return x
14
+
15
+ def np2pt(self, x):
16
+ x = torch.from_numpy(x) * 2 - 1.
17
+ return x
18
+
19
+ def __call__(self, sample):
20
+ # sample['jpg'] is tensor hwc in [-1, 1] at this point
21
+ x = self.pt2np(sample['jpg'])
22
+ x = self.transform({"image": x})["image"]
23
+ sample['midas_in'] = x
24
+ return sample
ControlNet/ldm/models/autoencoder.py ADDED
@@ -0,0 +1,219 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import pytorch_lightning as pl
3
+ import torch.nn.functional as F
4
+ from contextlib import contextmanager
5
+
6
+ from ldm.modules.diffusionmodules.model import Encoder, Decoder
7
+ from ldm.modules.distributions.distributions import DiagonalGaussianDistribution
8
+
9
+ from ldm.util import instantiate_from_config
10
+ from ldm.modules.ema import LitEma
11
+
12
+
13
+ class AutoencoderKL(pl.LightningModule):
14
+ def __init__(self,
15
+ ddconfig,
16
+ lossconfig,
17
+ embed_dim,
18
+ ckpt_path=None,
19
+ ignore_keys=[],
20
+ image_key="image",
21
+ colorize_nlabels=None,
22
+ monitor=None,
23
+ ema_decay=None,
24
+ learn_logvar=False
25
+ ):
26
+ super().__init__()
27
+ self.learn_logvar = learn_logvar
28
+ self.image_key = image_key
29
+ self.encoder = Encoder(**ddconfig)
30
+ self.decoder = Decoder(**ddconfig)
31
+ self.loss = instantiate_from_config(lossconfig)
32
+ assert ddconfig["double_z"]
33
+ self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1)
34
+ self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
35
+ self.embed_dim = embed_dim
36
+ if colorize_nlabels is not None:
37
+ assert type(colorize_nlabels)==int
38
+ self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
39
+ if monitor is not None:
40
+ self.monitor = monitor
41
+
42
+ self.use_ema = ema_decay is not None
43
+ if self.use_ema:
44
+ self.ema_decay = ema_decay
45
+ assert 0. < ema_decay < 1.
46
+ self.model_ema = LitEma(self, decay=ema_decay)
47
+ print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
48
+
49
+ if ckpt_path is not None:
50
+ self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
51
+
52
+ def init_from_ckpt(self, path, ignore_keys=list()):
53
+ sd = torch.load(path, map_location="cpu")["state_dict"]
54
+ keys = list(sd.keys())
55
+ for k in keys:
56
+ for ik in ignore_keys:
57
+ if k.startswith(ik):
58
+ print("Deleting key {} from state_dict.".format(k))
59
+ del sd[k]
60
+ self.load_state_dict(sd, strict=False)
61
+ print(f"Restored from {path}")
62
+
63
+ @contextmanager
64
+ def ema_scope(self, context=None):
65
+ if self.use_ema:
66
+ self.model_ema.store(self.parameters())
67
+ self.model_ema.copy_to(self)
68
+ if context is not None:
69
+ print(f"{context}: Switched to EMA weights")
70
+ try:
71
+ yield None
72
+ finally:
73
+ if self.use_ema:
74
+ self.model_ema.restore(self.parameters())
75
+ if context is not None:
76
+ print(f"{context}: Restored training weights")
77
+
78
+ def on_train_batch_end(self, *args, **kwargs):
79
+ if self.use_ema:
80
+ self.model_ema(self)
81
+
82
+ def encode(self, x):
83
+ h = self.encoder(x)
84
+ moments = self.quant_conv(h)
85
+ posterior = DiagonalGaussianDistribution(moments)
86
+ return posterior
87
+
88
+ def decode(self, z):
89
+ z = self.post_quant_conv(z)
90
+ dec = self.decoder(z)
91
+ return dec
92
+
93
+ def forward(self, input, sample_posterior=True):
94
+ posterior = self.encode(input)
95
+ if sample_posterior:
96
+ z = posterior.sample()
97
+ else:
98
+ z = posterior.mode()
99
+ dec = self.decode(z)
100
+ return dec, posterior
101
+
102
+ def get_input(self, batch, k):
103
+ x = batch[k]
104
+ if len(x.shape) == 3:
105
+ x = x[..., None]
106
+ x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
107
+ return x
108
+
109
+ def training_step(self, batch, batch_idx, optimizer_idx):
110
+ inputs = self.get_input(batch, self.image_key)
111
+ reconstructions, posterior = self(inputs)
112
+
113
+ if optimizer_idx == 0:
114
+ # train encoder+decoder+logvar
115
+ aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
116
+ last_layer=self.get_last_layer(), split="train")
117
+ self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
118
+ self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False)
119
+ return aeloss
120
+
121
+ if optimizer_idx == 1:
122
+ # train the discriminator
123
+ discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
124
+ last_layer=self.get_last_layer(), split="train")
125
+
126
+ self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
127
+ self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False)
128
+ return discloss
129
+
130
+ def validation_step(self, batch, batch_idx):
131
+ log_dict = self._validation_step(batch, batch_idx)
132
+ with self.ema_scope():
133
+ log_dict_ema = self._validation_step(batch, batch_idx, postfix="_ema")
134
+ return log_dict
135
+
136
+ def _validation_step(self, batch, batch_idx, postfix=""):
137
+ inputs = self.get_input(batch, self.image_key)
138
+ reconstructions, posterior = self(inputs)
139
+ aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step,
140
+ last_layer=self.get_last_layer(), split="val"+postfix)
141
+
142
+ discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step,
143
+ last_layer=self.get_last_layer(), split="val"+postfix)
144
+
145
+ self.log(f"val{postfix}/rec_loss", log_dict_ae[f"val{postfix}/rec_loss"])
146
+ self.log_dict(log_dict_ae)
147
+ self.log_dict(log_dict_disc)
148
+ return self.log_dict
149
+
150
+ def configure_optimizers(self):
151
+ lr = self.learning_rate
152
+ ae_params_list = list(self.encoder.parameters()) + list(self.decoder.parameters()) + list(
153
+ self.quant_conv.parameters()) + list(self.post_quant_conv.parameters())
154
+ if self.learn_logvar:
155
+ print(f"{self.__class__.__name__}: Learning logvar")
156
+ ae_params_list.append(self.loss.logvar)
157
+ opt_ae = torch.optim.Adam(ae_params_list,
158
+ lr=lr, betas=(0.5, 0.9))
159
+ opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
160
+ lr=lr, betas=(0.5, 0.9))
161
+ return [opt_ae, opt_disc], []
162
+
163
+ def get_last_layer(self):
164
+ return self.decoder.conv_out.weight
165
+
166
+ @torch.no_grad()
167
+ def log_images(self, batch, only_inputs=False, log_ema=False, **kwargs):
168
+ log = dict()
169
+ x = self.get_input(batch, self.image_key)
170
+ x = x.to(self.device)
171
+ if not only_inputs:
172
+ xrec, posterior = self(x)
173
+ if x.shape[1] > 3:
174
+ # colorize with random projection
175
+ assert xrec.shape[1] > 3
176
+ x = self.to_rgb(x)
177
+ xrec = self.to_rgb(xrec)
178
+ log["samples"] = self.decode(torch.randn_like(posterior.sample()))
179
+ log["reconstructions"] = xrec
180
+ if log_ema or self.use_ema:
181
+ with self.ema_scope():
182
+ xrec_ema, posterior_ema = self(x)
183
+ if x.shape[1] > 3:
184
+ # colorize with random projection
185
+ assert xrec_ema.shape[1] > 3
186
+ xrec_ema = self.to_rgb(xrec_ema)
187
+ log["samples_ema"] = self.decode(torch.randn_like(posterior_ema.sample()))
188
+ log["reconstructions_ema"] = xrec_ema
189
+ log["inputs"] = x
190
+ return log
191
+
192
+ def to_rgb(self, x):
193
+ assert self.image_key == "segmentation"
194
+ if not hasattr(self, "colorize"):
195
+ self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
196
+ x = F.conv2d(x, weight=self.colorize)
197
+ x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
198
+ return x
199
+
200
+
201
+ class IdentityFirstStage(torch.nn.Module):
202
+ def __init__(self, *args, vq_interface=False, **kwargs):
203
+ self.vq_interface = vq_interface
204
+ super().__init__()
205
+
206
+ def encode(self, x, *args, **kwargs):
207
+ return x
208
+
209
+ def decode(self, x, *args, **kwargs):
210
+ return x
211
+
212
+ def quantize(self, x, *args, **kwargs):
213
+ if self.vq_interface:
214
+ return x, None, [None, None, None]
215
+ return x
216
+
217
+ def forward(self, x, *args, **kwargs):
218
+ return x
219
+
ControlNet/ldm/models/diffusion/__init__.py ADDED
File without changes
ControlNet/ldm/models/diffusion/ddim.py ADDED
@@ -0,0 +1,336 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SAMPLING ONLY."""
2
+
3
+ import torch
4
+ import numpy as np
5
+ from tqdm import tqdm
6
+
7
+ from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like, extract_into_tensor
8
+
9
+
10
+ class DDIMSampler(object):
11
+ def __init__(self, model, schedule="linear", **kwargs):
12
+ super().__init__()
13
+ self.model = model
14
+ self.ddpm_num_timesteps = model.num_timesteps
15
+ self.schedule = schedule
16
+
17
+ def register_buffer(self, name, attr):
18
+ if type(attr) == torch.Tensor:
19
+ if attr.device != torch.device("cuda"):
20
+ attr = attr.to(torch.device("cuda"))
21
+ setattr(self, name, attr)
22
+
23
+ def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
24
+ self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
25
+ num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose)
26
+ alphas_cumprod = self.model.alphas_cumprod
27
+ assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
28
+ to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)
29
+
30
+ self.register_buffer('betas', to_torch(self.model.betas))
31
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
32
+ self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
33
+
34
+ # calculations for diffusion q(x_t | x_{t-1}) and others
35
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
36
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
37
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
38
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
39
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
40
+
41
+ # ddim sampling parameters
42
+ ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
43
+ ddim_timesteps=self.ddim_timesteps,
44
+ eta=ddim_eta,verbose=verbose)
45
+ self.register_buffer('ddim_sigmas', ddim_sigmas)
46
+ self.register_buffer('ddim_alphas', ddim_alphas)
47
+ self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
48
+ self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
49
+ sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
50
+ (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
51
+ 1 - self.alphas_cumprod / self.alphas_cumprod_prev))
52
+ self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)
53
+
54
+ @torch.no_grad()
55
+ def sample(self,
56
+ S,
57
+ batch_size,
58
+ shape,
59
+ conditioning=None,
60
+ callback=None,
61
+ normals_sequence=None,
62
+ img_callback=None,
63
+ quantize_x0=False,
64
+ eta=0.,
65
+ mask=None,
66
+ x0=None,
67
+ temperature=1.,
68
+ noise_dropout=0.,
69
+ score_corrector=None,
70
+ corrector_kwargs=None,
71
+ verbose=True,
72
+ x_T=None,
73
+ log_every_t=100,
74
+ unconditional_guidance_scale=1.,
75
+ unconditional_conditioning=None, # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
76
+ dynamic_threshold=None,
77
+ ucg_schedule=None,
78
+ **kwargs
79
+ ):
80
+ if conditioning is not None:
81
+ if isinstance(conditioning, dict):
82
+ ctmp = conditioning[list(conditioning.keys())[0]]
83
+ while isinstance(ctmp, list): ctmp = ctmp[0]
84
+ cbs = ctmp.shape[0]
85
+ if cbs != batch_size:
86
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
87
+
88
+ elif isinstance(conditioning, list):
89
+ for ctmp in conditioning:
90
+ if ctmp.shape[0] != batch_size:
91
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
92
+
93
+ else:
94
+ if conditioning.shape[0] != batch_size:
95
+ print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
96
+
97
+ self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
98
+ # sampling
99
+ C, H, W = shape
100
+ size = (batch_size, C, H, W)
101
+ print(f'Data shape for DDIM sampling is {size}, eta {eta}')
102
+
103
+ samples, intermediates = self.ddim_sampling(conditioning, size,
104
+ callback=callback,
105
+ img_callback=img_callback,
106
+ quantize_denoised=quantize_x0,
107
+ mask=mask, x0=x0,
108
+ ddim_use_original_steps=False,
109
+ noise_dropout=noise_dropout,
110
+ temperature=temperature,
111
+ score_corrector=score_corrector,
112
+ corrector_kwargs=corrector_kwargs,
113
+ x_T=x_T,
114
+ log_every_t=log_every_t,
115
+ unconditional_guidance_scale=unconditional_guidance_scale,
116
+ unconditional_conditioning=unconditional_conditioning,
117
+ dynamic_threshold=dynamic_threshold,
118
+ ucg_schedule=ucg_schedule
119
+ )
120
+ return samples, intermediates
121
+
122
+ @torch.no_grad()
123
+ def ddim_sampling(self, cond, shape,
124
+ x_T=None, ddim_use_original_steps=False,
125
+ callback=None, timesteps=None, quantize_denoised=False,
126
+ mask=None, x0=None, img_callback=None, log_every_t=100,
127
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
128
+ unconditional_guidance_scale=1., unconditional_conditioning=None, dynamic_threshold=None,
129
+ ucg_schedule=None):
130
+ device = self.model.betas.device
131
+ b = shape[0]
132
+ if x_T is None:
133
+ img = torch.randn(shape, device=device)
134
+ else:
135
+ img = x_T
136
+
137
+ if timesteps is None:
138
+ timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
139
+ elif timesteps is not None and not ddim_use_original_steps:
140
+ subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
141
+ timesteps = self.ddim_timesteps[:subset_end]
142
+
143
+ intermediates = {'x_inter': [img], 'pred_x0': [img]}
144
+ time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else np.flip(timesteps)
145
+ total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
146
+ print(f"Running DDIM Sampling with {total_steps} timesteps")
147
+
148
+ iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
149
+
150
+ for i, step in enumerate(iterator):
151
+ index = total_steps - i - 1
152
+ ts = torch.full((b,), step, device=device, dtype=torch.long)
153
+
154
+ if mask is not None:
155
+ assert x0 is not None
156
+ img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass?
157
+ img = img_orig * mask + (1. - mask) * img
158
+
159
+ if ucg_schedule is not None:
160
+ assert len(ucg_schedule) == len(time_range)
161
+ unconditional_guidance_scale = ucg_schedule[i]
162
+
163
+ outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
164
+ quantize_denoised=quantize_denoised, temperature=temperature,
165
+ noise_dropout=noise_dropout, score_corrector=score_corrector,
166
+ corrector_kwargs=corrector_kwargs,
167
+ unconditional_guidance_scale=unconditional_guidance_scale,
168
+ unconditional_conditioning=unconditional_conditioning,
169
+ dynamic_threshold=dynamic_threshold)
170
+ img, pred_x0 = outs
171
+ if callback: callback(i)
172
+ if img_callback: img_callback(pred_x0, i)
173
+
174
+ if index % log_every_t == 0 or index == total_steps - 1:
175
+ intermediates['x_inter'].append(img)
176
+ intermediates['pred_x0'].append(pred_x0)
177
+
178
+ return img, intermediates
179
+
180
+ @torch.no_grad()
181
+ def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
182
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
183
+ unconditional_guidance_scale=1., unconditional_conditioning=None,
184
+ dynamic_threshold=None):
185
+ b, *_, device = *x.shape, x.device
186
+
187
+ if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
188
+ model_output = self.model.apply_model(x, t, c)
189
+ else:
190
+ x_in = torch.cat([x] * 2)
191
+ t_in = torch.cat([t] * 2)
192
+ if isinstance(c, dict):
193
+ assert isinstance(unconditional_conditioning, dict)
194
+ c_in = dict()
195
+ for k in c:
196
+ if isinstance(c[k], list):
197
+ c_in[k] = [torch.cat([
198
+ unconditional_conditioning[k][i],
199
+ c[k][i]]) for i in range(len(c[k]))]
200
+ else:
201
+ c_in[k] = torch.cat([
202
+ unconditional_conditioning[k],
203
+ c[k]])
204
+ elif isinstance(c, list):
205
+ c_in = list()
206
+ assert isinstance(unconditional_conditioning, list)
207
+ for i in range(len(c)):
208
+ c_in.append(torch.cat([unconditional_conditioning[i], c[i]]))
209
+ else:
210
+ c_in = torch.cat([unconditional_conditioning, c])
211
+ model_uncond, model_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
212
+ model_output = model_uncond + unconditional_guidance_scale * (model_t - model_uncond)
213
+
214
+ if self.model.parameterization == "v":
215
+ e_t = self.model.predict_eps_from_z_and_v(x, t, model_output)
216
+ else:
217
+ e_t = model_output
218
+
219
+ if score_corrector is not None:
220
+ assert self.model.parameterization == "eps", 'not implemented'
221
+ e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
222
+
223
+ alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
224
+ alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
225
+ sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
226
+ sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
227
+ # select parameters corresponding to the currently considered timestep
228
+ a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
229
+ a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
230
+ sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
231
+ sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
232
+
233
+ # current prediction for x_0
234
+ if self.model.parameterization != "v":
235
+ pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
236
+ else:
237
+ pred_x0 = self.model.predict_start_from_z_and_v(x, t, model_output)
238
+
239
+ if quantize_denoised:
240
+ pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
241
+
242
+ if dynamic_threshold is not None:
243
+ raise NotImplementedError()
244
+
245
+ # direction pointing to x_t
246
+ dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
247
+ noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
248
+ if noise_dropout > 0.:
249
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
250
+ x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
251
+ return x_prev, pred_x0
252
+
253
+ @torch.no_grad()
254
+ def encode(self, x0, c, t_enc, use_original_steps=False, return_intermediates=None,
255
+ unconditional_guidance_scale=1.0, unconditional_conditioning=None, callback=None):
256
+ num_reference_steps = self.ddpm_num_timesteps if use_original_steps else self.ddim_timesteps.shape[0]
257
+
258
+ assert t_enc <= num_reference_steps
259
+ num_steps = t_enc
260
+
261
+ if use_original_steps:
262
+ alphas_next = self.alphas_cumprod[:num_steps]
263
+ alphas = self.alphas_cumprod_prev[:num_steps]
264
+ else:
265
+ alphas_next = self.ddim_alphas[:num_steps]
266
+ alphas = torch.tensor(self.ddim_alphas_prev[:num_steps])
267
+
268
+ x_next = x0
269
+ intermediates = []
270
+ inter_steps = []
271
+ for i in tqdm(range(num_steps), desc='Encoding Image'):
272
+ t = torch.full((x0.shape[0],), i, device=self.model.device, dtype=torch.long)
273
+ if unconditional_guidance_scale == 1.:
274
+ noise_pred = self.model.apply_model(x_next, t, c)
275
+ else:
276
+ assert unconditional_conditioning is not None
277
+ e_t_uncond, noise_pred = torch.chunk(
278
+ self.model.apply_model(torch.cat((x_next, x_next)), torch.cat((t, t)),
279
+ torch.cat((unconditional_conditioning, c))), 2)
280
+ noise_pred = e_t_uncond + unconditional_guidance_scale * (noise_pred - e_t_uncond)
281
+
282
+ xt_weighted = (alphas_next[i] / alphas[i]).sqrt() * x_next
283
+ weighted_noise_pred = alphas_next[i].sqrt() * (
284
+ (1 / alphas_next[i] - 1).sqrt() - (1 / alphas[i] - 1).sqrt()) * noise_pred
285
+ x_next = xt_weighted + weighted_noise_pred
286
+ if return_intermediates and i % (
287
+ num_steps // return_intermediates) == 0 and i < num_steps - 1:
288
+ intermediates.append(x_next)
289
+ inter_steps.append(i)
290
+ elif return_intermediates and i >= num_steps - 2:
291
+ intermediates.append(x_next)
292
+ inter_steps.append(i)
293
+ if callback: callback(i)
294
+
295
+ out = {'x_encoded': x_next, 'intermediate_steps': inter_steps}
296
+ if return_intermediates:
297
+ out.update({'intermediates': intermediates})
298
+ return x_next, out
299
+
300
+ @torch.no_grad()
301
+ def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):
302
+ # fast, but does not allow for exact reconstruction
303
+ # t serves as an index to gather the correct alphas
304
+ if use_original_steps:
305
+ sqrt_alphas_cumprod = self.sqrt_alphas_cumprod
306
+ sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod
307
+ else:
308
+ sqrt_alphas_cumprod = torch.sqrt(self.ddim_alphas)
309
+ sqrt_one_minus_alphas_cumprod = self.ddim_sqrt_one_minus_alphas
310
+
311
+ if noise is None:
312
+ noise = torch.randn_like(x0)
313
+ return (extract_into_tensor(sqrt_alphas_cumprod, t, x0.shape) * x0 +
314
+ extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x0.shape) * noise)
315
+
316
+ @torch.no_grad()
317
+ def decode(self, x_latent, cond, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None,
318
+ use_original_steps=False, callback=None):
319
+
320
+ timesteps = np.arange(self.ddpm_num_timesteps) if use_original_steps else self.ddim_timesteps
321
+ timesteps = timesteps[:t_start]
322
+
323
+ time_range = np.flip(timesteps)
324
+ total_steps = timesteps.shape[0]
325
+ print(f"Running DDIM Sampling with {total_steps} timesteps")
326
+
327
+ iterator = tqdm(time_range, desc='Decoding image', total=total_steps)
328
+ x_dec = x_latent
329
+ for i, step in enumerate(iterator):
330
+ index = total_steps - i - 1
331
+ ts = torch.full((x_latent.shape[0],), step, device=x_latent.device, dtype=torch.long)
332
+ x_dec, _ = self.p_sample_ddim(x_dec, cond, ts, index=index, use_original_steps=use_original_steps,
333
+ unconditional_guidance_scale=unconditional_guidance_scale,
334
+ unconditional_conditioning=unconditional_conditioning)
335
+ if callback: callback(i)
336
+ return x_dec
ControlNet/ldm/models/diffusion/ddpm.py ADDED
@@ -0,0 +1,1797 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ wild mixture of
3
+ https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
4
+ https://github.com/openai/improved-diffusion/blob/e94489283bb876ac1477d5dd7709bbbd2d9902ce/improved_diffusion/gaussian_diffusion.py
5
+ https://github.com/CompVis/taming-transformers
6
+ -- merci
7
+ """
8
+
9
+ import torch
10
+ import torch.nn as nn
11
+ import numpy as np
12
+ import pytorch_lightning as pl
13
+ from torch.optim.lr_scheduler import LambdaLR
14
+ from einops import rearrange, repeat
15
+ from contextlib import contextmanager, nullcontext
16
+ from functools import partial
17
+ import itertools
18
+ from tqdm import tqdm
19
+ from torchvision.utils import make_grid
20
+ from pytorch_lightning.utilities.distributed import rank_zero_only
21
+ from omegaconf import ListConfig
22
+
23
+ from ldm.util import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config
24
+ from ldm.modules.ema import LitEma
25
+ from ldm.modules.distributions.distributions import normal_kl, DiagonalGaussianDistribution
26
+ from ldm.models.autoencoder import IdentityFirstStage, AutoencoderKL
27
+ from ldm.modules.diffusionmodules.util import make_beta_schedule, extract_into_tensor, noise_like
28
+ from ldm.models.diffusion.ddim import DDIMSampler
29
+
30
+
31
+ __conditioning_keys__ = {'concat': 'c_concat',
32
+ 'crossattn': 'c_crossattn',
33
+ 'adm': 'y'}
34
+
35
+
36
+ def disabled_train(self, mode=True):
37
+ """Overwrite model.train with this function to make sure train/eval mode
38
+ does not change anymore."""
39
+ return self
40
+
41
+
42
+ def uniform_on_device(r1, r2, shape, device):
43
+ return (r1 - r2) * torch.rand(*shape, device=device) + r2
44
+
45
+
46
+ class DDPM(pl.LightningModule):
47
+ # classic DDPM with Gaussian diffusion, in image space
48
+ def __init__(self,
49
+ unet_config,
50
+ timesteps=1000,
51
+ beta_schedule="linear",
52
+ loss_type="l2",
53
+ ckpt_path=None,
54
+ ignore_keys=[],
55
+ load_only_unet=False,
56
+ monitor="val/loss",
57
+ use_ema=True,
58
+ first_stage_key="image",
59
+ image_size=256,
60
+ channels=3,
61
+ log_every_t=100,
62
+ clip_denoised=True,
63
+ linear_start=1e-4,
64
+ linear_end=2e-2,
65
+ cosine_s=8e-3,
66
+ given_betas=None,
67
+ original_elbo_weight=0.,
68
+ v_posterior=0., # weight for choosing posterior variance as sigma = (1-v) * beta_tilde + v * beta
69
+ l_simple_weight=1.,
70
+ conditioning_key=None,
71
+ parameterization="eps", # all assuming fixed variance schedules
72
+ scheduler_config=None,
73
+ use_positional_encodings=False,
74
+ learn_logvar=False,
75
+ logvar_init=0.,
76
+ make_it_fit=False,
77
+ ucg_training=None,
78
+ reset_ema=False,
79
+ reset_num_ema_updates=False,
80
+ ):
81
+ super().__init__()
82
+ assert parameterization in ["eps", "x0", "v"], 'currently only supporting "eps" and "x0" and "v"'
83
+ self.parameterization = parameterization
84
+ print(f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode")
85
+ self.cond_stage_model = None
86
+ self.clip_denoised = clip_denoised
87
+ self.log_every_t = log_every_t
88
+ self.first_stage_key = first_stage_key
89
+ self.image_size = image_size # try conv?
90
+ self.channels = channels
91
+ self.use_positional_encodings = use_positional_encodings
92
+ self.model = DiffusionWrapper(unet_config, conditioning_key)
93
+ count_params(self.model, verbose=True)
94
+ self.use_ema = use_ema
95
+ if self.use_ema:
96
+ self.model_ema = LitEma(self.model)
97
+ print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
98
+
99
+ self.use_scheduler = scheduler_config is not None
100
+ if self.use_scheduler:
101
+ self.scheduler_config = scheduler_config
102
+
103
+ self.v_posterior = v_posterior
104
+ self.original_elbo_weight = original_elbo_weight
105
+ self.l_simple_weight = l_simple_weight
106
+
107
+ if monitor is not None:
108
+ self.monitor = monitor
109
+ self.make_it_fit = make_it_fit
110
+ if reset_ema: assert exists(ckpt_path)
111
+ if ckpt_path is not None:
112
+ self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys, only_model=load_only_unet)
113
+ if reset_ema:
114
+ assert self.use_ema
115
+ print(f"Resetting ema to pure model weights. This is useful when restoring from an ema-only checkpoint.")
116
+ self.model_ema = LitEma(self.model)
117
+ if reset_num_ema_updates:
118
+ print(" +++++++++++ WARNING: RESETTING NUM_EMA UPDATES TO ZERO +++++++++++ ")
119
+ assert self.use_ema
120
+ self.model_ema.reset_num_updates()
121
+
122
+ self.register_schedule(given_betas=given_betas, beta_schedule=beta_schedule, timesteps=timesteps,
123
+ linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s)
124
+
125
+ self.loss_type = loss_type
126
+
127
+ self.learn_logvar = learn_logvar
128
+ logvar = torch.full(fill_value=logvar_init, size=(self.num_timesteps,))
129
+ if self.learn_logvar:
130
+ self.logvar = nn.Parameter(self.logvar, requires_grad=True)
131
+ else:
132
+ self.register_buffer('logvar', logvar)
133
+
134
+ self.ucg_training = ucg_training or dict()
135
+ if self.ucg_training:
136
+ self.ucg_prng = np.random.RandomState()
137
+
138
+ def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
139
+ linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
140
+ if exists(given_betas):
141
+ betas = given_betas
142
+ else:
143
+ betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end,
144
+ cosine_s=cosine_s)
145
+ alphas = 1. - betas
146
+ alphas_cumprod = np.cumprod(alphas, axis=0)
147
+ alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
148
+
149
+ timesteps, = betas.shape
150
+ self.num_timesteps = int(timesteps)
151
+ self.linear_start = linear_start
152
+ self.linear_end = linear_end
153
+ assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep'
154
+
155
+ to_torch = partial(torch.tensor, dtype=torch.float32)
156
+
157
+ self.register_buffer('betas', to_torch(betas))
158
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
159
+ self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
160
+
161
+ # calculations for diffusion q(x_t | x_{t-1}) and others
162
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
163
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
164
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
165
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
166
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
167
+
168
+ # calculations for posterior q(x_{t-1} | x_t, x_0)
169
+ posterior_variance = (1 - self.v_posterior) * betas * (1. - alphas_cumprod_prev) / (
170
+ 1. - alphas_cumprod) + self.v_posterior * betas
171
+ # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
172
+ self.register_buffer('posterior_variance', to_torch(posterior_variance))
173
+ # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
174
+ self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
175
+ self.register_buffer('posterior_mean_coef1', to_torch(
176
+ betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
177
+ self.register_buffer('posterior_mean_coef2', to_torch(
178
+ (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
179
+
180
+ if self.parameterization == "eps":
181
+ lvlb_weights = self.betas ** 2 / (
182
+ 2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod))
183
+ elif self.parameterization == "x0":
184
+ lvlb_weights = 0.5 * np.sqrt(torch.Tensor(alphas_cumprod)) / (2. * 1 - torch.Tensor(alphas_cumprod))
185
+ elif self.parameterization == "v":
186
+ lvlb_weights = torch.ones_like(self.betas ** 2 / (
187
+ 2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod)))
188
+ else:
189
+ raise NotImplementedError("mu not supported")
190
+ lvlb_weights[0] = lvlb_weights[1]
191
+ self.register_buffer('lvlb_weights', lvlb_weights, persistent=False)
192
+ assert not torch.isnan(self.lvlb_weights).all()
193
+
194
+ @contextmanager
195
+ def ema_scope(self, context=None):
196
+ if self.use_ema:
197
+ self.model_ema.store(self.model.parameters())
198
+ self.model_ema.copy_to(self.model)
199
+ if context is not None:
200
+ print(f"{context}: Switched to EMA weights")
201
+ try:
202
+ yield None
203
+ finally:
204
+ if self.use_ema:
205
+ self.model_ema.restore(self.model.parameters())
206
+ if context is not None:
207
+ print(f"{context}: Restored training weights")
208
+
209
+ @torch.no_grad()
210
+ def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
211
+ sd = torch.load(path, map_location="cpu")
212
+ if "state_dict" in list(sd.keys()):
213
+ sd = sd["state_dict"]
214
+ keys = list(sd.keys())
215
+ for k in keys:
216
+ for ik in ignore_keys:
217
+ if k.startswith(ik):
218
+ print("Deleting key {} from state_dict.".format(k))
219
+ del sd[k]
220
+ if self.make_it_fit:
221
+ n_params = len([name for name, _ in
222
+ itertools.chain(self.named_parameters(),
223
+ self.named_buffers())])
224
+ for name, param in tqdm(
225
+ itertools.chain(self.named_parameters(),
226
+ self.named_buffers()),
227
+ desc="Fitting old weights to new weights",
228
+ total=n_params
229
+ ):
230
+ if not name in sd:
231
+ continue
232
+ old_shape = sd[name].shape
233
+ new_shape = param.shape
234
+ assert len(old_shape) == len(new_shape)
235
+ if len(new_shape) > 2:
236
+ # we only modify first two axes
237
+ assert new_shape[2:] == old_shape[2:]
238
+ # assumes first axis corresponds to output dim
239
+ if not new_shape == old_shape:
240
+ new_param = param.clone()
241
+ old_param = sd[name]
242
+ if len(new_shape) == 1:
243
+ for i in range(new_param.shape[0]):
244
+ new_param[i] = old_param[i % old_shape[0]]
245
+ elif len(new_shape) >= 2:
246
+ for i in range(new_param.shape[0]):
247
+ for j in range(new_param.shape[1]):
248
+ new_param[i, j] = old_param[i % old_shape[0], j % old_shape[1]]
249
+
250
+ n_used_old = torch.ones(old_shape[1])
251
+ for j in range(new_param.shape[1]):
252
+ n_used_old[j % old_shape[1]] += 1
253
+ n_used_new = torch.zeros(new_shape[1])
254
+ for j in range(new_param.shape[1]):
255
+ n_used_new[j] = n_used_old[j % old_shape[1]]
256
+
257
+ n_used_new = n_used_new[None, :]
258
+ while len(n_used_new.shape) < len(new_shape):
259
+ n_used_new = n_used_new.unsqueeze(-1)
260
+ new_param /= n_used_new
261
+
262
+ sd[name] = new_param
263
+
264
+ missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict(
265
+ sd, strict=False)
266
+ print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
267
+ if len(missing) > 0:
268
+ print(f"Missing Keys:\n {missing}")
269
+ if len(unexpected) > 0:
270
+ print(f"\nUnexpected Keys:\n {unexpected}")
271
+
272
+ def q_mean_variance(self, x_start, t):
273
+ """
274
+ Get the distribution q(x_t | x_0).
275
+ :param x_start: the [N x C x ...] tensor of noiseless inputs.
276
+ :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
277
+ :return: A tuple (mean, variance, log_variance), all of x_start's shape.
278
+ """
279
+ mean = (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start)
280
+ variance = extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
281
+ log_variance = extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape)
282
+ return mean, variance, log_variance
283
+
284
+ def predict_start_from_noise(self, x_t, t, noise):
285
+ return (
286
+ extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
287
+ extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
288
+ )
289
+
290
+ def predict_start_from_z_and_v(self, x_t, t, v):
291
+ # self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
292
+ # self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
293
+ return (
294
+ extract_into_tensor(self.sqrt_alphas_cumprod, t, x_t.shape) * x_t -
295
+ extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_t.shape) * v
296
+ )
297
+
298
+ def predict_eps_from_z_and_v(self, x_t, t, v):
299
+ return (
300
+ extract_into_tensor(self.sqrt_alphas_cumprod, t, x_t.shape) * v +
301
+ extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_t.shape) * x_t
302
+ )
303
+
304
+ def q_posterior(self, x_start, x_t, t):
305
+ posterior_mean = (
306
+ extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start +
307
+ extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
308
+ )
309
+ posterior_variance = extract_into_tensor(self.posterior_variance, t, x_t.shape)
310
+ posterior_log_variance_clipped = extract_into_tensor(self.posterior_log_variance_clipped, t, x_t.shape)
311
+ return posterior_mean, posterior_variance, posterior_log_variance_clipped
312
+
313
+ def p_mean_variance(self, x, t, clip_denoised: bool):
314
+ model_out = self.model(x, t)
315
+ if self.parameterization == "eps":
316
+ x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
317
+ elif self.parameterization == "x0":
318
+ x_recon = model_out
319
+ if clip_denoised:
320
+ x_recon.clamp_(-1., 1.)
321
+
322
+ model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
323
+ return model_mean, posterior_variance, posterior_log_variance
324
+
325
+ @torch.no_grad()
326
+ def p_sample(self, x, t, clip_denoised=True, repeat_noise=False):
327
+ b, *_, device = *x.shape, x.device
328
+ model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=clip_denoised)
329
+ noise = noise_like(x.shape, device, repeat_noise)
330
+ # no noise when t == 0
331
+ nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
332
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
333
+
334
+ @torch.no_grad()
335
+ def p_sample_loop(self, shape, return_intermediates=False):
336
+ device = self.betas.device
337
+ b = shape[0]
338
+ img = torch.randn(shape, device=device)
339
+ intermediates = [img]
340
+ for i in tqdm(reversed(range(0, self.num_timesteps)), desc='Sampling t', total=self.num_timesteps):
341
+ img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long),
342
+ clip_denoised=self.clip_denoised)
343
+ if i % self.log_every_t == 0 or i == self.num_timesteps - 1:
344
+ intermediates.append(img)
345
+ if return_intermediates:
346
+ return img, intermediates
347
+ return img
348
+
349
+ @torch.no_grad()
350
+ def sample(self, batch_size=16, return_intermediates=False):
351
+ image_size = self.image_size
352
+ channels = self.channels
353
+ return self.p_sample_loop((batch_size, channels, image_size, image_size),
354
+ return_intermediates=return_intermediates)
355
+
356
+ def q_sample(self, x_start, t, noise=None):
357
+ noise = default(noise, lambda: torch.randn_like(x_start))
358
+ return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
359
+ extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)
360
+
361
+ def get_v(self, x, noise, t):
362
+ return (
363
+ extract_into_tensor(self.sqrt_alphas_cumprod, t, x.shape) * noise -
364
+ extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x.shape) * x
365
+ )
366
+
367
+ def get_loss(self, pred, target, mean=True):
368
+ if self.loss_type == 'l1':
369
+ loss = (target - pred).abs()
370
+ if mean:
371
+ loss = loss.mean()
372
+ elif self.loss_type == 'l2':
373
+ if mean:
374
+ loss = torch.nn.functional.mse_loss(target, pred)
375
+ else:
376
+ loss = torch.nn.functional.mse_loss(target, pred, reduction='none')
377
+ else:
378
+ raise NotImplementedError("unknown loss type '{loss_type}'")
379
+
380
+ return loss
381
+
382
+ def p_losses(self, x_start, t, noise=None):
383
+ noise = default(noise, lambda: torch.randn_like(x_start))
384
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
385
+ model_out = self.model(x_noisy, t)
386
+
387
+ loss_dict = {}
388
+ if self.parameterization == "eps":
389
+ target = noise
390
+ elif self.parameterization == "x0":
391
+ target = x_start
392
+ elif self.parameterization == "v":
393
+ target = self.get_v(x_start, noise, t)
394
+ else:
395
+ raise NotImplementedError(f"Parameterization {self.parameterization} not yet supported")
396
+
397
+ loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3])
398
+
399
+ log_prefix = 'train' if self.training else 'val'
400
+
401
+ loss_dict.update({f'{log_prefix}/loss_simple': loss.mean()})
402
+ loss_simple = loss.mean() * self.l_simple_weight
403
+
404
+ loss_vlb = (self.lvlb_weights[t] * loss).mean()
405
+ loss_dict.update({f'{log_prefix}/loss_vlb': loss_vlb})
406
+
407
+ loss = loss_simple + self.original_elbo_weight * loss_vlb
408
+
409
+ loss_dict.update({f'{log_prefix}/loss': loss})
410
+
411
+ return loss, loss_dict
412
+
413
+ def forward(self, x, *args, **kwargs):
414
+ # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size
415
+ # assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
416
+ t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long()
417
+ return self.p_losses(x, t, *args, **kwargs)
418
+
419
+ def get_input(self, batch, k):
420
+ x = batch[k]
421
+ if len(x.shape) == 3:
422
+ x = x[..., None]
423
+ x = rearrange(x, 'b h w c -> b c h w')
424
+ x = x.to(memory_format=torch.contiguous_format).float()
425
+ return x
426
+
427
+ def shared_step(self, batch):
428
+ x = self.get_input(batch, self.first_stage_key)
429
+ loss, loss_dict = self(x)
430
+ return loss, loss_dict
431
+
432
+ def training_step(self, batch, batch_idx):
433
+ for k in self.ucg_training:
434
+ p = self.ucg_training[k]["p"]
435
+ val = self.ucg_training[k]["val"]
436
+ if val is None:
437
+ val = ""
438
+ for i in range(len(batch[k])):
439
+ if self.ucg_prng.choice(2, p=[1 - p, p]):
440
+ batch[k][i] = val
441
+
442
+ loss, loss_dict = self.shared_step(batch)
443
+
444
+ self.log_dict(loss_dict, prog_bar=True,
445
+ logger=True, on_step=True, on_epoch=True)
446
+
447
+ self.log("global_step", self.global_step,
448
+ prog_bar=True, logger=True, on_step=True, on_epoch=False)
449
+
450
+ if self.use_scheduler:
451
+ lr = self.optimizers().param_groups[0]['lr']
452
+ self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False)
453
+
454
+ return loss
455
+
456
+ @torch.no_grad()
457
+ def validation_step(self, batch, batch_idx):
458
+ _, loss_dict_no_ema = self.shared_step(batch)
459
+ with self.ema_scope():
460
+ _, loss_dict_ema = self.shared_step(batch)
461
+ loss_dict_ema = {key + '_ema': loss_dict_ema[key] for key in loss_dict_ema}
462
+ self.log_dict(loss_dict_no_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True)
463
+ self.log_dict(loss_dict_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True)
464
+
465
+ def on_train_batch_end(self, *args, **kwargs):
466
+ if self.use_ema:
467
+ self.model_ema(self.model)
468
+
469
+ def _get_rows_from_list(self, samples):
470
+ n_imgs_per_row = len(samples)
471
+ denoise_grid = rearrange(samples, 'n b c h w -> b n c h w')
472
+ denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
473
+ denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
474
+ return denoise_grid
475
+
476
+ @torch.no_grad()
477
+ def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=None, **kwargs):
478
+ log = dict()
479
+ x = self.get_input(batch, self.first_stage_key)
480
+ N = min(x.shape[0], N)
481
+ n_row = min(x.shape[0], n_row)
482
+ x = x.to(self.device)[:N]
483
+ log["inputs"] = x
484
+
485
+ # get diffusion row
486
+ diffusion_row = list()
487
+ x_start = x[:n_row]
488
+
489
+ for t in range(self.num_timesteps):
490
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
491
+ t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
492
+ t = t.to(self.device).long()
493
+ noise = torch.randn_like(x_start)
494
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
495
+ diffusion_row.append(x_noisy)
496
+
497
+ log["diffusion_row"] = self._get_rows_from_list(diffusion_row)
498
+
499
+ if sample:
500
+ # get denoise row
501
+ with self.ema_scope("Plotting"):
502
+ samples, denoise_row = self.sample(batch_size=N, return_intermediates=True)
503
+
504
+ log["samples"] = samples
505
+ log["denoise_row"] = self._get_rows_from_list(denoise_row)
506
+
507
+ if return_keys:
508
+ if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
509
+ return log
510
+ else:
511
+ return {key: log[key] for key in return_keys}
512
+ return log
513
+
514
+ def configure_optimizers(self):
515
+ lr = self.learning_rate
516
+ params = list(self.model.parameters())
517
+ if self.learn_logvar:
518
+ params = params + [self.logvar]
519
+ opt = torch.optim.AdamW(params, lr=lr)
520
+ return opt
521
+
522
+
523
+ class LatentDiffusion(DDPM):
524
+ """main class"""
525
+
526
+ def __init__(self,
527
+ first_stage_config,
528
+ cond_stage_config,
529
+ num_timesteps_cond=None,
530
+ cond_stage_key="image",
531
+ cond_stage_trainable=False,
532
+ concat_mode=True,
533
+ cond_stage_forward=None,
534
+ conditioning_key=None,
535
+ scale_factor=1.0,
536
+ scale_by_std=False,
537
+ force_null_conditioning=False,
538
+ *args, **kwargs):
539
+ self.force_null_conditioning = force_null_conditioning
540
+ self.num_timesteps_cond = default(num_timesteps_cond, 1)
541
+ self.scale_by_std = scale_by_std
542
+ assert self.num_timesteps_cond <= kwargs['timesteps']
543
+ # for backwards compatibility after implementation of DiffusionWrapper
544
+ if conditioning_key is None:
545
+ conditioning_key = 'concat' if concat_mode else 'crossattn'
546
+ if cond_stage_config == '__is_unconditional__' and not self.force_null_conditioning:
547
+ conditioning_key = None
548
+ ckpt_path = kwargs.pop("ckpt_path", None)
549
+ reset_ema = kwargs.pop("reset_ema", False)
550
+ reset_num_ema_updates = kwargs.pop("reset_num_ema_updates", False)
551
+ ignore_keys = kwargs.pop("ignore_keys", [])
552
+ super().__init__(conditioning_key=conditioning_key, *args, **kwargs)
553
+ self.concat_mode = concat_mode
554
+ self.cond_stage_trainable = cond_stage_trainable
555
+ self.cond_stage_key = cond_stage_key
556
+ try:
557
+ self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1
558
+ except:
559
+ self.num_downs = 0
560
+ if not scale_by_std:
561
+ self.scale_factor = scale_factor
562
+ else:
563
+ self.register_buffer('scale_factor', torch.tensor(scale_factor))
564
+ self.instantiate_first_stage(first_stage_config)
565
+ self.instantiate_cond_stage(cond_stage_config)
566
+ self.cond_stage_forward = cond_stage_forward
567
+ self.clip_denoised = False
568
+ self.bbox_tokenizer = None
569
+
570
+ self.restarted_from_ckpt = False
571
+ if ckpt_path is not None:
572
+ self.init_from_ckpt(ckpt_path, ignore_keys)
573
+ self.restarted_from_ckpt = True
574
+ if reset_ema:
575
+ assert self.use_ema
576
+ print(
577
+ f"Resetting ema to pure model weights. This is useful when restoring from an ema-only checkpoint.")
578
+ self.model_ema = LitEma(self.model)
579
+ if reset_num_ema_updates:
580
+ print(" +++++++++++ WARNING: RESETTING NUM_EMA UPDATES TO ZERO +++++++++++ ")
581
+ assert self.use_ema
582
+ self.model_ema.reset_num_updates()
583
+
584
+ def make_cond_schedule(self, ):
585
+ self.cond_ids = torch.full(size=(self.num_timesteps,), fill_value=self.num_timesteps - 1, dtype=torch.long)
586
+ ids = torch.round(torch.linspace(0, self.num_timesteps - 1, self.num_timesteps_cond)).long()
587
+ self.cond_ids[:self.num_timesteps_cond] = ids
588
+
589
+ @rank_zero_only
590
+ @torch.no_grad()
591
+ def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
592
+ # only for very first batch
593
+ if self.scale_by_std and self.current_epoch == 0 and self.global_step == 0 and batch_idx == 0 and not self.restarted_from_ckpt:
594
+ assert self.scale_factor == 1., 'rather not use custom rescaling and std-rescaling simultaneously'
595
+ # set rescale weight to 1./std of encodings
596
+ print("### USING STD-RESCALING ###")
597
+ x = super().get_input(batch, self.first_stage_key)
598
+ x = x.to(self.device)
599
+ encoder_posterior = self.encode_first_stage(x)
600
+ z = self.get_first_stage_encoding(encoder_posterior).detach()
601
+ del self.scale_factor
602
+ self.register_buffer('scale_factor', 1. / z.flatten().std())
603
+ print(f"setting self.scale_factor to {self.scale_factor}")
604
+ print("### USING STD-RESCALING ###")
605
+
606
+ def register_schedule(self,
607
+ given_betas=None, beta_schedule="linear", timesteps=1000,
608
+ linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
609
+ super().register_schedule(given_betas, beta_schedule, timesteps, linear_start, linear_end, cosine_s)
610
+
611
+ self.shorten_cond_schedule = self.num_timesteps_cond > 1
612
+ if self.shorten_cond_schedule:
613
+ self.make_cond_schedule()
614
+
615
+ def instantiate_first_stage(self, config):
616
+ model = instantiate_from_config(config)
617
+ self.first_stage_model = model.eval()
618
+ self.first_stage_model.train = disabled_train
619
+ for param in self.first_stage_model.parameters():
620
+ param.requires_grad = False
621
+
622
+ def instantiate_cond_stage(self, config):
623
+ if not self.cond_stage_trainable:
624
+ if config == "__is_first_stage__":
625
+ print("Using first stage also as cond stage.")
626
+ self.cond_stage_model = self.first_stage_model
627
+ elif config == "__is_unconditional__":
628
+ print(f"Training {self.__class__.__name__} as an unconditional model.")
629
+ self.cond_stage_model = None
630
+ # self.be_unconditional = True
631
+ else:
632
+ model = instantiate_from_config(config)
633
+ self.cond_stage_model = model.eval()
634
+ self.cond_stage_model.train = disabled_train
635
+ for param in self.cond_stage_model.parameters():
636
+ param.requires_grad = False
637
+ else:
638
+ assert config != '__is_first_stage__'
639
+ assert config != '__is_unconditional__'
640
+ model = instantiate_from_config(config)
641
+ self.cond_stage_model = model
642
+
643
+ def _get_denoise_row_from_list(self, samples, desc='', force_no_decoder_quantization=False):
644
+ denoise_row = []
645
+ for zd in tqdm(samples, desc=desc):
646
+ denoise_row.append(self.decode_first_stage(zd.to(self.device),
647
+ force_not_quantize=force_no_decoder_quantization))
648
+ n_imgs_per_row = len(denoise_row)
649
+ denoise_row = torch.stack(denoise_row) # n_log_step, n_row, C, H, W
650
+ denoise_grid = rearrange(denoise_row, 'n b c h w -> b n c h w')
651
+ denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
652
+ denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
653
+ return denoise_grid
654
+
655
+ def get_first_stage_encoding(self, encoder_posterior):
656
+ if isinstance(encoder_posterior, DiagonalGaussianDistribution):
657
+ z = encoder_posterior.sample()
658
+ elif isinstance(encoder_posterior, torch.Tensor):
659
+ z = encoder_posterior
660
+ else:
661
+ raise NotImplementedError(f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented")
662
+ return self.scale_factor * z
663
+
664
+ def get_learned_conditioning(self, c):
665
+ if self.cond_stage_forward is None:
666
+ if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode):
667
+ c = self.cond_stage_model.encode(c)
668
+ if isinstance(c, DiagonalGaussianDistribution):
669
+ c = c.mode()
670
+ else:
671
+ c = self.cond_stage_model(c)
672
+ else:
673
+ assert hasattr(self.cond_stage_model, self.cond_stage_forward)
674
+ c = getattr(self.cond_stage_model, self.cond_stage_forward)(c)
675
+ return c
676
+
677
+ def meshgrid(self, h, w):
678
+ y = torch.arange(0, h).view(h, 1, 1).repeat(1, w, 1)
679
+ x = torch.arange(0, w).view(1, w, 1).repeat(h, 1, 1)
680
+
681
+ arr = torch.cat([y, x], dim=-1)
682
+ return arr
683
+
684
+ def delta_border(self, h, w):
685
+ """
686
+ :param h: height
687
+ :param w: width
688
+ :return: normalized distance to image border,
689
+ wtith min distance = 0 at border and max dist = 0.5 at image center
690
+ """
691
+ lower_right_corner = torch.tensor([h - 1, w - 1]).view(1, 1, 2)
692
+ arr = self.meshgrid(h, w) / lower_right_corner
693
+ dist_left_up = torch.min(arr, dim=-1, keepdims=True)[0]
694
+ dist_right_down = torch.min(1 - arr, dim=-1, keepdims=True)[0]
695
+ edge_dist = torch.min(torch.cat([dist_left_up, dist_right_down], dim=-1), dim=-1)[0]
696
+ return edge_dist
697
+
698
+ def get_weighting(self, h, w, Ly, Lx, device):
699
+ weighting = self.delta_border(h, w)
700
+ weighting = torch.clip(weighting, self.split_input_params["clip_min_weight"],
701
+ self.split_input_params["clip_max_weight"], )
702
+ weighting = weighting.view(1, h * w, 1).repeat(1, 1, Ly * Lx).to(device)
703
+
704
+ if self.split_input_params["tie_braker"]:
705
+ L_weighting = self.delta_border(Ly, Lx)
706
+ L_weighting = torch.clip(L_weighting,
707
+ self.split_input_params["clip_min_tie_weight"],
708
+ self.split_input_params["clip_max_tie_weight"])
709
+
710
+ L_weighting = L_weighting.view(1, 1, Ly * Lx).to(device)
711
+ weighting = weighting * L_weighting
712
+ return weighting
713
+
714
+ def get_fold_unfold(self, x, kernel_size, stride, uf=1, df=1): # todo load once not every time, shorten code
715
+ """
716
+ :param x: img of size (bs, c, h, w)
717
+ :return: n img crops of size (n, bs, c, kernel_size[0], kernel_size[1])
718
+ """
719
+ bs, nc, h, w = x.shape
720
+
721
+ # number of crops in image
722
+ Ly = (h - kernel_size[0]) // stride[0] + 1
723
+ Lx = (w - kernel_size[1]) // stride[1] + 1
724
+
725
+ if uf == 1 and df == 1:
726
+ fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
727
+ unfold = torch.nn.Unfold(**fold_params)
728
+
729
+ fold = torch.nn.Fold(output_size=x.shape[2:], **fold_params)
730
+
731
+ weighting = self.get_weighting(kernel_size[0], kernel_size[1], Ly, Lx, x.device).to(x.dtype)
732
+ normalization = fold(weighting).view(1, 1, h, w) # normalizes the overlap
733
+ weighting = weighting.view((1, 1, kernel_size[0], kernel_size[1], Ly * Lx))
734
+
735
+ elif uf > 1 and df == 1:
736
+ fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
737
+ unfold = torch.nn.Unfold(**fold_params)
738
+
739
+ fold_params2 = dict(kernel_size=(kernel_size[0] * uf, kernel_size[0] * uf),
740
+ dilation=1, padding=0,
741
+ stride=(stride[0] * uf, stride[1] * uf))
742
+ fold = torch.nn.Fold(output_size=(x.shape[2] * uf, x.shape[3] * uf), **fold_params2)
743
+
744
+ weighting = self.get_weighting(kernel_size[0] * uf, kernel_size[1] * uf, Ly, Lx, x.device).to(x.dtype)
745
+ normalization = fold(weighting).view(1, 1, h * uf, w * uf) # normalizes the overlap
746
+ weighting = weighting.view((1, 1, kernel_size[0] * uf, kernel_size[1] * uf, Ly * Lx))
747
+
748
+ elif df > 1 and uf == 1:
749
+ fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
750
+ unfold = torch.nn.Unfold(**fold_params)
751
+
752
+ fold_params2 = dict(kernel_size=(kernel_size[0] // df, kernel_size[0] // df),
753
+ dilation=1, padding=0,
754
+ stride=(stride[0] // df, stride[1] // df))
755
+ fold = torch.nn.Fold(output_size=(x.shape[2] // df, x.shape[3] // df), **fold_params2)
756
+
757
+ weighting = self.get_weighting(kernel_size[0] // df, kernel_size[1] // df, Ly, Lx, x.device).to(x.dtype)
758
+ normalization = fold(weighting).view(1, 1, h // df, w // df) # normalizes the overlap
759
+ weighting = weighting.view((1, 1, kernel_size[0] // df, kernel_size[1] // df, Ly * Lx))
760
+
761
+ else:
762
+ raise NotImplementedError
763
+
764
+ return fold, unfold, normalization, weighting
765
+
766
+ @torch.no_grad()
767
+ def get_input(self, batch, k, return_first_stage_outputs=False, force_c_encode=False,
768
+ cond_key=None, return_original_cond=False, bs=None, return_x=False):
769
+ x = super().get_input(batch, k)
770
+ if bs is not None:
771
+ x = x[:bs]
772
+ x = x.to(self.device)
773
+ encoder_posterior = self.encode_first_stage(x)
774
+ z = self.get_first_stage_encoding(encoder_posterior).detach()
775
+
776
+ if self.model.conditioning_key is not None and not self.force_null_conditioning:
777
+ if cond_key is None:
778
+ cond_key = self.cond_stage_key
779
+ if cond_key != self.first_stage_key:
780
+ if cond_key in ['caption', 'coordinates_bbox', "txt"]:
781
+ xc = batch[cond_key]
782
+ elif cond_key in ['class_label', 'cls']:
783
+ xc = batch
784
+ else:
785
+ xc = super().get_input(batch, cond_key).to(self.device)
786
+ else:
787
+ xc = x
788
+ if not self.cond_stage_trainable or force_c_encode:
789
+ if isinstance(xc, dict) or isinstance(xc, list):
790
+ c = self.get_learned_conditioning(xc)
791
+ else:
792
+ c = self.get_learned_conditioning(xc.to(self.device))
793
+ else:
794
+ c = xc
795
+ if bs is not None:
796
+ c = c[:bs]
797
+
798
+ if self.use_positional_encodings:
799
+ pos_x, pos_y = self.compute_latent_shifts(batch)
800
+ ckey = __conditioning_keys__[self.model.conditioning_key]
801
+ c = {ckey: c, 'pos_x': pos_x, 'pos_y': pos_y}
802
+
803
+ else:
804
+ c = None
805
+ xc = None
806
+ if self.use_positional_encodings:
807
+ pos_x, pos_y = self.compute_latent_shifts(batch)
808
+ c = {'pos_x': pos_x, 'pos_y': pos_y}
809
+ out = [z, c]
810
+ if return_first_stage_outputs:
811
+ xrec = self.decode_first_stage(z)
812
+ out.extend([x, xrec])
813
+ if return_x:
814
+ out.extend([x])
815
+ if return_original_cond:
816
+ out.append(xc)
817
+ return out
818
+
819
+ @torch.no_grad()
820
+ def decode_first_stage(self, z, predict_cids=False, force_not_quantize=False):
821
+ if predict_cids:
822
+ if z.dim() == 4:
823
+ z = torch.argmax(z.exp(), dim=1).long()
824
+ z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None)
825
+ z = rearrange(z, 'b h w c -> b c h w').contiguous()
826
+
827
+ z = 1. / self.scale_factor * z
828
+ return self.first_stage_model.decode(z)
829
+
830
+ @torch.no_grad()
831
+ def encode_first_stage(self, x):
832
+ return self.first_stage_model.encode(x)
833
+
834
+ def shared_step(self, batch, **kwargs):
835
+ x, c = self.get_input(batch, self.first_stage_key)
836
+ loss = self(x, c)
837
+ return loss
838
+
839
+ def forward(self, x, c, *args, **kwargs):
840
+ t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long()
841
+ if self.model.conditioning_key is not None:
842
+ assert c is not None
843
+ if self.cond_stage_trainable:
844
+ c = self.get_learned_conditioning(c)
845
+ if self.shorten_cond_schedule: # TODO: drop this option
846
+ tc = self.cond_ids[t].to(self.device)
847
+ c = self.q_sample(x_start=c, t=tc, noise=torch.randn_like(c.float()))
848
+ return self.p_losses(x, c, t, *args, **kwargs)
849
+
850
+ def apply_model(self, x_noisy, t, cond, return_ids=False):
851
+ if isinstance(cond, dict):
852
+ # hybrid case, cond is expected to be a dict
853
+ pass
854
+ else:
855
+ if not isinstance(cond, list):
856
+ cond = [cond]
857
+ key = 'c_concat' if self.model.conditioning_key == 'concat' else 'c_crossattn'
858
+ cond = {key: cond}
859
+
860
+ x_recon = self.model(x_noisy, t, **cond)
861
+
862
+ if isinstance(x_recon, tuple) and not return_ids:
863
+ return x_recon[0]
864
+ else:
865
+ return x_recon
866
+
867
+ def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
868
+ return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart) / \
869
+ extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
870
+
871
+ def _prior_bpd(self, x_start):
872
+ """
873
+ Get the prior KL term for the variational lower-bound, measured in
874
+ bits-per-dim.
875
+ This term can't be optimized, as it only depends on the encoder.
876
+ :param x_start: the [N x C x ...] tensor of inputs.
877
+ :return: a batch of [N] KL values (in bits), one per batch element.
878
+ """
879
+ batch_size = x_start.shape[0]
880
+ t = torch.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
881
+ qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
882
+ kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0)
883
+ return mean_flat(kl_prior) / np.log(2.0)
884
+
885
+ def p_losses(self, x_start, cond, t, noise=None):
886
+ noise = default(noise, lambda: torch.randn_like(x_start))
887
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
888
+ model_output = self.apply_model(x_noisy, t, cond)
889
+
890
+ loss_dict = {}
891
+ prefix = 'train' if self.training else 'val'
892
+
893
+ if self.parameterization == "x0":
894
+ target = x_start
895
+ elif self.parameterization == "eps":
896
+ target = noise
897
+ elif self.parameterization == "v":
898
+ target = self.get_v(x_start, noise, t)
899
+ else:
900
+ raise NotImplementedError()
901
+
902
+ loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3])
903
+ loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()})
904
+
905
+ logvar_t = self.logvar[t].to(self.device)
906
+ loss = loss_simple / torch.exp(logvar_t) + logvar_t
907
+ # loss = loss_simple / torch.exp(self.logvar) + self.logvar
908
+ if self.learn_logvar:
909
+ loss_dict.update({f'{prefix}/loss_gamma': loss.mean()})
910
+ loss_dict.update({'logvar': self.logvar.data.mean()})
911
+
912
+ loss = self.l_simple_weight * loss.mean()
913
+
914
+ loss_vlb = self.get_loss(model_output, target, mean=False).mean(dim=(1, 2, 3))
915
+ loss_vlb = (self.lvlb_weights[t] * loss_vlb).mean()
916
+ loss_dict.update({f'{prefix}/loss_vlb': loss_vlb})
917
+ loss += (self.original_elbo_weight * loss_vlb)
918
+ loss_dict.update({f'{prefix}/loss': loss})
919
+
920
+ return loss, loss_dict
921
+
922
+ def p_mean_variance(self, x, c, t, clip_denoised: bool, return_codebook_ids=False, quantize_denoised=False,
923
+ return_x0=False, score_corrector=None, corrector_kwargs=None):
924
+ t_in = t
925
+ model_out = self.apply_model(x, t_in, c, return_ids=return_codebook_ids)
926
+
927
+ if score_corrector is not None:
928
+ assert self.parameterization == "eps"
929
+ model_out = score_corrector.modify_score(self, model_out, x, t, c, **corrector_kwargs)
930
+
931
+ if return_codebook_ids:
932
+ model_out, logits = model_out
933
+
934
+ if self.parameterization == "eps":
935
+ x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
936
+ elif self.parameterization == "x0":
937
+ x_recon = model_out
938
+ else:
939
+ raise NotImplementedError()
940
+
941
+ if clip_denoised:
942
+ x_recon.clamp_(-1., 1.)
943
+ if quantize_denoised:
944
+ x_recon, _, [_, _, indices] = self.first_stage_model.quantize(x_recon)
945
+ model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
946
+ if return_codebook_ids:
947
+ return model_mean, posterior_variance, posterior_log_variance, logits
948
+ elif return_x0:
949
+ return model_mean, posterior_variance, posterior_log_variance, x_recon
950
+ else:
951
+ return model_mean, posterior_variance, posterior_log_variance
952
+
953
+ @torch.no_grad()
954
+ def p_sample(self, x, c, t, clip_denoised=False, repeat_noise=False,
955
+ return_codebook_ids=False, quantize_denoised=False, return_x0=False,
956
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None):
957
+ b, *_, device = *x.shape, x.device
958
+ outputs = self.p_mean_variance(x=x, c=c, t=t, clip_denoised=clip_denoised,
959
+ return_codebook_ids=return_codebook_ids,
960
+ quantize_denoised=quantize_denoised,
961
+ return_x0=return_x0,
962
+ score_corrector=score_corrector, corrector_kwargs=corrector_kwargs)
963
+ if return_codebook_ids:
964
+ raise DeprecationWarning("Support dropped.")
965
+ model_mean, _, model_log_variance, logits = outputs
966
+ elif return_x0:
967
+ model_mean, _, model_log_variance, x0 = outputs
968
+ else:
969
+ model_mean, _, model_log_variance = outputs
970
+
971
+ noise = noise_like(x.shape, device, repeat_noise) * temperature
972
+ if noise_dropout > 0.:
973
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
974
+ # no noise when t == 0
975
+ nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
976
+
977
+ if return_codebook_ids:
978
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, logits.argmax(dim=1)
979
+ if return_x0:
980
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, x0
981
+ else:
982
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
983
+
984
+ @torch.no_grad()
985
+ def progressive_denoising(self, cond, shape, verbose=True, callback=None, quantize_denoised=False,
986
+ img_callback=None, mask=None, x0=None, temperature=1., noise_dropout=0.,
987
+ score_corrector=None, corrector_kwargs=None, batch_size=None, x_T=None, start_T=None,
988
+ log_every_t=None):
989
+ if not log_every_t:
990
+ log_every_t = self.log_every_t
991
+ timesteps = self.num_timesteps
992
+ if batch_size is not None:
993
+ b = batch_size if batch_size is not None else shape[0]
994
+ shape = [batch_size] + list(shape)
995
+ else:
996
+ b = batch_size = shape[0]
997
+ if x_T is None:
998
+ img = torch.randn(shape, device=self.device)
999
+ else:
1000
+ img = x_T
1001
+ intermediates = []
1002
+ if cond is not None:
1003
+ if isinstance(cond, dict):
1004
+ cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
1005
+ list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
1006
+ else:
1007
+ cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
1008
+
1009
+ if start_T is not None:
1010
+ timesteps = min(timesteps, start_T)
1011
+ iterator = tqdm(reversed(range(0, timesteps)), desc='Progressive Generation',
1012
+ total=timesteps) if verbose else reversed(
1013
+ range(0, timesteps))
1014
+ if type(temperature) == float:
1015
+ temperature = [temperature] * timesteps
1016
+
1017
+ for i in iterator:
1018
+ ts = torch.full((b,), i, device=self.device, dtype=torch.long)
1019
+ if self.shorten_cond_schedule:
1020
+ assert self.model.conditioning_key != 'hybrid'
1021
+ tc = self.cond_ids[ts].to(cond.device)
1022
+ cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond))
1023
+
1024
+ img, x0_partial = self.p_sample(img, cond, ts,
1025
+ clip_denoised=self.clip_denoised,
1026
+ quantize_denoised=quantize_denoised, return_x0=True,
1027
+ temperature=temperature[i], noise_dropout=noise_dropout,
1028
+ score_corrector=score_corrector, corrector_kwargs=corrector_kwargs)
1029
+ if mask is not None:
1030
+ assert x0 is not None
1031
+ img_orig = self.q_sample(x0, ts)
1032
+ img = img_orig * mask + (1. - mask) * img
1033
+
1034
+ if i % log_every_t == 0 or i == timesteps - 1:
1035
+ intermediates.append(x0_partial)
1036
+ if callback: callback(i)
1037
+ if img_callback: img_callback(img, i)
1038
+ return img, intermediates
1039
+
1040
+ @torch.no_grad()
1041
+ def p_sample_loop(self, cond, shape, return_intermediates=False,
1042
+ x_T=None, verbose=True, callback=None, timesteps=None, quantize_denoised=False,
1043
+ mask=None, x0=None, img_callback=None, start_T=None,
1044
+ log_every_t=None):
1045
+
1046
+ if not log_every_t:
1047
+ log_every_t = self.log_every_t
1048
+ device = self.betas.device
1049
+ b = shape[0]
1050
+ if x_T is None:
1051
+ img = torch.randn(shape, device=device)
1052
+ else:
1053
+ img = x_T
1054
+
1055
+ intermediates = [img]
1056
+ if timesteps is None:
1057
+ timesteps = self.num_timesteps
1058
+
1059
+ if start_T is not None:
1060
+ timesteps = min(timesteps, start_T)
1061
+ iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed(
1062
+ range(0, timesteps))
1063
+
1064
+ if mask is not None:
1065
+ assert x0 is not None
1066
+ assert x0.shape[2:3] == mask.shape[2:3] # spatial size has to match
1067
+
1068
+ for i in iterator:
1069
+ ts = torch.full((b,), i, device=device, dtype=torch.long)
1070
+ if self.shorten_cond_schedule:
1071
+ assert self.model.conditioning_key != 'hybrid'
1072
+ tc = self.cond_ids[ts].to(cond.device)
1073
+ cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond))
1074
+
1075
+ img = self.p_sample(img, cond, ts,
1076
+ clip_denoised=self.clip_denoised,
1077
+ quantize_denoised=quantize_denoised)
1078
+ if mask is not None:
1079
+ img_orig = self.q_sample(x0, ts)
1080
+ img = img_orig * mask + (1. - mask) * img
1081
+
1082
+ if i % log_every_t == 0 or i == timesteps - 1:
1083
+ intermediates.append(img)
1084
+ if callback: callback(i)
1085
+ if img_callback: img_callback(img, i)
1086
+
1087
+ if return_intermediates:
1088
+ return img, intermediates
1089
+ return img
1090
+
1091
+ @torch.no_grad()
1092
+ def sample(self, cond, batch_size=16, return_intermediates=False, x_T=None,
1093
+ verbose=True, timesteps=None, quantize_denoised=False,
1094
+ mask=None, x0=None, shape=None, **kwargs):
1095
+ if shape is None:
1096
+ shape = (batch_size, self.channels, self.image_size, self.image_size)
1097
+ if cond is not None:
1098
+ if isinstance(cond, dict):
1099
+ cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
1100
+ list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
1101
+ else:
1102
+ cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
1103
+ return self.p_sample_loop(cond,
1104
+ shape,
1105
+ return_intermediates=return_intermediates, x_T=x_T,
1106
+ verbose=verbose, timesteps=timesteps, quantize_denoised=quantize_denoised,
1107
+ mask=mask, x0=x0)
1108
+
1109
+ @torch.no_grad()
1110
+ def sample_log(self, cond, batch_size, ddim, ddim_steps, **kwargs):
1111
+ if ddim:
1112
+ ddim_sampler = DDIMSampler(self)
1113
+ shape = (self.channels, self.image_size, self.image_size)
1114
+ samples, intermediates = ddim_sampler.sample(ddim_steps, batch_size,
1115
+ shape, cond, verbose=False, **kwargs)
1116
+
1117
+ else:
1118
+ samples, intermediates = self.sample(cond=cond, batch_size=batch_size,
1119
+ return_intermediates=True, **kwargs)
1120
+
1121
+ return samples, intermediates
1122
+
1123
+ @torch.no_grad()
1124
+ def get_unconditional_conditioning(self, batch_size, null_label=None):
1125
+ if null_label is not None:
1126
+ xc = null_label
1127
+ if isinstance(xc, ListConfig):
1128
+ xc = list(xc)
1129
+ if isinstance(xc, dict) or isinstance(xc, list):
1130
+ c = self.get_learned_conditioning(xc)
1131
+ else:
1132
+ if hasattr(xc, "to"):
1133
+ xc = xc.to(self.device)
1134
+ c = self.get_learned_conditioning(xc)
1135
+ else:
1136
+ if self.cond_stage_key in ["class_label", "cls"]:
1137
+ xc = self.cond_stage_model.get_unconditional_conditioning(batch_size, device=self.device)
1138
+ return self.get_learned_conditioning(xc)
1139
+ else:
1140
+ raise NotImplementedError("todo")
1141
+ if isinstance(c, list): # in case the encoder gives us a list
1142
+ for i in range(len(c)):
1143
+ c[i] = repeat(c[i], '1 ... -> b ...', b=batch_size).to(self.device)
1144
+ else:
1145
+ c = repeat(c, '1 ... -> b ...', b=batch_size).to(self.device)
1146
+ return c
1147
+
1148
+ @torch.no_grad()
1149
+ def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=50, ddim_eta=0., return_keys=None,
1150
+ quantize_denoised=True, inpaint=True, plot_denoise_rows=False, plot_progressive_rows=True,
1151
+ plot_diffusion_rows=True, unconditional_guidance_scale=1., unconditional_guidance_label=None,
1152
+ use_ema_scope=True,
1153
+ **kwargs):
1154
+ ema_scope = self.ema_scope if use_ema_scope else nullcontext
1155
+ use_ddim = ddim_steps is not None
1156
+
1157
+ log = dict()
1158
+ z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key,
1159
+ return_first_stage_outputs=True,
1160
+ force_c_encode=True,
1161
+ return_original_cond=True,
1162
+ bs=N)
1163
+ N = min(x.shape[0], N)
1164
+ n_row = min(x.shape[0], n_row)
1165
+ log["inputs"] = x
1166
+ log["reconstruction"] = xrec
1167
+ if self.model.conditioning_key is not None:
1168
+ if hasattr(self.cond_stage_model, "decode"):
1169
+ xc = self.cond_stage_model.decode(c)
1170
+ log["conditioning"] = xc
1171
+ elif self.cond_stage_key in ["caption", "txt"]:
1172
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch[self.cond_stage_key], size=x.shape[2] // 25)
1173
+ log["conditioning"] = xc
1174
+ elif self.cond_stage_key in ['class_label', "cls"]:
1175
+ try:
1176
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"], size=x.shape[2] // 25)
1177
+ log['conditioning'] = xc
1178
+ except KeyError:
1179
+ # probably no "human_label" in batch
1180
+ pass
1181
+ elif isimage(xc):
1182
+ log["conditioning"] = xc
1183
+ if ismap(xc):
1184
+ log["original_conditioning"] = self.to_rgb(xc)
1185
+
1186
+ if plot_diffusion_rows:
1187
+ # get diffusion row
1188
+ diffusion_row = list()
1189
+ z_start = z[:n_row]
1190
+ for t in range(self.num_timesteps):
1191
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
1192
+ t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
1193
+ t = t.to(self.device).long()
1194
+ noise = torch.randn_like(z_start)
1195
+ z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
1196
+ diffusion_row.append(self.decode_first_stage(z_noisy))
1197
+
1198
+ diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W
1199
+ diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w')
1200
+ diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w')
1201
+ diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0])
1202
+ log["diffusion_row"] = diffusion_grid
1203
+
1204
+ if sample:
1205
+ # get denoise row
1206
+ with ema_scope("Sampling"):
1207
+ samples, z_denoise_row = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,
1208
+ ddim_steps=ddim_steps, eta=ddim_eta)
1209
+ # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True)
1210
+ x_samples = self.decode_first_stage(samples)
1211
+ log["samples"] = x_samples
1212
+ if plot_denoise_rows:
1213
+ denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
1214
+ log["denoise_row"] = denoise_grid
1215
+
1216
+ if quantize_denoised and not isinstance(self.first_stage_model, AutoencoderKL) and not isinstance(
1217
+ self.first_stage_model, IdentityFirstStage):
1218
+ # also display when quantizing x0 while sampling
1219
+ with ema_scope("Plotting Quantized Denoised"):
1220
+ samples, z_denoise_row = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,
1221
+ ddim_steps=ddim_steps, eta=ddim_eta,
1222
+ quantize_denoised=True)
1223
+ # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True,
1224
+ # quantize_denoised=True)
1225
+ x_samples = self.decode_first_stage(samples.to(self.device))
1226
+ log["samples_x0_quantized"] = x_samples
1227
+
1228
+ if unconditional_guidance_scale > 1.0:
1229
+ uc = self.get_unconditional_conditioning(N, unconditional_guidance_label)
1230
+ if self.model.conditioning_key == "crossattn-adm":
1231
+ uc = {"c_crossattn": [uc], "c_adm": c["c_adm"]}
1232
+ with ema_scope("Sampling with classifier-free guidance"):
1233
+ samples_cfg, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,
1234
+ ddim_steps=ddim_steps, eta=ddim_eta,
1235
+ unconditional_guidance_scale=unconditional_guidance_scale,
1236
+ unconditional_conditioning=uc,
1237
+ )
1238
+ x_samples_cfg = self.decode_first_stage(samples_cfg)
1239
+ log[f"samples_cfg_scale_{unconditional_guidance_scale:.2f}"] = x_samples_cfg
1240
+
1241
+ if inpaint:
1242
+ # make a simple center square
1243
+ b, h, w = z.shape[0], z.shape[2], z.shape[3]
1244
+ mask = torch.ones(N, h, w).to(self.device)
1245
+ # zeros will be filled in
1246
+ mask[:, h // 4:3 * h // 4, w // 4:3 * w // 4] = 0.
1247
+ mask = mask[:, None, ...]
1248
+ with ema_scope("Plotting Inpaint"):
1249
+ samples, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim, eta=ddim_eta,
1250
+ ddim_steps=ddim_steps, x0=z[:N], mask=mask)
1251
+ x_samples = self.decode_first_stage(samples.to(self.device))
1252
+ log["samples_inpainting"] = x_samples
1253
+ log["mask"] = mask
1254
+
1255
+ # outpaint
1256
+ mask = 1. - mask
1257
+ with ema_scope("Plotting Outpaint"):
1258
+ samples, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim, eta=ddim_eta,
1259
+ ddim_steps=ddim_steps, x0=z[:N], mask=mask)
1260
+ x_samples = self.decode_first_stage(samples.to(self.device))
1261
+ log["samples_outpainting"] = x_samples
1262
+
1263
+ if plot_progressive_rows:
1264
+ with ema_scope("Plotting Progressives"):
1265
+ img, progressives = self.progressive_denoising(c,
1266
+ shape=(self.channels, self.image_size, self.image_size),
1267
+ batch_size=N)
1268
+ prog_row = self._get_denoise_row_from_list(progressives, desc="Progressive Generation")
1269
+ log["progressive_row"] = prog_row
1270
+
1271
+ if return_keys:
1272
+ if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
1273
+ return log
1274
+ else:
1275
+ return {key: log[key] for key in return_keys}
1276
+ return log
1277
+
1278
+ def configure_optimizers(self):
1279
+ lr = self.learning_rate
1280
+ params = list(self.model.parameters())
1281
+ if self.cond_stage_trainable:
1282
+ print(f"{self.__class__.__name__}: Also optimizing conditioner params!")
1283
+ params = params + list(self.cond_stage_model.parameters())
1284
+ if self.learn_logvar:
1285
+ print('Diffusion model optimizing logvar')
1286
+ params.append(self.logvar)
1287
+ opt = torch.optim.AdamW(params, lr=lr)
1288
+ if self.use_scheduler:
1289
+ assert 'target' in self.scheduler_config
1290
+ scheduler = instantiate_from_config(self.scheduler_config)
1291
+
1292
+ print("Setting up LambdaLR scheduler...")
1293
+ scheduler = [
1294
+ {
1295
+ 'scheduler': LambdaLR(opt, lr_lambda=scheduler.schedule),
1296
+ 'interval': 'step',
1297
+ 'frequency': 1
1298
+ }]
1299
+ return [opt], scheduler
1300
+ return opt
1301
+
1302
+ @torch.no_grad()
1303
+ def to_rgb(self, x):
1304
+ x = x.float()
1305
+ if not hasattr(self, "colorize"):
1306
+ self.colorize = torch.randn(3, x.shape[1], 1, 1).to(x)
1307
+ x = nn.functional.conv2d(x, weight=self.colorize)
1308
+ x = 2. * (x - x.min()) / (x.max() - x.min()) - 1.
1309
+ return x
1310
+
1311
+
1312
+ class DiffusionWrapper(pl.LightningModule):
1313
+ def __init__(self, diff_model_config, conditioning_key):
1314
+ super().__init__()
1315
+ self.sequential_cross_attn = diff_model_config.pop("sequential_crossattn", False)
1316
+ self.diffusion_model = instantiate_from_config(diff_model_config)
1317
+ self.conditioning_key = conditioning_key
1318
+ assert self.conditioning_key in [None, 'concat', 'crossattn', 'hybrid', 'adm', 'hybrid-adm', 'crossattn-adm']
1319
+
1320
+ def forward(self, x, t, c_concat: list = None, c_crossattn: list = None, c_adm=None):
1321
+ if self.conditioning_key is None:
1322
+ out = self.diffusion_model(x, t)
1323
+ elif self.conditioning_key == 'concat':
1324
+ xc = torch.cat([x] + c_concat, dim=1)
1325
+ out = self.diffusion_model(xc, t)
1326
+ elif self.conditioning_key == 'crossattn':
1327
+ if not self.sequential_cross_attn:
1328
+ cc = torch.cat(c_crossattn, 1)
1329
+ else:
1330
+ cc = c_crossattn
1331
+ out = self.diffusion_model(x, t, context=cc)
1332
+ elif self.conditioning_key == 'hybrid':
1333
+ xc = torch.cat([x] + c_concat, dim=1)
1334
+ cc = torch.cat(c_crossattn, 1)
1335
+ out = self.diffusion_model(xc, t, context=cc)
1336
+ elif self.conditioning_key == 'hybrid-adm':
1337
+ assert c_adm is not None
1338
+ xc = torch.cat([x] + c_concat, dim=1)
1339
+ cc = torch.cat(c_crossattn, 1)
1340
+ out = self.diffusion_model(xc, t, context=cc, y=c_adm)
1341
+ elif self.conditioning_key == 'crossattn-adm':
1342
+ assert c_adm is not None
1343
+ cc = torch.cat(c_crossattn, 1)
1344
+ out = self.diffusion_model(x, t, context=cc, y=c_adm)
1345
+ elif self.conditioning_key == 'adm':
1346
+ cc = c_crossattn[0]
1347
+ out = self.diffusion_model(x, t, y=cc)
1348
+ else:
1349
+ raise NotImplementedError()
1350
+
1351
+ return out
1352
+
1353
+
1354
+ class LatentUpscaleDiffusion(LatentDiffusion):
1355
+ def __init__(self, *args, low_scale_config, low_scale_key="LR", noise_level_key=None, **kwargs):
1356
+ super().__init__(*args, **kwargs)
1357
+ # assumes that neither the cond_stage nor the low_scale_model contain trainable params
1358
+ assert not self.cond_stage_trainable
1359
+ self.instantiate_low_stage(low_scale_config)
1360
+ self.low_scale_key = low_scale_key
1361
+ self.noise_level_key = noise_level_key
1362
+
1363
+ def instantiate_low_stage(self, config):
1364
+ model = instantiate_from_config(config)
1365
+ self.low_scale_model = model.eval()
1366
+ self.low_scale_model.train = disabled_train
1367
+ for param in self.low_scale_model.parameters():
1368
+ param.requires_grad = False
1369
+
1370
+ @torch.no_grad()
1371
+ def get_input(self, batch, k, cond_key=None, bs=None, log_mode=False):
1372
+ if not log_mode:
1373
+ z, c = super().get_input(batch, k, force_c_encode=True, bs=bs)
1374
+ else:
1375
+ z, c, x, xrec, xc = super().get_input(batch, self.first_stage_key, return_first_stage_outputs=True,
1376
+ force_c_encode=True, return_original_cond=True, bs=bs)
1377
+ x_low = batch[self.low_scale_key][:bs]
1378
+ x_low = rearrange(x_low, 'b h w c -> b c h w')
1379
+ x_low = x_low.to(memory_format=torch.contiguous_format).float()
1380
+ zx, noise_level = self.low_scale_model(x_low)
1381
+ if self.noise_level_key is not None:
1382
+ # get noise level from batch instead, e.g. when extracting a custom noise level for bsr
1383
+ raise NotImplementedError('TODO')
1384
+
1385
+ all_conds = {"c_concat": [zx], "c_crossattn": [c], "c_adm": noise_level}
1386
+ if log_mode:
1387
+ # TODO: maybe disable if too expensive
1388
+ x_low_rec = self.low_scale_model.decode(zx)
1389
+ return z, all_conds, x, xrec, xc, x_low, x_low_rec, noise_level
1390
+ return z, all_conds
1391
+
1392
+ @torch.no_grad()
1393
+ def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=200, ddim_eta=1., return_keys=None,
1394
+ plot_denoise_rows=False, plot_progressive_rows=True, plot_diffusion_rows=True,
1395
+ unconditional_guidance_scale=1., unconditional_guidance_label=None, use_ema_scope=True,
1396
+ **kwargs):
1397
+ ema_scope = self.ema_scope if use_ema_scope else nullcontext
1398
+ use_ddim = ddim_steps is not None
1399
+
1400
+ log = dict()
1401
+ z, c, x, xrec, xc, x_low, x_low_rec, noise_level = self.get_input(batch, self.first_stage_key, bs=N,
1402
+ log_mode=True)
1403
+ N = min(x.shape[0], N)
1404
+ n_row = min(x.shape[0], n_row)
1405
+ log["inputs"] = x
1406
+ log["reconstruction"] = xrec
1407
+ log["x_lr"] = x_low
1408
+ log[f"x_lr_rec_@noise_levels{'-'.join(map(lambda x: str(x), list(noise_level.cpu().numpy())))}"] = x_low_rec
1409
+ if self.model.conditioning_key is not None:
1410
+ if hasattr(self.cond_stage_model, "decode"):
1411
+ xc = self.cond_stage_model.decode(c)
1412
+ log["conditioning"] = xc
1413
+ elif self.cond_stage_key in ["caption", "txt"]:
1414
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch[self.cond_stage_key], size=x.shape[2] // 25)
1415
+ log["conditioning"] = xc
1416
+ elif self.cond_stage_key in ['class_label', 'cls']:
1417
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"], size=x.shape[2] // 25)
1418
+ log['conditioning'] = xc
1419
+ elif isimage(xc):
1420
+ log["conditioning"] = xc
1421
+ if ismap(xc):
1422
+ log["original_conditioning"] = self.to_rgb(xc)
1423
+
1424
+ if plot_diffusion_rows:
1425
+ # get diffusion row
1426
+ diffusion_row = list()
1427
+ z_start = z[:n_row]
1428
+ for t in range(self.num_timesteps):
1429
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
1430
+ t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
1431
+ t = t.to(self.device).long()
1432
+ noise = torch.randn_like(z_start)
1433
+ z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
1434
+ diffusion_row.append(self.decode_first_stage(z_noisy))
1435
+
1436
+ diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W
1437
+ diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w')
1438
+ diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w')
1439
+ diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0])
1440
+ log["diffusion_row"] = diffusion_grid
1441
+
1442
+ if sample:
1443
+ # get denoise row
1444
+ with ema_scope("Sampling"):
1445
+ samples, z_denoise_row = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,
1446
+ ddim_steps=ddim_steps, eta=ddim_eta)
1447
+ # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True)
1448
+ x_samples = self.decode_first_stage(samples)
1449
+ log["samples"] = x_samples
1450
+ if plot_denoise_rows:
1451
+ denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
1452
+ log["denoise_row"] = denoise_grid
1453
+
1454
+ if unconditional_guidance_scale > 1.0:
1455
+ uc_tmp = self.get_unconditional_conditioning(N, unconditional_guidance_label)
1456
+ # TODO explore better "unconditional" choices for the other keys
1457
+ # maybe guide away from empty text label and highest noise level and maximally degraded zx?
1458
+ uc = dict()
1459
+ for k in c:
1460
+ if k == "c_crossattn":
1461
+ assert isinstance(c[k], list) and len(c[k]) == 1
1462
+ uc[k] = [uc_tmp]
1463
+ elif k == "c_adm": # todo: only run with text-based guidance?
1464
+ assert isinstance(c[k], torch.Tensor)
1465
+ #uc[k] = torch.ones_like(c[k]) * self.low_scale_model.max_noise_level
1466
+ uc[k] = c[k]
1467
+ elif isinstance(c[k], list):
1468
+ uc[k] = [c[k][i] for i in range(len(c[k]))]
1469
+ else:
1470
+ uc[k] = c[k]
1471
+
1472
+ with ema_scope("Sampling with classifier-free guidance"):
1473
+ samples_cfg, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,
1474
+ ddim_steps=ddim_steps, eta=ddim_eta,
1475
+ unconditional_guidance_scale=unconditional_guidance_scale,
1476
+ unconditional_conditioning=uc,
1477
+ )
1478
+ x_samples_cfg = self.decode_first_stage(samples_cfg)
1479
+ log[f"samples_cfg_scale_{unconditional_guidance_scale:.2f}"] = x_samples_cfg
1480
+
1481
+ if plot_progressive_rows:
1482
+ with ema_scope("Plotting Progressives"):
1483
+ img, progressives = self.progressive_denoising(c,
1484
+ shape=(self.channels, self.image_size, self.image_size),
1485
+ batch_size=N)
1486
+ prog_row = self._get_denoise_row_from_list(progressives, desc="Progressive Generation")
1487
+ log["progressive_row"] = prog_row
1488
+
1489
+ return log
1490
+
1491
+
1492
+ class LatentFinetuneDiffusion(LatentDiffusion):
1493
+ """
1494
+ Basis for different finetunas, such as inpainting or depth2image
1495
+ To disable finetuning mode, set finetune_keys to None
1496
+ """
1497
+
1498
+ def __init__(self,
1499
+ concat_keys: tuple,
1500
+ finetune_keys=("model.diffusion_model.input_blocks.0.0.weight",
1501
+ "model_ema.diffusion_modelinput_blocks00weight"
1502
+ ),
1503
+ keep_finetune_dims=4,
1504
+ # if model was trained without concat mode before and we would like to keep these channels
1505
+ c_concat_log_start=None, # to log reconstruction of c_concat codes
1506
+ c_concat_log_end=None,
1507
+ *args, **kwargs
1508
+ ):
1509
+ ckpt_path = kwargs.pop("ckpt_path", None)
1510
+ ignore_keys = kwargs.pop("ignore_keys", list())
1511
+ super().__init__(*args, **kwargs)
1512
+ self.finetune_keys = finetune_keys
1513
+ self.concat_keys = concat_keys
1514
+ self.keep_dims = keep_finetune_dims
1515
+ self.c_concat_log_start = c_concat_log_start
1516
+ self.c_concat_log_end = c_concat_log_end
1517
+ if exists(self.finetune_keys): assert exists(ckpt_path), 'can only finetune from a given checkpoint'
1518
+ if exists(ckpt_path):
1519
+ self.init_from_ckpt(ckpt_path, ignore_keys)
1520
+
1521
+ def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
1522
+ sd = torch.load(path, map_location="cpu")
1523
+ if "state_dict" in list(sd.keys()):
1524
+ sd = sd["state_dict"]
1525
+ keys = list(sd.keys())
1526
+ for k in keys:
1527
+ for ik in ignore_keys:
1528
+ if k.startswith(ik):
1529
+ print("Deleting key {} from state_dict.".format(k))
1530
+ del sd[k]
1531
+
1532
+ # make it explicit, finetune by including extra input channels
1533
+ if exists(self.finetune_keys) and k in self.finetune_keys:
1534
+ new_entry = None
1535
+ for name, param in self.named_parameters():
1536
+ if name in self.finetune_keys:
1537
+ print(
1538
+ f"modifying key '{name}' and keeping its original {self.keep_dims} (channels) dimensions only")
1539
+ new_entry = torch.zeros_like(param) # zero init
1540
+ assert exists(new_entry), 'did not find matching parameter to modify'
1541
+ new_entry[:, :self.keep_dims, ...] = sd[k]
1542
+ sd[k] = new_entry
1543
+
1544
+ missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict(
1545
+ sd, strict=False)
1546
+ print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
1547
+ if len(missing) > 0:
1548
+ print(f"Missing Keys: {missing}")
1549
+ if len(unexpected) > 0:
1550
+ print(f"Unexpected Keys: {unexpected}")
1551
+
1552
+ @torch.no_grad()
1553
+ def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=200, ddim_eta=1., return_keys=None,
1554
+ quantize_denoised=True, inpaint=True, plot_denoise_rows=False, plot_progressive_rows=True,
1555
+ plot_diffusion_rows=True, unconditional_guidance_scale=1., unconditional_guidance_label=None,
1556
+ use_ema_scope=True,
1557
+ **kwargs):
1558
+ ema_scope = self.ema_scope if use_ema_scope else nullcontext
1559
+ use_ddim = ddim_steps is not None
1560
+
1561
+ log = dict()
1562
+ z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key, bs=N, return_first_stage_outputs=True)
1563
+ c_cat, c = c["c_concat"][0], c["c_crossattn"][0]
1564
+ N = min(x.shape[0], N)
1565
+ n_row = min(x.shape[0], n_row)
1566
+ log["inputs"] = x
1567
+ log["reconstruction"] = xrec
1568
+ if self.model.conditioning_key is not None:
1569
+ if hasattr(self.cond_stage_model, "decode"):
1570
+ xc = self.cond_stage_model.decode(c)
1571
+ log["conditioning"] = xc
1572
+ elif self.cond_stage_key in ["caption", "txt"]:
1573
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch[self.cond_stage_key], size=x.shape[2] // 25)
1574
+ log["conditioning"] = xc
1575
+ elif self.cond_stage_key in ['class_label', 'cls']:
1576
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"], size=x.shape[2] // 25)
1577
+ log['conditioning'] = xc
1578
+ elif isimage(xc):
1579
+ log["conditioning"] = xc
1580
+ if ismap(xc):
1581
+ log["original_conditioning"] = self.to_rgb(xc)
1582
+
1583
+ if not (self.c_concat_log_start is None and self.c_concat_log_end is None):
1584
+ log["c_concat_decoded"] = self.decode_first_stage(c_cat[:, self.c_concat_log_start:self.c_concat_log_end])
1585
+
1586
+ if plot_diffusion_rows:
1587
+ # get diffusion row
1588
+ diffusion_row = list()
1589
+ z_start = z[:n_row]
1590
+ for t in range(self.num_timesteps):
1591
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
1592
+ t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
1593
+ t = t.to(self.device).long()
1594
+ noise = torch.randn_like(z_start)
1595
+ z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
1596
+ diffusion_row.append(self.decode_first_stage(z_noisy))
1597
+
1598
+ diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W
1599
+ diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w')
1600
+ diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w')
1601
+ diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0])
1602
+ log["diffusion_row"] = diffusion_grid
1603
+
1604
+ if sample:
1605
+ # get denoise row
1606
+ with ema_scope("Sampling"):
1607
+ samples, z_denoise_row = self.sample_log(cond={"c_concat": [c_cat], "c_crossattn": [c]},
1608
+ batch_size=N, ddim=use_ddim,
1609
+ ddim_steps=ddim_steps, eta=ddim_eta)
1610
+ # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True)
1611
+ x_samples = self.decode_first_stage(samples)
1612
+ log["samples"] = x_samples
1613
+ if plot_denoise_rows:
1614
+ denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
1615
+ log["denoise_row"] = denoise_grid
1616
+
1617
+ if unconditional_guidance_scale > 1.0:
1618
+ uc_cross = self.get_unconditional_conditioning(N, unconditional_guidance_label)
1619
+ uc_cat = c_cat
1620
+ uc_full = {"c_concat": [uc_cat], "c_crossattn": [uc_cross]}
1621
+ with ema_scope("Sampling with classifier-free guidance"):
1622
+ samples_cfg, _ = self.sample_log(cond={"c_concat": [c_cat], "c_crossattn": [c]},
1623
+ batch_size=N, ddim=use_ddim,
1624
+ ddim_steps=ddim_steps, eta=ddim_eta,
1625
+ unconditional_guidance_scale=unconditional_guidance_scale,
1626
+ unconditional_conditioning=uc_full,
1627
+ )
1628
+ x_samples_cfg = self.decode_first_stage(samples_cfg)
1629
+ log[f"samples_cfg_scale_{unconditional_guidance_scale:.2f}"] = x_samples_cfg
1630
+
1631
+ return log
1632
+
1633
+
1634
+ class LatentInpaintDiffusion(LatentFinetuneDiffusion):
1635
+ """
1636
+ can either run as pure inpainting model (only concat mode) or with mixed conditionings,
1637
+ e.g. mask as concat and text via cross-attn.
1638
+ To disable finetuning mode, set finetune_keys to None
1639
+ """
1640
+
1641
+ def __init__(self,
1642
+ concat_keys=("mask", "masked_image"),
1643
+ masked_image_key="masked_image",
1644
+ *args, **kwargs
1645
+ ):
1646
+ super().__init__(concat_keys, *args, **kwargs)
1647
+ self.masked_image_key = masked_image_key
1648
+ assert self.masked_image_key in concat_keys
1649
+
1650
+ @torch.no_grad()
1651
+ def get_input(self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False):
1652
+ # note: restricted to non-trainable encoders currently
1653
+ assert not self.cond_stage_trainable, 'trainable cond stages not yet supported for inpainting'
1654
+ z, c, x, xrec, xc = super().get_input(batch, self.first_stage_key, return_first_stage_outputs=True,
1655
+ force_c_encode=True, return_original_cond=True, bs=bs)
1656
+
1657
+ assert exists(self.concat_keys)
1658
+ c_cat = list()
1659
+ for ck in self.concat_keys:
1660
+ cc = rearrange(batch[ck], 'b h w c -> b c h w').to(memory_format=torch.contiguous_format).float()
1661
+ if bs is not None:
1662
+ cc = cc[:bs]
1663
+ cc = cc.to(self.device)
1664
+ bchw = z.shape
1665
+ if ck != self.masked_image_key:
1666
+ cc = torch.nn.functional.interpolate(cc, size=bchw[-2:])
1667
+ else:
1668
+ cc = self.get_first_stage_encoding(self.encode_first_stage(cc))
1669
+ c_cat.append(cc)
1670
+ c_cat = torch.cat(c_cat, dim=1)
1671
+ all_conds = {"c_concat": [c_cat], "c_crossattn": [c]}
1672
+ if return_first_stage_outputs:
1673
+ return z, all_conds, x, xrec, xc
1674
+ return z, all_conds
1675
+
1676
+ @torch.no_grad()
1677
+ def log_images(self, *args, **kwargs):
1678
+ log = super(LatentInpaintDiffusion, self).log_images(*args, **kwargs)
1679
+ log["masked_image"] = rearrange(args[0]["masked_image"],
1680
+ 'b h w c -> b c h w').to(memory_format=torch.contiguous_format).float()
1681
+ return log
1682
+
1683
+
1684
+ class LatentDepth2ImageDiffusion(LatentFinetuneDiffusion):
1685
+ """
1686
+ condition on monocular depth estimation
1687
+ """
1688
+
1689
+ def __init__(self, depth_stage_config, concat_keys=("midas_in",), *args, **kwargs):
1690
+ super().__init__(concat_keys=concat_keys, *args, **kwargs)
1691
+ self.depth_model = instantiate_from_config(depth_stage_config)
1692
+ self.depth_stage_key = concat_keys[0]
1693
+
1694
+ @torch.no_grad()
1695
+ def get_input(self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False):
1696
+ # note: restricted to non-trainable encoders currently
1697
+ assert not self.cond_stage_trainable, 'trainable cond stages not yet supported for depth2img'
1698
+ z, c, x, xrec, xc = super().get_input(batch, self.first_stage_key, return_first_stage_outputs=True,
1699
+ force_c_encode=True, return_original_cond=True, bs=bs)
1700
+
1701
+ assert exists(self.concat_keys)
1702
+ assert len(self.concat_keys) == 1
1703
+ c_cat = list()
1704
+ for ck in self.concat_keys:
1705
+ cc = batch[ck]
1706
+ if bs is not None:
1707
+ cc = cc[:bs]
1708
+ cc = cc.to(self.device)
1709
+ cc = self.depth_model(cc)
1710
+ cc = torch.nn.functional.interpolate(
1711
+ cc,
1712
+ size=z.shape[2:],
1713
+ mode="bicubic",
1714
+ align_corners=False,
1715
+ )
1716
+
1717
+ depth_min, depth_max = torch.amin(cc, dim=[1, 2, 3], keepdim=True), torch.amax(cc, dim=[1, 2, 3],
1718
+ keepdim=True)
1719
+ cc = 2. * (cc - depth_min) / (depth_max - depth_min + 0.001) - 1.
1720
+ c_cat.append(cc)
1721
+ c_cat = torch.cat(c_cat, dim=1)
1722
+ all_conds = {"c_concat": [c_cat], "c_crossattn": [c]}
1723
+ if return_first_stage_outputs:
1724
+ return z, all_conds, x, xrec, xc
1725
+ return z, all_conds
1726
+
1727
+ @torch.no_grad()
1728
+ def log_images(self, *args, **kwargs):
1729
+ log = super().log_images(*args, **kwargs)
1730
+ depth = self.depth_model(args[0][self.depth_stage_key])
1731
+ depth_min, depth_max = torch.amin(depth, dim=[1, 2, 3], keepdim=True), \
1732
+ torch.amax(depth, dim=[1, 2, 3], keepdim=True)
1733
+ log["depth"] = 2. * (depth - depth_min) / (depth_max - depth_min) - 1.
1734
+ return log
1735
+
1736
+
1737
+ class LatentUpscaleFinetuneDiffusion(LatentFinetuneDiffusion):
1738
+ """
1739
+ condition on low-res image (and optionally on some spatial noise augmentation)
1740
+ """
1741
+ def __init__(self, concat_keys=("lr",), reshuffle_patch_size=None,
1742
+ low_scale_config=None, low_scale_key=None, *args, **kwargs):
1743
+ super().__init__(concat_keys=concat_keys, *args, **kwargs)
1744
+ self.reshuffle_patch_size = reshuffle_patch_size
1745
+ self.low_scale_model = None
1746
+ if low_scale_config is not None:
1747
+ print("Initializing a low-scale model")
1748
+ assert exists(low_scale_key)
1749
+ self.instantiate_low_stage(low_scale_config)
1750
+ self.low_scale_key = low_scale_key
1751
+
1752
+ def instantiate_low_stage(self, config):
1753
+ model = instantiate_from_config(config)
1754
+ self.low_scale_model = model.eval()
1755
+ self.low_scale_model.train = disabled_train
1756
+ for param in self.low_scale_model.parameters():
1757
+ param.requires_grad = False
1758
+
1759
+ @torch.no_grad()
1760
+ def get_input(self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False):
1761
+ # note: restricted to non-trainable encoders currently
1762
+ assert not self.cond_stage_trainable, 'trainable cond stages not yet supported for upscaling-ft'
1763
+ z, c, x, xrec, xc = super().get_input(batch, self.first_stage_key, return_first_stage_outputs=True,
1764
+ force_c_encode=True, return_original_cond=True, bs=bs)
1765
+
1766
+ assert exists(self.concat_keys)
1767
+ assert len(self.concat_keys) == 1
1768
+ # optionally make spatial noise_level here
1769
+ c_cat = list()
1770
+ noise_level = None
1771
+ for ck in self.concat_keys:
1772
+ cc = batch[ck]
1773
+ cc = rearrange(cc, 'b h w c -> b c h w')
1774
+ if exists(self.reshuffle_patch_size):
1775
+ assert isinstance(self.reshuffle_patch_size, int)
1776
+ cc = rearrange(cc, 'b c (p1 h) (p2 w) -> b (p1 p2 c) h w',
1777
+ p1=self.reshuffle_patch_size, p2=self.reshuffle_patch_size)
1778
+ if bs is not None:
1779
+ cc = cc[:bs]
1780
+ cc = cc.to(self.device)
1781
+ if exists(self.low_scale_model) and ck == self.low_scale_key:
1782
+ cc, noise_level = self.low_scale_model(cc)
1783
+ c_cat.append(cc)
1784
+ c_cat = torch.cat(c_cat, dim=1)
1785
+ if exists(noise_level):
1786
+ all_conds = {"c_concat": [c_cat], "c_crossattn": [c], "c_adm": noise_level}
1787
+ else:
1788
+ all_conds = {"c_concat": [c_cat], "c_crossattn": [c]}
1789
+ if return_first_stage_outputs:
1790
+ return z, all_conds, x, xrec, xc
1791
+ return z, all_conds
1792
+
1793
+ @torch.no_grad()
1794
+ def log_images(self, *args, **kwargs):
1795
+ log = super().log_images(*args, **kwargs)
1796
+ log["lr"] = rearrange(args[0]["lr"], 'b h w c -> b c h w')
1797
+ return log
ControlNet/ldm/models/diffusion/dpm_solver/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ from .sampler import DPMSolverSampler
ControlNet/ldm/models/diffusion/dpm_solver/dpm_solver.py ADDED
@@ -0,0 +1,1154 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn.functional as F
3
+ import math
4
+ from tqdm import tqdm
5
+
6
+
7
+ class NoiseScheduleVP:
8
+ def __init__(
9
+ self,
10
+ schedule='discrete',
11
+ betas=None,
12
+ alphas_cumprod=None,
13
+ continuous_beta_0=0.1,
14
+ continuous_beta_1=20.,
15
+ ):
16
+ """Create a wrapper class for the forward SDE (VP type).
17
+ ***
18
+ Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.
19
+ We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.
20
+ ***
21
+ The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
22
+ We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
23
+ Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:
24
+ log_alpha_t = self.marginal_log_mean_coeff(t)
25
+ sigma_t = self.marginal_std(t)
26
+ lambda_t = self.marginal_lambda(t)
27
+ Moreover, as lambda(t) is an invertible function, we also support its inverse function:
28
+ t = self.inverse_lambda(lambda_t)
29
+ ===============================================================
30
+ We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).
31
+ 1. For discrete-time DPMs:
32
+ For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:
33
+ t_i = (i + 1) / N
34
+ e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.
35
+ We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.
36
+ Args:
37
+ betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)
38
+ alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)
39
+ Note that we always have alphas_cumprod = cumprod(betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.
40
+ **Important**: Please pay special attention for the args for `alphas_cumprod`:
41
+ The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that
42
+ q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).
43
+ Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have
44
+ alpha_{t_n} = \sqrt{\hat{alpha_n}},
45
+ and
46
+ log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).
47
+ 2. For continuous-time DPMs:
48
+ We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise
49
+ schedule are the default settings in DDPM and improved-DDPM:
50
+ Args:
51
+ beta_min: A `float` number. The smallest beta for the linear schedule.
52
+ beta_max: A `float` number. The largest beta for the linear schedule.
53
+ cosine_s: A `float` number. The hyperparameter in the cosine schedule.
54
+ cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.
55
+ T: A `float` number. The ending time of the forward process.
56
+ ===============================================================
57
+ Args:
58
+ schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,
59
+ 'linear' or 'cosine' for continuous-time DPMs.
60
+ Returns:
61
+ A wrapper object of the forward SDE (VP type).
62
+
63
+ ===============================================================
64
+ Example:
65
+ # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):
66
+ >>> ns = NoiseScheduleVP('discrete', betas=betas)
67
+ # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):
68
+ >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
69
+ # For continuous-time DPMs (VPSDE), linear schedule:
70
+ >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)
71
+ """
72
+
73
+ if schedule not in ['discrete', 'linear', 'cosine']:
74
+ raise ValueError(
75
+ "Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(
76
+ schedule))
77
+
78
+ self.schedule = schedule
79
+ if schedule == 'discrete':
80
+ if betas is not None:
81
+ log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)
82
+ else:
83
+ assert alphas_cumprod is not None
84
+ log_alphas = 0.5 * torch.log(alphas_cumprod)
85
+ self.total_N = len(log_alphas)
86
+ self.T = 1.
87
+ self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1))
88
+ self.log_alpha_array = log_alphas.reshape((1, -1,))
89
+ else:
90
+ self.total_N = 1000
91
+ self.beta_0 = continuous_beta_0
92
+ self.beta_1 = continuous_beta_1
93
+ self.cosine_s = 0.008
94
+ self.cosine_beta_max = 999.
95
+ self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi) * 2. * (
96
+ 1. + self.cosine_s) / math.pi - self.cosine_s
97
+ self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1. + self.cosine_s) * math.pi / 2.))
98
+ self.schedule = schedule
99
+ if schedule == 'cosine':
100
+ # For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T.
101
+ # Note that T = 0.9946 may be not the optimal setting. However, we find it works well.
102
+ self.T = 0.9946
103
+ else:
104
+ self.T = 1.
105
+
106
+ def marginal_log_mean_coeff(self, t):
107
+ """
108
+ Compute log(alpha_t) of a given continuous-time label t in [0, T].
109
+ """
110
+ if self.schedule == 'discrete':
111
+ return interpolate_fn(t.reshape((-1, 1)), self.t_array.to(t.device),
112
+ self.log_alpha_array.to(t.device)).reshape((-1))
113
+ elif self.schedule == 'linear':
114
+ return -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
115
+ elif self.schedule == 'cosine':
116
+ log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.))
117
+ log_alpha_t = log_alpha_fn(t) - self.cosine_log_alpha_0
118
+ return log_alpha_t
119
+
120
+ def marginal_alpha(self, t):
121
+ """
122
+ Compute alpha_t of a given continuous-time label t in [0, T].
123
+ """
124
+ return torch.exp(self.marginal_log_mean_coeff(t))
125
+
126
+ def marginal_std(self, t):
127
+ """
128
+ Compute sigma_t of a given continuous-time label t in [0, T].
129
+ """
130
+ return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t)))
131
+
132
+ def marginal_lambda(self, t):
133
+ """
134
+ Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].
135
+ """
136
+ log_mean_coeff = self.marginal_log_mean_coeff(t)
137
+ log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff))
138
+ return log_mean_coeff - log_std
139
+
140
+ def inverse_lambda(self, lamb):
141
+ """
142
+ Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.
143
+ """
144
+ if self.schedule == 'linear':
145
+ tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
146
+ Delta = self.beta_0 ** 2 + tmp
147
+ return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0)
148
+ elif self.schedule == 'discrete':
149
+ log_alpha = -0.5 * torch.logaddexp(torch.zeros((1,)).to(lamb.device), -2. * lamb)
150
+ t = interpolate_fn(log_alpha.reshape((-1, 1)), torch.flip(self.log_alpha_array.to(lamb.device), [1]),
151
+ torch.flip(self.t_array.to(lamb.device), [1]))
152
+ return t.reshape((-1,))
153
+ else:
154
+ log_alpha = -0.5 * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
155
+ t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0)) * 2. * (
156
+ 1. + self.cosine_s) / math.pi - self.cosine_s
157
+ t = t_fn(log_alpha)
158
+ return t
159
+
160
+
161
+ def model_wrapper(
162
+ model,
163
+ noise_schedule,
164
+ model_type="noise",
165
+ model_kwargs={},
166
+ guidance_type="uncond",
167
+ condition=None,
168
+ unconditional_condition=None,
169
+ guidance_scale=1.,
170
+ classifier_fn=None,
171
+ classifier_kwargs={},
172
+ ):
173
+ """Create a wrapper function for the noise prediction model.
174
+ DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to
175
+ firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.
176
+ We support four types of the diffusion model by setting `model_type`:
177
+ 1. "noise": noise prediction model. (Trained by predicting noise).
178
+ 2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).
179
+ 3. "v": velocity prediction model. (Trained by predicting the velocity).
180
+ The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].
181
+ [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."
182
+ arXiv preprint arXiv:2202.00512 (2022).
183
+ [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."
184
+ arXiv preprint arXiv:2210.02303 (2022).
185
+
186
+ 4. "score": marginal score function. (Trained by denoising score matching).
187
+ Note that the score function and the noise prediction model follows a simple relationship:
188
+ ```
189
+ noise(x_t, t) = -sigma_t * score(x_t, t)
190
+ ```
191
+ We support three types of guided sampling by DPMs by setting `guidance_type`:
192
+ 1. "uncond": unconditional sampling by DPMs.
193
+ The input `model` has the following format:
194
+ ``
195
+ model(x, t_input, **model_kwargs) -> noise | x_start | v | score
196
+ ``
197
+ 2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.
198
+ The input `model` has the following format:
199
+ ``
200
+ model(x, t_input, **model_kwargs) -> noise | x_start | v | score
201
+ ``
202
+ The input `classifier_fn` has the following format:
203
+ ``
204
+ classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)
205
+ ``
206
+ [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"
207
+ in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.
208
+ 3. "classifier-free": classifier-free guidance sampling by conditional DPMs.
209
+ The input `model` has the following format:
210
+ ``
211
+ model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score
212
+ ``
213
+ And if cond == `unconditional_condition`, the model output is the unconditional DPM output.
214
+ [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."
215
+ arXiv preprint arXiv:2207.12598 (2022).
216
+
217
+ The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999)
218
+ or continuous-time labels (i.e. epsilon to T).
219
+ We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise:
220
+ ``
221
+ def model_fn(x, t_continuous) -> noise:
222
+ t_input = get_model_input_time(t_continuous)
223
+ return noise_pred(model, x, t_input, **model_kwargs)
224
+ ``
225
+ where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver.
226
+ ===============================================================
227
+ Args:
228
+ model: A diffusion model with the corresponding format described above.
229
+ noise_schedule: A noise schedule object, such as NoiseScheduleVP.
230
+ model_type: A `str`. The parameterization type of the diffusion model.
231
+ "noise" or "x_start" or "v" or "score".
232
+ model_kwargs: A `dict`. A dict for the other inputs of the model function.
233
+ guidance_type: A `str`. The type of the guidance for sampling.
234
+ "uncond" or "classifier" or "classifier-free".
235
+ condition: A pytorch tensor. The condition for the guided sampling.
236
+ Only used for "classifier" or "classifier-free" guidance type.
237
+ unconditional_condition: A pytorch tensor. The condition for the unconditional sampling.
238
+ Only used for "classifier-free" guidance type.
239
+ guidance_scale: A `float`. The scale for the guided sampling.
240
+ classifier_fn: A classifier function. Only used for the classifier guidance.
241
+ classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function.
242
+ Returns:
243
+ A noise prediction model that accepts the noised data and the continuous time as the inputs.
244
+ """
245
+
246
+ def get_model_input_time(t_continuous):
247
+ """
248
+ Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.
249
+ For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].
250
+ For continuous-time DPMs, we just use `t_continuous`.
251
+ """
252
+ if noise_schedule.schedule == 'discrete':
253
+ return (t_continuous - 1. / noise_schedule.total_N) * 1000.
254
+ else:
255
+ return t_continuous
256
+
257
+ def noise_pred_fn(x, t_continuous, cond=None):
258
+ if t_continuous.reshape((-1,)).shape[0] == 1:
259
+ t_continuous = t_continuous.expand((x.shape[0]))
260
+ t_input = get_model_input_time(t_continuous)
261
+ if cond is None:
262
+ output = model(x, t_input, **model_kwargs)
263
+ else:
264
+ output = model(x, t_input, cond, **model_kwargs)
265
+ if model_type == "noise":
266
+ return output
267
+ elif model_type == "x_start":
268
+ alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
269
+ dims = x.dim()
270
+ return (x - expand_dims(alpha_t, dims) * output) / expand_dims(sigma_t, dims)
271
+ elif model_type == "v":
272
+ alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
273
+ dims = x.dim()
274
+ return expand_dims(alpha_t, dims) * output + expand_dims(sigma_t, dims) * x
275
+ elif model_type == "score":
276
+ sigma_t = noise_schedule.marginal_std(t_continuous)
277
+ dims = x.dim()
278
+ return -expand_dims(sigma_t, dims) * output
279
+
280
+ def cond_grad_fn(x, t_input):
281
+ """
282
+ Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
283
+ """
284
+ with torch.enable_grad():
285
+ x_in = x.detach().requires_grad_(True)
286
+ log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs)
287
+ return torch.autograd.grad(log_prob.sum(), x_in)[0]
288
+
289
+ def model_fn(x, t_continuous):
290
+ """
291
+ The noise predicition model function that is used for DPM-Solver.
292
+ """
293
+ if t_continuous.reshape((-1,)).shape[0] == 1:
294
+ t_continuous = t_continuous.expand((x.shape[0]))
295
+ if guidance_type == "uncond":
296
+ return noise_pred_fn(x, t_continuous)
297
+ elif guidance_type == "classifier":
298
+ assert classifier_fn is not None
299
+ t_input = get_model_input_time(t_continuous)
300
+ cond_grad = cond_grad_fn(x, t_input)
301
+ sigma_t = noise_schedule.marginal_std(t_continuous)
302
+ noise = noise_pred_fn(x, t_continuous)
303
+ return noise - guidance_scale * expand_dims(sigma_t, dims=cond_grad.dim()) * cond_grad
304
+ elif guidance_type == "classifier-free":
305
+ if guidance_scale == 1. or unconditional_condition is None:
306
+ return noise_pred_fn(x, t_continuous, cond=condition)
307
+ else:
308
+ x_in = torch.cat([x] * 2)
309
+ t_in = torch.cat([t_continuous] * 2)
310
+ c_in = torch.cat([unconditional_condition, condition])
311
+ noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
312
+ return noise_uncond + guidance_scale * (noise - noise_uncond)
313
+
314
+ assert model_type in ["noise", "x_start", "v"]
315
+ assert guidance_type in ["uncond", "classifier", "classifier-free"]
316
+ return model_fn
317
+
318
+
319
+ class DPM_Solver:
320
+ def __init__(self, model_fn, noise_schedule, predict_x0=False, thresholding=False, max_val=1.):
321
+ """Construct a DPM-Solver.
322
+ We support both the noise prediction model ("predicting epsilon") and the data prediction model ("predicting x0").
323
+ If `predict_x0` is False, we use the solver for the noise prediction model (DPM-Solver).
324
+ If `predict_x0` is True, we use the solver for the data prediction model (DPM-Solver++).
325
+ In such case, we further support the "dynamic thresholding" in [1] when `thresholding` is True.
326
+ The "dynamic thresholding" can greatly improve the sample quality for pixel-space DPMs with large guidance scales.
327
+ Args:
328
+ model_fn: A noise prediction model function which accepts the continuous-time input (t in [epsilon, T]):
329
+ ``
330
+ def model_fn(x, t_continuous):
331
+ return noise
332
+ ``
333
+ noise_schedule: A noise schedule object, such as NoiseScheduleVP.
334
+ predict_x0: A `bool`. If true, use the data prediction model; else, use the noise prediction model.
335
+ thresholding: A `bool`. Valid when `predict_x0` is True. Whether to use the "dynamic thresholding" in [1].
336
+ max_val: A `float`. Valid when both `predict_x0` and `thresholding` are True. The max value for thresholding.
337
+
338
+ [1] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487, 2022b.
339
+ """
340
+ self.model = model_fn
341
+ self.noise_schedule = noise_schedule
342
+ self.predict_x0 = predict_x0
343
+ self.thresholding = thresholding
344
+ self.max_val = max_val
345
+
346
+ def noise_prediction_fn(self, x, t):
347
+ """
348
+ Return the noise prediction model.
349
+ """
350
+ return self.model(x, t)
351
+
352
+ def data_prediction_fn(self, x, t):
353
+ """
354
+ Return the data prediction model (with thresholding).
355
+ """
356
+ noise = self.noise_prediction_fn(x, t)
357
+ dims = x.dim()
358
+ alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
359
+ x0 = (x - expand_dims(sigma_t, dims) * noise) / expand_dims(alpha_t, dims)
360
+ if self.thresholding:
361
+ p = 0.995 # A hyperparameter in the paper of "Imagen" [1].
362
+ s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
363
+ s = expand_dims(torch.maximum(s, self.max_val * torch.ones_like(s).to(s.device)), dims)
364
+ x0 = torch.clamp(x0, -s, s) / s
365
+ return x0
366
+
367
+ def model_fn(self, x, t):
368
+ """
369
+ Convert the model to the noise prediction model or the data prediction model.
370
+ """
371
+ if self.predict_x0:
372
+ return self.data_prediction_fn(x, t)
373
+ else:
374
+ return self.noise_prediction_fn(x, t)
375
+
376
+ def get_time_steps(self, skip_type, t_T, t_0, N, device):
377
+ """Compute the intermediate time steps for sampling.
378
+ Args:
379
+ skip_type: A `str`. The type for the spacing of the time steps. We support three types:
380
+ - 'logSNR': uniform logSNR for the time steps.
381
+ - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
382
+ - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
383
+ t_T: A `float`. The starting time of the sampling (default is T).
384
+ t_0: A `float`. The ending time of the sampling (default is epsilon).
385
+ N: A `int`. The total number of the spacing of the time steps.
386
+ device: A torch device.
387
+ Returns:
388
+ A pytorch tensor of the time steps, with the shape (N + 1,).
389
+ """
390
+ if skip_type == 'logSNR':
391
+ lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
392
+ lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
393
+ logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)
394
+ return self.noise_schedule.inverse_lambda(logSNR_steps)
395
+ elif skip_type == 'time_uniform':
396
+ return torch.linspace(t_T, t_0, N + 1).to(device)
397
+ elif skip_type == 'time_quadratic':
398
+ t_order = 2
399
+ t = torch.linspace(t_T ** (1. / t_order), t_0 ** (1. / t_order), N + 1).pow(t_order).to(device)
400
+ return t
401
+ else:
402
+ raise ValueError(
403
+ "Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type))
404
+
405
+ def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device):
406
+ """
407
+ Get the order of each step for sampling by the singlestep DPM-Solver.
408
+ We combine both DPM-Solver-1,2,3 to use all the function evaluations, which is named as "DPM-Solver-fast".
409
+ Given a fixed number of function evaluations by `steps`, the sampling procedure by DPM-Solver-fast is:
410
+ - If order == 1:
411
+ We take `steps` of DPM-Solver-1 (i.e. DDIM).
412
+ - If order == 2:
413
+ - Denote K = (steps // 2). We take K or (K + 1) intermediate time steps for sampling.
414
+ - If steps % 2 == 0, we use K steps of DPM-Solver-2.
415
+ - If steps % 2 == 1, we use K steps of DPM-Solver-2 and 1 step of DPM-Solver-1.
416
+ - If order == 3:
417
+ - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
418
+ - If steps % 3 == 0, we use (K - 2) steps of DPM-Solver-3, and 1 step of DPM-Solver-2 and 1 step of DPM-Solver-1.
419
+ - If steps % 3 == 1, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-1.
420
+ - If steps % 3 == 2, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-2.
421
+ ============================================
422
+ Args:
423
+ order: A `int`. The max order for the solver (2 or 3).
424
+ steps: A `int`. The total number of function evaluations (NFE).
425
+ skip_type: A `str`. The type for the spacing of the time steps. We support three types:
426
+ - 'logSNR': uniform logSNR for the time steps.
427
+ - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
428
+ - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
429
+ t_T: A `float`. The starting time of the sampling (default is T).
430
+ t_0: A `float`. The ending time of the sampling (default is epsilon).
431
+ device: A torch device.
432
+ Returns:
433
+ orders: A list of the solver order of each step.
434
+ """
435
+ if order == 3:
436
+ K = steps // 3 + 1
437
+ if steps % 3 == 0:
438
+ orders = [3, ] * (K - 2) + [2, 1]
439
+ elif steps % 3 == 1:
440
+ orders = [3, ] * (K - 1) + [1]
441
+ else:
442
+ orders = [3, ] * (K - 1) + [2]
443
+ elif order == 2:
444
+ if steps % 2 == 0:
445
+ K = steps // 2
446
+ orders = [2, ] * K
447
+ else:
448
+ K = steps // 2 + 1
449
+ orders = [2, ] * (K - 1) + [1]
450
+ elif order == 1:
451
+ K = 1
452
+ orders = [1, ] * steps
453
+ else:
454
+ raise ValueError("'order' must be '1' or '2' or '3'.")
455
+ if skip_type == 'logSNR':
456
+ # To reproduce the results in DPM-Solver paper
457
+ timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device)
458
+ else:
459
+ timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[
460
+ torch.cumsum(torch.tensor([0, ] + orders)).to(device)]
461
+ return timesteps_outer, orders
462
+
463
+ def denoise_to_zero_fn(self, x, s):
464
+ """
465
+ Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization.
466
+ """
467
+ return self.data_prediction_fn(x, s)
468
+
469
+ def dpm_solver_first_update(self, x, s, t, model_s=None, return_intermediate=False):
470
+ """
471
+ DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`.
472
+ Args:
473
+ x: A pytorch tensor. The initial value at time `s`.
474
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
475
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
476
+ model_s: A pytorch tensor. The model function evaluated at time `s`.
477
+ If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
478
+ return_intermediate: A `bool`. If true, also return the model value at time `s`.
479
+ Returns:
480
+ x_t: A pytorch tensor. The approximated solution at time `t`.
481
+ """
482
+ ns = self.noise_schedule
483
+ dims = x.dim()
484
+ lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
485
+ h = lambda_t - lambda_s
486
+ log_alpha_s, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(t)
487
+ sigma_s, sigma_t = ns.marginal_std(s), ns.marginal_std(t)
488
+ alpha_t = torch.exp(log_alpha_t)
489
+
490
+ if self.predict_x0:
491
+ phi_1 = torch.expm1(-h)
492
+ if model_s is None:
493
+ model_s = self.model_fn(x, s)
494
+ x_t = (
495
+ expand_dims(sigma_t / sigma_s, dims) * x
496
+ - expand_dims(alpha_t * phi_1, dims) * model_s
497
+ )
498
+ if return_intermediate:
499
+ return x_t, {'model_s': model_s}
500
+ else:
501
+ return x_t
502
+ else:
503
+ phi_1 = torch.expm1(h)
504
+ if model_s is None:
505
+ model_s = self.model_fn(x, s)
506
+ x_t = (
507
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
508
+ - expand_dims(sigma_t * phi_1, dims) * model_s
509
+ )
510
+ if return_intermediate:
511
+ return x_t, {'model_s': model_s}
512
+ else:
513
+ return x_t
514
+
515
+ def singlestep_dpm_solver_second_update(self, x, s, t, r1=0.5, model_s=None, return_intermediate=False,
516
+ solver_type='dpm_solver'):
517
+ """
518
+ Singlestep solver DPM-Solver-2 from time `s` to time `t`.
519
+ Args:
520
+ x: A pytorch tensor. The initial value at time `s`.
521
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
522
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
523
+ r1: A `float`. The hyperparameter of the second-order solver.
524
+ model_s: A pytorch tensor. The model function evaluated at time `s`.
525
+ If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
526
+ return_intermediate: A `bool`. If true, also return the model value at time `s` and `s1` (the intermediate time).
527
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
528
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
529
+ Returns:
530
+ x_t: A pytorch tensor. The approximated solution at time `t`.
531
+ """
532
+ if solver_type not in ['dpm_solver', 'taylor']:
533
+ raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
534
+ if r1 is None:
535
+ r1 = 0.5
536
+ ns = self.noise_schedule
537
+ dims = x.dim()
538
+ lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
539
+ h = lambda_t - lambda_s
540
+ lambda_s1 = lambda_s + r1 * h
541
+ s1 = ns.inverse_lambda(lambda_s1)
542
+ log_alpha_s, log_alpha_s1, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(
543
+ s1), ns.marginal_log_mean_coeff(t)
544
+ sigma_s, sigma_s1, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(t)
545
+ alpha_s1, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_t)
546
+
547
+ if self.predict_x0:
548
+ phi_11 = torch.expm1(-r1 * h)
549
+ phi_1 = torch.expm1(-h)
550
+
551
+ if model_s is None:
552
+ model_s = self.model_fn(x, s)
553
+ x_s1 = (
554
+ expand_dims(sigma_s1 / sigma_s, dims) * x
555
+ - expand_dims(alpha_s1 * phi_11, dims) * model_s
556
+ )
557
+ model_s1 = self.model_fn(x_s1, s1)
558
+ if solver_type == 'dpm_solver':
559
+ x_t = (
560
+ expand_dims(sigma_t / sigma_s, dims) * x
561
+ - expand_dims(alpha_t * phi_1, dims) * model_s
562
+ - (0.5 / r1) * expand_dims(alpha_t * phi_1, dims) * (model_s1 - model_s)
563
+ )
564
+ elif solver_type == 'taylor':
565
+ x_t = (
566
+ expand_dims(sigma_t / sigma_s, dims) * x
567
+ - expand_dims(alpha_t * phi_1, dims) * model_s
568
+ + (1. / r1) * expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * (
569
+ model_s1 - model_s)
570
+ )
571
+ else:
572
+ phi_11 = torch.expm1(r1 * h)
573
+ phi_1 = torch.expm1(h)
574
+
575
+ if model_s is None:
576
+ model_s = self.model_fn(x, s)
577
+ x_s1 = (
578
+ expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x
579
+ - expand_dims(sigma_s1 * phi_11, dims) * model_s
580
+ )
581
+ model_s1 = self.model_fn(x_s1, s1)
582
+ if solver_type == 'dpm_solver':
583
+ x_t = (
584
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
585
+ - expand_dims(sigma_t * phi_1, dims) * model_s
586
+ - (0.5 / r1) * expand_dims(sigma_t * phi_1, dims) * (model_s1 - model_s)
587
+ )
588
+ elif solver_type == 'taylor':
589
+ x_t = (
590
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
591
+ - expand_dims(sigma_t * phi_1, dims) * model_s
592
+ - (1. / r1) * expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * (model_s1 - model_s)
593
+ )
594
+ if return_intermediate:
595
+ return x_t, {'model_s': model_s, 'model_s1': model_s1}
596
+ else:
597
+ return x_t
598
+
599
+ def singlestep_dpm_solver_third_update(self, x, s, t, r1=1. / 3., r2=2. / 3., model_s=None, model_s1=None,
600
+ return_intermediate=False, solver_type='dpm_solver'):
601
+ """
602
+ Singlestep solver DPM-Solver-3 from time `s` to time `t`.
603
+ Args:
604
+ x: A pytorch tensor. The initial value at time `s`.
605
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
606
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
607
+ r1: A `float`. The hyperparameter of the third-order solver.
608
+ r2: A `float`. The hyperparameter of the third-order solver.
609
+ model_s: A pytorch tensor. The model function evaluated at time `s`.
610
+ If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
611
+ model_s1: A pytorch tensor. The model function evaluated at time `s1` (the intermediate time given by `r1`).
612
+ If `model_s1` is None, we evaluate the model at `s1`; otherwise we directly use it.
613
+ return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
614
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
615
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
616
+ Returns:
617
+ x_t: A pytorch tensor. The approximated solution at time `t`.
618
+ """
619
+ if solver_type not in ['dpm_solver', 'taylor']:
620
+ raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
621
+ if r1 is None:
622
+ r1 = 1. / 3.
623
+ if r2 is None:
624
+ r2 = 2. / 3.
625
+ ns = self.noise_schedule
626
+ dims = x.dim()
627
+ lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
628
+ h = lambda_t - lambda_s
629
+ lambda_s1 = lambda_s + r1 * h
630
+ lambda_s2 = lambda_s + r2 * h
631
+ s1 = ns.inverse_lambda(lambda_s1)
632
+ s2 = ns.inverse_lambda(lambda_s2)
633
+ log_alpha_s, log_alpha_s1, log_alpha_s2, log_alpha_t = ns.marginal_log_mean_coeff(
634
+ s), ns.marginal_log_mean_coeff(s1), ns.marginal_log_mean_coeff(s2), ns.marginal_log_mean_coeff(t)
635
+ sigma_s, sigma_s1, sigma_s2, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(
636
+ s2), ns.marginal_std(t)
637
+ alpha_s1, alpha_s2, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_s2), torch.exp(log_alpha_t)
638
+
639
+ if self.predict_x0:
640
+ phi_11 = torch.expm1(-r1 * h)
641
+ phi_12 = torch.expm1(-r2 * h)
642
+ phi_1 = torch.expm1(-h)
643
+ phi_22 = torch.expm1(-r2 * h) / (r2 * h) + 1.
644
+ phi_2 = phi_1 / h + 1.
645
+ phi_3 = phi_2 / h - 0.5
646
+
647
+ if model_s is None:
648
+ model_s = self.model_fn(x, s)
649
+ if model_s1 is None:
650
+ x_s1 = (
651
+ expand_dims(sigma_s1 / sigma_s, dims) * x
652
+ - expand_dims(alpha_s1 * phi_11, dims) * model_s
653
+ )
654
+ model_s1 = self.model_fn(x_s1, s1)
655
+ x_s2 = (
656
+ expand_dims(sigma_s2 / sigma_s, dims) * x
657
+ - expand_dims(alpha_s2 * phi_12, dims) * model_s
658
+ + r2 / r1 * expand_dims(alpha_s2 * phi_22, dims) * (model_s1 - model_s)
659
+ )
660
+ model_s2 = self.model_fn(x_s2, s2)
661
+ if solver_type == 'dpm_solver':
662
+ x_t = (
663
+ expand_dims(sigma_t / sigma_s, dims) * x
664
+ - expand_dims(alpha_t * phi_1, dims) * model_s
665
+ + (1. / r2) * expand_dims(alpha_t * phi_2, dims) * (model_s2 - model_s)
666
+ )
667
+ elif solver_type == 'taylor':
668
+ D1_0 = (1. / r1) * (model_s1 - model_s)
669
+ D1_1 = (1. / r2) * (model_s2 - model_s)
670
+ D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
671
+ D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
672
+ x_t = (
673
+ expand_dims(sigma_t / sigma_s, dims) * x
674
+ - expand_dims(alpha_t * phi_1, dims) * model_s
675
+ + expand_dims(alpha_t * phi_2, dims) * D1
676
+ - expand_dims(alpha_t * phi_3, dims) * D2
677
+ )
678
+ else:
679
+ phi_11 = torch.expm1(r1 * h)
680
+ phi_12 = torch.expm1(r2 * h)
681
+ phi_1 = torch.expm1(h)
682
+ phi_22 = torch.expm1(r2 * h) / (r2 * h) - 1.
683
+ phi_2 = phi_1 / h - 1.
684
+ phi_3 = phi_2 / h - 0.5
685
+
686
+ if model_s is None:
687
+ model_s = self.model_fn(x, s)
688
+ if model_s1 is None:
689
+ x_s1 = (
690
+ expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x
691
+ - expand_dims(sigma_s1 * phi_11, dims) * model_s
692
+ )
693
+ model_s1 = self.model_fn(x_s1, s1)
694
+ x_s2 = (
695
+ expand_dims(torch.exp(log_alpha_s2 - log_alpha_s), dims) * x
696
+ - expand_dims(sigma_s2 * phi_12, dims) * model_s
697
+ - r2 / r1 * expand_dims(sigma_s2 * phi_22, dims) * (model_s1 - model_s)
698
+ )
699
+ model_s2 = self.model_fn(x_s2, s2)
700
+ if solver_type == 'dpm_solver':
701
+ x_t = (
702
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
703
+ - expand_dims(sigma_t * phi_1, dims) * model_s
704
+ - (1. / r2) * expand_dims(sigma_t * phi_2, dims) * (model_s2 - model_s)
705
+ )
706
+ elif solver_type == 'taylor':
707
+ D1_0 = (1. / r1) * (model_s1 - model_s)
708
+ D1_1 = (1. / r2) * (model_s2 - model_s)
709
+ D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
710
+ D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
711
+ x_t = (
712
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
713
+ - expand_dims(sigma_t * phi_1, dims) * model_s
714
+ - expand_dims(sigma_t * phi_2, dims) * D1
715
+ - expand_dims(sigma_t * phi_3, dims) * D2
716
+ )
717
+
718
+ if return_intermediate:
719
+ return x_t, {'model_s': model_s, 'model_s1': model_s1, 'model_s2': model_s2}
720
+ else:
721
+ return x_t
722
+
723
+ def multistep_dpm_solver_second_update(self, x, model_prev_list, t_prev_list, t, solver_type="dpm_solver"):
724
+ """
725
+ Multistep solver DPM-Solver-2 from time `t_prev_list[-1]` to time `t`.
726
+ Args:
727
+ x: A pytorch tensor. The initial value at time `s`.
728
+ model_prev_list: A list of pytorch tensor. The previous computed model values.
729
+ t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
730
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
731
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
732
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
733
+ Returns:
734
+ x_t: A pytorch tensor. The approximated solution at time `t`.
735
+ """
736
+ if solver_type not in ['dpm_solver', 'taylor']:
737
+ raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
738
+ ns = self.noise_schedule
739
+ dims = x.dim()
740
+ model_prev_1, model_prev_0 = model_prev_list
741
+ t_prev_1, t_prev_0 = t_prev_list
742
+ lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_1), ns.marginal_lambda(
743
+ t_prev_0), ns.marginal_lambda(t)
744
+ log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
745
+ sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
746
+ alpha_t = torch.exp(log_alpha_t)
747
+
748
+ h_0 = lambda_prev_0 - lambda_prev_1
749
+ h = lambda_t - lambda_prev_0
750
+ r0 = h_0 / h
751
+ D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)
752
+ if self.predict_x0:
753
+ if solver_type == 'dpm_solver':
754
+ x_t = (
755
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
756
+ - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
757
+ - 0.5 * expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * D1_0
758
+ )
759
+ elif solver_type == 'taylor':
760
+ x_t = (
761
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
762
+ - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
763
+ + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1_0
764
+ )
765
+ else:
766
+ if solver_type == 'dpm_solver':
767
+ x_t = (
768
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
769
+ - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
770
+ - 0.5 * expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * D1_0
771
+ )
772
+ elif solver_type == 'taylor':
773
+ x_t = (
774
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
775
+ - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
776
+ - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1_0
777
+ )
778
+ return x_t
779
+
780
+ def multistep_dpm_solver_third_update(self, x, model_prev_list, t_prev_list, t, solver_type='dpm_solver'):
781
+ """
782
+ Multistep solver DPM-Solver-3 from time `t_prev_list[-1]` to time `t`.
783
+ Args:
784
+ x: A pytorch tensor. The initial value at time `s`.
785
+ model_prev_list: A list of pytorch tensor. The previous computed model values.
786
+ t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
787
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
788
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
789
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
790
+ Returns:
791
+ x_t: A pytorch tensor. The approximated solution at time `t`.
792
+ """
793
+ ns = self.noise_schedule
794
+ dims = x.dim()
795
+ model_prev_2, model_prev_1, model_prev_0 = model_prev_list
796
+ t_prev_2, t_prev_1, t_prev_0 = t_prev_list
797
+ lambda_prev_2, lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_2), ns.marginal_lambda(
798
+ t_prev_1), ns.marginal_lambda(t_prev_0), ns.marginal_lambda(t)
799
+ log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
800
+ sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
801
+ alpha_t = torch.exp(log_alpha_t)
802
+
803
+ h_1 = lambda_prev_1 - lambda_prev_2
804
+ h_0 = lambda_prev_0 - lambda_prev_1
805
+ h = lambda_t - lambda_prev_0
806
+ r0, r1 = h_0 / h, h_1 / h
807
+ D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)
808
+ D1_1 = expand_dims(1. / r1, dims) * (model_prev_1 - model_prev_2)
809
+ D1 = D1_0 + expand_dims(r0 / (r0 + r1), dims) * (D1_0 - D1_1)
810
+ D2 = expand_dims(1. / (r0 + r1), dims) * (D1_0 - D1_1)
811
+ if self.predict_x0:
812
+ x_t = (
813
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
814
+ - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
815
+ + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1
816
+ - expand_dims(alpha_t * ((torch.exp(-h) - 1. + h) / h ** 2 - 0.5), dims) * D2
817
+ )
818
+ else:
819
+ x_t = (
820
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
821
+ - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
822
+ - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1
823
+ - expand_dims(sigma_t * ((torch.exp(h) - 1. - h) / h ** 2 - 0.5), dims) * D2
824
+ )
825
+ return x_t
826
+
827
+ def singlestep_dpm_solver_update(self, x, s, t, order, return_intermediate=False, solver_type='dpm_solver', r1=None,
828
+ r2=None):
829
+ """
830
+ Singlestep DPM-Solver with the order `order` from time `s` to time `t`.
831
+ Args:
832
+ x: A pytorch tensor. The initial value at time `s`.
833
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
834
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
835
+ order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
836
+ return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
837
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
838
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
839
+ r1: A `float`. The hyperparameter of the second-order or third-order solver.
840
+ r2: A `float`. The hyperparameter of the third-order solver.
841
+ Returns:
842
+ x_t: A pytorch tensor. The approximated solution at time `t`.
843
+ """
844
+ if order == 1:
845
+ return self.dpm_solver_first_update(x, s, t, return_intermediate=return_intermediate)
846
+ elif order == 2:
847
+ return self.singlestep_dpm_solver_second_update(x, s, t, return_intermediate=return_intermediate,
848
+ solver_type=solver_type, r1=r1)
849
+ elif order == 3:
850
+ return self.singlestep_dpm_solver_third_update(x, s, t, return_intermediate=return_intermediate,
851
+ solver_type=solver_type, r1=r1, r2=r2)
852
+ else:
853
+ raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))
854
+
855
+ def multistep_dpm_solver_update(self, x, model_prev_list, t_prev_list, t, order, solver_type='dpm_solver'):
856
+ """
857
+ Multistep DPM-Solver with the order `order` from time `t_prev_list[-1]` to time `t`.
858
+ Args:
859
+ x: A pytorch tensor. The initial value at time `s`.
860
+ model_prev_list: A list of pytorch tensor. The previous computed model values.
861
+ t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
862
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
863
+ order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
864
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
865
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
866
+ Returns:
867
+ x_t: A pytorch tensor. The approximated solution at time `t`.
868
+ """
869
+ if order == 1:
870
+ return self.dpm_solver_first_update(x, t_prev_list[-1], t, model_s=model_prev_list[-1])
871
+ elif order == 2:
872
+ return self.multistep_dpm_solver_second_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)
873
+ elif order == 3:
874
+ return self.multistep_dpm_solver_third_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)
875
+ else:
876
+ raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))
877
+
878
+ def dpm_solver_adaptive(self, x, order, t_T, t_0, h_init=0.05, atol=0.0078, rtol=0.05, theta=0.9, t_err=1e-5,
879
+ solver_type='dpm_solver'):
880
+ """
881
+ The adaptive step size solver based on singlestep DPM-Solver.
882
+ Args:
883
+ x: A pytorch tensor. The initial value at time `t_T`.
884
+ order: A `int`. The (higher) order of the solver. We only support order == 2 or 3.
885
+ t_T: A `float`. The starting time of the sampling (default is T).
886
+ t_0: A `float`. The ending time of the sampling (default is epsilon).
887
+ h_init: A `float`. The initial step size (for logSNR).
888
+ atol: A `float`. The absolute tolerance of the solver. For image data, the default setting is 0.0078, followed [1].
889
+ rtol: A `float`. The relative tolerance of the solver. The default setting is 0.05.
890
+ theta: A `float`. The safety hyperparameter for adapting the step size. The default setting is 0.9, followed [1].
891
+ t_err: A `float`. The tolerance for the time. We solve the diffusion ODE until the absolute error between the
892
+ current time and `t_0` is less than `t_err`. The default setting is 1e-5.
893
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
894
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
895
+ Returns:
896
+ x_0: A pytorch tensor. The approximated solution at time `t_0`.
897
+ [1] A. Jolicoeur-Martineau, K. Li, R. Piché-Taillefer, T. Kachman, and I. Mitliagkas, "Gotta go fast when generating data with score-based models," arXiv preprint arXiv:2105.14080, 2021.
898
+ """
899
+ ns = self.noise_schedule
900
+ s = t_T * torch.ones((x.shape[0],)).to(x)
901
+ lambda_s = ns.marginal_lambda(s)
902
+ lambda_0 = ns.marginal_lambda(t_0 * torch.ones_like(s).to(x))
903
+ h = h_init * torch.ones_like(s).to(x)
904
+ x_prev = x
905
+ nfe = 0
906
+ if order == 2:
907
+ r1 = 0.5
908
+ lower_update = lambda x, s, t: self.dpm_solver_first_update(x, s, t, return_intermediate=True)
909
+ higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1,
910
+ solver_type=solver_type,
911
+ **kwargs)
912
+ elif order == 3:
913
+ r1, r2 = 1. / 3., 2. / 3.
914
+ lower_update = lambda x, s, t: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1,
915
+ return_intermediate=True,
916
+ solver_type=solver_type)
917
+ higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_third_update(x, s, t, r1=r1, r2=r2,
918
+ solver_type=solver_type,
919
+ **kwargs)
920
+ else:
921
+ raise ValueError("For adaptive step size solver, order must be 2 or 3, got {}".format(order))
922
+ while torch.abs((s - t_0)).mean() > t_err:
923
+ t = ns.inverse_lambda(lambda_s + h)
924
+ x_lower, lower_noise_kwargs = lower_update(x, s, t)
925
+ x_higher = higher_update(x, s, t, **lower_noise_kwargs)
926
+ delta = torch.max(torch.ones_like(x).to(x) * atol, rtol * torch.max(torch.abs(x_lower), torch.abs(x_prev)))
927
+ norm_fn = lambda v: torch.sqrt(torch.square(v.reshape((v.shape[0], -1))).mean(dim=-1, keepdim=True))
928
+ E = norm_fn((x_higher - x_lower) / delta).max()
929
+ if torch.all(E <= 1.):
930
+ x = x_higher
931
+ s = t
932
+ x_prev = x_lower
933
+ lambda_s = ns.marginal_lambda(s)
934
+ h = torch.min(theta * h * torch.float_power(E, -1. / order).float(), lambda_0 - lambda_s)
935
+ nfe += order
936
+ print('adaptive solver nfe', nfe)
937
+ return x
938
+
939
+ def sample(self, x, steps=20, t_start=None, t_end=None, order=3, skip_type='time_uniform',
940
+ method='singlestep', lower_order_final=True, denoise_to_zero=False, solver_type='dpm_solver',
941
+ atol=0.0078, rtol=0.05,
942
+ ):
943
+ """
944
+ Compute the sample at time `t_end` by DPM-Solver, given the initial `x` at time `t_start`.
945
+ =====================================================
946
+ We support the following algorithms for both noise prediction model and data prediction model:
947
+ - 'singlestep':
948
+ Singlestep DPM-Solver (i.e. "DPM-Solver-fast" in the paper), which combines different orders of singlestep DPM-Solver.
949
+ We combine all the singlestep solvers with order <= `order` to use up all the function evaluations (steps).
950
+ The total number of function evaluations (NFE) == `steps`.
951
+ Given a fixed NFE == `steps`, the sampling procedure is:
952
+ - If `order` == 1:
953
+ - Denote K = steps. We use K steps of DPM-Solver-1 (i.e. DDIM).
954
+ - If `order` == 2:
955
+ - Denote K = (steps // 2) + (steps % 2). We take K intermediate time steps for sampling.
956
+ - If steps % 2 == 0, we use K steps of singlestep DPM-Solver-2.
957
+ - If steps % 2 == 1, we use (K - 1) steps of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.
958
+ - If `order` == 3:
959
+ - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
960
+ - If steps % 3 == 0, we use (K - 2) steps of singlestep DPM-Solver-3, and 1 step of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.
961
+ - If steps % 3 == 1, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of DPM-Solver-1.
962
+ - If steps % 3 == 2, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of singlestep DPM-Solver-2.
963
+ - 'multistep':
964
+ Multistep DPM-Solver with the order of `order`. The total number of function evaluations (NFE) == `steps`.
965
+ We initialize the first `order` values by lower order multistep solvers.
966
+ Given a fixed NFE == `steps`, the sampling procedure is:
967
+ Denote K = steps.
968
+ - If `order` == 1:
969
+ - We use K steps of DPM-Solver-1 (i.e. DDIM).
970
+ - If `order` == 2:
971
+ - We firstly use 1 step of DPM-Solver-1, then use (K - 1) step of multistep DPM-Solver-2.
972
+ - If `order` == 3:
973
+ - We firstly use 1 step of DPM-Solver-1, then 1 step of multistep DPM-Solver-2, then (K - 2) step of multistep DPM-Solver-3.
974
+ - 'singlestep_fixed':
975
+ Fixed order singlestep DPM-Solver (i.e. DPM-Solver-1 or singlestep DPM-Solver-2 or singlestep DPM-Solver-3).
976
+ We use singlestep DPM-Solver-`order` for `order`=1 or 2 or 3, with total [`steps` // `order`] * `order` NFE.
977
+ - 'adaptive':
978
+ Adaptive step size DPM-Solver (i.e. "DPM-Solver-12" and "DPM-Solver-23" in the paper).
979
+ We ignore `steps` and use adaptive step size DPM-Solver with a higher order of `order`.
980
+ You can adjust the absolute tolerance `atol` and the relative tolerance `rtol` to balance the computatation costs
981
+ (NFE) and the sample quality.
982
+ - If `order` == 2, we use DPM-Solver-12 which combines DPM-Solver-1 and singlestep DPM-Solver-2.
983
+ - If `order` == 3, we use DPM-Solver-23 which combines singlestep DPM-Solver-2 and singlestep DPM-Solver-3.
984
+ =====================================================
985
+ Some advices for choosing the algorithm:
986
+ - For **unconditional sampling** or **guided sampling with small guidance scale** by DPMs:
987
+ Use singlestep DPM-Solver ("DPM-Solver-fast" in the paper) with `order = 3`.
988
+ e.g.
989
+ >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=False)
990
+ >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3,
991
+ skip_type='time_uniform', method='singlestep')
992
+ - For **guided sampling with large guidance scale** by DPMs:
993
+ Use multistep DPM-Solver with `predict_x0 = True` and `order = 2`.
994
+ e.g.
995
+ >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=True)
996
+ >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=2,
997
+ skip_type='time_uniform', method='multistep')
998
+ We support three types of `skip_type`:
999
+ - 'logSNR': uniform logSNR for the time steps. **Recommended for low-resolutional images**
1000
+ - 'time_uniform': uniform time for the time steps. **Recommended for high-resolutional images**.
1001
+ - 'time_quadratic': quadratic time for the time steps.
1002
+ =====================================================
1003
+ Args:
1004
+ x: A pytorch tensor. The initial value at time `t_start`
1005
+ e.g. if `t_start` == T, then `x` is a sample from the standard normal distribution.
1006
+ steps: A `int`. The total number of function evaluations (NFE).
1007
+ t_start: A `float`. The starting time of the sampling.
1008
+ If `T` is None, we use self.noise_schedule.T (default is 1.0).
1009
+ t_end: A `float`. The ending time of the sampling.
1010
+ If `t_end` is None, we use 1. / self.noise_schedule.total_N.
1011
+ e.g. if total_N == 1000, we have `t_end` == 1e-3.
1012
+ For discrete-time DPMs:
1013
+ - We recommend `t_end` == 1. / self.noise_schedule.total_N.
1014
+ For continuous-time DPMs:
1015
+ - We recommend `t_end` == 1e-3 when `steps` <= 15; and `t_end` == 1e-4 when `steps` > 15.
1016
+ order: A `int`. The order of DPM-Solver.
1017
+ skip_type: A `str`. The type for the spacing of the time steps. 'time_uniform' or 'logSNR' or 'time_quadratic'.
1018
+ method: A `str`. The method for sampling. 'singlestep' or 'multistep' or 'singlestep_fixed' or 'adaptive'.
1019
+ denoise_to_zero: A `bool`. Whether to denoise to time 0 at the final step.
1020
+ Default is `False`. If `denoise_to_zero` is `True`, the total NFE is (`steps` + 1).
1021
+ This trick is firstly proposed by DDPM (https://arxiv.org/abs/2006.11239) and
1022
+ score_sde (https://arxiv.org/abs/2011.13456). Such trick can improve the FID
1023
+ for diffusion models sampling by diffusion SDEs for low-resolutional images
1024
+ (such as CIFAR-10). However, we observed that such trick does not matter for
1025
+ high-resolutional images. As it needs an additional NFE, we do not recommend
1026
+ it for high-resolutional images.
1027
+ lower_order_final: A `bool`. Whether to use lower order solvers at the final steps.
1028
+ Only valid for `method=multistep` and `steps < 15`. We empirically find that
1029
+ this trick is a key to stabilizing the sampling by DPM-Solver with very few steps
1030
+ (especially for steps <= 10). So we recommend to set it to be `True`.
1031
+ solver_type: A `str`. The taylor expansion type for the solver. `dpm_solver` or `taylor`. We recommend `dpm_solver`.
1032
+ atol: A `float`. The absolute tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
1033
+ rtol: A `float`. The relative tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
1034
+ Returns:
1035
+ x_end: A pytorch tensor. The approximated solution at time `t_end`.
1036
+ """
1037
+ t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end
1038
+ t_T = self.noise_schedule.T if t_start is None else t_start
1039
+ device = x.device
1040
+ if method == 'adaptive':
1041
+ with torch.no_grad():
1042
+ x = self.dpm_solver_adaptive(x, order=order, t_T=t_T, t_0=t_0, atol=atol, rtol=rtol,
1043
+ solver_type=solver_type)
1044
+ elif method == 'multistep':
1045
+ assert steps >= order
1046
+ timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
1047
+ assert timesteps.shape[0] - 1 == steps
1048
+ with torch.no_grad():
1049
+ vec_t = timesteps[0].expand((x.shape[0]))
1050
+ model_prev_list = [self.model_fn(x, vec_t)]
1051
+ t_prev_list = [vec_t]
1052
+ # Init the first `order` values by lower order multistep DPM-Solver.
1053
+ for init_order in tqdm(range(1, order), desc="DPM init order"):
1054
+ vec_t = timesteps[init_order].expand(x.shape[0])
1055
+ x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, init_order,
1056
+ solver_type=solver_type)
1057
+ model_prev_list.append(self.model_fn(x, vec_t))
1058
+ t_prev_list.append(vec_t)
1059
+ # Compute the remaining values by `order`-th order multistep DPM-Solver.
1060
+ for step in tqdm(range(order, steps + 1), desc="DPM multistep"):
1061
+ vec_t = timesteps[step].expand(x.shape[0])
1062
+ if lower_order_final and steps < 15:
1063
+ step_order = min(order, steps + 1 - step)
1064
+ else:
1065
+ step_order = order
1066
+ x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, step_order,
1067
+ solver_type=solver_type)
1068
+ for i in range(order - 1):
1069
+ t_prev_list[i] = t_prev_list[i + 1]
1070
+ model_prev_list[i] = model_prev_list[i + 1]
1071
+ t_prev_list[-1] = vec_t
1072
+ # We do not need to evaluate the final model value.
1073
+ if step < steps:
1074
+ model_prev_list[-1] = self.model_fn(x, vec_t)
1075
+ elif method in ['singlestep', 'singlestep_fixed']:
1076
+ if method == 'singlestep':
1077
+ timesteps_outer, orders = self.get_orders_and_timesteps_for_singlestep_solver(steps=steps, order=order,
1078
+ skip_type=skip_type,
1079
+ t_T=t_T, t_0=t_0,
1080
+ device=device)
1081
+ elif method == 'singlestep_fixed':
1082
+ K = steps // order
1083
+ orders = [order, ] * K
1084
+ timesteps_outer = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=K, device=device)
1085
+ for i, order in enumerate(orders):
1086
+ t_T_inner, t_0_inner = timesteps_outer[i], timesteps_outer[i + 1]
1087
+ timesteps_inner = self.get_time_steps(skip_type=skip_type, t_T=t_T_inner.item(), t_0=t_0_inner.item(),
1088
+ N=order, device=device)
1089
+ lambda_inner = self.noise_schedule.marginal_lambda(timesteps_inner)
1090
+ vec_s, vec_t = t_T_inner.tile(x.shape[0]), t_0_inner.tile(x.shape[0])
1091
+ h = lambda_inner[-1] - lambda_inner[0]
1092
+ r1 = None if order <= 1 else (lambda_inner[1] - lambda_inner[0]) / h
1093
+ r2 = None if order <= 2 else (lambda_inner[2] - lambda_inner[0]) / h
1094
+ x = self.singlestep_dpm_solver_update(x, vec_s, vec_t, order, solver_type=solver_type, r1=r1, r2=r2)
1095
+ if denoise_to_zero:
1096
+ x = self.denoise_to_zero_fn(x, torch.ones((x.shape[0],)).to(device) * t_0)
1097
+ return x
1098
+
1099
+
1100
+ #############################################################
1101
+ # other utility functions
1102
+ #############################################################
1103
+
1104
+ def interpolate_fn(x, xp, yp):
1105
+ """
1106
+ A piecewise linear function y = f(x), using xp and yp as keypoints.
1107
+ We implement f(x) in a differentiable way (i.e. applicable for autograd).
1108
+ The function f(x) is well-defined for all x-axis. (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.)
1109
+ Args:
1110
+ x: PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels (we use C = 1 for DPM-Solver).
1111
+ xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.
1112
+ yp: PyTorch tensor with shape [C, K].
1113
+ Returns:
1114
+ The function values f(x), with shape [N, C].
1115
+ """
1116
+ N, K = x.shape[0], xp.shape[1]
1117
+ all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)
1118
+ sorted_all_x, x_indices = torch.sort(all_x, dim=2)
1119
+ x_idx = torch.argmin(x_indices, dim=2)
1120
+ cand_start_idx = x_idx - 1
1121
+ start_idx = torch.where(
1122
+ torch.eq(x_idx, 0),
1123
+ torch.tensor(1, device=x.device),
1124
+ torch.where(
1125
+ torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
1126
+ ),
1127
+ )
1128
+ end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1)
1129
+ start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)
1130
+ end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)
1131
+ start_idx2 = torch.where(
1132
+ torch.eq(x_idx, 0),
1133
+ torch.tensor(0, device=x.device),
1134
+ torch.where(
1135
+ torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
1136
+ ),
1137
+ )
1138
+ y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)
1139
+ start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2)
1140
+ end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2)
1141
+ cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)
1142
+ return cand
1143
+
1144
+
1145
+ def expand_dims(v, dims):
1146
+ """
1147
+ Expand the tensor `v` to the dim `dims`.
1148
+ Args:
1149
+ `v`: a PyTorch tensor with shape [N].
1150
+ `dim`: a `int`.
1151
+ Returns:
1152
+ a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.
1153
+ """
1154
+ return v[(...,) + (None,) * (dims - 1)]
ControlNet/ldm/models/diffusion/dpm_solver/sampler.py ADDED
@@ -0,0 +1,87 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SAMPLING ONLY."""
2
+ import torch
3
+
4
+ from .dpm_solver import NoiseScheduleVP, model_wrapper, DPM_Solver
5
+
6
+
7
+ MODEL_TYPES = {
8
+ "eps": "noise",
9
+ "v": "v"
10
+ }
11
+
12
+
13
+ class DPMSolverSampler(object):
14
+ def __init__(self, model, **kwargs):
15
+ super().__init__()
16
+ self.model = model
17
+ to_torch = lambda x: x.clone().detach().to(torch.float32).to(model.device)
18
+ self.register_buffer('alphas_cumprod', to_torch(model.alphas_cumprod))
19
+
20
+ def register_buffer(self, name, attr):
21
+ if type(attr) == torch.Tensor:
22
+ if attr.device != torch.device("cuda"):
23
+ attr = attr.to(torch.device("cuda"))
24
+ setattr(self, name, attr)
25
+
26
+ @torch.no_grad()
27
+ def sample(self,
28
+ S,
29
+ batch_size,
30
+ shape,
31
+ conditioning=None,
32
+ callback=None,
33
+ normals_sequence=None,
34
+ img_callback=None,
35
+ quantize_x0=False,
36
+ eta=0.,
37
+ mask=None,
38
+ x0=None,
39
+ temperature=1.,
40
+ noise_dropout=0.,
41
+ score_corrector=None,
42
+ corrector_kwargs=None,
43
+ verbose=True,
44
+ x_T=None,
45
+ log_every_t=100,
46
+ unconditional_guidance_scale=1.,
47
+ unconditional_conditioning=None,
48
+ # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
49
+ **kwargs
50
+ ):
51
+ if conditioning is not None:
52
+ if isinstance(conditioning, dict):
53
+ cbs = conditioning[list(conditioning.keys())[0]].shape[0]
54
+ if cbs != batch_size:
55
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
56
+ else:
57
+ if conditioning.shape[0] != batch_size:
58
+ print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
59
+
60
+ # sampling
61
+ C, H, W = shape
62
+ size = (batch_size, C, H, W)
63
+
64
+ print(f'Data shape for DPM-Solver sampling is {size}, sampling steps {S}')
65
+
66
+ device = self.model.betas.device
67
+ if x_T is None:
68
+ img = torch.randn(size, device=device)
69
+ else:
70
+ img = x_T
71
+
72
+ ns = NoiseScheduleVP('discrete', alphas_cumprod=self.alphas_cumprod)
73
+
74
+ model_fn = model_wrapper(
75
+ lambda x, t, c: self.model.apply_model(x, t, c),
76
+ ns,
77
+ model_type=MODEL_TYPES[self.model.parameterization],
78
+ guidance_type="classifier-free",
79
+ condition=conditioning,
80
+ unconditional_condition=unconditional_conditioning,
81
+ guidance_scale=unconditional_guidance_scale,
82
+ )
83
+
84
+ dpm_solver = DPM_Solver(model_fn, ns, predict_x0=True, thresholding=False)
85
+ x = dpm_solver.sample(img, steps=S, skip_type="time_uniform", method="multistep", order=2, lower_order_final=True)
86
+
87
+ return x.to(device), None
ControlNet/ldm/models/diffusion/plms.py ADDED
@@ -0,0 +1,244 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SAMPLING ONLY."""
2
+
3
+ import torch
4
+ import numpy as np
5
+ from tqdm import tqdm
6
+ from functools import partial
7
+
8
+ from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
9
+ from ldm.models.diffusion.sampling_util import norm_thresholding
10
+
11
+
12
+ class PLMSSampler(object):
13
+ def __init__(self, model, schedule="linear", **kwargs):
14
+ super().__init__()
15
+ self.model = model
16
+ self.ddpm_num_timesteps = model.num_timesteps
17
+ self.schedule = schedule
18
+
19
+ def register_buffer(self, name, attr):
20
+ if type(attr) == torch.Tensor:
21
+ if attr.device != torch.device("cuda"):
22
+ attr = attr.to(torch.device("cuda"))
23
+ setattr(self, name, attr)
24
+
25
+ def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
26
+ if ddim_eta != 0:
27
+ raise ValueError('ddim_eta must be 0 for PLMS')
28
+ self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
29
+ num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose)
30
+ alphas_cumprod = self.model.alphas_cumprod
31
+ assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
32
+ to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)
33
+
34
+ self.register_buffer('betas', to_torch(self.model.betas))
35
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
36
+ self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
37
+
38
+ # calculations for diffusion q(x_t | x_{t-1}) and others
39
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
40
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
41
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
42
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
43
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
44
+
45
+ # ddim sampling parameters
46
+ ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
47
+ ddim_timesteps=self.ddim_timesteps,
48
+ eta=ddim_eta,verbose=verbose)
49
+ self.register_buffer('ddim_sigmas', ddim_sigmas)
50
+ self.register_buffer('ddim_alphas', ddim_alphas)
51
+ self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
52
+ self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
53
+ sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
54
+ (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
55
+ 1 - self.alphas_cumprod / self.alphas_cumprod_prev))
56
+ self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)
57
+
58
+ @torch.no_grad()
59
+ def sample(self,
60
+ S,
61
+ batch_size,
62
+ shape,
63
+ conditioning=None,
64
+ callback=None,
65
+ normals_sequence=None,
66
+ img_callback=None,
67
+ quantize_x0=False,
68
+ eta=0.,
69
+ mask=None,
70
+ x0=None,
71
+ temperature=1.,
72
+ noise_dropout=0.,
73
+ score_corrector=None,
74
+ corrector_kwargs=None,
75
+ verbose=True,
76
+ x_T=None,
77
+ log_every_t=100,
78
+ unconditional_guidance_scale=1.,
79
+ unconditional_conditioning=None,
80
+ # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
81
+ dynamic_threshold=None,
82
+ **kwargs
83
+ ):
84
+ if conditioning is not None:
85
+ if isinstance(conditioning, dict):
86
+ cbs = conditioning[list(conditioning.keys())[0]].shape[0]
87
+ if cbs != batch_size:
88
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
89
+ else:
90
+ if conditioning.shape[0] != batch_size:
91
+ print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
92
+
93
+ self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
94
+ # sampling
95
+ C, H, W = shape
96
+ size = (batch_size, C, H, W)
97
+ print(f'Data shape for PLMS sampling is {size}')
98
+
99
+ samples, intermediates = self.plms_sampling(conditioning, size,
100
+ callback=callback,
101
+ img_callback=img_callback,
102
+ quantize_denoised=quantize_x0,
103
+ mask=mask, x0=x0,
104
+ ddim_use_original_steps=False,
105
+ noise_dropout=noise_dropout,
106
+ temperature=temperature,
107
+ score_corrector=score_corrector,
108
+ corrector_kwargs=corrector_kwargs,
109
+ x_T=x_T,
110
+ log_every_t=log_every_t,
111
+ unconditional_guidance_scale=unconditional_guidance_scale,
112
+ unconditional_conditioning=unconditional_conditioning,
113
+ dynamic_threshold=dynamic_threshold,
114
+ )
115
+ return samples, intermediates
116
+
117
+ @torch.no_grad()
118
+ def plms_sampling(self, cond, shape,
119
+ x_T=None, ddim_use_original_steps=False,
120
+ callback=None, timesteps=None, quantize_denoised=False,
121
+ mask=None, x0=None, img_callback=None, log_every_t=100,
122
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
123
+ unconditional_guidance_scale=1., unconditional_conditioning=None,
124
+ dynamic_threshold=None):
125
+ device = self.model.betas.device
126
+ b = shape[0]
127
+ if x_T is None:
128
+ img = torch.randn(shape, device=device)
129
+ else:
130
+ img = x_T
131
+
132
+ if timesteps is None:
133
+ timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
134
+ elif timesteps is not None and not ddim_use_original_steps:
135
+ subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
136
+ timesteps = self.ddim_timesteps[:subset_end]
137
+
138
+ intermediates = {'x_inter': [img], 'pred_x0': [img]}
139
+ time_range = list(reversed(range(0,timesteps))) if ddim_use_original_steps else np.flip(timesteps)
140
+ total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
141
+ print(f"Running PLMS Sampling with {total_steps} timesteps")
142
+
143
+ iterator = tqdm(time_range, desc='PLMS Sampler', total=total_steps)
144
+ old_eps = []
145
+
146
+ for i, step in enumerate(iterator):
147
+ index = total_steps - i - 1
148
+ ts = torch.full((b,), step, device=device, dtype=torch.long)
149
+ ts_next = torch.full((b,), time_range[min(i + 1, len(time_range) - 1)], device=device, dtype=torch.long)
150
+
151
+ if mask is not None:
152
+ assert x0 is not None
153
+ img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass?
154
+ img = img_orig * mask + (1. - mask) * img
155
+
156
+ outs = self.p_sample_plms(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
157
+ quantize_denoised=quantize_denoised, temperature=temperature,
158
+ noise_dropout=noise_dropout, score_corrector=score_corrector,
159
+ corrector_kwargs=corrector_kwargs,
160
+ unconditional_guidance_scale=unconditional_guidance_scale,
161
+ unconditional_conditioning=unconditional_conditioning,
162
+ old_eps=old_eps, t_next=ts_next,
163
+ dynamic_threshold=dynamic_threshold)
164
+ img, pred_x0, e_t = outs
165
+ old_eps.append(e_t)
166
+ if len(old_eps) >= 4:
167
+ old_eps.pop(0)
168
+ if callback: callback(i)
169
+ if img_callback: img_callback(pred_x0, i)
170
+
171
+ if index % log_every_t == 0 or index == total_steps - 1:
172
+ intermediates['x_inter'].append(img)
173
+ intermediates['pred_x0'].append(pred_x0)
174
+
175
+ return img, intermediates
176
+
177
+ @torch.no_grad()
178
+ def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
179
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
180
+ unconditional_guidance_scale=1., unconditional_conditioning=None, old_eps=None, t_next=None,
181
+ dynamic_threshold=None):
182
+ b, *_, device = *x.shape, x.device
183
+
184
+ def get_model_output(x, t):
185
+ if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
186
+ e_t = self.model.apply_model(x, t, c)
187
+ else:
188
+ x_in = torch.cat([x] * 2)
189
+ t_in = torch.cat([t] * 2)
190
+ c_in = torch.cat([unconditional_conditioning, c])
191
+ e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
192
+ e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
193
+
194
+ if score_corrector is not None:
195
+ assert self.model.parameterization == "eps"
196
+ e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
197
+
198
+ return e_t
199
+
200
+ alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
201
+ alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
202
+ sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
203
+ sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
204
+
205
+ def get_x_prev_and_pred_x0(e_t, index):
206
+ # select parameters corresponding to the currently considered timestep
207
+ a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
208
+ a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
209
+ sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
210
+ sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
211
+
212
+ # current prediction for x_0
213
+ pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
214
+ if quantize_denoised:
215
+ pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
216
+ if dynamic_threshold is not None:
217
+ pred_x0 = norm_thresholding(pred_x0, dynamic_threshold)
218
+ # direction pointing to x_t
219
+ dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
220
+ noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
221
+ if noise_dropout > 0.:
222
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
223
+ x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
224
+ return x_prev, pred_x0
225
+
226
+ e_t = get_model_output(x, t)
227
+ if len(old_eps) == 0:
228
+ # Pseudo Improved Euler (2nd order)
229
+ x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index)
230
+ e_t_next = get_model_output(x_prev, t_next)
231
+ e_t_prime = (e_t + e_t_next) / 2
232
+ elif len(old_eps) == 1:
233
+ # 2nd order Pseudo Linear Multistep (Adams-Bashforth)
234
+ e_t_prime = (3 * e_t - old_eps[-1]) / 2
235
+ elif len(old_eps) == 2:
236
+ # 3nd order Pseudo Linear Multistep (Adams-Bashforth)
237
+ e_t_prime = (23 * e_t - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12
238
+ elif len(old_eps) >= 3:
239
+ # 4nd order Pseudo Linear Multistep (Adams-Bashforth)
240
+ e_t_prime = (55 * e_t - 59 * old_eps[-1] + 37 * old_eps[-2] - 9 * old_eps[-3]) / 24
241
+
242
+ x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index)
243
+
244
+ return x_prev, pred_x0, e_t
ControlNet/ldm/models/diffusion/sampling_util.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+
4
+
5
+ def append_dims(x, target_dims):
6
+ """Appends dimensions to the end of a tensor until it has target_dims dimensions.
7
+ From https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/utils.py"""
8
+ dims_to_append = target_dims - x.ndim
9
+ if dims_to_append < 0:
10
+ raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
11
+ return x[(...,) + (None,) * dims_to_append]
12
+
13
+
14
+ def norm_thresholding(x0, value):
15
+ s = append_dims(x0.pow(2).flatten(1).mean(1).sqrt().clamp(min=value), x0.ndim)
16
+ return x0 * (value / s)
17
+
18
+
19
+ def spatial_norm_thresholding(x0, value):
20
+ # b c h w
21
+ s = x0.pow(2).mean(1, keepdim=True).sqrt().clamp(min=value)
22
+ return x0 * (value / s)