r""" Wan diffusion backbone supporting both text-to-video and image-to-video.
| 429 | |
| 430 | |
| 431 | class WanModel(ModelMixin, ConfigMixin): |
| 432 | r""" |
| 433 | Wan diffusion backbone supporting both text-to-video and image-to-video. |
| 434 | """ |
| 435 | |
| 436 | ignore_for_config = [ |
| 437 | 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size' |
| 438 | ] |
| 439 | _no_split_modules = ['WanAttentionBlock'] |
| 440 | |
| 441 | @register_to_config |
| 442 | def __init__(self, |
| 443 | model_type='i2v', |
| 444 | patch_size=(1, 2, 2), |
| 445 | text_len=512, |
| 446 | in_dim=16, |
| 447 | dim=2048, |
| 448 | ffn_dim=8192, |
| 449 | freq_dim=256, |
| 450 | text_dim=4096, |
| 451 | out_dim=16, |
| 452 | num_heads=16, |
| 453 | num_layers=32, |
| 454 | window_size=(-1, -1), |
| 455 | qk_norm=True, |
| 456 | cross_attn_norm=True, |
| 457 | eps=1e-6, |
| 458 | # audio params |
| 459 | audio_window=5, |
| 460 | intermediate_dim=512, |
| 461 | output_dim=768, |
| 462 | context_tokens=32, |
| 463 | vae_scale=4, # vae timedownsample scale |
| 464 | |
| 465 | norm_input_visual=True, |
| 466 | norm_output_audio=True, |
| 467 | weight_init=True): |
| 468 | super().__init__() |
| 469 | |
| 470 | assert model_type == 'i2v', 'MultiTalk model requires your model_type is i2v.' |
| 471 | self.model_type = model_type |
| 472 | |
| 473 | self.patch_size = patch_size |
| 474 | self.text_len = text_len |
| 475 | self.in_dim = in_dim |
| 476 | self.dim = dim |
| 477 | self.ffn_dim = ffn_dim |
| 478 | self.freq_dim = freq_dim |
| 479 | self.text_dim = text_dim |
| 480 | self.out_dim = out_dim |
| 481 | self.num_heads = num_heads |
| 482 | self.num_layers = num_layers |
| 483 | self.window_size = window_size |
| 484 | self.qk_norm = qk_norm |
| 485 | self.cross_attn_norm = cross_attn_norm |
| 486 | self.eps = eps |
| 487 | |
| 488 |
nothing calls this directly
no outgoing calls
no test coverage detected