| 129 | |
| 130 | class SirenWrapper(nn.Module): |
| 131 | def __init__(self, net, image_width, image_height, latent_dim = None): |
| 132 | super().__init__() |
| 133 | assert isinstance(net, SirenNet), 'SirenWrapper must receive a Siren network' |
| 134 | |
| 135 | self.net = net |
| 136 | self.image_width = image_width |
| 137 | self.image_height = image_height |
| 138 | |
| 139 | self.modulator = None |
| 140 | if exists(latent_dim): |
| 141 | self.modulator = Modulator( |
| 142 | dim_in = latent_dim, |
| 143 | dim_hidden = net.dim_hidden, |
| 144 | num_layers = net.num_layers |
| 145 | ) |
| 146 | |
| 147 | tensors = [torch.linspace(-1, 1, steps = image_height), torch.linspace(-1, 1, steps = image_width)] |
| 148 | mgrid = torch.stack(torch.meshgrid(*tensors, indexing = 'ij'), dim=-1) |
| 149 | mgrid = rearrange(mgrid, 'h w c -> (h w) c') |
| 150 | self.register_buffer('grid', mgrid) |
| 151 | |
| 152 | def forward(self, img = None, *, latent = None): |
| 153 | modulate = exists(self.modulator) |