| 225 | |
| 226 | |
| 227 | class StrandGeneratorSiren(nn.Module): |
| 228 | # a siren network which predicts various direction vectors along the strand similar |
| 229 | def __init__(self, in_channels, modulation_hidden_dim, siren_hidden_dim, scale_init, decode_type, decode_random_verts, nr_verts_per_strand=256, nr_values_to_decode=256, dim_per_value_decoded=3): |
| 230 | super(StrandGeneratorSiren, self).__init__() |
| 231 | |
| 232 | self.nr_verts_per_strand = nr_verts_per_strand |
| 233 | self.nr_values_to_decode=nr_values_to_decode |
| 234 | |
| 235 | self.decode_type = decode_type |
| 236 | self.decode_random_verts = decode_random_verts |
| 237 | |
| 238 | |
| 239 | |
| 240 | if self.decode_type=="xyz": |
| 241 | nr_verts_to_create=self.nr_verts_per_strand |
| 242 | elif self.decode_type=="dir": |
| 243 | nr_verts_to_create = self.nr_verts_per_strand - 1 # we create only 99 because the frist one is just the origin |
| 244 | else: |
| 245 | raise ValueError("Unkown decode type: ", self.decode_type) |
| 246 | |
| 247 | if self.decode_random_verts: |
| 248 | nr_verts_to_create = 1 |
| 249 | |
| 250 | self.nr_verts_to_create=nr_verts_to_create |
| 251 | |
| 252 | self.activ = torch.nn.SiLU() |
| 253 | |
| 254 | |
| 255 | self.nr_layers = 3 |
| 256 | cur_nr_channels = in_channels |
| 257 | cur_nr_channels += 1 |
| 258 | # cur_nr_channels+=1 #+1 for the time t |
| 259 | self.modulation_layers = torch.nn.ModuleList([]) |
| 260 | self.gain_per_layer = torch.nn.ParameterList([]) |
| 261 | # self.w_per_layer = torch.nn.ParameterList([]) |
| 262 | for i in range(self.nr_layers): |
| 263 | self.modulation_layers.append(LinearWN_v2(cur_nr_channels, modulation_hidden_dim)) |
| 264 | cur_nr_channels = modulation_hidden_dim+in_channels +1 # at the end we concatenate the input z and a t |
| 265 | self.gain_per_layer.append(torch.nn.Parameter(torch.ones([]))) |
| 266 | |
| 267 | #not actually used during the forward pass but I have them so that the checkpoint still can load correctly |
| 268 | self.second_modulation_layers = torch.nn.ModuleList([]) |
| 269 | for i in range(self.nr_layers): |
| 270 | self.second_modulation_layers.append(LinearDummy(modulation_hidden_dim, modulation_hidden_dim)) |
| 271 | |
| 272 | |
| 273 | |
| 274 | self.decode_val = LinearWN_v2(siren_hidden_dim, dim_per_value_decoded) |
| 275 | self.gain_val = torch.nn.Parameter(torch.ones([siren_hidden_dim])) |
| 276 | |
| 277 | |
| 278 | self.apply(lambda x: kaiming_init(x, False, nonlinearity="silu")) |
| 279 | kaiming_init(self.decode_val, True) |
| 280 | |
| 281 | self.siren_layers = torch.nn.ModuleList([]) |
| 282 | self.siren_layers.append(BlockSiren(in_channels=1, out_channels=siren_hidden_dim, is_first_layer=True, scale_init=scale_init)) |
| 283 | for i in range(self.nr_layers-1): |
| 284 | self.siren_layers.append(BlockSiren(in_channels=siren_hidden_dim, out_channels=siren_hidden_dim, scale_init=scale_init )) |