import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, PNDMScheduler, StableDiffusionInpaintPipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





class UpperCAmelCase						( __A ,	__A ,	__A ,	unittest.TestCase  ):


				'''simple docstring'''






				lowerCamelCase_    = StableDiffusionInpaintPipeline
				lowerCamelCase_    = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
				lowerCamelCase_    = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
				lowerCamelCase_    = frozenset(
				    []  )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
				lowerCamelCase_    = frozenset([]  )
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						torch.manual_seed(0 )
						A_		:    Optional[Any]         =				UNetaDConditionModel(
						    block_out_channels=(3_2, 6_4)		,     layers_per_block=2		,     sample_size=3_2		,     in_channels=9		,     out_channels=4		,     down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')		,     up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')		,     cross_attention_dim=3_2		,     attention_head_dim=(2, 4)		,     use_linear_projection=lowercase		,     )
						A_		:    str         =				PNDMScheduler(skip_prk_steps=lowercase )
						torch.manual_seed(0 )
						A_		:    Dict         =				AutoencoderKL(
						    block_out_channels=[3_2, 6_4]		,     in_channels=3		,     out_channels=3		,     down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']		,     up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']		,     latent_channels=4		,     sample_size=1_2_8		,     )
						torch.manual_seed(0 )
						A_		:    Dict         =				CLIPTextConfig(
						    bos_token_id=0		,     eos_token_id=2		,     hidden_size=3_2		,     intermediate_size=3_7		,     layer_norm_eps=1E-05		,     num_attention_heads=4		,     num_hidden_layers=5		,     pad_token_id=1		,     vocab_size=1_0_0_0		,     hidden_act='gelu'		,     projection_dim=5_1_2		,     )
						A_		:    Dict         =				CLIPTextModel(lowercase )
						A_		:    Dict         =				CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

						A_		:    Any         =				{
						    'unet': unet,
						    'scheduler': scheduler,
						    'vae': vae,
						    'text_encoder': text_encoder,
						    'tokenizer': tokenizer,
						    'safety_checker': None,
						    'feature_extractor': None,
						}
						return components
				def   lowerCAmelCase_       (						self		,     lowercase		,     lowercase=0 ):



						"""simple docstring"""





						A_		:    int         =				floats_tensor((1, 3, 3_2, 3_2)		,     rng=random.Random(lowercase ) ).to(lowercase )
						A_		:    Any         =				image.cpu().permute(0		,     2		,     3		,     1 )[0]
						A_		:    int         =				Image.fromarray(np.uinta(lowercase ) ).convert('RGB' ).resize((6_4, 6_4) )
						A_		:    int         =				Image.fromarray(np.uinta(image + 4 ) ).convert('RGB' ).resize((6_4, 6_4) )
						if str(lowercase ).startswith('mps' ):
								A_		:    Tuple         =				torch.manual_seed(lowercase )
						else:
								A_		:    List[str]         =				torch.Generator(device=lowercase ).manual_seed(lowercase )
						A_		:    List[Any]         =				{
						    'prompt': 'A painting of a squirrel eating a burger',
						    'image': init_image,
						    'mask_image': mask_image,
						    'generator': generator,
						    'num_inference_steps': 2,
						    'guidance_scale': 6.0,
						    'output_type': 'numpy',
						}
						return inputs
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    Any         =				'cpu'  # ensure determinism for the device-dependent torch.Generator
						A_		:    str         =				self.get_dummy_components()
						A_		:    Any         =				StableDiffusionInpaintPipeline(**lowercase )
						A_		:    Tuple         =				sd_pipe.to(lowercase )
						sd_pipe.set_progress_bar_config(disable=lowercase )

						A_		:    Optional[Any]         =				self.get_dummy_inputs(lowercase )
						A_		:    Optional[int]         =				sd_pipe(**lowercase ).images
						A_		:    int         =				image[0, -3:, -3:, -1]

						assert image.shape == (1, 6_4, 6_4, 3)
						A_		:    str         =				np.array([0.4727, 0.5735, 0.3941, 0.5446, 0.5926, 0.4394, 0.5062, 0.4654, 0.4476] )

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







				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						super().test_inference_batch_single_identical(expected_max_diff=3E-3 )





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


				'''simple docstring'''






				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    Dict         =				load_image(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
						    '/sd2-inpaint/init_image.png' )
						A_		:    Dict         =				load_image(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png' )
						A_		:    List[str]         =				load_numpy(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint'
						    '/yellow_cat_sitting_on_a_park_bench.npy' )

						A_		:    Optional[Any]         =				'stabilityai/stable-diffusion-2-inpainting'
						A_		:    List[Any]         =				StableDiffusionInpaintPipeline.from_pretrained(lowercase		,     safety_checker=lowercase )
						pipe.to(lowercase )
						pipe.set_progress_bar_config(disable=lowercase )
						pipe.enable_attention_slicing()

						A_		:    Optional[int]         =				'Face of a yellow cat, high resolution, sitting on a park bench'

						A_		:    Optional[Any]         =				torch.manual_seed(0 )
						A_		:    str         =				pipe(
						    prompt=lowercase		,     image=lowercase		,     mask_image=lowercase		,     generator=lowercase		,     output_type='np'		,     )
						A_		:    int         =				output.images[0]

						assert image.shape == (5_1_2, 5_1_2, 3)
						assert np.abs(expected_image - image ).max() < 9E-3
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    int         =				load_image(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
						    '/sd2-inpaint/init_image.png' )
						A_		:    Tuple         =				load_image(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png' )
						A_		:    int         =				load_numpy(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint'
						    '/yellow_cat_sitting_on_a_park_bench_fp16.npy' )

						A_		:    List[Any]         =				'stabilityai/stable-diffusion-2-inpainting'
						A_		:    Union[str, Any]         =				StableDiffusionInpaintPipeline.from_pretrained(
						    lowercase		,     torch_dtype=torch.floataa		,     safety_checker=lowercase		,     )
						pipe.to(lowercase )
						pipe.set_progress_bar_config(disable=lowercase )
						pipe.enable_attention_slicing()

						A_		:    str         =				'Face of a yellow cat, high resolution, sitting on a park bench'

						A_		:    Dict         =				torch.manual_seed(0 )
						A_		:    List[str]         =				pipe(
						    prompt=lowercase		,     image=lowercase		,     mask_image=lowercase		,     generator=lowercase		,     output_type='np'		,     )
						A_		:    List[Any]         =				output.images[0]

						assert image.shape == (5_1_2, 5_1_2, 3)
						assert np.abs(expected_image - image ).max() < 5E-1







				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						torch.cuda.empty_cache()
						torch.cuda.reset_max_memory_allocated()
						torch.cuda.reset_peak_memory_stats()

						A_		:    Optional[int]         =				load_image(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
						    '/sd2-inpaint/init_image.png' )
						A_		:    Optional[int]         =				load_image(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png' )

						A_		:    Any         =				'stabilityai/stable-diffusion-2-inpainting'
						A_		:    Tuple         =				PNDMScheduler.from_pretrained(lowercase		,     subfolder='scheduler' )
						A_		:    List[Any]         =				StableDiffusionInpaintPipeline.from_pretrained(
						    lowercase		,     safety_checker=lowercase		,     scheduler=lowercase		,     torch_dtype=torch.floataa		,     )
						pipe.to(lowercase )
						pipe.set_progress_bar_config(disable=lowercase )
						pipe.enable_attention_slicing(1 )
						pipe.enable_sequential_cpu_offload()

						A_		:    Union[str, Any]         =				'Face of a yellow cat, high resolution, sitting on a park bench'

						A_		:    Tuple         =				torch.manual_seed(0 )
						A_		:    Tuple         =				pipe(
						    prompt=lowercase		,     image=lowercase		,     mask_image=lowercase		,     generator=lowercase		,     num_inference_steps=2		,     output_type='np'		,     )

						A_		:    Optional[Any]         =				torch.cuda.max_memory_allocated()
						# make sure that less than 2.65 GB is allocated
						assert mem_bytes < 2.65 * 1_0**9







import argparse

import fairseq
import torch
from torch import nn

from transformers import (
    MBartaaTokenizer,
    MBartConfig,
    MBartForCausalLM,
    SpeechEncoderDecoderConfig,
    SpeechEncoderDecoderModel,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaModel,
    logging,
)


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

_UpperCAmelCase            =		{
    """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""": """lm_head""",
    """mask_emb""": """masked_spec_embed""",
}
_UpperCAmelCase            =		[
    """lm_head""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
]




def     UpperCamelCase						(	__lowercase						:			Optional[Any]						,__lowercase						:			Tuple						,__lowercase						:			str						,__lowercase						:			List[Any]						,__lowercase						:			Optional[int]     ):
		'''simple docstring'''




		for attribute in key.split('.'     ):
				A_		:    List[Any]         =				getattr(__lowercase						,__lowercase     )

		if weight_type is not None:
				A_		:    Dict         =				getattr(__lowercase						,__lowercase     ).shape
		else:
				A_		:    Union[str, Any]         =				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":
				A_		:    Dict         =				value
		elif weight_type == "weight_g":
				A_		:    str         =				value
		elif weight_type == "weight_v":
				A_		:    int         =				value
		elif weight_type == "bias":
				A_		:    int         =				value
		else:
				A_		:    List[Any]         =				value

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




def     UpperCamelCase						(	__lowercase						:			Optional[Any]						,__lowercase						:			Union[str, Any]     ):
		'''simple docstring'''




		A_		:    List[Any]         =				[]
		A_		:    int         =				fairseq_model.state_dict()

		A_		:    Optional[Any]         =				hf_model.feature_extractor
		A_		:    List[str]         =				hf_model.adapter

		for name, value in fairseq_dict.items():
				A_		:    Tuple         =				False
				if "conv_layers" in name:
						load_conv_layer(
						    __lowercase						,__lowercase						,__lowercase						,__lowercase						,hf_model.config.feat_extract_norm == 'group'						,)
						A_		:    Optional[Any]         =				True
				elif any(x in name for x in ['adaptor', 'w2v_encoder.proj.', 'w2v_proj_ln.']     ):
						load_adapter(__lowercase						,__lowercase						,__lowercase						,__lowercase     )
						A_		:    Union[str, Any]         =				True
				else:
						for key, mapped_key in MAPPING.items():
								if key in name or key.split('w2v_model.'     )[-1] == name.split('.'     )[0]:
										A_		:    Tuple         =				True
										if "*" in mapped_key:
												A_		:    Optional[Any]         =				name.split(__lowercase     )[0].split('.'     )[-2]
												A_		:    List[Any]         =				mapped_key.replace('*'						,__lowercase     )
										if "weight_g" in name:
												A_		:    Optional[int]         =				'weight_g'
										elif "weight_v" in name:
												A_		:    Union[str, Any]         =				'weight_v'
										elif "bias" in name:
												A_		:    Any         =				'bias'
										elif "weight" in name:
												A_		:    str         =				'weight'
										else:
												A_		:    Optional[Any]         =				None
										set_recursively(__lowercase						,__lowercase						,__lowercase						,__lowercase						,__lowercase     )
								continue
				if not is_used:
						unused_weights.append(__lowercase     )

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




def     UpperCamelCase						(	__lowercase						:			Dict						,__lowercase						:			List[Any]						,__lowercase						:			Tuple						,__lowercase						:			Dict						,__lowercase						:			Any     ):
		'''simple docstring'''




		A_		:    List[Any]         =				full_name.split('conv_layers.'     )[-1]
		A_		:    Optional[int]         =				name.split('.'     )
		A_		:    Tuple         =				int(items[0]     )
		A_		:    int         =				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.'''
						)
						A_		:    Union[str, Any]         =				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.'''
						)
						A_		:    str         =				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."
						)
						A_		:    List[str]         =				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.'''
						)
						A_		:    List[Any]         =				value
						logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''     )
		else:
				unused_weights.append(__lowercase     )




def     UpperCamelCase						(	__lowercase						:			Any						,__lowercase						:			Tuple						,__lowercase						:			Optional[Any]						,__lowercase						:			Union[str, Any]     ):
		'''simple docstring'''




		A_		:    Union[str, Any]         =				full_name.split('adaptor.'     )[-1]
		A_		:    List[Any]         =				name.split('.'     )

		if items[1].isdigit():
				A_		:    Union[str, Any]         =				int(items[1]     )
		else:
				A_		:    Tuple         =				None

		if "adaptor" not in full_name:
				if "proj_ln" in full_name:
						# has to be layer norm
						if "bias" in name:
								assert (
								    value.shape == adapter.proj_layer_norm.bias.data.shape
								), f'''{full_name} has size {value.shape}, but {adapter.proj_layer_norm.bias.data.shape} was found.'''
								A_		:    Dict         =				value
								logger.info(f'''Adapter proj layer norm bias was initialized from {full_name}.'''     )
						if "weight" in name:
								assert (
								    value.shape == adapter.proj_layer_norm.weight.data.shape
								), f'''{full_name} has size {value.shape}, but {adapter.proj_layer_norm.weight.data.shape} was found.'''
								A_		:    int         =				value
				else:
						# has to be projection layer
						if "bias" in name:
								assert (
								    value.shape == adapter.proj.bias.data.shape
								), f'''{full_name} has size {value.shape}, but {adapter.proj.bias.data.shape} was found.'''
								A_		:    Dict         =				value
								logger.info(f'''Adapter proj layer bias was initialized from {full_name}.'''     )
						if "weight" in name:
								assert (
								    value.shape == adapter.proj.weight.data.shape
								), f'''{full_name} has size {value.shape}, but {adapter.proj.weight.data.shape} was found.'''
								A_		:    Tuple         =				value
								logger.info(f'''Adapter proj layer weight was initialized from {full_name}.'''     )
		elif isinstance(__lowercase						,__lowercase     ):
				if "bias" in name:
						assert (
						    value.shape == adapter.layers[layer_id].conv.bias.data.shape
						), f'''{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.bias.data.shape} was found.'''
						A_		:    Tuple         =				value
						logger.info(f'''Adapter layer {layer_id} bias was initialized from {full_name}.'''     )
				elif "weight" in name:
						assert (
						    value.shape == adapter.layers[layer_id].conv.weight.data.shape
						), f'''{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.weight.data.shape} was found.'''
						A_		:    str         =				value
						logger.info(f'''Adapter layer {layer_id} bias was initialized from {full_name}.'''     )
		else:
				unused_weights.append(__lowercase     )




def     UpperCamelCase						(	__lowercase						:			List[Any]     ):
		'''simple docstring'''




		A_       ,      A_		:    Any         =				emb.weight.shape
		A_		:    Tuple         =				nn.Linear(__lowercase						,__lowercase						,bias=__lowercase     )
		A_		:    Optional[Any]         =				emb.weight.data
		return lin_layer




@torch.no_grad()
def     UpperCamelCase						(	__lowercase						:			Any						,__lowercase						:			Optional[int]						,__lowercase						:			Any						,__lowercase						:			str						,__lowercase						:			Dict						,__lowercase						:			Dict						,__lowercase						:			Tuple						,__lowercase						:			Optional[int]						,__lowercase						:			List[str]						,__lowercase						:			List[Any]						,__lowercase						:			str						,):
		'''simple docstring'''




		A_		:    Optional[int]         =				WavaVecaConfig.from_pretrained(
		    __lowercase						,add_adapter=__lowercase						,adapter_stride=__lowercase						,adapter_kernel_size=__lowercase						,use_auth_token=__lowercase						,output_hidden_size=__lowercase						,)
		A_		:    Any         =				MBartConfig.from_pretrained(__lowercase     )

		# load model
		A_       ,      A_       ,      A_		:    Union[str, Any]         =				fairseq.checkpoint_utils.load_model_ensemble_and_task(
		    [checkpoint_path]						,arg_overrides={
		        'config_yaml': config_yaml_path,
		        'data': '/'.join(dict_path.split('/'     )[:-1]     ),
		        'w2v_path': checkpoint_path,
		        'load_pretrained_decoder_from': None,
		    }						,)
		A_		:    Union[str, Any]         =				model[0].eval()

		# load feature extractor
		A_		:    Any         =				WavaVecaFeatureExtractor.from_pretrained(__lowercase						,use_auth_token=__lowercase     )

		# set weights for wav2vec2 encoder
		A_		:    Optional[Any]         =				WavaVecaModel(__lowercase     )

		recursively_load_weights_wavaveca(model.encoder						,__lowercase     )

		# load decoder weights
		A_		:    Dict         =				MBartForCausalLM(__lowercase     )
		A_       ,      A_		:    Union[str, Any]         =				hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()						,strict=__lowercase     )
		logger.warning(f'''The following keys are missing when loading the decoder weights: {missing_keys}'''     )
		logger.warning(f'''The following keys are unexpected when loading the decoder weights: {unexpected_keys}'''     )

		A_		:    Optional[int]         =				SpeechEncoderDecoderModel(encoder=__lowercase						,decoder=__lowercase     )
		A_		:    Any         =				False

		A_		:    List[Any]         =				MBartaaTokenizer(__lowercase     )
		tokenizer.save_pretrained(__lowercase     )

		A_		:    Dict         =				hf_wavavec.config.to_dict()
		A_		:    Any         =				tokenizer.pad_token_id
		A_		:    Optional[Any]         =				tokenizer.bos_token_id
		A_		:    Union[str, Any]         =				tokenizer.eos_token_id
		A_		:    Dict         =				'mbart50'
		A_		:    str         =				'wav2vec2'

		A_		:    int         =				tokenizer.eos_token_id
		A_		:    List[str]         =				25_00_04
		A_		:    int         =				tokenizer.eos_token_id

		A_		:    Optional[Any]         =				SpeechEncoderDecoderConfig.from_dict(__lowercase     )

		hf_wavavec.save_pretrained(__lowercase     )
		feature_extractor.save_pretrained(__lowercase     )


if __name__ == "__main__":
	_UpperCAmelCase            =		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_yaml_path""", default=None, type=str, help="""Path to yaml file of fine-tuned model""")
	parser.add_argument(
	    """--encoder_config_path""",
	    default="""facebook/wav2vec2-xls-r-1b""",
	    type=str,
	    help="""Path to hf encoder wav2vec2 checkpoint config""",
	)
	parser.add_argument(
	    """--decoder_config_path""",
	    default="""facebook/mbart-large-50-one-to-many-mmt""",
	    type=str,
	    help="""Path to hf decoder checkpoint config""",
	)
	parser.add_argument("""--add_adapter""", default=True, type=bool, help="""whethere to add model adapter layers""")
	parser.add_argument("""--adapter_stride""", default=2, type=int, help="""stride of adapter layers""")
	parser.add_argument("""--adapter_kernel_size""", default=3, type=int, help="""kernel size of adapter layers""")
	parser.add_argument("""--encoder_output_dim""", default=1024, type=int, help="""encoder output dim""")
	parser.add_argument("""--start_token_id""", default=250004, type=int, help="""`decoder_start_token_id` of model config""")

	_UpperCAmelCase            =		parser.parse_args()
	convert_wavaveca_checkpoint(
	    args.checkpoint_path,
	    args.pytorch_dump_folder_path,
	    args.dict_path,
	    args.config_yaml_path,
	    encoder_config_path=args.encoder_config_path,
	    decoder_config_path=args.decoder_config_path,
	    add_adapter=args.add_adapter,
	    adapter_kernel_size=args.adapter_kernel_size,
	    adapter_stride=args.adapter_stride,
	    decoder_start_token_id=args.start_token_id,
	    encoder_output_dim=args.encoder_output_dim,
	)






