(self, batch, N=8, n_row=4, sample=True, ddim_steps=200, ddim_eta=1., return_keys=None,
plot_denoise_rows=False, plot_progressive_rows=True, plot_diffusion_rows=True,
unconditional_guidance_scale=1., unconditional_guidance_label=None, use_ema_scope=True,
**kwargs)
| 1399 | |
| 1400 | @torch.no_grad() |
| 1401 | def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=200, ddim_eta=1., return_keys=None, |
| 1402 | plot_denoise_rows=False, plot_progressive_rows=True, plot_diffusion_rows=True, |
| 1403 | unconditional_guidance_scale=1., unconditional_guidance_label=None, use_ema_scope=True, |
| 1404 | **kwargs): |
| 1405 | ema_scope = self.ema_scope if use_ema_scope else nullcontext |
| 1406 | use_ddim = ddim_steps is not None |
| 1407 | |
| 1408 | log = dict() |
| 1409 | z, c, x, xrec, xc, x_low, x_low_rec, noise_level = self.get_input(batch, self.first_stage_key, bs=N, |
| 1410 | log_mode=True) |
| 1411 | N = min(x.shape[0], N) |
| 1412 | n_row = min(x.shape[0], n_row) |
| 1413 | log["inputs"] = x |
| 1414 | log["reconstruction"] = xrec |
| 1415 | log["x_lr"] = x_low |
| 1416 | log[f"x_lr_rec_@noise_levels{'-'.join(map(lambda x: str(x), list(noise_level.cpu().numpy())))}"] = x_low_rec |
| 1417 | if self.model.conditioning_key is not None: |
| 1418 | if hasattr(self.cond_stage_model, "decode"): |
| 1419 | xc = self.cond_stage_model.decode(c) |
| 1420 | log["conditioning"] = xc |
| 1421 | elif self.cond_stage_key in ["caption", "txt"]: |
| 1422 | xc = log_txt_as_img((x.shape[2], x.shape[3]), batch[self.cond_stage_key], size=x.shape[2] // 25) |
| 1423 | log["conditioning"] = xc |
| 1424 | elif self.cond_stage_key in ['class_label', 'cls']: |
| 1425 | xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"], size=x.shape[2] // 25) |
| 1426 | log['conditioning'] = xc |
| 1427 | elif isimage(xc): |
| 1428 | log["conditioning"] = xc |
| 1429 | if ismap(xc): |
| 1430 | log["original_conditioning"] = self.to_rgb(xc) |
| 1431 | |
| 1432 | if plot_diffusion_rows: |
| 1433 | # get diffusion row |
| 1434 | diffusion_row = list() |
| 1435 | z_start = z[:n_row] |
| 1436 | for t in range(self.num_timesteps): |
| 1437 | if t % self.log_every_t == 0 or t == self.num_timesteps - 1: |
| 1438 | t = repeat(torch.tensor([t]), '1 -> b', b=n_row) |
| 1439 | t = t.to(self.device).long() |
| 1440 | noise = torch.randn_like(z_start) |
| 1441 | z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise) |
| 1442 | diffusion_row.append(self.decode_first_stage(z_noisy)) |
| 1443 | |
| 1444 | diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W |
| 1445 | diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w') |
| 1446 | diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w') |
| 1447 | diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0]) |
| 1448 | log["diffusion_row"] = diffusion_grid |
| 1449 | |
| 1450 | if sample: |
| 1451 | # get denoise row |
| 1452 | with ema_scope("Sampling"): |
| 1453 | samples, z_denoise_row = self.sample_log(cond=c, batch_size=N, ddim=use_ddim, |
| 1454 | ddim_steps=ddim_steps, eta=ddim_eta) |
| 1455 | # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True) |
| 1456 | x_samples = self.decode_first_stage(samples) |
| 1457 | log["samples"] = x_samples |
| 1458 | if plot_denoise_rows: |
nothing calls this directly
no test coverage detected