FilesExpand file tree

main

/

transformer_temporal.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

176 lines (152 loc) · 7.27 KB

main

/

transformer_temporal.py

Copy path

Top

File metadata and controls

• Code
• Blame

176 lines (152 loc) · 7.27 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

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

# Copyright 2023 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.

from dataclasses import dataclass

from typing import Optional

import torch

from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config

from ..utils import BaseOutput

from .attention import BasicTransformerBlock

from .modeling_utils import ModelMixin

@dataclass

class TransformerTemporalModelOutput(BaseOutput):

"""

Args:

sample (`torch.FloatTensor` of shape `(batch_size x num_frames, num_channels, height, width)`)

Hidden states conditioned on `encoder_hidden_states` input.

"""

sample: torch.FloatTensor

class TransformerTemporalModel(ModelMixin, ConfigMixin):

"""

Transformer model for video-like data.

Parameters:

num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.

attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.

in_channels (`int`, *optional*):

Pass if the input is continuous. The number of channels in the input and output.

num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.

dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.

cross_attention_dim (`int`, *optional*): The number of encoder_hidden_states dimensions to use.

sample_size (`int`, *optional*): Pass if the input is discrete. The width of the latent images.

Note that this is fixed at training time as it is used for learning a number of position embeddings. See

`ImagePositionalEmbeddings`.

activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.

attention_bias (`bool`, *optional*):

Configure if the TransformerBlocks' attention should contain a bias parameter.

double_self_attention (`bool`, *optional*):

Configure if each TransformerBlock should contain two self-attention layers

"""

@register_to_config

def __init__(

self,

num_attention_heads: int = 16,

attention_head_dim: int = 88,

in_channels: Optional[int] = None,

out_channels: Optional[int] = None,

num_layers: int = 1,

dropout: float = 0.0,

norm_num_groups: int = 32,

cross_attention_dim: Optional[int] = None,

attention_bias: bool = False,

sample_size: Optional[int] = None,

activation_fn: str = "geglu",

norm_elementwise_affine: bool = True,

double_self_attention: bool = True,

):

super().__init__()

self.num_attention_heads = num_attention_heads

self.attention_head_dim = attention_head_dim

inner_dim = num_attention_heads * attention_head_dim

self.in_channels = in_channels

self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)

self.proj_in = nn.Linear(in_channels, inner_dim)

# 3. Define transformers blocks

self.transformer_blocks = nn.ModuleList(

[

BasicTransformerBlock(

inner_dim,

num_attention_heads,

attention_head_dim,

dropout=dropout,

cross_attention_dim=cross_attention_dim,

activation_fn=activation_fn,

attention_bias=attention_bias,

double_self_attention=double_self_attention,

norm_elementwise_affine=norm_elementwise_affine,

)

for d in range(num_layers)

]

)

self.proj_out = nn.Linear(inner_dim, in_channels)

def forward(

self,

hidden_states,

encoder_hidden_states=None,

timestep=None,

class_labels=None,

num_frames=1,

cross_attention_kwargs=None,

return_dict: bool = True,

):

"""

Args:

hidden_states ( When discrete, `torch.LongTensor` of shape `(batch size, num latent pixels)`.

When continous, `torch.FloatTensor` of shape `(batch size, channel, height, width)`): Input

hidden_states

encoder_hidden_states ( `torch.LongTensor` of shape `(batch size, encoder_hidden_states dim)`, *optional*):

Conditional embeddings for cross attention layer. If not given, cross-attention defaults to

self-attention.

timestep ( `torch.long`, *optional*):

Optional timestep to be applied as an embedding in AdaLayerNorm's. Used to indicate denoising step.

class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):

Optional class labels to be applied as an embedding in AdaLayerZeroNorm. Used to indicate class labels

conditioning.

return_dict (`bool`, *optional*, defaults to `True`):

Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.

Returns:

[`~models.transformer_2d.TransformerTemporalModelOutput`] or `tuple`:

[`~models.transformer_2d.TransformerTemporalModelOutput`] if `return_dict` is True, otherwise a `tuple`.

When returning a tuple, the first element is the sample tensor.

"""

# 1. Input

batch_frames, channel, height, width = hidden_states.shape

batch_size = batch_frames // num_frames

residual = hidden_states

hidden_states = hidden_states[None, :].reshape(batch_size, num_frames, channel, height, width)

hidden_states = hidden_states.permute(0, 2, 1, 3, 4)

hidden_states = self.norm(hidden_states)

hidden_states = hidden_states.permute(0, 3, 4, 2, 1).reshape(batch_size * height * width, num_frames, channel)

hidden_states = self.proj_in(hidden_states)

# 2. Blocks

for block in self.transformer_blocks:

hidden_states = block(

hidden_states,

encoder_hidden_states=encoder_hidden_states,

timestep=timestep,

cross_attention_kwargs=cross_attention_kwargs,

class_labels=class_labels,

)

# 3. Output

hidden_states = self.proj_out(hidden_states)

hidden_states = (

hidden_states[None, None, :]

.reshape(batch_size, height, width, channel, num_frames)

.permute(0, 3, 4, 1, 2)

.contiguous()

)

hidden_states = hidden_states.reshape(batch_frames, channel, height, width)

output = hidden_states + residual

if not return_dict:

return (output,)

return TransformerTemporalModelOutput(sample=output)