
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()
__lowerCAmelCase					:			Any			      =		logging.get_logger(__name__)

__lowerCAmelCase					:			Optional[int]			      =		{
    "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",
}
__lowerCAmelCase					:			str			      =		[
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
]


def     UpperCAmelCase_  (						__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	) ->						str:







    for attribute in key.split('''.'''	):
        __lowercase		:			int	    =       getattr(_lowercase	,					_lowercase	)

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

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


def     UpperCAmelCase_  (						__lowerCAmelCase	,					__lowerCAmelCase	) ->						List[str]:
    __lowercase		:			Any	    =       []
    __lowercase		:			List[str]	    =       fairseq_model.state_dict()

    __lowercase		:			Optional[Any]	    =       hf_model.feature_extractor

    # if encoder has different dim to decoder -> use proj_weight
    __lowercase		:			Optional[Any]	    =       None

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

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

    return proj_weight


def     UpperCAmelCase_  (						__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	) ->						Optional[int]:
    __lowercase		:			List[str]	    =       full_name.split('''conv_layers.'''	)[-1]
    __lowercase		:			Optional[Any]	    =       name.split('''.'''	)
    __lowercase		:			Tuple	    =       int(items[0]	)
    __lowercase		:			Any	    =       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		:			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.'
            )
            __lowercase		:			List[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."
            )
            __lowercase		:			Optional[Any]	    =       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		:			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_  (						__lowerCAmelCase	) ->						int:
    __lowercase    ,      __lowercase		:			Optional[int]	    =       emb.weight.shape
    __lowercase		:			Tuple	    =       nn.Linear(_lowercase	,					_lowercase	,					bias=_lowercase	)
    __lowercase		:			List[Any]	    =       emb.weight.data
    return lin_layer


def     UpperCAmelCase_  (						__lowerCAmelCase	) ->						Optional[Any]:







    with open(_lowercase	,					'''r'''	,					encoding='''utf-8'''	) as f:
        __lowercase		:			str	    =       f.readlines()
        __lowercase		:			Union[str, Any]	    =       [line.split(''' '''	)[0] for line in lines]

    __lowercase		:			int	    =       len(_lowercase	)

    __lowercase		:			Any	    =       {
        '''<s>''': 0,
        '''<pad>''': 1,
        '''</s>''': 2,
        '''<unk>''': 3,
    }

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


@torch.no_grad()
def     UpperCAmelCase_  (						__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	,					) ->						int:
    __lowercase		:			List[str]	    =       WavaVecaConfig.from_pretrained(_lowercase	)
    __lowercase		:			Dict	    =       SpeechaTextaConfig.from_pretrained(
        _lowercase	,					vocab_size=_lowercase	,					decoder_layers=_lowercase	,					do_stable_layer_norm=_lowercase	)

    __lowercase		:			Optional[Any]	    =       WavaVecaFeatureExtractor(
        feature_size=1	,					sampling_rate=16_000	,					padding_value=0	,					do_normalize=_lowercase	,					return_attention_mask=_lowercase	,					)

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

    # set weights for wav2vec2 encoder
    __lowercase		:			int	    =       WavaVecaModel(_lowercase	)
    __lowercase		:			Tuple	    =       recursively_load_weights_wavaveca(model.encoder	,					_lowercase	)

    __lowercase		:			Optional[Any]	    =       SpeechaTextaForCausalLM(_lowercase	)
    __lowercase    ,      __lowercase		:			List[str]	    =       hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()	,					strict=_lowercase	)

    # set output linear layer
    unexpected_keys.remove('''embed_out'''	)
    __lowercase		:			Dict	    =       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		:			List[Any]	    =       SpeechEncoderDecoderModel(encoder=_lowercase	,					decoder=_lowercase	)
    __lowercase		:			Optional[int]	    =       False

    # add projection layer
    __lowercase		:			str	    =       nn.Parameter(projection_layer.weight	)
    __lowercase		:			Optional[int]	    =       nn.Parameter(projection_layer.bias	)

    __lowercase		:			Tuple	    =       create_vocab_dict(_lowercase	)

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

    __lowercase		:			List[str]	    =       SpeechaTextaTokenizer(os.path.join(_lowercase	,					'''vocab.json'''	)	)
    tokenizer.save_pretrained(_lowercase	)

    __lowercase		:			Dict	    =       hf_wavavec.config.to_dict()
    __lowercase		:			Optional[Any]	    =       tokenizer.pad_token_id
    __lowercase		:			List[str]	    =       tokenizer.bos_token_id
    __lowercase		:			Optional[Any]	    =       tokenizer.eos_token_id
    __lowercase		:			Any	    =       '''speech_to_text_2'''
    __lowercase		:			List[Any]	    =       '''wav2vec2'''

    __lowercase		:			Any	    =       SpeechEncoderDecoderConfig.from_dict(_lowercase	)

    hf_wavavec.save_pretrained(_lowercase	)
    feature_extractor.save_pretrained(_lowercase	)


if __name__ == "__main__":
      __lowerCAmelCase					:			Union[str, Any]			      =		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=10_224, type=int, help="Vocab size of decoder")
      parser.add_argument("--num_decoder_layers", default=7, type=int, help="Number of decoder layers")

      __lowerCAmelCase					:			Dict			      =		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,
      )




import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class 						__lowerCAmelCase	(	lowerCAmelCase_ , lowerCAmelCase_ , unittest.TestCase      ):



      """simple docstring"""







      A__	:			int        =    StableDiffusionSAGPipeline
      A__	:			Union[str, Any]        =    TEXT_TO_IMAGE_PARAMS
      A__	:			Dict        =    TEXT_TO_IMAGE_BATCH_PARAMS
      A__	:			List[str]        =    TEXT_TO_IMAGE_IMAGE_PARAMS
      A__	:			int        =    TEXT_TO_IMAGE_IMAGE_PARAMS
      A__	:			List[str]        =    False




      def 		snake_case_			(      self		:       Optional[Any]				):
          torch.manual_seed(0				)
          __lowercase		:			Optional[Any]	    =       UNetaDConditionModel(
              block_out_channels=(32, 64)				,       layers_per_block=2				,       sample_size=32				,       in_channels=4				,       out_channels=4				,       down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')				,       up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')				,       cross_attention_dim=32				,       )
          __lowercase		:			int	    =       DDIMScheduler(
              beta_start=0.0_00_85				,       beta_end=0.0_12				,       beta_schedule='''scaled_linear'''				,       clip_sample=_snake_case				,       set_alpha_to_one=_snake_case				,       )
          torch.manual_seed(0				)
          __lowercase		:			List[Any]	    =       AutoencoderKL(
              block_out_channels=[32, 64]				,       in_channels=3				,       out_channels=3				,       down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']				,       up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']				,       latent_channels=4				,       )
          torch.manual_seed(0				)
          __lowercase		:			Tuple	    =       CLIPTextConfig(
              bos_token_id=0				,       eos_token_id=2				,       hidden_size=32				,       intermediate_size=37				,       layer_norm_eps=1E-0_5				,       num_attention_heads=4				,       num_hidden_layers=5				,       pad_token_id=1				,       vocab_size=1000				,       )
          __lowercase		:			List[Any]	    =       CLIPTextModel(_snake_case				)
          __lowercase		:			List[Any]	    =       CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				)

          __lowercase		:			int	    =       {
              '''unet''': unet,
              '''scheduler''': scheduler,
              '''vae''': vae,
              '''text_encoder''': text_encoder,
              '''tokenizer''': tokenizer,
              '''safety_checker''': None,
              '''feature_extractor''': None,
          }
          return components




      def 		snake_case_			(      self		:       Tuple				,       _snake_case		:       List[str]				,       _snake_case		:       List[str]=0				):
          if str(_snake_case				).startswith('''mps'''				):
              __lowercase		:			Optional[int]	    =       torch.manual_seed(_snake_case				)
          else:
              __lowercase		:			Tuple	    =       torch.Generator(device=_snake_case				).manual_seed(_snake_case				)
          __lowercase		:			Tuple	    =       {
              '''prompt''': '''.''',
              '''generator''': generator,
              '''num_inference_steps''': 2,
              '''guidance_scale''': 1.0,
              '''sag_scale''': 1.0,
              '''output_type''': '''numpy''',
          }
          return inputs




      def 		snake_case_			(      self		:       Any				):
          super().test_inference_batch_single_identical(expected_max_diff=3E-3				)






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



      """simple docstring"""







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




      def 		snake_case_			(      self		:       Dict				):
          __lowercase		:			Optional[int]	    =       StableDiffusionSAGPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''				)
          __lowercase		:			Optional[int]	    =       sag_pipe.to(_snake_case				)
          sag_pipe.set_progress_bar_config(disable=_snake_case				)

          __lowercase		:			List[str]	    =       '''.'''
          __lowercase		:			List[str]	    =       torch.manual_seed(0				)
          __lowercase		:			List[str]	    =       sag_pipe(
              [prompt]				,       generator=_snake_case				,       guidance_scale=7.5				,       sag_scale=1.0				,       num_inference_steps=20				,       output_type='''np'''				)

          __lowercase		:			Optional[Any]	    =       output.images

          __lowercase		:			Dict	    =       image[0, -3:, -3:, -1]

          assert image.shape == (1, 512, 512, 3)
          __lowercase		:			int	    =       np.array([0.15_68, 0.17_38, 0.16_95, 0.16_93, 0.15_07, 0.17_05, 0.15_47, 0.17_51, 0.19_49]				)

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




      def 		snake_case_			(      self		:       Optional[int]				):
          __lowercase		:			Dict	    =       StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''				)
          __lowercase		:			Any	    =       sag_pipe.to(_snake_case				)
          sag_pipe.set_progress_bar_config(disable=_snake_case				)

          __lowercase		:			List[str]	    =       '''.'''
          __lowercase		:			List[str]	    =       torch.manual_seed(0				)
          __lowercase		:			int	    =       sag_pipe(
              [prompt]				,       generator=_snake_case				,       guidance_scale=7.5				,       sag_scale=1.0				,       num_inference_steps=20				,       output_type='''np'''				)

          __lowercase		:			str	    =       output.images

          __lowercase		:			Tuple	    =       image[0, -3:, -3:, -1]

          assert image.shape == (1, 512, 512, 3)
          __lowercase		:			Any	    =       np.array([0.34_59, 0.28_76, 0.25_37, 0.30_02, 0.26_71, 0.21_60, 0.30_26, 0.22_62, 0.23_71]				)

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




      def 		snake_case_			(      self		:       int				):
          __lowercase		:			List[str]	    =       StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''				)
          __lowercase		:			Tuple	    =       sag_pipe.to(_snake_case				)
          sag_pipe.set_progress_bar_config(disable=_snake_case				)

          __lowercase		:			Union[str, Any]	    =       '''.'''
          __lowercase		:			Tuple	    =       torch.manual_seed(0				)
          __lowercase		:			Optional[int]	    =       sag_pipe(
              [prompt]				,       width=768				,       height=512				,       generator=_snake_case				,       guidance_scale=7.5				,       sag_scale=1.0				,       num_inference_steps=20				,       output_type='''np'''				,       )

          __lowercase		:			Dict	    =       output.images

          assert image.shape == (1, 512, 768, 3)


