(
self,
dim: int,
in_dim: int,
ffn_dim: int,
out_dim: int,
text_dim: int,
freq_dim: int,
eps: float,
patch_size: Tuple[int, int, int],
num_heads: int,
num_layers: int,
# init_context: torch.Tensor, # <<<< 必填：在 __init__ 里用它生成 cross-attn KV 缓存
has_image_input: bool = False,
)
| 515 | # ---------------------------- |
| 516 | class WanModel(torch.nn.Module): |
| 517 | def __init__( |
| 518 | self, |
| 519 | dim: int, |
| 520 | in_dim: int, |
| 521 | ffn_dim: int, |
| 522 | out_dim: int, |
| 523 | text_dim: int, |
| 524 | freq_dim: int, |
| 525 | eps: float, |
| 526 | patch_size: Tuple[int, int, int], |
| 527 | num_heads: int, |
| 528 | num_layers: int, |
| 529 | # init_context: torch.Tensor, # <<<< 必填：在 __init__ 里用它生成 cross-attn KV 缓存 |
| 530 | has_image_input: bool = False, |
| 531 | ): |
| 532 | super().__init__() |
| 533 | self.dim = dim |
| 534 | self.freq_dim = freq_dim |
| 535 | self.patch_size = patch_size |
| 536 | |
| 537 | # patch embed |
| 538 | self.patch_embedding = nn.Conv3d( |
| 539 | in_dim, dim, kernel_size=patch_size, stride=patch_size) |
| 540 | |
| 541 | # text / time embed |
| 542 | self.text_embedding = nn.Sequential( |
| 543 | nn.Linear(text_dim, dim), |
| 544 | nn.GELU(approximate='tanh'), |
| 545 | nn.Linear(dim, dim) |
| 546 | ) |
| 547 | self.time_embedding = nn.Sequential( |
| 548 | nn.Linear(freq_dim, dim), |
| 549 | nn.SiLU(), |
| 550 | nn.Linear(dim, dim) |
| 551 | ) |
| 552 | self.time_projection = nn.Sequential( |
| 553 | nn.SiLU(), nn.Linear(dim, dim * 6)) |
| 554 | |
| 555 | # blocks |
| 556 | self.blocks = nn.ModuleList([ |
| 557 | DiTBlock(dim, num_heads, ffn_dim, eps) |
| 558 | for _ in range(num_layers) |
| 559 | ]) |
| 560 | self.head = Head(dim, out_dim, patch_size, eps) |
| 561 | |
| 562 | head_dim = dim // num_heads |
| 563 | self.freqs = precompute_freqs_cis_3d(head_dim) |
| 564 | |
| 565 | self._cross_kv_initialized = False |
| 566 | |
| 567 | # 可选：手动清空 / 重新初始化 |
| 568 | def clear_cross_kv(self): |
nothing calls this directly
no test coverage detected