




import random
import unittest

import numpy as np
import torch

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionUpscalePipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
					import onnxruntime as ort







class     __lowerCAmelCase (       __magic_name__			,				unittest.TestCase ):

			"""simple docstring"""







			snake_case_   		=       '''ssube/stable-diffusion-x4-upscaler-onnx'''


			def       lowercase_    (      self			,						lowerCamelCase__=0							)    ->      List[Any]:



							'''simple docstring'''
							__lowerCamelCase		       =    floats_tensor((1, 3, 128, 128)			,						rng=random.Random(lowerCamelCase__							)							)
							__lowerCamelCase		       =    torch.manual_seed(lowerCamelCase__							)
							__lowerCamelCase		       =    {
							    'prompt': 'A painting of a squirrel eating a burger',
							    'image': image,
							    'generator': generator,
							    'num_inference_steps': 3,
							    'guidance_scale': 7.5,
							    'output_type': 'numpy',
							}
							return inputs


			def       lowercase_    (      self							)    ->      List[Any]:



							'''simple docstring'''
							__lowerCamelCase		       =    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint			,						provider='CPUExecutionProvider'							)
							pipe.set_progress_bar_config(disable=lowerCamelCase__							)

							__lowerCamelCase		       =    self.get_dummy_inputs()
							__lowerCamelCase		       =    pipe(**lowerCamelCase__							).images
							__lowerCamelCase		       =    image[0, -3:, -3:, -1].flatten()

							# started as 128, should now be 512
							assert image.shape == (1, 512, 512, 3)
							__lowerCamelCase		       =    np.array(
							    [0.6_97_47_82, 0.68_90_20_93, 0.70_13_58_85, 0.7_58_36_18, 0.7_80_45_45, 0.7_85_49_12, 0.78_66_74_26, 0.78_74_38_63, 0.78_07_02_23]							)
							assert np.abs(image_slice - expected_slice							).max() < 1e-1


			def       lowercase_    (      self							)    ->      Optional[Any]:



							'''simple docstring'''
							__lowerCamelCase		       =    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint			,						provider='CPUExecutionProvider'							)
							__lowerCamelCase		       =    PNDMScheduler.from_config(pipe.scheduler.config			,						skip_prk_steps=lowerCamelCase__							)
							pipe.set_progress_bar_config(disable=lowerCamelCase__							)

							__lowerCamelCase		       =    self.get_dummy_inputs()
							__lowerCamelCase		       =    pipe(**lowerCamelCase__							).images
							__lowerCamelCase		       =    image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							__lowerCamelCase		       =    np.array(
							    [0.6_89_88_92, 0.59_24_05_56, 0.52_49_95_27, 0.58_86_62_15, 0.52_25_82_35, 0.52_57_27_15, 0.62_41_44_73, 0.6_17_43_87, 0.6_21_49_64]							)
							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-1


			def       lowercase_    (      self							)    ->      str:



							'''simple docstring'''
							__lowerCamelCase		       =    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint			,						provider='CPUExecutionProvider'							)
							__lowerCamelCase		       =    DPMSolverMultistepScheduler.from_config(pipe.scheduler.config							)
							pipe.set_progress_bar_config(disable=lowerCamelCase__							)

							__lowerCamelCase		       =    self.get_dummy_inputs()
							__lowerCamelCase		       =    pipe(**lowerCamelCase__							).images
							__lowerCamelCase		       =    image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							__lowerCamelCase		       =    np.array(
							    [0.7_65_92_78, 0.76_43_76_64, 0.75_57_91_07, 0.7_69_11_16, 0.77_66_69_86, 0.7_72_76_72, 0.7_75_86_64, 0.7_81_22_26, 0.76_94_25_15]							)

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


			def       lowercase_    (      self							)    ->      int:



							'''simple docstring'''
							__lowerCamelCase		       =    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint			,						provider='CPUExecutionProvider'							)
							__lowerCamelCase		       =    EulerDiscreteScheduler.from_config(pipe.scheduler.config							)
							pipe.set_progress_bar_config(disable=lowerCamelCase__							)

							__lowerCamelCase		       =    self.get_dummy_inputs()
							__lowerCamelCase		       =    pipe(**lowerCamelCase__							).images
							__lowerCamelCase		       =    image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							__lowerCamelCase		       =    np.array(
							    [0.6_97_47_82, 0.68_90_20_93, 0.70_13_58_85, 0.7_58_36_18, 0.7_80_45_45, 0.7_85_49_12, 0.78_66_74_26, 0.78_74_38_63, 0.78_07_02_23]							)
							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-1






			def       lowercase_    (      self							)    ->      List[str]:



							'''simple docstring'''
							__lowerCamelCase		       =    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint			,						provider='CPUExecutionProvider'							)
							__lowerCamelCase		       =    EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config							)
							pipe.set_progress_bar_config(disable=lowerCamelCase__							)

							__lowerCamelCase		       =    self.get_dummy_inputs()
							__lowerCamelCase		       =    pipe(**lowerCamelCase__							).images
							__lowerCamelCase		       =    image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							__lowerCamelCase		       =    np.array(
							    [0.77_42_44_96, 0.77_36_01, 0.7_64_52_88, 0.7_76_95_98, 0.7_77_27_39, 0.7_73_86_88, 0.78_18_72_33, 0.77_87_95_84, 0.76_70_43]							)

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







@nightly
@require_onnxruntime
@require_torch_gpu
class     __lowerCAmelCase (       unittest.TestCase ):

			"""simple docstring"""







			@property
			def       lowercase_    (      self							)    ->      Any:



							'''simple docstring'''
							return (
							    "CUDAExecutionProvider",
							    {
							        "gpu_mem_limit": "15000000000",  # 15GB
							        "arena_extend_strategy": "kSameAsRequested",
							    },
							)


			@property
			def       lowercase_    (      self							)    ->      Any:



							'''simple docstring'''
							__lowerCamelCase		       =    ort.SessionOptions()
							__lowerCamelCase		       =    False
							return options


			def       lowercase_    (      self							)    ->      Tuple:



							'''simple docstring'''
							__lowerCamelCase		       =    load_image(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
							    '/img2img/sketch-mountains-input.jpg'							)
							__lowerCamelCase		       =    init_image.resize((128, 128)							)
							# using the PNDM scheduler by default
							__lowerCamelCase		       =    OnnxStableDiffusionUpscalePipeline.from_pretrained(
							    'ssube/stable-diffusion-x4-upscaler-onnx'			,						provider=self.gpu_provider			,						sess_options=self.gpu_options			,						)
							pipe.set_progress_bar_config(disable=lowerCamelCase__							)

							__lowerCamelCase		       =    'A fantasy landscape, trending on artstation'

							__lowerCamelCase		       =    torch.manual_seed(0							)
							__lowerCamelCase		       =    pipe(
							    prompt=lowerCamelCase__			,						image=lowerCamelCase__			,						guidance_scale=7.5			,						num_inference_steps=10			,						generator=lowerCamelCase__			,						output_type='np'			,						)
							__lowerCamelCase		       =    output.images
							__lowerCamelCase		       =    images[0, 255:258, 383:386, -1]

							assert images.shape == (1, 512, 512, 3)
							__lowerCamelCase		       =    np.array([0.48_83, 0.49_47, 0.49_80, 0.49_75, 0.49_82, 0.49_80, 0.50_00, 0.50_06, 0.49_72]							)
							# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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






			def       lowercase_    (      self							)    ->      Tuple:



							'''simple docstring'''
							__lowerCamelCase		       =    load_image(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
							    '/img2img/sketch-mountains-input.jpg'							)
							__lowerCamelCase		       =    init_image.resize((128, 128)							)
							__lowerCamelCase		       =    LMSDiscreteScheduler.from_pretrained(
							    'ssube/stable-diffusion-x4-upscaler-onnx'			,						subfolder='scheduler'							)
							__lowerCamelCase		       =    OnnxStableDiffusionUpscalePipeline.from_pretrained(
							    'ssube/stable-diffusion-x4-upscaler-onnx'			,						scheduler=lowerCamelCase__			,						provider=self.gpu_provider			,						sess_options=self.gpu_options			,						)
							pipe.set_progress_bar_config(disable=lowerCamelCase__							)

							__lowerCamelCase		       =    'A fantasy landscape, trending on artstation'

							__lowerCamelCase		       =    torch.manual_seed(0							)
							__lowerCamelCase		       =    pipe(
							    prompt=lowerCamelCase__			,						image=lowerCamelCase__			,						guidance_scale=7.5			,						num_inference_steps=20			,						generator=lowerCamelCase__			,						output_type='np'			,						)
							__lowerCamelCase		       =    output.images
							__lowerCamelCase		       =    images[0, 255:258, 383:386, -1]

							assert images.shape == (1, 512, 512, 3)
							__lowerCamelCase		       =    np.array(
							    [0.50_17_37_53, 0.50_22_33_56, 0.50_20_39, 0.50_23_30_36, 0.5_02_37_25, 0.5_02_26_01, 0.5_01_87_58, 0.50_23_40_85, 0.50_24_15_66]							)
							# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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








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()
__A        =    logging.get_logger(__name__)

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







def 			lowerCamelCase_		(       UpperCamelCase__				:      Any ,     UpperCamelCase__				:      Any ,     UpperCamelCase__				:      Optional[Any] ,     UpperCamelCase__				:      Union[str, Any] ,     UpperCamelCase__				:      Optional[int]      )     ->       Tuple:





				"""simple docstring"""


				for attribute in key.split('.'      ):
								__lowerCamelCase		       =    getattr(UpperCamelCase__ ,     UpperCamelCase__      )

				if weight_type is not None:
								__lowerCamelCase		       =    getattr(UpperCamelCase__ ,     UpperCamelCase__      ).shape
				else:
								__lowerCamelCase		       =    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":
								__lowerCamelCase		       =    value
				elif weight_type == "weight_g":
								__lowerCamelCase		       =    value
				elif weight_type == "weight_v":
								__lowerCamelCase		       =    value
				elif weight_type == "bias":
								__lowerCamelCase		       =    value
				else:
								__lowerCamelCase		       =    value

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







def 			lowerCamelCase_		(       UpperCamelCase__				:      Optional[Any] ,     UpperCamelCase__				:      Tuple      )     ->       Optional[Any]:





				"""simple docstring"""


				__lowerCamelCase		       =    []
				__lowerCamelCase		       =    fairseq_model.state_dict()

				__lowerCamelCase		       =    hf_model.feature_extractor
				__lowerCamelCase		       =    hf_model.adapter

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

				logger.warning(F"""Unused weights: {unused_weights}"""      )







def 			lowerCamelCase_		(       UpperCamelCase__				:      Dict ,     UpperCamelCase__				:      Any ,     UpperCamelCase__				:      str ,     UpperCamelCase__				:      Tuple ,     UpperCamelCase__				:      Tuple      )     ->       int:





				"""simple docstring"""


				__lowerCamelCase		       =    full_name.split('conv_layers.'      )[-1]
				__lowerCamelCase		       =    name.split('.'      )
				__lowerCamelCase		       =    int(items[0]      )
				__lowerCamelCase		       =    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."""
												)
												__lowerCamelCase		       =    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."""
												)
												__lowerCamelCase		       =    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."
												)
												__lowerCamelCase		       =    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."""
												)
												__lowerCamelCase		       =    value
												logger.info(F"""Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."""      )
				else:
								unused_weights.append(UpperCamelCase__      )







def 			lowerCamelCase_		(       UpperCamelCase__				:      Tuple ,     UpperCamelCase__				:      Optional[Any] ,     UpperCamelCase__				:      Any ,     UpperCamelCase__				:      int      )     ->       Union[str, Any]:





				"""simple docstring"""


				__lowerCamelCase		       =    full_name.split('adaptor.'      )[-1]
				__lowerCamelCase		       =    name.split('.'      )

				if items[1].isdigit():
								__lowerCamelCase		       =    int(items[1]      )
				else:
								__lowerCamelCase		       =    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."""
																__lowerCamelCase		       =    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."""
																__lowerCamelCase		       =    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."""
																__lowerCamelCase		       =    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."""
																__lowerCamelCase		       =    value
																logger.info(F"""Adapter proj layer weight was initialized from {full_name}."""      )
				elif isinstance(UpperCamelCase__ ,     UpperCamelCase__      ):
								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."""
												__lowerCamelCase		       =    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."""
												__lowerCamelCase		       =    value
												logger.info(F"""Adapter layer {layer_id} bias was initialized from {full_name}."""      )
				else:
								unused_weights.append(UpperCamelCase__      )







def 			lowerCamelCase_		(       UpperCamelCase__				:      Tuple      )     ->       Tuple:





				"""simple docstring"""


				__lowerCamelCase       ,						__lowerCamelCase		       =    emb.weight.shape
				__lowerCamelCase		       =    nn.Linear(UpperCamelCase__ ,     UpperCamelCase__ ,     bias=UpperCamelCase__      )
				__lowerCamelCase		       =    emb.weight.data
				return lin_layer







@torch.no_grad()
def 			lowerCamelCase_		(       UpperCamelCase__				:      List[str] ,     UpperCamelCase__				:      List[Any] ,     UpperCamelCase__				:      Union[str, Any] ,     UpperCamelCase__				:      Dict ,     UpperCamelCase__				:      List[str] ,     UpperCamelCase__				:      Union[str, Any] ,     UpperCamelCase__				:      Optional[int] ,     UpperCamelCase__				:      Union[str, Any] ,     UpperCamelCase__				:      Optional[Any] ,     UpperCamelCase__				:      str ,     UpperCamelCase__				:      Optional[Any] ,     )     ->       str:





				"""simple docstring"""


				__lowerCamelCase		       =    WavaVecaConfig.from_pretrained(
				    UpperCamelCase__ ,     add_adapter=UpperCamelCase__ ,     adapter_stride=UpperCamelCase__ ,     adapter_kernel_size=UpperCamelCase__ ,     use_auth_token=UpperCamelCase__ ,     output_hidden_size=UpperCamelCase__ ,     )
				__lowerCamelCase		       =    MBartConfig.from_pretrained(UpperCamelCase__      )

				# load model
				__lowerCamelCase       ,						__lowerCamelCase       ,						__lowerCamelCase		       =    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,
				    } ,     )
				__lowerCamelCase		       =    model[0].eval()

				# load feature extractor
				__lowerCamelCase		       =    WavaVecaFeatureExtractor.from_pretrained(UpperCamelCase__ ,     use_auth_token=UpperCamelCase__      )

				# set weights for wav2vec2 encoder
				__lowerCamelCase		       =    WavaVecaModel(UpperCamelCase__      )

				recursively_load_weights_wavaveca(model.encoder ,     UpperCamelCase__      )

				# load decoder weights
				__lowerCamelCase		       =    MBartForCausalLM(UpperCamelCase__      )
				__lowerCamelCase       ,						__lowerCamelCase		       =    hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict() ,     strict=UpperCamelCase__      )
				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}"""      )

				__lowerCamelCase		       =    SpeechEncoderDecoderModel(encoder=UpperCamelCase__ ,     decoder=UpperCamelCase__      )
				__lowerCamelCase		       =    False

				__lowerCamelCase		       =    MBartaaTokenizer(UpperCamelCase__      )
				tokenizer.save_pretrained(UpperCamelCase__      )

				__lowerCamelCase		       =    hf_wavavec.config.to_dict()
				__lowerCamelCase		       =    tokenizer.pad_token_id
				__lowerCamelCase		       =    tokenizer.bos_token_id
				__lowerCamelCase		       =    tokenizer.eos_token_id
				__lowerCamelCase		       =    'mbart50'
				__lowerCamelCase		       =    'wav2vec2'

				__lowerCamelCase		       =    tokenizer.eos_token_id
				__lowerCamelCase		       =    25_0004
				__lowerCamelCase		       =    tokenizer.eos_token_id

				__lowerCamelCase		       =    SpeechEncoderDecoderConfig.from_dict(UpperCamelCase__      )

				hf_wavavec.save_pretrained(UpperCamelCase__      )
				feature_extractor.save_pretrained(UpperCamelCase__      )


if __name__ == "__main__":
					__A        =    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=10_24, type=int, help="encoder output dim")
					parser.add_argument("--start_token_id", default=25_00_04, type=int, help="`decoder_start_token_id` of model config")

					__A        =    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,
					)


