



import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


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

lowercase_:  List[str]     						=					[
    # (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_:  int     						=					[
    # (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_:  Union[str, 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_:  Dict     						=					F"""down_blocks.{i}.resnets.{j}."""
  lowercase_:  Optional[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_:  str     						=					F"""down_blocks.{i}.attentions.{j}."""
   lowercase_:  Any     						=					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_:  str     						=					F"""up_blocks.{i}.resnets.{j}."""
  lowercase_:  Dict     						=					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[Any]     						=					F"""up_blocks.{i}.attentions.{j}."""
   lowercase_:  Optional[Any]     						=					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_:  List[str]     						=					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_:  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_:  str     						=					'mid_block.attentions.0.'
lowercase_:  Optional[int]     						=					'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 	_lowercase	(     UpperCAmelCase_):




 """simple docstring"""




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


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

lowercase_:  List[str]     						=					[
    # (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_:  int     						=					F"""encoder.down_blocks.{i}.resnets.{j}."""
  lowercase_:  List[str]     						=					F"""encoder.down.{i}.block.{j}."""
  vae_conversion_map.append((sd_down_prefix, hf_down_prefix))

 if i < 3:
  lowercase_:  Any     						=					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_:  List[Any]     						=					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_:  Union[str, Any]     						=					F"""decoder.up_blocks.{i}.resnets.{j}."""
  lowercase_:  Optional[int]     						=					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_:  Union[str, Any]     						=					F"""mid_block.resnets.{i}."""
 lowercase_:  List[str]     						=					F"""mid.block_{i+1}."""
 vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))


lowercase_:  List[str]     						=					[
    # (stable-diffusion, HF Diffusers)
    ('norm.', 'group_norm.'),
    ('q.', 'query.'),
    ('k.', 'key.'),
    ('v.', 'value.'),
    ('proj_out.', 'proj_attn.'),
]
def 	_lowercase	(     UpperCAmelCase_):




 """simple docstring"""




 return w.reshape(*w.shape						,   1						,   1)
def 	_lowercase	(     UpperCAmelCase_):




 """simple docstring"""




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


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


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

# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
lowercase_:  Union[str, Any]     						=					{'q': 0, 'k': 1, 'v': 2}
def 	_lowercase	(     UpperCAmelCase_):




 """simple docstring"""




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

  snake_case__				:							str      							= textenc_pattern.sub(lambda UpperCAmelCase_: protected[re.escape(m.group(0))]						,   _lowerCamelCase)
  snake_case__				:							Tuple      							= 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__				:							Tuple      							= textenc_pattern.sub(lambda UpperCAmelCase_: protected[re.escape(m.group(0))]						,   _lowerCamelCase)
  snake_case__				:							int      							= 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""")
  snake_case__				:							Any      							= textenc_pattern.sub(lambda UpperCAmelCase_: protected[re.escape(m.group(0))]						,   _lowerCamelCase)
  snake_case__				:							Optional[int]      							= torch.cat(_lowerCamelCase)

 return new_state_dict
def 	_lowercase	(     UpperCAmelCase_):




 """simple docstring"""




 return text_enc_dict


if __name__ == "__main__":
 lowercase_:  List[Any]     						=					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_:  Tuple     						=					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_:  Union[str, Any]     						=					osp.join(args.model_path, 'unet', 'diffusion_pytorch_model.safetensors')
 lowercase_:  Optional[int]     						=					osp.join(args.model_path, 'vae', 'diffusion_pytorch_model.safetensors')
 lowercase_:  Optional[int]     						=					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_:  Union[str, Any]     						=					load_file(unet_path, device='cpu')
 else:
  lowercase_:  Dict     						=					osp.join(args.model_path, 'unet', 'diffusion_pytorch_model.bin')
  lowercase_:  List[str]     						=					torch.load(unet_path, map_location='cpu')

 if osp.exists(vae_path):
  lowercase_:  Dict     						=					load_file(vae_path, device='cpu')
 else:
  lowercase_:  int     						=					osp.join(args.model_path, 'vae', 'diffusion_pytorch_model.bin')
  lowercase_:  Optional[Any]     						=					torch.load(vae_path, map_location='cpu')

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

 # Convert the UNet model
 lowercase_:  List[str]     						=					convert_unet_state_dict(unet_state_dict)
 lowercase_:  str     						=					{'model.diffusion_model.' + k: v for k, v in unet_state_dict.items()}

 # Convert the VAE model
 lowercase_:  Union[str, Any]     						=					convert_vae_state_dict(vae_state_dict)
 lowercase_:  List[str]     						=					{'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[int]     						=					'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_:  List[Any]     						=					{'transformer.' + k: v for k, v in text_enc_dict.items()}
  lowercase_:  Tuple     						=					convert_text_enc_state_dict_vaa(text_enc_dict)
  lowercase_:  Union[str, Any]     						=					{'cond_stage_model.model.' + k: v for k, v in text_enc_dict.items()}
 else:
  lowercase_:  Optional[Any]     						=					convert_text_enc_state_dict(text_enc_dict)
  lowercase_:  Optional[int]     						=					{'cond_stage_model.transformer.' + k: v for k, v in text_enc_dict.items()}

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

 if args.use_safetensors:
  save_file(state_dict, args.checkpoint_path)
 else:
  lowercase_:  Optional[int]     						=					{'state_dict': state_dict}
  torch.save(state_dict, args.checkpoint_path)











'''simple docstring'''






import argparse
import os
import re

import packaging.version


__UpperCamelCase					    =	"examples/"
__UpperCamelCase					    =	{
    "examples": (re.compile(R"^check_min_version\(\"[^\"]+\"\)\s*$", re.MULTILINE), "check_min_version(\"VERSION\")\n"),
    "init": (re.compile(R"^__version__\s+=\s+\"([^\"]+)\"\s*$", re.MULTILINE), "__version__ = \"VERSION\"\n"),
    "setup": (re.compile(R"^(\s*)version\s*=\s*\"[^\"]+\",", re.MULTILINE), R"\1version=\"VERSION\","),
    "doc": (re.compile(R"^(\s*)release\s*=\s*\"[^\"]+\"$", re.MULTILINE), "release = \"VERSION\"\n"),
}
__UpperCamelCase					    =	{
    "init": "src/transformers/__init__.py",
    "setup": "setup.py",
}
__UpperCamelCase					    =	"README.md"







def 					_a						(       _lowerCamelCase    ,	_lowerCamelCase    ,	_lowerCamelCase ) ->				Tuple:



			"""simple docstring"""


			with open(_lowerCamelCase    ,	"""r"""    ,	encoding="""utf-8"""    ,	newline="""\n""" ) as f:
						__snake_case					:      Union[str, Any]							      =     f.read()
			__snake_case      ,     __snake_case					:      List[Any]							      =     REPLACE_PATTERNS[pattern]
			__snake_case					:      Optional[Any]							      =     replace.replace("""VERSION"""    ,	_lowerCamelCase )
			__snake_case					:      Optional[Any]							      =     re_pattern.sub(_lowerCamelCase    ,	_lowerCamelCase )
			with open(_lowerCamelCase    ,	"""w"""    ,	encoding="""utf-8"""    ,	newline="""\n""" ) as f:
						f.write(_lowerCamelCase )







def 					_a						(       _lowerCamelCase ) ->				Union[str, Any]:



			"""simple docstring"""


			for folder, directories, fnames in os.walk(_lowerCamelCase ):
						# Removing some of the folders with non-actively maintained examples from the walk
						if "research_projects" in directories:
									directories.remove("""research_projects""" )
						if "legacy" in directories:
									directories.remove("""legacy""" )
						for fname in fnames:
									if fname.endswith(""".py""" ):
												update_version_in_file(os.path.join(_lowerCamelCase    ,	_lowerCamelCase )    ,	_lowerCamelCase    ,	pattern="""examples""" )







def 					_a						(       _lowerCamelCase    ,	_lowerCamelCase=False ) ->				str:



			"""simple docstring"""


			for pattern, fname in REPLACE_FILES.items():
						update_version_in_file(_lowerCamelCase    ,	_lowerCamelCase    ,	_lowerCamelCase )
			if not patch:
						update_version_in_examples(_lowerCamelCase )







def 					_a						(       ) ->				Optional[int]:



			"""simple docstring"""


			__snake_case					:      str							      =     """🤗 Transformers currently provides the following architectures"""
			__snake_case					:      List[Any]							      =     """1. Want to contribute a new model?"""
			with open(_lowerCamelCase    ,	"""r"""    ,	encoding="""utf-8"""    ,	newline="""\n""" ) as f:
						__snake_case					:      List[str]							      =     f.readlines()

			# Find the start of the list.
			__snake_case					:      Optional[Any]							      =     0
			while not lines[start_index].startswith(_start_prompt ):
						start_index += 1
			start_index += 1

			__snake_case					:      int							      =     start_index
			# Update the lines in the model list.
			while not lines[index].startswith(_end_prompt ):
						if lines[index].startswith("""1.""" ):
									__snake_case					:      Optional[Any]							      =     lines[index].replace(
									    """https://huggingface.co/docs/transformers/main/model_doc"""    ,	"""https://huggingface.co/docs/transformers/model_doc"""    ,	)
						index += 1

			with open(_lowerCamelCase    ,	"""w"""    ,	encoding="""utf-8"""    ,	newline="""\n""" ) as f:
						f.writelines(_lowerCamelCase )







def 					_a						(       ) ->				Union[str, Any]:



			"""simple docstring"""


			with open(REPLACE_FILES["""init"""]    ,	"""r""" ) as f:
						__snake_case					:      List[Any]							      =     f.read()
			__snake_case					:      str							      =     REPLACE_PATTERNS["""init"""][0].search(_lowerCamelCase ).groups()[0]
			return packaging.version.parse(_lowerCamelCase )







def 					_a						(       _lowerCamelCase=False ) ->				int:



			"""simple docstring"""


			__snake_case					:      List[Any]							      =     get_version()
			if patch and default_version.is_devrelease:
						raise ValueError("""Can't create a patch version from the dev branch, checkout a released version!""" )
			if default_version.is_devrelease:
						__snake_case					:      str							      =     default_version.base_version
			elif patch:
						__snake_case					:      Optional[int]							      =     F'''{default_version.major}.{default_version.minor}.{default_version.micro + 1}'''
			else:
						__snake_case					:      Dict							      =     F'''{default_version.major}.{default_version.minor + 1}.0'''

			# Now let's ask nicely if that's the right one.
			__snake_case					:      Dict							      =     input(F'''Which version are you releasing? [{default_version}]''' )
			if len(_lowerCamelCase ) == 0:
						__snake_case					:      Any							      =     default_version

			print(F'''Updating version to {version}.''' )
			global_version_update(_lowerCamelCase    ,	patch=_lowerCamelCase )
			if not patch:
						print("""Cleaning main README, don't forget to run `make fix-copies`.""" )
						clean_main_ref_in_model_list()







def 					_a						(       ) ->				Tuple:



			"""simple docstring"""


			__snake_case					:      Optional[Any]							      =     get_version()
			__snake_case					:      Tuple							      =     F'''{current_version.major}.{current_version.minor + 1}.0.dev0'''
			__snake_case					:      Union[str, Any]							      =     current_version.base_version

			# Check with the user we got that right.
			__snake_case					:      int							      =     input(F'''Which version are we developing now? [{dev_version}]''' )
			if len(_lowerCamelCase ) == 0:
						__snake_case					:      Optional[int]							      =     dev_version

			print(F'''Updating version to {version}.''' )
			global_version_update(_lowerCamelCase )
			print("""Cleaning main README, don't forget to run `make fix-copies`.""" )
			clean_main_ref_in_model_list()


if __name__ == "__main__":
	__UpperCamelCase					    =	argparse.ArgumentParser()
	parser.add_argument("--post_release", action="store_true", help="Whether this is pre or post release.")
	parser.add_argument("--patch", action="store_true", help="Whether or not this is a patch release.")
	__UpperCamelCase					    =	parser.parse_args()
	if not args.post_release:
		pre_release_work(patch=args.patch)
	elif args.patch:
		print("Nothing to do after a patch :-)")
	else:
		post_release_work()


