()
| 644 | |
| 645 | |
| 646 | def convert_vae(): |
| 647 | vae_ckpt_path = hf_hub_download("Wan-AI/Wan2.1-T2V-14B", "Wan2.1_VAE.pth") |
| 648 | old_state_dict = torch.load(vae_ckpt_path, weights_only=True) |
| 649 | new_state_dict = {} |
| 650 | |
| 651 | # Create mappings for specific components |
| 652 | middle_key_mapping = { |
| 653 | # Encoder middle block |
| 654 | "encoder.middle.0.residual.0.gamma": "encoder.mid_block.resnets.0.norm1.gamma", |
| 655 | "encoder.middle.0.residual.2.bias": "encoder.mid_block.resnets.0.conv1.bias", |
| 656 | "encoder.middle.0.residual.2.weight": "encoder.mid_block.resnets.0.conv1.weight", |
| 657 | "encoder.middle.0.residual.3.gamma": "encoder.mid_block.resnets.0.norm2.gamma", |
| 658 | "encoder.middle.0.residual.6.bias": "encoder.mid_block.resnets.0.conv2.bias", |
| 659 | "encoder.middle.0.residual.6.weight": "encoder.mid_block.resnets.0.conv2.weight", |
| 660 | "encoder.middle.2.residual.0.gamma": "encoder.mid_block.resnets.1.norm1.gamma", |
| 661 | "encoder.middle.2.residual.2.bias": "encoder.mid_block.resnets.1.conv1.bias", |
| 662 | "encoder.middle.2.residual.2.weight": "encoder.mid_block.resnets.1.conv1.weight", |
| 663 | "encoder.middle.2.residual.3.gamma": "encoder.mid_block.resnets.1.norm2.gamma", |
| 664 | "encoder.middle.2.residual.6.bias": "encoder.mid_block.resnets.1.conv2.bias", |
| 665 | "encoder.middle.2.residual.6.weight": "encoder.mid_block.resnets.1.conv2.weight", |
| 666 | # Decoder middle block |
| 667 | "decoder.middle.0.residual.0.gamma": "decoder.mid_block.resnets.0.norm1.gamma", |
| 668 | "decoder.middle.0.residual.2.bias": "decoder.mid_block.resnets.0.conv1.bias", |
| 669 | "decoder.middle.0.residual.2.weight": "decoder.mid_block.resnets.0.conv1.weight", |
| 670 | "decoder.middle.0.residual.3.gamma": "decoder.mid_block.resnets.0.norm2.gamma", |
| 671 | "decoder.middle.0.residual.6.bias": "decoder.mid_block.resnets.0.conv2.bias", |
| 672 | "decoder.middle.0.residual.6.weight": "decoder.mid_block.resnets.0.conv2.weight", |
| 673 | "decoder.middle.2.residual.0.gamma": "decoder.mid_block.resnets.1.norm1.gamma", |
| 674 | "decoder.middle.2.residual.2.bias": "decoder.mid_block.resnets.1.conv1.bias", |
| 675 | "decoder.middle.2.residual.2.weight": "decoder.mid_block.resnets.1.conv1.weight", |
| 676 | "decoder.middle.2.residual.3.gamma": "decoder.mid_block.resnets.1.norm2.gamma", |
| 677 | "decoder.middle.2.residual.6.bias": "decoder.mid_block.resnets.1.conv2.bias", |
| 678 | "decoder.middle.2.residual.6.weight": "decoder.mid_block.resnets.1.conv2.weight", |
| 679 | } |
| 680 | |
| 681 | # Create a mapping for attention blocks |
| 682 | attention_mapping = { |
| 683 | # Encoder middle attention |
| 684 | "encoder.middle.1.norm.gamma": "encoder.mid_block.attentions.0.norm.gamma", |
| 685 | "encoder.middle.1.to_qkv.weight": "encoder.mid_block.attentions.0.to_qkv.weight", |
| 686 | "encoder.middle.1.to_qkv.bias": "encoder.mid_block.attentions.0.to_qkv.bias", |
| 687 | "encoder.middle.1.proj.weight": "encoder.mid_block.attentions.0.proj.weight", |
| 688 | "encoder.middle.1.proj.bias": "encoder.mid_block.attentions.0.proj.bias", |
| 689 | # Decoder middle attention |
| 690 | "decoder.middle.1.norm.gamma": "decoder.mid_block.attentions.0.norm.gamma", |
| 691 | "decoder.middle.1.to_qkv.weight": "decoder.mid_block.attentions.0.to_qkv.weight", |
| 692 | "decoder.middle.1.to_qkv.bias": "decoder.mid_block.attentions.0.to_qkv.bias", |
| 693 | "decoder.middle.1.proj.weight": "decoder.mid_block.attentions.0.proj.weight", |
| 694 | "decoder.middle.1.proj.bias": "decoder.mid_block.attentions.0.proj.bias", |
| 695 | } |
| 696 | |
| 697 | # Create a mapping for the head components |
| 698 | head_mapping = { |
| 699 | # Encoder head |
| 700 | "encoder.head.0.gamma": "encoder.norm_out.gamma", |
| 701 | "encoder.head.2.bias": "encoder.conv_out.bias", |
| 702 | "encoder.head.2.weight": "encoder.conv_out.weight", |
| 703 | # Decoder head |
no test coverage detected
searching dependent graphs…