()
| 961 | |
| 962 | |
| 963 | def convert_vae_22(): |
| 964 | vae_ckpt_path = hf_hub_download("Wan-AI/Wan2.2-TI2V-5B", "Wan2.2_VAE.pth") |
| 965 | old_state_dict = torch.load(vae_ckpt_path, weights_only=True) |
| 966 | new_state_dict = {} |
| 967 | |
| 968 | # Create mappings for specific components |
| 969 | middle_key_mapping = { |
| 970 | # Encoder middle block |
| 971 | "encoder.middle.0.residual.0.gamma": "encoder.mid_block.resnets.0.norm1.gamma", |
| 972 | "encoder.middle.0.residual.2.bias": "encoder.mid_block.resnets.0.conv1.bias", |
| 973 | "encoder.middle.0.residual.2.weight": "encoder.mid_block.resnets.0.conv1.weight", |
| 974 | "encoder.middle.0.residual.3.gamma": "encoder.mid_block.resnets.0.norm2.gamma", |
| 975 | "encoder.middle.0.residual.6.bias": "encoder.mid_block.resnets.0.conv2.bias", |
| 976 | "encoder.middle.0.residual.6.weight": "encoder.mid_block.resnets.0.conv2.weight", |
| 977 | "encoder.middle.2.residual.0.gamma": "encoder.mid_block.resnets.1.norm1.gamma", |
| 978 | "encoder.middle.2.residual.2.bias": "encoder.mid_block.resnets.1.conv1.bias", |
| 979 | "encoder.middle.2.residual.2.weight": "encoder.mid_block.resnets.1.conv1.weight", |
| 980 | "encoder.middle.2.residual.3.gamma": "encoder.mid_block.resnets.1.norm2.gamma", |
| 981 | "encoder.middle.2.residual.6.bias": "encoder.mid_block.resnets.1.conv2.bias", |
| 982 | "encoder.middle.2.residual.6.weight": "encoder.mid_block.resnets.1.conv2.weight", |
| 983 | # Decoder middle block |
| 984 | "decoder.middle.0.residual.0.gamma": "decoder.mid_block.resnets.0.norm1.gamma", |
| 985 | "decoder.middle.0.residual.2.bias": "decoder.mid_block.resnets.0.conv1.bias", |
| 986 | "decoder.middle.0.residual.2.weight": "decoder.mid_block.resnets.0.conv1.weight", |
| 987 | "decoder.middle.0.residual.3.gamma": "decoder.mid_block.resnets.0.norm2.gamma", |
| 988 | "decoder.middle.0.residual.6.bias": "decoder.mid_block.resnets.0.conv2.bias", |
| 989 | "decoder.middle.0.residual.6.weight": "decoder.mid_block.resnets.0.conv2.weight", |
| 990 | "decoder.middle.2.residual.0.gamma": "decoder.mid_block.resnets.1.norm1.gamma", |
| 991 | "decoder.middle.2.residual.2.bias": "decoder.mid_block.resnets.1.conv1.bias", |
| 992 | "decoder.middle.2.residual.2.weight": "decoder.mid_block.resnets.1.conv1.weight", |
| 993 | "decoder.middle.2.residual.3.gamma": "decoder.mid_block.resnets.1.norm2.gamma", |
| 994 | "decoder.middle.2.residual.6.bias": "decoder.mid_block.resnets.1.conv2.bias", |
| 995 | "decoder.middle.2.residual.6.weight": "decoder.mid_block.resnets.1.conv2.weight", |
| 996 | } |
| 997 | |
| 998 | # Create a mapping for attention blocks |
| 999 | attention_mapping = { |
| 1000 | # Encoder middle attention |
| 1001 | "encoder.middle.1.norm.gamma": "encoder.mid_block.attentions.0.norm.gamma", |
| 1002 | "encoder.middle.1.to_qkv.weight": "encoder.mid_block.attentions.0.to_qkv.weight", |
| 1003 | "encoder.middle.1.to_qkv.bias": "encoder.mid_block.attentions.0.to_qkv.bias", |
| 1004 | "encoder.middle.1.proj.weight": "encoder.mid_block.attentions.0.proj.weight", |
| 1005 | "encoder.middle.1.proj.bias": "encoder.mid_block.attentions.0.proj.bias", |
| 1006 | # Decoder middle attention |
| 1007 | "decoder.middle.1.norm.gamma": "decoder.mid_block.attentions.0.norm.gamma", |
| 1008 | "decoder.middle.1.to_qkv.weight": "decoder.mid_block.attentions.0.to_qkv.weight", |
| 1009 | "decoder.middle.1.to_qkv.bias": "decoder.mid_block.attentions.0.to_qkv.bias", |
| 1010 | "decoder.middle.1.proj.weight": "decoder.mid_block.attentions.0.proj.weight", |
| 1011 | "decoder.middle.1.proj.bias": "decoder.mid_block.attentions.0.proj.bias", |
| 1012 | } |
| 1013 | |
| 1014 | # Create a mapping for the head components |
| 1015 | head_mapping = { |
| 1016 | # Encoder head |
| 1017 | "encoder.head.0.gamma": "encoder.norm_out.gamma", |
| 1018 | "encoder.head.2.bias": "encoder.conv_out.bias", |
| 1019 | "encoder.head.2.weight": "encoder.conv_out.weight", |
| 1020 | # Decoder head |
no test coverage detected
searching dependent graphs…