


import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    WavaVecaConfig,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaForCTC,
    WavaVecaForPreTraining,
    WavaVecaProcessor,
    logging,
)
from transformers.models.wavaveca.modeling_wavaveca import WavaVecaForSequenceClassification


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''',
    '''adapter_layer''': '''encoder.layers.*.adapter_layer''',
    '''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''': '''lm_head''',
    '''mask_emb''': '''masked_spec_embed''',
    '''pooling_layer.linear''': '''projector''',
    '''pooling_layer.projection''': '''classifier''',
}
lowerCamelCase						=		[
    '''lm_head''',
    '''quantizer.weight_proj''',
    '''quantizer.codevectors''',
    '''project_q''',
    '''project_hid''',
    '''projector''',
    '''classifier''',
]






def 					lowerCamelCase_   (	_a	):




		"""simple docstring"""






		lowerCAmelCase__       :      Optional[Any]   					= {}
		with open(_a     ,       '''r'''	) as file:
				for line_number, line in enumerate(_a	):
						lowerCAmelCase__       :      Union[str, Any]   					= line.strip()
						if line:
								lowerCAmelCase__       :      Dict   					= line.split()
								lowerCAmelCase__       :      Tuple   					= line_number
								lowerCAmelCase__       :      List[str]   					= words[0]
								lowerCAmelCase__       :      List[str]   					= value
		return result






def 					lowerCamelCase_   (	_a     ,       _a     ,       _a     ,       _a     ,       _a	):




		"""simple docstring"""






		for attribute in key.split('''.'''	):
				lowerCAmelCase__       :      List[str]   					= getattr(_a     ,       _a	)

		lowerCAmelCase__       :      Dict   					= None
		for param_key in PARAM_MAPPING.keys():
				if full_name.endswith(_a	):
						lowerCAmelCase__       :      List[Any]   					= PARAM_MAPPING[full_name.split('''.'''	)[-1]]
						lowerCAmelCase__       :      Optional[Any]   					= '''param'''

		if weight_type is not None and weight_type != "param":
				lowerCAmelCase__       :      Dict   					= getattr(_a     ,       _a	).shape
		elif weight_type is not None and weight_type == "param":
				lowerCAmelCase__       :      Dict   					= hf_pointer
				for attribute in hf_param_name.split('''.'''	):
						lowerCAmelCase__       :      Union[str, Any]   					= getattr(_a     ,       _a	)
				lowerCAmelCase__       :      List[str]   					= shape_pointer.shape

				# let's reduce dimension
				lowerCAmelCase__       :      Dict   					= value[0]
		else:
				lowerCAmelCase__       :      Dict   					= hf_pointer.shape

		if hf_shape != value.shape:
				raise ValueError(
				    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":
				lowerCAmelCase__       :      int   					= value
		elif weight_type == "weight_g":
				lowerCAmelCase__       :      List[str]   					= value
		elif weight_type == "weight_v":
				lowerCAmelCase__       :      Any   					= value
		elif weight_type == "bias":
				lowerCAmelCase__       :      Dict   					= value
		elif weight_type == "param":
				for attribute in hf_param_name.split('''.'''	):
						lowerCAmelCase__       :      Union[str, Any]   					= getattr(_a     ,       _a	)
				lowerCAmelCase__       :      Optional[int]   					= value
		else:
				lowerCAmelCase__       :      int   					= value

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






def 					lowerCamelCase_   (	_a     ,       _a     ,       _a     ,       _a     ,       _a	):




		"""simple docstring"""






		lowerCAmelCase__       :      Optional[int]   					= None
		for param_key in PARAM_MAPPING.keys():
				if full_name.endswith(_a	):
						lowerCAmelCase__       :      List[str]   					= PARAM_MAPPING[full_name.split('''.'''	)[-1]]
						lowerCAmelCase__       :      List[str]   					= '''param'''

		if weight_type is not None and weight_type != "param":
				lowerCAmelCase__       :      int   					= '''.'''.join([key, weight_type]	)
		elif weight_type is not None and weight_type == "param":
				lowerCAmelCase__       :      Dict   					= '''.'''.join([key, hf_param_name]	)
		else:
				lowerCAmelCase__       :      Optional[int]   					= key

		lowerCAmelCase__       :      int   					= value if '''lm_head''' in full_key else value[0]


lowerCamelCase						=		{
    '''W_a''': '''linear_1.weight''',
    '''W_b''': '''linear_2.weight''',
    '''b_a''': '''linear_1.bias''',
    '''b_b''': '''linear_2.bias''',
    '''ln_W''': '''norm.weight''',
    '''ln_b''': '''norm.bias''',
}






def 					lowerCamelCase_   (	_a     ,       _a     ,       _a=None     ,       _a=None	):




		"""simple docstring"""






		lowerCAmelCase__       :      Any   					= False
		for key, mapped_key in MAPPING.items():
				lowerCAmelCase__       :      int   					= '''wav2vec2.''' + 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]:
						lowerCAmelCase__       :      Union[str, Any]   					= True
						if "*" in mapped_key:
								lowerCAmelCase__       :      List[Any]   					= name.split(_a	)[0].split('''.'''	)[-2]
								lowerCAmelCase__       :      Optional[Any]   					= mapped_key.replace('''*'''     ,       _a	)
						if "weight_g" in name:
								lowerCAmelCase__       :      Dict   					= '''weight_g'''
						elif "weight_v" in name:
								lowerCAmelCase__       :      int   					= '''weight_v'''
						elif "bias" in name:
								lowerCAmelCase__       :      int   					= '''bias'''
						elif "weight" in name:
								# TODO: don't match quantizer.weight_proj
								lowerCAmelCase__       :      Dict   					= '''weight'''
						else:
								lowerCAmelCase__       :      Optional[int]   					= None
						if hf_dict is not None:
								rename_dict(_a     ,       _a     ,       _a     ,       _a     ,       _a	)
						else:
								set_recursively(_a     ,       _a     ,       _a     ,       _a     ,       _a	)
						return is_used
		return is_used






def 					lowerCamelCase_   (	_a     ,       _a     ,       _a	):




		"""simple docstring"""






		lowerCAmelCase__       :      Dict   					= []
		lowerCAmelCase__       :      Dict   					= fairseq_model.state_dict()

		lowerCAmelCase__       :      Optional[Any]   					= hf_model.wavaveca.feature_extractor

		for name, value in fairseq_dict.items():
				lowerCAmelCase__       :      int   					= False
				if "conv_layers" in name:
						load_conv_layer(
						    _a     ,       _a     ,       _a     ,       _a     ,       hf_model.config.feat_extract_norm == '''group'''     ,       )
						lowerCAmelCase__       :      Dict   					= True
				else:
						lowerCAmelCase__       :      Dict   					= load_wavaveca_layer(_a     ,       _a     ,       _a	)
				if not is_used:
						unused_weights.append(_a	)

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






def 					lowerCamelCase_   (	_a     ,       _a     ,       _a     ,       _a     ,       _a	):




		"""simple docstring"""






		lowerCAmelCase__       :      int   					= full_name.split('''conv_layers.'''	)[-1]
		lowerCAmelCase__       :      Dict   					= name.split('''.'''	)
		lowerCAmelCase__       :      Dict   					= int(items[0]	)
		lowerCAmelCase__       :      str   					= int(items[1]	)

		if type_id == 0:
				if "bias" in name:
						if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
								raise ValueError(
								    f'{full_name} has size {value.shape}, but'
								    f' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'	)
						lowerCAmelCase__       :      Any   					= value
						logger.info(f'Feat extract conv layer {layer_id} was initialized from {full_name}.'	)
				elif "weight" in name:
						if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
								raise ValueError(
								    f'{full_name} has size {value.shape}, but'
								    f' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'	)
						lowerCAmelCase__       :      Union[str, Any]   					= 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:
						if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
								raise ValueError(
								    f'{full_name} has size {value.shape}, but'
								    f' {feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape} was found.'	)
						lowerCAmelCase__       :      str   					= value
						logger.info(f'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'	)
				elif "weight" in name:
						if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
								raise ValueError(
								    f'{full_name} has size {value.shape}, but'
								    f' {feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape} was found.'	)
						lowerCAmelCase__       :      Union[str, Any]   					= value
						logger.info(f'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'	)
		else:
				unused_weights.append(_a	)






@torch.no_grad()
def 					lowerCamelCase_   (	_a     ,       _a     ,       _a=None     ,       _a=None     ,       _a=True     ,       _a=False	):




		"""simple docstring"""






		if config_path is not None:
				lowerCAmelCase__       :      int   					= WavaVecaConfig.from_pretrained(_a	)
		else:
				lowerCAmelCase__       :      List[str]   					= WavaVecaConfig()

		if is_seq_class:
				lowerCAmelCase__       :      List[Any]   					= read_txt_into_dict(_a	)
				lowerCAmelCase__       :      Dict   					= idalabel
				lowerCAmelCase__       :      List[Any]   					= WavaVecaForSequenceClassification(_a	)
				lowerCAmelCase__       :      int   					= WavaVecaFeatureExtractor(
				    feature_size=1     ,       sampling_rate=16_000     ,       padding_value=0     ,       do_normalize=_a     ,       return_attention_mask=_a     ,       )
				feature_extractor.save_pretrained(_a	)

		elif is_finetuned:
				if dict_path:
						lowerCAmelCase__       :      Dict   					= Dictionary.load(_a	)

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

						# fairseq has the <pad> and <s> switched
						lowerCAmelCase__       :      Optional[Any]   					= 0
						lowerCAmelCase__       :      Tuple   					= 1
						with open(_a     ,       '''w'''     ,       encoding='''utf-8'''	) as vocab_handle:
								json.dump(_a     ,       _a	)
						lowerCAmelCase__       :      Dict   					= WavaVecaCTCTokenizer(
						    _a     ,       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=_a     ,       )
						lowerCAmelCase__       :      str   					= True if config.feat_extract_norm == '''layer''' else False
						lowerCAmelCase__       :      Dict   					= WavaVecaFeatureExtractor(
						    feature_size=1     ,       sampling_rate=16_000     ,       padding_value=0     ,       do_normalize=_a     ,       return_attention_mask=_a     ,       )
						lowerCAmelCase__       :      List[str]   					= WavaVecaProcessor(feature_extractor=_a     ,       tokenizer=_a	)
						processor.save_pretrained(_a	)

				lowerCAmelCase__       :      Union[str, Any]   					= WavaVecaForCTC(_a	)
		else:
				lowerCAmelCase__       :      Union[str, Any]   					= WavaVecaForPreTraining(_a	)

		if is_finetuned or is_seq_class:
				lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__       :      Optional[int]   					= fairseq.checkpoint_utils.load_model_ensemble_and_task(
				    [checkpoint_path]     ,       arg_overrides={'''data''': '''/'''.join(dict_path.split('''/'''	)[:-1]	)}	)
		else:
				lowerCAmelCase__       :      Optional[Any]   					= argparse.Namespace(task='''audio_pretraining'''	)
				lowerCAmelCase__       :      Optional[Any]   					= fairseq.tasks.setup_task(_a	)

				lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__       :      Optional[Any]   					= fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]     ,       task=_a	)

		lowerCAmelCase__       :      List[str]   					= model[0].eval()

		recursively_load_weights(_a     ,       _a     ,       not is_finetuned	)

		hf_wavavec.save_pretrained(_a	)


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'''
		)
		parser.add_argument(
		    '''--is_seq_class''',
		    action='''store_true''',
		    help='''Whether the model to convert is a fine-tuned sequence classification model or not''',
		)
		lowerCamelCase						=		parser.parse_args()

		lowerCamelCase						=		not args.not_finetuned and not args.is_seq_class
		convert_wavaveca_checkpoint(
		    args.checkpoint_path,
		    args.pytorch_dump_folder_path,
		    args.config_path,
		    args.dict_path,
		    is_finetuned,
		    args.is_seq_class,
		)








import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionPipeline
from diffusers.utils.testing_utils import load_image, nightly, require_torch_gpu, torch_device


lowerCamelCase						=		False





class  _a							(    unittest.TestCase):
	pass





@nightly
@require_torch_gpu
class  _a							(    unittest.TestCase):




	def 							UpperCAmelCase__( self						:		int		)->	int:
			# clean up the VRAM after each test
			super().tearDown()
			gc.collect()
			torch.cuda.empty_cache()




	def 							UpperCAmelCase__( self						:		Any		)->	Dict:
			lowerCAmelCase__       :      int   					= VersatileDiffusionPipeline.from_pretrained('''shi-labs/versatile-diffusion'''		, torch_dtype=torch.floataa		)
			pipe.to(_SCREAMING_SNAKE_CASE		)
			pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

			lowerCAmelCase__       :      List[str]   					= load_image(
			    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'''		)

			lowerCAmelCase__       :      List[str]   					= torch.manual_seed(0		)
			lowerCAmelCase__       :      List[str]   					= pipe.dual_guided(
			    prompt='''first prompt'''		, image=_SCREAMING_SNAKE_CASE		, text_to_image_strength=0.75		, generator=_SCREAMING_SNAKE_CASE		, guidance_scale=7.5		, num_inference_steps=2		, output_type='''numpy'''		, ).images

			with tempfile.TemporaryDirectory() as tmpdirname:
					pipe.save_pretrained(_SCREAMING_SNAKE_CASE		)
					lowerCAmelCase__       :      List[Any]   					= VersatileDiffusionPipeline.from_pretrained(_SCREAMING_SNAKE_CASE		, torch_dtype=torch.floataa		)
			pipe.to(_SCREAMING_SNAKE_CASE		)
			pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

			lowerCAmelCase__       :      List[str]   					= generator.manual_seed(0		)
			lowerCAmelCase__       :      str   					= pipe.dual_guided(
			    prompt='''first prompt'''		, image=_SCREAMING_SNAKE_CASE		, text_to_image_strength=0.75		, generator=_SCREAMING_SNAKE_CASE		, guidance_scale=7.5		, num_inference_steps=2		, output_type='''numpy'''		, ).images

			assert np.abs(image - new_image		).sum() < 1E-5, "Models don't have the same forward pass"




	def 							UpperCAmelCase__( self						:		Tuple		)->	List[Any]:
			lowerCAmelCase__       :      Optional[int]   					= VersatileDiffusionPipeline.from_pretrained('''shi-labs/versatile-diffusion'''		, torch_dtype=torch.floataa		)
			pipe.to(_SCREAMING_SNAKE_CASE		)
			pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

			lowerCAmelCase__       :      int   					= '''cyberpunk 2077'''
			lowerCAmelCase__       :      List[str]   					= load_image(
			    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'''		)
			lowerCAmelCase__       :      List[Any]   					= torch.manual_seed(0		)
			lowerCAmelCase__       :      Union[str, Any]   					= pipe.dual_guided(
			    prompt=_SCREAMING_SNAKE_CASE		, image=_SCREAMING_SNAKE_CASE		, text_to_image_strength=0.75		, generator=_SCREAMING_SNAKE_CASE		, guidance_scale=7.5		, num_inference_steps=50		, output_type='''numpy'''		, ).images

			lowerCAmelCase__       :      List[Any]   					= image[0, 253:256, 253:256, -1]

			assert image.shape == (1, 512, 512, 3)
			lowerCAmelCase__       :      List[str]   					= np.array([0.1448, 0.1619, 0.1741, 0.1086, 0.1147, 0.1128, 0.1199, 0.1165, 0.1001]		)

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

			lowerCAmelCase__       :      int   					= '''A painting of a squirrel eating a burger '''
			lowerCAmelCase__       :      Optional[Any]   					= torch.manual_seed(0		)
			lowerCAmelCase__       :      Union[str, Any]   					= pipe.text_to_image(
			    prompt=_SCREAMING_SNAKE_CASE		, generator=_SCREAMING_SNAKE_CASE		, guidance_scale=7.5		, num_inference_steps=50		, output_type='''numpy'''		).images

			lowerCAmelCase__       :      List[Any]   					= image[0, 253:256, 253:256, -1]

			assert image.shape == (1, 512, 512, 3)
			lowerCAmelCase__       :      Optional[Any]   					= np.array([0.3367, 0.3169, 0.2656, 0.3870, 0.4790, 0.3796, 0.4009, 0.4878, 0.4778]		)

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

			lowerCAmelCase__       :      Optional[Any]   					= pipe.image_variation(_SCREAMING_SNAKE_CASE		, generator=_SCREAMING_SNAKE_CASE		, output_type='''numpy'''		).images

			lowerCAmelCase__       :      List[Any]   					= image[0, 253:256, 253:256, -1]

			assert image.shape == (1, 512, 512, 3)
			lowerCAmelCase__       :      Any   					= np.array([0.3076, 0.3123, 0.3284, 0.3782, 0.3770, 0.3894, 0.4297, 0.4331, 0.4456]		)

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




