




import sys
from pathlib import Path


lowerCAmelCase__					  =  Path(__file__).resolve().parents[3] / '''src'''
sys.path.insert(1, str(git_repo_path))

import dataclasses  # noqa
import io  # noqa
import itertools  # noqa
import json  # noqa
import os  # noqa
import unittest  # noqa
from copy import deepcopy  # noqa

from parameterized import parameterized  # noqa
from transformers import TrainingArguments, is_torch_available  # noqa
from transformers.deepspeed import is_deepspeed_available  # noqa
from transformers.file_utils import WEIGHTS_NAME  # noqa
from transformers.testing_utils import (  # noqa
    CaptureLogger,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    mockenv_context,
    require_deepspeed,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)
from transformers.trainer_utils import set_seed  # noqa


set_seed(4_2)

lowerCAmelCase__					  =  {'''base''': '''patrickvonplaten/wav2vec2_tiny_random''', '''robust''': '''patrickvonplaten/wav2vec2_tiny_random_robust'''}

lowerCAmelCase__					  =  '''zero2'''
lowerCAmelCase__					  =  '''zero3'''
lowerCAmelCase__					  =  [ZEROa, ZEROa]


def  __lowerCamelCase (       lowerCamelCase__					,			lowerCamelCase__					,			lowerCamelCase__						):

				"""simple docstring"""


				lowercase__					:   List[str]													=					parameterized.to_safe_name("_".join(str(snake_case_						) for x in param.args						)						)
				return F"""{func.__name__}_{param_based_name}"""


# Cartesian-product of zero stages with models to test
lowerCAmelCase__					  =  list(itertools.product(stages, models.keys()))

@slow
@require_deepspeed
@require_torch_gpu
class 				snake_case__(_UpperCAmelCase      ):

					"""simple docstring"""






					@parameterized.expand(SCREAMING_SNAKE_CASE_				,       name_func=SCREAMING_SNAKE_CASE_					)
					def 			snake_case					(							self :     int				,       SCREAMING_SNAKE_CASE :     Optional[int]				,       SCREAMING_SNAKE_CASE :     Optional[int]					):
									self.run_and_check(
									    stage=SCREAMING_SNAKE_CASE_				,       model=SCREAMING_SNAKE_CASE_				,       distributed=SCREAMING_SNAKE_CASE_				,       fpaa=SCREAMING_SNAKE_CASE_				,       )







					@require_torch_multi_gpu
					@parameterized.expand(SCREAMING_SNAKE_CASE_				,       name_func=SCREAMING_SNAKE_CASE_					)
					def 			snake_case					(							self :     Optional[int]				,       SCREAMING_SNAKE_CASE :     List[Any]				,       SCREAMING_SNAKE_CASE :     List[str]					):
									self.run_and_check(
									    stage=SCREAMING_SNAKE_CASE_				,       model=SCREAMING_SNAKE_CASE_				,       distributed=SCREAMING_SNAKE_CASE_				,       fpaa=SCREAMING_SNAKE_CASE_				,       )







					@parameterized.expand(SCREAMING_SNAKE_CASE_				,       name_func=SCREAMING_SNAKE_CASE_					)
					def 			snake_case					(							self :     Optional[int]				,       SCREAMING_SNAKE_CASE :     List[Any]				,       SCREAMING_SNAKE_CASE :     Dict					):
									self.run_and_check(
									    stage=SCREAMING_SNAKE_CASE_				,       model=SCREAMING_SNAKE_CASE_				,       distributed=SCREAMING_SNAKE_CASE_				,       fpaa=SCREAMING_SNAKE_CASE_				,       )







					@require_torch_multi_gpu
					@parameterized.expand(SCREAMING_SNAKE_CASE_				,       name_func=SCREAMING_SNAKE_CASE_					)
					def 			snake_case					(							self :     Tuple				,       SCREAMING_SNAKE_CASE :     str				,       SCREAMING_SNAKE_CASE :     Tuple					):
									self.run_and_check(
									    stage=SCREAMING_SNAKE_CASE_				,       model=SCREAMING_SNAKE_CASE_				,       distributed=SCREAMING_SNAKE_CASE_				,       fpaa=SCREAMING_SNAKE_CASE_				,       )







					def 			snake_case					(							self :     List[Any]				,       SCREAMING_SNAKE_CASE :     Optional[int]					):
									# XXX: run_asr is premature and doesn't save any results
									# so all we check for now is that the process didn't fail
									pass







					def 			snake_case					(							self :     Optional[int]				,       SCREAMING_SNAKE_CASE :     List[str]				,       SCREAMING_SNAKE_CASE :     List[Any]				,       SCREAMING_SNAKE_CASE :     str = 10				,       SCREAMING_SNAKE_CASE :     List[Any] = True				,       SCREAMING_SNAKE_CASE :     str = True				,       SCREAMING_SNAKE_CASE :     List[str] = True				,       ):
									lowercase__					:   int													=					models[model]

									lowercase__					:   Optional[int]													=					self.run_trainer(
									    stage=SCREAMING_SNAKE_CASE_				,       model_name=SCREAMING_SNAKE_CASE_				,       eval_steps=SCREAMING_SNAKE_CASE_				,       num_train_epochs=1				,       distributed=SCREAMING_SNAKE_CASE_				,       fpaa=SCREAMING_SNAKE_CASE_				,       )

									self.do_checks(SCREAMING_SNAKE_CASE_					)

									return output_dir







					def 			snake_case					(							self :     Union[str, Any]				,       SCREAMING_SNAKE_CASE :     List[Any]				,       SCREAMING_SNAKE_CASE :     Tuple				,       SCREAMING_SNAKE_CASE :     Any = 10				,       SCREAMING_SNAKE_CASE :     str = 1				,       SCREAMING_SNAKE_CASE :     Optional[Any] = True				,       SCREAMING_SNAKE_CASE :     Union[str, Any] = True				,       ):
									lowercase__					:   Optional[int]													=					self.get_auto_remove_tmp_dir("./xxx"				,       after=SCREAMING_SNAKE_CASE_					)
									lowercase__					:   Optional[Any]													=					f"""\n            --model_name_or_path {model_name}\n            --dataset_name hf-internal-testing/librispeech_asr_dummy\n            --dataset_config_name clean\n            --train_split_name validation\n            --validation_split_name validation\n            --output_dir {output_dir}\n            --num_train_epochs {str(SCREAMING_SNAKE_CASE_					)}\n            --per_device_train_batch_size 2\n            --per_device_eval_batch_size 2\n            --evaluation_strategy steps\n            --learning_rate 5e-4\n            --warmup_steps 8\n            --orthography timit\n            --preprocessing_num_workers 1\n            --group_by_length\n            --freeze_feature_extractor\n            --report_to none\n            --save_steps 0\n            --eval_steps {eval_steps}\n            --report_to none\n        """.split()

									if fpaa:
													args.extend(["--fp16"]					)

									# currently ds_config_wav2vec2_zero.json requires "zero_optimization.find_unused_parameters": true,
									# hence the separate config files
									lowercase__					:   int													=					f"""--deepspeed {self.test_file_dir_str}/ds_config_wav2vec2_{stage}.json""".split()
									lowercase__					:   Union[str, Any]													=					[f"""{self.examples_dir_str}/research_projects/wav2vec2/run_asr.py"""]
									lowercase__					:   Dict													=					self.get_launcher(SCREAMING_SNAKE_CASE_					)

									lowercase__					:   int													=					launcher + script + args + ds_args
									# keep for quick debug
									# print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
									execute_subprocess_async(SCREAMING_SNAKE_CASE_				,       env=self.get_env()					)

									return output_dir







					def 			snake_case					(							self :     Union[str, Any]				,       SCREAMING_SNAKE_CASE :     str=False					):
									# 1. explicitly set --num_nodes=1 just in case these tests end up run on a multi-node setup
									# - it won't be able to handle that
									# 2. for now testing with just 2 gpus max (since some quality tests may give different
									# results with mode gpus because we use very little data)
									lowercase__					:   Optional[Any]													=					min(2				,       get_gpu_count()					) if distributed else 1
									return f"""deepspeed --num_nodes 1 --num_gpus {num_gpus}""".split()









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)




