FilesExpand file tree

main

/

dual_transformer_2d.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

151 lines (135 loc) · 7.13 KB

main

/

dual_transformer_2d.py

Copy path

Top

File metadata and controls

• Code
• Blame

151 lines (135 loc) · 7.13 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

# 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 typing import Optional

from torch import nn

from .transformer_2d import Transformer2DModel, Transformer2DModelOutput

class DualTransformer2DModel(nn.Module):

"""

Dual transformer wrapper that combines two `Transformer2DModel`s for mixed inference.

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.1): 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`.

num_vector_embeds (`int`, *optional*):

Pass if the input is discrete. The number of classes of the vector embeddings of the latent pixels.

Includes the class for the masked latent pixel.

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

num_embeds_ada_norm ( `int`, *optional*): Pass if at least one of the norm_layers is `AdaLayerNorm`.

The number of diffusion steps used during training. Note that this is fixed at training time as it is used

to learn a number of embeddings that are added to the hidden states. During inference, you can denoise for

up to but not more than steps than `num_embeds_ada_norm`.

attention_bias (`bool`, *optional*):

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

"""

def __init__(

self,

num_attention_heads: int = 16,

attention_head_dim: int = 88,

in_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,

num_vector_embeds: Optional[int] = None,

activation_fn: str = "geglu",

num_embeds_ada_norm: Optional[int] = None,

):

super().__init__()

self.transformers = nn.ModuleList(

[

Transformer2DModel(

num_attention_heads=num_attention_heads,

attention_head_dim=attention_head_dim,

in_channels=in_channels,

num_layers=num_layers,

dropout=dropout,

norm_num_groups=norm_num_groups,

cross_attention_dim=cross_attention_dim,

attention_bias=attention_bias,

sample_size=sample_size,

num_vector_embeds=num_vector_embeds,

activation_fn=activation_fn,

num_embeds_ada_norm=num_embeds_ada_norm,

)

for _ in range(2)

]

)

# Variables that can be set by a pipeline:

# The ratio of transformer1 to transformer2's output states to be combined during inference

self.mix_ratio = 0.5

# The shape of `encoder_hidden_states` is expected to be

# `(batch_size, condition_lengths[0]+condition_lengths[1], num_features)`

self.condition_lengths = [77, 257]

# Which transformer to use to encode which condition.

# E.g. `(1, 0)` means that we'll use `transformers[1](conditions[0])` and `transformers[0](conditions[1])`

self.transformer_index_for_condition = [1, 0]

def forward(

self,

hidden_states,

encoder_hidden_states,

timestep=None,

attention_mask=None,

cross_attention_kwargs=None,

return_dict: bool = True,

):

"""

Args:

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

When continuous, `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.

attention_mask (`torch.FloatTensor`, *optional*):

Optional attention mask to be applied in Attention

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.Transformer2DModelOutput`] or `tuple`:

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

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

"""

input_states = hidden_states

encoded_states = []

tokens_start = 0

# attention_mask is not used yet

for i in range(2):

# for each of the two transformers, pass the corresponding condition tokens

condition_state = encoder_hidden_states[:, tokens_start : tokens_start + self.condition_lengths[i]]

transformer_index = self.transformer_index_for_condition[i]

encoded_state = self.transformers[transformer_index](

input_states,

encoder_hidden_states=condition_state,

timestep=timestep,

cross_attention_kwargs=cross_attention_kwargs,

return_dict=False,

)[0]

encoded_states.append(encoded_state - input_states)

tokens_start += self.condition_lengths[i]

output_states = encoded_states[0] * self.mix_ratio + encoded_states[1] * (1 - self.mix_ratio)

output_states = output_states + input_states

if not return_dict:

return (output_states,)

return Transformer2DModelOutput(sample=output_states)