(self, batch, N=8, n_row=2, sample=True, return_keys=None, **kwargs)
| 482 | |
| 483 | @torch.no_grad() |
| 484 | def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=None, **kwargs): |
| 485 | log = dict() |
| 486 | x = self.get_input(batch, self.first_stage_key) |
| 487 | N = min(x.shape[0], N) |
| 488 | n_row = min(x.shape[0], n_row) |
| 489 | x = x.to(self.device)[:N] |
| 490 | log["inputs"] = x |
| 491 | |
| 492 | # get diffusion row |
| 493 | diffusion_row = list() |
| 494 | x_start = x[:n_row] |
| 495 | |
| 496 | for t in range(self.num_timesteps): |
| 497 | if t % self.log_every_t == 0 or t == self.num_timesteps - 1: |
| 498 | t = repeat(torch.tensor([t]), '1 -> b', b=n_row) |
| 499 | t = t.to(self.device).long() |
| 500 | noise = torch.randn_like(x_start) |
| 501 | x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) |
| 502 | diffusion_row.append(x_noisy) |
| 503 | |
| 504 | log["diffusion_row"] = self._get_rows_from_list(diffusion_row) |
| 505 | |
| 506 | if sample: |
| 507 | # get denoise row |
| 508 | with self.ema_scope("Plotting"): |
| 509 | samples, denoise_row = self.sample(batch_size=N, return_intermediates=True) |
| 510 | |
| 511 | log["samples"] = samples |
| 512 | log["denoise_row"] = self._get_rows_from_list(denoise_row) |
| 513 | |
| 514 | if return_keys: |
| 515 | if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0: |
| 516 | return log |
| 517 | else: |
| 518 | return {key: log[key] for key in return_keys} |
| 519 | return log |
| 520 | |
| 521 | def configure_optimizers(self): |
| 522 | lr = self.learning_rate |
no test coverage detected