| 221 | |
| 222 | # PositionalEncoding Source： https://github.com/lmnt-com/wavegrad/blob/master/src/wavegrad/model.py |
| 223 | class PositionalEncoding(nn.Module): |
| 224 | def __init__(self, dim): |
| 225 | super().__init__() |
| 226 | self.dim = dim |
| 227 | |
| 228 | def forward(self, noise_level): |
| 229 | count = self.dim // 2 |
| 230 | step = ( |
| 231 | torch.arange(count, dtype=noise_level.dtype, device=noise_level.device) |
| 232 | / count |
| 233 | ) |
| 234 | encoding = noise_level.unsqueeze(1) * torch.exp( |
| 235 | -math.log(1e4) * step.unsqueeze(0) |
| 236 | ) |
| 237 | encoding = torch.cat([torch.sin(encoding), torch.cos(encoding)], dim=-1) |
| 238 | return encoding |
| 239 | |
| 240 | |
| 241 | class FeatureWiseAffine(nn.Module): |