Takes a state dict and a config, and returns a converted checkpoint.
(
checkpoint, config, path=None, extract_ema=False, controlnet=False, skip_extract_state_dict=False
)
| 248 | |
| 249 | |
| 250 | def convert_ldm_unet_checkpoint( |
| 251 | checkpoint, config, path=None, extract_ema=False, controlnet=False, skip_extract_state_dict=False |
| 252 | ): |
| 253 | """ |
| 254 | Takes a state dict and a config, and returns a converted checkpoint. |
| 255 | """ |
| 256 | |
| 257 | if skip_extract_state_dict: |
| 258 | unet_state_dict = checkpoint |
| 259 | else: |
| 260 | # extract state_dict for UNet |
| 261 | unet_state_dict = {} |
| 262 | keys = list(checkpoint.keys()) |
| 263 | |
| 264 | if controlnet: |
| 265 | unet_key = "control_model." |
| 266 | else: |
| 267 | unet_key = "model.diffusion_model." |
| 268 | |
| 269 | # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA |
| 270 | if sum(k.startswith("model_ema") for k in keys) > 100 and extract_ema: |
| 271 | logger.warning(f"Checkpoint {path} has both EMA and non-EMA weights.") |
| 272 | logger.warning( |
| 273 | "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA" |
| 274 | " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag." |
| 275 | ) |
| 276 | for key in keys: |
| 277 | if key.startswith("model.diffusion_model"): |
| 278 | flat_ema_key = "model_ema." + "".join(key.split(".")[1:]) |
| 279 | unet_state_dict[key.replace(unet_key, "")] = checkpoint[flat_ema_key] |
| 280 | else: |
| 281 | if sum(k.startswith("model_ema") for k in keys) > 100: |
| 282 | logger.warning( |
| 283 | "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA" |
| 284 | " weights (usually better for inference), please make sure to add the `--extract_ema` flag." |
| 285 | ) |
| 286 | |
| 287 | for key in keys: |
| 288 | if key.startswith(unet_key): |
| 289 | unet_state_dict[key.replace(unet_key, "")] = checkpoint[key] |
| 290 | |
| 291 | new_checkpoint = {} |
| 292 | |
| 293 | new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"] |
| 294 | new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"] |
| 295 | new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"] |
| 296 | new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"] |
| 297 | |
| 298 | if config["class_embed_type"] is None: |
| 299 | # No parameters to port |
| 300 | ... |
| 301 | elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection": |
| 302 | new_checkpoint["class_embedding.linear_1.weight"] = unet_state_dict["label_emb.0.0.weight"] |
| 303 | new_checkpoint["class_embedding.linear_1.bias"] = unet_state_dict["label_emb.0.0.bias"] |
| 304 | new_checkpoint["class_embedding.linear_2.weight"] = unet_state_dict["label_emb.0.2.weight"] |
| 305 | new_checkpoint["class_embedding.linear_2.bias"] = unet_state_dict["label_emb.0.2.bias"] |
| 306 | else: |
| 307 | raise NotImplementedError(f"Not implemented `class_embed_type`: {config['class_embed_type']}") |
no test coverage detected