A DiT block with parallel linear layers as described in https://arxiv.org/abs/2302.05442 and adapted modulation interface.
| 296 | |
| 297 | |
| 298 | class SingleStreamBlock(nn.Module): |
| 299 | """ |
| 300 | A DiT block with parallel linear layers as described in |
| 301 | https://arxiv.org/abs/2302.05442 and adapted modulation interface. |
| 302 | """ |
| 303 | |
| 304 | def __init__( |
| 305 | self, |
| 306 | hidden_size: int, |
| 307 | num_heads: int, |
| 308 | mlp_ratio: float = 4.0, |
| 309 | qk_scale: float | None = None, |
| 310 | backend='pytorch' |
| 311 | ): |
| 312 | super().__init__() |
| 313 | self.hidden_dim = hidden_size |
| 314 | self.num_heads = num_heads |
| 315 | head_dim = hidden_size // num_heads |
| 316 | self.scale = qk_scale or head_dim**-0.5 |
| 317 | |
| 318 | self.mlp_hidden_dim = int(hidden_size * mlp_ratio) |
| 319 | # qkv and mlp_in |
| 320 | self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim) |
| 321 | # proj and mlp_out |
| 322 | self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size) |
| 323 | |
| 324 | self.norm = QKNorm(head_dim) |
| 325 | |
| 326 | self.hidden_size = hidden_size |
| 327 | self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 328 | |
| 329 | self.mlp_act = nn.GELU(approximate="tanh") |
| 330 | self.modulation = Modulation(hidden_size, double=False) |
| 331 | self.backend = backend |
| 332 | |
| 333 | def forward(self, x: Tensor, vec: Tensor, pe: Tensor, mask: Tensor = None) -> Tensor: |
| 334 | mod, _ = self.modulation(vec) |
| 335 | x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift |
| 336 | qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1) |
| 337 | |
| 338 | q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) |
| 339 | q, k = self.norm(q, k, v) |
| 340 | if mask is not None: |
| 341 | mask = repeat(mask, 'B L S-> B H L S', H=self.num_heads) |
| 342 | # compute attention |
| 343 | attn = attention(q, k, v, pe=pe, mask = mask, backend=self.backend) |
| 344 | # compute activation in mlp stream, cat again and run second linear layer |
| 345 | output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)) |
| 346 | return x + mod.gate * output |
| 347 | |
| 348 | |
| 349 | class DoubleStreamBlockC(DoubleStreamBlock): |