(new_checkpoint, unet_state_dict, old_path, new_path, config)
| 599 | |
| 600 | |
| 601 | def assign_attention_to_checkpoint(new_checkpoint, unet_state_dict, old_path, new_path, config): |
| 602 | qkv_weight = unet_state_dict.pop(f"{old_path}.qkv.weight") |
| 603 | qkv_weight = qkv_weight[:, :, 0] |
| 604 | |
| 605 | qkv_bias = unet_state_dict.pop(f"{old_path}.qkv.bias") |
| 606 | |
| 607 | is_cross_attn_only = "only_cross_attention" in config and config["only_cross_attention"] |
| 608 | |
| 609 | split = 1 if is_cross_attn_only else 3 |
| 610 | |
| 611 | weights, bias = split_attentions( |
| 612 | weight=qkv_weight, |
| 613 | bias=qkv_bias, |
| 614 | split=split, |
| 615 | chunk_size=config["attention_head_dim"], |
| 616 | ) |
| 617 | |
| 618 | if is_cross_attn_only: |
| 619 | query_weight, q_bias = weights, bias |
| 620 | new_checkpoint[f"{new_path}.to_q.weight"] = query_weight[0] |
| 621 | new_checkpoint[f"{new_path}.to_q.bias"] = q_bias[0] |
| 622 | else: |
| 623 | [query_weight, key_weight, value_weight], [q_bias, k_bias, v_bias] = weights, bias |
| 624 | new_checkpoint[f"{new_path}.to_q.weight"] = query_weight |
| 625 | new_checkpoint[f"{new_path}.to_q.bias"] = q_bias |
| 626 | new_checkpoint[f"{new_path}.to_k.weight"] = key_weight |
| 627 | new_checkpoint[f"{new_path}.to_k.bias"] = k_bias |
| 628 | new_checkpoint[f"{new_path}.to_v.weight"] = value_weight |
| 629 | new_checkpoint[f"{new_path}.to_v.bias"] = v_bias |
| 630 | |
| 631 | encoder_kv_weight = unet_state_dict.pop(f"{old_path}.encoder_kv.weight") |
| 632 | encoder_kv_weight = encoder_kv_weight[:, :, 0] |
| 633 | |
| 634 | encoder_kv_bias = unet_state_dict.pop(f"{old_path}.encoder_kv.bias") |
| 635 | |
| 636 | [encoder_k_weight, encoder_v_weight], [encoder_k_bias, encoder_v_bias] = split_attentions( |
| 637 | weight=encoder_kv_weight, |
| 638 | bias=encoder_kv_bias, |
| 639 | split=2, |
| 640 | chunk_size=config["attention_head_dim"], |
| 641 | ) |
| 642 | |
| 643 | new_checkpoint[f"{new_path}.add_k_proj.weight"] = encoder_k_weight |
| 644 | new_checkpoint[f"{new_path}.add_k_proj.bias"] = encoder_k_bias |
| 645 | new_checkpoint[f"{new_path}.add_v_proj.weight"] = encoder_v_weight |
| 646 | new_checkpoint[f"{new_path}.add_v_proj.bias"] = encoder_v_bias |
| 647 | |
| 648 | |
| 649 | def assign_to_checkpoint(paths, checkpoint, old_checkpoint, additional_replacements=None, config=None): |
no test coverage detected