| 1701 | |
| 1702 | @torch.no_grad() |
| 1703 | def get_input(self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False): |
| 1704 | # note: restricted to non-trainable encoders currently |
| 1705 | assert not self.cond_stage_trainable, 'trainable cond stages not yet supported for depth2img' |
| 1706 | z, c, x, xrec, xc = super().get_input(batch, self.first_stage_key, return_first_stage_outputs=True, |
| 1707 | force_c_encode=True, return_original_cond=True, bs=bs) |
| 1708 | |
| 1709 | assert exists(self.concat_keys) |
| 1710 | assert len(self.concat_keys) == 1 |
| 1711 | c_cat = list() |
| 1712 | for ck in self.concat_keys: |
| 1713 | cc = batch[ck] |
| 1714 | if bs is not None: |
| 1715 | cc = cc[:bs] |
| 1716 | cc = cc.to(self.device) |
| 1717 | cc = self.depth_model(cc) |
| 1718 | cc = torch.nn.functional.interpolate( |
| 1719 | cc, |
| 1720 | size=z.shape[2:], |
| 1721 | mode="bicubic", |
| 1722 | align_corners=False, |
| 1723 | ) |
| 1724 | |
| 1725 | depth_min, depth_max = torch.amin(cc, dim=[1, 2, 3], keepdim=True), torch.amax(cc, dim=[1, 2, 3], |
| 1726 | keepdim=True) |
| 1727 | cc = 2. * (cc - depth_min) / (depth_max - depth_min + 0.001) - 1. |
| 1728 | c_cat.append(cc) |
| 1729 | c_cat = torch.cat(c_cat, dim=1) |
| 1730 | all_conds = {"c_concat": [c_cat], "c_crossattn": [c]} |
| 1731 | if return_first_stage_outputs: |
| 1732 | return z, all_conds, x, xrec, xc |
| 1733 | return z, all_conds |
| 1734 | |
| 1735 | @torch.no_grad() |
| 1736 | def log_images(self, *args, **kwargs): |