FilesExpand file tree

main

/

embeddings.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

117 lines (92 loc) · 3.79 KB

main

/

embeddings.py

Copy path

Top

File metadata and controls

• Code
• Blame

117 lines (92 loc) · 3.79 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

# Copyright 2022 The HuggingFace Team. All rights reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import math

import numpy as np

import torch

from torch import nn

def get_timestep_embedding(

timesteps: torch.Tensor,

embedding_dim: int,

flip_sin_to_cos: bool = False,

downscale_freq_shift: float = 1,

scale: float = 1,

max_period: int = 10000,

):

"""

This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.

:param timesteps: a 1-D Tensor of N indices, one per batch element.

These may be fractional.

:param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the

embeddings. :return: an [N x dim] Tensor of positional embeddings.

"""

assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"

half_dim = embedding_dim // 2

exponent = -math.log(max_period) * torch.arange(

start=0, end=half_dim, dtype=torch.float32, device=timesteps.device

)

exponent = exponent / (half_dim - downscale_freq_shift)

emb = torch.exp(exponent)

emb = timesteps[:, None].float() * emb[None, :]

# scale embeddings

emb = scale * emb

# concat sine and cosine embeddings

emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)

# flip sine and cosine embeddings

if flip_sin_to_cos:

emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)

# zero pad

if embedding_dim % 2 == 1:

emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))

return emb

class TimestepEmbedding(nn.Module):

def __init__(self, channel: int, time_embed_dim: int, act_fn: str = "silu"):

super().__init__()

self.linear_1 = nn.Linear(channel, time_embed_dim)

self.act = None

if act_fn == "silu":

self.act = nn.SiLU()

self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim)

def forward(self, sample):

sample = self.linear_1(sample)

if self.act is not None:

sample = self.act(sample)

sample = self.linear_2(sample)

return sample

class Timesteps(nn.Module):

def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float):

super().__init__()

self.num_channels = num_channels

self.flip_sin_to_cos = flip_sin_to_cos

self.downscale_freq_shift = downscale_freq_shift

def forward(self, timesteps):

t_emb = get_timestep_embedding(

timesteps,

self.num_channels,

flip_sin_to_cos=self.flip_sin_to_cos,

downscale_freq_shift=self.downscale_freq_shift,

)

return t_emb

class GaussianFourierProjection(nn.Module):

"""Gaussian Fourier embeddings for noise levels."""

def __init__(self, embedding_size: int = 256, scale: float = 1.0):

super().__init__()

self.weight = nn.Parameter(torch.randn(embedding_size) * scale, requires_grad=False)

# to delete later

self.W = nn.Parameter(torch.randn(embedding_size) * scale, requires_grad=False)

self.weight = self.W

def forward(self, x):

x = torch.log(x)

x_proj = x[:, None] * self.weight[None, :] * 2 * np.pi

out = torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1)

return out