(self, hidden_size: int, num_heads: int, mlp_ratio: float,
qkv_bias: bool = False, backend='pytorch')
| 412 | """ |
| 413 | |
| 414 | def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, |
| 415 | qkv_bias: bool = False, backend='pytorch'): |
| 416 | super().__init__(hidden_size, num_heads, mlp_ratio, |
| 417 | qkv_bias, backend) |
| 418 | mlp_hidden_dim = int(hidden_size * mlp_ratio) |
| 419 | self.edit_mod = Modulation(hidden_size, double=True) |
| 420 | self.edit_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 421 | self.edit_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) |
| 422 | |
| 423 | self.edit_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 424 | self.edit_mlp = nn.Sequential( |
| 425 | nn.Linear(hidden_size, mlp_hidden_dim, bias=True), |
| 426 | nn.GELU(approximate="tanh"), |
| 427 | nn.Linear(mlp_hidden_dim, hidden_size, bias=True), |
| 428 | ) |
| 429 | |
| 430 | def forward(self, x: Tensor, vec: Tensor, |
| 431 | pe: Tensor, mask: Tensor = None, |
nothing calls this directly
no test coverage detected