Function convert_ldm_unet_checkpoint

scripts/convert_i2vgen_to_diffusers.py:135–445

Takes a state dict and a config, and returns a converted checkpoint.

(checkpoint, config, path=None, extract_ema=False)

def convert_ldm_unet_checkpoint(checkpoint, config, path=None, extract_ema=False):
    """
    Takes a state dict and a config, and returns a converted checkpoint.
    """

    # extract state_dict for UNet
    unet_state_dict = {}
    keys = list(checkpoint.keys())

    unet_key = "model.diffusion_model."

    # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA
    if sum(k.startswith("model_ema") for k in keys) > 100 and extract_ema:
        print(f"Checkpoint {path} has both EMA and non-EMA weights.")
        print(
            "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA"
            " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag."
        )
        for key in keys:
            if key.startswith("model.diffusion_model"):
                flat_ema_key = "model_ema." + "".join(key.split(".")[1:])
                unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(flat_ema_key)
    else:
        if sum(k.startswith("model_ema") for k in keys) > 100:
            print(
                "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA"
                " weights (usually better for inference), please make sure to add the `--extract_ema` flag."
            )

        for key in keys:
            unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key)

    new_checkpoint = {}

    new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"]
    new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"]
    new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"]
    new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"]

    additional_embedding_substrings = [
        "local_image_concat",
        "context_embedding",
        "local_image_embedding",
        "fps_embedding",
    ]
    for k in unet_state_dict:
        if any(substring in k for substring in additional_embedding_substrings):
            diffusers_key = k.replace("local_image_concat", "image_latents_proj_in").replace(
                "local_image_embedding", "image_latents_context_embedding"
            )
            new_checkpoint[diffusers_key] = unet_state_dict[k]

    # temporal encoder.
    new_checkpoint["image_latents_temporal_encoder.norm1.weight"] = unet_state_dict[
        "local_temporal_encoder.layers.0.0.norm.weight"
    ]
    new_checkpoint["image_latents_temporal_encoder.norm1.bias"] = unet_state_dict[
        "local_temporal_encoder.layers.0.0.norm.bias"