(
self,
vace_layers=(0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28),
vace_in_dim=96,
patch_size=(1, 2, 2),
has_image_input=False,
dim=1536,
num_heads=12,
ffn_dim=8960,
eps=1e-6,
)
| 26 | |
| 27 | class VaceWanModel(torch.nn.Module): |
| 28 | def __init__( |
| 29 | self, |
| 30 | vace_layers=(0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28), |
| 31 | vace_in_dim=96, |
| 32 | patch_size=(1, 2, 2), |
| 33 | has_image_input=False, |
| 34 | dim=1536, |
| 35 | num_heads=12, |
| 36 | ffn_dim=8960, |
| 37 | eps=1e-6, |
| 38 | ): |
| 39 | super().__init__() |
| 40 | self.vace_layers = vace_layers |
| 41 | self.vace_in_dim = vace_in_dim |
| 42 | self.vace_layers_mapping = {i: n for n, i in enumerate(self.vace_layers)} |
| 43 | |
| 44 | # vace blocks |
| 45 | self.vace_blocks = torch.nn.ModuleList([ |
| 46 | VaceWanAttentionBlock(has_image_input, dim, num_heads, ffn_dim, eps, block_id=i) |
| 47 | for i in self.vace_layers |
| 48 | ]) |
| 49 | |
| 50 | # vace patch embeddings |
| 51 | self.vace_patch_embedding = torch.nn.Conv3d(vace_in_dim, dim, kernel_size=patch_size, stride=patch_size) |
| 52 | |
| 53 | def forward(self, x, vace_context, context, t_mod, freqs): |
| 54 | c = [self.vace_patch_embedding(u.unsqueeze(0)) for u in vace_context] |
no test coverage detected