


"""simple docstring"""





import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


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

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

lowercase  							=					"mid_block.attentions.0."
lowercase  							=					"middle_block.1."
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))

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




def 	UpperCAmelCase     (		A	:						Any       ):

		'''simple docstring'''


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


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

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

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

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


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




def 	UpperCAmelCase     (		A	:						Optional[int]       ):

		'''simple docstring'''


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




def 	UpperCAmelCase     (		A	:						List[Any]       ):

		'''simple docstring'''


		_UpperCAmelCase							 =		{k: k for k in vae_state_dict.keys()}
		for k, v in mapping.items():
				for sd_part, hf_part in vae_conversion_map:
						_UpperCAmelCase							 =		v.replace(_lowerCamelCase     ,				_lowerCamelCase       )
				_UpperCAmelCase							 =		v
		for k, v in mapping.items():
				if "attentions" in k:
						for sd_part, hf_part in vae_conversion_map_attn:
								_UpperCAmelCase							 =		v.replace(_lowerCamelCase     ,				_lowerCamelCase       )
						_UpperCAmelCase							 =		v
		_UpperCAmelCase							 =		{v: vae_state_dict[k] for k, v in mapping.items()}
		_UpperCAmelCase							 =		["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'       )
								_UpperCAmelCase							 =		reshape_weight_for_sd(_lowerCamelCase       )
		return new_state_dict


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


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

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




def 	UpperCAmelCase     (		A	:						List[str]       ):

		'''simple docstring'''


		_UpperCAmelCase							 =		{}
		_UpperCAmelCase							 =		{}
		_UpperCAmelCase							 =		{}
		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'       )
				):
						_UpperCAmelCase							 =		k[: -len('.q_proj.weight'       )]
						_UpperCAmelCase							 =		k[-len('q_proj.weight'       )]
						if k_pre not in capture_qkv_weight:
								_UpperCAmelCase							 =		[None, None, None]
						_UpperCAmelCase							 =		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'       )
				):
						_UpperCAmelCase							 =		k[: -len('.q_proj.bias'       )]
						_UpperCAmelCase							 =		k[-len('q_proj.bias'       )]
						if k_pre not in capture_qkv_bias:
								_UpperCAmelCase							 =		[None, None, None]
						_UpperCAmelCase							 =		v
						continue

				_UpperCAmelCase							 =		textenc_pattern.sub(lambda A       : protected[re.escape(m.group(0       )       )]     ,				_lowerCamelCase       )
				_UpperCAmelCase							 =		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'       )
				_UpperCAmelCase							 =		textenc_pattern.sub(lambda A       : protected[re.escape(m.group(0       )       )]     ,				_lowerCamelCase       )
				_UpperCAmelCase							 =		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'       )
				_UpperCAmelCase							 =		textenc_pattern.sub(lambda A       : protected[re.escape(m.group(0       )       )]     ,				_lowerCamelCase       )
				_UpperCAmelCase							 =		torch.cat(_lowerCamelCase       )

		return new_state_dict




def 	UpperCAmelCase     (		A	:						Any       ):

		'''simple docstring'''


		return text_enc_dict


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

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

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

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

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

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

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







"""simple docstring"""





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


lowercase  							=					{
    '''configuration_mvp''': ['''MVP_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''MvpConfig''', '''MvpOnnxConfig'''],
    '''tokenization_mvp''': ['''MvpTokenizer'''],
}

try:
					if not is_tokenizers_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowercase  							=					['''MvpTokenizerFast''']

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowercase  							=					[
					    '''MVP_PRETRAINED_MODEL_ARCHIVE_LIST''',
					    '''MvpForCausalLM''',
					    '''MvpForConditionalGeneration''',
					    '''MvpForQuestionAnswering''',
					    '''MvpForSequenceClassification''',
					    '''MvpModel''',
					    '''MvpPreTrainedModel''',
					]

if TYPE_CHECKING:
					from .configuration_mvp import MVP_PRETRAINED_CONFIG_ARCHIVE_MAP, MvpConfig, MvpOnnxConfig
					from .tokenization_mvp import MvpTokenizer

					try:
										if not is_tokenizers_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .tokenization_mvp_fast import MvpTokenizerFast

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_mvp import (
										    MVP_PRETRAINED_MODEL_ARCHIVE_LIST,
										    MvpForCausalLM,
										    MvpForConditionalGeneration,
										    MvpForQuestionAnswering,
										    MvpForSequenceClassification,
										    MvpModel,
										    MvpPreTrainedModel,
										)

else:
					import sys

					lowercase  							=					_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)


