(ckpt_path: str)
| 159 | |
| 160 | # Function to convert a transformer checkpoint to the CogVideoX format |
| 161 | def convert_transformer(ckpt_path: str): |
| 162 | PREFIX_KEY = "model.diffusion_model." |
| 163 | |
| 164 | original_state_dict = get_state_dict(torch.load(ckpt_path, map_location="cpu", mmap=True)) |
| 165 | transformer = CogVideoXTransformer3DModel() |
| 166 | |
| 167 | for key in list(original_state_dict.keys()): |
| 168 | new_key = key[len(PREFIX_KEY) :] |
| 169 | for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items(): |
| 170 | new_key = new_key.replace(replace_key, rename_key) |
| 171 | update_state_dict_inplace(original_state_dict, key, new_key) |
| 172 | |
| 173 | for key in list(original_state_dict.keys()): |
| 174 | for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items(): |
| 175 | if special_key not in key: |
| 176 | continue |
| 177 | handler_fn_inplace(key, original_state_dict) |
| 178 | |
| 179 | transformer.load_state_dict(original_state_dict, strict=True) |
| 180 | return transformer |
| 181 | |
| 182 | |
| 183 | # Function to convert a VAE checkpoint to the CogVideoX format |
no test coverage detected