import argparse

import tempfile

import torch

from accelerate import load_checkpoint_and_dispatch

from transformers import CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import UnCLIPPipeline, UNet2DConditionModel, UNet2DModel

from diffusers.models.transformers.prior_transformer import PriorTransformer

from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel

from diffusers.schedulers.scheduling_unclip import UnCLIPScheduler

r"""

Example - From the diffusers root directory:

Download weights:

```sh

$ wget https://arena.kakaocdn.net/brainrepo/models/karlo-public/v1.0.0.alpha/efdf6206d8ed593961593dc029a8affa/decoder-ckpt-step%3D01000000-of-01000000.ckpt

$ wget https://arena.kakaocdn.net/brainrepo/models/karlo-public/v1.0.0.alpha/4226b831ae0279020d134281f3c31590/improved-sr-ckpt-step%3D1.2M.ckpt

$ wget https://arena.kakaocdn.net/brainrepo/models/karlo-public/v1.0.0.alpha/85626483eaca9f581e2a78d31ff905ca/prior-ckpt-step%3D01000000-of-01000000.ckpt

$ wget https://arena.kakaocdn.net/brainrepo/models/karlo-public/v1.0.0.alpha/0b62380a75e56f073e2844ab5199153d/ViT-L-14_stats.th

```

Convert the model:

```sh

$ python scripts/convert_kakao_brain_unclip_to_diffusers.py \

--decoder_checkpoint_path ./decoder-ckpt-step\=01000000-of-01000000.ckpt \

--super_res_unet_checkpoint_path ./improved-sr-ckpt-step\=1.2M.ckpt \

--prior_checkpoint_path ./prior-ckpt-step\=01000000-of-01000000.ckpt \

--clip_stat_path ./ViT-L-14_stats.th \

--dump_path <path where to save model>

```

"""

# prior

PRIOR_ORIGINAL_PREFIX = "model"

# Uses default arguments

PRIOR_CONFIG = {}

def prior_model_from_original_config():

model = PriorTransformer(**PRIOR_CONFIG)

return model

def prior_original_checkpoint_to_diffusers_checkpoint(model, checkpoint, clip_stats_checkpoint):

diffusers_checkpoint = {}

# <original>.time_embed.0 -> <diffusers>.time_embedding.linear_1

diffusers_checkpoint.update(

{

"time_embedding.linear_1.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.time_embed.0.weight"],

"time_embedding.linear_1.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.time_embed.0.bias"],

}

)

# <original>.clip_img_proj -> <diffusers>.proj_in

diffusers_checkpoint.update(

{

"proj_in.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.clip_img_proj.weight"],

"proj_in.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.clip_img_proj.bias"],

}

)

# <original>.text_emb_proj -> <diffusers>.embedding_proj

diffusers_checkpoint.update(

{

"embedding_proj.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.text_emb_proj.weight"],

"embedding_proj.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.text_emb_proj.bias"],

}

)

# <original>.text_enc_proj -> <diffusers>.encoder_hidden_states_proj

diffusers_checkpoint.update(

{

"encoder_hidden_states_proj.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.text_enc_proj.weight"],

"encoder_hidden_states_proj.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.text_enc_proj.bias"],

}

)

# <original>.positional_embedding -> <diffusers>.positional_embedding

diffusers_checkpoint.update({"positional_embedding": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.positional_embedding"]})

# <original>.prd_emb -> <diffusers>.prd_embedding

diffusers_checkpoint.update({"prd_embedding": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.prd_emb"]})

# <original>.time_embed.2 -> <diffusers>.time_embedding.linear_2

diffusers_checkpoint.update(

{

"time_embedding.linear_2.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.time_embed.2.weight"],

"time_embedding.linear_2.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.time_embed.2.bias"],

}

)

# <original>.resblocks.<x> -> <diffusers>.transformer_blocks.<x>

for idx in range(len(model.transformer_blocks)):

diffusers_transformer_prefix = f"transformer_blocks.{idx}"

original_transformer_prefix = f"{PRIOR_ORIGINAL_PREFIX}.transformer.resblocks.{idx}"

# <original>.attn -> <diffusers>.attn1

diffusers_attention_prefix = f"{diffusers_transformer_prefix}.attn1"

original_attention_prefix = f"{original_transformer_prefix}.attn"

diffusers_checkpoint.update(

prior_attention_to_diffusers(

checkpoint,

diffusers_attention_prefix=diffusers_attention_prefix,

original_attention_prefix=original_attention_prefix,

attention_head_dim=model.attention_head_dim,

)

)

# <original>.mlp -> <diffusers>.ff

diffusers_ff_prefix = f"{diffusers_transformer_prefix}.ff"

original_ff_prefix = f"{original_transformer_prefix}.mlp"

diffusers_checkpoint.update(

prior_ff_to_diffusers(

checkpoint, diffusers_ff_prefix=diffusers_ff_prefix, original_ff_prefix=original_ff_prefix

)

)

# <original>.ln_1 -> <diffusers>.norm1

diffusers_checkpoint.update(

{

f"{diffusers_transformer_prefix}.norm1.weight": checkpoint[

f"{original_transformer_prefix}.ln_1.weight"

],

f"{diffusers_transformer_prefix}.norm1.bias": checkpoint[f"{original_transformer_prefix}.ln_1.bias"],

}

)

# <original>.ln_2 -> <diffusers>.norm3

diffusers_checkpoint.update(

{

f"{diffusers_transformer_prefix}.norm3.weight": checkpoint[

f"{original_transformer_prefix}.ln_2.weight"

],

f"{diffusers_transformer_prefix}.norm3.bias": checkpoint[f"{original_transformer_prefix}.ln_2.bias"],

}

)

# <original>.final_ln -> <diffusers>.norm_out

diffusers_checkpoint.update(

{

"norm_out.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.final_ln.weight"],

"norm_out.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.final_ln.bias"],

}

)

# <original>.out_proj -> <diffusers>.proj_to_clip_embeddings

diffusers_checkpoint.update(

{

"proj_to_clip_embeddings.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.out_proj.weight"],

"proj_to_clip_embeddings.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.out_proj.bias"],

}

)

# clip stats

clip_mean, clip_std = clip_stats_checkpoint

clip_mean = clip_mean[None, :]

clip_std = clip_std[None, :]

diffusers_checkpoint.update({"clip_mean": clip_mean, "clip_std": clip_std})

return diffusers_checkpoint

def prior_attention_to_diffusers(

checkpoint, *, diffusers_attention_prefix, original_attention_prefix, attention_head_dim

):

diffusers_checkpoint = {}

# <original>.c_qkv -> <diffusers>.{to_q, to_k, to_v}

[q_weight, k_weight, v_weight], [q_bias, k_bias, v_bias] = split_attentions(

weight=checkpoint[f"{original_attention_prefix}.c_qkv.weight"],

bias=checkpoint[f"{original_attention_prefix}.c_qkv.bias"],

split=3,

chunk_size=attention_head_dim,

)

diffusers_checkpoint.update(

{

f"{diffusers_attention_prefix}.to_q.weight": q_weight,

f"{diffusers_attention_prefix}.to_q.bias": q_bias,

f"{diffusers_attention_prefix}.to_k.weight": k_weight,

f"{diffusers_attention_prefix}.to_k.bias": k_bias,

f"{diffusers_attention_prefix}.to_v.weight": v_weight,

f"{diffusers_attention_prefix}.to_v.bias": v_bias,

}

)

# <original>.c_proj -> <diffusers>.to_out.0

diffusers_checkpoint.update(

{

f"{diffusers_attention_prefix}.to_out.0.weight": checkpoint[f"{original_attention_prefix}.c_proj.weight"],

f"{diffusers_attention_prefix}.to_out.0.bias": checkpoint[f"{original_attention_prefix}.c_proj.bias"],

}

)

return diffusers_checkpoint

def prior_ff_to_diffusers(checkpoint, *, diffusers_ff_prefix, original_ff_prefix):

diffusers_checkpoint = {

# <original>.c_fc -> <diffusers>.net.0.proj

f"{diffusers_ff_prefix}.net.{0}.proj.weight": checkpoint[f"{original_ff_prefix}.c_fc.weight"],

f"{diffusers_ff_prefix}.net.{0}.proj.bias": checkpoint[f"{original_ff_prefix}.c_fc.bias"],

# <original>.c_proj -> <diffusers>.net.2

f"{diffusers_ff_prefix}.net.{2}.weight": checkpoint[f"{original_ff_prefix}.c_proj.weight"],

f"{diffusers_ff_prefix}.net.{2}.bias": checkpoint[f"{original_ff_prefix}.c_proj.bias"],

}

return diffusers_checkpoint

# done prior

# decoder

DECODER_ORIGINAL_PREFIX = "model"

# We are hardcoding the model configuration for now. If we need to generalize to more model configurations, we can

# update then.

DECODER_CONFIG = {

"sample_size": 64,

"layers_per_block": 3,

"down_block_types": (

"ResnetDownsampleBlock2D",

"SimpleCrossAttnDownBlock2D",

"SimpleCrossAttnDownBlock2D",

"SimpleCrossAttnDownBlock2D",

),

"up_block_types": (

"SimpleCrossAttnUpBlock2D",

"SimpleCrossAttnUpBlock2D",

"SimpleCrossAttnUpBlock2D",

"ResnetUpsampleBlock2D",

),

"mid_block_type": "UNetMidBlock2DSimpleCrossAttn",

"block_out_channels": (320, 640, 960, 1280),

"in_channels": 3,

"out_channels": 6,

"cross_attention_dim": 1536,

"class_embed_type": "identity",

"attention_head_dim": 64,

"resnet_time_scale_shift": "scale_shift",

}

def decoder_model_from_original_config():

model = UNet2DConditionModel(**DECODER_CONFIG)

return model

def decoder_original_checkpoint_to_diffusers_checkpoint(model, checkpoint):

diffusers_checkpoint = {}

original_unet_prefix = DECODER_ORIGINAL_PREFIX

num_head_channels = DECODER_CONFIG["attention_head_dim"]

diffusers_checkpoint.update(unet_time_embeddings(checkpoint, original_unet_prefix))

diffusers_checkpoint.update(unet_conv_in(checkpoint, original_unet_prefix))

# <original>.input_blocks -> <diffusers>.down_blocks

original_down_block_idx = 1

for diffusers_down_block_idx in range(len(model.down_blocks)):

checkpoint_update, num_original_down_blocks = unet_downblock_to_diffusers_checkpoint(

model,

checkpoint,

diffusers_down_block_idx=diffusers_down_block_idx,

original_down_block_idx=original_down_block_idx,

original_unet_prefix=original_unet_prefix,

num_head_channels=num_head_channels,

)

original_down_block_idx += num_original_down_blocks

diffusers_checkpoint.update(checkpoint_update)

# done <original>.input_blocks -> <diffusers>.down_blocks

diffusers_checkpoint.update(

unet_midblock_to_diffusers_checkpoint(

model,

checkpoint,

original_unet_prefix=original_unet_prefix,

num_head_channels=num_head_channels,

)

)

# <original>.output_blocks -> <diffusers>.up_blocks

original_up_block_idx = 0

for diffusers_up_block_idx in range(len(model.up_blocks)):

checkpoint_update, num_original_up_blocks = unet_upblock_to_diffusers_checkpoint(

model,

checkpoint,

diffusers_up_block_idx=diffusers_up_block_idx,

original_up_block_idx=original_up_block_idx,

original_unet_prefix=original_unet_prefix,

num_head_channels=num_head_channels,

)

original_up_block_idx += num_original_up_blocks

diffusers_checkpoint.update(checkpoint_update)

# done <original>.output_blocks -> <diffusers>.up_blocks

diffusers_checkpoint.update(unet_conv_norm_out(checkpoint, original_unet_prefix))

diffusers_checkpoint.update(unet_conv_out(checkpoint, original_unet_prefix))

return diffusers_checkpoint

# done decoder

# text proj

def text_proj_from_original_config():

# From the conditional unet constructor where the dimension of the projected time embeddings is

# constructed

time_embed_dim = DECODER_CONFIG["block_out_channels"][0] * 4

cross_attention_dim = DECODER_CONFIG["cross_attention_dim"]

model = UnCLIPTextProjModel(time_embed_dim=time_embed_dim, cross_attention_dim=cross_attention_dim)

return model

# Note that the input checkpoint is the original decoder checkpoint

def text_proj_original_checkpoint_to_diffusers_checkpoint(checkpoint):

diffusers_checkpoint = {

# <original>.text_seq_proj.0 -> <diffusers>.encoder_hidden_states_proj

"encoder_hidden_states_proj.weight": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_seq_proj.0.weight"],

"encoder_hidden_states_proj.bias": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_seq_proj.0.bias"],

# <original>.text_seq_proj.1 -> <diffusers>.text_encoder_hidden_states_norm

"text_encoder_hidden_states_norm.weight": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_seq_proj.1.weight"],

"text_encoder_hidden_states_norm.bias": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_seq_proj.1.bias"],

# <original>.clip_tok_proj -> <diffusers>.clip_extra_context_tokens_proj

"clip_extra_context_tokens_proj.weight": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.clip_tok_proj.weight"],

"clip_extra_context_tokens_proj.bias": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.clip_tok_proj.bias"],

# <original>.text_feat_proj -> <diffusers>.embedding_proj

"embedding_proj.weight": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_feat_proj.weight"],

"embedding_proj.bias": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_feat_proj.bias"],

# <original>.cf_param -> <diffusers>.learned_classifier_free_guidance_embeddings

"learned_classifier_free_guidance_embeddings": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.cf_param"],

# <original>.clip_emb -> <diffusers>.clip_image_embeddings_project_to_time_embeddings

"clip_image_embeddings_project_to_time_embeddings.weight": checkpoint[

f"{DECODER_ORIGINAL_PREFIX}.clip_emb.weight"

],

"clip_image_embeddings_project_to_time_embeddings.bias": checkpoint[

f"{DECODER_ORIGINAL_PREFIX}.clip_emb.bias"

],

}

return diffusers_checkpoint

# done text proj

# super res unet first steps

SUPER_RES_UNET_FIRST_STEPS_PREFIX = "model_first_steps"

SUPER_RES_UNET_FIRST_STEPS_CONFIG = {

"sample_size": 256,

"layers_per_block": 3,

"down_block_types": (

"ResnetDownsampleBlock2D",

"ResnetDownsampleBlock2D",

"ResnetDownsampleBlock2D",

"ResnetDownsampleBlock2D",

),

"up_block_types": (

"ResnetUpsampleBlock2D",

"ResnetUpsampleBlock2D",

"ResnetUpsampleBlock2D",

"ResnetUpsampleBlock2D",

),

"block_out_channels": (320, 640, 960, 1280),

"in_channels": 6,

"out_channels": 3,

"add_attention": False,

}

def super_res_unet_first_steps_model_from_original_config():

model = UNet2DModel(**SUPER_RES_UNET_FIRST_STEPS_CONFIG)

return model

def super_res_unet_first_steps_original_checkpoint_to_diffusers_checkpoint(model, checkpoint):

diffusers_checkpoint = {}

original_unet_prefix = SUPER_RES_UNET_FIRST_STEPS_PREFIX

diffusers_checkpoint.update(unet_time_embeddings(checkpoint, original_unet_prefix))

diffusers_checkpoint.update(unet_conv_in(checkpoint, original_unet_prefix))

# <original>.input_blocks -> <diffusers>.down_blocks

original_down_block_idx = 1

for diffusers_down_block_idx in range(len(model.down_blocks)):

checkpoint_update, num_original_down_blocks = unet_downblock_to_diffusers_checkpoint(

model,

checkpoint,

diffusers_down_block_idx=diffusers_down_block_idx,

original_down_block_idx=original_down_block_idx,

original_unet_prefix=original_unet_prefix,

num_head_channels=None,

)

original_down_block_idx += num_original_down_blocks

diffusers_checkpoint.update(checkpoint_update)

diffusers_checkpoint.update(

unet_midblock_to_diffusers_checkpoint(

model,

checkpoint,

original_unet_prefix=original_unet_prefix,

num_head_channels=None,

)

)

# <original>.output_blocks -> <diffusers>.up_blocks

original_up_block_idx = 0

for diffusers_up_block_idx in range(len(model.up_blocks)):

checkpoint_update, num_original_up_blocks = unet_upblock_to_diffusers_checkpoint(

model,

checkpoint,

diffusers_up_block_idx=diffusers_up_block_idx,

original_up_block_idx=original_up_block_idx,

original_unet_prefix=original_unet_prefix,

num_head_channels=None,

)

original_up_block_idx += num_original_up_blocks

diffusers_checkpoint.update(checkpoint_update)

# done <original>.output_blocks -> <diffusers>.up_blocks

diffusers_checkpoint.update(unet_conv_norm_out(checkpoint, original_unet_prefix))

diffusers_checkpoint.update(unet_conv_out(checkpoint, original_unet_prefix))

return diffusers_checkpoint

# done super res unet first steps

# super res unet last step

SUPER_RES_UNET_LAST_STEP_PREFIX = "model_last_step"

SUPER_RES_UNET_LAST_STEP_CONFIG = {

"sample_size": 256,

"layers_per_block": 3,

"down_block_types": (

"ResnetDownsampleBlock2D",

"ResnetDownsampleBlock2D",

"ResnetDownsampleBlock2D",

"ResnetDownsampleBlock2D",

),

"up_block_types": (

"ResnetUpsampleBlock2D",

"ResnetUpsampleBlock2D",

"ResnetUpsampleBlock2D",

"ResnetUpsampleBlock2D",

),

"block_out_channels": (320, 640, 960, 1280),

"in_channels": 6,

"out_channels": 3,

"add_attention": False,

}

def super_res_unet_last_step_model_from_original_config():

model = UNet2DModel(**SUPER_RES_UNET_LAST_STEP_CONFIG)

return model

def super_res_unet_last_step_original_checkpoint_to_diffusers_checkpoint(model, checkpoint):

diffusers_checkpoint = {}

original_unet_prefix = SUPER_RES_UNET_LAST_STEP_PREFIX

diffusers_checkpoint.update(unet_time_embeddings(checkpoint, original_unet_prefix))

diffusers_checkpoint.update(unet_conv_in(checkpoint, original_unet_prefix))

# <original>.input_blocks -> <diffusers>.down_blocks

original_down_block_idx = 1

for diffusers_down_block_idx in range(len(model.down_blocks)):

checkpoint_update, num_original_down_blocks = unet_downblock_to_diffusers_checkpoint(

model,

checkpoint,

diffusers_down_block_idx=diffusers_down_block_idx,

original_down_block_idx=original_down_block_idx,

original_unet_prefix=original_unet_prefix,

num_head_channels=None,

)

original_down_block_idx += num_original_down_blocks

diffusers_checkpoint.update(checkpoint_update)

diffusers_checkpoint.update(

unet_midblock_to_diffusers_checkpoint(

model,

checkpoint,

original_unet_prefix=original_unet_prefix,

num_head_channels=None,

)

)

# <original>.output_blocks -> <diffusers>.up_blocks

original_up_block_idx = 0

for diffusers_up_block_idx in range(len(model.up_blocks)):

checkpoint_update, num_original_up_blocks = unet_upblock_to_diffusers_checkpoint(

model,

checkpoint,

diffusers_up_block_idx=diffusers_up_block_idx,

original_up_block_idx=original_up_block_idx,

original_unet_prefix=original_unet_prefix,

num_head_channels=None,

)

original_up_block_idx += num_original_up_blocks

diffusers_checkpoint.update(checkpoint_update)

# done <original>.output_blocks -> <diffusers>.up_blocks

diffusers_checkpoint.update(unet_conv_norm_out(checkpoint, original_unet_prefix))

diffusers_checkpoint.update(unet_conv_out(checkpoint, original_unet_prefix))

return diffusers_checkpoint

# done super res unet last step

# unet utils

# <original>.time_embed -> <diffusers>.time_embedding

def unet_time_embeddings(checkpoint, original_unet_prefix):

diffusers_checkpoint = {}

diffusers_checkpoint.update(

{

"time_embedding.linear_1.weight": checkpoint[f"{original_unet_prefix}.time_embed.0.weight"],

"time_embedding.linear_1.bias": checkpoint[f"{original_unet_prefix}.time_embed.0.bias"],

"time_embedding.linear_2.weight": checkpoint[f"{original_unet_prefix}.time_embed.2.weight"],

"time_embedding.linear_2.bias": checkpoint[f"{original_unet_prefix}.time_embed.2.bias"],

}

)

return diffusers_checkpoint

# <original>.input_blocks.0 -> <diffusers>.conv_in

def unet_conv_in(checkpoint, original_unet_prefix):

diffusers_checkpoint = {}

diffusers_checkpoint.update(

{

"conv_in.weight": checkpoint[f"{original_unet_prefix}.input_blocks.0.0.weight"],

"conv_in.bias": checkpoint[f"{original_unet_prefix}.input_blocks.0.0.bias"],

}

)

return diffusers_checkpoint

# <original>.out.0 -> <diffusers>.conv_norm_out

def unet_conv_norm_out(checkpoint, original_unet_prefix):

diffusers_checkpoint = {}

diffusers_checkpoint.update(

{

"conv_norm_out.weight": checkpoint[f"{original_unet_prefix}.out.0.weight"],

"conv_norm_out.bias": checkpoint[f"{original_unet_prefix}.out.0.bias"],

}

)

return diffusers_checkpoint

# <original>.out.2 -> <diffusers>.conv_out

def unet_conv_out(checkpoint, original_unet_prefix):

diffusers_checkpoint = {}

diffusers_checkpoint.update(

{

"conv_out.weight": checkpoint[f"{original_unet_prefix}.out.2.weight"],

"conv_out.bias": checkpoint[f"{original_unet_prefix}.out.2.bias"],

}

)

return diffusers_checkpoint

# <original>.input_blocks -> <diffusers>.down_blocks

def unet_downblock_to_diffusers_checkpoint(

model, checkpoint, *, diffusers_down_block_idx, original_down_block_idx, original_unet_prefix, num_head_channels

):

diffusers_checkpoint = {}

diffusers_resnet_prefix = f"down_blocks.{diffusers_down_block_idx}.resnets"

original_down_block_prefix = f"{original_unet_prefix}.input_blocks"

down_block = model.down_blocks[diffusers_down_block_idx]

num_resnets = len(down_block.resnets)

if down_block.downsamplers is None:

downsampler = False

else:

assert len(down_block.downsamplers) == 1

downsampler = True

# The downsample block is also a resnet

num_resnets += 1

for resnet_idx_inc in range(num_resnets):

full_resnet_prefix = f"{original_down_block_prefix}.{original_down_block_idx + resnet_idx_inc}.0"

if downsampler and 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(

resnet_to_diffusers_checkpoint(

checkpoint, resnet_prefix=full_resnet_prefix, diffusers_resnet_prefix=full_diffusers_resnet_prefix

)

)

if hasattr(down_block, "attentions"):

num_attentions = len(down_block.attentions)

diffusers_attention_prefix = f"down_blocks.{diffusers_down_block_idx}.attentions"

for attention_idx_inc in range(num_attentions):

full_attention_prefix = f"{original_down_block_prefix}.{original_down_block_idx + attention_idx_inc}.1"

full_diffusers_attention_prefix = f"{diffusers_attention_prefix}.{attention_idx_inc}"

diffusers_checkpoint.update(

attention_to_diffusers_checkpoint(

checkpoint,

attention_prefix=full_attention_prefix,

diffusers_attention_prefix=full_diffusers_attention_prefix,

num_head_channels=num_head_channels,

)

)

num_original_down_blocks = num_resnets

return diffusers_checkpoint, num_original_down_blocks

# <original>.middle_block -> <diffusers>.mid_block

def unet_midblock_to_diffusers_checkpoint(model, checkpoint, *, original_unet_prefix, num_head_channels):

diffusers_checkpoint = {}

# block 0

original_block_idx = 0

diffusers_checkpoint.update(

resnet_to_diffusers_checkpoint(

checkpoint,

diffusers_resnet_prefix="mid_block.resnets.0",

resnet_prefix=f"{original_unet_prefix}.middle_block.{original_block_idx}",

)

)

original_block_idx += 1

# optional block 1

if hasattr(model.mid_block, "attentions") and model.mid_block.attentions[0] is not None:

diffusers_checkpoint.update(

attention_to_diffusers_checkpoint(

checkpoint,

diffusers_attention_prefix="mid_block.attentions.0",

attention_prefix=f"{original_unet_prefix}.middle_block.{original_block_idx}",

num_head_channels=num_head_channels,

)

)

original_block_idx += 1

# block 1 or block 2

diffusers_checkpoint.update(

resnet_to_diffusers_checkpoint(

checkpoint,

diffusers_resnet_prefix="mid_block.resnets.1",

resnet_prefix=f"{original_unet_prefix}.middle_block.{original_block_idx}",

)

)

return diffusers_checkpoint

# <original>.output_blocks -> <diffusers>.up_blocks

def unet_upblock_to_diffusers_checkpoint(

model, checkpoint, *, diffusers_up_block_idx, original_up_block_idx, original_unet_prefix, num_head_channels

):

diffusers_checkpoint = {}

diffusers_resnet_prefix = f"up_blocks.{diffusers_up_block_idx}.resnets"

original_up_block_prefix = f"{original_unet_prefix}.output_blocks"

up_block = model.up_blocks[diffusers_up_block_idx]

num_resnets = len(up_block.resnets)

if up_block.upsamplers is None:

upsampler = False

else:

assert len(up_block.upsamplers) == 1

upsampler = True

# The upsample block is also a resnet

num_resnets += 1

has_attentions = hasattr(up_block, "attentions")

for resnet_idx_inc in range(num_resnets):

if upsampler and resnet_idx_inc == num_resnets - 1:

# this is an upsample block

if has_attentions:

# There is a middle attention block that we skip

original_resnet_block_idx = 2

else:

original_resnet_block_idx = 1

# we add the `minus 1` because the last two resnets are stuck together in the same output block

full_resnet_prefix = (

f"{original_up_block_prefix}.{original_up_block_idx + resnet_idx_inc - 1}.{original_resnet_block_idx}"

)

full_diffusers_resnet_prefix = f"up_blocks.{diffusers_up_block_idx}.upsamplers.0"

else:

# this is a regular resnet block

full_resnet_prefix = f"{original_up_block_prefix}.{original_up_block_idx + resnet_idx_inc}.0"

full_diffusers_resnet_prefix = f"{diffusers_resnet_prefix}.{resnet_idx_inc}"

diffusers_checkpoint.update(

resnet_to_diffusers_checkpoint(

checkpoint, resnet_prefix=full_resnet_prefix, diffusers_resnet_prefix=full_diffusers_resnet_prefix

)

)

if has_attentions:

num_attentions = len(up_block.attentions)

diffusers_attention_prefix = f"up_blocks.{diffusers_up_block_idx}.attentions"

for attention_idx_inc in range(num_attentions):

full_attention_prefix = f"{original_up_block_prefix}.{original_up_block_idx + attention_idx_inc}.1"

full_diffusers_attention_prefix = f"{diffusers_attention_prefix}.{attention_idx_inc}"

diffusers_checkpoint.update(

attention_to_diffusers_checkpoint(

checkpoint,

attention_prefix=full_attention_prefix,

diffusers_attention_prefix=full_diffusers_attention_prefix,

num_head_channels=num_head_channels,

)

)

num_original_down_blocks = num_resnets - 1 if upsampler else num_resnets

return diffusers_checkpoint, num_original_down_blocks

def 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"],

}

skip_connection_prefix = f"{resnet_prefix}.skip_connection"

if f"{skip_connection_prefix}.weight" in checkpoint:

diffusers_checkpoint.update(

{

f"{diffusers_resnet_prefix}.conv_shortcut.weight": checkpoint[f"{skip_connection_prefix}.weight"],

f"{diffusers_resnet_prefix}.conv_shortcut.bias": checkpoint[f"{skip_connection_prefix}.bias"],

}

)

return diffusers_checkpoint

def attention_to_diffusers_checkpoint(checkpoint, *, diffusers_attention_prefix, attention_prefix, num_head_channels):

diffusers_checkpoint = {}

# <original>.norm -> <diffusers>.group_norm

diffusers_checkpoint.update(

{

f"{diffusers_attention_prefix}.group_norm.weight": checkpoint[f"{attention_prefix}.norm.weight"],

f"{diffusers_attention_prefix}.group_norm.bias": checkpoint[f"{attention_prefix}.norm.bias"],

}

)

# <original>.qkv -> <diffusers>.{query, key, value}

[q_weight, k_weight, v_weight], [q_bias, k_bias, v_bias] = split_attentions(

weight=checkpoint[f"{attention_prefix}.qkv.weight"][:, :, 0],

bias=checkpoint[f"{attention_prefix}.qkv.bias"],

split=3,

chunk_size=num_head_channels,

)

diffusers_checkpoint.update(

{

f"{diffusers_attention_prefix}.to_q.weight": q_weight,

f"{diffusers_attention_prefix}.to_q.bias": q_bias,

f"{diffusers_attention_prefix}.to_k.weight": k_weight,

f"{diffusers_attention_prefix}.to_k.bias": k_bias,

f"{diffusers_attention_prefix}.to_v.weight": v_weight,

f"{diffusers_attention_prefix}.to_v.bias": v_bias,

}

)

# <original>.encoder_kv -> <diffusers>.{context_key, context_value}

[encoder_k_weight, encoder_v_weight], [encoder_k_bias, encoder_v_bias] = split_attentions(

weight=checkpoint[f"{attention_prefix}.encoder_kv.weight"][:, :, 0],

bias=checkpoint[f"{attention_prefix}.encoder_kv.bias"],

split=2,

chunk_size=num_head_channels,

)

diffusers_checkpoint.update(

{

f"{diffusers_attention_prefix}.add_k_proj.weight": encoder_k_weight,

f"{diffusers_attention_prefix}.add_k_proj.bias": encoder_k_bias,

f"{diffusers_attention_prefix}.add_v_proj.weight": encoder_v_weight,

f"{diffusers_attention_prefix}.add_v_proj.bias": encoder_v_bias,

}

)

# <original>.proj_out (1d conv) -> <diffusers>.proj_attn (linear)

diffusers_checkpoint.update(

{

f"{diffusers_attention_prefix}.to_out.0.weight": checkpoint[f"{attention_prefix}.proj_out.weight"][

:, :, 0

],

f"{diffusers_attention_prefix}.to_out.0.bias": checkpoint[f"{attention_prefix}.proj_out.bias"],

}

)

return diffusers_checkpoint

# TODO maybe document and/or can do more efficiently (build indices in for loop and extract once for each split?)

def split_attentions(*, weight, bias, split, chunk_size):

weights = [None] * split

biases = [None] * split

weights_biases_idx = 0

for starting_row_index in range(0, weight.shape[0], chunk_size):

row_indices = torch.arange(starting_row_index, starting_row_index + chunk_size)

weight_rows = weight[row_indices, :]

bias_rows = bias[row_indices]

if weights[weights_biases_idx] is None:

assert weights[weights_biases_idx] is None

weights[weights_biases_idx] = weight_rows

biases[weights_biases_idx] = bias_rows

else:

assert weights[weights_biases_idx] is not None

weights[weights_biases_idx] = torch.concat([weights[weights_biases_idx], weight_rows])

biases[weights_biases_idx] = torch.concat([biases[weights_biases_idx], bias_rows])

weights_biases_idx = (weights_biases_idx + 1) % split

return weights, biases

# done unet utils

# Driver functions

def text_encoder():

print("loading CLIP text encoder")

clip_name = "openai/clip-vit-large-patch14"

# sets pad_value to 0

pad_token = "!"

tokenizer_model = CLIPTokenizer.from_pretrained(clip_name, pad_token=pad_token, device_map="auto")

assert tokenizer_model.convert_tokens_to_ids(pad_token) == 0

text_encoder_model = CLIPTextModelWithProjection.from_pretrained(

clip_name,

# `CLIPTextModel` does not support device_map="auto"

# device_map="auto"

)

print("done loading CLIP text encoder")

return text_encoder_model, tokenizer_model

def prior(*, args, checkpoint_map_location):

print("loading prior")

prior_checkpoint = torch.load(args.prior_checkpoint_path, map_location=checkpoint_map_location)

prior_checkpoint = prior_checkpoint["state_dict"]

clip_stats_checkpoint = torch.load(args.clip_stat_path, map_location=checkpoint_map_location)

prior_model = prior_model_from_original_config()

prior_diffusers_checkpoint = prior_original_checkpoint_to_diffusers_checkpoint(

prior_model, prior_checkpoint, clip_stats_checkpoint

)

del prior_checkpoint

del clip_stats_checkpoint

load_checkpoint_to_model(prior_diffusers_checkpoint, prior_model, strict=True)

print("done loading prior")

return prior_model

def decoder(*, args, checkpoint_map_location):

print("loading decoder")

decoder_checkpoint = torch.load(args.decoder_checkpoint_path, map_location=checkpoint_map_location)

decoder_checkpoint = decoder_checkpoint["state_dict"]

decoder_model = decoder_model_from_original_config()

decoder_diffusers_checkpoint = decoder_original_checkpoint_to_diffusers_checkpoint(

decoder_model, decoder_checkpoint

)

# text proj interlude

# The original decoder implementation includes a set of parameters that are used

# for creating the `encoder_hidden_states` which are what the U-net is conditioned

# on. The diffusers conditional unet directly takes the encoder_hidden_states. We pull

# the parameters into the UnCLIPTextProjModel class

text_proj_model = text_proj_from_original_config()

text_proj_checkpoint = text_proj_original_checkpoint_to_diffusers_checkpoint(decoder_checkpoint)

load_checkpoint_to_model(text_proj_checkpoint, text_proj_model, strict=True)

# done text proj interlude

del decoder_checkpoint

load_checkpoint_to_model(decoder_diffusers_checkpoint, decoder_model, strict=True)

print("done loading decoder")

return decoder_model, text_proj_model