(self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False)
| 1657 | |
| 1658 | @torch.no_grad() |
| 1659 | def get_input(self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False): |
| 1660 | # note: restricted to non-trainable encoders currently |
| 1661 | assert not self.cond_stage_trainable, 'trainable cond stages not yet supported for inpainting' |
| 1662 | z, c, x, xrec, xc = super().get_input(batch, self.first_stage_key, return_first_stage_outputs=True, |
| 1663 | force_c_encode=True, return_original_cond=True, bs=bs) |
| 1664 | |
| 1665 | assert exists(self.concat_keys) |
| 1666 | c_cat = list() |
| 1667 | for ck in self.concat_keys: |
| 1668 | cc = rearrange(batch[ck], 'b h w c -> b c h w').to(memory_format=torch.contiguous_format).float() |
| 1669 | if bs is not None: |
| 1670 | cc = cc[:bs] |
| 1671 | cc = cc.to(self.device) |
| 1672 | bchw = z.shape |
| 1673 | if ck != self.masked_image_key: |
| 1674 | cc = torch.nn.functional.interpolate(cc, size=bchw[-2:]) |
| 1675 | else: |
| 1676 | cc = self.get_first_stage_encoding(self.encode_first_stage(cc)) |
| 1677 | c_cat.append(cc) |
| 1678 | c_cat = torch.cat(c_cat, dim=1) |
| 1679 | all_conds = {"c_concat": [c_cat], "c_crossattn": [c]} |
| 1680 | if return_first_stage_outputs: |
| 1681 | return z, all_conds, x, xrec, xc |
| 1682 | return z, all_conds |
| 1683 | |
| 1684 | @torch.no_grad() |
| 1685 | def log_images(self, *args, **kwargs): |
nothing calls this directly
no test coverage detected