ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-02-27 07:43:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into development

Browse Source
This commit is contained in:
Jaret Burkett
2023-08-28 16:21:51 -06:00
parent 2ecb5cf024 bee0b6a235
commit b79ced3e10
9 changed files with 5120 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files

208
testing/generate_weight_mappings.py
View File

@@ -2,6 +2,11 @@ import argparse
import gc
import os
import re
import os
# add project root to sys path
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
from diffusers.loaders import LoraLoaderMixin
@@ -90,90 +95,134 @@ matched_diffusers_keys = []
error_margin = 1e-4
te_suffix = ''
proj_pattern_weight = None
proj_pattern_bias = None
text_proj_layer = None
if args.sdxl:
te_suffix = '1'
ldm_res_block_prefix = "conditioner.embedders.1.model.transformer.resblocks"
proj_pattern_weight = r"conditioner\.embedders\.1\.model\.transformer\.resblocks\.(\d+)\.attn\.in_proj_weight"
proj_pattern_bias = r"conditioner\.embedders\.1\.model\.transformer\.resblocks\.(\d+)\.attn\.in_proj_bias"
text_proj_layer = "conditioner.embedders.1.model.text_projection"
if args.sd2:
te_suffix = ''
ldm_res_block_prefix = "cond_stage_model.model.transformer.resblocks"
proj_pattern_weight = r"cond_stage_model\.model\.transformer\.resblocks\.(\d+)\.attn\.in_proj_weight"
proj_pattern_bias = r"cond_stage_model\.model\.transformer\.resblocks\.(\d+)\.attn\.in_proj_bias"
text_proj_layer = "cond_stage_model.model.text_projection"
if args.sdxl or args.sd2:
if "conditioner.embedders.1.model.text_projection" in ldm_dict_keys:
# d_model = int(checkpoint[prefix + "text_projection"].shape[0]))
d_model = int(ldm_state_dict["conditioner.embedders.1.model.text_projection"].shape[0])
else:
d_model = 1024
# do pre known merging
for ldm_key in ldm_dict_keys:
pattern = r"conditioner\.embedders\.1\.model\.transformer\.resblocks\.(\d+)\.attn\.in_proj_weight"
match = re.match(pattern, ldm_key)
if match:
number = int(match.group(1))
new_val = torch.cat([
diffusers_state_dict[f"te1_text_model.encoder.layers.{number}.self_attn.q_proj.weight"],
diffusers_state_dict[f"te1_text_model.encoder.layers.{number}.self_attn.k_proj.weight"],
diffusers_state_dict[f"te1_text_model.encoder.layers.{number}.self_attn.v_proj.weight"],
], dim=0)
# add to matched so we dont check them
matched_diffusers_keys.append(f"te1_text_model.encoder.layers.{number}.self_attn.q_proj.weight")
matched_diffusers_keys.append(f"te1_text_model.encoder.layers.{number}.self_attn.k_proj.weight")
matched_diffusers_keys.append(f"te1_text_model.encoder.layers.{number}.self_attn.v_proj.weight")
# make diffusers convertable_dict
diffusers_state_dict[f"te1_text_model.encoder.layers.{number}.self_attn.MERGED.weight"] = new_val
try:
match = re.match(proj_pattern_weight, ldm_key)
if match:
number = int(match.group(1))
new_val = torch.cat([
diffusers_state_dict[f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.q_proj.weight"],
diffusers_state_dict[f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.k_proj.weight"],
diffusers_state_dict[f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.v_proj.weight"],
], dim=0)
# add to matched so we dont check them
matched_diffusers_keys.append(
f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.q_proj.weight")
matched_diffusers_keys.append(
f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.k_proj.weight")
matched_diffusers_keys.append(
f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.v_proj.weight")
# make diffusers convertable_dict
diffusers_state_dict[
f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.MERGED.weight"] = new_val
# add operator
ldm_operator_map[ldm_key] = {
"cat": [
f"te1_text_model.encoder.layers.{number}.self_attn.q_proj.weight",
f"te1_text_model.encoder.layers.{number}.self_attn.k_proj.weight",
f"te1_text_model.encoder.layers.{number}.self_attn.v_proj.weight",
],
"target": f"te1_text_model.encoder.layers.{number}.self_attn.MERGED.weight"
}
# add operator
ldm_operator_map[ldm_key] = {
"cat": [
f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.q_proj.weight",
f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.k_proj.weight",
f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.v_proj.weight",
],
"target": f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.MERGED.weight"
}
if "conditioner.embedders.1.model.text_projection" in ldm_dict_keys:
# d_model = int(checkpoint[prefix + "text_projection"].shape[0]))
d_model = int(ldm_state_dict["conditioner.embedders.1.model.text_projection"].shape[0])
else:
d_model = 1024
# text_model_dict[new_key + ".q_proj.weight"] = checkpoint[key][:d_model, :]
# text_model_dict[new_key + ".k_proj.weight"] = checkpoint[key][d_model: d_model * 2, :]
# text_model_dict[new_key + ".v_proj.weight"] = checkpoint[key][d_model * 2:, :]
# text_model_dict[new_key + ".q_proj.weight"] = checkpoint[key][:d_model, :]
# text_model_dict[new_key + ".k_proj.weight"] = checkpoint[key][d_model: d_model * 2, :]
# text_model_dict[new_key + ".v_proj.weight"] = checkpoint[key][d_model * 2:, :]
# add diffusers operators
diffusers_operator_map[f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.q_proj.weight"] = {
"slice": [
f"{ldm_res_block_prefix}.{number}.attn.in_proj_weight",
f"0:{d_model}, :"
]
}
diffusers_operator_map[f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.k_proj.weight"] = {
"slice": [
f"{ldm_res_block_prefix}.{number}.attn.in_proj_weight",
f"{d_model}:{d_model * 2}, :"
]
}
diffusers_operator_map[f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.v_proj.weight"] = {
"slice": [
f"{ldm_res_block_prefix}.{number}.attn.in_proj_weight",
f"{d_model * 2}:, :"
]
}
# add diffusers operators
diffusers_operator_map[f"te1_text_model.encoder.layers.{number}.self_attn.q_proj.weight"] = {
"slice": [
f"conditioner.embedders.1.model.transformer.resblocks.{number}.attn.in_proj_weight",
f"0:{d_model}, :"
]
}
diffusers_operator_map[f"te1_text_model.encoder.layers.{number}.self_attn.k_proj.weight"] = {
"slice": [
f"conditioner.embedders.1.model.transformer.resblocks.{number}.attn.in_proj_weight",
f"{d_model}:{d_model * 2}, :"
]
}
diffusers_operator_map[f"te1_text_model.encoder.layers.{number}.self_attn.v_proj.weight"] = {
"slice": [
f"conditioner.embedders.1.model.transformer.resblocks.{number}.attn.in_proj_weight",
f"{d_model * 2}:, :"
]
}
match = re.match(proj_pattern_bias, ldm_key)
if match:
number = int(match.group(1))
new_val = torch.cat([
diffusers_state_dict[f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.q_proj.bias"],
diffusers_state_dict[f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.k_proj.bias"],
diffusers_state_dict[f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.v_proj.bias"],
], dim=0)
# add to matched so we dont check them
matched_diffusers_keys.append(f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.q_proj.bias")
matched_diffusers_keys.append(f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.k_proj.bias")
matched_diffusers_keys.append(f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.v_proj.bias")
# make diffusers convertable_dict
diffusers_state_dict[
f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.MERGED.bias"] = new_val
pattern = r"conditioner\.embedders\.1\.model\.transformer\.resblocks\.(\d+)\.attn\.in_proj_bias"
match = re.match(pattern, ldm_key)
if match:
number = int(match.group(1))
new_val = torch.cat([
diffusers_state_dict[f"te1_text_model.encoder.layers.{number}.self_attn.q_proj.bias"],
diffusers_state_dict[f"te1_text_model.encoder.layers.{number}.self_attn.k_proj.bias"],
diffusers_state_dict[f"te1_text_model.encoder.layers.{number}.self_attn.v_proj.bias"],
], dim=0)
# add to matched so we dont check them
matched_diffusers_keys.append(f"te1_text_model.encoder.layers.{number}.self_attn.q_proj.bias")
matched_diffusers_keys.append(f"te1_text_model.encoder.layers.{number}.self_attn.k_proj.bias")
matched_diffusers_keys.append(f"te1_text_model.encoder.layers.{number}.self_attn.v_proj.bias")
# make diffusers convertable_dict
diffusers_state_dict[f"te1_text_model.encoder.layers.{number}.self_attn.MERGED.bias"] = new_val
# add operator
ldm_operator_map[ldm_key] = {
"cat": [
f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.q_proj.bias",
f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.k_proj.bias",
f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.v_proj.bias",
],
# "target": f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.MERGED.bias"
}
# add operator
ldm_operator_map[ldm_key] = {
"cat": [
f"te1_text_model.encoder.layers.{number}.self_attn.q_proj.bias",
f"te1_text_model.encoder.layers.{number}.self_attn.k_proj.bias",
f"te1_text_model.encoder.layers.{number}.self_attn.v_proj.bias",
],
"target": f"te1_text_model.encoder.layers.{number}.self_attn.MERGED.bias"
}
# add diffusers operators
diffusers_operator_map[f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.q_proj.bias"] = {
"slice": [
f"{ldm_res_block_prefix}.{number}.attn.in_proj_bias",
f"0:{d_model}, :"
]
}
diffusers_operator_map[f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.k_proj.bias"] = {
"slice": [
f"{ldm_res_block_prefix}.{number}.attn.in_proj_bias",
f"{d_model}:{d_model * 2}, :"
]
}
diffusers_operator_map[f"te{te_suffix}_text_model.encoder.layers.{number}.self_attn.v_proj.bias"] = {
"slice": [
f"{ldm_res_block_prefix}.{number}.attn.in_proj_bias",
f"{d_model * 2}:, :"
]
}
except Exception as e:
print(f"Error on key {ldm_key}")
print(e)
# update keys
diffusers_dict_keys = list(diffusers_state_dict.keys())
@@ -275,14 +324,10 @@ if has_unmatched_keys:
weight = ldm_state_dict[key]
weight_min = weight.min().item()
weight_max = weight.max().item()
weight_mean = weight.mean().item()
weight_std = weight.std().item()
unmatched_obj['ldm'][key] = {
'shape': weight.shape,
"min": weight_min,
"max": weight_max,
"mean": weight_mean,
"std": weight_std,
}
del weight
flush()
@@ -292,14 +337,10 @@ if has_unmatched_keys:
weight = diffusers_state_dict[key]
weight_min = weight.min().item()
weight_max = weight.max().item()
weight_mean = weight.mean().item()
weight_std = weight.std().item()
unmatched_obj['diffusers'][key] = {
"shape": weight.shape,
"min": weight_min,
"max": weight_max,
"mean": weight_mean,
"std": weight_std,
}
del weight
flush()
@@ -318,7 +359,6 @@ for key in unmatched_ldm_keys:
save_file(remaining_ldm_values, os.path.join(KEYMAPS_FOLDER, f'{name}_ldm_base.safetensors'))
print(f'Saved remaining ldm values to {os.path.join(KEYMAPS_FOLDER, f"{name}_ldm_base.safetensors")}')
dest_path = os.path.join(KEYMAPS_FOLDER, f'{name}.json')
save_obj = OrderedDict()
save_obj["ldm_diffusers_keymap"] = ldm_diffusers_keymap

21
testing/test_model_load_save.py
View File

@@ -13,7 +13,7 @@ import json
from toolkit.config_modules import ModelConfig
from toolkit.paths import KEYMAPS_ROOT
from toolkit.saving import convert_state_dict_to_ldm_with_mapping
from toolkit.saving import convert_state_dict_to_ldm_with_mapping, get_ldm_state_dict_from_diffusers
from toolkit.stable_diffusion_model import StableDiffusion
# this was just used to match the vae keys to the diffusers keys
@@ -39,6 +39,12 @@ parser.add_argument(
help='Is the model an XL model'
)
parser.add_argument(
'--is_v2',
action='store_true',
help='Is the model a v2 model'
)
args = parser.parse_args()
find_matches = False
@@ -58,19 +64,20 @@ sd = StableDiffusion(
)
sd.load_model()
if not args.is_xl:
# not supported yet
raise NotImplementedError("Only SDXL is supported at this time with this method")
# load our base
base_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sdxl_ldm_base.safetensors')
mapping_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sdxl.json')
print("Converting model back to LDM")
version_string = '1'
if args.is_v2:
version_string = '2'
if args.is_xl:
version_string = 'sdxl'
# convert the state dict
state_dict_file_2 = convert_state_dict_to_ldm_with_mapping(
state_dict_file_2 = get_ldm_state_dict_from_diffusers(
sd.state_dict(),
mapping_path,
base_path,
version_string,
device='cpu',
dtype=dtype
)

1962
toolkit/keymaps/stable_diffusion_sd1.json Normal file
View File

File diff suppressed because it is too large Load Diff

BIN
toolkit/keymaps/stable_diffusion_sd1_ldm_base.safetensors Normal file
View File

Binary file not shown.

2773
toolkit/keymaps/stable_diffusion_sd2.json Normal file
View File

File diff suppressed because it is too large Load Diff

BIN
toolkit/keymaps/stable_diffusion_sd2_ldm_base.safetensors Normal file
View File

Binary file not shown.

200
toolkit/keymaps/stable_diffusion_sd2_unmatched.json Normal file
View File

@@ -0,0 +1,200 @@
{
"ldm": {
"alphas_cumprod": {
"shape": [
1000
],
"min": 0.00466156005859375,
"max": 0.9990234375
},
"alphas_cumprod_prev": {
"shape": [
1000
],
"min": 0.0047149658203125,
"max": 1.0
},
"betas": {
"shape": [
1000
],
"min": 0.0008502006530761719,
"max": 0.01200103759765625
},
"cond_stage_model.model.logit_scale": {
"shape": [],
"min": 4.60546875,
"max": 4.60546875
},
"cond_stage_model.model.text_projection": {
"shape": [
1024,
1024
],
"min": -0.109130859375,
"max": 0.09271240234375
},
"cond_stage_model.model.transformer.resblocks.23.attn.in_proj_bias": {
"shape": [
3072
],
"min": -2.525390625,
"max": 2.591796875
},
"cond_stage_model.model.transformer.resblocks.23.attn.in_proj_weight": {
"shape": [
3072,
1024
],
"min": -0.12261962890625,
"max": 0.1258544921875
},
"cond_stage_model.model.transformer.resblocks.23.attn.out_proj.bias": {
"shape": [
1024
],
"min": -0.422607421875,
"max": 1.17578125
},
"cond_stage_model.model.transformer.resblocks.23.attn.out_proj.weight": {
"shape": [
1024,
1024
],
"min": -0.0738525390625,
"max": 0.08673095703125
},
"cond_stage_model.model.transformer.resblocks.23.ln_1.bias": {
"shape": [
1024
],
"min": -3.392578125,
"max": 0.90625
},
"cond_stage_model.model.transformer.resblocks.23.ln_1.weight": {
"shape": [
1024
],
"min": 0.379638671875,
"max": 2.02734375
},
"cond_stage_model.model.transformer.resblocks.23.ln_2.bias": {
"shape": [
1024
],
"min": -0.833984375,
"max": 2.525390625
},
"cond_stage_model.model.transformer.resblocks.23.ln_2.weight": {
"shape": [
1024
],
"min": 1.17578125,
"max": 2.037109375
},
"cond_stage_model.model.transformer.resblocks.23.mlp.c_fc.bias": {
"shape": [
4096
],
"min": -1.619140625,
"max": 0.5595703125
},
"cond_stage_model.model.transformer.resblocks.23.mlp.c_fc.weight": {
"shape": [
4096,
1024
],
"min": -0.08953857421875,
"max": 0.13232421875
},
"cond_stage_model.model.transformer.resblocks.23.mlp.c_proj.bias": {
"shape": [
1024
],
"min": -1.8662109375,
"max": 0.74658203125
},
"cond_stage_model.model.transformer.resblocks.23.mlp.c_proj.weight": {
"shape": [
1024,
4096
],
"min": -0.12939453125,
"max": 0.1009521484375
},
"log_one_minus_alphas_cumprod": {
"shape": [
1000
],
"min": -7.0703125,
"max": -0.004669189453125
},
"model_ema.decay": {
"shape": [],
"min": 1.0,
"max": 1.0
},
"model_ema.num_updates": {
"shape": [],
"min": 219996,
"max": 219996
},
"posterior_log_variance_clipped": {
"shape": [
1000
],
"min": -46.0625,
"max": -4.421875
},
"posterior_mean_coef1": {
"shape": [
1000
],
"min": 0.000827789306640625,
"max": 1.0
},
"posterior_mean_coef2": {
"shape": [
1000
],
"min": 0.0,
"max": 0.99560546875
},
"posterior_variance": {
"shape": [
1000
],
"min": 0.0,
"max": 0.01200103759765625
},
"sqrt_alphas_cumprod": {
"shape": [
1000
],
"min": 0.0682373046875,
"max": 0.99951171875
},
"sqrt_one_minus_alphas_cumprod": {
"shape": [
1000
],
"min": 0.0291595458984375,
"max": 0.99755859375
},
"sqrt_recip_alphas_cumprod": {
"shape": [
1000
],
"min": 1.0,
"max": 14.6484375
},
"sqrt_recipm1_alphas_cumprod": {
"shape": [
1000
],
"min": 0.0291595458984375,
"max": 14.6171875
}
},
"diffusers": {}
}

45
toolkit/saving.py
View File

@@ -56,7 +56,7 @@ def convert_state_dict_to_ldm_with_mapping(
tensor_to_slice = diffusers_state_dict[ldm_diffusers_operator_map[ldm_key]['slice'][0]]
slice_text = diffusers_state_dict[ldm_diffusers_operator_map[ldm_key]['slice'][1]]
converted_state_dict[ldm_key] = tensor_to_slice[get_slices_from_string(slice_text)].detach().to(device,
dtype=dtype)
dtype=dtype)
# process the rest of the keys
for ldm_key in ldm_diffusers_keymap:
@@ -71,6 +71,35 @@ def convert_state_dict_to_ldm_with_mapping(
return converted_state_dict
def get_ldm_state_dict_from_diffusers(
state_dict: 'OrderedDict',
sd_version: Literal['1', '2', 'sdxl'] = '2',
device='cpu',
dtype=get_torch_dtype('fp32'),
):
if sd_version == '1':
base_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sd1_ldm_base.safetensors')
mapping_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sd1.json')
elif sd_version == '2':
base_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sd2_ldm_base.safetensors')
mapping_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sd2.json')
elif sd_version == 'sdxl':
# load our base
base_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sdxl_ldm_base.safetensors')
mapping_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sdxl.json')
else:
raise ValueError(f"Invalid sd_version {sd_version}")
# convert the state dict
return convert_state_dict_to_ldm_with_mapping(
state_dict,
mapping_path,
base_path,
device=device,
dtype=dtype
)
def save_ldm_model_from_diffusers(
sd: 'StableDiffusion',
output_file: str,
@@ -78,21 +107,13 @@ def save_ldm_model_from_diffusers(
save_dtype=get_torch_dtype('fp16'),
sd_version: Literal['1', '2', 'sdxl'] = '2'
):
if sd_version != 'sdxl':
# not supported yet
raise NotImplementedError("Only SDXL is supported at this time with this method")
# load our base
base_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sdxl_ldm_base.safetensors')
mapping_path = os.path.join(KEYMAPS_ROOT, 'stable_diffusion_sdxl.json')
# convert the state dict
converted_state_dict = convert_state_dict_to_ldm_with_mapping(
converted_state_dict = get_ldm_state_dict_from_diffusers(
sd.state_dict(),
mapping_path,
base_path,
sd_version,
device='cpu',
dtype=save_dtype
)
# make sure parent folder exists
os.makedirs(os.path.dirname(output_file), exist_ok=True)
save_file(converted_state_dict, output_file, metadata=meta)

61
toolkit/stable_diffusion_model.py
View File

@@ -709,61 +709,28 @@ class StableDiffusion:
for i, encoder in enumerate(self.text_encoder):
for k, v in encoder.state_dict().items():
new_key = k if k.startswith(
f"{SD_PREFIX_TEXT_ENCODER}{i}") else f"{SD_PREFIX_TEXT_ENCODER}{i}_{k}"
f"{SD_PREFIX_TEXT_ENCODER}{i}_") else f"{SD_PREFIX_TEXT_ENCODER}{i}_{k}"
state_dict[new_key] = v
else:
for k, v in self.text_encoder.state_dict().items():
new_key = k if k.startswith(f"{SD_PREFIX_TEXT_ENCODER}") else f"{SD_PREFIX_TEXT_ENCODER}_{k}"
new_key = k if k.startswith(f"{SD_PREFIX_TEXT_ENCODER}_") else f"{SD_PREFIX_TEXT_ENCODER}_{k}"
state_dict[new_key] = v
if unet:
for k, v in self.unet.state_dict().items():
new_key = k if k.startswith(f"{SD_PREFIX_UNET}") else f"{SD_PREFIX_UNET}_{k}"
new_key = k if k.startswith(f"{SD_PREFIX_UNET}_") else f"{SD_PREFIX_UNET}_{k}"
state_dict[new_key] = v
return state_dict
def save(self, output_file: str, meta: OrderedDict, save_dtype=get_torch_dtype('fp16'), logit_scale=None):
state_dict = {}
# prepare metadata
meta = get_meta_for_safetensors(meta)
def update_sd(prefix, sd):
for k, v in sd.items():
key = prefix + k
v = v.detach().clone()
state_dict[key] = v.to("cpu", dtype=get_torch_dtype(save_dtype))
# make sure there are not nan values
if torch.isnan(state_dict[key]).any():
raise ValueError(f"NaN value in state dict: {key}")
# todo see what logit scale is
version_string = '1'
if self.is_v2:
version_string = '2'
if self.is_xl:
save_ldm_model_from_diffusers(
sd=self,
output_file=output_file,
meta=meta,
save_dtype=save_dtype,
sd_version='sdxl',
)
else:
# Convert the UNet model
unet_state_dict = convert_unet_state_dict_to_sd(self.is_v2, self.unet.state_dict())
update_sd("model.diffusion_model.", unet_state_dict)
# Convert the text encoder model
if self.is_v2:
make_dummy = True
text_enc_dict = convert_text_encoder_state_dict_to_sd_v2(self.text_encoder.state_dict(), make_dummy)
update_sd("cond_stage_model.model.", text_enc_dict)
else:
text_enc_dict = self.text_encoder.state_dict()
update_sd("cond_stage_model.transformer.", text_enc_dict)
# Convert the VAE
if self.vae is not None:
vae_dict = model_util.convert_vae_state_dict(self.vae.state_dict())
update_sd("first_stage_model.", vae_dict)
# make sure parent folder exists
os.makedirs(os.path.dirname(output_file), exist_ok=True)
save_file(state_dict, output_file, metadata=meta)
version_string = 'sdxl'
save_ldm_model_from_diffusers(
sd=self,
output_file=output_file,
meta=meta,
save_dtype=save_dtype,
sd_version=version_string,
)