| 201 | @persistence.persistent_class |
| 202 | class GenInput(nn.Module): |
| 203 | def __init__(self, cfg: DictConfig, channel_dim: int, motion_v_dim: int=None): |
| 204 | super().__init__() |
| 205 | |
| 206 | self.cfg = cfg |
| 207 | |
| 208 | if self.cfg.input.type == 'const': |
| 209 | self.input = torch.nn.Parameter(torch.randn([channel_dim, 4, 4])) |
| 210 | self.total_dim = channel_dim |
| 211 | elif self.cfg.input.type == 'temporal': |
| 212 | self.input = TemporalInput(self.cfg, channel_dim, motion_v_dim=motion_v_dim) |
| 213 | self.total_dim = self.input.get_dim() |
| 214 | else: |
| 215 | raise NotImplementedError(f'Unkown input type: {self.cfg.input.type}') |
| 216 | |
| 217 | def forward(self, batch_size: int, motion_v: Optional[torch.Tensor]=None, dtype=None, memory_format=None) -> torch.Tensor: |
| 218 | if self.cfg.input.type == 'const': |