57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
import argparse
|
|
|
|
import OmegaConf
|
|
import torch
|
|
|
|
from diffusers import UNetLDMModel, VQModel, LDMPipeline, DDIMScheduler
|
|
|
|
def convert_ldm_original(checkpoint_path, config_path, output_path):
|
|
config = OmegaConf.load(config_path)
|
|
state_dict = torch.load(checkpoint_path, map_location="cpu")["model"]
|
|
keys = list(state_dict.keys())
|
|
|
|
# extract state_dict for VQVAE
|
|
first_stage_dict = {}
|
|
first_stage_key = "first_stage_model."
|
|
for key in keys:
|
|
if key.startswith(first_stage_key):
|
|
first_stage_dict[key.replace(first_stage_key, "")] = state_dict[key]
|
|
|
|
# extract state_dict for UNetLDM
|
|
unet_state_dict = {}
|
|
unet_key = "model.diffusion_model."
|
|
for key in keys:
|
|
if key.startswith(unet_key):
|
|
unet_state_dict[key.replace(unet_key, "")] = state_dict[key]
|
|
|
|
vqvae_init_args = config.model.params.first_stage_config.params
|
|
unet_init_args = config.model.params.unet_config.params
|
|
|
|
vqvae = VQModel(**vqvae_init_args).eval()
|
|
vqvae.load_state_dict(first_stage_dict)
|
|
|
|
unet = UNetLDMModel(**unet_init_args).eval()
|
|
unet.load_state_dict(unet_state_dict)
|
|
|
|
noise_scheduler = DDIMScheduler(
|
|
timesteps=config.model.params.timesteps,
|
|
beta_schedule="scaled_linear",
|
|
beta_start=config.model.params.linear_start,
|
|
beta_end=config.model.params.linear_end,
|
|
clip_sample=False,
|
|
)
|
|
|
|
pipeline = LDMPipeline(vqvae, unet, noise_scheduler)
|
|
pipeline.save_pretrained(output_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--checkpoint_path", type=str, required=True)
|
|
parser.add_argument("--config_path", type=str, required=True)
|
|
parser.add_argument("--output_path", type=str, required=True)
|
|
args = parser.parse_args()
|
|
|
|
convert_ldm_original(args.checkpoint_path, args.config_path, args.output_path)
|
|
|