Coverage for widitapp/diffusion/gaussian_diffusion.py: 99.38%

1# Modified from OpenAI's diffusion repos 

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

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

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

5 

6 

7import inspect 

8import math 

9 

10import numpy as np 

11import torch as th 

12import enum 

13 

14from .diffusion_utils import discretized_gaussian_log_likelihood, normal_kl 

15 

16 

17def mean_flat(tensor): 

18 """ 

19 Take the mean over all non-batch dimensions. 

20 """ 

21 return tensor.mean(dim=list(range(1, len(tensor.shape)))) 

22 

23 

24class ModelMeanType(enum.Enum): 

25 """ 

26 Which type of output the model predicts. 

27 """ 

28 

29 PREVIOUS_X = enum.auto() # the model predicts x_{t-1} 

30 START_X = enum.auto() # the model predicts x_0 

31 EPSILON = enum.auto() # the model predicts epsilon 

32 

33 

34class ModelVarType(enum.Enum): 

35 """ 

36 What is used as the model's output variance. 

37 The LEARNED_RANGE option has been added to allow the model to predict 

38 values between FIXED_SMALL and FIXED_LARGE, making its job easier. 

39 """ 

40 

41 LEARNED = enum.auto() 

42 FIXED_SMALL = enum.auto() 

43 FIXED_LARGE = enum.auto() 

44 LEARNED_RANGE = enum.auto() 

45 

46 

47class LossType(enum.Enum): 

48 MSE = enum.auto() # use raw MSE loss (and KL when learning variances) 

49 RESCALED_MSE = ( 

50 enum.auto() 

51 ) # use raw MSE loss (with RESCALED_KL when learning variances) 

52 KL = enum.auto() # use the variational lower-bound 

53 RESCALED_KL = enum.auto() # like KL, but rescale to estimate the full VLB 

54 

55 def is_vb(self): 

56 return self == LossType.KL or self == LossType.RESCALED_KL 

57 

58 

59def _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, warmup_frac): 

60 betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64) 

61 warmup_time = int(num_diffusion_timesteps * warmup_frac) 

62 betas[:warmup_time] = np.linspace(beta_start, beta_end, warmup_time, dtype=np.float64) 

63 return betas 

64 

65 

66def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps): 

67 """ 

68 This is the deprecated API for creating beta schedules. 

69 See get_named_beta_schedule() for the new library of schedules. 

70 """ 

71 if beta_schedule == "quad": 

72 betas = ( 

73 np.linspace( 

74 beta_start ** 0.5, 

75 beta_end ** 0.5, 

76 num_diffusion_timesteps, 

77 dtype=np.float64, 

78 ) 

79 ** 2 

80 ) 

81 elif beta_schedule == "linear": 

82 betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64) 

83 elif beta_schedule == "warmup10": 

84 betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.1) 

85 elif beta_schedule == "warmup50": 

86 betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.5) 

87 elif beta_schedule == "const": 

88 betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64) 

89 elif beta_schedule == "jsd": # 1/T, 1/(T-1), 1/(T-2), ..., 1 

90 betas = 1.0 / np.linspace( 

91 num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64 

92 ) 

93 else: 

94 raise NotImplementedError(beta_schedule) 

95 assert betas.shape == (num_diffusion_timesteps,), ( 

96 f"betas.shape must be ({num_diffusion_timesteps},), got {betas.shape}" 

97 ) 

98 return betas 

99 

100 

101def get_named_beta_schedule(schedule_name, num_diffusion_timesteps): 

102 """ 

103 Get a pre-defined beta schedule for the given name. 

104 The beta schedule library consists of beta schedules which remain similar 

105 in the limit of num_diffusion_timesteps. 

106 Beta schedules may be added, but should not be removed or changed once 

107 they are committed to maintain backwards compatibility. 

108 """ 

109 if schedule_name == "linear": 

110 # Linear schedule from Ho et al, extended to work for any number of 

111 # diffusion steps. 

112 scale = 1000 / num_diffusion_timesteps 

113 return get_beta_schedule( 

114 "linear", 

115 beta_start=scale * 0.0001, 

116 beta_end=scale * 0.02, 

117 num_diffusion_timesteps=num_diffusion_timesteps, 

118 ) 

119 elif schedule_name == "squaredcos_cap_v2": 

120 return betas_for_alpha_bar( 

121 num_diffusion_timesteps, 

122 lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2, 

123 ) 

124 else: 

125 raise NotImplementedError(f"unknown beta schedule: {schedule_name}") 

126 

127 

128def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): 

129 """ 

130 Create a beta schedule that discretizes the given alpha_t_bar function, 

131 which defines the cumulative product of (1-beta) over time from t = [0,1]. 

132 :param num_diffusion_timesteps: the number of betas to produce. 

133 :param alpha_bar: a lambda that takes an argument t from 0 to 1 and 

134 produces the cumulative product of (1-beta) up to that 

135 part of the diffusion process. 

136 :param max_beta: the maximum beta to use; use values lower than 1 to 

137 prevent singularities. 

138 """ 

139 betas = [] 

140 for i in range(num_diffusion_timesteps): 

141 t1 = i / num_diffusion_timesteps 

142 t2 = (i + 1) / num_diffusion_timesteps 

143 betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) 

144 return np.array(betas) 

145 

146 

147class GaussianDiffusion: 

148 """ 

149 Utilities for training and sampling diffusion models. 

150 Original ported from this codebase: 

151 https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42 

152 :param betas: a 1-D numpy array of betas for each diffusion timestep, 

153 starting at T and going to 1. 

154 """ 

155 

156 def __init__( 

157 self, 

158 *, 

159 betas, 

160 model_mean_type, 

161 model_var_type, 

162 loss_type 

163 ): 

164 

165 self.model_mean_type = model_mean_type 

166 self.model_var_type = model_var_type 

167 self.loss_type = loss_type 

168 

169 # Use float64 for accuracy. 

170 betas = np.array(betas, dtype=np.float64) 

171 self.betas = betas 

172 assert len(betas.shape) == 1, f"betas must be 1-D, got shape {betas.shape}" 

173 assert (betas > 0).all() and (betas <= 1).all(), ( 

174 "betas must be in (0, 1]; " 

175 f"min={betas.min()} max={betas.max()}" 

176 ) 

177 

178 self.num_timesteps = int(betas.shape[0]) 

179 

180 alphas = 1.0 - betas 

181 self.alphas_cumprod = np.cumprod(alphas, axis=0) 

182 self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1]) 

183 self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0) 

184 assert self.alphas_cumprod_prev.shape == (self.num_timesteps,), ( 

185 f"alphas_cumprod_prev shape must be ({self.num_timesteps},), " 

186 f"got {self.alphas_cumprod_prev.shape}" 

187 ) 

188 

189 # calculations for diffusion q(x_t | x_{t-1}) and others 

190 self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod) 

191 self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod) 

192 self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod) 

193 self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod) 

194 self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1) 

195 

196 # calculations for posterior q(x_{t-1} | x_t, x_0) 

197 self.posterior_variance = ( 

198 betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) 

199 ) 

200 # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain 

201 self.posterior_log_variance_clipped = np.log( 

202 np.append(self.posterior_variance[1], self.posterior_variance[1:]) 

203 ) if len(self.posterior_variance) > 1 else np.array([]) 

204 

205 self.posterior_mean_coef1 = ( 

206 betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) 

207 ) 

208 self.posterior_mean_coef2 = ( 

209 (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - self.alphas_cumprod) 

210 ) 

211 

212 def q_mean_variance(self, x_start, t): 

213 """ 

214 Get the distribution q(x_t | x_0). 

215 :param x_start: the [N x C x ...] tensor of noiseless inputs. 

216 :param t: the number of diffusion steps (minus 1). Here, 0 means one step. 

217 :return: A tuple (mean, variance, log_variance), all of x_start's shape. 

218 """ 

219 mean = _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start 

220 variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape) 

221 log_variance = _extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape) 

222 return mean, variance, log_variance 

223 

224 def q_sample(self, x_start, t, noise=None): 

225 """ 

226 Diffuse the data for a given number of diffusion steps. 

227 In other words, sample from q(x_t | x_0). 

228 :param x_start: the initial data batch. 

229 :param t: the number of diffusion steps (minus 1). Here, 0 means one step. 

230 :param noise: if specified, the split-out normal noise. 

231 :return: A noisy version of x_start. 

232 """ 

233 if noise is None: 

234 noise = th.randn_like(x_start) 

235 assert noise.shape == x_start.shape, ( 

236 f"noise shape {noise.shape} must match x_start shape {x_start.shape}" 

237 ) 

238 return ( 

239 _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start 

240 + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise 

241 ) 

242 

243 def q_posterior_mean_variance(self, x_start, x_t, t): 

244 """ 

245 Compute the mean and variance of the diffusion posterior: 

246 q(x_{t-1} | x_t, x_0) 

247 """ 

248 assert x_start.shape == x_t.shape, ( 

249 f"x_start shape {x_start.shape} must match x_t shape {x_t.shape}" 

250 ) 

251 posterior_mean = ( 

252 _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start 

253 + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t 

254 ) 

255 posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape) 

256 posterior_log_variance_clipped = _extract_into_tensor( 

257 self.posterior_log_variance_clipped, t, x_t.shape 

258 ) 

259 assert ( 

260 posterior_mean.shape[0] 

261 == posterior_variance.shape[0] 

262 == posterior_log_variance_clipped.shape[0] 

263 == x_start.shape[0] 

264 ), ( 

265 "posterior batch size mismatch: " 

266 f"mean={posterior_mean.shape[0]} " 

267 f"var={posterior_variance.shape[0]} " 

268 f"logvar={posterior_log_variance_clipped.shape[0]} " 

269 f"x_start={x_start.shape[0]}" 

270 ) 

271 return posterior_mean, posterior_variance, posterior_log_variance_clipped 

272 

273 def p_mean_variance(self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None): 

274 """ 

275 Apply the model to get p(x_{t-1} | x_t), as well as a prediction of 

276 the initial x, x_0. 

277 :param model: the model, which takes a signal and a batch of timesteps 

278 as input. 

279 :param x: the [N x C x ...] tensor at time t. 

280 :param t: a 1-D Tensor of timesteps. 

281 :param clip_denoised: if True, clip the denoised signal into [-1, 1]. 

282 :param denoised_fn: if not None, a function which applies to the 

283 x_start prediction before it is used to sample. Applies before 

284 clip_denoised. 

285 :param model_kwargs: if not None, a dict of extra keyword arguments to 

286 pass to the model. This can be used for conditioning. 

287 :return: a dict with the following keys: 

288 - 'mean': the model mean output. 

289 - 'variance': the model variance output. 

290 - 'log_variance': the log of 'variance'. 

291 - 'pred_xstart': the prediction for x_0. 

292 """ 

293 if model_kwargs is None: 

294 model_kwargs = {} 

295 

296 B, C = x.shape[:2] 

297 assert t.shape == (B,), f"t shape must be ({B},), got {t.shape}" 

298 model_output = model(x, t, **model_kwargs) 

299 if isinstance(model_output, tuple): 

300 model_output, extra = model_output 

301 else: 

302 extra = None 

303 

304 if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]: 

305 assert model_output.shape == (B, C * 2, *x.shape[2:]), ( 

306 f"model_output shape must be {(B, C * 2, *x.shape[2:])}, " 

307 f"got {model_output.shape}" 

308 ) 

309 model_output, model_var_values = th.split(model_output, C, dim=1) 

310 min_log = _extract_into_tensor(self.posterior_log_variance_clipped, t, x.shape) 

311 max_log = _extract_into_tensor(np.log(self.betas), t, x.shape) 

312 # The model_var_values is [-1, 1] for [min_var, max_var]. 

313 frac = (model_var_values + 1) / 2 

314 model_log_variance = frac * max_log + (1 - frac) * min_log 

315 model_variance = th.exp(model_log_variance) 

316 else: 

317 model_variance, model_log_variance = { 

318 # for fixedlarge, we set the initial (log-)variance like so 

319 # to get a better decoder log likelihood. 

320 ModelVarType.FIXED_LARGE: ( 

321 np.append(self.posterior_variance[1], self.betas[1:]), 

322 np.log(np.append(self.posterior_variance[1], self.betas[1:])), 

323 ), 

324 ModelVarType.FIXED_SMALL: ( 

325 self.posterior_variance, 

326 self.posterior_log_variance_clipped, 

327 ), 

328 }[self.model_var_type] 

329 model_variance = _extract_into_tensor(model_variance, t, x.shape) 

330 model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape) 

331 

332 def process_xstart(x): 

333 if denoised_fn is not None: 

334 x = denoised_fn(x) 

335 if clip_denoised: 

336 return x.clamp(-1, 1) 

337 return x 

338 

339 if self.model_mean_type == ModelMeanType.START_X: 

340 pred_xstart = process_xstart(model_output) 

341 else: 

342 pred_xstart = process_xstart( 

343 self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output) 

344 ) 

345 model_mean, _, _ = self.q_posterior_mean_variance(x_start=pred_xstart, x_t=x, t=t) 

346 

347 assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape, ( 

348 "shape mismatch in p_mean_variance: " 

349 f"model_mean={model_mean.shape} " 

350 f"log_variance={model_log_variance.shape} " 

351 f"pred_xstart={pred_xstart.shape} " 

352 f"x={x.shape}" 

353 ) 

354 return { 

355 "mean": model_mean, 

356 "variance": model_variance, 

357 "log_variance": model_log_variance, 

358 "pred_xstart": pred_xstart, 

359 "extra": extra, 

360 } 

361 

362 def _predict_xstart_from_eps(self, x_t, t, eps): 

363 assert x_t.shape == eps.shape, ( 

364 f"x_t shape {x_t.shape} must match eps shape {eps.shape}" 

365 ) 

366 return ( 

367 _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t 

368 - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps 

369 ) 

370 

371 def _predict_eps_from_xstart(self, x_t, t, pred_xstart): 

372 return ( 

373 _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart 

374 ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) 

375 

376 def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None): 

377 """ 

378 Compute the mean for the previous step, given a function cond_fn that 

379 computes the gradient of a conditional log probability with respect to 

380 x. In particular, cond_fn computes grad(log(p(y|x))), and we want to 

381 condition on y. 

382 This uses the conditioning strategy from Sohl-Dickstein et al. (2015). 

383 """ 

384 gradient = cond_fn(x, t, **model_kwargs) 

385 new_mean = p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float() 

386 return new_mean 

387 

388 def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None): 

389 """ 

390 Compute what the p_mean_variance output would have been, should the 

391 model's score function be conditioned by cond_fn. 

392 See condition_mean() for details on cond_fn. 

393 Unlike condition_mean(), this instead uses the conditioning strategy 

394 from Song et al (2020). 

395 """ 

396 alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) 

397 

398 eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"]) 

399 eps = eps - (1 - alpha_bar).sqrt() * cond_fn(x, t, **model_kwargs) 

400 

401 out = p_mean_var.copy() 

402 out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps) 

403 out["mean"], _, _ = self.q_posterior_mean_variance(x_start=out["pred_xstart"], x_t=x, t=t) 

404 return out 

405 

406 def p_sample( 

407 self, 

408 model, 

409 x, 

410 t, 

411 clip_denoised=True, 

412 denoised_fn=None, 

413 cond_fn=None, 

414 model_kwargs=None, 

415 ): 

416 """ 

417 Sample x_{t-1} from the model at the given timestep. 

418 :param model: the model to sample from. 

419 :param x: the current tensor at x_{t-1}. 

420 :param t: the value of t, starting at 0 for the first diffusion step. 

421 :param clip_denoised: if True, clip the x_start prediction to [-1, 1]. 

422 :param denoised_fn: if not None, a function which applies to the 

423 x_start prediction before it is used to sample. 

424 :param cond_fn: if not None, this is a gradient function that acts 

425 similarly to the model. 

426 :param model_kwargs: if not None, a dict of extra keyword arguments to 

427 pass to the model. This can be used for conditioning. 

428 :return: a dict containing the following keys: 

429 - 'sample': a random sample from the model. 

430 - 'pred_xstart': a prediction of x_0. 

431 """ 

432 out = self.p_mean_variance( 

433 model, 

434 x, 

435 t, 

436 clip_denoised=clip_denoised, 

437 denoised_fn=denoised_fn, 

438 model_kwargs=model_kwargs, 

439 ) 

440 noise = th.randn_like(x) 

441 nonzero_mask = ( 

442 (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) 

443 ) # no noise when t == 0 

444 if cond_fn is not None: 

445 out["mean"] = self.condition_mean(cond_fn, out, x, t, model_kwargs=model_kwargs) 

446 sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise 

447 return {"sample": sample, "pred_xstart": out["pred_xstart"]} 

448 

449 def p_sample_loop( 

450 self, 

451 model, 

452 shape, 

453 noise=None, 

454 clip_denoised=True, 

455 denoised_fn=None, 

456 cond_fn=None, 

457 model_kwargs=None, 

458 device=None, 

459 progress=False, 

460 ): 

461 """ 

462 Generate samples from the model. 

463 :param model: the model module. 

464 :param shape: the shape of the samples, (N, C, H, W). 

465 :param noise: if specified, the noise from the encoder to sample. 

466 Should be of the same shape as `shape`. 

467 :param clip_denoised: if True, clip x_start predictions to [-1, 1]. 

468 :param denoised_fn: if not None, a function which applies to the 

469 x_start prediction before it is used to sample. 

470 :param cond_fn: if not None, this is a gradient function that acts 

471 similarly to the model. 

472 :param model_kwargs: if not None, a dict of extra keyword arguments to 

473 pass to the model. This can be used for conditioning. 

474 :param device: if specified, the device to create the samples on. 

475 If not specified, use a model parameter's device. 

476 :param progress: if True, show a tqdm progress bar. 

477 :return: a non-differentiable batch of samples. 

478 """ 

479 final = None 

480 for sample in self.p_sample_loop_progressive( 

481 model, 

482 shape, 

483 noise=noise, 

484 clip_denoised=clip_denoised, 

485 denoised_fn=denoised_fn, 

486 cond_fn=cond_fn, 

487 model_kwargs=model_kwargs, 

488 device=device, 

489 progress=progress, 

490 ): 

491 final = sample 

492 return final["sample"] 

493 

494 def p_sample_loop_progressive( 

495 self, 

496 model, 

497 shape, 

498 noise=None, 

499 clip_denoised=True, 

500 denoised_fn=None, 

501 cond_fn=None, 

502 model_kwargs=None, 

503 device=None, 

504 progress=False, 

505 ): 

506 """ 

507 Generate samples from the model and yield intermediate samples from 

508 each timestep of diffusion. 

509 Arguments are the same as p_sample_loop(). 

510 Returns a generator over dicts, where each dict is the return value of 

511 p_sample(). 

512 """ 

513 if device is None: 

514 device = next(model.parameters()).device 

515 assert isinstance(shape, (tuple, list)), ( 

516 f"shape must be tuple or list, got {type(shape)}" 

517 ) 

518 if noise is not None: 

519 img = noise 

520 else: 

521 img = th.randn(*shape, device=device) 

522 indices = list(range(self.num_timesteps))[::-1] 

523 

524 if progress: 

525 # Lazy import so that we don't depend on tqdm. 

526 from tqdm.auto import tqdm 

527 

528 indices = tqdm(indices) 

529 

530 for i in indices: 

531 t = th.tensor([i] * shape[0], device=device) 

532 with th.no_grad(): 

533 out = self.p_sample( 

534 model, 

535 img, 

536 t, 

537 clip_denoised=clip_denoised, 

538 denoised_fn=denoised_fn, 

539 cond_fn=cond_fn, 

540 model_kwargs=model_kwargs, 

541 ) 

542 yield out 

543 img = out["sample"] 

544 

545 def ddim_sample( 

546 self, 

547 model, 

548 x, 

549 t, 

550 clip_denoised=True, 

551 denoised_fn=None, 

552 cond_fn=None, 

553 model_kwargs=None, 

554 eta=0.0, 

555 ): 

556 """ 

557 Sample x_{t-1} from the model using DDIM. 

558 Same usage as p_sample(). 

559 """ 

560 out = self.p_mean_variance( 

561 model, 

562 x, 

563 t, 

564 clip_denoised=clip_denoised, 

565 denoised_fn=denoised_fn, 

566 model_kwargs=model_kwargs, 

567 ) 

568 if cond_fn is not None: 

569 out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs) 

570 

571 # Usually our model outputs epsilon, but we re-derive it 

572 # in case we used x_start or x_prev prediction. 

573 eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"]) 

574 

575 alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) 

576 alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape) 

577 sigma = ( 

578 eta 

579 * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) 

580 * th.sqrt(1 - alpha_bar / alpha_bar_prev) 

581 ) 

582 # Equation 12. 

583 noise = th.randn_like(x) 

584 mean_pred = ( 

585 out["pred_xstart"] * th.sqrt(alpha_bar_prev) 

586 + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps 

587 ) 

588 nonzero_mask = ( 

589 (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) 

590 ) # no noise when t == 0 

591 sample = mean_pred + nonzero_mask * sigma * noise 

592 return {"sample": sample, "pred_xstart": out["pred_xstart"]} 

593 

594 def ddim_reverse_sample( 

595 self, 

596 model, 

597 x, 

598 t, 

599 clip_denoised=True, 

600 denoised_fn=None, 

601 cond_fn=None, 

602 model_kwargs=None, 

603 eta=0.0, 

604 ): 

605 """ 

606 Sample x_{t+1} from the model using DDIM reverse ODE. 

607 """ 

608 assert eta == 0.0, f"Reverse ODE only for deterministic path, got eta={eta}" 

609 out = self.p_mean_variance( 

610 model, 

611 x, 

612 t, 

613 clip_denoised=clip_denoised, 

614 denoised_fn=denoised_fn, 

615 model_kwargs=model_kwargs, 

616 ) 

617 if cond_fn is not None: 

618 out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs) 

619 # Usually our model outputs epsilon, but we re-derive it 

620 # in case we used x_start or x_prev prediction. 

621 eps = ( 

622 _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x 

623 - out["pred_xstart"] 

624 ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape) 

625 alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape) 

626 

627 # Equation 12. reversed 

628 mean_pred = out["pred_xstart"] * th.sqrt(alpha_bar_next) + th.sqrt(1 - alpha_bar_next) * eps 

629 

630 return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]} 

631 

632 def ddim_sample_loop( 

633 self, 

634 model, 

635 shape, 

636 noise=None, 

637 clip_denoised=True, 

638 denoised_fn=None, 

639 cond_fn=None, 

640 model_kwargs=None, 

641 device=None, 

642 progress=False, 

643 eta=0.0, 

644 ): 

645 """ 

646 Generate samples from the model using DDIM. 

647 Same usage as p_sample_loop(). 

648 """ 

649 final = None 

650 for sample in self.ddim_sample_loop_progressive( 

651 model, 

652 shape, 

653 noise=noise, 

654 clip_denoised=clip_denoised, 

655 denoised_fn=denoised_fn, 

656 cond_fn=cond_fn, 

657 model_kwargs=model_kwargs, 

658 device=device, 

659 progress=progress, 

660 eta=eta, 

661 ): 

662 final = sample 

663 return final["sample"] 

664 

665 def ddim_sample_loop_progressive( 

666 self, 

667 model, 

668 shape, 

669 noise=None, 

670 clip_denoised=True, 

671 denoised_fn=None, 

672 cond_fn=None, 

673 model_kwargs=None, 

674 device=None, 

675 progress=False, 

676 eta=0.0, 

677 ): 

678 """ 

679 Use DDIM to sample from the model and yield intermediate samples from 

680 each timestep of DDIM. 

681 Same usage as p_sample_loop_progressive(). 

682 """ 

683 if device is None: 

684 device = next(model.parameters()).device 

685 assert isinstance(shape, (tuple, list)), ( 

686 f"shape must be tuple or list, got {type(shape)}" 

687 ) 

688 if noise is not None: 

689 img = noise 

690 else: 

691 img = th.randn(*shape, device=device) 

692 indices = list(range(self.num_timesteps))[::-1] 

693 

694 if progress: 

695 # Lazy import so that we don't depend on tqdm. 

696 from tqdm.auto import tqdm 

697 

698 indices = tqdm(indices) 

699 

700 for i in indices: 

701 t = th.tensor([i] * shape[0], device=device) 

702 with th.no_grad(): 

703 out = self.ddim_sample( 

704 model, 

705 img, 

706 t, 

707 clip_denoised=clip_denoised, 

708 denoised_fn=denoised_fn, 

709 cond_fn=cond_fn, 

710 model_kwargs=model_kwargs, 

711 eta=eta, 

712 ) 

713 yield out 

714 img = out["sample"] 

715 

716 def _vb_terms_bpd( 

717 self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None 

718 ): 

719 """ 

720 Get a term for the variational lower-bound. 

721 The resulting units are bits (rather than nats, as one might expect). 

722 This allows for comparison to other papers. 

723 :return: a dict with the following keys: 

724 - 'output': a shape [N] tensor of NLLs or KLs. 

725 - 'pred_xstart': the x_0 predictions. 

726 """ 

727 true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance( 

728 x_start=x_start, x_t=x_t, t=t 

729 ) 

730 out = self.p_mean_variance( 

731 model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs 

732 ) 

733 kl = normal_kl( 

734 true_mean, true_log_variance_clipped, out["mean"], out["log_variance"] 

735 ) 

736 kl = mean_flat(kl) / np.log(2.0) 

737 

738 decoder_nll = -discretized_gaussian_log_likelihood( 

739 x_start, means=out["mean"], log_scales=0.5 * out["log_variance"] 

740 ) 

741 assert decoder_nll.shape == x_start.shape, ( 

742 f"decoder_nll shape {decoder_nll.shape} must match x_start shape {x_start.shape}" 

743 ) 

744 decoder_nll = mean_flat(decoder_nll) / np.log(2.0) 

745 

746 # At the first timestep return the decoder NLL, 

747 # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t)) 

748 output = th.where((t == 0), decoder_nll, kl) 

749 return {"output": output, "pred_xstart": out["pred_xstart"]} 

750 

751 def training_losses( 

752 self, model, x_start, t, model_kwargs=None, noise=None, 

753 image_loss_fn=None, return_pred_xstart=False, 

754 ): 

755 """ 

756 Compute training losses for a single timestep. 

757 :param model: the model to evaluate loss on. 

758 :param x_start: the [N x C x ...] tensor of inputs. 

759 :param t: a batch of timestep indices. 

760 :param model_kwargs: if not None, a dict of extra keyword arguments to 

761 pass to the model. This can be used for conditioning. 

762 :param noise: if specified, the specific Gaussian noise to try to remove. 

763 :param image_loss_fn: optional extra loss on (unclipped predicted x_0, 

764 x_start), returning a scalar batch mean or one loss per example. 

765 Supported for MSE and RESCALED_MSE objectives. The callable may 

766 also accept a keyword-only/optional ``alpha_bar`` argument, which 

767 receives a [N] tensor of the cumulative product of alphas at the 

768 sampled timesteps, so that the loss can compensate for the 

769 timestep-dependent scaling of its gradient. Callables that do not 

770 accept ``alpha_bar`` are still called with two arguments. 

771 :param return_pred_xstart: include the unclipped clean-image estimate 

772 for metrics, without another model forward pass. 

773 :return: a dict with the key "loss" containing a tensor of shape [N]. 

774 Some mean or variance settings may also have other keys. 

775 """ 

776 if image_loss_fn is not None and self.loss_type not in (LossType.MSE, LossType.RESCALED_MSE): 

777 raise ValueError("image_loss_fn requires an MSE diffusion objective.") 

778 if model_kwargs is None: 

779 model_kwargs = {} 

780 if noise is None: 

781 noise = th.randn_like(x_start) 

782 x_t = self.q_sample(x_start, t, noise=noise) 

783 

784 terms = {} 

785 

786 if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL: 

787 vb_terms = self._vb_terms_bpd( 

788 model=model, 

789 x_start=x_start, 

790 x_t=x_t, 

791 t=t, 

792 clip_denoised=False, 

793 model_kwargs=model_kwargs, 

794 ) 

795 terms["loss"] = vb_terms["output"] 

796 if return_pred_xstart: 

797 terms["pred_xstart"] = vb_terms["pred_xstart"] 

798 if self.loss_type == LossType.RESCALED_KL: 

799 terms["loss"] *= self.num_timesteps 

800 elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE: 

801 model_output = model(x_t, t, **model_kwargs) 

802 

803 if self.model_var_type in [ 

804 ModelVarType.LEARNED, 

805 ModelVarType.LEARNED_RANGE, 

806 ]: 

807 B, C = x_t.shape[:2] 

808 assert model_output.shape == (B, C * 2, *x_t.shape[2:]), ( 

809 f"model_output shape must be {(B, C * 2, *x_t.shape[2:])}, " 

810 f"got {model_output.shape}" 

811 ) 

812 model_output, model_var_values = th.split(model_output, C, dim=1) 

813 # Learn the variance using the variational bound, but don't let 

814 # it affect our mean prediction. 

815 frozen_out = th.cat([model_output.detach(), model_var_values], dim=1) 

816 terms["vb"] = self._vb_terms_bpd( 

817 model=lambda *args, r=frozen_out: r, 

818 x_start=x_start, 

819 x_t=x_t, 

820 t=t, 

821 clip_denoised=False, 

822 )["output"] 

823 if self.loss_type == LossType.RESCALED_MSE: 

824 # Divide by 1000 for equivalence with initial implementation. 

825 # Without a factor of 1/1000, the VB term hurts the MSE term. 

826 terms["vb"] *= self.num_timesteps / 1000.0 

827 

828 target = { 

829 ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance( 

830 x_start=x_start, x_t=x_t, t=t 

831 )[0],