

"""simple docstring"""






import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    UniSpeechConfig,
    UniSpeechForCTC,
    UniSpeechForPreTraining,
    WavaVecaFeatureExtractor,
    WavaVecaPhonemeCTCTokenizer,
    WavaVecaProcessor,
    logging,
)


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

_lowerCAmelCase						=		{
    """post_extract_proj""": """feature_projection.projection""",
    """encoder.pos_conv.0""": """encoder.pos_conv_embed.conv""",
    """self_attn.k_proj""": """encoder.layers.*.attention.k_proj""",
    """self_attn.v_proj""": """encoder.layers.*.attention.v_proj""",
    """self_attn.q_proj""": """encoder.layers.*.attention.q_proj""",
    """self_attn.out_proj""": """encoder.layers.*.attention.out_proj""",
    """self_attn_layer_norm""": """encoder.layers.*.layer_norm""",
    """fc1""": """encoder.layers.*.feed_forward.intermediate_dense""",
    """fc2""": """encoder.layers.*.feed_forward.output_dense""",
    """final_layer_norm""": """encoder.layers.*.final_layer_norm""",
    """encoder.layer_norm""": """encoder.layer_norm""",
    """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""": """ctc_proj""",
    """mask_emb""": """masked_spec_embed""",
}
_lowerCAmelCase						=		[
    """ctc_proj""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
]

def 		lowerCamelCase__		(					_lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase   ):

    '''simple docstring'''



    for attribute in key.split('.'   ):
        if is_finetuned:
            if attribute in ["quantizer", "project_q", "project_hid"]:
                # those layers are only relevant for pretraining and should be dropped
                return

            if attribute == "ctc_proj":
                # we should rename `ctc_proj` to `lm_head` for fine-tuned phoneme models
                _lowerCAmelCase       :   str  =    'lm_head'

        _lowerCAmelCase       :   Optional[Any]  =    getattr(lowerCamelCase_ ,   lowerCamelCase_   )

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

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

def 		lowerCamelCase__		(					_lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase   ):

    '''simple docstring'''



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

    _lowerCAmelCase       :   int  =    hf_model.unispeech.feature_extractor

    for name, value in fairseq_dict.items():
        _lowerCAmelCase       :   Union[str, Any]  =    False
        if "conv_layers" in name:
            load_conv_layer(
                lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_ ,   hf_model.config.feat_extract_norm == 'group' ,   )
            _lowerCAmelCase       :   List[Any]  =    True
        else:
            for key, mapped_key in MAPPING.items():
                _lowerCAmelCase       :   List[Any]  =    'unispeech.' + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
                if key in name or key.split('w2v_model.'   )[-1] == name.split('.'   )[0]:
                    _lowerCAmelCase       :   Any  =    True
                    if "*" in mapped_key:
                        _lowerCAmelCase       :   Any  =    name.split(lowerCamelCase_   )[0].split('.'   )[-2]
                        _lowerCAmelCase       :   Union[str, Any]  =    mapped_key.replace('*' ,   lowerCamelCase_   )
                    if "weight_g" in name:
                        _lowerCAmelCase       :   int  =    'weight_g'
                    elif "weight_v" in name:
                        _lowerCAmelCase       :   Any  =    'weight_v'
                    elif "bias" in name:
                        _lowerCAmelCase       :   Union[str, Any]  =    'bias'
                    elif "weight" in name:
                        # TODO: don't match quantizer.weight_proj
                        _lowerCAmelCase       :   Any  =    'weight'
                    else:
                        _lowerCAmelCase       :   Tuple  =    None
                    set_recursively(lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_   )
                continue
        if not is_used:
            unused_weights.append(lowerCamelCase_   )

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

def 		lowerCamelCase__		(					_lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase   ):

    '''simple docstring'''



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

@torch.no_grad()
def 		lowerCamelCase__		(					_lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase=None ,   _lowerCamelCase=None ,   _lowerCamelCase=True   ):

    '''simple docstring'''



    if config_path is not None:
        _lowerCAmelCase       :   Optional[Any]  =    UniSpeechConfig.from_pretrained(lowerCamelCase_   )
    else:
        _lowerCAmelCase       :   int  =    UniSpeechConfig()

    if is_finetuned:
        if dict_path:
            _lowerCAmelCase       :   Union[str, Any]  =    Dictionary.load_from_json(lowerCamelCase_   )

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

            # fairseq has the <pad> and <s> switched
            _lowerCAmelCase       :   Any  =    42
            _lowerCAmelCase       :   List[str]  =    43
            with open(lowerCamelCase_ ,   'w' ,   encoding='utf-8'   ) as vocab_handle:
                json.dump(lowerCamelCase_ ,   lowerCamelCase_   )
            _lowerCAmelCase       :   Optional[int]  =    WavaVecaPhonemeCTCTokenizer(
                lowerCamelCase_ ,   unk_token=target_dict.unk_word ,   pad_token=target_dict.pad_word ,   bos_token=target_dict.bos_word ,   eos_token=target_dict.eos_word ,   word_delimiter_token='|' ,   do_lower_case=lowerCamelCase_ ,   )
            _lowerCAmelCase       :   Optional[Any]  =    True if config.feat_extract_norm == 'layer' else False
            _lowerCAmelCase       :   Union[str, Any]  =    WavaVecaFeatureExtractor(
                feature_size=1 ,   sampling_rate=16000 ,   padding_value=0 ,   do_normalize=lowerCamelCase_ ,   return_attention_mask=lowerCamelCase_ ,   )
            _lowerCAmelCase       :   Tuple  =    WavaVecaProcessor(feature_extractor=lowerCamelCase_ ,   tokenizer=lowerCamelCase_   )
            processor.save_pretrained(lowerCamelCase_   )

        _lowerCAmelCase       :   Dict  =    UniSpeechForCTC(lowerCamelCase_   )
    else:
        _lowerCAmelCase       :   List[Any]  =    UniSpeechForPreTraining(lowerCamelCase_   )

    if is_finetuned:
        _lowerCAmelCase       :   int  =    fairseq.checkpoint_utils.load_model_ensemble_and_task(
            [checkpoint_path] ,   arg_overrides={'data': '/'.join(dict_path.split('/'   )[:-1]   ), 'w2v_path': checkpoint_path}   )
    else:
        _lowerCAmelCase       :   str  =    fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]   )

    _lowerCAmelCase       :   int  =    model[0].eval()

    recursively_load_weights(lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_   )

    hf_unispeech.save_pretrained(lowerCamelCase_   )


if __name__ == "__main__":
 _lowerCAmelCase						=		argparse.ArgumentParser()
 parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
 parser.add_argument("""--checkpoint_path""", default=None, type=str, help="""Path to fairseq checkpoint""")
 parser.add_argument("""--dict_path""", default=None, type=str, help="""Path to dict of fine-tuned model""")
 parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
 parser.add_argument(
     """--not_finetuned""", action="""store_true""", help="""Whether the model to convert is a fine-tuned model or not"""
 )
 _lowerCAmelCase						=		parser.parse_args()
 convert_unispeech_checkpoint(
     args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
 )








"""simple docstring"""






import gc
import unittest

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

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

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


enable_full_determinism()







class 	__UpperCamelCase    (    a__							,							unittest.TestCase						):
  _UpperCAmelCase          =		LDMTextToImagePipeline
  _UpperCAmelCase          =		TEXT_TO_IMAGE_PARAMS - {
      "negative_prompt",
      "negative_prompt_embeds",
      "cross_attention_kwargs",
      "prompt_embeds",
  }
  _UpperCAmelCase          =		PipelineTesterMixin.required_optional_params - {
      "num_images_per_prompt",
      "callback",
      "callback_steps",
  }
  _UpperCAmelCase          =		TEXT_TO_IMAGE_BATCH_PARAMS
  _UpperCAmelCase          =		False

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      torch.manual_seed(0      )
      _lowerCAmelCase       :   Optional[int]  =    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    ,)
      _lowerCAmelCase       :   Union[str, Any]  =    DDIMScheduler(
          beta_start=0.0_0_0_8_5    ,beta_end=0.0_1_2    ,beta_schedule='scaled_linear'    ,clip_sample=_A    ,set_alpha_to_one=_A    ,)
      torch.manual_seed(0      )
      _lowerCAmelCase       :   Union[str, 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      )
      _lowerCAmelCase       :   Dict  =    CLIPTextConfig(
          bos_token_id=0    ,eos_token_id=2    ,hidden_size=32    ,intermediate_size=37    ,layer_norm_eps=1E-05    ,num_attention_heads=4    ,num_hidden_layers=5    ,pad_token_id=1    ,vocab_size=1000    ,)
      _lowerCAmelCase       :   Tuple  =    CLIPTextModel(_A      )
      _lowerCAmelCase       :   Any  =    CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'      )

      _lowerCAmelCase       :   List[str]  =    {
          'unet': unet,
          'scheduler': scheduler,
          'vqvae': vae,
          'bert': text_encoder,
          'tokenizer': tokenizer,
      }
      return components

  def __lowerCamelCase		(		self    ,_A    ,_A=0      ):
      '''simple docstring'''

      if str(_A      ).startswith('mps'      ):
          _lowerCAmelCase       :   int  =    torch.manual_seed(_A      )
      else:
          _lowerCAmelCase       :   Optional[Any]  =    torch.Generator(device=_A      ).manual_seed(_A      )
      _lowerCAmelCase       :   List[Any]  =    {
          'prompt': 'A painting of a squirrel eating a burger',
          'generator': generator,
          'num_inference_steps': 2,
          'guidance_scale': 6.0,
          'output_type': 'numpy',
      }
      return inputs







  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[int]  =    'cpu'  # ensure determinism for the device-dependent torch.Generator

      _lowerCAmelCase       :   int  =    self.get_dummy_components()
      _lowerCAmelCase       :   str  =    LDMTextToImagePipeline(**_A      )
      pipe.to(_A      )
      pipe.set_progress_bar_config(disable=_A      )

      _lowerCAmelCase       :   List[Any]  =    self.get_dummy_inputs(_A      )
      _lowerCAmelCase       :   Any  =    pipe(**_A      ).images
      _lowerCAmelCase       :   Optional[int]  =    image[0, -3:, -3:, -1]

      assert image.shape == (1, 16, 16, 3)
      _lowerCAmelCase       :   Tuple  =    np.array([0.6_1_0_1, 0.6_1_5_6, 0.5_6_2_2, 0.4_8_9_5, 0.6_6_6_1, 0.3_8_0_4, 0.5_7_4_8, 0.6_1_3_6, 0.5_0_1_4]      )

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







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

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

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

  def __lowerCamelCase		(		self    ,_A    ,_A=torch.floataa    ,_A=0      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[Any]  =    torch.manual_seed(_A      )
      _lowerCAmelCase       :   Union[str, Any]  =    np.random.RandomState(_A      ).standard_normal((1, 4, 32, 32)      )
      _lowerCAmelCase       :   Optional[Any]  =    torch.from_numpy(_A      ).to(device=_A    ,dtype=_A      )
      _lowerCAmelCase       :   List[str]  =    {
          'prompt': 'A painting of a squirrel eating a burger',
          'latents': latents,
          'generator': generator,
          'num_inference_steps': 3,
          'guidance_scale': 6.0,
          'output_type': 'numpy',
      }
      return inputs







  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    LDMTextToImagePipeline.from_pretrained('CompVis/ldm-text2im-large-256'      ).to(_A      )
      pipe.set_progress_bar_config(disable=_A      )

      _lowerCAmelCase       :   Optional[Any]  =    self.get_inputs(_A      )
      _lowerCAmelCase       :   List[Any]  =    pipe(**_A      ).images
      _lowerCAmelCase       :   Dict  =    image[0, -3:, -3:, -1].flatten()

      assert image.shape == (1, 256, 256, 3)
      _lowerCAmelCase       :   str  =    np.array([0.5_1_8_2_5, 0.5_2_8_5_0, 0.5_2_5_4_3, 0.5_4_2_5_8, 0.5_2_3_0_4, 0.5_2_5_6_9, 0.5_4_3_6_3, 0.5_5_2_7_6, 0.5_6_8_7_8]      )
      _lowerCAmelCase       :   Dict  =    np.abs(expected_slice - image_slice      ).max()
      assert max_diff < 1E-3







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

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

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

  def __lowerCamelCase		(		self    ,_A    ,_A=torch.floataa    ,_A=0      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    torch.manual_seed(_A      )
      _lowerCAmelCase       :   Optional[int]  =    np.random.RandomState(_A      ).standard_normal((1, 4, 32, 32)      )
      _lowerCAmelCase       :   List[Any]  =    torch.from_numpy(_A      ).to(device=_A    ,dtype=_A      )
      _lowerCAmelCase       :   int  =    {
          'prompt': 'A painting of a squirrel eating a burger',
          'latents': latents,
          'generator': generator,
          'num_inference_steps': 50,
          'guidance_scale': 6.0,
          'output_type': 'numpy',
      }
      return inputs







  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    LDMTextToImagePipeline.from_pretrained('CompVis/ldm-text2im-large-256'      ).to(_A      )
      pipe.set_progress_bar_config(disable=_A      )

      _lowerCAmelCase       :   str  =    self.get_inputs(_A      )
      _lowerCAmelCase       :   Union[str, Any]  =    pipe(**_A      ).images[0]

      _lowerCAmelCase       :   int  =    load_numpy(
          'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/ldm_text2img/ldm_large_256_ddim.npy'      )
      _lowerCAmelCase       :   List[str]  =    np.abs(expected_image - image      ).max()
      assert max_diff < 1E-3





