(self,
model_type='t2v',
patch_size=(1, 2, 2),
text_len=512,
in_dim=16,
dim=2048,
ffn_dim=8192,
freq_dim=256,
text_dim=4096,
out_dim=16,
num_heads=16,
num_layers=32,
window_size=(-1, -1),
qk_norm=True,
cross_attn_norm=True,
eps=1e-6,
ref_motion_dim=16,
**kwargs)
| 555 | class WanModelTM2M(nn.Module): |
| 556 | |
| 557 | def __init__(self, |
| 558 | model_type='t2v', |
| 559 | patch_size=(1, 2, 2), |
| 560 | text_len=512, |
| 561 | in_dim=16, |
| 562 | dim=2048, |
| 563 | ffn_dim=8192, |
| 564 | freq_dim=256, |
| 565 | text_dim=4096, |
| 566 | out_dim=16, |
| 567 | num_heads=16, |
| 568 | num_layers=32, |
| 569 | window_size=(-1, -1), |
| 570 | qk_norm=True, |
| 571 | cross_attn_norm=True, |
| 572 | eps=1e-6, |
| 573 | ref_motion_dim=16, |
| 574 | **kwargs): |
| 575 | super().__init__() |
| 576 | |
| 577 | assert model_type in ['t2v', 'i2v'] |
| 578 | self.model_type = model_type |
| 579 | |
| 580 | self.patch_size = patch_size |
| 581 | self.text_len = text_len |
| 582 | self.in_dim = in_dim |
| 583 | self.dim = dim |
| 584 | self.ffn_dim = ffn_dim |
| 585 | self.freq_dim = freq_dim |
| 586 | self.text_dim = text_dim |
| 587 | self.out_dim = out_dim |
| 588 | self.num_heads = num_heads |
| 589 | self.num_layers = num_layers |
| 590 | self.window_size = window_size |
| 591 | self.qk_norm = qk_norm |
| 592 | self.cross_attn_norm = cross_attn_norm |
| 593 | self.eps = eps |
| 594 | |
| 595 | # embeddings |
| 596 | self.motion_embedding = nn.Linear(in_dim, dim) |
| 597 | self.text_embedding = nn.Sequential( |
| 598 | nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'), |
| 599 | nn.Linear(dim, dim)) |
| 600 | |
| 601 | self.time_embedding = nn.Sequential( |
| 602 | nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) |
| 603 | self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) |
| 604 | |
| 605 | self.motion_pos_embedding = PositionalEncoding(d_model=dim, max_len=5000) |
| 606 | |
| 607 | # ref motion |
| 608 | self.ref_motion_embedding = nn.Linear(ref_motion_dim, dim) |
| 609 | # self.ref_motion_embedding = nn.Sequential( |
| 610 | # nn.Linear(ref_motion_dim, dim), |
| 611 | # nn.SiLU(), nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim), |
| 612 | # nn.SiLU(), nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim)) |
| 613 | self.ref_motion_pos_embedding = PositionalEncoding(d_model=dim, max_len=5000) |
| 614 | self.ref_motion_encoder = nn.ModuleList([ |
nothing calls this directly
no test coverage detected