



import unittest

from transformers import AlbertTokenizer, AlbertTokenizerFast
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


lowerCAmelCase__					  =  get_tests_dir('''fixtures/spiece.model''')

@require_sentencepiece
@require_tokenizers
class 				snake_case__(_UpperCamelCase    ,      unittest.TestCase      ):

					"""simple docstring"""






					lowercase_								=		AlbertTokenizer
					lowercase_								=		AlbertTokenizerFast
					lowercase_								=		True
					lowercase_								=		True
					lowercase_								=		True







					def 			snake_case					(							self :     Any					):
									super().setUp()

									# We have a SentencePiece fixture for testing
									lowercase__					:   Tuple													=					AlbertTokenizer(SCREAMING_SNAKE_CASE					)
									tokenizer.save_pretrained(self.tmpdirname					)







					def 			snake_case					(							self :     Union[str, Any]				,       SCREAMING_SNAKE_CASE :     int					):
									lowercase__					:   List[str]													=					"this is a test"
									lowercase__					:   Tuple													=					"this is a test"
									return input_text, output_text







					def 			snake_case					(							self :     List[Any]					):
									lowercase__					:   int													=					"<pad>"
									lowercase__					:   List[Any]													=					0

									self.assertEqual(self.get_tokenizer()._convert_token_to_id(SCREAMING_SNAKE_CASE					)				,       SCREAMING_SNAKE_CASE					)
									self.assertEqual(self.get_tokenizer()._convert_id_to_token(SCREAMING_SNAKE_CASE					)				,       SCREAMING_SNAKE_CASE					)







					def 			snake_case					(							self :     str					):
									lowercase__					:   Tuple													=					list(self.get_tokenizer().get_vocab().keys()					)

									self.assertEqual(vocab_keys[0]				,       "<pad>"					)
									self.assertEqual(vocab_keys[1]				,       "<unk>"					)
									self.assertEqual(vocab_keys[-1]				,       "▁eloquent"					)
									self.assertEqual(len(SCREAMING_SNAKE_CASE					)				,       30_000					)







					def 			snake_case					(							self :     str					):
									self.assertEqual(self.get_tokenizer().vocab_size				,       30_000					)







					def 			snake_case					(							self :     int					):
									if not self.test_rust_tokenizer:
													return

									lowercase__					:   Any													=					self.get_tokenizer()
									lowercase__					:   int													=					self.get_rust_tokenizer()

									lowercase__					:   List[Any]													=					"I was born in 92000, and this is falsé."

									lowercase__					:   Tuple													=					tokenizer.tokenize(SCREAMING_SNAKE_CASE					)
									lowercase__					:   Tuple													=					rust_tokenizer.tokenize(SCREAMING_SNAKE_CASE					)
									self.assertListEqual(SCREAMING_SNAKE_CASE				,       SCREAMING_SNAKE_CASE					)

									lowercase__					:   List[Any]													=					tokenizer.encode(SCREAMING_SNAKE_CASE				,       add_special_tokens=SCREAMING_SNAKE_CASE					)
									lowercase__					:   Optional[Any]													=					rust_tokenizer.encode(SCREAMING_SNAKE_CASE				,       add_special_tokens=SCREAMING_SNAKE_CASE					)
									self.assertListEqual(SCREAMING_SNAKE_CASE				,       SCREAMING_SNAKE_CASE					)

									lowercase__					:   str													=					self.get_rust_tokenizer()
									lowercase__					:   List[Any]													=					tokenizer.encode(SCREAMING_SNAKE_CASE					)
									lowercase__					:   List[Any]													=					rust_tokenizer.encode(SCREAMING_SNAKE_CASE					)
									self.assertListEqual(SCREAMING_SNAKE_CASE				,       SCREAMING_SNAKE_CASE					)







					def 			snake_case					(							self :     Optional[Any]					):
									lowercase__					:   Any													=					AlbertTokenizer(SCREAMING_SNAKE_CASE				,       keep_accents=SCREAMING_SNAKE_CASE					)

									lowercase__					:   str													=					tokenizer.tokenize("This is a test"					)
									self.assertListEqual(SCREAMING_SNAKE_CASE				,       ["▁this", "▁is", "▁a", "▁test"]					)

									self.assertListEqual(tokenizer.convert_tokens_to_ids(SCREAMING_SNAKE_CASE					)				,       [48, 25, 21, 1_289]					)

									lowercase__					:   Tuple													=					tokenizer.tokenize("I was born in 92000, and this is falsé."					)
									self.assertListEqual(
									    SCREAMING_SNAKE_CASE				,       ["▁i", "▁was", "▁born", "▁in", "▁9", "2000", ",", "▁and", "▁this", "▁is", "▁fal", "s", "é", "."]					)
									lowercase__					:   Union[str, Any]													=					tokenizer.convert_tokens_to_ids(SCREAMING_SNAKE_CASE					)
									self.assertListEqual(SCREAMING_SNAKE_CASE				,       [31, 23, 386, 19, 561, 3_050, 15, 17, 48, 25, 8_256, 18, 1, 9]					)

									lowercase__					:   List[str]													=					tokenizer.convert_ids_to_tokens(SCREAMING_SNAKE_CASE					)
									self.assertListEqual(
									    SCREAMING_SNAKE_CASE				,       ["▁i", "▁was", "▁born", "▁in", "▁9", "2000", ",", "▁and", "▁this", "▁is", "▁fal", "s", "<unk>", "."]				,       )







					def 			snake_case					(							self :     Any					):
									lowercase__					:   int													=					AlbertTokenizer(SCREAMING_SNAKE_CASE					)

									lowercase__					:   List[Any]													=					tokenizer.encode("sequence builders"					)
									lowercase__					:   List[Any]													=					tokenizer.encode("multi-sequence build"					)

									lowercase__					:   Dict													=					tokenizer.build_inputs_with_special_tokens(SCREAMING_SNAKE_CASE					)
									lowercase__					:   Union[str, Any]													=					tokenizer.build_inputs_with_special_tokens(SCREAMING_SNAKE_CASE				,       SCREAMING_SNAKE_CASE					)

									assert encoded_sentence == [tokenizer.cls_token_id] + text + [tokenizer.sep_token_id]
									assert encoded_pair == [tokenizer.cls_token_id] + text + [tokenizer.sep_token_id] + text_a + [
									    tokenizer.sep_token_id
									]







					@slow
					def 			snake_case					(							self :     Any					):
									# fmt: off
									lowercase__					:   int													=					{"attention_mask": [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], "input_ids": [[2, 21_970, 13, 5, 6_092, 167, 28, 7_103, 2_153, 673, 8, 7_028, 12_051, 18, 17, 7_103, 2_153, 673, 8, 3_515, 18_684, 8, 4_461, 6, 1_927, 297, 8, 12_060, 2_607, 18, 13, 5, 4_461, 15, 10_538, 38, 8, 135, 15, 822, 58, 15, 993, 10_363, 15, 1_460, 8_005, 4_461, 15, 993, 255, 2_328, 9, 9, 9, 6, 26, 1_112, 816, 3_260, 13, 5, 103, 2_377, 6, 17, 1_112, 816, 2_782, 13, 5, 103, 10_641, 6, 29, 84, 2_512, 2_430, 782, 18_684, 2_761, 19, 808, 2_430, 2_556, 17, 855, 1_480, 9_477, 4_091, 128, 11_712, 15, 7_103, 2_153, 673, 17, 24_883, 9_990, 9, 3], [2, 11_502, 25, 1_006, 20, 782, 8, 11_809, 855, 1_732, 19_393, 18_667, 37, 367, 21_018, 69, 1_854, 34, 11_860, 19_124, 27, 156, 225, 17, 193, 4_141, 19, 65, 9_124, 9, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [2, 14, 2_231, 886, 2_385, 17_659, 84, 14, 16_792, 1_952, 9, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], "token_type_ids": [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
									# fmt: on

									self.tokenizer_integration_test_util(
									    expected_encoding=SCREAMING_SNAKE_CASE				,       model_name="albert-base-v2"				,       revision="6b6560eaf5ff2e250b00c50f380c5389a9c2d82e"				,       )









import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


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

lowerCAmelCase__					  =  [
    # (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'''),
]

lowerCAmelCase__					  =  [
    # (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'''),
]

lowerCAmelCase__					  =  []
# 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
								lowerCAmelCase__					  =  f'''down_blocks.{i}.resnets.{j}.'''
								lowerCAmelCase__					  =  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
												lowerCAmelCase__					  =  f'''down_blocks.{i}.attentions.{j}.'''
												lowerCAmelCase__					  =  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
								lowerCAmelCase__					  =  f'''up_blocks.{i}.resnets.{j}.'''
								lowerCAmelCase__					  =  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
												lowerCAmelCase__					  =  f'''up_blocks.{i}.attentions.{j}.'''
												lowerCAmelCase__					  =  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
								lowerCAmelCase__					  =  f'''down_blocks.{i}.downsamplers.0.conv.'''
								lowerCAmelCase__					  =  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
								lowerCAmelCase__					  =  f'''up_blocks.{i}.upsamplers.0.'''
								lowerCAmelCase__					  =  f'''output_blocks.{3*i + 2}.{1 if i == 0 else 2}.'''
								unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))

lowerCAmelCase__					  =  '''mid_block.attentions.0.'''
lowerCAmelCase__					  =  '''middle_block.1.'''
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))

for j in range(2):
				lowerCAmelCase__					  =  f'''mid_block.resnets.{j}.'''
				lowerCAmelCase__					  =  f'''middle_block.{2*j}.'''
				unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))


def  __lowerCamelCase (       lowerCamelCase__						):

				"""simple docstring"""


				lowercase__					:   List[str]													=					{k: k for k in unet_state_dict.keys()}
				for sd_name, hf_name in unet_conversion_map:
								lowercase__					:   str													=					sd_name
				for k, v in mapping.items():
								if "resnets" in k:
												for sd_part, hf_part in unet_conversion_map_resnet:
																lowercase__					:   List[str]													=					v.replace(lowerCamelCase__					,			lowerCamelCase__						)
												lowercase__					:   str													=					v
				for k, v in mapping.items():
								for sd_part, hf_part in unet_conversion_map_layer:
												lowercase__					:   int													=					v.replace(lowerCamelCase__					,			lowerCamelCase__						)
								lowercase__					:   Optional[Any]													=					v
				lowercase__					:   Union[str, Any]													=					{v: unet_state_dict[k] for k, v in mapping.items()}
				return new_state_dict


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

lowerCAmelCase__					  =  [
    # (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):
								lowerCAmelCase__					  =  f'''encoder.down_blocks.{i}.resnets.{j}.'''
								lowerCAmelCase__					  =  f'''encoder.down.{i}.block.{j}.'''
								vae_conversion_map.append((sd_down_prefix, hf_down_prefix))

				if i < 3:
								lowerCAmelCase__					  =  f'''down_blocks.{i}.downsamplers.0.'''
								lowerCAmelCase__					  =  f'''down.{i}.downsample.'''
								vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))

								lowerCAmelCase__					  =  f'''up_blocks.{i}.upsamplers.0.'''
								lowerCAmelCase__					  =  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):
								lowerCAmelCase__					  =  f'''decoder.up_blocks.{i}.resnets.{j}.'''
								lowerCAmelCase__					  =  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):
				lowerCAmelCase__					  =  f'''mid_block.resnets.{i}.'''
				lowerCAmelCase__					  =  f'''mid.block_{i+1}.'''
				vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))


lowerCAmelCase__					  =  [
    # (stable-diffusion, HF Diffusers)
    ('''norm.''', '''group_norm.'''),
    ('''q.''', '''query.'''),
    ('''k.''', '''key.'''),
    ('''v.''', '''value.'''),
    ('''proj_out.''', '''proj_attn.'''),
]


def  __lowerCamelCase (       lowerCamelCase__						):

				"""simple docstring"""


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


def  __lowerCamelCase (       lowerCamelCase__						):

				"""simple docstring"""


				lowercase__					:   str													=					{k: k for k in vae_state_dict.keys()}
				for k, v in mapping.items():
								for sd_part, hf_part in vae_conversion_map:
												lowercase__					:   Optional[int]													=					v.replace(lowerCamelCase__					,			lowerCamelCase__						)
								lowercase__					:   Dict													=					v
				for k, v in mapping.items():
								if "attentions" in k:
												for sd_part, hf_part in vae_conversion_map_attn:
																lowercase__					:   List[str]													=					v.replace(lowerCamelCase__					,			lowerCamelCase__						)
												lowercase__					:   int													=					v
				lowercase__					:   Union[str, Any]													=					{v: vae_state_dict[k] for k, v in mapping.items()}
				lowercase__					:   Optional[int]													=					["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"""						)
																lowercase__					:   Dict													=					reshape_weight_for_sd(lowerCamelCase__						)
				return new_state_dict


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


lowerCAmelCase__					  =  [
    # (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'''),
]
lowerCAmelCase__					  =  {re.escape(x[1]): x[0] for x in textenc_conversion_lst}
lowerCAmelCase__					  =  re.compile('''|'''.join(protected.keys()))

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


def  __lowerCamelCase (       lowerCamelCase__						):

				"""simple docstring"""


				lowercase__					:   Union[str, Any]													=					{}
				lowercase__					:   List[Any]													=					{}
				lowercase__					:   List[Any]													=					{}
				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"						)
								):
												lowercase__					:   int													=					k[: -len(".q_proj.weight"						)]
												lowercase__					:   Optional[Any]													=					k[-len("q_proj.weight"						)]
												if k_pre not in capture_qkv_weight:
																lowercase__					:   Dict													=					[None, None, None]
												lowercase__					:   Any													=					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"						)
								):
												lowercase__					:   Optional[int]													=					k[: -len(".q_proj.bias"						)]
												lowercase__					:   Any													=					k[-len("q_proj.bias"						)]
												if k_pre not in capture_qkv_bias:
																lowercase__					:   str													=					[None, None, None]
												lowercase__					:   str													=					v
												continue

								lowercase__					:   Union[str, Any]													=					textenc_pattern.sub(lambda lowerCamelCase__						: protected[re.escape(m.group(0						)						)]					,			lowerCamelCase__						)
								lowercase__					:   List[Any]													=					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"						)
								lowercase__					:   str													=					textenc_pattern.sub(lambda lowerCamelCase__						: protected[re.escape(m.group(0						)						)]					,			lowerCamelCase__						)
								lowercase__					:   Any													=					torch.cat(lowerCamelCase__						)

				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"						)
								lowercase__					:   List[str]													=					textenc_pattern.sub(lambda lowerCamelCase__						: protected[re.escape(m.group(0						)						)]					,			lowerCamelCase__						)
								lowercase__					:   Tuple													=					torch.cat(lowerCamelCase__						)

				return new_state_dict


def  __lowerCamelCase (       lowerCamelCase__						):

				"""simple docstring"""


				return text_enc_dict


if __name__ == "__main__":
				lowerCAmelCase__					  =  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.'''
				)

				lowerCAmelCase__					  =  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
				lowerCAmelCase__					  =  osp.join(args.model_path, '''unet''', '''diffusion_pytorch_model.safetensors''')
				lowerCAmelCase__					  =  osp.join(args.model_path, '''vae''', '''diffusion_pytorch_model.safetensors''')
				lowerCAmelCase__					  =  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):
								lowerCAmelCase__					  =  load_file(unet_path, device='''cpu''')
				else:
								lowerCAmelCase__					  =  osp.join(args.model_path, '''unet''', '''diffusion_pytorch_model.bin''')
								lowerCAmelCase__					  =  torch.load(unet_path, map_location='''cpu''')

				if osp.exists(vae_path):
								lowerCAmelCase__					  =  load_file(vae_path, device='''cpu''')
				else:
								lowerCAmelCase__					  =  osp.join(args.model_path, '''vae''', '''diffusion_pytorch_model.bin''')
								lowerCAmelCase__					  =  torch.load(vae_path, map_location='''cpu''')

				if osp.exists(text_enc_path):
								lowerCAmelCase__					  =  load_file(text_enc_path, device='''cpu''')
				else:
								lowerCAmelCase__					  =  osp.join(args.model_path, '''text_encoder''', '''pytorch_model.bin''')
								lowerCAmelCase__					  =  torch.load(text_enc_path, map_location='''cpu''')

				# Convert the UNet model
				lowerCAmelCase__					  =  convert_unet_state_dict(unet_state_dict)
				lowerCAmelCase__					  =  {'''model.diffusion_model.''' + k: v for k, v in unet_state_dict.items()}

				# Convert the VAE model
				lowerCAmelCase__					  =  convert_vae_state_dict(vae_state_dict)
				lowerCAmelCase__					  =  {'''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
				lowerCAmelCase__					  =  '''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
								lowerCAmelCase__					  =  {'''transformer.''' + k: v for k, v in text_enc_dict.items()}
								lowerCAmelCase__					  =  convert_text_enc_state_dict_vaa(text_enc_dict)
								lowerCAmelCase__					  =  {'''cond_stage_model.model.''' + k: v for k, v in text_enc_dict.items()}
				else:
								lowerCAmelCase__					  =  convert_text_enc_state_dict(text_enc_dict)
								lowerCAmelCase__					  =  {'''cond_stage_model.transformer.''' + k: v for k, v in text_enc_dict.items()}

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

				if args.use_safetensors:
								save_file(state_dict, args.checkpoint_path)
				else:
								lowerCAmelCase__					  =  {'''state_dict''': state_dict}
								torch.save(state_dict, args.checkpoint_path)




