Allex21/LT
0
1# v1: split from train_db_fixed.py.2# v2: support safetensors3 4import math5import os6 7import torch8from library.device_utils import init_ipex9init_ipex()10 11import diffusers12from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig, logging13from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline # , UNet2DConditionModel14from safetensors.torch import load_file, save_file15from library.original_unet import UNet2DConditionModel16from library.utils import setup_logging17setup_logging()18import logging19logger = logging.getLogger(__name__)20 21# DiffUsers版StableDiffusionのモデルパラメータ22NUM_TRAIN_TIMESTEPS = 100023BETA_START = 0.0008524BETA_END = 0.012025 26UNET_PARAMS_MODEL_CHANNELS = 32027UNET_PARAMS_CHANNEL_MULT = [1, 2, 4, 4]28UNET_PARAMS_ATTENTION_RESOLUTIONS = [4, 2, 1]29UNET_PARAMS_IMAGE_SIZE = 64 # fixed from old invalid value `32`30UNET_PARAMS_IN_CHANNELS = 431UNET_PARAMS_OUT_CHANNELS = 432UNET_PARAMS_NUM_RES_BLOCKS = 233UNET_PARAMS_CONTEXT_DIM = 76834UNET_PARAMS_NUM_HEADS = 835# UNET_PARAMS_USE_LINEAR_PROJECTION = False36 37VAE_PARAMS_Z_CHANNELS = 438VAE_PARAMS_RESOLUTION = 25639VAE_PARAMS_IN_CHANNELS = 340VAE_PARAMS_OUT_CH = 341VAE_PARAMS_CH = 12842VAE_PARAMS_CH_MULT = [1, 2, 4, 4]43VAE_PARAMS_NUM_RES_BLOCKS = 244 45# V246V2_UNET_PARAMS_ATTENTION_HEAD_DIM = [5, 10, 20, 20]47V2_UNET_PARAMS_CONTEXT_DIM = 102448# V2_UNET_PARAMS_USE_LINEAR_PROJECTION = True49 50# Diffusersの設定を読み込むための参照モデル51DIFFUSERS_REF_MODEL_ID_V1 = "runwayml/stable-diffusion-v1-5"52DIFFUSERS_REF_MODEL_ID_V2 = "stabilityai/stable-diffusion-2-1"53 54 55# region StableDiffusion->Diffusersの変換コード56# convert_original_stable_diffusion_to_diffusers をコピーして修正している（ASL 2.0）57 58 59def shave_segments(path, n_shave_prefix_segments=1):60 """61 Removes segments. Positive values shave the first segments, negative shave the last segments.62 """63 if n_shave_prefix_segments >= 0:64 return ".".join(path.split(".")[n_shave_prefix_segments:])65 else:66 return ".".join(path.split(".")[:n_shave_prefix_segments])67 68 69def renew_resnet_paths(old_list, n_shave_prefix_segments=0):70 """71 Updates paths inside resnets to the new naming scheme (local renaming)72 """73 mapping = []74 for old_item in old_list:75 new_item = old_item.replace("in_layers.0", "norm1")76 new_item = new_item.replace("in_layers.2", "conv1")77 78 new_item = new_item.replace("out_layers.0", "norm2")79 new_item = new_item.replace("out_layers.3", "conv2")80 81 new_item = new_item.replace("emb_layers.1", "time_emb_proj")82 new_item = new_item.replace("skip_connection", "conv_shortcut")83 84 new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)85 86 mapping.append({"old": old_item, "new": new_item})87 88 return mapping89 90 91def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0):92 """93 Updates paths inside resnets to the new naming scheme (local renaming)94 """95 mapping = []96 for old_item in old_list:97 new_item = old_item98 99 new_item = new_item.replace("nin_shortcut", "conv_shortcut")100 new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)101 102 mapping.append({"old": old_item, "new": new_item})103 104 return mapping105 106 107def renew_attention_paths(old_list, n_shave_prefix_segments=0):108 """109 Updates paths inside attentions to the new naming scheme (local renaming)110 """111 mapping = []112 for old_item in old_list:113 new_item = old_item114 115 # new_item = new_item.replace('norm.weight', 'group_norm.weight')116 # new_item = new_item.replace('norm.bias', 'group_norm.bias')117 118 # new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')119 # new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')120 121 # new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)122 123 mapping.append({"old": old_item, "new": new_item})124 125 return mapping126 127 128def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0):129 """130 Updates paths inside attentions to the new naming scheme (local renaming)131 """132 mapping = []133 for old_item in old_list:134 new_item = old_item135 136 new_item = new_item.replace("norm.weight", "group_norm.weight")137 new_item = new_item.replace("norm.bias", "group_norm.bias")138 139 if diffusers.__version__ < "0.17.0":140 new_item = new_item.replace("q.weight", "query.weight")141 new_item = new_item.replace("q.bias", "query.bias")142 143 new_item = new_item.replace("k.weight", "key.weight")144 new_item = new_item.replace("k.bias", "key.bias")145 146 new_item = new_item.replace("v.weight", "value.weight")147 new_item = new_item.replace("v.bias", "value.bias")148 149 new_item = new_item.replace("proj_out.weight", "proj_attn.weight")150 new_item = new_item.replace("proj_out.bias", "proj_attn.bias")151 else:152 new_item = new_item.replace("q.weight", "to_q.weight")153 new_item = new_item.replace("q.bias", "to_q.bias")154 155 new_item = new_item.replace("k.weight", "to_k.weight")156 new_item = new_item.replace("k.bias", "to_k.bias")157 158 new_item = new_item.replace("v.weight", "to_v.weight")159 new_item = new_item.replace("v.bias", "to_v.bias")160 161 new_item = new_item.replace("proj_out.weight", "to_out.0.weight")162 new_item = new_item.replace("proj_out.bias", "to_out.0.bias")163 164 new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)165 166 mapping.append({"old": old_item, "new": new_item})167 168 return mapping169 170 171def assign_to_checkpoint(172 paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None173):174 """175 This does the final conversion step: take locally converted weights and apply a global renaming176 to them. It splits attention layers, and takes into account additional replacements177 that may arise.178 179 Assigns the weights to the new checkpoint.180 """181 assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys."182 183 # Splits the attention layers into three variables.184 if attention_paths_to_split is not None:185 for path, path_map in attention_paths_to_split.items():186 old_tensor = old_checkpoint[path]187 channels = old_tensor.shape[0] // 3188 189 target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1)190 191 num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3192 193 old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:])194 query, key, value = old_tensor.split(channels // num_heads, dim=1)195 196 checkpoint[path_map["query"]] = query.reshape(target_shape)197 checkpoint[path_map["key"]] = key.reshape(target_shape)198 checkpoint[path_map["value"]] = value.reshape(target_shape)199 200 for path in paths:201 new_path = path["new"]202 203 # These have already been assigned204 if attention_paths_to_split is not None and new_path in attention_paths_to_split:205 continue206 207 # Global renaming happens here208 new_path = new_path.replace("middle_block.0", "mid_block.resnets.0")209 new_path = new_path.replace("middle_block.1", "mid_block.attentions.0")210 new_path = new_path.replace("middle_block.2", "mid_block.resnets.1")211 212 if additional_replacements is not None:213 for replacement in additional_replacements:214 new_path = new_path.replace(replacement["old"], replacement["new"])215 216 # proj_attn.weight has to be converted from conv 1D to linear217 reshaping = False218 if diffusers.__version__ < "0.17.0":219 if "proj_attn.weight" in new_path:220 reshaping = True221 else:222 if ".attentions." in new_path and ".0.to_" in new_path and old_checkpoint[path["old"]].ndim > 2:223 reshaping = True224 225 if reshaping:226 checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0, 0]227 else:228 checkpoint[new_path] = old_checkpoint[path["old"]]229 230 231def conv_attn_to_linear(checkpoint):232 keys = list(checkpoint.keys())233 attn_keys = ["query.weight", "key.weight", "value.weight"]234 for key in keys:235 if ".".join(key.split(".")[-2:]) in attn_keys:236 if checkpoint[key].ndim > 2:237 checkpoint[key] = checkpoint[key][:, :, 0, 0]238 elif "proj_attn.weight" in key:239 if checkpoint[key].ndim > 2:240 checkpoint[key] = checkpoint[key][:, :, 0]241 242 243def linear_transformer_to_conv(checkpoint):244 keys = list(checkpoint.keys())245 tf_keys = ["proj_in.weight", "proj_out.weight"]246 for key in keys:247 if ".".join(key.split(".")[-2:]) in tf_keys:248 if checkpoint[key].ndim == 2:249 checkpoint[key] = checkpoint[key].unsqueeze(2).unsqueeze(2)250 251 252def convert_ldm_unet_checkpoint(v2, checkpoint, config):253 """254 Takes a state dict and a config, and returns a converted checkpoint.255 """256 257 # extract state_dict for UNet258 unet_state_dict = {}259 unet_key = "model.diffusion_model."260 keys = list(checkpoint.keys())261 for key in keys:262 if key.startswith(unet_key):263 unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key)264 265 new_checkpoint = {}266 267 new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"]268 new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"]269 new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"]270 new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"]271 272 new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"]273 new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"]274 275 new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"]276 new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"]277 new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"]278 new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"]279 280 # Retrieves the keys for the input blocks only281 num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer})282 input_blocks = {283 layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}." in key] for layer_id in range(num_input_blocks)284 }285 286 # Retrieves the keys for the middle blocks only287 num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer})288 middle_blocks = {289 layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}." in key] for layer_id in range(num_middle_blocks)290 }291 292 # Retrieves the keys for the output blocks only293 num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer})294 output_blocks = {295 layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}." in key] for layer_id in range(num_output_blocks)296 }297 298 for i in range(1, num_input_blocks):299 block_id = (i - 1) // (config["layers_per_block"] + 1)300 layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1)301 302 resnets = [key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key]303 attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key]304 305 if f"input_blocks.{i}.0.op.weight" in unet_state_dict:306 new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop(307 f"input_blocks.{i}.0.op.weight"308 )309 new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop(f"input_blocks.{i}.0.op.bias")310 311 paths = renew_resnet_paths(resnets)312 meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"}313 assign_to_checkpoint(paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config)314 315 if len(attentions):316 paths = renew_attention_paths(attentions)317 meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"}318 assign_to_checkpoint(paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config)319 320 resnet_0 = middle_blocks[0]321 attentions = middle_blocks[1]322 resnet_1 = middle_blocks[2]323 324 resnet_0_paths = renew_resnet_paths(resnet_0)325 assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config)326 327 resnet_1_paths = renew_resnet_paths(resnet_1)328 assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=config)329 330 attentions_paths = renew_attention_paths(attentions)331 meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"}332 assign_to_checkpoint(attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config)333 334 for i in range(num_output_blocks):335 block_id = i // (config["layers_per_block"] + 1)336 layer_in_block_id = i % (config["layers_per_block"] + 1)337 output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]]338 output_block_list = {}339 340 for layer in output_block_layers:341 layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1)342 if layer_id in output_block_list:343 output_block_list[layer_id].append(layer_name)344 else:345 output_block_list[layer_id] = [layer_name]346 347 if len(output_block_list) > 1:348 resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key]349 attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key]350 351 resnet_0_paths = renew_resnet_paths(resnets)352 paths = renew_resnet_paths(resnets)353 354 meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"}355 assign_to_checkpoint(paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config)356 357 # オリジナル：358 # if ["conv.weight", "conv.bias"] in output_block_list.values():359 # index = list(output_block_list.values()).index(["conv.weight", "conv.bias"])360 361 # biasとweightの順番に依存しないようにする：もっといいやり方がありそうだが362 for l in output_block_list.values():363 l.sort()364 365 if ["conv.bias", "conv.weight"] in output_block_list.values():366 index = list(output_block_list.values()).index(["conv.bias", "conv.weight"])367 new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[368 f"output_blocks.{i}.{index}.conv.bias"369 ]370 new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[371 f"output_blocks.{i}.{index}.conv.weight"372 ]373 374 # Clear attentions as they have been attributed above.375 if len(attentions) == 2:376 attentions = []377 378 if len(attentions):379 paths = renew_attention_paths(attentions)380 meta_path = {381 "old": f"output_blocks.{i}.1",382 "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}",383 }384 assign_to_checkpoint(paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config)385 else:386 resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1)387 for path in resnet_0_paths:388 old_path = ".".join(["output_blocks", str(i), path["old"]])389 new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]])390 391 new_checkpoint[new_path] = unet_state_dict[old_path]392 393 # SDのv2では1*1のconv2dがlinearに変わっている394 # 誤って Diffusers 側を conv2d のままにしてしまったので、変換必要395 if v2 and not config.get("use_linear_projection", False):396 linear_transformer_to_conv(new_checkpoint)397 398 return new_checkpoint399 400 401def convert_ldm_vae_checkpoint(checkpoint, config):402 # extract state dict for VAE403 vae_state_dict = {}404 vae_key = "first_stage_model."405 keys = list(checkpoint.keys())406 for key in keys:407 if key.startswith(vae_key):408 vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key)409 # if len(vae_state_dict) == 0:410 # # 渡されたcheckpointは.ckptから読み込んだcheckpointではなくvaeのstate_dict411 # vae_state_dict = checkpoint412 413 new_checkpoint = {}414 415 new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"]416 new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"]417 new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"]418 new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"]419 new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"]420 new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"]421 422 new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"]423 new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"]424 new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"]425 new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"]426 new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"]427 new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"]428 429 new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"]430 new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"]431 new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"]432 new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"]433 434 # Retrieves the keys for the encoder down blocks only435 num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer})436 down_blocks = {layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks)}437 438 # Retrieves the keys for the decoder up blocks only439 num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer})440 up_blocks = {layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks)}441 442 for i in range(num_down_blocks):443 resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key]444 445 if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict:446 new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop(447 f"encoder.down.{i}.downsample.conv.weight"448 )449 new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop(450 f"encoder.down.{i}.downsample.conv.bias"451 )452 453 paths = renew_vae_resnet_paths(resnets)454 meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"}455 assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)456 457 mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key]458 num_mid_res_blocks = 2459 for i in range(1, num_mid_res_blocks + 1):460 resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key]461 462 paths = renew_vae_resnet_paths(resnets)463 meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}464 assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)465 466 mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key]467 paths = renew_vae_attention_paths(mid_attentions)468 meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}469 assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)470 conv_attn_to_linear(new_checkpoint)471 472 for i in range(num_up_blocks):473 block_id = num_up_blocks - 1 - i474 resnets = [key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key]475 476 if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict:477 new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[478 f"decoder.up.{block_id}.upsample.conv.weight"479 ]480 new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[481 f"decoder.up.{block_id}.upsample.conv.bias"482 ]483 484 paths = renew_vae_resnet_paths(resnets)485 meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"}486 assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)487 488 mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key]489 num_mid_res_blocks = 2490 for i in range(1, num_mid_res_blocks + 1):491 resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key]492 493 paths = renew_vae_resnet_paths(resnets)494 meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}495 assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)496 497 mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key]498 paths = renew_vae_attention_paths(mid_attentions)499 meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}500 assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)501 conv_attn_to_linear(new_checkpoint)502 return new_checkpoint503 504 505def create_unet_diffusers_config(v2, use_linear_projection_in_v2=False):506 """507 Creates a config for the diffusers based on the config of the LDM model.508 """509 # unet_params = original_config.model.params.unet_config.params510 511 block_out_channels = [UNET_PARAMS_MODEL_CHANNELS * mult for mult in UNET_PARAMS_CHANNEL_MULT]512 513 down_block_types = []514 resolution = 1515 for i in range(len(block_out_channels)):516 block_type = "CrossAttnDownBlock2D" if resolution in UNET_PARAMS_ATTENTION_RESOLUTIONS else "DownBlock2D"517 down_block_types.append(block_type)518 if i != len(block_out_channels) - 1:519 resolution *= 2520 521 up_block_types = []522 for i in range(len(block_out_channels)):523 block_type = "CrossAttnUpBlock2D" if resolution in UNET_PARAMS_ATTENTION_RESOLUTIONS else "UpBlock2D"524 up_block_types.append(block_type)525 resolution //= 2526 527 config = dict(528 sample_size=UNET_PARAMS_IMAGE_SIZE,529 in_channels=UNET_PARAMS_IN_CHANNELS,530 out_channels=UNET_PARAMS_OUT_CHANNELS,531 down_block_types=tuple(down_block_types),532 up_block_types=tuple(up_block_types),533 block_out_channels=tuple(block_out_channels),534 layers_per_block=UNET_PARAMS_NUM_RES_BLOCKS,535 cross_attention_dim=UNET_PARAMS_CONTEXT_DIM if not v2 else V2_UNET_PARAMS_CONTEXT_DIM,536 attention_head_dim=UNET_PARAMS_NUM_HEADS if not v2 else V2_UNET_PARAMS_ATTENTION_HEAD_DIM,537 # use_linear_projection=UNET_PARAMS_USE_LINEAR_PROJECTION if not v2 else V2_UNET_PARAMS_USE_LINEAR_PROJECTION,538 )539 if v2 and use_linear_projection_in_v2:540 config["use_linear_projection"] = True541 542 return config543 544 545def create_vae_diffusers_config():546 """547 Creates a config for the diffusers based on the config of the LDM model.548 """549 # vae_params = original_config.model.params.first_stage_config.params.ddconfig550 # _ = original_config.model.params.first_stage_config.params.embed_dim551 block_out_channels = [VAE_PARAMS_CH * mult for mult in VAE_PARAMS_CH_MULT]552 down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels)553 up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels)554 555 config = dict(556 sample_size=VAE_PARAMS_RESOLUTION,557 in_channels=VAE_PARAMS_IN_CHANNELS,558 out_channels=VAE_PARAMS_OUT_CH,559 down_block_types=tuple(down_block_types),560 up_block_types=tuple(up_block_types),561 block_out_channels=tuple(block_out_channels),562 latent_channels=VAE_PARAMS_Z_CHANNELS,563 layers_per_block=VAE_PARAMS_NUM_RES_BLOCKS,564 )565 return config566 567 568def convert_ldm_clip_checkpoint_v1(checkpoint):569 keys = list(checkpoint.keys())570 text_model_dict = {}571 for key in keys:572 if key.startswith("cond_stage_model.transformer"):573 text_model_dict[key[len("cond_stage_model.transformer.") :]] = checkpoint[key]574 575 # remove position_ids for newer transformer, which causes error :(576 if "text_model.embeddings.position_ids" in text_model_dict:577 text_model_dict.pop("text_model.embeddings.position_ids")578 579 return text_model_dict580 581 582def convert_ldm_clip_checkpoint_v2(checkpoint, max_length):583 # 嫌になるくらい違うぞ！584 def convert_key(key):585 if not key.startswith("cond_stage_model"):586 return None587 588 # common conversion589 key = key.replace("cond_stage_model.model.transformer.", "text_model.encoder.")590 key = key.replace("cond_stage_model.model.", "text_model.")591 592 if "resblocks" in key:593 # resblocks conversion594 key = key.replace(".resblocks.", ".layers.")595 if ".ln_" in key:596 key = key.replace(".ln_", ".layer_norm")597 elif ".mlp." in key:598 key = key.replace(".c_fc.", ".fc1.")599 key = key.replace(".c_proj.", ".fc2.")600 elif ".attn.out_proj" in key:601 key = key.replace(".attn.out_proj.", ".self_attn.out_proj.")602 elif ".attn.in_proj" in key:603 key = None # 特殊なので後で処理する604 else:605 raise ValueError(f"unexpected key in SD: {key}")606 elif ".positional_embedding" in key:607 key = key.replace(".positional_embedding", ".embeddings.position_embedding.weight")608 elif ".text_projection" in key:609 key = None # 使われない???610 elif ".logit_scale" in key:611 key = None # 使われない???612 elif ".token_embedding" in key:613 key = key.replace(".token_embedding.weight", ".embeddings.token_embedding.weight")614 elif ".ln_final" in key:615 key = key.replace(".ln_final", ".final_layer_norm")616 return key617 618 keys = list(checkpoint.keys())619 new_sd = {}620 for key in keys:621 # remove resblocks 23622 if ".resblocks.23." in key:623 continue624 new_key = convert_key(key)625 if new_key is None:626 continue627 new_sd[new_key] = checkpoint[key]628 629 # attnの変換630 for key in keys:631 if ".resblocks.23." in key:632 continue633 if ".resblocks" in key and ".attn.in_proj_" in key:634 # 三つに分割635 values = torch.chunk(checkpoint[key], 3)636 637 key_suffix = ".weight" if "weight" in key else ".bias"638 key_pfx = key.replace("cond_stage_model.model.transformer.resblocks.", "text_model.encoder.layers.")639 key_pfx = key_pfx.replace("_weight", "")640 key_pfx = key_pfx.replace("_bias", "")641 key_pfx = key_pfx.replace(".attn.in_proj", ".self_attn.")642 new_sd[key_pfx + "q_proj" + key_suffix] = values[0]643 new_sd[key_pfx + "k_proj" + key_suffix] = values[1]644 new_sd[key_pfx + "v_proj" + key_suffix] = values[2]645 646 # remove position_ids for newer transformer, which causes error :(647 ANOTHER_POSITION_IDS_KEY = "text_model.encoder.text_model.embeddings.position_ids"648 if ANOTHER_POSITION_IDS_KEY in new_sd:649 # waifu diffusion v1.4650 del new_sd[ANOTHER_POSITION_IDS_KEY]651 652 if "text_model.embeddings.position_ids" in new_sd:653 del new_sd["text_model.embeddings.position_ids"]654 655 return new_sd656 657 658# endregion659 660 661# region Diffusers->StableDiffusion の変換コード662# convert_diffusers_to_original_stable_diffusion をコピーして修正している（ASL 2.0）663 664 665def conv_transformer_to_linear(checkpoint):666 keys = list(checkpoint.keys())667 tf_keys = ["proj_in.weight", "proj_out.weight"]668 for key in keys:669 if ".".join(key.split(".")[-2:]) in tf_keys:670 if checkpoint[key].ndim > 2:671 checkpoint[key] = checkpoint[key][:, :, 0, 0]672 673 674def convert_unet_state_dict_to_sd(v2, unet_state_dict):675 unet_conversion_map = [676 # (stable-diffusion, HF Diffusers)677 ("time_embed.0.weight", "time_embedding.linear_1.weight"),678 ("time_embed.0.bias", "time_embedding.linear_1.bias"),679 ("time_embed.2.weight", "time_embedding.linear_2.weight"),680 ("time_embed.2.bias", "time_embedding.linear_2.bias"),681 ("input_blocks.0.0.weight", "conv_in.weight"),682 ("input_blocks.0.0.bias", "conv_in.bias"),683 ("out.0.weight", "conv_norm_out.weight"),684 ("out.0.bias", "conv_norm_out.bias"),685 ("out.2.weight", "conv_out.weight"),686 ("out.2.bias", "conv_out.bias"),687 ]688 689 unet_conversion_map_resnet = [690 # (stable-diffusion, HF Diffusers)691 ("in_layers.0", "norm1"),692 ("in_layers.2", "conv1"),693 ("out_layers.0", "norm2"),694 ("out_layers.3", "conv2"),695 ("emb_layers.1", "time_emb_proj"),696 ("skip_connection", "conv_shortcut"),697 ]698 699 unet_conversion_map_layer = []700 for i in range(4):701 # loop over downblocks/upblocks702 703 for j in range(2):704 # loop over resnets/attentions for downblocks705 hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."706 sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."707 unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))708 709 if i < 3:710 # no attention layers in down_blocks.3711 hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."712 sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."713 unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))714 715 for j in range(3):716 # loop over resnets/attentions for upblocks717 hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."718 sd_up_res_prefix = f"output_blocks.{3*i + j}.0."719 unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))720 721 if i > 0:722 # no attention layers in up_blocks.0723 hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."724 sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."725 unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))726 727 if i < 3:728 # no downsample in down_blocks.3729 hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."730 sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."731 unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))732 733 # no upsample in up_blocks.3734 hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."735 sd_upsample_prefix = f"output_blocks.{3*i + 2}.{1 if i == 0 else 2}."736 unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))737 738 hf_mid_atn_prefix = "mid_block.attentions.0."739 sd_mid_atn_prefix = "middle_block.1."740 unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))741 742 for j in range(2):743 hf_mid_res_prefix = f"mid_block.resnets.{j}."744 sd_mid_res_prefix = f"middle_block.{2*j}."745 unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))746 747 # buyer beware: this is a *brittle* function,748 # and correct output requires that all of these pieces interact in749 # the exact order in which I have arranged them.750 mapping = {k: k for k in unet_state_dict.keys()}751 for sd_name, hf_name in unet_conversion_map:752 mapping[hf_name] = sd_name753 for k, v in mapping.items():754 if "resnets" in k:755 for sd_part, hf_part in unet_conversion_map_resnet:756 v = v.replace(hf_part, sd_part)757 mapping[k] = v758 for k, v in mapping.items():759 for sd_part, hf_part in unet_conversion_map_layer:760 v = v.replace(hf_part, sd_part)761 mapping[k] = v762 new_state_dict = {v: unet_state_dict[k] for k, v in mapping.items()}763 764 if v2:765 conv_transformer_to_linear(new_state_dict)766 767 return new_state_dict768 769 770def controlnet_conversion_map():771 unet_conversion_map = [772 ("time_embed.0.weight", "time_embedding.linear_1.weight"),773 ("time_embed.0.bias", "time_embedding.linear_1.bias"),774 ("time_embed.2.weight", "time_embedding.linear_2.weight"),775 ("time_embed.2.bias", "time_embedding.linear_2.bias"),776 ("input_blocks.0.0.weight", "conv_in.weight"),777 ("input_blocks.0.0.bias", "conv_in.bias"),778 ("middle_block_out.0.weight", "controlnet_mid_block.weight"),779 ("middle_block_out.0.bias", "controlnet_mid_block.bias"),780 ]781 782 unet_conversion_map_resnet = [783 ("in_layers.0", "norm1"),784 ("in_layers.2", "conv1"),785 ("out_layers.0", "norm2"),786 ("out_layers.3", "conv2"),787 ("emb_layers.1", "time_emb_proj"),788 ("skip_connection", "conv_shortcut"),789 ]790 791 unet_conversion_map_layer = []792 for i in range(4):793 for j in range(2):794 hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."795 sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."796 unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))797 798 if i < 3:799 hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."800 sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."801 unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))802 803 if i < 3:804 hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."805 sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."806 unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))807 808 hf_mid_atn_prefix = "mid_block.attentions.0."809 sd_mid_atn_prefix = "middle_block.1."810 unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))811 812 for j in range(2):813 hf_mid_res_prefix = f"mid_block.resnets.{j}."814 sd_mid_res_prefix = f"middle_block.{2*j}."815 unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))816 817 controlnet_cond_embedding_names = ["conv_in"] + [f"blocks.{i}" for i in range(6)] + ["conv_out"]818 for i, hf_prefix in enumerate(controlnet_cond_embedding_names):819 hf_prefix = f"controlnet_cond_embedding.{hf_prefix}."820 sd_prefix = f"input_hint_block.{i*2}."821 unet_conversion_map_layer.append((sd_prefix, hf_prefix))822 823 for i in range(12):824 hf_prefix = f"controlnet_down_blocks.{i}."825 sd_prefix = f"zero_convs.{i}.0."826 unet_conversion_map_layer.append((sd_prefix, hf_prefix))827 828 return unet_conversion_map, unet_conversion_map_resnet, unet_conversion_map_layer829 830 831def convert_controlnet_state_dict_to_sd(controlnet_state_dict):832 unet_conversion_map, unet_conversion_map_resnet, unet_conversion_map_layer = controlnet_conversion_map()833 834 mapping = {k: k for k in controlnet_state_dict.keys()}835 for sd_name, diffusers_name in unet_conversion_map:836 mapping[diffusers_name] = sd_name837 for k, v in mapping.items():838 if "resnets" in k:839 for sd_part, diffusers_part in unet_conversion_map_resnet:840 v = v.replace(diffusers_part, sd_part)841 mapping[k] = v842 for k, v in mapping.items():843 for sd_part, diffusers_part in unet_conversion_map_layer:844 v = v.replace(diffusers_part, sd_part)845 mapping[k] = v846 new_state_dict = {v: controlnet_state_dict[k] for k, v in mapping.items()}847 return new_state_dict848 849 850def convert_controlnet_state_dict_to_diffusers(controlnet_state_dict):851 unet_conversion_map, unet_conversion_map_resnet, unet_conversion_map_layer = controlnet_conversion_map()852 853 mapping = {k: k for k in controlnet_state_dict.keys()}854 for sd_name, diffusers_name in unet_conversion_map:855 mapping[sd_name] = diffusers_name856 for k, v in mapping.items():857 for sd_part, diffusers_part in unet_conversion_map_layer:858 v = v.replace(sd_part, diffusers_part)859 mapping[k] = v860 for k, v in mapping.items():861 if "resnets" in v:862 for sd_part, diffusers_part in unet_conversion_map_resnet:863 v = v.replace(sd_part, diffusers_part)864 mapping[k] = v865 new_state_dict = {v: controlnet_state_dict[k] for k, v in mapping.items()}866 return new_state_dict867 868 869# ================#870# VAE Conversion #871# ================#872 873 874def reshape_weight_for_sd(w):875 # convert HF linear weights to SD conv2d weights876 return w.reshape(*w.shape, 1, 1)877 878 879def convert_vae_state_dict(vae_state_dict):880 vae_conversion_map = [881 # (stable-diffusion, HF Diffusers)882 ("nin_shortcut", "conv_shortcut"),883 ("norm_out", "conv_norm_out"),884 ("mid.attn_1.", "mid_block.attentions.0."),885 ]886 887 for i in range(4):888 # down_blocks have two resnets889 for j in range(2):890 hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."891 sd_down_prefix = f"encoder.down.{i}.block.{j}."892 vae_conversion_map.append((sd_down_prefix, hf_down_prefix))893 894 if i < 3:895 hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."896 sd_downsample_prefix = f"down.{i}.downsample."897 vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))898 899 hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."900 sd_upsample_prefix = f"up.{3-i}.upsample."901 vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))902 903 # up_blocks have three resnets904 # also, up blocks in hf are numbered in reverse from sd905 for j in range(3):906 hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."907 sd_up_prefix = f"decoder.up.{3-i}.block.{j}."908 vae_conversion_map.append((sd_up_prefix, hf_up_prefix))909 910 # this part accounts for mid blocks in both the encoder and the decoder911 for i in range(2):912 hf_mid_res_prefix = f"mid_block.resnets.{i}."913 sd_mid_res_prefix = f"mid.block_{i+1}."914 vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))915 916 if diffusers.__version__ < "0.17.0":917 vae_conversion_map_attn = [918 # (stable-diffusion, HF Diffusers)919 ("norm.", "group_norm."),920 ("q.", "query."),921 ("k.", "key."),922 ("v.", "value."),923 ("proj_out.", "proj_attn."),924 ]925 else:926 vae_conversion_map_attn = [927 # (stable-diffusion, HF Diffusers)928 ("norm.", "group_norm."),929 ("q.", "to_q."),930 ("k.", "to_k."),931 ("v.", "to_v."),932 ("proj_out.", "to_out.0."),933 ]934 935 mapping = {k: k for k in vae_state_dict.keys()}936 for k, v in mapping.items():937 for sd_part, hf_part in vae_conversion_map:938 v = v.replace(hf_part, sd_part)939 mapping[k] = v940 for k, v in mapping.items():941 if "attentions" in k:942 for sd_part, hf_part in vae_conversion_map_attn:943 v = v.replace(hf_part, sd_part)944 mapping[k] = v945 new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}946 weights_to_convert = ["q", "k", "v", "proj_out"]947 for k, v in new_state_dict.items():948 for weight_name in weights_to_convert:949 if f"mid.attn_1.{weight_name}.weight" in k:950 # logger.info(f"Reshaping {k} for SD format: shape {v.shape} -> {v.shape} x 1 x 1")951 new_state_dict[k] = reshape_weight_for_sd(v)952 953 return new_state_dict954 955 956# endregion957 958# region 自作のモデル読み書きなど959 960 961def is_safetensors(path):962 return os.path.splitext(path)[1].lower() == ".safetensors"963 964 965def load_checkpoint_with_text_encoder_conversion(ckpt_path, device="cpu"):966 # text encoderの格納形式が違うモデルに対応する ('text_model'がない)967 TEXT_ENCODER_KEY_REPLACEMENTS = [968 ("cond_stage_model.transformer.embeddings.", "cond_stage_model.transformer.text_model.embeddings."),969 ("cond_stage_model.transformer.encoder.", "cond_stage_model.transformer.text_model.encoder."),970 ("cond_stage_model.transformer.final_layer_norm.", "cond_stage_model.transformer.text_model.final_layer_norm."),971 ]972 973 if is_safetensors(ckpt_path):974 checkpoint = None975 state_dict = load_file(ckpt_path) # , device) # may causes error976 else:977 checkpoint = torch.load(ckpt_path, map_location=device)978 if "state_dict" in checkpoint:979 state_dict = checkpoint["state_dict"]980 else:981 state_dict = checkpoint982 checkpoint = None983 984 key_reps = []985 for rep_from, rep_to in TEXT_ENCODER_KEY_REPLACEMENTS:986 for key in state_dict.keys():987 if key.startswith(rep_from):988 new_key = rep_to + key[len(rep_from) :]989 key_reps.append((key, new_key))990 991 for key, new_key in key_reps:992 state_dict[new_key] = state_dict[key]993 del state_dict[key]994 995 return checkpoint, state_dict996 997 998# TODO dtype指定の動作が怪しいので確認する text_encoderを指定形式で作れるか未確認999def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, device="cpu", dtype=None, unet_use_linear_projection_in_v2=True):1000 _, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path, device)1001 1002 # Convert the UNet2DConditionModel model.1003 unet_config = create_unet_diffusers_config(v2, unet_use_linear_projection_in_v2)1004 converted_unet_checkpoint = convert_ldm_unet_checkpoint(v2, state_dict, unet_config)1005 1006 unet = UNet2DConditionModel(**unet_config).to(device)1007 info = unet.load_state_dict(converted_unet_checkpoint)1008 logger.info(f"loading u-net: {info}")1009 1010 # Convert the VAE model.1011 vae_config = create_vae_diffusers_config()1012 converted_vae_checkpoint = convert_ldm_vae_checkpoint(state_dict, vae_config)1013 1014 vae = AutoencoderKL(**vae_config).to(device)1015 info = vae.load_state_dict(converted_vae_checkpoint)1016 logger.info(f"loading vae: {info}")1017 1018 # convert text_model1019 if v2:1020 converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v2(state_dict, 77)1021 cfg = CLIPTextConfig(1022 vocab_size=49408,1023 hidden_size=1024,1024 intermediate_size=4096,1025 num_hidden_layers=23,1026 num_attention_heads=16,1027 max_position_embeddings=77,1028 hidden_act="gelu",1029 layer_norm_eps=1e-05,1030 dropout=0.0,1031 attention_dropout=0.0,1032 initializer_range=0.02,1033 initializer_factor=1.0,1034 pad_token_id=1,1035 bos_token_id=0,1036 eos_token_id=2,1037 model_type="clip_text_model",1038 projection_dim=512,1039 torch_dtype="float32",1040 transformers_version="4.25.0.dev0",1041 )1042 text_model = CLIPTextModel._from_config(cfg)1043 info = text_model.load_state_dict(converted_text_encoder_checkpoint)1044 else:1045 converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v1(state_dict)1046 1047 # logging.set_verbosity_error() # don't show annoying warning1048 # text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to(device)1049 # logging.set_verbosity_warning()1050 # logger.info(f"config: {text_model.config}")1051 cfg = CLIPTextConfig(1052 vocab_size=49408,1053 hidden_size=768,1054 intermediate_size=3072,1055 num_hidden_layers=12,1056 num_attention_heads=12,1057 max_position_embeddings=77,1058 hidden_act="quick_gelu",1059 layer_norm_eps=1e-05,1060 dropout=0.0,1061 attention_dropout=0.0,1062 initializer_range=0.02,1063 initializer_factor=1.0,1064 pad_token_id=1,1065 bos_token_id=0,1066 eos_token_id=2,1067 model_type="clip_text_model",1068 projection_dim=768,1069 torch_dtype="float32",1070 )1071 text_model = CLIPTextModel._from_config(cfg)1072 info = text_model.load_state_dict(converted_text_encoder_checkpoint)1073 logger.info(f"loading text encoder: {info}")1074 1075 return text_model, vae, unet1076 1077 1078def get_model_version_str_for_sd1_sd2(v2, v_parameterization):1079 # only for reference1080 version_str = "sd"1081 if v2:1082 version_str += "_v2"1083 else:1084 version_str += "_v1"1085 if v_parameterization:1086 version_str += "_v"1087 return version_str1088 1089 1090def convert_text_encoder_state_dict_to_sd_v2(checkpoint, make_dummy_weights=False):1091 def convert_key(key):1092 # position_idsの除去1093 if ".position_ids" in key:1094 return None1095 1096 # common1097 key = key.replace("text_model.encoder.", "transformer.")1098 key = key.replace("text_model.", "")1099 if "layers" in key:1100 # resblocks conversion1101 key = key.replace(".layers.", ".resblocks.")1102 if ".layer_norm" in key:1103 key = key.replace(".layer_norm", ".ln_")1104 elif ".mlp." in key:1105 key = key.replace(".fc1.", ".c_fc.")1106 key = key.replace(".fc2.", ".c_proj.")1107 elif ".self_attn.out_proj" in key:1108 key = key.replace(".self_attn.out_proj.", ".attn.out_proj.")1109 elif ".self_attn." in key:1110 key = None # 特殊なので後で処理する1111 else:1112 raise ValueError(f"unexpected key in DiffUsers model: {key}")1113 elif ".position_embedding" in key:1114 key = key.replace("embeddings.position_embedding.weight", "positional_embedding")1115 elif ".token_embedding" in key:1116 key = key.replace("embeddings.token_embedding.weight", "token_embedding.weight")1117 elif "final_layer_norm" in key:1118 key = key.replace("final_layer_norm", "ln_final")1119 return key1120 1121 keys = list(checkpoint.keys())1122 new_sd = {}1123 for key in keys:1124 new_key = convert_key(key)1125 if new_key is None:1126 continue1127 new_sd[new_key] = checkpoint[key]1128 1129 # attnの変換1130 for key in keys:1131 if "layers" in key and "q_proj" in key:1132 # 三つを結合1133 key_q = key1134 key_k = key.replace("q_proj", "k_proj")1135 key_v = key.replace("q_proj", "v_proj")1136 1137 value_q = checkpoint[key_q]1138 value_k = checkpoint[key_k]1139 value_v = checkpoint[key_v]1140 value = torch.cat([value_q, value_k, value_v])1141 1142 new_key = key.replace("text_model.encoder.layers.", "transformer.resblocks.")1143 new_key = new_key.replace(".self_attn.q_proj.", ".attn.in_proj_")1144 new_sd[new_key] = value1145 1146 # 最後の層などを捏造するか1147 if make_dummy_weights:1148 logger.info("make dummy weights for resblock.23, text_projection and logit scale.")1149 keys = list(new_sd.keys())1150 for key in keys:1151 if key.startswith("transformer.resblocks.22."):1152 new_sd[key.replace(".22.", ".23.")] = new_sd[key].clone() # copyしないとsafetensorsの保存で落ちる1153 1154 # Diffusersに含まれない重みを作っておく1155 new_sd["text_projection"] = torch.ones((1024, 1024), dtype=new_sd[keys[0]].dtype, device=new_sd[keys[0]].device)1156 new_sd["logit_scale"] = torch.tensor(1)1157 1158 return new_sd1159 1160 1161def save_stable_diffusion_checkpoint(1162 v2, output_file, text_encoder, unet, ckpt_path, epochs, steps, metadata, save_dtype=None, vae=None1163):1164 if ckpt_path is not None:1165 # epoch/stepを参照する。またVAEがメモリ上にないときなど、もう一度VAEを含めて読み込む1166 checkpoint, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path)1167 if checkpoint is None: # safetensors または state_dictのckpt1168 checkpoint = {}1169 strict = False1170 else:1171 strict = True1172 if "state_dict" in state_dict:1173 del state_dict["state_dict"]1174 else:1175 # 新しく作る1176 assert vae is not None, "VAE is required to save a checkpoint without a given checkpoint"1177 checkpoint = {}1178 state_dict = {}1179 strict = False1180 1181 def update_sd(prefix, sd):1182 for k, v in sd.items():1183 key = prefix + k1184 assert not strict or key in state_dict, f"Illegal key in save SD: {key}"1185 if save_dtype is not None:1186 v = v.detach().clone().to("cpu").to(save_dtype)1187 state_dict[key] = v1188 1189 # Convert the UNet model1190 unet_state_dict = convert_unet_state_dict_to_sd(v2, unet.state_dict())1191 update_sd("model.diffusion_model.", unet_state_dict)1192 1193 # Convert the text encoder model1194 if v2:1195 make_dummy = ckpt_path is None # 参照元のcheckpointがない場合は最後の層を前の層から複製して作るなどダミーの重みを入れる1196 text_enc_dict = convert_text_encoder_state_dict_to_sd_v2(text_encoder.state_dict(), make_dummy)1197 update_sd("cond_stage_model.model.", text_enc_dict)1198 else:1199 text_enc_dict = text_encoder.state_dict()1200 update_sd("cond_stage_model.transformer.", text_enc_dict)