import argparse
import torch
from safetensors.torch import load_file, save_file

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--sd15", default=None, type=str, required=True, help="Path to the original sd15.")
    parser.add_argument("--control", default=None, type=str, required=True, help="Path to the sd15 with control.")
    parser.add_argument("--dst", default=None, type=str, required=True, help="Path to the output difference model.")
    parser.add_argument("--fp16", action="store_true", help="Save as fp16.")
    parser.add_argument("--bf16", action="store_true", help="Save as bf16.")
    args = parser.parse_args()

    assert args.sd15 is not None, "Must provide a original sd15 model path!"
    assert args.control is not None, "Must provide a sd15 with control model path!"
    assert args.dst is not None, "Must provide a output path!"

    # make differences: copy from https://github.com/lllyasviel/ControlNet/blob/main/tool_transfer_control.py

    def get_node_name(name, parent_name):
        if len(name) <= len(parent_name):
            return False, ''
        p = name[:len(parent_name)]
        if p != parent_name:
            return False, ''
        return True, name[len(parent_name):]

    # remove first/cond stage from sd to reduce memory usage
    def remove_first_and_cond(sd):
        keys = list(sd.keys())
        for key in keys:
            is_first_stage, _ = get_node_name(key, 'first_stage_model')
            is_cond_stage, _ = get_node_name(key, 'cond_stage_model')
            if is_first_stage or is_cond_stage:
                sd.pop(key, None)
        return sd
    
    print(f"loading: {args.sd15}")
    if args.sd15.endswith(".safetensors"):
        sd15_state_dict = load_file(args.sd15)
    else:
        sd15_state_dict = torch.load(args.sd15)
        sd15_state_dict = sd15_state_dict.pop("state_dict", sd15_state_dict)
    sd15_state_dict = remove_first_and_cond(sd15_state_dict)

    print(f"loading: {args.control}")
    if args.control.endswith(".safetensors"):
        control_state_dict = load_file(args.control)
    else:
        control_state_dict = torch.load(args.control)
    control_state_dict = remove_first_and_cond(control_state_dict)

    # make diff of original and control
    print(f"create difference")
    keys = list(control_state_dict.keys())
    final_state_dict = {"difference": torch.tensor(1.0)}                  # indicates difference
    for key in keys:
        p = control_state_dict.pop(key)

        is_control, node_name = get_node_name(key, 'control_')
        if not is_control:
            continue

        sd15_key_name = 'model.diffusion_' + node_name
        if sd15_key_name in sd15_state_dict:                              # part of U-Net
            # print("in sd15", key, sd15_key_name)
            p_new = p - sd15_state_dict.pop(sd15_key_name)
            if torch.max(torch.abs(p_new)) < 1e-6:                        # no difference?
                print("no diff", key, sd15_key_name)
                continue
        else:
            # print("not in sd15", key, sd15_key_name)
            p_new = p                                                     # hint or zero_conv

        final_state_dict[key] = p_new

    save_dtype = None
    if args.fp16:
        save_dtype = torch.float16
    elif args.bf16:
        save_dtype = torch.bfloat16
    if save_dtype is not None:
        for key in final_state_dict.keys():
            final_state_dict[key] = final_state_dict[key].to(save_dtype)

    print("saving difference.")
    if args.dst.endswith(".safetensors"):
        save_file(final_state_dict, args.dst)
    else:
        torch.save({"state_dict": final_state_dict}, args.dst)
    print("done!")