| 207 | |
| 208 | |
| 209 | class Modulation(nn.Module): |
| 210 | def __init__(self, dim: int, double: bool): |
| 211 | super().__init__() |
| 212 | self.is_double = double |
| 213 | self.multiplier = 6 if double else 3 |
| 214 | self.lin = nn.Linear(dim, self.multiplier * dim, bias=True) |
| 215 | |
| 216 | def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]: |
| 217 | out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1) |
| 218 | |
| 219 | return ( |
| 220 | ModulationOut(*out[:3]), |
| 221 | ModulationOut(*out[3:]) if self.is_double else None, |
| 222 | ) |
| 223 | |
| 224 | |
| 225 | class DoubleStreamBlock(nn.Module): |