| 241 | |
| 242 | |
| 243 | class FinalLayer(nn.Module): |
| 244 | def __init__(self, model): |
| 245 | super().__init__() |
| 246 | self.final_layer = model.final_layer |
| 247 | self.model = [model] |
| 248 | |
| 249 | def __getattr__(self, name): |
| 250 | return getattr(self.model[0], name) |
| 251 | |
| 252 | def forward(self, inputs): |
| 253 | h, attn_mask, adaln_input, sizes, *freqs_cis = inputs |
| 254 | out = self.final_layer(h, adaln_input) |
| 255 | L_text, gh, gw = sizes |
| 256 | out = -self._tokens_to_img(out[:, L_text:], gh, gw) |
| 257 | return out |
| 258 | |
| 259 | |
| 260 | class Wrapper(nn.Module): |