| 344 | |
| 345 | |
| 346 | class MLPProj(torch.nn.Module): |
| 347 | |
| 348 | def __init__(self, in_dim, out_dim, flf_pos_emb=False): |
| 349 | super().__init__() |
| 350 | |
| 351 | self.proj = torch.nn.Sequential( |
| 352 | torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim), |
| 353 | torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim), |
| 354 | torch.nn.LayerNorm(out_dim)) |
| 355 | if flf_pos_emb: # NOTE: we only use this for `flf2v` |
| 356 | self.emb_pos = nn.Parameter( |
| 357 | torch.zeros(1, FIRST_LAST_FRAME_CONTEXT_TOKEN_NUMBER, 1280)) |
| 358 | |
| 359 | def forward(self, image_embeds): |
| 360 | if hasattr(self, 'emb_pos'): |
| 361 | bs, n, d = image_embeds.shape |
| 362 | image_embeds = image_embeds.view(-1, 2 * n, d) |
| 363 | image_embeds = image_embeds + self.emb_pos |
| 364 | clip_extra_context_tokens = self.proj(image_embeds) |
| 365 | return clip_extra_context_tokens |
| 366 | |
| 367 | |
| 368 | class WanModel(ModelMixin, ConfigMixin): |