Gaussian Fourier embeddings for noise levels.
| 782 | |
| 783 | |
| 784 | class GaussianFourierProjection(nn.Module): |
| 785 | """Gaussian Fourier embeddings for noise levels.""" |
| 786 | |
| 787 | def __init__( |
| 788 | self, embedding_size: int = 256, scale: float = 1.0, set_W_to_weight=True, log=True, flip_sin_to_cos=False |
| 789 | ): |
| 790 | super().__init__() |
| 791 | self.weight = nn.Parameter(torch.randn(embedding_size) * scale, requires_grad=False) |
| 792 | self.log = log |
| 793 | self.flip_sin_to_cos = flip_sin_to_cos |
| 794 | |
| 795 | if set_W_to_weight: |
| 796 | # to delete later |
| 797 | del self.weight |
| 798 | self.W = nn.Parameter(torch.randn(embedding_size) * scale, requires_grad=False) |
| 799 | self.weight = self.W |
| 800 | del self.W |
| 801 | |
| 802 | def forward(self, x): |
| 803 | if self.log: |
| 804 | x = torch.log(x) |
| 805 | |
| 806 | x_proj = x[:, None] * self.weight[None, :] * 2 * np.pi |
| 807 | |
| 808 | if self.flip_sin_to_cos: |
| 809 | out = torch.cat([torch.cos(x_proj), torch.sin(x_proj)], dim=-1) |
| 810 | else: |
| 811 | out = torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1) |
| 812 | return out |
| 813 | |
| 814 | |
| 815 | class SinusoidalPositionalEmbedding(nn.Module): |
no outgoing calls
no test coverage detected