(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False)
| 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) |
nothing calls this directly
no test coverage detected