(self, batch, k, cond_key=None, bs=None, log_mode=False)
| 1377 | |
| 1378 | @torch.no_grad() |
| 1379 | def get_input(self, batch, k, cond_key=None, bs=None, log_mode=False): |
| 1380 | if not log_mode: |
| 1381 | z, c = super().get_input(batch, k, force_c_encode=True, bs=bs) |
| 1382 | else: |
| 1383 | z, c, x, xrec, xc = super().get_input(batch, self.first_stage_key, return_first_stage_outputs=True, |
| 1384 | force_c_encode=True, return_original_cond=True, bs=bs) |
| 1385 | x_low = batch[self.low_scale_key][:bs] |
| 1386 | x_low = rearrange(x_low, 'b h w c -> b c h w') |
| 1387 | x_low = x_low.to(memory_format=torch.contiguous_format).float() |
| 1388 | zx, noise_level = self.low_scale_model(x_low) |
| 1389 | if self.noise_level_key is not None: |
| 1390 | # get noise level from batch instead, e.g. when extracting a custom noise level for bsr |
| 1391 | raise NotImplementedError('TODO') |
| 1392 | |
| 1393 | all_conds = {"c_concat": [zx], "c_crossattn": [c], "c_adm": noise_level} |
| 1394 | if log_mode: |
| 1395 | # TODO: maybe disable if too expensive |
| 1396 | x_low_rec = self.low_scale_model.decode(zx) |
| 1397 | return z, all_conds, x, xrec, xc, x_low, x_low_rec, noise_level |
| 1398 | return z, all_conds |
| 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, |
no test coverage detected