condition on low-res image (and optionally on some spatial noise augmentation)
| 1743 | |
| 1744 | |
| 1745 | class LatentUpscaleFinetuneDiffusion(LatentFinetuneDiffusion): |
| 1746 | """ |
| 1747 | condition on low-res image (and optionally on some spatial noise augmentation) |
| 1748 | """ |
| 1749 | def __init__(self, concat_keys=("lr",), reshuffle_patch_size=None, |
| 1750 | low_scale_config=None, low_scale_key=None, *args, **kwargs): |
| 1751 | super().__init__(concat_keys=concat_keys, *args, **kwargs) |
| 1752 | self.reshuffle_patch_size = reshuffle_patch_size |
| 1753 | self.low_scale_model = None |
| 1754 | if low_scale_config is not None: |
| 1755 | print("Initializing a low-scale model") |
| 1756 | assert exists(low_scale_key) |
| 1757 | self.instantiate_low_stage(low_scale_config) |
| 1758 | self.low_scale_key = low_scale_key |
| 1759 | |
| 1760 | def instantiate_low_stage(self, config): |
| 1761 | model = instantiate_from_config(config) |
| 1762 | self.low_scale_model = model.eval() |
| 1763 | self.low_scale_model.train = disabled_train |
| 1764 | for param in self.low_scale_model.parameters(): |
| 1765 | param.requires_grad = False |
| 1766 | |
| 1767 | @torch.no_grad() |
| 1768 | def get_input(self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False): |
| 1769 | # note: restricted to non-trainable encoders currently |
| 1770 | assert not self.cond_stage_trainable, 'trainable cond stages not yet supported for upscaling-ft' |
| 1771 | z, c, x, xrec, xc = super().get_input(batch, self.first_stage_key, return_first_stage_outputs=True, |
| 1772 | force_c_encode=True, return_original_cond=True, bs=bs) |
| 1773 | |
| 1774 | assert exists(self.concat_keys) |
| 1775 | assert len(self.concat_keys) == 1 |
| 1776 | # optionally make spatial noise_level here |
| 1777 | c_cat = list() |
| 1778 | noise_level = None |
| 1779 | for ck in self.concat_keys: |
| 1780 | cc = batch[ck] |
| 1781 | cc = rearrange(cc, 'b h w c -> b c h w') |
| 1782 | if exists(self.reshuffle_patch_size): |
| 1783 | assert isinstance(self.reshuffle_patch_size, int) |
| 1784 | cc = rearrange(cc, 'b c (p1 h) (p2 w) -> b (p1 p2 c) h w', |
| 1785 | p1=self.reshuffle_patch_size, p2=self.reshuffle_patch_size) |
| 1786 | if bs is not None: |
| 1787 | cc = cc[:bs] |
| 1788 | cc = cc.to(self.device) |
| 1789 | if exists(self.low_scale_model) and ck == self.low_scale_key: |
| 1790 | cc, noise_level = self.low_scale_model(cc) |
| 1791 | c_cat.append(cc) |
| 1792 | c_cat = torch.cat(c_cat, dim=1) |
| 1793 | if exists(noise_level): |
| 1794 | all_conds = {"c_concat": [c_cat], "c_crossattn": [c], "c_adm": noise_level} |
| 1795 | else: |
| 1796 | all_conds = {"c_concat": [c_cat], "c_crossattn": [c]} |
| 1797 | if return_first_stage_outputs: |
| 1798 | return z, all_conds, x, xrec, xc |
| 1799 | return z, all_conds |
| 1800 | |
| 1801 | @torch.no_grad() |
| 1802 | def log_images(self, *args, **kwargs): |
nothing calls this directly
no outgoing calls
no test coverage detected