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
)
| 140 | |
| 141 | |
| 142 | def assign_to_checkpoint( |
| 143 | paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None |
| 144 | ): |
| 145 | """ |
| 146 | This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits |
| 147 | attention layers, and takes into account additional replacements that may arise. |
| 148 | |
| 149 | Assigns the weights to the new checkpoint. |
| 150 | """ |
| 151 | assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." |
| 152 | |
| 153 | # Splits the attention layers into three variables. |
| 154 | if attention_paths_to_split is not None: |
| 155 | for path, path_map in attention_paths_to_split.items(): |
| 156 | old_tensor = old_checkpoint[path] |
| 157 | channels = old_tensor.shape[0] // 3 |
| 158 | |
| 159 | target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) |
| 160 | |
| 161 | num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3 |
| 162 | |
| 163 | old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) |
| 164 | query, key, value = old_tensor.split(channels // num_heads, dim=1) |
| 165 | |
| 166 | checkpoint[path_map["query"]] = query.reshape(target_shape) |
| 167 | checkpoint[path_map["key"]] = key.reshape(target_shape) |
| 168 | checkpoint[path_map["value"]] = value.reshape(target_shape) |
| 169 | |
| 170 | for path in paths: |
| 171 | new_path = path["new"] |
| 172 | |
| 173 | # These have already been assigned |
| 174 | if attention_paths_to_split is not None and new_path in attention_paths_to_split: |
| 175 | continue |
| 176 | |
| 177 | # Global renaming happens here |
| 178 | new_path = new_path.replace("middle_block.0", "mid_block.resnets.0") |
| 179 | new_path = new_path.replace("middle_block.1", "mid_block.attentions.0") |
| 180 | new_path = new_path.replace("middle_block.2", "mid_block.resnets.1") |
| 181 | |
| 182 | if additional_replacements is not None: |
| 183 | for replacement in additional_replacements: |
| 184 | new_path = new_path.replace(replacement["old"], replacement["new"]) |
| 185 | |
| 186 | # proj_attn.weight has to be converted from conv 1D to linear |
| 187 | is_attn_weight = "proj_attn.weight" in new_path or ("attentions" in new_path and "to_" in new_path) |
| 188 | shape = old_checkpoint[path["old"]].shape |
| 189 | if is_attn_weight and len(shape) == 3: |
| 190 | checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0] |
| 191 | elif is_attn_weight and len(shape) == 4: |
| 192 | checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0, 0] |
| 193 | else: |
| 194 | checkpoint[new_path] = old_checkpoint[path["old"]] |
| 195 | |
| 196 | |
| 197 | def shave_segments(path, n_shave_prefix_segments=1): |
no outgoing calls
no test coverage detected