convert_k_upscaler_to_diffusers.py

import argparse

import huggingface_hub
import k_diffusion as K
import torch

from diffusers import UNet2DConditionModel


UPSCALER_REPO = "pcuenq/k-upscaler"


def resnet_to_diffusers_checkpoint(resnet, checkpoint, *, diffusers_resnet_prefix, resnet_prefix):
    rv = {
        # norm1
        f"{diffusers_resnet_prefix}.norm1.linear.weight": checkpoint[f"{resnet_prefix}.main.0.mapper.weight"],
        f"{diffusers_resnet_prefix}.norm1.linear.bias": checkpoint[f"{resnet_prefix}.main.0.mapper.bias"],
        # conv1
        f"{diffusers_resnet_prefix}.conv1.weight": checkpoint[f"{resnet_prefix}.main.2.weight"],
        f"{diffusers_resnet_prefix}.conv1.bias": checkpoint[f"{resnet_prefix}.main.2.bias"],
        # norm2
        f"{diffusers_resnet_prefix}.norm2.linear.weight": checkpoint[f"{resnet_prefix}.main.4.mapper.weight"],
        f"{diffusers_resnet_prefix}.norm2.linear.bias": checkpoint[f"{resnet_prefix}.main.4.mapper.bias"],
        # conv2
        f"{diffusers_resnet_prefix}.conv2.weight": checkpoint[f"{resnet_prefix}.main.6.weight"],
        f"{diffusers_resnet_prefix}.conv2.bias": checkpoint[f"{resnet_prefix}.main.6.bias"],
    }

    if resnet.conv_shortcut is not None:
        rv.update(
            {
                f"{diffusers_resnet_prefix}.conv_shortcut.weight": checkpoint[f"{resnet_prefix}.skip.weight"],
            }
        )

    return rv


def self_attn_to_diffusers_checkpoint(checkpoint, *, diffusers_attention_prefix, attention_prefix):
    weight_q, weight_k, weight_v = checkpoint[f"{attention_prefix}.qkv_proj.weight"].chunk(3, dim=0)
    bias_q, bias_k, bias_v = checkpoint[f"{attention_prefix}.qkv_proj.bias"].chunk(3, dim=0)
    rv = {
        # norm
        f"{diffusers_attention_prefix}.norm1.linear.weight": checkpoint[f"{attention_prefix}.norm_in.mapper.weight"],
        f"{diffusers_attention_prefix}.norm1.linear.bias": checkpoint[f"{attention_prefix}.norm_in.mapper.bias"],
        # to_q
        f"{diffusers_attention_prefix}.attn1.to_q.weight": weight_q.squeeze(-1).squeeze(-1),
        f"{diffusers_attention_prefix}.attn1.to_q.bias": bias_q,
        # to_k
        f"{diffusers_attention_prefix}.attn1.to_k.weight": weight_k.squeeze(-1).squeeze(-1),
        f"{diffusers_attention_prefix}.attn1.to_k.bias": bias_k,
        # to_v
        f"{diffusers_attention_prefix}.attn1.to_v.weight": weight_v.squeeze(-1).squeeze(-1),
        f"{diffusers_attention_prefix}.attn1.to_v.bias": bias_v,
        # to_out
        f"{diffusers_attention_prefix}.attn1.to_out.0.weight": checkpoint[f"{attention_prefix}.out_proj.weight"]
        .squeeze(-1)
        .squeeze(-1),
        f"{diffusers_attention_prefix}.attn1.to_out.0.bias": checkpoint[f"{attention_prefix}.out_proj.bias"],
    }

    return rv


def cross_attn_to_diffusers_checkpoint(
    checkpoint, *, diffusers_attention_prefix, diffusers_attention_index, attention_prefix
):
    weight_k, weight_v = checkpoint[f"{attention_prefix}.kv_proj.weight"].chunk(2, dim=0)
    bias_k, bias_v = checkpoint[f"{attention_prefix}.kv_proj.bias"].chunk(2, dim=0)

    rv = {
        # norm2 (ada groupnorm)
        f"{diffusers_attention_prefix}.norm{diffusers_attention_index}.linear.weight": checkpoint[
            f"{attention_prefix}.norm_dec.mapper.weight"
        ],
        f"{diffusers_attention_prefix}.norm{diffusers_attention_index}.linear.bias": checkpoint[
            f"{attention_prefix}.norm_dec.mapper.bias"
        ],
        # layernorm on encoder_hidden_state
        f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.norm_cross.weight": checkpoint[
            f"{attention_prefix}.norm_enc.weight"
        ],
        f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.norm_cross.bias": checkpoint[
            f"{attention_prefix}.norm_enc.bias"
        ],
        # to_q
        f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.to_q.weight": checkpoint[
            f"{attention_prefix}.q_proj.weight"
        ]
        .squeeze(-1)
        .squeeze(-1),
        f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.to_q.bias": checkpoint[
            f"{attention_prefix}.q_proj.bias"
        ],
        # to_k
        f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.to_k.weight": weight_k.squeeze(-1).squeeze(-1),
        f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.to_k.bias": bias_k,
        # to_v
        f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.to_v.weight": weight_v.squeeze(-1).squeeze(-1),
        f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.to_v.bias": bias_v,
        # to_out
        f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.to_out.0.weight": checkpoint[
            f"{attention_prefix}.out_proj.weight"
        ]
        .squeeze(-1)
        .squeeze(-1),
        f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.to_out.0.bias": checkpoint[
            f"{attention_prefix}.out_proj.bias"
        ],
    }

    return rv


def block_to_diffusers_checkpoint(block, checkpoint, block_idx, block_type):
    block_prefix = "inner_model.u_net.u_blocks" if block_type == "up" else "inner_model.u_net.d_blocks"
    block_prefix = f"{block_prefix}.{block_idx}"

    diffusers_checkpoint = {}

    if not hasattr(block, "attentions"):
        n = 1  # resnet only
    elif not block.attentions[0].add_self_attention:
        n = 2  # resnet -> cross-attention
    else:
        n = 3  # resnet -> self-attention -> cross-attention)

    for resnet_idx, resnet in enumerate(block.resnets):
        # diffusers_resnet_prefix = f"{diffusers_up_block_prefix}.resnets.{resnet_idx}"
        diffusers_resnet_prefix = f"{block_type}_blocks.{block_idx}.resnets.{resnet_idx}"
        idx = n * resnet_idx if block_type == "up" else n * resnet_idx + 1
        resnet_prefix = f"{block_prefix}.{idx}" if block_type == "up" else f"{block_prefix}.{idx}"

        diffusers_checkpoint.update(
            resnet_to_diffusers_checkpoint(
                resnet, checkpoint, diffusers_resnet_prefix=diffusers_resnet_prefix, resnet_prefix=resnet_prefix
            )
        )

    if hasattr(block, "attentions"):
        for attention_idx, attention in enumerate(block.attentions):
            diffusers_attention_prefix = f"{block_type}_blocks.{block_idx}.attentions.{attention_idx}"
            idx = n * attention_idx + 1 if block_type == "up" else n * attention_idx + 2
            self_attention_prefix = f"{block_prefix}.{idx}"
            cross_attention_prefix = f"{block_prefix}.{idx }"
            cross_attention_index = 1 if not attention.add_self_attention else 2
            idx = (
                n * attention_idx + cross_attention_index
                if block_type == "up"
                else n * attention_idx + cross_attention_index + 1
            )
            cross_attention_prefix = f"{block_prefix}.{idx }"

            diffusers_checkpoint.update(
                cross_attn_to_diffusers_checkpoint(
                    checkpoint,
                    diffusers_attention_prefix=diffusers_attention_prefix,
                    diffusers_attention_index=2,
                    attention_prefix=cross_attention_prefix,
                )
            )

            if attention.add_self_attention is True:
                diffusers_checkpoint.update(
                    self_attn_to_diffusers_checkpoint(
                        checkpoint,
                        diffusers_attention_prefix=diffusers_attention_prefix,
                        attention_prefix=self_attention_prefix,
                    )
                )

    return diffusers_checkpoint


def unet_to_diffusers_checkpoint(model, checkpoint):
    diffusers_checkpoint = {}

    # pre-processing
    diffusers_checkpoint.update(
        {
            "conv_in.weight": checkpoint["inner_model.proj_in.weight"],
            "conv_in.bias": checkpoint["inner_model.proj_in.bias"],
        }
    )

    # timestep and class embedding
    diffusers_checkpoint.update(
        {
            "time_proj.weight": checkpoint["inner_model.timestep_embed.weight"].squeeze(-1),
            "time_embedding.linear_1.weight": checkpoint["inner_model.mapping.0.weight"],
            "time_embedding.linear_1.bias": checkpoint["inner_model.mapping.0.bias"],
            "time_embedding.linear_2.weight": checkpoint["inner_model.mapping.2.weight"],
            "time_embedding.linear_2.bias": checkpoint["inner_model.mapping.2.bias"],
            "time_embedding.cond_proj.weight": checkpoint["inner_model.mapping_cond.weight"],
        }
    )

    # down_blocks
    for down_block_idx, down_block in enumerate(model.down_blocks):
        diffusers_checkpoint.update(block_to_diffusers_checkpoint(down_block, checkpoint, down_block_idx, "down"))

    # up_blocks
    for up_block_idx, up_block in enumerate(model.up_blocks):
        diffusers_checkpoint.update(block_to_diffusers_checkpoint(up_block, checkpoint, up_block_idx, "up"))

    # post-processing
    diffusers_checkpoint.update(
        {
            "conv_out.weight": checkpoint["inner_model.proj_out.weight"],
            "conv_out.bias": checkpoint["inner_model.proj_out.bias"],
        }
    )

    return diffusers_checkpoint


def unet_model_from_original_config(original_config):
    in_channels = original_config["input_channels"] + original_config["unet_cond_dim"]
    out_channels = original_config["input_channels"] + (1 if original_config["has_variance"] else 0)

    block_out_channels = original_config["channels"]

    assert (
        len(set(original_config["depths"])) == 1
    ), "UNet2DConditionModel currently do not support blocks with different number of layers"
    layers_per_block = original_config["depths"][0]

    class_labels_dim = original_config["mapping_cond_dim"]
    cross_attention_dim = original_config["cross_cond_dim"]

    attn1_types = []
    attn2_types = []
    for s, c in zip(original_config["self_attn_depths"], original_config["cross_attn_depths"]):
        if s:
            a1 = "self"
            a2 = "cross" if c else None
        elif c:
            a1 = "cross"
            a2 = None
        else:
            a1 = None
            a2 = None
        attn1_types.append(a1)
        attn2_types.append(a2)

    unet = UNet2DConditionModel(
        in_channels=in_channels,
        out_channels=out_channels,
        down_block_types=("KDownBlock2D", "KCrossAttnDownBlock2D", "KCrossAttnDownBlock2D", "KCrossAttnDownBlock2D"),
        mid_block_type=None,
        up_block_types=("KCrossAttnUpBlock2D", "KCrossAttnUpBlock2D", "KCrossAttnUpBlock2D", "KUpBlock2D"),
        block_out_channels=block_out_channels,
        layers_per_block=layers_per_block,
        act_fn="gelu",
        norm_num_groups=None,
        cross_attention_dim=cross_attention_dim,
        attention_head_dim=64,
        time_cond_proj_dim=class_labels_dim,
        resnet_time_scale_shift="scale_shift",
        time_embedding_type="fourier",
        timestep_post_act="gelu",
        conv_in_kernel=1,
        conv_out_kernel=1,
    )

    return unet


def main(args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    orig_config_path = huggingface_hub.hf_hub_download(UPSCALER_REPO, "config_laion_text_cond_latent_upscaler_2.json")
    orig_weights_path = huggingface_hub.hf_hub_download(
        UPSCALER_REPO, "laion_text_cond_latent_upscaler_2_1_00470000_slim.pth"
    )
    print(f"loading original model configuration from {orig_config_path}")
    print(f"loading original model checkpoint from {orig_weights_path}")

    print("converting to diffusers unet")
    orig_config = K.config.load_config(open(orig_config_path))["model"]
    model = unet_model_from_original_config(orig_config)

    orig_checkpoint = torch.load(orig_weights_path, map_location=device)["model_ema"]
    converted_checkpoint = unet_to_diffusers_checkpoint(model, orig_checkpoint)

    model.load_state_dict(converted_checkpoint, strict=True)
    model.save_pretrained(args.dump_path)
    print(f"saving converted unet model in {args.dump_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.")
    args = parser.parse_args()

    main(args)