| 150 | self.register_buffer('grid', mgrid) |
| 151 | |
| 152 | def forward(self, img = None, *, latent = None): |
| 153 | modulate = exists(self.modulator) |
| 154 | assert not (modulate ^ exists(latent)), 'latent vector must be only supplied if `latent_dim` was passed in on instantiation' |
| 155 | |
| 156 | mods = self.modulator(latent) if modulate else None |
| 157 | |
| 158 | coords = self.grid.clone().detach().requires_grad_() |
| 159 | out = self.net(coords, mods) |
| 160 | out = rearrange(out, '(h w) c -> () c h w', h = self.image_height, w = self.image_width) |
| 161 | |
| 162 | if exists(img): |
| 163 | return F.mse_loss(img, out) |
| 164 | |
| 165 | return out |