|
import os, sys
|
|
|
|
sys.path.insert(0, os.getcwd())
|
|
import argparse
|
|
|
|
|
|
def get_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"base_model", help="The model you want to merge with loha", default="", type=str
|
|
)
|
|
parser.add_argument(
|
|
"lycoris_model",
|
|
help="the lyco model you want to merge into sd model",
|
|
default="",
|
|
type=str,
|
|
)
|
|
parser.add_argument(
|
|
"output_name", help="the output model", default="./out.pt", type=str
|
|
)
|
|
parser.add_argument(
|
|
"--is_v2",
|
|
help="Your base model is sd v2 or not",
|
|
default=False,
|
|
action="store_true",
|
|
)
|
|
parser.add_argument(
|
|
"--is_sdxl",
|
|
help="Your base/db model is sdxl or not",
|
|
default=False,
|
|
action="store_true",
|
|
)
|
|
parser.add_argument(
|
|
"--device",
|
|
help="Which device you want to use to merge the weight",
|
|
default="cpu",
|
|
type=str,
|
|
)
|
|
parser.add_argument("--dtype", help="dtype to save", default="float", type=str)
|
|
parser.add_argument(
|
|
"--weight", help="weight for the lyco model to merge", default="1.0", type=float
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
args = ARGS = get_args()
|
|
|
|
|
|
from lycoris.utils import merge
|
|
from lycoris.kohya.model_utils import (
|
|
load_models_from_stable_diffusion_checkpoint,
|
|
save_stable_diffusion_checkpoint,
|
|
load_file,
|
|
)
|
|
from lycoris.kohya.sdxl_model_util import (
|
|
load_models_from_sdxl_checkpoint,
|
|
save_stable_diffusion_checkpoint as save_sdxl_checkpoint,
|
|
)
|
|
|
|
import torch
|
|
|
|
|
|
@torch.no_grad()
|
|
def main():
|
|
if args.is_sdxl:
|
|
base = load_models_from_sdxl_checkpoint(
|
|
None, args.base_model, map_location=args.device
|
|
)
|
|
else:
|
|
base = load_models_from_stable_diffusion_checkpoint(args.is_v2, args.base_model)
|
|
if ARGS.lycoris_model.rsplit(".", 1)[-1] == "safetensors":
|
|
lyco = load_file(ARGS.lycoris_model)
|
|
else:
|
|
lyco = torch.load(ARGS.lycoris_model)
|
|
|
|
dtype_str = ARGS.dtype.replace("fp", "float").replace("bf", "bfloat")
|
|
dtype = {
|
|
"float": torch.float,
|
|
"float16": torch.float16,
|
|
"float32": torch.float32,
|
|
"float64": torch.float64,
|
|
"bfloat": torch.bfloat16,
|
|
"bfloat16": torch.bfloat16,
|
|
}.get(dtype_str, None)
|
|
if dtype is None:
|
|
raise ValueError(f'Cannot Find the dtype "{dtype}"')
|
|
|
|
if args.is_sdxl:
|
|
base_tes = [base[0], base[1]]
|
|
base_unet = base[3]
|
|
else:
|
|
base_tes = [base[0]]
|
|
base_unet = base[2]
|
|
|
|
merge(base_tes, base_unet, lyco, ARGS.weight, ARGS.device)
|
|
|
|
if args.is_sdxl:
|
|
save_sdxl_checkpoint(
|
|
ARGS.output_name,
|
|
base[0].cpu(),
|
|
base[1].cpu(),
|
|
base[3].cpu(),
|
|
0,
|
|
0,
|
|
None,
|
|
base[2],
|
|
getattr(base[1], "logit_scale", None),
|
|
dtype,
|
|
)
|
|
else:
|
|
save_stable_diffusion_checkpoint(
|
|
ARGS.is_v2,
|
|
ARGS.output_name,
|
|
base[0].cpu(),
|
|
base[2].cpu(),
|
|
None,
|
|
0,
|
|
0,
|
|
dtype,
|
|
base[1],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|
