| import argparse |
| from typing import Dict |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from diffusers import SparseControlNetModel |
|
|
|
|
| KEYS_RENAME_MAPPING = { |
| ".attention_blocks.0": ".attn1", |
| ".attention_blocks.1": ".attn2", |
| ".attn1.pos_encoder": ".pos_embed", |
| ".ff_norm": ".norm3", |
| ".norms.0": ".norm1", |
| ".norms.1": ".norm2", |
| ".temporal_transformer": "", |
| } |
|
|
|
|
| def convert(original_state_dict: Dict[str, nn.Module]) -> Dict[str, nn.Module]: |
| converted_state_dict = {} |
|
|
| for key in list(original_state_dict.keys()): |
| renamed_key = key |
| for new_name, old_name in KEYS_RENAME_MAPPING.items(): |
| renamed_key = renamed_key.replace(new_name, old_name) |
| converted_state_dict[renamed_key] = original_state_dict.pop(key) |
|
|
| return converted_state_dict |
|
|
|
|
| def get_args(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--ckpt_path", type=str, required=True, help="Path to checkpoint") |
| parser.add_argument("--output_path", type=str, required=True, help="Path to output directory") |
| parser.add_argument( |
| "--max_motion_seq_length", |
| type=int, |
| default=32, |
| help="Max motion sequence length supported by the motion adapter", |
| ) |
| parser.add_argument( |
| "--conditioning_channels", type=int, default=4, help="Number of channels in conditioning input to controlnet" |
| ) |
| parser.add_argument( |
| "--use_simplified_condition_embedding", |
| action="store_true", |
| default=False, |
| help="Whether or not to use simplified condition embedding. When `conditioning_channels==4` i.e. latent inputs, set this to `True`. When `conditioning_channels==3` i.e. image inputs, set this to `False`", |
| ) |
| parser.add_argument( |
| "--save_fp16", |
| action="store_true", |
| default=False, |
| help="Whether or not to save model in fp16 precision along with fp32", |
| ) |
| parser.add_argument( |
| "--push_to_hub", action="store_true", default=False, help="Whether or not to push saved model to the HF hub" |
| ) |
| return parser.parse_args() |
|
|
|
|
| if __name__ == "__main__": |
| args = get_args() |
|
|
| state_dict = torch.load(args.ckpt_path, map_location="cpu") |
| if "state_dict" in state_dict.keys(): |
| state_dict: dict = state_dict["state_dict"] |
|
|
| controlnet = SparseControlNetModel( |
| conditioning_channels=args.conditioning_channels, |
| motion_max_seq_length=args.max_motion_seq_length, |
| use_simplified_condition_embedding=args.use_simplified_condition_embedding, |
| ) |
|
|
| state_dict = convert(state_dict) |
| controlnet.load_state_dict(state_dict, strict=True) |
|
|
| controlnet.save_pretrained(args.output_path, push_to_hub=args.push_to_hub) |
| if args.save_fp16: |
| controlnet = controlnet.to(dtype=torch.float16) |
| controlnet.save_pretrained(args.output_path, variant="fp16", push_to_hub=args.push_to_hub) |
|
|
