| 302 | |
| 303 | @torch.no_grad() |
| 304 | def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=200, ddim_eta=1., return_keys=None, |
| 305 | quantize_denoised=True, inpaint=True, plot_denoise_rows=False, plot_progressive_rows=True, |
| 306 | plot_diffusion_rows=True, **kwargs): |
| 307 | |
| 308 | use_ddim = ddim_steps is not None |
| 309 | |
| 310 | log = dict() |
| 311 | if isinstance(batch, list): |
| 312 | batch = batch[0] |
| 313 | z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key, |
| 314 | return_first_stage_outputs=True, |
| 315 | force_c_encode=True, |
| 316 | return_original_cond=True, |
| 317 | bs=N) |
| 318 | N = min(x.shape[0], N) |
| 319 | n_row = min(x.shape[0], n_row) |
| 320 | log["inputs"] = x |
| 321 | log["reconstruction"] = xrec |
| 322 | if self.model.conditioning_key is not None: |
| 323 | if hasattr(self.cond_stage_model, "decode"): |
| 324 | xc = self.cond_stage_model.decode(c) |
| 325 | log["conditioning"] = xc |
| 326 | elif self.cond_stage_key in ["caption"]: |
| 327 | xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["caption"]) |
| 328 | log["conditioning"] = xc |
| 329 | elif self.cond_stage_key == 'class_label': |
| 330 | xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"]) |
| 331 | log['conditioning'] = xc |
| 332 | elif isimage(xc): |
| 333 | log["conditioning"] = xc |
| 334 | if ismap(xc): |
| 335 | log["original_conditioning"] = self.to_rgb(xc) |
| 336 | |
| 337 | if plot_diffusion_rows: |
| 338 | # get diffusion row |
| 339 | diffusion_row = list() |
| 340 | z_start = z[:n_row] |
| 341 | for t in range(self.num_timesteps): |
| 342 | if t % self.log_every_t == 0 or t == self.num_timesteps - 1: |
| 343 | t = repeat(torch.tensor([t]), '1 -> b', b=n_row) |
| 344 | t = t.to(self.device).long() |
| 345 | noise = torch.randn_like(z_start) |
| 346 | z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise) |
| 347 | diffusion_row.append(self.decode_first_stage(z_noisy)) |
| 348 | |
| 349 | diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W |
| 350 | diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w') |
| 351 | diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w') |
| 352 | diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0]) |
| 353 | log["diffusion_row"] = diffusion_grid |
| 354 | |
| 355 | if sample: |
| 356 | # get denoise row |
| 357 | with self.ema_scope("Plotting"): |
| 358 | unconditional_guidance_scale=6. |
| 359 | unconditional_conditioning = self.get_learned_conditioning(len(c) * [""]) |
| 360 | samples, z_denoise_row = self.sample_log(cond=c,batch_size=N,ddim=use_ddim, |
| 361 | ddim_steps=ddim_steps,eta=ddim_eta, |