| 18 | |
| 19 | |
| 20 | class AttnProcsLayers(torch.nn.Module): |
| 21 | def __init__(self, state_dict: Dict[str, torch.Tensor]): |
| 22 | super().__init__() |
| 23 | self.layers = torch.nn.ModuleList(state_dict.values()) |
| 24 | self.mapping = dict(enumerate(state_dict.keys())) |
| 25 | self.rev_mapping = {v: k for k, v in enumerate(state_dict.keys())} |
| 26 | |
| 27 | # .processor for unet, .self_attn for text encoder |
| 28 | self.split_keys = [".processor", ".self_attn"] |
| 29 | |
| 30 | # we add a hook to state_dict() and load_state_dict() so that the |
| 31 | # naming fits with `unet.attn_processors` |
| 32 | def map_to(module, state_dict, *args, **kwargs): |
| 33 | new_state_dict = {} |
| 34 | for key, value in state_dict.items(): |
| 35 | num = int(key.split(".")[1]) # 0 is always "layers" |
| 36 | new_key = key.replace(f"layers.{num}", module.mapping[num]) |
| 37 | new_state_dict[new_key] = value |
| 38 | |
| 39 | return new_state_dict |
| 40 | |
| 41 | def remap_key(key, state_dict): |
| 42 | for k in self.split_keys: |
| 43 | if k in key: |
| 44 | return key.split(k)[0] + k |
| 45 | |
| 46 | raise ValueError( |
| 47 | f"There seems to be a problem with the state_dict: {set(state_dict.keys())}. {key} has to have one of {self.split_keys}." |
| 48 | ) |
| 49 | |
| 50 | def map_from(module, state_dict, *args, **kwargs): |
| 51 | all_keys = list(state_dict.keys()) |
| 52 | for key in all_keys: |
| 53 | replace_key = remap_key(key, state_dict) |
| 54 | new_key = key.replace(replace_key, f"layers.{module.rev_mapping[replace_key]}") |
| 55 | state_dict[new_key] = state_dict[key] |
| 56 | del state_dict[key] |
| 57 | |
| 58 | self._register_state_dict_hook(map_to) |
| 59 | self._register_load_state_dict_pre_hook(map_from, with_module=True) |
no outgoing calls
no test coverage detected