Updates paths inside attentions to the new naming scheme (local renaming)
(old_list, n_shave_prefix_segments=0)
| 569 | |
| 570 | |
| 571 | def renew_attention_paths(old_list, n_shave_prefix_segments=0): |
| 572 | """ |
| 573 | Updates paths inside attentions to the new naming scheme (local renaming) |
| 574 | """ |
| 575 | mapping = [] |
| 576 | for old_item in old_list: |
| 577 | new_item = old_item |
| 578 | |
| 579 | if "qkv" in new_item: |
| 580 | continue |
| 581 | |
| 582 | if "encoder_kv" in new_item: |
| 583 | continue |
| 584 | |
| 585 | new_item = new_item.replace("norm.weight", "group_norm.weight") |
| 586 | new_item = new_item.replace("norm.bias", "group_norm.bias") |
| 587 | |
| 588 | new_item = new_item.replace("proj_out.weight", "to_out.0.weight") |
| 589 | new_item = new_item.replace("proj_out.bias", "to_out.0.bias") |
| 590 | |
| 591 | new_item = new_item.replace("norm_encoder.weight", "norm_cross.weight") |
| 592 | new_item = new_item.replace("norm_encoder.bias", "norm_cross.bias") |
| 593 | |
| 594 | new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) |
| 595 | |
| 596 | mapping.append({"old": old_item, "new": new_item}) |
| 597 | |
| 598 | return mapping |
| 599 | |
| 600 | |
| 601 | def assign_attention_to_checkpoint(new_checkpoint, unet_state_dict, old_path, new_path, config): |
no test coverage detected