This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits attention layers, and takes into account additional replacements that may arise. Assigns the weights to the new checkpoint.
(
paths,
checkpoint,
old_checkpoint,
attention_paths_to_split=None,
additional_replacements=None,
config=None,
mid_block_suffix="",
)
| 115 | |
| 116 | |
| 117 | def assign_to_checkpoint( |
| 118 | paths, |
| 119 | checkpoint, |
| 120 | old_checkpoint, |
| 121 | attention_paths_to_split=None, |
| 122 | additional_replacements=None, |
| 123 | config=None, |
| 124 | mid_block_suffix="", |
| 125 | ): |
| 126 | """ |
| 127 | This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits |
| 128 | attention layers, and takes into account additional replacements that may arise. |
| 129 | |
| 130 | Assigns the weights to the new checkpoint. |
| 131 | """ |
| 132 | assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." |
| 133 | |
| 134 | # Splits the attention layers into three variables. |
| 135 | if attention_paths_to_split is not None: |
| 136 | for path, path_map in attention_paths_to_split.items(): |
| 137 | old_tensor = old_checkpoint[path] |
| 138 | channels = old_tensor.shape[0] // 3 |
| 139 | |
| 140 | target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) |
| 141 | |
| 142 | num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3 |
| 143 | |
| 144 | old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) |
| 145 | query, key, value = old_tensor.split(channels // num_heads, dim=1) |
| 146 | |
| 147 | checkpoint[path_map["query"]] = query.reshape(target_shape) |
| 148 | checkpoint[path_map["key"]] = key.reshape(target_shape) |
| 149 | checkpoint[path_map["value"]] = value.reshape(target_shape) |
| 150 | |
| 151 | if mid_block_suffix is not None: |
| 152 | mid_block_suffix = f".{mid_block_suffix}" |
| 153 | else: |
| 154 | mid_block_suffix = "" |
| 155 | |
| 156 | for path in paths: |
| 157 | new_path = path["new"] |
| 158 | |
| 159 | # These have already been assigned |
| 160 | if attention_paths_to_split is not None and new_path in attention_paths_to_split: |
| 161 | continue |
| 162 | |
| 163 | # Global renaming happens here |
| 164 | new_path = new_path.replace("middle_block.0", f"mid_block.resnets.0{mid_block_suffix}") |
| 165 | new_path = new_path.replace("middle_block.1", "mid_block.attentions.0") |
| 166 | new_path = new_path.replace("middle_block.2", f"mid_block.resnets.1{mid_block_suffix}") |
| 167 | |
| 168 | if additional_replacements is not None: |
| 169 | for replacement in additional_replacements: |
| 170 | new_path = new_path.replace(replacement["old"], replacement["new"]) |
| 171 | |
| 172 | if new_path == "mid_block.resnets.0.spatial_res_block.norm1.weight": |
| 173 | print("yeyy") |
| 174 |
no outgoing calls
no test coverage detected