|
|
|
|
|
import torch |
|
from diffusers import AutoencoderKL |
|
|
|
|
|
def shave_segments(path, n_shave_prefix_segments=1): |
|
""" |
|
Removes segments. Positive values shave the first segments, negative shave the last segments. |
|
""" |
|
if n_shave_prefix_segments >= 0: |
|
return ".".join(path.split(".")[n_shave_prefix_segments:]) |
|
else: |
|
return ".".join(path.split(".")[:n_shave_prefix_segments]) |
|
|
|
|
|
def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0): |
|
""" |
|
Updates paths inside resnets to the new naming scheme (local renaming) |
|
""" |
|
mapping = [] |
|
for old_item in old_list: |
|
new_item = old_item |
|
|
|
new_item = new_item.replace("nin_shortcut", "conv_shortcut") |
|
new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) |
|
|
|
mapping.append({"old": old_item, "new": new_item}) |
|
|
|
return mapping |
|
|
|
|
|
def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0): |
|
""" |
|
Updates paths inside attentions to the new naming scheme (local renaming) |
|
""" |
|
mapping = [] |
|
for old_item in old_list: |
|
new_item = old_item |
|
|
|
new_item = new_item.replace("norm.weight", "group_norm.weight") |
|
new_item = new_item.replace("norm.bias", "group_norm.bias") |
|
|
|
new_item = new_item.replace("q.weight", "query.weight") |
|
new_item = new_item.replace("q.bias", "query.bias") |
|
|
|
new_item = new_item.replace("k.weight", "key.weight") |
|
new_item = new_item.replace("k.bias", "key.bias") |
|
|
|
new_item = new_item.replace("v.weight", "value.weight") |
|
new_item = new_item.replace("v.bias", "value.bias") |
|
|
|
new_item = new_item.replace("proj_out.weight", "proj_attn.weight") |
|
new_item = new_item.replace("proj_out.bias", "proj_attn.bias") |
|
|
|
new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) |
|
|
|
mapping.append({"old": old_item, "new": new_item}) |
|
|
|
return mapping |
|
|
|
|
|
def assign_to_checkpoint( |
|
paths, |
|
checkpoint, |
|
old_checkpoint, |
|
attention_paths_to_split=None, |
|
additional_replacements=None, |
|
config=None, |
|
): |
|
""" |
|
This does the final conversion step: take locally converted weights and apply a global renaming |
|
to them. It splits attention layers, and takes into account additional replacements |
|
that may arise. |
|
|
|
Assigns the weights to the new checkpoint. |
|
""" |
|
assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." |
|
|
|
|
|
if attention_paths_to_split is not None: |
|
for path, path_map in attention_paths_to_split.items(): |
|
old_tensor = old_checkpoint[path] |
|
channels = old_tensor.shape[0] // 3 |
|
|
|
target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) |
|
|
|
num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3 |
|
|
|
old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) |
|
query, key, value = old_tensor.split(channels // num_heads, dim=1) |
|
|
|
checkpoint[path_map["query"]] = query.reshape(target_shape) |
|
checkpoint[path_map["key"]] = key.reshape(target_shape) |
|
checkpoint[path_map["value"]] = value.reshape(target_shape) |
|
|
|
for path in paths: |
|
new_path = path["new"] |
|
|
|
|
|
if attention_paths_to_split is not None and new_path in attention_paths_to_split: |
|
continue |
|
|
|
|
|
new_path = new_path.replace("middle_block.0", "mid_block.resnets.0") |
|
new_path = new_path.replace("middle_block.1", "mid_block.attentions.0") |
|
new_path = new_path.replace("middle_block.2", "mid_block.resnets.1") |
|
|
|
if additional_replacements is not None: |
|
for replacement in additional_replacements: |
|
new_path = new_path.replace(replacement["old"], replacement["new"]) |
|
|
|
|
|
if "proj_attn.weight" in new_path: |
|
checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0] |
|
else: |
|
checkpoint[new_path] = old_checkpoint[path["old"]] |
|
|
|
|
|
def conv_attn_to_linear(checkpoint): |
|
keys = list(checkpoint.keys()) |
|
attn_keys = ["query.weight", "key.weight", "value.weight"] |
|
for key in keys: |
|
if ".".join(key.split(".")[-2:]) in attn_keys: |
|
if checkpoint[key].ndim > 2: |
|
checkpoint[key] = checkpoint[key][:, :, 0, 0] |
|
elif "proj_attn.weight" in key: |
|
if checkpoint[key].ndim > 2: |
|
checkpoint[key] = checkpoint[key][:, :, 0] |
|
|
|
|
|
def create_vae_diffusers_config(original_config): |
|
""" |
|
Creates a config for the diffusers based on the config of the LDM model. |
|
""" |
|
vae_params = original_config.model.params.ddconfig |
|
_ = original_config.model.params.embed_dim |
|
|
|
block_out_channels = [vae_params.ch * mult for mult in vae_params.ch_mult] |
|
down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels) |
|
up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels) |
|
|
|
config = dict( |
|
sample_size=tuple(vae_params.resolution), |
|
in_channels=vae_params.in_channels, |
|
out_channels=vae_params.out_ch, |
|
down_block_types=tuple(down_block_types), |
|
up_block_types=tuple(up_block_types), |
|
block_out_channels=tuple(block_out_channels), |
|
latent_channels=vae_params.z_channels, |
|
layers_per_block=vae_params.num_res_blocks, |
|
) |
|
return config |
|
|
|
|
|
def convert_ldm_vae_checkpoint(checkpoint, config): |
|
|
|
vae_state_dict = checkpoint |
|
|
|
new_checkpoint = {} |
|
|
|
new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"] |
|
new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"] |
|
new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"] |
|
new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"] |
|
new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"] |
|
new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"] |
|
|
|
new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"] |
|
new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"] |
|
new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"] |
|
new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"] |
|
new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"] |
|
new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"] |
|
|
|
new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"] |
|
new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"] |
|
new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"] |
|
new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"] |
|
|
|
|
|
num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer}) |
|
down_blocks = { |
|
layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) |
|
} |
|
|
|
|
|
num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer}) |
|
up_blocks = { |
|
layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) |
|
} |
|
|
|
for i in range(num_down_blocks): |
|
resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] |
|
|
|
if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: |
|
new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop( |
|
f"encoder.down.{i}.downsample.conv.weight" |
|
) |
|
new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop( |
|
f"encoder.down.{i}.downsample.conv.bias" |
|
) |
|
|
|
paths = renew_vae_resnet_paths(resnets) |
|
meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"} |
|
assign_to_checkpoint( |
|
paths, |
|
new_checkpoint, |
|
vae_state_dict, |
|
additional_replacements=[meta_path], |
|
config=config, |
|
) |
|
|
|
mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] |
|
num_mid_res_blocks = 2 |
|
for i in range(1, num_mid_res_blocks + 1): |
|
resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] |
|
|
|
paths = renew_vae_resnet_paths(resnets) |
|
meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} |
|
assign_to_checkpoint( |
|
paths, |
|
new_checkpoint, |
|
vae_state_dict, |
|
additional_replacements=[meta_path], |
|
config=config, |
|
) |
|
|
|
mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] |
|
paths = renew_vae_attention_paths(mid_attentions) |
|
meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} |
|
assign_to_checkpoint( |
|
paths, |
|
new_checkpoint, |
|
vae_state_dict, |
|
additional_replacements=[meta_path], |
|
config=config, |
|
) |
|
conv_attn_to_linear(new_checkpoint) |
|
|
|
for i in range(num_up_blocks): |
|
block_id = num_up_blocks - 1 - i |
|
resnets = [ |
|
key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key |
|
] |
|
|
|
if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: |
|
new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ |
|
f"decoder.up.{block_id}.upsample.conv.weight" |
|
] |
|
new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ |
|
f"decoder.up.{block_id}.upsample.conv.bias" |
|
] |
|
|
|
paths = renew_vae_resnet_paths(resnets) |
|
meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"} |
|
assign_to_checkpoint( |
|
paths, |
|
new_checkpoint, |
|
vae_state_dict, |
|
additional_replacements=[meta_path], |
|
config=config, |
|
) |
|
|
|
mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] |
|
num_mid_res_blocks = 2 |
|
for i in range(1, num_mid_res_blocks + 1): |
|
resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] |
|
|
|
paths = renew_vae_resnet_paths(resnets) |
|
meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} |
|
assign_to_checkpoint( |
|
paths, |
|
new_checkpoint, |
|
vae_state_dict, |
|
additional_replacements=[meta_path], |
|
config=config, |
|
) |
|
|
|
mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] |
|
paths = renew_vae_attention_paths(mid_attentions) |
|
meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} |
|
assign_to_checkpoint( |
|
paths, |
|
new_checkpoint, |
|
vae_state_dict, |
|
additional_replacements=[meta_path], |
|
config=config, |
|
) |
|
conv_attn_to_linear(new_checkpoint) |
|
return new_checkpoint |
|
|
|
|
|
def convert_ldm_to_hf_vae(ldm_checkpoint, ldm_config, hf_checkpoint, sample_size): |
|
checkpoint = torch.load(ldm_checkpoint)["state_dict"] |
|
|
|
|
|
vae_config = create_vae_diffusers_config(ldm_config) |
|
converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config) |
|
|
|
vae = AutoencoderKL(**vae_config) |
|
vae.load_state_dict(converted_vae_checkpoint) |
|
vae.save_pretrained(hf_checkpoint) |
|
|