





import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    UniSpeechConfig,
    UniSpeechForCTC,
    UniSpeechForPreTraining,
    WavaVecaFeatureExtractor,
    WavaVecaPhonemeCTCTokenizer,
    WavaVecaProcessor,
    logging,
)


logging.set_verbosity_info()
_lowerCamelCase 					=      logging.get_logger(__name__)

_lowerCamelCase 					=      {
    'post_extract_proj': 'feature_projection.projection',
    'encoder.pos_conv.0': 'encoder.pos_conv_embed.conv',
    'self_attn.k_proj': 'encoder.layers.*.attention.k_proj',
    'self_attn.v_proj': 'encoder.layers.*.attention.v_proj',
    'self_attn.q_proj': 'encoder.layers.*.attention.q_proj',
    'self_attn.out_proj': 'encoder.layers.*.attention.out_proj',
    'self_attn_layer_norm': 'encoder.layers.*.layer_norm',
    'fc1': 'encoder.layers.*.feed_forward.intermediate_dense',
    'fc2': 'encoder.layers.*.feed_forward.output_dense',
    'final_layer_norm': 'encoder.layers.*.final_layer_norm',
    'encoder.layer_norm': 'encoder.layer_norm',
    'w2v_model.layer_norm': 'feature_projection.layer_norm',
    'quantizer.weight_proj': 'quantizer.weight_proj',
    'quantizer.vars': 'quantizer.codevectors',
    'project_q': 'project_q',
    'final_proj': 'project_hid',
    'w2v_encoder.proj': 'ctc_proj',
    'mask_emb': 'masked_spec_embed',
}
_lowerCamelCase 					=      [
    'ctc_proj',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
]






def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				Union[str, Any]       , UpperCamelCase__:				List[Any]       , UpperCamelCase__:				List[Any]       , UpperCamelCase__:				Dict       , UpperCamelCase__:				Union[str, Any]       , UpperCamelCase__:				Dict     ):
		for attribute in key.split("""."""     ):
				if is_finetuned:
						if attribute in ["quantizer", "project_q", "project_hid"]:
								# those layers are only relevant for pretraining and should be dropped
								return

						if attribute == "ctc_proj":
								# we should rename `ctc_proj` to `lm_head` for fine-tuned phoneme models
								SCREAMING_SNAKE_CASE__								=   """lm_head"""

				SCREAMING_SNAKE_CASE__								=   getattr(UpperCamelCase__       , UpperCamelCase__     )

		if weight_type is not None:
				SCREAMING_SNAKE_CASE__								=   getattr(UpperCamelCase__       , UpperCamelCase__     ).shape
		else:
				SCREAMING_SNAKE_CASE__								=   hf_pointer.shape

		assert hf_shape == value.shape, (
		    f'''Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be'''
		    f''' {value.shape} for {full_name}'''
		)

		if weight_type == "weight":
				SCREAMING_SNAKE_CASE__								=   value
		elif weight_type == "weight_g":
				SCREAMING_SNAKE_CASE__								=   value
		elif weight_type == "weight_v":
				SCREAMING_SNAKE_CASE__								=   value
		elif weight_type == "bias":
				SCREAMING_SNAKE_CASE__								=   value
		else:
				SCREAMING_SNAKE_CASE__								=   value

		logger.info(f'''{key + '.' + weight_type if weight_type is not None else ''} was initialized from {full_name}.'''     )






def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				Any       , UpperCamelCase__:				int       , UpperCamelCase__:				str     ):
		SCREAMING_SNAKE_CASE__								=   []
		SCREAMING_SNAKE_CASE__								=   fairseq_model.state_dict()

		SCREAMING_SNAKE_CASE__								=   hf_model.unispeech.feature_extractor

		for name, value in fairseq_dict.items():
				SCREAMING_SNAKE_CASE__								=   False
				if "conv_layers" in name:
						load_conv_layer(
						    UpperCamelCase__       , UpperCamelCase__       , UpperCamelCase__       , UpperCamelCase__       , hf_model.config.feat_extract_norm == """group"""       , )
						SCREAMING_SNAKE_CASE__								=   True
				else:
						for key, mapped_key in MAPPING.items():
								SCREAMING_SNAKE_CASE__								=   """unispeech.""" + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
								if key in name or key.split("""w2v_model."""     )[-1] == name.split("""."""     )[0]:
										SCREAMING_SNAKE_CASE__								=   True
										if "*" in mapped_key:
												SCREAMING_SNAKE_CASE__								=   name.split(UpperCamelCase__     )[0].split("""."""     )[-2]
												SCREAMING_SNAKE_CASE__								=   mapped_key.replace("""*"""       , UpperCamelCase__     )
										if "weight_g" in name:
												SCREAMING_SNAKE_CASE__								=   """weight_g"""
										elif "weight_v" in name:
												SCREAMING_SNAKE_CASE__								=   """weight_v"""
										elif "bias" in name:
												SCREAMING_SNAKE_CASE__								=   """bias"""
										elif "weight" in name:
												# TODO: don't match quantizer.weight_proj
												SCREAMING_SNAKE_CASE__								=   """weight"""
										else:
												SCREAMING_SNAKE_CASE__								=   None
										set_recursively(UpperCamelCase__       , UpperCamelCase__       , UpperCamelCase__       , UpperCamelCase__       , UpperCamelCase__       , UpperCamelCase__     )
								continue
				if not is_used:
						unused_weights.append(UpperCamelCase__     )

		logger.warning(f'''Unused weights: {unused_weights}'''     )






def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				List[str]       , UpperCamelCase__:				Optional[int]       , UpperCamelCase__:				int       , UpperCamelCase__:				Optional[int]       , UpperCamelCase__:				Optional[int]     ):
		SCREAMING_SNAKE_CASE__								=   full_name.split("""conv_layers."""     )[-1]
		SCREAMING_SNAKE_CASE__								=   name.split("""."""     )
		SCREAMING_SNAKE_CASE__								=   int(items[0]     )
		SCREAMING_SNAKE_CASE__								=   int(items[1]     )

		if type_id == 0:
				if "bias" in name:
						assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
						    f'''{full_name} has size {value.shape}, but'''
						    f''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'''
						)
						SCREAMING_SNAKE_CASE__								=   value
						logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''     )
				elif "weight" in name:
						assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
						    f'''{full_name} has size {value.shape}, but'''
						    f''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'''
						)
						SCREAMING_SNAKE_CASE__								=   value
						logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''     )
		elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
				if "bias" in name:
						assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
						    f'''{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was'''
						    " found."
						)
						SCREAMING_SNAKE_CASE__								=   value
						logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''     )
				elif "weight" in name:
						assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
						    f'''{full_name} has size {value.shape}, but'''
						    f''' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.'''
						)
						SCREAMING_SNAKE_CASE__								=   value
						logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''     )
		else:
				unused_weights.append(UpperCamelCase__     )






@torch.no_grad()
def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				List[Any]       , UpperCamelCase__:				Tuple       , UpperCamelCase__:				Optional[Any]=None       , UpperCamelCase__:				Any=None       , UpperCamelCase__:				Optional[int]=True     ):
		if config_path is not None:
				SCREAMING_SNAKE_CASE__								=   UniSpeechConfig.from_pretrained(UpperCamelCase__     )
		else:
				SCREAMING_SNAKE_CASE__								=   UniSpeechConfig()

		if is_finetuned:
				if dict_path:
						SCREAMING_SNAKE_CASE__								=   Dictionary.load_from_json(UpperCamelCase__     )

						# important change bos & pad token id since CTC symbol is <pad> and
						# not <s> as in fairseq
						SCREAMING_SNAKE_CASE__								=   target_dict.pad_index
						SCREAMING_SNAKE_CASE__								=   target_dict.bos_index
						SCREAMING_SNAKE_CASE__								=   target_dict.eos_index
						SCREAMING_SNAKE_CASE__								=   len(target_dict.symbols     )
						SCREAMING_SNAKE_CASE__								=   os.path.join(UpperCamelCase__       , """vocab.json"""     )
						if not os.path.isdir(UpperCamelCase__     ):
								logger.error("""--pytorch_dump_folder_path ({}) should be a directory""".format(UpperCamelCase__     )     )
								return
						os.makedirs(UpperCamelCase__       , exist_ok=UpperCamelCase__     )
						SCREAMING_SNAKE_CASE__								=   target_dict.indices

						# fairseq has the <pad> and <s> switched
						SCREAMING_SNAKE_CASE__								=   42
						SCREAMING_SNAKE_CASE__								=   43
						with open(UpperCamelCase__       , """w"""       , encoding="""utf-8"""     ) as vocab_handle:
								json.dump(UpperCamelCase__       , UpperCamelCase__     )
						SCREAMING_SNAKE_CASE__								=   WavaVecaPhonemeCTCTokenizer(
						    UpperCamelCase__       , unk_token=target_dict.unk_word       , pad_token=target_dict.pad_word       , bos_token=target_dict.bos_word       , eos_token=target_dict.eos_word       , word_delimiter_token="""|"""       , do_lower_case=UpperCamelCase__       , )
						SCREAMING_SNAKE_CASE__								=   True if config.feat_extract_norm == """layer""" else False
						SCREAMING_SNAKE_CASE__								=   WavaVecaFeatureExtractor(
						    feature_size=1       , sampling_rate=16_000       , padding_value=0       , do_normalize=UpperCamelCase__       , return_attention_mask=UpperCamelCase__       , )
						SCREAMING_SNAKE_CASE__								=   WavaVecaProcessor(feature_extractor=UpperCamelCase__       , tokenizer=UpperCamelCase__     )
						processor.save_pretrained(UpperCamelCase__     )

				SCREAMING_SNAKE_CASE__								=   UniSpeechForCTC(UpperCamelCase__     )
		else:
				SCREAMING_SNAKE_CASE__								=   UniSpeechForPreTraining(UpperCamelCase__     )

		if is_finetuned:
				SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   fairseq.checkpoint_utils.load_model_ensemble_and_task(
				    [checkpoint_path]       , arg_overrides={"""data""": """/""".join(dict_path.split("""/"""     )[:-1]     ), """w2v_path""": checkpoint_path}     )
		else:
				SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]     )

		SCREAMING_SNAKE_CASE__								=   model[0].eval()

		recursively_load_weights(UpperCamelCase__       , UpperCamelCase__       , UpperCamelCase__     )

		hf_unispeech.save_pretrained(UpperCamelCase__     )


if __name__ == "__main__":
	_lowerCamelCase 					=      argparse.ArgumentParser()
	parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
	parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to fairseq checkpoint')
	parser.add_argument('--dict_path', default=None, type=str, help='Path to dict of fine-tuned model')
	parser.add_argument('--config_path', default=None, type=str, help='Path to hf config.json of model to convert')
	parser.add_argument(
	    '--not_finetuned', action='store_true', help='Whether the model to convert is a fine-tuned model or not'
	)
	_lowerCamelCase 					=      parser.parse_args()
	convert_unispeech_checkpoint(
	    args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
	)






import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()




@slow
@require_torch_gpu
class 	UpperCamelCase_							(  unittest.TestCase     ):


							def   _snake_case	(				self     :Union[str, Any]		) -> List[str]:



									"""simple docstring"""



									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()


							def   _snake_case	(				self     :Any		) -> Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   StableDiffusionKDiffusionPipeline.from_pretrained("""CompVis/stable-diffusion-v1-4"""		)
									SCREAMING_SNAKE_CASE__								=   sd_pipe.to(__A		)
									sd_pipe.set_progress_bar_config(disable=__A		)

									sd_pipe.set_scheduler("""sample_euler"""		)

									SCREAMING_SNAKE_CASE__								=   """A painting of a squirrel eating a burger"""
									SCREAMING_SNAKE_CASE__								=   torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   sd_pipe([prompt]				,					generator=__A				,					guidance_scale=9.0				,					num_inference_steps=20				,					output_type="""np"""		)

									SCREAMING_SNAKE_CASE__								=   output.images

									SCREAMING_SNAKE_CASE__								=   image[0, -3:, -3:, -1]

									assert image.shape == (1, 512, 512, 3)
									SCREAMING_SNAKE_CASE__								=   np.array([0.0_4_4_7, 0.0_4_9_2, 0.0_4_6_8, 0.0_4_0_8, 0.0_3_8_3, 0.0_4_0_8, 0.0_3_5_4, 0.0_3_8_0, 0.0_3_3_9]		)

									assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   StableDiffusionKDiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""		)
									SCREAMING_SNAKE_CASE__								=   sd_pipe.to(__A		)
									sd_pipe.set_progress_bar_config(disable=__A		)

									sd_pipe.set_scheduler("""sample_euler"""		)

									SCREAMING_SNAKE_CASE__								=   """A painting of a squirrel eating a burger"""
									SCREAMING_SNAKE_CASE__								=   torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   sd_pipe([prompt]				,					generator=__A				,					guidance_scale=9.0				,					num_inference_steps=20				,					output_type="""np"""		)

									SCREAMING_SNAKE_CASE__								=   output.images

									SCREAMING_SNAKE_CASE__								=   image[0, -3:, -3:, -1]

									assert image.shape == (1, 512, 512, 3)
									SCREAMING_SNAKE_CASE__								=   np.array([0.1_2_3_7, 0.1_3_2_0, 0.1_4_3_8, 0.1_3_5_9, 0.1_3_9_0, 0.1_1_3_2, 0.1_2_7_7, 0.1_1_7_5, 0.1_1_1_2]		)

									assert np.abs(image_slice.flatten() - expected_slice		).max() < 5E-1


							def   _snake_case	(				self     :Tuple		) -> Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   StableDiffusionKDiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""		)
									SCREAMING_SNAKE_CASE__								=   sd_pipe.to(__A		)
									sd_pipe.set_progress_bar_config(disable=__A		)

									sd_pipe.set_scheduler("""sample_dpmpp_2m"""		)

									SCREAMING_SNAKE_CASE__								=   """A painting of a squirrel eating a burger"""
									SCREAMING_SNAKE_CASE__								=   torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   sd_pipe(
									    [prompt]				,					generator=__A				,					guidance_scale=7.5				,					num_inference_steps=15				,					output_type="""np"""				,					use_karras_sigmas=__A				,					)

									SCREAMING_SNAKE_CASE__								=   output.images

									SCREAMING_SNAKE_CASE__								=   image[0, -3:, -3:, -1]

									assert image.shape == (1, 512, 512, 3)
									SCREAMING_SNAKE_CASE__								=   np.array(
									    [0.1_1_3_8_1_6_8_9, 0.1_2_1_1_2_9_2_1, 0.1_3_8_9_4_5_7, 0.1_2_5_4_9_6_0_6, 0.1_2_4_4_9_6_4, 0.1_0_8_3_1_5_1_7, 0.1_1_5_6_2_8_6_6, 0.1_0_8_6_7_8_1_6, 0.1_0_4_9_9_0_4_8]		)

									assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2