(module, state_dict, *args, **kwargs)
| 238 | # we add a hook to state_dict() and load_state_dict() so that the |
| 239 | # naming fits with `unet.attn_processors` |
| 240 | def map_to(module, state_dict, *args, **kwargs): |
| 241 | new_state_dict = {} |
| 242 | for key, value in state_dict.items(): |
| 243 | num = int(key.split(".")[1]) # 0 is always "layers" |
| 244 | new_key = key.replace(f"layers.{num}", module.mapping[num]) |
| 245 | new_state_dict[new_key] = value |
| 246 | |
| 247 | return new_state_dict |
| 248 | |
| 249 | def remap_key(key, state_dict): |
| 250 | for k in self.split_keys: |
nothing calls this directly
no outgoing calls
no test coverage detected