The final decoder of FlowMatchingTransformer.
| 170 | return x |
| 171 | |
| 172 | class Decoder(nn.Module): |
| 173 | """ |
| 174 | The final decoder of FlowMatchingTransformer. |
| 175 | """ |
| 176 | def __init__(self, hidden_size, dim_w): |
| 177 | super().__init__() |
| 178 | self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 179 | self.adaLN_modulation = nn.Sequential( |
| 180 | nn.SiLU(), |
| 181 | nn.Linear(hidden_size, 2 * hidden_size, bias=True) |
| 182 | ) |
| 183 | self.linear = nn.Linear(hidden_size, dim_w, bias=True) |
| 184 | |
| 185 | def framewise_modulate(self, x, shift, scale) -> torch.Tensor: |
| 186 | return x * (1 + scale) + shift |
| 187 | |
| 188 | def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor: |
| 189 | shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1) |
| 190 | x = self.framewise_modulate(self.norm_final(x), shift, scale) |
| 191 | return self.linear(x) |
| 192 | |
| 193 | |
| 194 | class FlowMatchingTransformer(BaseModel): |