(self, batch, N=8, n_row=4, sample=True, ddim_steps=50, ddim_eta=0., return_keys=None,
quantize_denoised=True, inpaint=True, 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)
| 1155 | |
| 1156 | @torch.no_grad() |
| 1157 | def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=50, ddim_eta=0., return_keys=None, |
| 1158 | quantize_denoised=True, inpaint=True, plot_denoise_rows=False, plot_progressive_rows=True, |
| 1159 | plot_diffusion_rows=True, unconditional_guidance_scale=1., unconditional_guidance_label=None, |
| 1160 | use_ema_scope=True, |
| 1161 | **kwargs): |
| 1162 | ema_scope = self.ema_scope if use_ema_scope else nullcontext |
| 1163 | use_ddim = ddim_steps is not None |
| 1164 | |
| 1165 | log = dict() |
| 1166 | z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key, |
| 1167 | return_first_stage_outputs=True, |
| 1168 | force_c_encode=True, |
| 1169 | return_original_cond=True, |
| 1170 | bs=N) |
| 1171 | N = min(x.shape[0], N) |
| 1172 | n_row = min(x.shape[0], n_row) |
| 1173 | log["inputs"] = x |
| 1174 | log["reconstruction"] = xrec |
| 1175 | if self.model.conditioning_key is not None: |
| 1176 | if hasattr(self.cond_stage_model, "decode"): |
| 1177 | xc = self.cond_stage_model.decode(c) |
| 1178 | log["conditioning"] = xc |
| 1179 | elif self.cond_stage_key in ["caption", "txt"]: |
| 1180 | xc = log_txt_as_img((x.shape[2], x.shape[3]), batch[self.cond_stage_key], size=x.shape[2] // 25) |
| 1181 | log["conditioning"] = xc |
| 1182 | elif self.cond_stage_key in ['class_label', "cls"]: |
| 1183 | try: |
| 1184 | xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"], size=x.shape[2] // 25) |
| 1185 | log['conditioning'] = xc |
| 1186 | except KeyError: |
| 1187 | # probably no "human_label" in batch |
| 1188 | pass |
| 1189 | elif isimage(xc): |
| 1190 | log["conditioning"] = xc |
| 1191 | if ismap(xc): |
| 1192 | log["original_conditioning"] = self.to_rgb(xc) |
| 1193 | |
| 1194 | if plot_diffusion_rows: |
| 1195 | # get diffusion row |
| 1196 | diffusion_row = list() |
| 1197 | z_start = z[:n_row] |
| 1198 | for t in range(self.num_timesteps): |
| 1199 | if t % self.log_every_t == 0 or t == self.num_timesteps - 1: |
| 1200 | t = repeat(torch.tensor([t]), '1 -> b', b=n_row) |
| 1201 | t = t.to(self.device).long() |
| 1202 | noise = torch.randn_like(z_start) |
| 1203 | z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise) |
| 1204 | diffusion_row.append(self.decode_first_stage(z_noisy)) |
| 1205 | |
| 1206 | diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W |
| 1207 | diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w') |
| 1208 | diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w') |
| 1209 | diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0]) |
| 1210 | log["diffusion_row"] = diffusion_grid |
| 1211 | |
| 1212 | if sample: |
| 1213 | # get denoise row |
| 1214 | with ema_scope("Sampling"): |
nothing calls this directly
no test coverage detected