(checkpoint, config)
| 514 | |
| 515 | |
| 516 | def convert_ldm_vae_checkpoint(checkpoint, config): |
| 517 | # extract state dict for VAE |
| 518 | vae_state_dict = {} |
| 519 | vae_key = "first_stage_model." |
| 520 | keys = list(checkpoint.keys()) |
| 521 | for key in keys: |
| 522 | if key.startswith(vae_key): |
| 523 | vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key) |
| 524 | |
| 525 | new_checkpoint = {} |
| 526 | |
| 527 | new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"] |
| 528 | new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"] |
| 529 | new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"] |
| 530 | new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"] |
| 531 | new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"] |
| 532 | new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"] |
| 533 | |
| 534 | new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"] |
| 535 | new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"] |
| 536 | new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"] |
| 537 | new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"] |
| 538 | new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"] |
| 539 | new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"] |
| 540 | |
| 541 | new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"] |
| 542 | new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"] |
| 543 | new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"] |
| 544 | new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"] |
| 545 | |
| 546 | # Retrieves the keys for the encoder down blocks only |
| 547 | num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer}) |
| 548 | down_blocks = { |
| 549 | layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) |
| 550 | } |
| 551 | |
| 552 | # Retrieves the keys for the decoder up blocks only |
| 553 | num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer}) |
| 554 | up_blocks = { |
| 555 | layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) |
| 556 | } |
| 557 | |
| 558 | for i in range(num_down_blocks): |
| 559 | resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] |
| 560 | |
| 561 | if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: |
| 562 | new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop( |
| 563 | f"encoder.down.{i}.downsample.conv.weight" |
| 564 | ) |
| 565 | new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop( |
| 566 | f"encoder.down.{i}.downsample.conv.bias" |
| 567 | ) |
| 568 | |
| 569 | paths = renew_vae_resnet_paths(resnets) |
| 570 | meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"} |
| 571 | assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) |
| 572 | |
| 573 | mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] |
no test coverage detected