| 127 | |
| 128 | |
| 129 | class DoubleStreamBlock(nn.Module): |
| 130 | def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False): |
| 131 | super().__init__() |
| 132 | |
| 133 | mlp_hidden_dim = int(hidden_size * mlp_ratio) |
| 134 | self.num_heads = num_heads |
| 135 | self.hidden_size = hidden_size |
| 136 | self.img_mod = Modulation(hidden_size, double=True) |
| 137 | self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 138 | self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) |
| 139 | |
| 140 | self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 141 | self.img_mlp = nn.Sequential( |
| 142 | nn.Linear(hidden_size, mlp_hidden_dim, bias=True), |
| 143 | nn.GELU(approximate="tanh"), |
| 144 | nn.Linear(mlp_hidden_dim, hidden_size, bias=True), |
| 145 | ) |
| 146 | |
| 147 | self.txt_mod = Modulation(hidden_size, double=True) |
| 148 | self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 149 | self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) |
| 150 | |
| 151 | self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 152 | self.txt_mlp = nn.Sequential( |
| 153 | nn.Linear(hidden_size, mlp_hidden_dim, bias=True), |
| 154 | nn.GELU(approximate="tanh"), |
| 155 | nn.Linear(mlp_hidden_dim, hidden_size, bias=True), |
| 156 | ) |
| 157 | |
| 158 | def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]: |
| 159 | img_mod1, img_mod2 = self.img_mod(vec) |
| 160 | txt_mod1, txt_mod2 = self.txt_mod(vec) |
| 161 | |
| 162 | # prepare image for attention |
| 163 | img_modulated = self.img_norm1(img) |
| 164 | img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift |
| 165 | img_qkv = self.img_attn.qkv(img_modulated) |
| 166 | img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) |
| 167 | img_q, img_k = self.img_attn.norm(img_q, img_k, img_v) |
| 168 | |
| 169 | # prepare txt for attention |
| 170 | txt_modulated = self.txt_norm1(txt) |
| 171 | txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift |
| 172 | txt_qkv = self.txt_attn.qkv(txt_modulated) |
| 173 | txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) |
| 174 | txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v) |
| 175 | |
| 176 | # run actual attention |
| 177 | q = torch.cat((txt_q, img_q), dim=2) |
| 178 | k = torch.cat((txt_k, img_k), dim=2) |
| 179 | v = torch.cat((txt_v, img_v), dim=2) |
| 180 | |
| 181 | attn = attention(q, k, v, pe=pe) |
| 182 | txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :] |
| 183 | |
| 184 | # calculate the img bloks |
| 185 | img = img + img_mod1.gate * self.img_attn.proj(img_attn) |
| 186 | img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift) |