| 511 | return x |
| 512 | |
| 513 | class WanModelT2M(nn.Module): |
| 514 | |
| 515 | def __init__(self, |
| 516 | model_type='t2v', |
| 517 | patch_size=(1, 2, 2), |
| 518 | text_len=512, |
| 519 | in_dim=16, |
| 520 | dim=2048, |
| 521 | ffn_dim=8192, |
| 522 | freq_dim=256, |
| 523 | text_dim=4096, |
| 524 | out_dim=16, |
| 525 | num_heads=16, |
| 526 | num_layers=32, |
| 527 | window_size=(-1, -1), |
| 528 | qk_norm=True, |
| 529 | cross_attn_norm=True, |
| 530 | eps=1e-6, |
| 531 | split_head=False, |
| 532 | **kwargs): |
| 533 | super().__init__() |
| 534 | |
| 535 | assert model_type in ['t2v', 'i2v'] |
| 536 | self.model_type = model_type |
| 537 | |
| 538 | self.patch_size = patch_size |
| 539 | self.text_len = text_len |
| 540 | self.in_dim = in_dim |
| 541 | self.dim = dim |
| 542 | self.ffn_dim = ffn_dim |
| 543 | self.freq_dim = freq_dim |
| 544 | self.text_dim = text_dim |
| 545 | self.out_dim = out_dim |
| 546 | self.num_heads = num_heads |
| 547 | self.num_layers = num_layers |
| 548 | self.window_size = window_size |
| 549 | self.qk_norm = qk_norm |
| 550 | self.cross_attn_norm = cross_attn_norm |
| 551 | self.eps = eps |
| 552 | |
| 553 | # embeddings |
| 554 | self.motion_embedding = nn.Linear(in_dim, dim) |
| 555 | self.text_embedding = nn.Sequential( |
| 556 | nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'), |
| 557 | nn.Linear(dim, dim)) |
| 558 | |
| 559 | self.time_embedding = nn.Sequential( |
| 560 | nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) |
| 561 | self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) |
| 562 | |
| 563 | self.motion_pos_embedding = PositionalEncoding(d_model=dim, max_len=5000) |
| 564 | |
| 565 | # blocks |
| 566 | cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn' |
| 567 | self.blocks = nn.ModuleList([ |
| 568 | WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads, |
| 569 | window_size, qk_norm, cross_attn_norm, eps) |
| 570 | for _ in range(num_layers) |