


import gc
import unittest

import torch
from parameterized import parameterized

from diffusers import AutoencoderKL
from diffusers.utils import floats_tensor, load_hf_numpy, require_torch_gpu, slow, torch_all_close, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()







class        UpperCamelCase__   (	lowercase__							,				lowercase__							,				unittest.TestCase   ):
					__SCREAMING_SNAKE_CASE       :  Any						=   AutoencoderKL
					__SCREAMING_SNAKE_CASE       :  Any						=   """sample"""
					__SCREAMING_SNAKE_CASE       :  str						=   1e-2

					@property
					def   UpperCAmelCase__       (     self       :   int       ):






										'''simple docstring'''






										lowercase_  							=			4
										lowercase_  							=			3
										lowercase_  							=			(32, 32)

										lowercase_  							=			floats_tensor((batch_size, num_channels) + sizes       ).to(__lowercase       )

										return {"sample": image}

					@property
					def   UpperCAmelCase__       (     self       :   int       ):






										'''simple docstring'''






										return (3, 32, 32)

					@property
					def   UpperCAmelCase__       (     self       :   Optional[int]       ):






										'''simple docstring'''






										return (3, 32, 32)

					def   UpperCAmelCase__       (     self       :   int       ):






										'''simple docstring'''






										lowercase_  							=			{
										    """block_out_channels""": [32, 64],
										    """in_channels""": 3,
										    """out_channels""": 3,
										    """down_block_types""": ["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],
										    """up_block_types""": ["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],
										    """latent_channels""": 4,
										}
										lowercase_  							=			self.dummy_input
										return init_dict, inputs_dict

					def   UpperCAmelCase__       (     self       :   Union[str, Any]       ):






										'''simple docstring'''






										pass

					def   UpperCAmelCase__       (     self       :   Dict       ):






										'''simple docstring'''






										pass

					@unittest.skipIf(torch_device == """mps"""							,					"""Gradient checkpointing skipped on MPS"""       )
					def   UpperCAmelCase__       (     self       :   List[Any]       ):






										'''simple docstring'''






										lowercase_							,					lowercase_  							=			self.prepare_init_args_and_inputs_for_common()
										lowercase_  							=			self.model_class(**__lowercase       )
										model.to(__lowercase       )

										assert not model.is_gradient_checkpointing and model.training

										lowercase_  							=			model(**__lowercase       ).sample
										# run the backwards pass on the model. For backwards pass, for simplicity purpose,
										# we won't calculate the loss and rather backprop on out.sum()
										model.zero_grad()

										lowercase_  							=			torch.randn_like(__lowercase       )
										lowercase_  							=			(out - labels).mean()
										loss.backward()

										# re-instantiate the model now enabling gradient checkpointing
										lowercase_  							=			self.model_class(**__lowercase       )
										# clone model
										model_a.load_state_dict(model.state_dict()       )
										model_a.to(__lowercase       )
										model_a.enable_gradient_checkpointing()

										assert model_a.is_gradient_checkpointing and model_a.training

										lowercase_  							=			model_a(**__lowercase       ).sample
										# run the backwards pass on the model. For backwards pass, for simplicity purpose,
										# we won't calculate the loss and rather backprop on out.sum()
										model_a.zero_grad()
										lowercase_  							=			(out_a - labels).mean()
										loss_a.backward()

										# compare the output and parameters gradients
										self.assertTrue((loss - loss_a).abs() < 1e-5       )
										lowercase_  							=			dict(model.named_parameters()       )
										lowercase_  							=			dict(model_a.named_parameters()       )
										for name, param in named_params.items():
															self.assertTrue(torch_all_close(param.grad.data							,					named_params_a[name].grad.data							,					atol=5e-5       )       )

					def   UpperCAmelCase__       (     self       :   Union[str, Any]       ):






										'''simple docstring'''






										lowercase_							,					lowercase_  							=			AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""							,					output_loading_info=__lowercase       )
										self.assertIsNotNone(__lowercase       )
										self.assertEqual(len(loading_info["""missing_keys"""]       )							,					0       )

										model.to(__lowercase       )
										lowercase_  							=			model(**self.dummy_input       )

										assert image is not None, "Make sure output is not None"







					def   UpperCAmelCase__       (     self       :   Any       ):






										'''simple docstring'''






										lowercase_  							=			AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""       )
										lowercase_  							=			model.to(__lowercase       )
										model.eval()

										if torch_device == "mps":
															lowercase_  							=			torch.manual_seed(0       )
										else:
															lowercase_  							=			torch.Generator(device=__lowercase       ).manual_seed(0       )

										lowercase_  							=			torch.randn(
										    1							,					model.config.in_channels							,					model.config.sample_size							,					model.config.sample_size							,					generator=torch.manual_seed(0       )							,					)
										lowercase_  							=			image.to(__lowercase       )
										with torch.no_grad():
															lowercase_  							=			model(__lowercase							,					sample_posterior=__lowercase							,					generator=__lowercase       ).sample

										lowercase_  							=			output[0, -1, -3:, -3:].flatten().cpu()

										# Since the VAE Gaussian prior's generator is seeded on the appropriate device,
										# the expected output slices are not the same for CPU and GPU.
										if torch_device == "mps":
															lowercase_  							=			torch.tensor(
															    [
															        -4.0078e-01,
															        -3.8323e-04,
															        -1.2681e-01,
															        -1.1462e-01,
															        2.0095e-01,
															        1.0893e-01,
															        -8.8247e-02,
															        -3.0361e-01,
															        -9.8644e-03,
															    ]       )
										elif torch_device == "cpu":
															lowercase_  							=			torch.tensor(
															    [-0.1_352, 0.0_878, 0.0_419, -0.0_818, -0.1_069, 0.0_688, -0.1_458, -0.4_446, -0.0_026]       )
										else:
															lowercase_  							=			torch.tensor(
															    [-0.2_421, 0.4_642, 0.2_507, -0.0_438, 0.0_682, 0.3_160, -0.2_018, -0.0_727, 0.2_485]       )

										self.assertTrue(torch_all_close(__lowercase							,					__lowercase							,					rtol=1e-2       )       )







@slow
class        UpperCamelCase__   (	unittest.TestCase   ):

					def   UpperCAmelCase__       (     self       :   Dict							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   str       ):






										'''simple docstring'''






										return F'''gaussian_noise_s={seed}_shape={"_".join([str(__lowercase       ) for s in shape]       )}.npy'''

					def   UpperCAmelCase__       (     self       :   Optional[Any]       ):






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   Optional[Any]							,					UpperCamelCase__       :   Optional[Any]=0							,					UpperCamelCase__       :   List[str]=(4, 3, 512, 512)							,					UpperCamelCase__       :   List[Any]=False       ):






										'''simple docstring'''






										lowercase_  							=			torch.floataa if fpaa else torch.floataa
										lowercase_  							=			torch.from_numpy(load_hf_numpy(self.get_file_format(__lowercase							,					__lowercase       )       )       ).to(__lowercase       ).to(__lowercase       )
										return image

					def   UpperCAmelCase__       (     self       :   Optional[int]							,					UpperCamelCase__       :   Any="CompVis/stable-diffusion-v1-4"							,					UpperCamelCase__       :   List[Any]=False       ):






										'''simple docstring'''






										lowercase_  							=			"""fp16""" if fpaa else None
										lowercase_  							=			torch.floataa if fpaa else torch.floataa

										lowercase_  							=			AutoencoderKL.from_pretrained(
										    __lowercase							,					subfolder="""vae"""							,					torch_dtype=__lowercase							,					revision=__lowercase							,					)
										model.to(__lowercase       ).eval()

										return model

					def   UpperCAmelCase__       (     self       :   str							,					UpperCamelCase__       :   Optional[int]=0       ):






										'''simple docstring'''






										if torch_device == "mps":
															return torch.manual_seed(__lowercase       )
										return torch.Generator(device=__lowercase       ).manual_seed(__lowercase       )

					@parameterized.expand(
					    [
					        # fmt: off
					        [33, [-0.1_603, 0.9_878, -0.0_495, -0.0_790, -0.2_709, 0.8_375, -0.2_060, -0.0_824], [-0.2_395, 0.0_098, 0.0_102, -0.0_709, -0.2_840, -0.0_274, -0.0_718, -0.1_824]],
					        [47, [-0.2_376, 0.1_168, 0.1_332, -0.4_840, -0.2_508, -0.0_791, -0.0_493, -0.4_089], [0.0_350, 0.0_847, 0.0_467, 0.0_344, -0.0_842, -0.0_547, -0.0_633, -0.1_131]],
					        # fmt: on
					    ]       )
					def   UpperCAmelCase__       (     self       :   int							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   Union[str, Any]							,					UpperCamelCase__       :   List[Any]       ):






										'''simple docstring'''






										lowercase_  							=			self.get_sd_vae_model()
										lowercase_  							=			self.get_sd_image(__lowercase       )
										lowercase_  							=			self.get_generator(__lowercase       )

										with torch.no_grad():
															lowercase_  							=			model(__lowercase							,					generator=__lowercase							,					sample_posterior=__lowercase       ).sample

										assert sample.shape == image.shape

										lowercase_  							=			sample[-1, -2:, -2:, :2].flatten().float().cpu()
										lowercase_  							=			torch.tensor(expected_slice_mps if torch_device == """mps""" else expected_slice       )

										assert torch_all_close(__lowercase							,					__lowercase							,					atol=3e-3       )

					@parameterized.expand(
					    [
					        # fmt: off
					        [33, [-0.0_513, 0.0_289, 1.3_799, 0.2_166, -0.2_573, -0.0_871, 0.5_103, -0.0_999]],
					        [47, [-0.4_128, -0.1_320, -0.3_704, 0.1_965, -0.4_116, -0.2_332, -0.3_340, 0.2_247]],
					        # fmt: on
					    ]       )
					@require_torch_gpu
					def   UpperCAmelCase__       (     self       :   int							,					UpperCamelCase__       :   Optional[int]							,					UpperCamelCase__       :   Tuple       ):






										'''simple docstring'''






										lowercase_  							=			self.get_sd_vae_model(fpaa=__lowercase       )
										lowercase_  							=			self.get_sd_image(__lowercase							,					fpaa=__lowercase       )
										lowercase_  							=			self.get_generator(__lowercase       )

										with torch.no_grad():
															lowercase_  							=			model(__lowercase							,					generator=__lowercase							,					sample_posterior=__lowercase       ).sample

										assert sample.shape == image.shape

										lowercase_  							=			sample[-1, -2:, :2, -2:].flatten().float().cpu()
										lowercase_  							=			torch.tensor(__lowercase       )

										assert torch_all_close(__lowercase							,					__lowercase							,					atol=1e-2       )

					@parameterized.expand(
					    [
					        # fmt: off
					        [33, [-0.1_609, 0.9_866, -0.0_487, -0.0_777, -0.2_716, 0.8_368, -0.2_055, -0.0_814], [-0.2_395, 0.0_098, 0.0_102, -0.0_709, -0.2_840, -0.0_274, -0.0_718, -0.1_824]],
					        [47, [-0.2_377, 0.1_147, 0.1_333, -0.4_841, -0.2_506, -0.0_805, -0.0_491, -0.4_085], [0.0_350, 0.0_847, 0.0_467, 0.0_344, -0.0_842, -0.0_547, -0.0_633, -0.1_131]],
					        # fmt: on
					    ]       )
					def   UpperCAmelCase__       (     self       :   Any							,					UpperCamelCase__       :   Optional[Any]							,					UpperCamelCase__       :   Tuple							,					UpperCamelCase__       :   List[Any]       ):






										'''simple docstring'''






										lowercase_  							=			self.get_sd_vae_model()
										lowercase_  							=			self.get_sd_image(__lowercase       )

										with torch.no_grad():
															lowercase_  							=			model(__lowercase       ).sample

										assert sample.shape == image.shape

										lowercase_  							=			sample[-1, -2:, -2:, :2].flatten().float().cpu()
										lowercase_  							=			torch.tensor(expected_slice_mps if torch_device == """mps""" else expected_slice       )

										assert torch_all_close(__lowercase							,					__lowercase							,					atol=3e-3       )

					@parameterized.expand(
					    [
					        # fmt: off
					        [13, [-0.2_051, -0.1_803, -0.2_311, -0.2_114, -0.3_292, -0.3_574, -0.2_953, -0.3_323]],
					        [37, [-0.2_632, -0.2_625, -0.2_199, -0.2_741, -0.4_539, -0.4_990, -0.3_720, -0.4_925]],
					        # fmt: on
					    ]       )
					@require_torch_gpu
					def   UpperCAmelCase__       (     self       :   Optional[Any]							,					UpperCamelCase__       :   Union[str, Any]							,					UpperCamelCase__       :   Tuple       ):






										'''simple docstring'''






										lowercase_  							=			self.get_sd_vae_model()
										lowercase_  							=			self.get_sd_image(__lowercase							,					shape=(3, 4, 64, 64)       )

										with torch.no_grad():
															lowercase_  							=			model.decode(__lowercase       ).sample

										assert list(sample.shape       ) == [3, 3, 512, 512]

										lowercase_  							=			sample[-1, -2:, :2, -2:].flatten().cpu()
										lowercase_  							=			torch.tensor(__lowercase       )

										assert torch_all_close(__lowercase							,					__lowercase							,					atol=1e-3       )

					@parameterized.expand(
					    [
					        # fmt: off
					        [27, [-0.0_369, 0.0_207, -0.0_776, -0.0_682, -0.1_747, -0.1_930, -0.1_465, -0.2_039]],
					        [16, [-0.1_628, -0.2_134, -0.2_747, -0.2_642, -0.3_774, -0.4_404, -0.3_687, -0.4_277]],
					        # fmt: on
					    ]       )
					@require_torch_gpu
					def   UpperCAmelCase__       (     self       :   Optional[int]							,					UpperCamelCase__       :   List[Any]							,					UpperCamelCase__       :   Any       ):






										'''simple docstring'''






										lowercase_  							=			self.get_sd_vae_model(fpaa=__lowercase       )
										lowercase_  							=			self.get_sd_image(__lowercase							,					shape=(3, 4, 64, 64)							,					fpaa=__lowercase       )

										with torch.no_grad():
															lowercase_  							=			model.decode(__lowercase       ).sample

										assert list(sample.shape       ) == [3, 3, 512, 512]

										lowercase_  							=			sample[-1, -2:, :2, -2:].flatten().float().cpu()
										lowercase_  							=			torch.tensor(__lowercase       )

										assert torch_all_close(__lowercase							,					__lowercase							,					atol=5e-3       )

					@parameterized.expand([(13,), (16,), (27,)]       )
					@require_torch_gpu
					@unittest.skipIf(not is_xformers_available()							,					reason="""xformers is not required when using PyTorch 2.0."""       )
					def   UpperCAmelCase__       (     self       :   Tuple							,					UpperCamelCase__       :   Union[str, Any]       ):






										'''simple docstring'''






										lowercase_  							=			self.get_sd_vae_model(fpaa=__lowercase       )
										lowercase_  							=			self.get_sd_image(__lowercase							,					shape=(3, 4, 64, 64)							,					fpaa=__lowercase       )

										with torch.no_grad():
															lowercase_  							=			model.decode(__lowercase       ).sample

										model.enable_xformers_memory_efficient_attention()
										with torch.no_grad():
															lowercase_  							=			model.decode(__lowercase       ).sample

										assert list(sample.shape       ) == [3, 3, 512, 512]

										assert torch_all_close(__lowercase							,					__lowercase							,					atol=1e-1       )

					@parameterized.expand([(13,), (16,), (37,)]       )
					@require_torch_gpu
					@unittest.skipIf(not is_xformers_available()							,					reason="""xformers is not required when using PyTorch 2.0."""       )
					def   UpperCAmelCase__       (     self       :   int							,					UpperCamelCase__       :   Any       ):






										'''simple docstring'''






										lowercase_  							=			self.get_sd_vae_model()
										lowercase_  							=			self.get_sd_image(__lowercase							,					shape=(3, 4, 64, 64)       )

										with torch.no_grad():
															lowercase_  							=			model.decode(__lowercase       ).sample

										model.enable_xformers_memory_efficient_attention()
										with torch.no_grad():
															lowercase_  							=			model.decode(__lowercase       ).sample

										assert list(sample.shape       ) == [3, 3, 512, 512]

										assert torch_all_close(__lowercase							,					__lowercase							,					atol=1e-2       )







					@parameterized.expand(
					    [
					        # fmt: off
					        [33, [-0.3_001, 0.0_918, -2.6_984, -3.9_720, -3.2_099, -5.0_353, 1.7_338, -0.2_065, 3.4_267]],
					        [47, [-1.5_030, -4.3_871, -6.0_355, -9.1_157, -1.6_661, -2.7_853, 2.1_607, -5.0_823, 2.5_633]],
					        # fmt: on
					    ]       )
					def   UpperCAmelCase__       (     self       :   Optional[Any]							,					UpperCamelCase__       :   Union[str, Any]							,					UpperCamelCase__       :   Optional[Any]       ):






										'''simple docstring'''






										lowercase_  							=			self.get_sd_vae_model()
										lowercase_  							=			self.get_sd_image(__lowercase       )
										lowercase_  							=			self.get_generator(__lowercase       )

										with torch.no_grad():
															lowercase_  							=			model.encode(__lowercase       ).latent_dist
															lowercase_  							=			dist.sample(generator=__lowercase       )

										assert list(sample.shape       ) == [image.shape[0], 4] + [i // 8 for i in image.shape[2:]]

										lowercase_  							=			sample[0, -1, -3:, -3:].flatten().cpu()
										lowercase_  							=			torch.tensor(__lowercase       )

										lowercase_  							=			3e-3 if torch_device != """mps""" else 1e-2
										assert torch_all_close(__lowercase							,					__lowercase							,					atol=__lowercase       )






import argparse
import json
import os

import fairseq
import torch
from torch import nn

from transformers import (
    SpeechaTextaConfig,
    SpeechaTextaForCausalLM,
    SpeechaTextaTokenizer,
    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 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__  ):
					for attribute in key.split("""."""  ):
										lowercase_  							=			getattr(UpperCAmelCase__						,							UpperCAmelCase__  )

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

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


def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__  ):
					lowercase_  							=			[]
					lowercase_  							=			fairseq_model.state_dict()

					lowercase_  							=			hf_model.feature_extractor

					# if encoder has different dim to decoder -> use proj_weight
					lowercase_  							=			None

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

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

					return proj_weight


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


def 			UpperCAmelCase_				(      UpperCAmelCase__  ):
					lowercase_							,					lowercase_  							=			emb.weight.shape
					lowercase_  							=			nn.Linear(UpperCAmelCase__						,							UpperCAmelCase__						,							bias=UpperCAmelCase__  )
					lowercase_  							=			emb.weight.data
					return lin_layer


def 			UpperCAmelCase_				(      UpperCAmelCase__  ):
					with open(UpperCAmelCase__						,							"""r"""						,							encoding="""utf-8"""  ) as f:
										lowercase_  							=			f.readlines()
										lowercase_  							=			[line.split(""" """  )[0] for line in lines]

					lowercase_  							=			len(UpperCAmelCase__  )

					lowercase_  							=			{
					    """<s>""": 0,
					    """<pad>""": 1,
					    """</s>""": 2,
					    """<unk>""": 3,
					}

					vocab_dict.update(dict(zip(UpperCAmelCase__						,							range(4						,							num_words + 4  )  )  )  )
					return vocab_dict


@torch.no_grad()
def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__						,							):
					lowercase_  							=			WavaVecaConfig.from_pretrained(UpperCAmelCase__  )
					lowercase_  							=			SpeechaTextaConfig.from_pretrained(
					    UpperCAmelCase__						,							vocab_size=UpperCAmelCase__						,							decoder_layers=UpperCAmelCase__						,							do_stable_layer_norm=UpperCAmelCase__  )

					lowercase_  							=			WavaVecaFeatureExtractor(
					    feature_size=1						,							sampling_rate=1_6_0_0_0						,							padding_value=0						,							do_normalize=UpperCAmelCase__						,							return_attention_mask=UpperCAmelCase__						,							)

					lowercase_							,					lowercase_							,					lowercase_  							=			fairseq.checkpoint_utils.load_model_ensemble_and_task(
					    [checkpoint_path]						,							arg_overrides={"""data""": """/""".join(dict_path.split("""/"""  )[:-1]  )}  )
					lowercase_  							=			model[0].eval()

					# set weights for wav2vec2 encoder
					lowercase_  							=			WavaVecaModel(UpperCAmelCase__  )
					lowercase_  							=			recursively_load_weights_wavaveca(model.encoder						,							UpperCAmelCase__  )

					lowercase_  							=			SpeechaTextaForCausalLM(UpperCAmelCase__  )
					lowercase_							,					lowercase_  							=			hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()						,							strict=UpperCAmelCase__  )

					# set output linear layer
					unexpected_keys.remove("""embed_out"""  )
					lowercase_  							=			nn.Parameter(model.decoder.embed_out.detach()  )

					# layer norm is init to identity matrix so leaving it is fine
					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}'''  )

					lowercase_  							=			SpeechEncoderDecoderModel(encoder=UpperCAmelCase__						,							decoder=UpperCAmelCase__  )
					lowercase_  							=			False

					# add projection layer
					lowercase_  							=			nn.Parameter(projection_layer.weight  )
					lowercase_  							=			nn.Parameter(projection_layer.bias  )

					lowercase_  							=			create_vocab_dict(UpperCAmelCase__  )

					with open(os.path.join(UpperCAmelCase__						,							"""vocab.json"""  )						,							"""w"""  ) as fp:
										json.dump(UpperCAmelCase__						,							UpperCAmelCase__  )

					lowercase_  							=			SpeechaTextaTokenizer(os.path.join(UpperCAmelCase__						,							"""vocab.json"""  )  )
					tokenizer.save_pretrained(UpperCAmelCase__  )

					lowercase_  							=			hf_wavavec.config.to_dict()
					lowercase_  							=			tokenizer.pad_token_id
					lowercase_  							=			tokenizer.bos_token_id
					lowercase_  							=			tokenizer.eos_token_id
					lowercase_  							=			"""speech_to_text_2"""
					lowercase_  							=			"""wav2vec2"""

					lowercase_  							=			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(
							    '--encoder_config_path',
							    default='facebook/wav2vec2-large-lv60',
							    type=str,
							    help='Path to hf encoder wav2vec2 checkpoint config',
							)
							parser.add_argument(
							    '--decoder_config_path',
							    default='facebook/s2t-small-mustc-en-fr-st',
							    type=str,
							    help='Path to hf decoder s2t checkpoint config',
							)
							parser.add_argument('--vocab_size', default=1_0_2_2_4, type=int, help='Vocab size of decoder')
							parser.add_argument('--num_decoder_layers', default=7, type=int, help='Number of decoder layers')

							a					  =     parser.parse_args()
							convert_wavaveca_checkpoint(
							    args.checkpoint_path,
							    args.pytorch_dump_folder_path,
							    args.dict_path,
							    encoder_config_path=args.encoder_config_path,
							    decoder_config_path=args.decoder_config_path,
							    vocab_size=args.vocab_size,
							    num_decoder_layers=args.num_decoder_layers,
							)


