()
| 366 | |
| 367 | |
| 368 | def convert_vae(): |
| 369 | vae_ckpt_path = hf_hub_download("Wan-AI/Wan2.1-T2V-14B", "Wan2.1_VAE.pth") |
| 370 | old_state_dict = torch.load(vae_ckpt_path, weights_only=True) |
| 371 | new_state_dict = {} |
| 372 | |
| 373 | # Create mappings for specific components |
| 374 | middle_key_mapping = { |
| 375 | # Encoder middle block |
| 376 | "encoder.middle.0.residual.0.gamma": "encoder.mid_block.resnets.0.norm1.gamma", |
| 377 | "encoder.middle.0.residual.2.bias": "encoder.mid_block.resnets.0.conv1.bias", |
| 378 | "encoder.middle.0.residual.2.weight": "encoder.mid_block.resnets.0.conv1.weight", |
| 379 | "encoder.middle.0.residual.3.gamma": "encoder.mid_block.resnets.0.norm2.gamma", |
| 380 | "encoder.middle.0.residual.6.bias": "encoder.mid_block.resnets.0.conv2.bias", |
| 381 | "encoder.middle.0.residual.6.weight": "encoder.mid_block.resnets.0.conv2.weight", |
| 382 | "encoder.middle.2.residual.0.gamma": "encoder.mid_block.resnets.1.norm1.gamma", |
| 383 | "encoder.middle.2.residual.2.bias": "encoder.mid_block.resnets.1.conv1.bias", |
| 384 | "encoder.middle.2.residual.2.weight": "encoder.mid_block.resnets.1.conv1.weight", |
| 385 | "encoder.middle.2.residual.3.gamma": "encoder.mid_block.resnets.1.norm2.gamma", |
| 386 | "encoder.middle.2.residual.6.bias": "encoder.mid_block.resnets.1.conv2.bias", |
| 387 | "encoder.middle.2.residual.6.weight": "encoder.mid_block.resnets.1.conv2.weight", |
| 388 | # Decoder middle block |
| 389 | "decoder.middle.0.residual.0.gamma": "decoder.mid_block.resnets.0.norm1.gamma", |
| 390 | "decoder.middle.0.residual.2.bias": "decoder.mid_block.resnets.0.conv1.bias", |
| 391 | "decoder.middle.0.residual.2.weight": "decoder.mid_block.resnets.0.conv1.weight", |
| 392 | "decoder.middle.0.residual.3.gamma": "decoder.mid_block.resnets.0.norm2.gamma", |
| 393 | "decoder.middle.0.residual.6.bias": "decoder.mid_block.resnets.0.conv2.bias", |
| 394 | "decoder.middle.0.residual.6.weight": "decoder.mid_block.resnets.0.conv2.weight", |
| 395 | "decoder.middle.2.residual.0.gamma": "decoder.mid_block.resnets.1.norm1.gamma", |
| 396 | "decoder.middle.2.residual.2.bias": "decoder.mid_block.resnets.1.conv1.bias", |
| 397 | "decoder.middle.2.residual.2.weight": "decoder.mid_block.resnets.1.conv1.weight", |
| 398 | "decoder.middle.2.residual.3.gamma": "decoder.mid_block.resnets.1.norm2.gamma", |
| 399 | "decoder.middle.2.residual.6.bias": "decoder.mid_block.resnets.1.conv2.bias", |
| 400 | "decoder.middle.2.residual.6.weight": "decoder.mid_block.resnets.1.conv2.weight", |
| 401 | } |
| 402 | |
| 403 | # Create a mapping for attention blocks |
| 404 | attention_mapping = { |
| 405 | # Encoder middle attention |
| 406 | "encoder.middle.1.norm.gamma": "encoder.mid_block.attentions.0.norm.gamma", |
| 407 | "encoder.middle.1.to_qkv.weight": "encoder.mid_block.attentions.0.to_qkv.weight", |
| 408 | "encoder.middle.1.to_qkv.bias": "encoder.mid_block.attentions.0.to_qkv.bias", |
| 409 | "encoder.middle.1.proj.weight": "encoder.mid_block.attentions.0.proj.weight", |
| 410 | "encoder.middle.1.proj.bias": "encoder.mid_block.attentions.0.proj.bias", |
| 411 | # Decoder middle attention |
| 412 | "decoder.middle.1.norm.gamma": "decoder.mid_block.attentions.0.norm.gamma", |
| 413 | "decoder.middle.1.to_qkv.weight": "decoder.mid_block.attentions.0.to_qkv.weight", |
| 414 | "decoder.middle.1.to_qkv.bias": "decoder.mid_block.attentions.0.to_qkv.bias", |
| 415 | "decoder.middle.1.proj.weight": "decoder.mid_block.attentions.0.proj.weight", |
| 416 | "decoder.middle.1.proj.bias": "decoder.mid_block.attentions.0.proj.bias", |
| 417 | } |
| 418 | |
| 419 | # Create a mapping for the head components |
| 420 | head_mapping = { |
| 421 | # Encoder head |
| 422 | "encoder.head.0.gamma": "encoder.norm_out.gamma", |
| 423 | "encoder.head.2.bias": "encoder.conv_out.bias", |
| 424 | "encoder.head.2.weight": "encoder.conv_out.weight", |
| 425 | # Decoder head |
no test coverage detected
searching dependent graphs…