|
import argparse |
|
import time |
|
from pathlib import Path |
|
from typing import Any, Dict, Literal |
|
|
|
import torch |
|
|
|
from diffusers import AsymmetricAutoencoderKL |
|
|
|
|
|
ASYMMETRIC_AUTOENCODER_KL_x_1_5_CONFIG = { |
|
"in_channels": 3, |
|
"out_channels": 3, |
|
"down_block_types": [ |
|
"DownEncoderBlock2D", |
|
"DownEncoderBlock2D", |
|
"DownEncoderBlock2D", |
|
"DownEncoderBlock2D", |
|
], |
|
"down_block_out_channels": [128, 256, 512, 512], |
|
"layers_per_down_block": 2, |
|
"up_block_types": [ |
|
"UpDecoderBlock2D", |
|
"UpDecoderBlock2D", |
|
"UpDecoderBlock2D", |
|
"UpDecoderBlock2D", |
|
], |
|
"up_block_out_channels": [192, 384, 768, 768], |
|
"layers_per_up_block": 3, |
|
"act_fn": "silu", |
|
"latent_channels": 4, |
|
"norm_num_groups": 32, |
|
"sample_size": 256, |
|
"scaling_factor": 0.18215, |
|
} |
|
|
|
ASYMMETRIC_AUTOENCODER_KL_x_2_CONFIG = { |
|
"in_channels": 3, |
|
"out_channels": 3, |
|
"down_block_types": [ |
|
"DownEncoderBlock2D", |
|
"DownEncoderBlock2D", |
|
"DownEncoderBlock2D", |
|
"DownEncoderBlock2D", |
|
], |
|
"down_block_out_channels": [128, 256, 512, 512], |
|
"layers_per_down_block": 2, |
|
"up_block_types": [ |
|
"UpDecoderBlock2D", |
|
"UpDecoderBlock2D", |
|
"UpDecoderBlock2D", |
|
"UpDecoderBlock2D", |
|
], |
|
"up_block_out_channels": [256, 512, 1024, 1024], |
|
"layers_per_up_block": 5, |
|
"act_fn": "silu", |
|
"latent_channels": 4, |
|
"norm_num_groups": 32, |
|
"sample_size": 256, |
|
"scaling_factor": 0.18215, |
|
} |
|
|
|
|
|
def convert_asymmetric_autoencoder_kl_state_dict(original_state_dict: Dict[str, Any]) -> Dict[str, Any]: |
|
converted_state_dict = {} |
|
for k, v in original_state_dict.items(): |
|
if k.startswith("encoder."): |
|
converted_state_dict[ |
|
k.replace("encoder.down.", "encoder.down_blocks.") |
|
.replace("encoder.mid.", "encoder.mid_block.") |
|
.replace("encoder.norm_out.", "encoder.conv_norm_out.") |
|
.replace(".downsample.", ".downsamplers.0.") |
|
.replace(".nin_shortcut.", ".conv_shortcut.") |
|
.replace(".block.", ".resnets.") |
|
.replace(".block_1.", ".resnets.0.") |
|
.replace(".block_2.", ".resnets.1.") |
|
.replace(".attn_1.k.", ".attentions.0.to_k.") |
|
.replace(".attn_1.q.", ".attentions.0.to_q.") |
|
.replace(".attn_1.v.", ".attentions.0.to_v.") |
|
.replace(".attn_1.proj_out.", ".attentions.0.to_out.0.") |
|
.replace(".attn_1.norm.", ".attentions.0.group_norm.") |
|
] = v |
|
elif k.startswith("decoder.") and "up_layers" not in k: |
|
converted_state_dict[ |
|
k.replace("decoder.encoder.", "decoder.condition_encoder.") |
|
.replace(".norm_out.", ".conv_norm_out.") |
|
.replace(".up.0.", ".up_blocks.3.") |
|
.replace(".up.1.", ".up_blocks.2.") |
|
.replace(".up.2.", ".up_blocks.1.") |
|
.replace(".up.3.", ".up_blocks.0.") |
|
.replace(".block.", ".resnets.") |
|
.replace("mid", "mid_block") |
|
.replace(".0.upsample.", ".0.upsamplers.0.") |
|
.replace(".1.upsample.", ".1.upsamplers.0.") |
|
.replace(".2.upsample.", ".2.upsamplers.0.") |
|
.replace(".nin_shortcut.", ".conv_shortcut.") |
|
.replace(".block_1.", ".resnets.0.") |
|
.replace(".block_2.", ".resnets.1.") |
|
.replace(".attn_1.k.", ".attentions.0.to_k.") |
|
.replace(".attn_1.q.", ".attentions.0.to_q.") |
|
.replace(".attn_1.v.", ".attentions.0.to_v.") |
|
.replace(".attn_1.proj_out.", ".attentions.0.to_out.0.") |
|
.replace(".attn_1.norm.", ".attentions.0.group_norm.") |
|
] = v |
|
elif k.startswith("quant_conv."): |
|
converted_state_dict[k] = v |
|
elif k.startswith("post_quant_conv."): |
|
converted_state_dict[k] = v |
|
else: |
|
print(f" skipping key `{k}`") |
|
|
|
for k, v in converted_state_dict.items(): |
|
if ( |
|
(k.startswith("encoder.mid_block.attentions.0") or k.startswith("decoder.mid_block.attentions.0")) |
|
and k.endswith("weight") |
|
and ("to_q" in k or "to_k" in k or "to_v" in k or "to_out" in k) |
|
): |
|
converted_state_dict[k] = converted_state_dict[k][:, :, 0, 0] |
|
|
|
return converted_state_dict |
|
|
|
|
|
def get_asymmetric_autoencoder_kl_from_original_checkpoint( |
|
scale: Literal["1.5", "2"], original_checkpoint_path: str, map_location: torch.device |
|
) -> AsymmetricAutoencoderKL: |
|
print("Loading original state_dict") |
|
original_state_dict = torch.load(original_checkpoint_path, map_location=map_location) |
|
original_state_dict = original_state_dict["state_dict"] |
|
print("Converting state_dict") |
|
converted_state_dict = convert_asymmetric_autoencoder_kl_state_dict(original_state_dict) |
|
kwargs = ASYMMETRIC_AUTOENCODER_KL_x_1_5_CONFIG if scale == "1.5" else ASYMMETRIC_AUTOENCODER_KL_x_2_CONFIG |
|
print("Initializing AsymmetricAutoencoderKL model") |
|
asymmetric_autoencoder_kl = AsymmetricAutoencoderKL(**kwargs) |
|
print("Loading weight from converted state_dict") |
|
asymmetric_autoencoder_kl.load_state_dict(converted_state_dict) |
|
asymmetric_autoencoder_kl.eval() |
|
print("AsymmetricAutoencoderKL successfully initialized") |
|
return asymmetric_autoencoder_kl |
|
|
|
|
|
if __name__ == "__main__": |
|
start = time.time() |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument( |
|
"--scale", |
|
default=None, |
|
type=str, |
|
required=True, |
|
help="Asymmetric VQGAN scale: `1.5` or `2`", |
|
) |
|
parser.add_argument( |
|
"--original_checkpoint_path", |
|
default=None, |
|
type=str, |
|
required=True, |
|
help="Path to the original Asymmetric VQGAN checkpoint", |
|
) |
|
parser.add_argument( |
|
"--output_path", |
|
default=None, |
|
type=str, |
|
required=True, |
|
help="Path to save pretrained AsymmetricAutoencoderKL model", |
|
) |
|
parser.add_argument( |
|
"--map_location", |
|
default="cpu", |
|
type=str, |
|
required=False, |
|
help="The device passed to `map_location` when loading the checkpoint", |
|
) |
|
args = parser.parse_args() |
|
|
|
assert args.scale in ["1.5", "2"], f"{args.scale} should be `1.5` of `2`" |
|
assert Path(args.original_checkpoint_path).is_file() |
|
|
|
asymmetric_autoencoder_kl = get_asymmetric_autoencoder_kl_from_original_checkpoint( |
|
scale=args.scale, |
|
original_checkpoint_path=args.original_checkpoint_path, |
|
map_location=torch.device(args.map_location), |
|
) |
|
print("Saving pretrained AsymmetricAutoencoderKL") |
|
asymmetric_autoencoder_kl.save_pretrained(args.output_path) |
|
print(f"Done in {time.time() - start:.2f} seconds") |
|
|
