(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False, backend = 'pytorch')
| 224 | |
| 225 | class DoubleStreamBlock(nn.Module): |
| 226 | def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False, backend = 'pytorch'): |
| 227 | super().__init__() |
| 228 | |
| 229 | mlp_hidden_dim = int(hidden_size * mlp_ratio) |
| 230 | self.num_heads = num_heads |
| 231 | self.hidden_size = hidden_size |
| 232 | self.img_mod = Modulation(hidden_size, double=True) |
| 233 | self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 234 | self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) |
| 235 | |
| 236 | self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 237 | self.img_mlp = nn.Sequential( |
| 238 | nn.Linear(hidden_size, mlp_hidden_dim, bias=True), |
| 239 | nn.GELU(approximate="tanh"), |
| 240 | nn.Linear(mlp_hidden_dim, hidden_size, bias=True), |
| 241 | ) |
| 242 | |
| 243 | self.backend = backend |
| 244 | |
| 245 | self.txt_mod = Modulation(hidden_size, double=True) |
| 246 | self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 247 | self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) |
| 248 | |
| 249 | self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 250 | self.txt_mlp = nn.Sequential( |
| 251 | nn.Linear(hidden_size, mlp_hidden_dim, bias=True), |
| 252 | nn.GELU(approximate="tanh"), |
| 253 | nn.Linear(mlp_hidden_dim, hidden_size, bias=True), |
| 254 | ) |
| 255 | |
| 256 | |
| 257 |
nothing calls this directly
no test coverage detected