diffusers形式のInstructPix2Pixを通常モデルとマージ

main.py

# 動作未確認

import torch
from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler, StableDiffusionPipeline
import gc

# load models
ip2p_pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained("/content/drive/MyDrive/models/diffusers/ip2p", torch_dtype=torch.float16)
base_pipe = StableDiffusionPipeline.from_pretrained("/content/drive/MyDrive/models/diffusers/sd15_fp16", torch_dtype=torch.float16)
target_pipe = StableDiffusionPipeline.from_pretrained("/content/drive/MyDrive/models/diffusers/dedede_fp16", torch_dtype=torch.float16)

# get state dict
ip2p_dict = ip2p_pipe.unet.state_dict()
base_dict = base_pipe.unet.state_dict()
target_dict = target_pipe.unet.state_dict()

# merged = ip2p + target - base
for key in ip2p_dict.keys():
    # skip conv_in
    if key == "conv_in.weight":
        print("pass")
    else:
        ip2p_dict[key] = ip2p_dict[key] + target_dict[key] - base_dict[key]
        
# save コンフィグファイルが元のip2pのままなので、何か不具合が出るかもしれない
ip2p_pipe.unet.load_state_dict(ip2p_dict)
ip2p_pipe.vae = target_pipe.vae
ip2p_pipe.save_pretrained("/content/drive/MyDrive/models/diffusers/dedede_ip2p", safe_serialization=True)