[wip] kakaobrain unCLIP convert script

Author: williamberman

scripts/convert_kakao_brain_unclip_to_diffusers.py

@@ -0,0 +1,213 @@
+import argparse
+import tempfile
+
+import torch
+
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
+from diffusers import UnCLIPPipeline, UNet2DConditionModel
+
+
+# decoder model
+
+
+def decoder_model_from_original_config():
+    # We are hardcoding the model configuration for now. If we need to generalize to more model configurations, we can
+    # update then.
+    model = UNet2DConditionModel(
+        down_block_types=("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D"),
+        layers_per_block=3,
+        resnet_time_scale_shift="scale_shift",
+        block_out_channels=(320, 640, 960, 1280),
+        downsample_resnet=True,
+        up_block_types=("CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D"),
+        upsample_resnet=True,
+        up_block_layers_per_block=3,
+        in_channels=3,
+        out_channels=6,
+    )
+
+    return model
+
+
+# done decoder model
+
+# decoder checkpoint
+
+DECODER_ORIGINAL_PREFIX = "model"
+
+
+def decoder_original_checkpoint_to_diffusers_checkpoint(model, checkpoint):
+    diffusers_checkpoint = {}
+
+    # TODO
+    # Linear
+    checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_seq_proj.0.weight"]
+    checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_seq_proj.0.bias"]
+    # Norm
+    checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_seq_proj.1.weight"]
+    checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_seq_proj.1.bias"]
+
+    # TODO
+    # Linear
+    checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_feat_proj.weight"]
+    checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_feat_proj.bias"]
+
+    # TODO
+    # Linear
+    checkpoint[f"{DECODER_ORIGINAL_PREFIX}.clip_tok_proj.weight"]
+    checkpoint[f"{DECODER_ORIGINAL_PREFIX}.clip_tok_proj.bias"]
+
+    # TODO There's also a `clip_emb_mult` that's a scalar and not a model parameter
+
+    # input_blocks.0 -> conv_in
+
+    diffusers_checkpoint.update(
+        {
+            "conv_in.weight": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.input_blocks.0.0.weight"],
+            "conv_in.bias": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.input_blocks.0.0.bias"],
+        }
+    )
+
+    # DownBlock2D
+    # input_blocks.[1, 2, 3, 4] -> down_blocks.0
+
+    # CrossAttnDownBlock2D
+    # input_blocks.[5, 6, 7, 8] -> down_blocks.1
+
+    # CrossAttnDownBlock2D
+    # 9, 10, 11, 12 -> 2
+
+    # TODO bug here
+    # 13, 14, 15 -> 3
+
+    resnets_per_down_block = len(model.down_blocks[0].resnets)
+    resnets_per_down_block += 1
+
+    for down_block_idx in range(len(model.down_blocks)):
+        original_resnet_idx = 1 + resnets_per_down_block * down_block_idx
+
+        diffusers_checkpoint.update(
+            decoder_downblock_to_diffusers_checkpoint(
+                model, checkpoint, diffusers_down_block_idx=down_block_idx, original_resnet_idx=original_resnet_idx
+            )
+        )
+
+    # middle_block -> mid_block
+
+    # output_blocks -> up_blocks
+
+    return diffusers_checkpoint
+
+
+# TODO add transformers
+def decoder_downblock_to_diffusers_checkpoint(model, checkpoint, *, diffusers_down_block_idx, original_resnet_idx):
+    diffusers_checkpoint = {}
+
+    diffusers_resnet_prefix = f"down_blocks.{diffusers_down_block_idx}.resnets"
+    resnet_prefix = f"{DECODER_ORIGINAL_PREFIX}.input_blocks"
+
+    num_resnets = len(model.down_blocks[diffusers_down_block_idx].resnets)
+
+    # The last downsample block is also a resnet
+    num_resnets = num_resnets + 1
+
+    for resnet_idx_inc in range(num_resnets):
+        full_resnet_prefix = f"{resnet_prefix}.{original_resnet_idx + resnet_idx_inc}.0"
+
+        if resnet_idx_inc == num_resnets - 1:
+            # this is a downsample block
+            full_diffusers_resnet_prefix = f"down_blocks.{diffusers_down_block_idx}.downsamplers.0"
+        else:
+            # this is a regular resnet block
+            full_diffusers_resnet_prefix = f"{diffusers_resnet_prefix}.{resnet_idx_inc}"
+
+        diffusers_checkpoint.update(
+            decoder_resnet_to_diffusers_checkpoint(
+                checkpoint, resnet_prefix=full_resnet_prefix, diffusers_resnet_prefix=full_diffusers_resnet_prefix
+            )
+        )
+
+    return diffusers_checkpoint
+
+
+def decoder_resnet_to_diffusers_checkpoint(checkpoint, *, diffusers_resnet_prefix, resnet_prefix):
+    diffusers_checkpoint = {
+        f"{diffusers_resnet_prefix}.norm1.weight": checkpoint[f"{resnet_prefix}.in_layers.0.weight"],
+        f"{diffusers_resnet_prefix}.norm1.bias": checkpoint[f"{resnet_prefix}.in_layers.0.bias"],
+        f"{diffusers_resnet_prefix}.conv1.weight": checkpoint[f"{resnet_prefix}.in_layers.2.weight"],
+        f"{diffusers_resnet_prefix}.conv1.bias": checkpoint[f"{resnet_prefix}.in_layers.2.bias"],
+        f"{diffusers_resnet_prefix}.time_emb_proj.weight": checkpoint[f"{resnet_prefix}.emb_layers.1.weight"],
+        f"{diffusers_resnet_prefix}.time_emb_proj.bias": checkpoint[f"{resnet_prefix}.emb_layers.1.bias"],
+        f"{diffusers_resnet_prefix}.norm2.weight": checkpoint[f"{resnet_prefix}.out_layers.0.weight"],
+        f"{diffusers_resnet_prefix}.norm2.bias": checkpoint[f"{resnet_prefix}.out_layers.0.bias"],
+        f"{diffusers_resnet_prefix}.conv2.weight": checkpoint[f"{resnet_prefix}.out_layers.3.weight"],
+        f"{diffusers_resnet_prefix}.conv2.bias": checkpoint[f"{resnet_prefix}.out_layers.3.bias"],
+    }
+
+    return diffusers_checkpoint
+
+
+# done decoder checkpoint
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.")
+
+    parser.add_argument(
+        "--decoder_checkpoint_path",
+        default=None,
+        type=str,
+        required=True,
+        help="Path to the decoder checkpoint to convert.",
+    )
+
+    parser.add_argument(
+        "--checkpoint_load_device",
+        default="cpu",
+        type=str,
+        required=False,
+        help="The device passed to `map_location` when loading checkpoints.",
+    )
+
+    args = parser.parse_args()
+
+    print(f"loading checkpoints to {args.checkpoint_load_device}")
+
+    checkpoint_map_location = torch.device(args.checkpoint_load_device)
+
+    # decoder_model
+
+    print("loading decoder")
+
+    decoder_checkpoint = torch.load(args.decoder_checkpoint_path, map_location=checkpoint_map_location)
+    decoder_checkpoint = decoder_checkpoint["state_dict"]
+
+    with init_empty_weights():
+        decoder_model = decoder_model_from_original_config()
+
+    decoder_diffusers_checkpoint = decoder_original_checkpoint_to_diffusers_checkpoint(
+        decoder_model, decoder_checkpoint
+    )
+
+    with tempfile.NamedTemporaryFile() as decoder_diffusers_checkpoint_file:
+        torch.save(decoder_diffusers_checkpoint, decoder_diffusers_checkpoint_file.name)
+        del decoder_diffusers_checkpoint
+        del decoder_checkpoint
+        # TODO use load_checkpoint_and_dispatch
+        # load_checkpoint_and_dispatch(decoder_model, decoder_diffusers_checkpoint_file.name, device_map="auto")
+        decoder_model.load_state_dict(
+            torch.load(decoder_diffusers_checkpoint_file.name, map_location=checkpoint_map_location), strict=False
+        )
+
+    print("done loading decoder")
+
+    # done decoder_model
+
+    print(f"saving Kakao Brain unCLIP to {args.dump_path}")
+
+    pipe = UnCLIPPipeline(decoder=decoder_model)
+    pipe.save_pretrained(args.dump_path)
+
+    print("done writing Kakao Brain unCLIP")