




# Note: if you intend to run this script make sure you look under scripts/fsmt/
# to locate the appropriate script to do the work correctly. There is a set of scripts to:
# - download and prepare data and run the conversion script
# - perform eval to get the best hparam into the config
# - generate model_cards - useful if you have multiple models from the same paper

import argparse
import json
import os
import re
from collections import OrderedDict
from os.path import basename, dirname

import fairseq
import torch
from fairseq import hub_utils
from fairseq.data.dictionary import Dictionary

from transformers import FSMTConfig, FSMTForConditionalGeneration
from transformers.models.fsmt.tokenization_fsmt import VOCAB_FILES_NAMES
from transformers.tokenization_utils_base import TOKENIZER_CONFIG_FILE
from transformers.utils import WEIGHTS_NAME, logging


logging.set_verbosity_warning()

lowercase__				:     Optional[Any] 		=						2

# based on the results of a search on a range of `num_beams`, `length_penalty` and `early_stopping`
# values against wmt19 test data to obtain the best BLEU scores, we will use the following defaults:
#
# * `num_beams`: 5 (higher scores better, but requires more memory/is slower, can be adjusted by users)
# * `early_stopping`: `False` consistently scored better
# * `length_penalty` varied, so will assign the best one depending on the model
lowercase__				:     Union[str, Any] 		=						{
    # fairseq:
    "wmt19-ru-en": {"length_penalty": 1.1},
    "wmt19-en-ru": {"length_penalty": 1.15},
    "wmt19-en-de": {"length_penalty": 1.0},
    "wmt19-de-en": {"length_penalty": 1.1},
    # allenai:
    "wmt16-en-de-dist-12-1": {"length_penalty": 0.6},
    "wmt16-en-de-dist-6-1": {"length_penalty": 0.6},
    "wmt16-en-de-12-1": {"length_penalty": 0.8},
    "wmt19-de-en-6-6-base": {"length_penalty": 0.6},
    "wmt19-de-en-6-6-big": {"length_penalty": 0.6},
}

# this remaps the different models to their organization names
lowercase__				:     Optional[Any] 		=						{}
for m in ["wmt19-ru-en", "wmt19-en-ru", "wmt19-en-de", "wmt19-de-en"]:
		lowercase__				:     Optional[int] 		=						"facebook"
for m in [
    "wmt16-en-de-dist-12-1",
    "wmt16-en-de-dist-6-1",
    "wmt16-en-de-12-1",
    "wmt19-de-en-6-6-base",
    "wmt19-de-en-6-6-big",
]:
		lowercase__				:     List[str] 		=						"allenai"


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	snake_case_		    =     dict((re.sub(R"@@$"     ,	""     ,	_A							), v) if k.endswith("@@"							) else (re.sub(R"$"     ,	"</w>"     ,	_A							), v) for k, v in d.items()							)
	snake_case_		    =     "<s> <pad> </s> <unk>".split()
	# restore the special tokens
	for k in keep_keys:
		del da[f"{k}</w>"]
		snake_case_		    =     d[k]  # restore
	return da


def 						lowerCamelCase__    (		_A     ,	_A							):







	'''simple docstring'''





	assert os.path.exists(_A							)
	os.makedirs(_A     ,	exist_ok=_A							)
	print(f"Writing results to {pytorch_dump_folder_path}"							)

	# handle various types of models

	snake_case_		    =     basename(_A							)
	snake_case_		    =     dirname(_A							)

	snake_case_		    =     fairseq.model_parallel.models.transformer.ModelParallelTransformerModel
	snake_case_		    =     cls.hub_models()
	snake_case_		    =     {"bpe": "fastbpe", "tokenizer": "moses"}
	snake_case_		    =     "."
	# note: since the model dump is old, fairseq has upgraded its model some
	# time later, and it does a whole lot of rewrites and splits on the saved
	# weights, therefore we can't use torch.load() directly on the model file.
	# see: upgrade_state_dict(state_dict) in fairseq_model.py
	print(f"using checkpoint {checkpoint_file}"							)
	snake_case_		    =     hub_utils.from_pretrained(
	    _A     ,	_A     ,	_A     ,	archive_map=_A     ,	**_A							)

	snake_case_		    =     vars(chkpt["args"]["model"]							)

	snake_case_		    =     args["source_lang"]
	snake_case_		    =     args["target_lang"]

	snake_case_		    =     dirname(_A							)
	snake_case_		    =     basename(_A							)

	# dicts
	snake_case_		    =     os.path.join(_A     ,	f"dict.{src_lang}.txt"							)
	snake_case_		    =     os.path.join(_A     ,	f"dict.{tgt_lang}.txt"							)

	snake_case_		    =     Dictionary.load(_A							)
	snake_case_		    =     rewrite_dict_keys(src_dict.indices							)
	snake_case_		    =     len(_A							)
	snake_case_		    =     os.path.join(_A     ,	"vocab-src.json"							)
	print(f"Generating {src_vocab_file} of {src_vocab_size} of {src_lang} records"							)
	with open(_A     ,	"w"     ,	encoding="utf-8"							) as f:
		f.write(json.dumps(_A     ,	ensure_ascii=_A     ,	indent=_A							)							)

	# detect whether this is a do_lower_case situation, which can be derived by checking whether we
	# have at least one uppercase letter in the source vocab
	snake_case_		    =     True
	for k in src_vocab.keys():
		if not k.islower():
			snake_case_		    =     False
			break

	snake_case_		    =     Dictionary.load(_A							)
	snake_case_		    =     rewrite_dict_keys(tgt_dict.indices							)
	snake_case_		    =     len(_A							)
	snake_case_		    =     os.path.join(_A     ,	"vocab-tgt.json"							)
	print(f"Generating {tgt_vocab_file} of {tgt_vocab_size} of {tgt_lang} records"							)
	with open(_A     ,	"w"     ,	encoding="utf-8"							) as f:
		f.write(json.dumps(_A     ,	ensure_ascii=_A     ,	indent=_A							)							)

	# merges_file (bpecodes)
	snake_case_		    =     os.path.join(_A     ,	VOCAB_FILES_NAMES["merges_file"]							)
	for fn in ["bpecodes", "code"]:  # older fairseq called the merges file "code"
		snake_case_		    =     os.path.join(_A     ,	_A							)
		if os.path.exists(_A							):
			break
	with open(_A     ,	encoding="utf-8"							) as fin:
		snake_case_		    =     fin.read()
	snake_case_		    =     re.sub(R" \d+$"     ,	""     ,	_A     ,	0     ,	re.M							)  # remove frequency number
	print(f"Generating {merges_file}"							)
	with open(_A     ,	"w"     ,	encoding="utf-8"							) as fout:
		fout.write(_A							)

	# model config
	snake_case_		    =     os.path.join(_A     ,	"config.json"							)

	# validate bpe/tokenizer config, as currently it's hardcoded to moses+fastbpe -
	# may have to modify the tokenizer if a different type is used by a future model
	assert args["bpe"] == "fastbpe", f"need to extend tokenizer to support bpe={args['bpe']}"
	assert args["tokenizer"] == "moses", f"need to extend tokenizer to support bpe={args['tokenizer']}"

	snake_case_		    =     {
	    "architectures": ["FSMTForConditionalGeneration"],
	    "model_type": "fsmt",
	    "activation_dropout": args["activation_dropout"],
	    "activation_function": "relu",
	    "attention_dropout": args["attention_dropout"],
	    "d_model": args["decoder_embed_dim"],
	    "dropout": args["dropout"],
	    "init_std": 0.02,
	    "max_position_embeddings": args["max_source_positions"],
	    "num_hidden_layers": args["encoder_layers"],
	    "src_vocab_size": src_vocab_size,
	    "tgt_vocab_size": tgt_vocab_size,
	    "langs": [src_lang, tgt_lang],
	    "encoder_attention_heads": args["encoder_attention_heads"],
	    "encoder_ffn_dim": args["encoder_ffn_embed_dim"],
	    "encoder_layerdrop": args["encoder_layerdrop"],
	    "encoder_layers": args["encoder_layers"],
	    "decoder_attention_heads": args["decoder_attention_heads"],
	    "decoder_ffn_dim": args["decoder_ffn_embed_dim"],
	    "decoder_layerdrop": args["decoder_layerdrop"],
	    "decoder_layers": args["decoder_layers"],
	    "bos_token_id": 0,
	    "pad_token_id": 1,
	    "eos_token_id": 2,
	    "is_encoder_decoder": True,
	    "scale_embedding": not args["no_scale_embedding"],
	    "tie_word_embeddings": args["share_all_embeddings"],
	}

	# good hparam defaults to start with
	snake_case_		    =     5
	snake_case_		    =     False
	if model_dir in best_score_hparams and "length_penalty" in best_score_hparams[model_dir]:
		snake_case_		    =     best_score_hparams[model_dir]["length_penalty"]
	else:
		snake_case_		    =     1.0

	print(f"Generating {fsmt_model_config_file}"							)
	with open(_A     ,	"w"     ,	encoding="utf-8"							) as f:
		f.write(json.dumps(_A     ,	ensure_ascii=_A     ,	indent=_A							)							)

	# tokenizer config
	snake_case_		    =     os.path.join(_A     ,	_A							)

	snake_case_		    =     {
	    "langs": [src_lang, tgt_lang],
	    "model_max_length": 1024,
	    "do_lower_case": do_lower_case,
	}

	print(f"Generating {fsmt_tokenizer_config_file}"							)
	with open(_A     ,	"w"     ,	encoding="utf-8"							) as f:
		f.write(json.dumps(_A     ,	ensure_ascii=_A     ,	indent=_A							)							)

	# model
	snake_case_		    =     chkpt["models"][0]
	snake_case_		    =     model.state_dict()

	# rename keys to start with 'model.'
	snake_case_		    =     OrderedDict(("model." + k, v) for k, v in model_state_dict.items()							)

	# remove unneeded keys
	snake_case_		    =     [
	    "model.model",
	    "model.encoder.version",
	    "model.decoder.version",
	    "model.encoder_embed_tokens.weight",
	    "model.decoder_embed_tokens.weight",
	    "model.encoder.embed_positions._float_tensor",
	    "model.decoder.embed_positions._float_tensor",
	]
	for k in ignore_keys:
		model_state_dict.pop(_A     ,	_A							)

	snake_case_		    =     FSMTConfig.from_pretrained(_A							)
	snake_case_		    =     FSMTForConditionalGeneration(_A							)

	# check that it loads ok
	model_new.load_state_dict(_A     ,	strict=_A							)

	# save
	snake_case_		    =     os.path.join(_A     ,	_A							)
	print(f"Generating {pytorch_weights_dump_path}"							)
	torch.save(_A     ,	_A							)

	print("Conversion is done!"							)
	print("\nLast step is to upload the files to s3"							)
	print(f"cd {data_root}"							)
	print(f"transformers-cli upload {model_dir}"							)


if __name__ == "__main__":
		lowercase__				:     Optional[int] 		=						argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    "--fsmt_checkpoint_path",
		    default=None,
		    type=str,
		    required=True,
		    help=(
		        "Path to the official PyTorch checkpoint file which is expected to reside in the dump dir with dicts,"
		        " bpecodes, etc."
		    ),
		)
		parser.add_argument(
		    "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
		)
		lowercase__				:     Optional[Any] 		=						parser.parse_args()
		convert_fsmt_checkpoint_to_pytorch(args.fsmt_checkpoint_path, args.pytorch_dump_folder_path)






import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


# =================#
# UNet Conversion #
# =================#

lowercase__				:     Optional[int] 		=						[
    # (stable-diffusion, HF Diffusers)
    ("time_embed.0.weight", "time_embedding.linear_1.weight"),
    ("time_embed.0.bias", "time_embedding.linear_1.bias"),
    ("time_embed.2.weight", "time_embedding.linear_2.weight"),
    ("time_embed.2.bias", "time_embedding.linear_2.bias"),
    ("input_blocks.0.0.weight", "conv_in.weight"),
    ("input_blocks.0.0.bias", "conv_in.bias"),
    ("out.0.weight", "conv_norm_out.weight"),
    ("out.0.bias", "conv_norm_out.bias"),
    ("out.2.weight", "conv_out.weight"),
    ("out.2.bias", "conv_out.bias"),
]

lowercase__				:     List[Any] 		=						[
    # (stable-diffusion, HF Diffusers)
    ("in_layers.0", "norm1"),
    ("in_layers.2", "conv1"),
    ("out_layers.0", "norm2"),
    ("out_layers.3", "conv2"),
    ("emb_layers.1", "time_emb_proj"),
    ("skip_connection", "conv_shortcut"),
]

lowercase__				:     Optional[Any] 		=						[]
# hardcoded number of downblocks and resnets/attentions...
# would need smarter logic for other networks.
for i in range(4):
		# loop over downblocks/upblocks

		for j in range(2):
				# loop over resnets/attentions for downblocks
				lowercase__				:     str 		=						f'''down_blocks.{i}.resnets.{j}.'''
				lowercase__				:     Union[str, Any] 		=						f'''input_blocks.{3*i + j + 1}.0.'''
				unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))

				if i < 3:
						# no attention layers in down_blocks.3
						lowercase__				:     Tuple 		=						f'''down_blocks.{i}.attentions.{j}.'''
						lowercase__				:     Dict 		=						f'''input_blocks.{3*i + j + 1}.1.'''
						unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))

		for j in range(3):
				# loop over resnets/attentions for upblocks
				lowercase__				:     List[Any] 		=						f'''up_blocks.{i}.resnets.{j}.'''
				lowercase__				:     int 		=						f'''output_blocks.{3*i + j}.0.'''
				unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))

				if i > 0:
						# no attention layers in up_blocks.0
						lowercase__				:     List[str] 		=						f'''up_blocks.{i}.attentions.{j}.'''
						lowercase__				:     Tuple 		=						f'''output_blocks.{3*i + j}.1.'''
						unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))

		if i < 3:
				# no downsample in down_blocks.3
				lowercase__				:     List[str] 		=						f'''down_blocks.{i}.downsamplers.0.conv.'''
				lowercase__				:     Any 		=						f'''input_blocks.{3*(i+1)}.0.op.'''
				unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))

				# no upsample in up_blocks.3
				lowercase__				:     Optional[int] 		=						f'''up_blocks.{i}.upsamplers.0.'''
				lowercase__				:     int 		=						f'''output_blocks.{3*i + 2}.{1 if i == 0 else 2}.'''
				unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))

lowercase__				:     Union[str, Any] 		=						"mid_block.attentions.0."
lowercase__				:     List[Any] 		=						"middle_block.1."
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))

for j in range(2):
		lowercase__				:     Tuple 		=						f'''mid_block.resnets.{j}.'''
		lowercase__				:     List[str] 		=						f'''middle_block.{2*j}.'''
		unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	snake_case_		    =     {k: k for k in unet_state_dict.keys()}
	for sd_name, hf_name in unet_conversion_map:
		snake_case_		    =     sd_name
	for k, v in mapping.items():
		if "resnets" in k:
			for sd_part, hf_part in unet_conversion_map_resnet:
				snake_case_		    =     v.replace(_A     ,	_A							)
			snake_case_		    =     v
	for k, v in mapping.items():
		for sd_part, hf_part in unet_conversion_map_layer:
			snake_case_		    =     v.replace(_A     ,	_A							)
		snake_case_		    =     v
	snake_case_		    =     {v: unet_state_dict[k] for k, v in mapping.items()}
	return new_state_dict


# ================#
# VAE Conversion #
# ================#

lowercase__				:     Dict 		=						[
    # (stable-diffusion, HF Diffusers)
    ("nin_shortcut", "conv_shortcut"),
    ("norm_out", "conv_norm_out"),
    ("mid.attn_1.", "mid_block.attentions.0."),
]

for i in range(4):
		# down_blocks have two resnets
		for j in range(2):
				lowercase__				:     Any 		=						f'''encoder.down_blocks.{i}.resnets.{j}.'''
				lowercase__				:     List[Any] 		=						f'''encoder.down.{i}.block.{j}.'''
				vae_conversion_map.append((sd_down_prefix, hf_down_prefix))

		if i < 3:
				lowercase__				:     Optional[int] 		=						f'''down_blocks.{i}.downsamplers.0.'''
				lowercase__				:     Tuple 		=						f'''down.{i}.downsample.'''
				vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))

				lowercase__				:     List[str] 		=						f'''up_blocks.{i}.upsamplers.0.'''
				lowercase__				:     Optional[int] 		=						f'''up.{3-i}.upsample.'''
				vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))

		# up_blocks have three resnets
		# also, up blocks in hf are numbered in reverse from sd
		for j in range(3):
				lowercase__				:     int 		=						f'''decoder.up_blocks.{i}.resnets.{j}.'''
				lowercase__				:     Union[str, Any] 		=						f'''decoder.up.{3-i}.block.{j}.'''
				vae_conversion_map.append((sd_up_prefix, hf_up_prefix))

# this part accounts for mid blocks in both the encoder and the decoder
for i in range(2):
		lowercase__				:     Dict 		=						f'''mid_block.resnets.{i}.'''
		lowercase__				:     int 		=						f'''mid.block_{i+1}.'''
		vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))


lowercase__				:     Optional[Any] 		=						[
    # (stable-diffusion, HF Diffusers)
    ("norm.", "group_norm."),
    ("q.", "query."),
    ("k.", "key."),
    ("v.", "value."),
    ("proj_out.", "proj_attn."),
]


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	return w.reshape(*w.shape     ,	1     ,	1							)


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	snake_case_		    =     {k: k for k in vae_state_dict.keys()}
	for k, v in mapping.items():
		for sd_part, hf_part in vae_conversion_map:
			snake_case_		    =     v.replace(_A     ,	_A							)
		snake_case_		    =     v
	for k, v in mapping.items():
		if "attentions" in k:
			for sd_part, hf_part in vae_conversion_map_attn:
				snake_case_		    =     v.replace(_A     ,	_A							)
			snake_case_		    =     v
	snake_case_		    =     {v: vae_state_dict[k] for k, v in mapping.items()}
	snake_case_		    =     ["q", "k", "v", "proj_out"]
	for k, v in new_state_dict.items():
		for weight_name in weights_to_convert:
			if f"mid.attn_1.{weight_name}.weight" in k:
				print(f"Reshaping {k} for SD format"							)
				snake_case_		    =     reshape_weight_for_sd(_A							)
	return new_state_dict


# =========================#
# Text Encoder Conversion #
# =========================#


lowercase__				:     int 		=						[
    # (stable-diffusion, HF Diffusers)
    ("resblocks.", "text_model.encoder.layers."),
    ("ln_1", "layer_norm1"),
    ("ln_2", "layer_norm2"),
    (".c_fc.", ".fc1."),
    (".c_proj.", ".fc2."),
    (".attn", ".self_attn"),
    ("ln_final.", "transformer.text_model.final_layer_norm."),
    ("token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"),
    ("positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"),
]
lowercase__				:     Dict 		=						{re.escape(x[1]): x[0] for x in textenc_conversion_lst}
lowercase__				:     Tuple 		=						re.compile("|".join(protected.keys()))

# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
lowercase__				:     Dict 		=						{"q": 0, "k": 1, "v": 2}


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	snake_case_		    =     {}
	snake_case_		    =     {}
	snake_case_		    =     {}
	for k, v in text_enc_dict.items():
		if (
		    k.endswith(".self_attn.q_proj.weight"							)
		    or k.endswith(".self_attn.k_proj.weight"							)
		    or k.endswith(".self_attn.v_proj.weight"							)
		):
			snake_case_		    =     k[: -len(".q_proj.weight"							)]
			snake_case_		    =     k[-len("q_proj.weight"							)]
			if k_pre not in capture_qkv_weight:
				snake_case_		    =     [None, None, None]
			snake_case_		    =     v
			continue

		if (
		    k.endswith(".self_attn.q_proj.bias"							)
		    or k.endswith(".self_attn.k_proj.bias"							)
		    or k.endswith(".self_attn.v_proj.bias"							)
		):
			snake_case_		    =     k[: -len(".q_proj.bias"							)]
			snake_case_		    =     k[-len("q_proj.bias"							)]
			if k_pre not in capture_qkv_bias:
				snake_case_		    =     [None, None, None]
			snake_case_		    =     v
			continue

		snake_case_		    =     textenc_pattern.sub(lambda _A							: protected[re.escape(m.group(0							)							)]     ,	_A							)
		snake_case_		    =     v

	for k_pre, tensors in capture_qkv_weight.items():
		if None in tensors:
			raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing"							)
		snake_case_		    =     textenc_pattern.sub(lambda _A							: protected[re.escape(m.group(0							)							)]     ,	_A							)
		snake_case_		    =     torch.cat(_A							)

	for k_pre, tensors in capture_qkv_bias.items():
		if None in tensors:
			raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing"							)
		snake_case_		    =     textenc_pattern.sub(lambda _A							: protected[re.escape(m.group(0							)							)]     ,	_A							)
		snake_case_		    =     torch.cat(_A							)

	return new_state_dict


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	return text_enc_dict


if __name__ == "__main__":
		lowercase__				:     int 		=						argparse.ArgumentParser()

		parser.add_argument("--model_path", default=None, type=str, required=True, help="Path to the model to convert.")
		parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the output model.")
		parser.add_argument("--half", action="store_true", help="Save weights in half precision.")
		parser.add_argument(
		    "--use_safetensors", action="store_true", help="Save weights use safetensors, default is ckpt."
		)

		lowercase__				:     Dict 		=						parser.parse_args()

		assert args.model_path is not None, "Must provide a model path!"

		assert args.checkpoint_path is not None, "Must provide a checkpoint path!"

		# Path for safetensors
		lowercase__				:     Tuple 		=						osp.join(args.model_path, "unet", "diffusion_pytorch_model.safetensors")
		lowercase__				:     int 		=						osp.join(args.model_path, "vae", "diffusion_pytorch_model.safetensors")
		lowercase__				:     Any 		=						osp.join(args.model_path, "text_encoder", "model.safetensors")

		# Load models from safetensors if it exists, if it doesn't pytorch
		if osp.exists(unet_path):
				lowercase__				:     str 		=						load_file(unet_path, device="cpu")
		else:
				lowercase__				:     Optional[Any] 		=						osp.join(args.model_path, "unet", "diffusion_pytorch_model.bin")
				lowercase__				:     Any 		=						torch.load(unet_path, map_location="cpu")

		if osp.exists(vae_path):
				lowercase__				:     Tuple 		=						load_file(vae_path, device="cpu")
		else:
				lowercase__				:     Any 		=						osp.join(args.model_path, "vae", "diffusion_pytorch_model.bin")
				lowercase__				:     Dict 		=						torch.load(vae_path, map_location="cpu")

		if osp.exists(text_enc_path):
				lowercase__				:     Union[str, Any] 		=						load_file(text_enc_path, device="cpu")
		else:
				lowercase__				:     Union[str, Any] 		=						osp.join(args.model_path, "text_encoder", "pytorch_model.bin")
				lowercase__				:     Optional[int] 		=						torch.load(text_enc_path, map_location="cpu")

		# Convert the UNet model
		lowercase__				:     Dict 		=						convert_unet_state_dict(unet_state_dict)
		lowercase__				:     Any 		=						{"model.diffusion_model." + k: v for k, v in unet_state_dict.items()}

		# Convert the VAE model
		lowercase__				:     Dict 		=						convert_vae_state_dict(vae_state_dict)
		lowercase__				:     Union[str, Any] 		=						{"first_stage_model." + k: v for k, v in vae_state_dict.items()}

		# Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper
		lowercase__				:     Optional[Any] 		=						"text_model.encoder.layers.22.layer_norm2.bias" in text_enc_dict

		if is_vaa_model:
				# Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm
				lowercase__				:     Any 		=						{"transformer." + k: v for k, v in text_enc_dict.items()}
				lowercase__				:     List[Any] 		=						convert_text_enc_state_dict_vaa(text_enc_dict)
				lowercase__				:     List[str] 		=						{"cond_stage_model.model." + k: v for k, v in text_enc_dict.items()}
		else:
				lowercase__				:     Tuple 		=						convert_text_enc_state_dict(text_enc_dict)
				lowercase__				:     Any 		=						{"cond_stage_model.transformer." + k: v for k, v in text_enc_dict.items()}

		# Put together new checkpoint
		lowercase__				:     Tuple 		=						{**unet_state_dict, **vae_state_dict, **text_enc_dict}
		if args.half:
				lowercase__				:     Any 		=						{k: v.half() for k, v in state_dict.items()}

		if args.use_safetensors:
				save_file(state_dict, args.checkpoint_path)
		else:
				lowercase__				:     Union[str, Any] 		=						{"state_dict": state_dict}
				torch.save(state_dict, args.checkpoint_path)
